Repository: tikv/pd (go)

Issue:
region is not balanced during br restore
## Bug Report

<!-- Please answer these questions before submitting your issue. Thanks! -->

### What did you do?
1、command="export AWS_ACCESS_KEY_ID=minioadmin;export AWS_SECRET_ACCESS_KEY=minioadmin;/br restore db  --send-credentials-to-tikv=true --db tpcc20000 --pd http://tc-pd.endless-ha-test-with-some-feature-tps-2760236-1-101:2379 --storage s3://benchmark/tpcc20000 --s3.endpoint  http://minio.pingcap.net:9000 --check-requirements=false

### What did you expect to see?
region is balanced during br restore

### What did you see instead?
region is not balanced during br restore
![697e5e61-9453-4c4b-93b1-176e7eec52e5](https://github.com/tikv/pd/assets/84712107/2456acc5-c7bb-4260-a548-2210e7a9c7b0)

### What version of PD are you using (`pd-server -V`)?
./pd-server -V
 Release Version: v7.4.0-alpha
Edition: Community
Git Commit Hash: a21fd58d9953eeb8a4ff2657c04de2f44ee29e0e
Git Branch: heads/refs/tags/v7.4.0-alpha
UTC Build Time:  2023-09-22 11:36:18
2023-09-23T01:25:02.465+0800


Pull request:
Revert "statistics: fix empty region count when resuming (#7009)"
This reverts commit 74ead5cbdf17df955ad7075810be3a724033ca0f.

<!--

Thank you for working on PD! Please read PD's [CONTRIBUTING](https://github.com/tikv/pd/blob/master/CONTRIBUTING.md) document **BEFORE** filing this PR.

PR Title Format:
1. pkg [, pkg2, pkg3]: what's changed
2. *: what's changed

-->

### What problem does this PR solve?

<!--

Please create an issue first to describe the problem.
There MUST be one line starting with "Issue Number:  " and 
linking the relevant issues via the "close" or "ref".
For more info, check https://github.com/tikv/pd/blob/master/CONTRIBUTING.md#linking-issues.

-->
Issue Number: Close #7148

### What is changed and how does it work?

<!--

You could use the "commit message" code block to add more description to the final commit message.
For more info, check https://github.com/tikv/pd/blob/master/CONTRIBUTING.md#format-of-the-commit-message.

-->

```commit-message
```

### Check List

<!-- Remove the items that are not applicable. -->

Tests

<!-- At least one of these tests must be included. -->

- Unit test
- Integration test
- Manual test (add detailed scripts or steps below)
- No code

Code changes

- Has the configuration change
- Has HTTP API interfaces changed (Don't forget to [add the declarative for the new API](https://github.com/tikv/pd/blob/master/docs/development.md#updating-api-documentation))
- Has persistent data change

Side effects

- Possible performance regression
- Increased code complexity
- Breaking backward compatibility

Related changes

- PR to update [`pingcap/docs`](https://github.com/pingcap/docs)/[`pingcap/docs-cn`](https://github.com/pingcap/docs-cn):
- PR to update [`pingcap/tiup`](https://github.com/pingcap/tiup):
- Need to cherry-pick to the release branch

### Release note

<!--

A bugfix or a new feature needs a release note. If there is no need to give a release note, just leave it with the `None`.

Please refer to [Release Notes Language Style Guide](https://pingcap.github.io/tidb-dev-guide/contribute-to-tidb/release-notes-style-guide.html) to write a quality release note.

-->

```release-note
None.
```


Patch:
diff --git a/pkg/core/region.go b/pkg/core/region.go
index 4540f7aaf..9e32bf7c2 100644
--- a/pkg/core/region.go
+++ b/pkg/core/region.go
@@ -170,9 +170,8 @@ type RegionHeartbeatRequest interface {
 func RegionFromHeartbeat(heartbeat RegionHeartbeatRequest, opts ...RegionCreateOption) *RegionInfo {
 	// Convert unit to MB.
 	// 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() / units.MiB
-	// Due to https://github.com/tikv/tikv/pull/11170, if region size is not initialized,
-	// approximate size will be zero, and region size is zero not EmptyRegionApproximateSize
 	if heartbeat.GetApproximateSize() > 0 && regionSize < EmptyRegionApproximateSize {
 		regionSize = EmptyRegionApproximateSize
 	}
@@ -220,9 +219,19 @@ func RegionFromHeartbeat(heartbeat RegionHeartbeatRequest, opts ...RegionCreateO
 	return region
 }
 
-// InheritBuckets inherits the buckets from the parent region if bucket enabled.
-func (r *RegionInfo) InheritBuckets(origin *RegionInfo) {
-	if origin != nil && r.buckets == nil {
+// Inherit inherits the buckets and region size from the parent region if bucket enabled.
+// correct approximate size and buckets by the previous size if here exists a reported RegionInfo.
+// See https://github.com/tikv/tikv/issues/11114
+func (r *RegionInfo) Inherit(origin *RegionInfo, bucketEnable bool) {
+	// regionSize should not be zero if region is not empty.
+	if r.GetApproximateSize() == 0 {
+		if origin != nil {
+			r.approximateSize = origin.approximateSize
+		} else {
+			r.approximateSize = EmptyRegionApproximateSize
+		}
+	}
+	if bucketEnable && origin != nil && r.buckets == nil {
 		r.buckets = origin.buckets
 	}
 }
@@ -532,13 +541,6 @@ func (r *RegionInfo) GetApproximateSize() int64 {
 	return r.approximateSize
 }
 
-// IsEmptyRegion returns whether the region is empty.
-func (r *RegionInfo) IsEmptyRegion() bool {
-	// When cluster resumes, the region size may be not initialized, but region heartbeat is send.
-	// So use `==` here.
-	return r.approximateSize == EmptyRegionApproximateSize
-}
-
 // GetStorePeerApproximateKeys returns the approximate keys of the peer on the specified store.
 func (r *RegionInfo) GetStorePeerApproximateKeys(storeID uint64) int64 {
 	peer := r.GetStorePeer(storeID)
diff --git a/pkg/mcs/scheduling/server/cluster.go b/pkg/mcs/scheduling/server/cluster.go
index 20af077d2..b5e42b40a 100644
--- a/pkg/mcs/scheduling/server/cluster.go
+++ b/pkg/mcs/scheduling/server/cluster.go
@@ -423,9 +423,7 @@ func (c *Cluster) processRegionHeartbeat(region *core.RegionInfo) error {
 	if err != nil {
 		return err
 	}
-	if c.GetStoreConfig().IsEnableRegionBucket() {
-		region.InheritBuckets(origin)
-	}
+	region.Inherit(origin, c.GetStoreConfig().IsEnableRegionBucket())
 
 	cluster.HandleStatsAsync(c, region)
 
diff --git a/pkg/statistics/region.go b/pkg/statistics/region.go
index f78c8c6ad..f39c58ed8 100644
--- a/pkg/statistics/region.go
+++ b/pkg/statistics/region.go
@@ -59,12 +59,10 @@ func (s *RegionStats) Observe(r *core.RegionInfo) {
 	approximateKeys := r.GetApproximateKeys()
 	approximateSize := r.GetApproximateSize()
 	approximateKvSize := r.GetApproximateKvSize()
-	if approximateSize == core.EmptyRegionApproximateSize {
+	if approximateSize <= core.EmptyRegionApproximateSize {
 		s.EmptyCount++
 	}
-	if !r.IsEmptyRegion() {
-		s.StorageSize += approximateSize
-	}
+	s.StorageSize += approximateSize
 	s.UserStorageSize += approximateKvSize
 	s.StorageKeys += approximateKeys
 	leader := r.GetLeader()
diff --git a/pkg/statistics/region_collection.go b/pkg/statistics/region_collection.go
index c24a1581d..6d3ac3f17 100644
--- a/pkg/statistics/region_collection.go
+++ b/pkg/statistics/region_collection.go
@@ -198,7 +198,7 @@ func (r *RegionStatistics) Observe(region *core.RegionInfo, stores []*core.Store
 			return false
 		}(),
 		LearnerPeer: len(region.GetLearners()) > 0,
-		EmptyRegion: region.IsEmptyRegion(),
+		EmptyRegion: region.GetApproximateSize() <= core.EmptyRegionApproximateSize,
 		OversizedRegion: region.IsOversized(
 			int64(r.conf.GetRegionMaxSize()),
 			int64(r.conf.GetRegionMaxKeys()),
@@ -206,7 +206,7 @@ func (r *RegionStatistics) Observe(region *core.RegionInfo, stores []*core.Store
 		UndersizedRegion: region.NeedMerge(
 			int64(r.conf.GetMaxMergeRegionSize()),
 			int64(r.conf.GetMaxMergeRegionKeys()),
-		) && region.GetApproximateSize() >= core.EmptyRegionApproximateSize,
+		),
 		WitnessLeader: region.GetLeader().GetIsWitness(),
 	}
 	// Check if the region meets any of the conditions and update the corresponding info.
diff --git a/server/cluster/cluster.go b/server/cluster/cluster.go
index 22e1b16d8..30ab50f0c 100644
--- a/server/cluster/cluster.go
+++ b/server/cluster/cluster.go
@@ -1105,9 +1105,7 @@ func (c *RaftCluster) processRegionHeartbeat(region *core.RegionInfo) error {
 	if err != nil {
 		return err
 	}
-	if c.GetStoreConfig().IsEnableRegionBucket() {
-		region.InheritBuckets(origin)
-	}
+	region.Inherit(origin, c.GetStoreConfig().IsEnableRegionBucket())
 
 	if !c.isAPIServiceMode {
 		cluster.HandleStatsAsync(c, region)


Test patch:
diff --git a/pkg/core/region_test.go b/pkg/core/region_test.go
index 1e6b43fbf..5588d9190 100644
--- a/pkg/core/region_test.go
+++ b/pkg/core/region_test.go
@@ -186,9 +186,35 @@ func TestSortedEqual(t *testing.T) {
 	}
 }
 
-func TestInheritBuckets(t *testing.T) {
+func TestInherit(t *testing.T) {
 	re := require.New(t)
+	// size in MB
+	// case for approximateSize
+	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 _, testCase := range testCases {
+		var origin *RegionInfo
+		if testCase.originExists {
+			origin = NewRegionInfo(&metapb.Region{Id: 100}, nil)
+			origin.approximateSize = int64(testCase.originSize)
+		}
+		r := NewRegionInfo(&metapb.Region{Id: 100}, nil)
+		r.approximateSize = int64(testCase.size)
+		r.Inherit(origin, false)
+		re.Equal(int64(testCase.expect), r.approximateSize)
+	}
 
+	// bucket
 	data := []struct {
 		originBuckets *metapb.Buckets
 		buckets       *metapb.Buckets
@@ -201,11 +227,12 @@ func TestInheritBuckets(t *testing.T) {
 	for _, d := range data {
 		origin := NewRegionInfo(&metapb.Region{Id: 100}, nil, SetBuckets(d.originBuckets))
 		r := NewRegionInfo(&metapb.Region{Id: 100}, nil)
-		r.InheritBuckets(origin)
+		r.Inherit(origin, true)
 		re.Equal(d.originBuckets, r.GetBuckets())
 		// region will not inherit bucket keys.
 		if origin.GetBuckets() != nil {
 			newRegion := NewRegionInfo(&metapb.Region{Id: 100}, nil)
+			newRegion.Inherit(origin, false)
 			re.NotEqual(d.originBuckets, newRegion.GetBuckets())
 		}
 	}
diff --git a/pkg/statistics/region_collection_test.go b/pkg/statistics/region_collection_test.go
index 2706ffeb0..f0df9ce6e 100644
--- a/pkg/statistics/region_collection_test.go
+++ b/pkg/statistics/region_collection_test.go
@@ -63,7 +63,7 @@ func TestRegionStatistics(t *testing.T) {
 	stores[3] = store3
 	r1 := &metapb.Region{Id: 1, Peers: peers, StartKey: []byte("aa"), EndKey: []byte("bb")}
 	r2 := &metapb.Region{Id: 2, Peers: peers[0:2], StartKey: []byte("cc"), EndKey: []byte("dd")}
-	region1 := core.NewRegionInfo(r1, peers[0], core.SetApproximateSize(1))
+	region1 := core.NewRegionInfo(r1, peers[0])
 	region2 := core.NewRegionInfo(r2, peers[0])
 	regionStats := NewRegionStatistics(nil, opt, manager)
 	regionStats.Observe(region1, stores)
@@ -97,8 +97,7 @@ func TestRegionStatistics(t *testing.T) {
 	re.Len(regionStats.stats[PendingPeer], 1)
 	re.Len(regionStats.stats[LearnerPeer], 1)
 	re.Len(regionStats.stats[OversizedRegion], 1)
-	re.Len(regionStats.stats[UndersizedRegion], 0)
-	re.Len(regionStats.stats[EmptyRegion], 0)
+	re.Len(regionStats.stats[UndersizedRegion], 1)
 	re.Len(regionStats.stats[OfflinePeer], 1)
 
 	region1 = region1.Clone(core.WithRemoveStorePeer(7))
diff --git a/server/api/stats_test.go b/server/api/stats_test.go
index 7281f8ec9..4003f53ed 100644
--- a/server/api/stats_test.go
+++ b/server/api/stats_test.go
@@ -142,7 +142,7 @@ func (suite *statsTestSuite) TestRegionStats() {
 	statsAll := &statistics.RegionStats{
 		Count:            4,
 		EmptyCount:       1,
-		StorageSize:      350,
+		StorageSize:      351,
 		UserStorageSize:  291,
 		StorageKeys:      221,
 		StoreLeaderCount: map[uint64]int{1: 1, 4: 2, 5: 1},
@@ -156,7 +156,7 @@ func (suite *statsTestSuite) TestRegionStats() {
 	stats23 := &statistics.RegionStats{
 		Count:            2,
 		EmptyCount:       1,
-		StorageSize:      200,
+		StorageSize:      201,
 		UserStorageSize:  181,
 		StorageKeys:      151,
 		StoreLeaderCount: map[uint64]int{4: 1, 5: 1},