server/cluster: keep full range preparing size current

Signed-off-by: lhy1024 <19542290+lhy1024@users.noreply.github.com>

Author: lhy1024

server/cluster/cluster.go

@@ -1946,12 +1946,19 @@ func (c regionSizeCache) getRegionSize(startKey, endKey []byte) int64 {
 	return size
 }
 
+func (c *RaftCluster) getPreparingRegionSize(startKey, endKey []byte) int64 {
+	if len(startKey) == 0 && len(endKey) == 0 {
+		return c.GetRegionSizeByRange(startKey, endKey)
+	}
+	return c.GetRegionSizeByRangeFromSubTree(startKey, endKey)
+}
+
 func (c *RaftCluster) checkStores() {
 	var (
 		offlineStores []*metapb.Store
 		upStoreCount  int
 		stores        = c.GetStores()
-		regionSizes   = newRegionSizeCache(c.GetRegionSizeByRangeFromSubTree)
+		regionSizes   = newRegionSizeCache(c.getPreparingRegionSize)
 	)
 
 	for _, store := range stores {

server/cluster/cluster_test.go

@@ -1941,6 +1941,28 @@ func TestRegionSizeCacheAcrossStoresAndRules(t *testing.T) {
 	re.Equal(2, loadCounts[regionSizeCacheKey{startKey: "a", endKey: "m"}])
 }
 
+func TestPreparingRegionSizeUsesRootForFullRange(t *testing.T) {
+	re := require.New(t)
+	cluster := &RaftCluster{BasicCluster: core.NewBasicCluster()}
+	peer := &metapb.Peer{Id: 1, StoreId: 1}
+	region := core.NewRegionInfo(&metapb.Region{
+		Id:       1,
+		StartKey: []byte("a"),
+		EndKey:   []byte("z"),
+		Peers:    []*metapb.Peer{peer},
+	}, peer, core.SetApproximateSize(10))
+	cluster.PutRegion(region)
+
+	updated := region.Clone(core.SetApproximateSize(40))
+	_, err := cluster.CheckAndPutRootTree(core.ContextTODO(), updated)
+	re.NoError(err)
+	re.Equal(int64(40), cluster.GetRegionSizeByRange(nil, nil))
+	re.Equal(int64(10), cluster.GetRegionSizeByRangeFromSubTree(nil, nil))
+
+	re.Equal(int64(40), cluster.getPreparingRegionSize(nil, nil))
+	re.Equal(int64(10), cluster.getPreparingRegionSize([]byte("a"), []byte("z")))
+}
+
 func TestStatsRegions(t *testing.T) {
 	re := require.New(t)
 	ctx, cancel := context.WithCancel(context.Background())