batch get region size

Signed-off-by: Ryan Leung <rleungx@gmail.com>

pkg/core/region.go

@@ -1609,17 +1609,34 @@ func (r *RegionsInfo) ScanRegionWithIterator(startKey []byte, iterator func(regi
 }
 
 // GetRegionSizeByRange scans regions intersecting [start key, end key), returns the total region size of this range.
-func (r *RegionsInfo) GetRegionSizeByRange(startKey, endKey []byte) int64 {
-	r.t.RLock()
-	defer r.t.RUnlock()
+func (r *RegionsInfo) GetRegionSizeByRange(startKey, endKey []byte, limit int) int64 {
 	var size int64
-	r.tree.scanRange(startKey, func(region *RegionInfo) bool {
-		if len(endKey) > 0 && bytes.Compare(region.GetStartKey(), endKey) >= 0 {
-			return false
+	for {
+		r.t.RLock()
+		var cnt int
+		r.tree.scanRange(startKey, func(region *RegionInfo) bool {
+			if len(endKey) > 0 && bytes.Compare(region.GetStartKey(), endKey) >= 0 {
+				startKey = region.GetEndKey()
+				return false
+			}
+			if limit > 0 && cnt >= limit {
+				startKey = region.GetEndKey()
+				return false
+			}
+			cnt++
+			size += region.GetApproximateSize()
+			return true
+		})
+		r.t.RUnlock()
+		if cnt == 0 {
+			break
 		}
-		size += region.GetApproximateSize()
-		return true
-	})
+
+		if len(startKey) == 0 {
+			break
+		}
+	}
+
 	return size
 }
 

server/cluster/cluster.go

@@ -108,6 +108,7 @@ const (
 	// minSnapshotDurationSec is the minimum duration that a store can tolerate.
 	// It should enlarge the limiter if the snapshot's duration is less than this value.
 	minSnapshotDurationSec = 5
+	scanRegionLimit        = 1000
 )
 
 // Server is the interface for cluster.
@@ -1908,7 +1909,7 @@ func (c *RaftCluster) checkStores() {
 func (c *RaftCluster) getThreshold(stores []*core.StoreInfo, store *core.StoreInfo) float64 {
 	start := time.Now()
 	if !c.opt.IsPlacementRulesEnabled() {
-		regionSize := c.core.GetRegionSizeByRange([]byte(""), []byte("")) * int64(c.opt.GetMaxReplicas())
+		regionSize := c.core.GetRegionSizeByRange([]byte(""), []byte(""), scanRegionLimit) * int64(c.opt.GetMaxReplicas())
 		weight := getStoreTopoWeight(store, stores, c.opt.GetLocationLabels(), c.opt.GetMaxReplicas())
 		return float64(regionSize) * weight * 0.9
 	}
@@ -1948,7 +1949,7 @@ func (c *RaftCluster) calculateRange(stores []*core.StoreInfo, store *core.Store
 				matchStores = append(matchStores, s)
 			}
 		}
-		regionSize := c.core.GetRegionSizeByRange(startKey, endKey) * int64(rule.Count)
+		regionSize := c.core.GetRegionSizeByRange(startKey, endKey, scanRegionLimit) * int64(rule.Count)
 		weight := getStoreTopoWeight(store, matchStores, rule.LocationLabels, rule.Count)
 		storeSize += float64(regionSize) * weight
 		log.Debug("calculate range result",