Skip to content

server/cluster: avoid repeated region size scans #11072

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
lhy1024 wants to merge 4 commits into
base: master
Choose a base branch
from
Open

server/cluster: avoid repeated region size scans #11072

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
2cc3a5a
server/cluster: avoid repeated region size scans
lhy1024 Jul 31, 2026
1bacbb0
server/cluster: reuse store snapshot during state checks
lhy1024 Jul 31, 2026
8835b97
Revert "server/cluster: reuse store snapshot during state checks"
lhy1024 Jul 31, 2026
ccc20b6
server/cluster: tighten range size cache tests
lhy1024 Jul 31, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
79 changes: 66 additions & 13 deletions server/cluster/cluster.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 &regionSizeCache{
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++
}
Expand All @@ -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))

Expand All @@ -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))
Expand Down Expand Up @@ -2013,39 +2048,57 @@ 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)
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 {
Expand All @@ -2056,15 +2109,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),
)
Expand Down
73 changes: 68 additions & 5 deletions server/cluster/cluster_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1872,13 +1872,73 @@ 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"})
regionSizes = newRegionSizeCache(cluster.GetRegionSizeByRange)
// 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", "m")
otherKR := keyutil.NewKeyRange("m", "z")
re.Len(cluster.ruleManager.GetRulesForApplyRange(kr.StartKey, kr.EndKey), 2)
loadCounts := make(map[regionSizeCacheKey]int)
loader := func(startKey, endKey []byte) int64 {
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)
// (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, loadCounts[regionSizeCacheKey{startKey: "a", endKey: "m"}])
}

func TestStatsRegions(t *testing.T) {
Expand Down Expand Up @@ -1987,8 +2047,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) {
Expand Down Expand Up @@ -4183,8 +4244,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++
}
Expand All @@ -4203,8 +4265,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++
}
Expand Down
Loading