From fc7f774a65fb506ddbfbaec88feee97fa2170efc Mon Sep 17 00:00:00 2001
From: p4tr1ck <patrick.li@pingcap.com>
Date: Thu, 28 Oct 2021 11:02:12 +0800
Subject: [PATCH] core: use previous size if heartbeat's size is zero

ref tikv/tikv#11114

Signed-off-by: p4tr1ck <patrick.li@pingcap.com>
---
 server/cluster/cluster.go  |  2 ++
 server/core/region.go      | 25 +++++++++++++++++--------
 server/core/region_test.go | 29 +++++++++++++++++++++++++++++
 3 files changed, 48 insertions(+), 8 deletions(-)

diff --git a/server/cluster/cluster.go b/server/cluster/cluster.go
index c2e62236d692..2d8d42b2e5ee 100644
--- a/server/cluster/cluster.go
+++ b/server/cluster/cluster.go
@@ -622,6 +622,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()
diff --git a/server/core/region.go b/server/core/region.go
index febdf40dd628..df21547a5fba 100644
--- a/server/core/region.go
+++ b/server/core/region.go
@@ -106,13 +106,6 @@ 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.
-	regionSize := heartbeat.GetApproximateSize() / (1 << 20)
-	if regionSize < EmptyRegionApproximateSize {
-		regionSize = EmptyRegionApproximateSize
-	}
-
 	region := &RegionInfo{
 		term:              heartbeat.GetTerm(),
 		meta:              heartbeat.GetRegion(),
@@ -123,7 +116,7 @@ func RegionFromHeartbeat(heartbeat *pdpb.RegionHeartbeatRequest, opts ...RegionC
 		writtenKeys:       heartbeat.GetKeysWritten(),
 		readBytes:         heartbeat.GetBytesRead(),
 		readKeys:          heartbeat.GetKeysRead(),
-		approximateSize:   int64(regionSize),
+		approximateSize:   int64(heartbeat.GetApproximateSize()),
 		approximateKeys:   int64(heartbeat.GetApproximateKeys()),
 		interval:          heartbeat.GetInterval(),
 		replicationStatus: heartbeat.GetReplicationStatus(),
@@ -150,6 +143,22 @@ func RegionFromHeartbeat(heartbeat *pdpb.RegionHeartbeatRequest, opts ...RegionC
 	return region
 }
 
+// AdjustApproximateSize convert approximate size unit to MB, and
+// - if region is empty, and there exists an reported RegionInfo, use the previous size;
+// - otherwise if region size less than 1MB, use 1MB instead.
+//
+// See https://github.com/tikv/tikv/issues/11114
+func (r *RegionInfo) CorrectApproximateSize(origin *RegionInfo) {
+	if r.approximateSize == 0 && origin != nil {
+		r.approximateSize = origin.approximateSize
+	} else {
+		r.approximateSize = r.approximateSize / (1 << 20)
+		if r.approximateSize < EmptyRegionApproximateSize {
+			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))
diff --git a/server/core/region_test.go b/server/core/region_test.go
index af7ea2dfb5da..71d7e5123e37 100644
--- a/server/core/region_test.go
+++ b/server/core/region_test.go
@@ -157,6 +157,35 @@ 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 << 20)
+		fmt.Printf("before %d\n", r.approximateSize)
+		r.CorrectApproximateSize(origin)
+		fmt.Printf("after %d\n", r.approximateSize)
+		c.Assert(r.approximateSize, Equals, int64(t.expect))
+	}
+}
+
 func (s *testRegionInfoSuite) TestRegionRoundingFlow(c *C) {
 	testcases := []struct {
 		flow   uint64