core: use previous size if heartbeat's size is zero (#4243)

* core: use previous size if heartbeat's size is zero

ref tikv/tikv#11114

Signed-off-by: p4tr1ck <patrick.li@pingcap.com>

* add issue link (#4258)

Signed-off-by: p4tr1ck <patrick.li@pingcap.com>

Co-authored-by: Ti Chi Robot <ti-community-prow-bot@tidb.io>

server/cluster/cluster.go

@@ -630,6 +630,8 @@ func (c *RaftCluster) processRegionHeartbeat(region *core.RegionInfo) error {
 	if err != nil {
 		return err
 	}
+	region.CorrectApproximateSize(origin)
+
 	hotStat.CheckWriteAsync(statistics.NewCheckExpiredItemTask(region))
 	hotStat.CheckReadAsync(statistics.NewCheckExpiredItemTask(region))
 	reportInterval := region.GetInterval()

server/core/region.go

@@ -107,9 +107,10 @@ const (
 // RegionFromHeartbeat constructs a Region from region heartbeat.
 func RegionFromHeartbeat(heartbeat *pdpb.RegionHeartbeatRequest, opts ...RegionCreateOption) *RegionInfo {
 	// Convert unit to MB.
-	// If region is empty or less than 1MB, use 1MB instead.
+	// If region isn't empty and less than 1MB, use 1MB instead.
+	// The size of empty region will be correct by the previous RegionInfo.
 	regionSize := heartbeat.GetApproximateSize() / (1 << 20)
-	if regionSize < EmptyRegionApproximateSize {
+	if heartbeat.GetApproximateSize() > 0 && regionSize < EmptyRegionApproximateSize {
 		regionSize = EmptyRegionApproximateSize
 	}
 
@@ -150,6 +151,21 @@ func RegionFromHeartbeat(heartbeat *pdpb.RegionHeartbeatRequest, opts ...RegionC
 	return region
 }
 
+// CorrectApproximateSize correct approximate size by the previous size if here exists an reported RegionInfo.
+//
+// See https://github.com/tikv/tikv/issues/11114
+func (r *RegionInfo) CorrectApproximateSize(origin *RegionInfo) {
+	if r.approximateSize != 0 {
+		return
+	}
+
+	if origin != nil {
+		r.approximateSize = origin.approximateSize
+	} else {
+		r.approximateSize = EmptyRegionApproximateSize
+	}
+}
+
 // Clone returns a copy of current regionInfo.
 func (r *RegionInfo) Clone(opts ...RegionCreateOption) *RegionInfo {
 	downPeers := make([]*pdpb.PeerStats, 0, len(r.downPeers))

server/core/region_test.go

@@ -157,6 +157,33 @@ func (s *testRegionInfoSuite) TestSortedEqual(c *C) {
 	}
 }
 
+func (s *testRegionInfoSuite) TestCorrectRegionApproximateSize(c *C) {
+	// size in MB
+	testcases := []struct {
+		originExists bool
+		originSize   uint64
+		size         uint64
+		expect       uint64
+	}{
+		{false, 0, 0, 1},
+		{false, 0, 2, 2},
+		{true, 0, 2, 2},
+		{true, 1, 2, 2},
+		{true, 2, 0, 2},
+	}
+	for _, t := range testcases {
+		var origin *RegionInfo
+		if t.originExists {
+			origin = NewRegionInfo(&metapb.Region{Id: 100}, nil)
+			origin.approximateSize = int64(t.originSize)
+		}
+		r := NewRegionInfo(&metapb.Region{Id: 100}, nil)
+		r.approximateSize = int64(t.size)
+		r.CorrectApproximateSize(origin)
+		c.Assert(r.approximateSize, Equals, int64(t.expect))
+	}
+}
+
 func (s *testRegionInfoSuite) TestRegionRoundingFlow(c *C) {
 	testcases := []struct {
 		flow   uint64