scheduler: extract storage operations to schedulerConfig interface

pkg/schedule/schedulers/balance_leader.go

@@ -32,7 +32,6 @@ import (
 	"github.com/tikv/pd/pkg/schedule/operator"
 	"github.com/tikv/pd/pkg/schedule/plan"
 	types "github.com/tikv/pd/pkg/schedule/type"
-	"github.com/tikv/pd/pkg/storage/endpoint"
 	"github.com/tikv/pd/pkg/utils/reflectutil"
 	"github.com/tikv/pd/pkg/utils/syncutil"
 	"github.com/tikv/pd/pkg/utils/typeutil"
@@ -56,8 +55,9 @@ const (
 
 type balanceLeaderSchedulerConfig struct {
 	syncutil.RWMutex
-	storage endpoint.ConfigStorage
-	Ranges  []core.KeyRange `json:"ranges"`
+	schedulerConfig
+
+	Ranges []core.KeyRange `json:"ranges"`
 	// Batch is used to generate multiple operators by one scheduling
 	Batch int `json:"batch"`
 }
@@ -79,7 +79,7 @@ func (conf *balanceLeaderSchedulerConfig) update(data []byte) (int, any) {
 			}
 			return http.StatusBadRequest, "invalid batch size which should be an integer between 1 and 10"
 		}
-		if err := conf.persistLocked(); err != nil {
+		if err := conf.schedulerConfig.save(conf); err != nil {
 			log.Warn("failed to save balance-leader-scheduler config", errs.ZapError(err))
 		}
 		log.Info("balance-leader-scheduler config is updated", zap.ByteString("old", oldConfig), zap.ByteString("new", newConfig))
@@ -111,14 +111,6 @@ func (conf *balanceLeaderSchedulerConfig) clone() *balanceLeaderSchedulerConfig
 	}
 }
 
-func (conf *balanceLeaderSchedulerConfig) persistLocked() error {
-	data, err := EncodeConfig(conf)
-	if err != nil {
-		return err
-	}
-	return conf.storage.SaveSchedulerConfig(BalanceLeaderName, data)
-}
-
 func (conf *balanceLeaderSchedulerConfig) getBatch() int {
 	conf.RLock()
 	defer conf.RUnlock()
@@ -216,15 +208,9 @@ func (l *balanceLeaderScheduler) EncodeConfig() ([]byte, error) {
 func (l *balanceLeaderScheduler) ReloadConfig() error {
 	l.conf.Lock()
 	defer l.conf.Unlock()
-	cfgData, err := l.conf.storage.LoadSchedulerConfig(l.GetName())
-	if err != nil {
-		return err
-	}
-	if len(cfgData) == 0 {
-		return nil
-	}
+
 	newCfg := &balanceLeaderSchedulerConfig{}
-	if err = DecodeConfig([]byte(cfgData), newCfg); err != nil {
+	if err := l.conf.load(newCfg); err != nil {
 		return err
 	}
 	l.conf.Ranges = newCfg.Ranges

pkg/schedule/schedulers/balance_witness.go

@@ -33,7 +33,6 @@ import (
 	"github.com/tikv/pd/pkg/schedule/operator"
 	"github.com/tikv/pd/pkg/schedule/plan"
 	types "github.com/tikv/pd/pkg/schedule/type"
-	"github.com/tikv/pd/pkg/storage/endpoint"
 	"github.com/tikv/pd/pkg/utils/reflectutil"
 	"github.com/tikv/pd/pkg/utils/syncutil"
 	"github.com/unrolled/render"
@@ -53,8 +52,9 @@ const (
 
 type balanceWitnessSchedulerConfig struct {
 	syncutil.RWMutex
-	storage endpoint.ConfigStorage
-	Ranges  []core.KeyRange `json:"ranges"`
+	schedulerConfig
+
+	Ranges []core.KeyRange `json:"ranges"`
 	// Batch is used to generate multiple operators by one scheduling
 	Batch int `json:"batch"`
 }
@@ -76,7 +76,7 @@ func (conf *balanceWitnessSchedulerConfig) update(data []byte) (int, any) {
 			}
 			return http.StatusBadRequest, "invalid batch size which should be an integer between 1 and 10"
 		}
-		if err := conf.persistLocked(); err != nil {
+		if err := conf.schedulerConfig.save(conf); err != nil {
 			log.Warn("failed to persist config", zap.Error(err))
 		}
 		log.Info("balance-witness-scheduler config is updated", zap.ByteString("old", oldc), zap.ByteString("new", newc))
@@ -108,14 +108,6 @@ func (conf *balanceWitnessSchedulerConfig) clone() *balanceWitnessSchedulerConfi
 	}
 }
 
-func (conf *balanceWitnessSchedulerConfig) persistLocked() error {
-	data, err := EncodeConfig(conf)
-	if err != nil {
-		return err
-	}
-	return conf.storage.SaveSchedulerConfig(BalanceWitnessName, data)
-}
-
 func (conf *balanceWitnessSchedulerConfig) getBatch() int {
 	conf.RLock()
 	defer conf.RUnlock()
@@ -215,15 +207,9 @@ func (b *balanceWitnessScheduler) EncodeConfig() ([]byte, error) {
 func (b *balanceWitnessScheduler) ReloadConfig() error {
 	b.conf.Lock()
 	defer b.conf.Unlock()
-	cfgData, err := b.conf.storage.LoadSchedulerConfig(b.GetName())
-	if err != nil {
-		return err
-	}
-	if len(cfgData) == 0 {
-		return nil
-	}
+
 	newCfg := &balanceWitnessSchedulerConfig{}
-	if err = DecodeConfig([]byte(cfgData), newCfg); err != nil {
+	if err := b.conf.load(newCfg); err != nil {
 		return err
 	}
 	b.conf.Ranges = newCfg.Ranges

pkg/schedule/schedulers/config.go

@@ -0,0 +1,56 @@
+// Copyright 2024 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 schedulers
+
+import (
+	"github.com/pingcap/errors"
+	"github.com/pingcap/failpoint"
+	"github.com/tikv/pd/pkg/storage/endpoint"
+)
+
+type schedulerConfig interface {
+	save(any) error
+	load(any) error
+	init(name string, storage endpoint.ConfigStorage)
+}
+
+type baseSchedulerConfig struct {
+	name    string
+	storage endpoint.ConfigStorage
+}
+
+func (b *baseSchedulerConfig) init(name string, storage endpoint.ConfigStorage) {
+	b.name = name
+	b.storage = storage
+}
+
+func (b *baseSchedulerConfig) save(v any) error {
+	data, err := EncodeConfig(v)
+	failpoint.Inject("persistFail", func() {
+		err = errors.New("fail to persist")
+	})
+	if err != nil {
+		return err
+	}
+	return b.storage.SaveSchedulerConfig(b.name, data)
+}
+
+func (b *baseSchedulerConfig) load(v any) error {
+	data, err := b.storage.LoadSchedulerConfig(b.name)
+	if err != nil {
+		return err
+	}
+	return DecodeConfig([]byte(data), v)
+}

pkg/schedule/schedulers/config_test.go

@@ -0,0 +1,46 @@
+// Copyright 2024 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 schedulers
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/require"
+	"github.com/tikv/pd/pkg/storage"
+)
+
+func TestSchedulerConfig(t *testing.T) {
+	s := storage.NewStorageWithMemoryBackend()
+	cfg := &baseSchedulerConfig{}
+	cfg.init("test", s)
+
+	type testConfig struct {
+		Value string `json:"value"`
+	}
+
+	tc := &testConfig{
+		Value: "test",
+	}
+	require.NoError(t, cfg.save(tc))
+	newTc := &testConfig{}
+	require.NoError(t, cfg.load(newTc))
+	require.Equal(t, tc.Value, newTc.Value)
+
+	// config with another name cannot loaded the previous config
+	cfg2 := &baseSchedulerConfig{}
+	cfg2.init("test2", s)
+	// report error because the config is empty and cannot be decoded
+	require.Error(t, cfg2.load(newTc))
+}

pkg/schedule/schedulers/evict_leader.go

@@ -30,7 +30,6 @@ import (
 	"github.com/tikv/pd/pkg/schedule/operator"
 	"github.com/tikv/pd/pkg/schedule/plan"
 	types "github.com/tikv/pd/pkg/schedule/type"
-	"github.com/tikv/pd/pkg/storage/endpoint"
 	"github.com/tikv/pd/pkg/utils/apiutil"
 	"github.com/tikv/pd/pkg/utils/syncutil"
 	"github.com/unrolled/render"
@@ -48,7 +47,8 @@ const (
 
 type evictLeaderSchedulerConfig struct {
 	syncutil.RWMutex
-	storage           endpoint.ConfigStorage
+	schedulerConfig
+
 	StoreIDWithRanges map[uint64][]core.KeyRange `json:"store-id-ranges"`
 	// Batch is used to generate multiple operators by one scheduling
 	Batch             int `json:"batch"`
@@ -85,17 +85,6 @@ func (conf *evictLeaderSchedulerConfig) clone() *evictLeaderSchedulerConfig {
 	}
 }
 
-func (conf *evictLeaderSchedulerConfig) persistLocked() error {
-	data, err := EncodeConfig(conf)
-	failpoint.Inject("persistFail", func() {
-		err = errors.New("fail to persist")
-	})
-	if err != nil {
-		return err
-	}
-	return conf.storage.SaveSchedulerConfig(types.EvictLeaderScheduler.String(), data)
-}
-
 func (conf *evictLeaderSchedulerConfig) getRanges(id uint64) []string {
 	conf.RLock()
 	defer conf.RUnlock()
@@ -145,18 +134,11 @@ func (conf *evictLeaderSchedulerConfig) encodeConfig() ([]byte, error) {
 	return EncodeConfig(conf)
 }
 
-func (conf *evictLeaderSchedulerConfig) reloadConfig(name string) error {
+func (conf *evictLeaderSchedulerConfig) reloadConfig() error {
 	conf.Lock()
 	defer conf.Unlock()
-	cfgData, err := conf.storage.LoadSchedulerConfig(name)
-	if err != nil {
-		return err
-	}
-	if len(cfgData) == 0 {
-		return nil
-	}
 	newCfg := &evictLeaderSchedulerConfig{}
-	if err = DecodeConfig([]byte(cfgData), newCfg); err != nil {
+	if err := conf.load(newCfg); err != nil {
 		return err
 	}
 	pauseAndResumeLeaderTransfer(conf.cluster, conf.StoreIDWithRanges, newCfg.StoreIDWithRanges)
@@ -203,7 +185,11 @@ func (conf *evictLeaderSchedulerConfig) update(id uint64, newRanges []core.KeyRa
 		conf.StoreIDWithRanges[id] = newRanges
 	}
 	conf.Batch = batch
-	err := conf.persistLocked()
+	var err error
+	failpoint.Inject("persistFail", func() {
+		err = errors.New("fail to persist")
+	})
+	err = conf.save(conf)
 	if err != nil && id != 0 {
 		_, _ = conf.removeStoreLocked(id)
 	}
@@ -220,7 +206,7 @@ func (conf *evictLeaderSchedulerConfig) delete(id uint64) (any, error) {
 	}
 
 	keyRanges := conf.StoreIDWithRanges[id]
-	err = conf.persistLocked()
+	err = conf.save(conf)
 	if err != nil {
 		conf.resetStoreLocked(id, keyRanges)
 		conf.Unlock()
@@ -275,7 +261,7 @@ func (s *evictLeaderScheduler) EncodeConfig() ([]byte, error) {
 
 // ReloadConfig reloads the config from the storage.
 func (s *evictLeaderScheduler) ReloadConfig() error {
-	return s.conf.reloadConfig(s.GetName())
+	return s.conf.reloadConfig()
 }
 
 // PrepareConfig implements the Scheduler interface.