diff --git a/pkg/mcs/scheduling/server/config/config.go b/pkg/mcs/scheduling/server/config/config.go index e04bea49e2e..6f6b94facf5 100644 --- a/pkg/mcs/scheduling/server/config/config.go +++ b/pkg/mcs/scheduling/server/config/config.go @@ -499,10 +499,7 @@ func (o *PersistConfig) GetStoreLimit(storeID uint64) (returnSC sc.StoreLimitCon return limit } cfg := o.GetScheduleConfig().Clone() - limitCfg := sc.StoreLimitConfig{ - AddPeer: sc.DefaultStoreLimit.GetDefaultStoreLimit(storelimit.AddPeer), - RemovePeer: sc.DefaultStoreLimit.GetDefaultStoreLimit(storelimit.RemovePeer), - } + limitCfg := cfg.GetDefaultStoreLimit() cfg.StoreLimit[storeID] = limitCfg o.SetScheduleConfig(cfg) return o.GetScheduleConfig().StoreLimit[storeID] @@ -579,12 +576,14 @@ func (o *PersistConfig) SetAllStoresLimit(typ storelimit.Type, ratePerMin float6 v := o.GetScheduleConfig().Clone() switch typ { case storelimit.AddPeer: + v.DefaultStoreLimit.AddPeer = ratePerMin sc.DefaultStoreLimit.SetDefaultStoreLimit(storelimit.AddPeer, ratePerMin) for storeID := range v.StoreLimit { sc := sc.StoreLimitConfig{AddPeer: ratePerMin, RemovePeer: v.StoreLimit[storeID].RemovePeer} v.StoreLimit[storeID] = sc } case storelimit.RemovePeer: + v.DefaultStoreLimit.RemovePeer = ratePerMin sc.DefaultStoreLimit.SetDefaultStoreLimit(storelimit.RemovePeer, ratePerMin) for storeID := range v.StoreLimit { sc := sc.StoreLimitConfig{AddPeer: v.StoreLimit[storeID].AddPeer, RemovePeer: ratePerMin} diff --git a/pkg/mcs/scheduling/server/config/config_test.go b/pkg/mcs/scheduling/server/config/config_test.go new file mode 100644 index 00000000000..30752ae685f --- /dev/null +++ b/pkg/mcs/scheduling/server/config/config_test.go @@ -0,0 +1,50 @@ +// Copyright 2026 TiKV Project Authors. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package config + +import ( + "testing" + + "github.com/stretchr/testify/require" + + "github.com/tikv/pd/pkg/core/storelimit" + sc "github.com/tikv/pd/pkg/schedule/config" +) + +func TestPersistConfigDefaultStoreLimit(t *testing.T) { + re := require.New(t) + oldAddPeer := sc.DefaultStoreLimit.GetDefaultStoreLimit(storelimit.AddPeer) + oldRemovePeer := sc.DefaultStoreLimit.GetDefaultStoreLimit(storelimit.RemovePeer) + defer func() { + sc.DefaultStoreLimit.SetDefaultStoreLimit(storelimit.AddPeer, oldAddPeer) + sc.DefaultStoreLimit.SetDefaultStoreLimit(storelimit.RemovePeer, oldRemovePeer) + }() + sc.DefaultStoreLimit.SetDefaultStoreLimit(storelimit.AddPeer, 15) + sc.DefaultStoreLimit.SetDefaultStoreLimit(storelimit.RemovePeer, 15) + + cfg := NewConfig() + re.NoError(cfg.adjust(nil)) + persistConfig := NewPersistConfig(cfg, nil) + persistConfig.GetScheduleConfig().StoreLimit[1] = sc.StoreLimitConfig{AddPeer: 10, RemovePeer: 20} + + persistConfig.SetAllStoresLimit(storelimit.AddPeer, 60) + re.Equal(sc.StoreLimitConfig{AddPeer: 60, RemovePeer: 15}, persistConfig.GetScheduleConfig().DefaultStoreLimit) + re.Equal(sc.StoreLimitConfig{AddPeer: 60, RemovePeer: 20}, persistConfig.GetStoreLimit(1)) + + // Simulate a restarted process whose package-level default goes back to the built-in value. + sc.DefaultStoreLimit.SetDefaultStoreLimit(storelimit.AddPeer, 15) + sc.DefaultStoreLimit.SetDefaultStoreLimit(storelimit.RemovePeer, 15) + re.Equal(sc.StoreLimitConfig{AddPeer: 60, RemovePeer: 15}, persistConfig.GetStoreLimit(2)) +} diff --git a/pkg/schedule/config/config.go b/pkg/schedule/config/config.go index eacfa1a4847..0a76b513bb9 100644 --- a/pkg/schedule/config/config.go +++ b/pkg/schedule/config/config.go @@ -227,6 +227,8 @@ type ScheduleConfig struct { // StoreBalanceRate is the maximum of balance rate for each store. // WARN: StoreBalanceRate is deprecated. StoreBalanceRate float64 `toml:"store-balance-rate" json:"store-balance-rate,omitempty"` + // DefaultStoreLimit is the default limit of scheduling for stores. + DefaultStoreLimit StoreLimitConfig `toml:"default-store-limit" json:"default-store-limit"` // StoreLimit is the limit of scheduling for stores. StoreLimit map[uint64]StoreLimitConfig `toml:"store-limit" json:"store-limit"` // TolerantSizeRatio is the ratio of buffer size for balance scheduler. @@ -457,6 +459,9 @@ func (c *ScheduleConfig) Adjust(meta *configutil.ConfigMetaData, reloading bool) } adjustSchedulers(&c.Schedulers, DefaultSchedulers) + defaultStoreLimitMeta := meta.Child("default-store-limit") + c.migrateStoreBalanceRate(defaultStoreLimitMeta) + c.adjustDefaultStoreLimit(defaultStoreLimitMeta) for k, b := range c.migrateConfigurationMap() { v, err := parseDeprecatedFlag(meta, k, *b[0], *b[1]) @@ -466,11 +471,6 @@ func (c *ScheduleConfig) Adjust(meta *configutil.ConfigMetaData, reloading bool) *b[0], *b[1] = false, v // reset old flag false to make it ignored when marshal to JSON } - if c.StoreBalanceRate != 0 { - DefaultStoreLimit = StoreLimit{AddPeer: c.StoreBalanceRate, RemovePeer: c.StoreBalanceRate} - c.StoreBalanceRate = 0 - } - if c.StoreLimit == nil { c.StoreLimit = make(map[uint64]StoreLimitConfig) } @@ -489,6 +489,38 @@ func (c *ScheduleConfig) Adjust(meta *configutil.ConfigMetaData, reloading bool) return c.Validate() } +func (c *ScheduleConfig) adjustDefaultStoreLimit(meta *configutil.ConfigMetaData) { + defaultStoreLimit := DefaultStoreLimitConfig() + if !meta.IsDefined("add-peer") { + configutil.AdjustFloat64(&c.DefaultStoreLimit.AddPeer, defaultStoreLimit.AddPeer) + } + if !meta.IsDefined("remove-peer") { + configutil.AdjustFloat64(&c.DefaultStoreLimit.RemovePeer, defaultStoreLimit.RemovePeer) + } +} + +func (c *ScheduleConfig) migrateStoreBalanceRate(defaultStoreLimitMeta *configutil.ConfigMetaData) { + if c.StoreBalanceRate == 0 { + return + } + defaultStoreLimit := StoreLimitConfig{AddPeer: c.StoreBalanceRate, RemovePeer: c.StoreBalanceRate} + DefaultStoreLimit = StoreLimit{AddPeer: defaultStoreLimit.AddPeer, RemovePeer: defaultStoreLimit.RemovePeer} + if defaultStoreLimitMeta == nil { + if c.DefaultStoreLimit == (StoreLimitConfig{}) { + c.DefaultStoreLimit = defaultStoreLimit + } + c.StoreBalanceRate = 0 + return + } + if !defaultStoreLimitMeta.IsDefined("add-peer") { + c.DefaultStoreLimit.AddPeer = defaultStoreLimit.AddPeer + } + if !defaultStoreLimitMeta.IsDefined("remove-peer") { + c.DefaultStoreLimit.RemovePeer = defaultStoreLimit.RemovePeer + } + c.StoreBalanceRate = 0 +} + func (c *ScheduleConfig) migrateConfigurationMap() map[string][2]*bool { return map[string][2]*bool{ "remove-down-replica": {&c.DisableRemoveDownReplica, &c.EnableRemoveDownReplica}, @@ -536,10 +568,7 @@ func parseDeprecatedFlag(meta *configutil.ConfigMetaData, name string, old, new // MigrateDeprecatedFlags updates new flags according to deprecated flags. func (c *ScheduleConfig) MigrateDeprecatedFlags() { c.DisableLearner = false - if c.StoreBalanceRate != 0 { - DefaultStoreLimit = StoreLimit{AddPeer: c.StoreBalanceRate, RemovePeer: c.StoreBalanceRate} - c.StoreBalanceRate = 0 - } + c.migrateStoreBalanceRate(nil) for _, b := range c.migrateConfigurationMap() { // If old=false (previously disabled), set both old and new to false. if *b[0] { @@ -606,6 +635,19 @@ type StoreLimitConfig struct { RemovePeer float64 `toml:"remove-peer" json:"remove-peer"` } +// DefaultStoreLimitConfig returns the current process default store limit config. +func DefaultStoreLimitConfig() StoreLimitConfig { + return StoreLimitConfig{ + AddPeer: DefaultStoreLimit.GetDefaultStoreLimit(storelimit.AddPeer), + RemovePeer: DefaultStoreLimit.GetDefaultStoreLimit(storelimit.RemovePeer), + } +} + +// GetDefaultStoreLimit returns the default store limit config. +func (c *ScheduleConfig) GetDefaultStoreLimit() StoreLimitConfig { + return c.DefaultStoreLimit +} + // SchedulerConfigs is a slice of customized scheduler configuration. type SchedulerConfigs []SchedulerConfig diff --git a/server/cluster/cluster.go b/server/cluster/cluster.go index f1ce644643a..0855d109c3d 100644 --- a/server/cluster/cluster.go +++ b/server/cluster/cluster.go @@ -2314,10 +2314,7 @@ func (c *RaftCluster) AddStoreLimit(store *metapb.Store) { return } - slc := sc.StoreLimitConfig{ - AddPeer: sc.DefaultStoreLimit.GetDefaultStoreLimit(storelimit.AddPeer), - RemovePeer: sc.DefaultStoreLimit.GetDefaultStoreLimit(storelimit.RemovePeer), - } + slc := cfg.GetDefaultStoreLimit() if core.IsStoreContainLabel(store, core.EngineKey, core.EngineTiFlash) { slc = sc.StoreLimitConfig{ AddPeer: sc.DefaultTiFlashStoreLimit.GetDefaultStoreLimit(storelimit.AddPeer), @@ -2516,14 +2513,13 @@ func (c *RaftCluster) SetStoreLimit(storeID uint64, typ storelimit.Type, ratePer // SetAllStoresLimit sets all store limit for a given type and rate. func (c *RaftCluster) SetAllStoresLimit(typ storelimit.Type, ratePerMin float64) error { old := c.opt.GetScheduleConfig().Clone() - oldAdd := sc.DefaultStoreLimit.GetDefaultStoreLimit(storelimit.AddPeer) - oldRemove := sc.DefaultStoreLimit.GetDefaultStoreLimit(storelimit.RemovePeer) c.opt.SetAllStoresLimit(typ, ratePerMin) if err := c.opt.Persist(c.storage); err != nil { // roll back the store limit c.opt.SetScheduleConfig(old) - sc.DefaultStoreLimit.SetDefaultStoreLimit(storelimit.AddPeer, oldAdd) - sc.DefaultStoreLimit.SetDefaultStoreLimit(storelimit.RemovePeer, oldRemove) + oldDefaultStoreLimit := old.GetDefaultStoreLimit() + sc.DefaultStoreLimit.SetDefaultStoreLimit(storelimit.AddPeer, oldDefaultStoreLimit.AddPeer) + sc.DefaultStoreLimit.SetDefaultStoreLimit(storelimit.RemovePeer, oldDefaultStoreLimit.RemovePeer) log.Error("persist store limit meet error", errs.ZapError(err)) return err } diff --git a/server/cluster/cluster_test.go b/server/cluster/cluster_test.go index 09f554f7851..e78a5f91ffb 100644 --- a/server/cluster/cluster_test.go +++ b/server/cluster/cluster_test.go @@ -3150,6 +3150,41 @@ func TestStoreLimitChangeRefreshLimiter(t *testing.T) { re.True(store.IsAvailable(storelimit.AddPeer, constant.Low)) } +func TestAddStoreLimitUsesPersistedDefaultStoreLimit(t *testing.T) { + re := require.New(t) + oldAddPeer := sc.DefaultStoreLimit.GetDefaultStoreLimit(storelimit.AddPeer) + oldRemovePeer := sc.DefaultStoreLimit.GetDefaultStoreLimit(storelimit.RemovePeer) + defer func() { + sc.DefaultStoreLimit.SetDefaultStoreLimit(storelimit.AddPeer, oldAddPeer) + sc.DefaultStoreLimit.SetDefaultStoreLimit(storelimit.RemovePeer, oldRemovePeer) + }() + + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + _, opt, err := newTestScheduleConfig() + re.NoError(err) + rc := newTestRaftCluster(ctx, mockid.NewIDAllocator(), opt, storage.NewStorageWithMemoryBackend()) + opt.SetAllStoresLimit(storelimit.AddPeer, 60) + + // Simulate a restarted process whose package-level default goes back to the built-in value. + sc.DefaultStoreLimit.SetDefaultStoreLimit(storelimit.AddPeer, 15) + sc.DefaultStoreLimit.SetDefaultStoreLimit(storelimit.RemovePeer, 15) + + rc.AddStoreLimit(&metapb.Store{Id: 1}) + re.Equal(sc.StoreLimitConfig{AddPeer: 60, RemovePeer: 15}, opt.GetScheduleConfig().StoreLimit[1]) + + opt.SetStoreLimit(2, storelimit.RemovePeer, 80) + rc.AddStoreLimit(&metapb.Store{Id: 2}) + re.Equal(sc.StoreLimitConfig{AddPeer: 60, RemovePeer: 80}, opt.GetScheduleConfig().StoreLimit[2]) + + rc.AddStoreLimit(&metapb.Store{ + Id: 3, + Labels: []*metapb.StoreLabel{{Key: core.EngineKey, Value: core.EngineTiFlash}}, + }) + re.Equal(sc.StoreLimitConfig{AddPeer: 30, RemovePeer: 30}, opt.GetScheduleConfig().StoreLimit[3]) +} + func TestPatrolRegionConcurrency(t *testing.T) { re := require.New(t) diff --git a/server/config/config_test.go b/server/config/config_test.go index cb15d982435..59c60c60ba8 100644 --- a/server/config/config_test.go +++ b/server/config/config_test.go @@ -28,6 +28,7 @@ import ( "github.com/stretchr/testify/require" "go.uber.org/goleak" + "github.com/tikv/pd/pkg/core/storelimit" "github.com/tikv/pd/pkg/ratelimit" sc "github.com/tikv/pd/pkg/schedule/config" "github.com/tikv/pd/pkg/storage" @@ -82,6 +83,127 @@ func TestReloadConfig(t *testing.T) { re.Equal(int64(512), newOpt.GetMaxMovableHotPeerSize()) } +func TestReloadDefaultStoreLimit(t *testing.T) { + re := require.New(t) + oldAddPeer := sc.DefaultStoreLimit.GetDefaultStoreLimit(storelimit.AddPeer) + oldRemovePeer := sc.DefaultStoreLimit.GetDefaultStoreLimit(storelimit.RemovePeer) + defer func() { + sc.DefaultStoreLimit.SetDefaultStoreLimit(storelimit.AddPeer, oldAddPeer) + sc.DefaultStoreLimit.SetDefaultStoreLimit(storelimit.RemovePeer, oldRemovePeer) + }() + sc.DefaultStoreLimit.SetDefaultStoreLimit(storelimit.AddPeer, 15) + sc.DefaultStoreLimit.SetDefaultStoreLimit(storelimit.RemovePeer, 15) + + opt, err := newTestScheduleOption() + re.NoError(err) + opt.SetAllStoresLimit(storelimit.AddPeer, 60) + re.Equal(sc.StoreLimitConfig{AddPeer: 60, RemovePeer: 15}, opt.GetScheduleConfig().DefaultStoreLimit) + + storage := storage.NewStorageWithMemoryBackend() + re.NoError(opt.Persist(storage)) + + // Simulate a restarted process whose package-level default goes back to the built-in value. + sc.DefaultStoreLimit.SetDefaultStoreLimit(storelimit.AddPeer, 15) + sc.DefaultStoreLimit.SetDefaultStoreLimit(storelimit.RemovePeer, 15) + newOpt, err := newTestScheduleOption() + re.NoError(err) + re.NoError(newOpt.Reload(storage)) + + expected := sc.StoreLimitConfig{AddPeer: 60, RemovePeer: 15} + re.Equal(expected, newOpt.GetScheduleConfig().DefaultStoreLimit) + re.Equal(expected, newOpt.GetStoreLimit(100)) + + newOpt.SetStoreLimit(101, storelimit.RemovePeer, 70) + re.Equal(sc.StoreLimitConfig{AddPeer: 60, RemovePeer: 70}, newOpt.GetStoreLimit(101)) + + cfg := newOpt.GetScheduleConfig().Clone() + cfg.DefaultStoreLimit.AddPeer = 0 + newOpt.SetScheduleConfig(cfg) + re.NoError(newOpt.Persist(storage)) + + sc.DefaultStoreLimit.SetDefaultStoreLimit(storelimit.AddPeer, 15) + sc.DefaultStoreLimit.SetDefaultStoreLimit(storelimit.RemovePeer, 15) + reloadedOpt, err := newTestScheduleOption() + re.NoError(err) + re.NoError(reloadedOpt.Reload(storage)) + re.Equal(sc.StoreLimitConfig{AddPeer: 0, RemovePeer: 15}, reloadedOpt.GetScheduleConfig().DefaultStoreLimit) + re.Equal(sc.StoreLimitConfig{AddPeer: 0, RemovePeer: 15}, reloadedOpt.GetStoreLimit(102)) +} + +func TestDefaultStoreLimitAdjust(t *testing.T) { + re := require.New(t) + oldAddPeer := sc.DefaultStoreLimit.GetDefaultStoreLimit(storelimit.AddPeer) + oldRemovePeer := sc.DefaultStoreLimit.GetDefaultStoreLimit(storelimit.RemovePeer) + defer func() { + sc.DefaultStoreLimit.SetDefaultStoreLimit(storelimit.AddPeer, oldAddPeer) + sc.DefaultStoreLimit.SetDefaultStoreLimit(storelimit.RemovePeer, oldRemovePeer) + }() + sc.DefaultStoreLimit.SetDefaultStoreLimit(storelimit.AddPeer, 15) + sc.DefaultStoreLimit.SetDefaultStoreLimit(storelimit.RemovePeer, 15) + + cases := []struct { + name string + config string + expect sc.StoreLimitConfig + }{ + { + name: "preserve explicit zero", + config: ` +[schedule.default-store-limit] +add-peer = 0 +remove-peer = 60 +`, + expect: sc.StoreLimitConfig{AddPeer: 0, RemovePeer: 60}, + }, + { + name: "store balance rate backfills undefined field", + config: ` +[schedule] +store-balance-rate = 50 + +[schedule.default-store-limit] +add-peer = 0 +`, + expect: sc.StoreLimitConfig{AddPeer: 0, RemovePeer: 50}, + }, + { + name: "explicit default store limit wins over store balance rate", + config: ` +[schedule] +store-balance-rate = 50 + +[schedule.default-store-limit] +add-peer = 60 +remove-peer = 70 +`, + expect: sc.StoreLimitConfig{AddPeer: 60, RemovePeer: 70}, + }, + } + for _, testCase := range cases { + t.Run(testCase.name, func(t *testing.T) { + cfg := NewConfig() + meta, err := toml.Decode(testCase.config, cfg) + require.NoError(t, err) + require.NoError(t, cfg.Adjust(&meta, false)) + require.Equal(t, testCase.expect, cfg.Schedule.DefaultStoreLimit) + }) + } + + schedule := &sc.ScheduleConfig{StoreBalanceRate: 50} + schedule.MigrateDeprecatedFlags() + re.Equal(sc.StoreLimitConfig{AddPeer: 50, RemovePeer: 50}, schedule.DefaultStoreLimit) + + schedule = &sc.ScheduleConfig{ + StoreBalanceRate: 50, + DefaultStoreLimit: sc.StoreLimitConfig{ + AddPeer: 0, + RemovePeer: 60, + }, + } + schedule.MigrateDeprecatedFlags() + re.Equal(sc.StoreLimitConfig{AddPeer: 0, RemovePeer: 60}, schedule.DefaultStoreLimit) +} + func TestReloadUpgrade(t *testing.T) { re := require.New(t) opt, err := newTestScheduleOption() diff --git a/server/config/persist_options.go b/server/config/persist_options.go index 15cda11b139..a8609ea33e9 100644 --- a/server/config/persist_options.go +++ b/server/config/persist_options.go @@ -390,19 +390,20 @@ func (o *PersistOptions) SetMaxMergeRegionKeys(maxMergeRegionKeys uint64) { // SetStoreLimit sets a store limit for a given type and rate. func (o *PersistOptions) SetStoreLimit(storeID uint64, typ storelimit.Type, ratePerMin float64) { v := o.GetScheduleConfig().Clone() + defaultStoreLimit := v.GetDefaultStoreLimit() var slc sc.StoreLimitConfig var rate float64 switch typ { case storelimit.AddPeer: if _, ok := v.StoreLimit[storeID]; !ok { - rate = sc.DefaultStoreLimit.GetDefaultStoreLimit(storelimit.RemovePeer) + rate = defaultStoreLimit.RemovePeer } else { rate = v.StoreLimit[storeID].RemovePeer } slc = sc.StoreLimitConfig{AddPeer: ratePerMin, RemovePeer: rate} case storelimit.RemovePeer: if _, ok := v.StoreLimit[storeID]; !ok { - rate = sc.DefaultStoreLimit.GetDefaultStoreLimit(storelimit.AddPeer) + rate = defaultStoreLimit.AddPeer } else { rate = v.StoreLimit[storeID].AddPeer } @@ -417,12 +418,14 @@ func (o *PersistOptions) SetAllStoresLimit(typ storelimit.Type, ratePerMin float v := o.GetScheduleConfig().Clone() switch typ { case storelimit.AddPeer: + v.DefaultStoreLimit.AddPeer = ratePerMin sc.DefaultStoreLimit.SetDefaultStoreLimit(storelimit.AddPeer, ratePerMin) for storeID := range v.StoreLimit { sc := sc.StoreLimitConfig{AddPeer: ratePerMin, RemovePeer: v.StoreLimit[storeID].RemovePeer} v.StoreLimit[storeID] = sc } case storelimit.RemovePeer: + v.DefaultStoreLimit.RemovePeer = ratePerMin sc.DefaultStoreLimit.SetDefaultStoreLimit(storelimit.RemovePeer, ratePerMin) for storeID := range v.StoreLimit { sc := sc.StoreLimitConfig{AddPeer: v.StoreLimit[storeID].AddPeer, RemovePeer: ratePerMin} @@ -494,10 +497,7 @@ func (o *PersistOptions) GetStoreLimit(storeID uint64) (returnSC sc.StoreLimitCo return limit } cfg := o.GetScheduleConfig().Clone() - limitCfg := sc.StoreLimitConfig{ - AddPeer: sc.DefaultStoreLimit.GetDefaultStoreLimit(storelimit.AddPeer), - RemovePeer: sc.DefaultStoreLimit.GetDefaultStoreLimit(storelimit.RemovePeer), - } + limitCfg := cfg.GetDefaultStoreLimit() cfg.StoreLimit[storeID] = limitCfg o.SetScheduleConfig(cfg) return o.GetScheduleConfig().StoreLimit[storeID]