From 2cc3a5aaf29090fb9fc7dc5168c783b02d108f73 Mon Sep 17 00:00:00 2001 From: lhy1024 <19542290+lhy1024@users.noreply.github.com> Date: Fri, 31 Jul 2026 10:43:33 +0800 Subject: [PATCH 1/4] server/cluster: avoid repeated region size scans Cache each key range size for one checkStores round so preparing stores share the same calculation. Move the range size lookup out of the placement rule processing loop. Ref #9574 Signed-off-by: lhy1024 <19542290+lhy1024@users.noreply.github.com> --- server/cluster/cluster.go | 88 ++++++++++++++++++++++++++++------ server/cluster/cluster_test.go | 59 +++++++++++++++++++++-- 2 files changed, 127 insertions(+), 20 deletions(-) diff --git a/server/cluster/cluster.go b/server/cluster/cluster.go index 5cdb02cb23..33d8eb319f 100644 --- a/server/cluster/cluster.go +++ b/server/cluster/cluster.go @@ -1914,16 +1914,51 @@ func (c *RaftCluster) isStorePrepared() bool { return true } +type regionSizeCacheKey struct { + startKey string + endKey string +} + +// regionSizeCache is scoped to one checkStores round and is refreshed on the next round. +type regionSizeCache struct { + loader func(startKey, endKey []byte) int64 + sizes map[regionSizeCacheKey]int64 +} + +func newRegionSizeCache(loader func(startKey, endKey []byte) int64) *regionSizeCache { + return ®ionSizeCache{ + loader: loader, + } +} + +func (c *regionSizeCache) get(startKey, endKey []byte) int64 { + key := regionSizeCacheKey{ + startKey: string(startKey), + endKey: string(endKey), + } + if size, ok := c.sizes[key]; ok { + return size + } + + size := c.loader(startKey, endKey) + if c.sizes == nil { + c.sizes = make(map[regionSizeCacheKey]int64) + } + c.sizes[key] = size + return size +} + func (c *RaftCluster) checkStores() { var ( offlineStores []*metapb.Store upStoreCount int stores = c.GetStores() + regionSizes = newRegionSizeCache(c.GetRegionSizeByRange) ) for _, store := range stores { storeID := store.GetID() - isInUp, isInOffline := c.checkStore(storeID) + isInUp, isInOffline := c.checkStore(storeID, regionSizes) if isInUp { upStoreCount++ } @@ -1940,7 +1975,7 @@ func (c *RaftCluster) checkStores() { } } -func (c *RaftCluster) checkStore(storeID uint64) (isInUp, isInOffline bool) { +func (c *RaftCluster) checkStore(storeID uint64, regionSizes *regionSizeCache) (isInUp, isInOffline bool) { c.storeStateLock.Lock(uint32(storeID)) defer c.storeStateLock.Unlock(uint32(storeID)) @@ -1960,7 +1995,7 @@ func (c *RaftCluster) checkStore(storeID uint64) (isInUp, isInOffline bool) { c.GetTotalRegionCount() < core.InitClusterRegionThreshold if !readyToServe && (c.IsPrepared() || (c.IsServiceIndependent(constant.SchedulingServiceName) && c.isStorePrepared())) { kr := keyutil.NewKeyRange("", "") - threshold = c.getThreshold(c.GetStores(), store, &kr) + threshold = c.getThreshold(c.GetStores(), store, &kr, regionSizes) log.Debug("store preparing threshold", zap.Uint64("store-id", storeID), zap.Float64("threshold", threshold), zap.Float64("region-size", regionSize)) @@ -2013,37 +2048,60 @@ func (c *RaftCluster) checkStore(storeID uint64) (isInUp, isInOffline bool) { return isInUp, isInOffline } -func (c *RaftCluster) getThreshold(stores []*core.StoreInfo, store *core.StoreInfo, kr *keyutil.KeyRange) float64 { +func (c *RaftCluster) getThreshold( + stores []*core.StoreInfo, + store *core.StoreInfo, + kr *keyutil.KeyRange, + regionSizes *regionSizeCache, +) float64 { start := time.Now() if !c.opt.IsPlacementRulesEnabled() { - regionSize := c.GetRegionSizeByRange(kr.StartKey, kr.EndKey) * int64(c.opt.GetMaxReplicas()) + regionSize := regionSizes.get(kr.StartKey, kr.EndKey) * int64(c.opt.GetMaxReplicas()) weight := core.GetStoreTopoWeight(store, stores, c.opt.GetLocationLabels(), c.opt.GetMaxReplicas()) return float64(regionSize) * weight * 0.9 } keys := c.ruleManager.GetSplitKeys(kr.StartKey, kr.EndKey) if len(keys) == 0 { - return c.calculateRange(stores, store, kr.StartKey, kr.EndKey) * 0.9 + return c.calculateRange(stores, store, kr.StartKey, kr.EndKey, regionSizes) * 0.9 } storeSize := 0.0 startKey := kr.StartKey for _, key := range keys { endKey := key - storeSize += c.calculateRange(stores, store, startKey, endKey) + storeSize += c.calculateRange(stores, store, startKey, endKey, regionSizes) startKey = endKey } // the range from the last split key to the last key - storeSize += c.calculateRange(stores, store, startKey, kr.EndKey) + storeSize += c.calculateRange(stores, store, startKey, kr.EndKey, regionSizes) log.Debug("threshold calculation time", zap.Duration("cost", time.Since(start))) return storeSize * 0.9 } -func (c *RaftCluster) calculateRange(stores []*core.StoreInfo, store *core.StoreInfo, startKey, endKey []byte) float64 { - var storeSize float64 +func (c *RaftCluster) calculateRange( + stores []*core.StoreInfo, + store *core.StoreInfo, + startKey, endKey []byte, + regionSizes *regionSizeCache, +) float64 { rules := c.ruleManager.GetRulesForApplyRange(startKey, endKey) - for _, rule := range rules { - if !placement.MatchLabelConstraints(store, rule.LabelConstraints) { + firstMatchedRule := -1 + for i, rule := range rules { + if placement.MatchLabelConstraints(store, rule.LabelConstraints) { + firstMatchedRule = i + break + } + } + if firstMatchedRule == -1 { + return 0 + } + + regionSize := regionSizes.get(startKey, endKey) + var storeSize float64 + for i := firstMatchedRule; i < len(rules); i++ { + rule := rules[i] + if i != firstMatchedRule && !placement.MatchLabelConstraints(store, rule.LabelConstraints) { continue } @@ -2056,15 +2114,15 @@ func (c *RaftCluster) calculateRange(stores []*core.StoreInfo, store *core.Store matchStores = append(matchStores, s) } } - regionSize := c.GetRegionSizeByRange(startKey, endKey) * int64(rule.Count) + ruleRegionSize := regionSize * int64(rule.Count) weight := core.GetStoreTopoWeight(store, matchStores, rule.LocationLabels, rule.Count) - storeSize += float64(regionSize) * weight + storeSize += float64(ruleRegionSize) * weight log.Debug("calculate range result", logutil.ZapRedactString("start-key", string(core.HexRegionKey(startKey))), logutil.ZapRedactString("end-key", string(core.HexRegionKey(endKey))), zap.Uint64("store-id", store.GetID()), zap.String("rule", rule.String()), - zap.Int64("region-size", regionSize), + zap.Int64("region-size", ruleRegionSize), zap.Float64("weight", weight), zap.Float64("store-size", storeSize), ) diff --git a/server/cluster/cluster_test.go b/server/cluster/cluster_test.go index b8080d7a45..f4bfe033b9 100644 --- a/server/cluster/cluster_test.go +++ b/server/cluster/cluster_test.go @@ -1872,13 +1872,59 @@ func TestCalculateStoreSize1(t *testing.T) { stores := cluster.GetStores() store := cluster.GetStore(1) kr := keyutil.NewKeyRange("", "") + regionSizes := newRegionSizeCache(cluster.GetRegionSizeByRange) // 100 * 100 * 2 (placement rule) / 4 (host) * 0.9 = 4500 - re.Equal(4500.0, cluster.getThreshold(stores, store, &kr)) + re.Equal(4500.0, cluster.getThreshold(stores, store, &kr, regionSizes)) cluster.opt.SetPlacementRuleEnabled(false) cluster.opt.SetLocationLabels([]string{"zone", "rack", "host"}) // 30000 (total region size) / 3 (zone) / 4 (host) * 0.9 = 2250 - re.Equal(2250.0, cluster.getThreshold(stores, store, &kr)) + re.Equal(2250.0, cluster.getThreshold(stores, store, &kr, regionSizes)) +} + +func TestRegionSizeCacheAcrossStoresAndRules(t *testing.T) { + re := require.New(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + _, opt, err := newTestScheduleConfig() + re.NoError(err) + cfg := opt.GetReplicationConfig() + cfg.EnablePlacementRules = true + opt.SetReplicationConfig(cfg) + cluster := newTestRaftCluster(ctx, mockid.NewIDAllocator(), opt, storage.NewStorageWithMemoryBackend()) + + for _, store := range newTestStores(2, "6.0.0") { + re.NoError(cluster.PutMetaStore(store.GetMeta())) + } + re.NoError(cluster.ruleManager.SetRule(&placement.Rule{ + GroupID: placement.DefaultGroupID, + ID: "learner", + Role: placement.Learner, + Count: 1, + })) + + kr := keyutil.NewKeyRange("a", "z") + re.Len(cluster.ruleManager.GetRulesForApplyRange(kr.StartKey, kr.EndKey), 2) + loadCount := 0 + loader := func(startKey, endKey []byte) int64 { + loadCount++ + re.Equal(kr.StartKey, startKey) + re.Equal(kr.EndKey, endKey) + return 100 + } + regionSizes := newRegionSizeCache(loader) + + stores := cluster.GetStores() + threshold1 := cluster.getThreshold(stores, cluster.GetStore(1), &kr, regionSizes) + threshold2 := cluster.getThreshold(stores, cluster.GetStore(2), &kr, regionSizes) + re.Positive(threshold1) + re.Equal(threshold1, threshold2) + re.Equal(1, loadCount) + + nextRoundRegionSizes := newRegionSizeCache(loader) + re.Equal(threshold1, cluster.getThreshold(stores, cluster.GetStore(1), &kr, nextRoundRegionSizes)) + re.Equal(2, loadCount) } func TestStatsRegions(t *testing.T) { @@ -1987,8 +2033,9 @@ func TestCalculateStoreSize2(t *testing.T) { stores := cluster.GetStores() store := cluster.GetStore(1) kr := keyutil.NewKeyRange("", "") + regionSizes := newRegionSizeCache(cluster.GetRegionSizeByRange) // 100 * 100 * 4 (total region size) / 2 (dc) / 2 (logic) / 3 (host) * 0.9 = 3000 - re.Equal(3000.0, cluster.getThreshold(stores, store, &kr)) + re.Equal(3000.0, cluster.getThreshold(stores, store, &kr, regionSizes)) } func TestStores(t *testing.T) { @@ -4183,8 +4230,9 @@ func TestCheckStoresUpCountWithLowSpace(t *testing.T) { re.Equal(int(storeCount), cluster.GetStoreCount()) upStoreCount := 0 + regionSizes := newRegionSizeCache(cluster.GetRegionSizeByRange) for _, s := range cluster.GetStores() { - isUp, _ := cluster.checkStore(s.GetID()) + isUp, _ := cluster.checkStore(s.GetID(), regionSizes) if isUp { upStoreCount++ } @@ -4203,8 +4251,9 @@ func TestCheckStoresUpCountWithLowSpace(t *testing.T) { } re.NoError(cluster.HandleStoreHeartbeat(req, resp)) } + regionSizes = newRegionSizeCache(cluster.GetRegionSizeByRange) for _, s := range cluster.GetStores() { - isUp, _ := cluster.checkStore(s.GetID()) + isUp, _ := cluster.checkStore(s.GetID(), regionSizes) if isUp { upStoreCount++ } From 1bacbb04cb975887e0031e37174abe56e59a49bf Mon Sep 17 00:00:00 2001 From: lhy1024 <19542290+lhy1024@users.noreply.github.com> Date: Fri, 31 Jul 2026 10:51:46 +0800 Subject: [PATCH 2/4] server/cluster: reuse store snapshot during state checks Pass the store snapshot collected by checkStores into each store state check instead of rebuilding the complete store slice for every preparing store. Ref #9574 Signed-off-by: lhy1024 <19542290+lhy1024@users.noreply.github.com> --- server/cluster/cluster.go | 10 +++++++--- server/cluster/cluster_test.go | 10 ++++++---- 2 files changed, 13 insertions(+), 7 deletions(-) diff --git a/server/cluster/cluster.go b/server/cluster/cluster.go index 33d8eb319f..b82f2ab049 100644 --- a/server/cluster/cluster.go +++ b/server/cluster/cluster.go @@ -1958,7 +1958,7 @@ func (c *RaftCluster) checkStores() { for _, store := range stores { storeID := store.GetID() - isInUp, isInOffline := c.checkStore(storeID, regionSizes) + isInUp, isInOffline := c.checkStore(storeID, stores, regionSizes) if isInUp { upStoreCount++ } @@ -1975,7 +1975,11 @@ func (c *RaftCluster) checkStores() { } } -func (c *RaftCluster) checkStore(storeID uint64, regionSizes *regionSizeCache) (isInUp, isInOffline bool) { +func (c *RaftCluster) checkStore( + storeID uint64, + stores []*core.StoreInfo, + regionSizes *regionSizeCache, +) (isInUp, isInOffline bool) { c.storeStateLock.Lock(uint32(storeID)) defer c.storeStateLock.Unlock(uint32(storeID)) @@ -1995,7 +1999,7 @@ func (c *RaftCluster) checkStore(storeID uint64, regionSizes *regionSizeCache) ( c.GetTotalRegionCount() < core.InitClusterRegionThreshold if !readyToServe && (c.IsPrepared() || (c.IsServiceIndependent(constant.SchedulingServiceName) && c.isStorePrepared())) { kr := keyutil.NewKeyRange("", "") - threshold = c.getThreshold(c.GetStores(), store, &kr, regionSizes) + threshold = c.getThreshold(stores, store, &kr, regionSizes) log.Debug("store preparing threshold", zap.Uint64("store-id", storeID), zap.Float64("threshold", threshold), zap.Float64("region-size", regionSize)) diff --git a/server/cluster/cluster_test.go b/server/cluster/cluster_test.go index f4bfe033b9..b48905f625 100644 --- a/server/cluster/cluster_test.go +++ b/server/cluster/cluster_test.go @@ -4230,9 +4230,10 @@ func TestCheckStoresUpCountWithLowSpace(t *testing.T) { re.Equal(int(storeCount), cluster.GetStoreCount()) upStoreCount := 0 + stores := cluster.GetStores() regionSizes := newRegionSizeCache(cluster.GetRegionSizeByRange) - for _, s := range cluster.GetStores() { - isUp, _ := cluster.checkStore(s.GetID(), regionSizes) + for _, s := range stores { + isUp, _ := cluster.checkStore(s.GetID(), stores, regionSizes) if isUp { upStoreCount++ } @@ -4251,9 +4252,10 @@ func TestCheckStoresUpCountWithLowSpace(t *testing.T) { } re.NoError(cluster.HandleStoreHeartbeat(req, resp)) } + stores = cluster.GetStores() regionSizes = newRegionSizeCache(cluster.GetRegionSizeByRange) - for _, s := range cluster.GetStores() { - isUp, _ := cluster.checkStore(s.GetID(), regionSizes) + for _, s := range stores { + isUp, _ := cluster.checkStore(s.GetID(), stores, regionSizes) if isUp { upStoreCount++ } From 8835b97a4ab171654e58a8a75188be71d2ad2965 Mon Sep 17 00:00:00 2001 From: lhy1024 <19542290+lhy1024@users.noreply.github.com> Date: Fri, 31 Jul 2026 14:59:16 +0800 Subject: [PATCH 3/4] Revert "server/cluster: reuse store snapshot during state checks" This reverts commit 1bacbb04cb975887e0031e37174abe56e59a49bf. Signed-off-by: lhy1024 <19542290+lhy1024@users.noreply.github.com> --- server/cluster/cluster.go | 10 +++------- server/cluster/cluster_test.go | 10 ++++------ 2 files changed, 7 insertions(+), 13 deletions(-) diff --git a/server/cluster/cluster.go b/server/cluster/cluster.go index b82f2ab049..33d8eb319f 100644 --- a/server/cluster/cluster.go +++ b/server/cluster/cluster.go @@ -1958,7 +1958,7 @@ func (c *RaftCluster) checkStores() { for _, store := range stores { storeID := store.GetID() - isInUp, isInOffline := c.checkStore(storeID, stores, regionSizes) + isInUp, isInOffline := c.checkStore(storeID, regionSizes) if isInUp { upStoreCount++ } @@ -1975,11 +1975,7 @@ func (c *RaftCluster) checkStores() { } } -func (c *RaftCluster) checkStore( - storeID uint64, - stores []*core.StoreInfo, - regionSizes *regionSizeCache, -) (isInUp, isInOffline bool) { +func (c *RaftCluster) checkStore(storeID uint64, regionSizes *regionSizeCache) (isInUp, isInOffline bool) { c.storeStateLock.Lock(uint32(storeID)) defer c.storeStateLock.Unlock(uint32(storeID)) @@ -1999,7 +1995,7 @@ func (c *RaftCluster) checkStore( c.GetTotalRegionCount() < core.InitClusterRegionThreshold if !readyToServe && (c.IsPrepared() || (c.IsServiceIndependent(constant.SchedulingServiceName) && c.isStorePrepared())) { kr := keyutil.NewKeyRange("", "") - threshold = c.getThreshold(stores, store, &kr, regionSizes) + threshold = c.getThreshold(c.GetStores(), store, &kr, regionSizes) log.Debug("store preparing threshold", zap.Uint64("store-id", storeID), zap.Float64("threshold", threshold), zap.Float64("region-size", regionSize)) diff --git a/server/cluster/cluster_test.go b/server/cluster/cluster_test.go index b48905f625..f4bfe033b9 100644 --- a/server/cluster/cluster_test.go +++ b/server/cluster/cluster_test.go @@ -4230,10 +4230,9 @@ func TestCheckStoresUpCountWithLowSpace(t *testing.T) { re.Equal(int(storeCount), cluster.GetStoreCount()) upStoreCount := 0 - stores := cluster.GetStores() regionSizes := newRegionSizeCache(cluster.GetRegionSizeByRange) - for _, s := range stores { - isUp, _ := cluster.checkStore(s.GetID(), stores, regionSizes) + for _, s := range cluster.GetStores() { + isUp, _ := cluster.checkStore(s.GetID(), regionSizes) if isUp { upStoreCount++ } @@ -4252,10 +4251,9 @@ func TestCheckStoresUpCountWithLowSpace(t *testing.T) { } re.NoError(cluster.HandleStoreHeartbeat(req, resp)) } - stores = cluster.GetStores() regionSizes = newRegionSizeCache(cluster.GetRegionSizeByRange) - for _, s := range stores { - isUp, _ := cluster.checkStore(s.GetID(), stores, regionSizes) + for _, s := range cluster.GetStores() { + isUp, _ := cluster.checkStore(s.GetID(), regionSizes) if isUp { upStoreCount++ } From ccc20b601abf57c950216a5359b3fd13cb939a6e Mon Sep 17 00:00:00 2001 From: lhy1024 <19542290+lhy1024@users.noreply.github.com> Date: Fri, 31 Jul 2026 15:04:29 +0800 Subject: [PATCH 4/4] server/cluster: tighten range size cache tests Signed-off-by: lhy1024 <19542290+lhy1024@users.noreply.github.com> --- server/cluster/cluster.go | 27 +++++++++++---------------- server/cluster/cluster_test.go | 34 ++++++++++++++++++++++++---------- 2 files changed, 35 insertions(+), 26 deletions(-) diff --git a/server/cluster/cluster.go b/server/cluster/cluster.go index 33d8eb319f..40c08d8f26 100644 --- a/server/cluster/cluster.go +++ b/server/cluster/cluster.go @@ -2086,24 +2086,19 @@ func (c *RaftCluster) calculateRange( regionSizes *regionSizeCache, ) float64 { rules := c.ruleManager.GetRulesForApplyRange(startKey, endKey) - firstMatchedRule := -1 - for i, rule := range rules { - if placement.MatchLabelConstraints(store, rule.LabelConstraints) { - firstMatchedRule = i - break - } - } - if firstMatchedRule == -1 { - return 0 - } - - regionSize := regionSizes.get(startKey, endKey) - var storeSize float64 - for i := firstMatchedRule; i < len(rules); i++ { - rule := rules[i] - if i != firstMatchedRule && !placement.MatchLabelConstraints(store, rule.LabelConstraints) { + var ( + regionSize int64 + regionSizeLoaded bool + storeSize float64 + ) + for _, rule := range rules { + if !placement.MatchLabelConstraints(store, rule.LabelConstraints) { continue } + if !regionSizeLoaded { + regionSize = regionSizes.get(startKey, endKey) + regionSizeLoaded = true + } var matchStores []*core.StoreInfo for _, s := range stores { diff --git a/server/cluster/cluster_test.go b/server/cluster/cluster_test.go index f4bfe033b9..b3e341bbae 100644 --- a/server/cluster/cluster_test.go +++ b/server/cluster/cluster_test.go @@ -1878,6 +1878,7 @@ func TestCalculateStoreSize1(t *testing.T) { cluster.opt.SetPlacementRuleEnabled(false) cluster.opt.SetLocationLabels([]string{"zone", "rack", "host"}) + regionSizes = newRegionSizeCache(cluster.GetRegionSizeByRange) // 30000 (total region size) / 3 (zone) / 4 (host) * 0.9 = 2250 re.Equal(2250.0, cluster.getThreshold(stores, store, &kr, regionSizes)) } @@ -1904,27 +1905,40 @@ func TestRegionSizeCacheAcrossStoresAndRules(t *testing.T) { Count: 1, })) - kr := keyutil.NewKeyRange("a", "z") + kr := keyutil.NewKeyRange("a", "m") + otherKR := keyutil.NewKeyRange("m", "z") re.Len(cluster.ruleManager.GetRulesForApplyRange(kr.StartKey, kr.EndKey), 2) - loadCount := 0 + loadCounts := make(map[regionSizeCacheKey]int) loader := func(startKey, endKey []byte) int64 { - loadCount++ - re.Equal(kr.StartKey, startKey) - re.Equal(kr.EndKey, endKey) - return 100 + key := regionSizeCacheKey{startKey: string(startKey), endKey: string(endKey)} + loadCounts[key]++ + switch key { + case regionSizeCacheKey{startKey: "a", endKey: "m"}: + return 100 + case regionSizeCacheKey{startKey: "m", endKey: "z"}: + return 200 + default: + re.FailNow("unexpected range", "start-key: %q, end-key: %q", startKey, endKey) + return 0 + } } regionSizes := newRegionSizeCache(loader) stores := cluster.GetStores() threshold1 := cluster.getThreshold(stores, cluster.GetStore(1), &kr, regionSizes) threshold2 := cluster.getThreshold(stores, cluster.GetStore(2), &kr, regionSizes) - re.Positive(threshold1) - re.Equal(threshold1, threshold2) - re.Equal(1, loadCount) + // (100 * 3 replicas / 2 stores + 100 * 1 learner / 2 stores) * 0.9 = 180. + re.Equal(180.0, threshold1) + re.Equal(180.0, threshold2) + re.Equal(1, loadCounts[regionSizeCacheKey{startKey: "a", endKey: "m"}]) + + // A different range is loaded separately, then shared by all rules. + re.Equal(360.0, cluster.getThreshold(stores, cluster.GetStore(1), &otherKR, regionSizes)) + re.Equal(1, loadCounts[regionSizeCacheKey{startKey: "m", endKey: "z"}]) nextRoundRegionSizes := newRegionSizeCache(loader) re.Equal(threshold1, cluster.getThreshold(stores, cluster.GetStore(1), &kr, nextRoundRegionSizes)) - re.Equal(2, loadCount) + re.Equal(2, loadCounts[regionSizeCacheKey{startKey: "a", endKey: "m"}]) } func TestStatsRegions(t *testing.T) {