api: add a new scheduler to balance the regions of the given key range

pkg/mcs/scheduling/server/cluster.go

@@ -314,7 +314,11 @@
 		)
 		// Create the newly added schedulers.
 		for _, scheduler := range latestSchedulersConfig {
-			schedulerType := types.ConvertOldStrToType[scheduler.Type]
+			schedulerType, ok := types.ConvertOldStrToType[scheduler.Type]
+			if !ok {
+				log.Error("scheduler not found  ", zap.String("type", scheduler.Type))
+				continue
+			}
 			s, err := schedulers.CreateScheduler(
 				schedulerType,
 				c.coordinator.GetOperatorController(),

pkg/schedule/operator/kind.go

@@ -35,6 +35,8 @@ const (
 	OpMerge
 	// Initiated by range scheduler.
 	OpRange
+	// Initiated by key range scheduler.
+	OpKeyRange
 	// Initiated by replica checker.
 	OpReplica
 	// Include region split. Initiated by rule checker if `kind & OpAdmin == 0`.
@@ -60,6 +62,7 @@ var flagToName = map[OpKind]string{
 	OpHotRegion:     "hot-region",
 	OpReplica:       "replica",
 	OpMerge:         "merge",
+	OpKeyRange:      "key-range",
 	OpRange:         "range",
 	OpWitness:       "witness",
 	OpWitnessLeader: "witness-leader",
@@ -74,6 +77,7 @@ var nameToFlag = map[string]OpKind{
 	"replica":        OpReplica,
 	"merge":          OpMerge,
 	"range":          OpRange,
+	"key-range":      OpKeyRange,
 	"witness-leader": OpWitnessLeader,
 }
 

pkg/schedule/operator/operator_test.go

@@ -476,6 +476,9 @@ func (suite *operatorTestSuite) TestSchedulerKind() {
 		}, {
 			op:     NewTestOperator(1, &metapb.RegionEpoch{}, OpLeader),
 			expect: OpLeader,
+		}, {
+			op:     NewTestOperator(1, &metapb.RegionEpoch{}, OpKeyRange|OpLeader),
+			expect: OpKeyRange,
 		},
 	}
 	for _, v := range testData {

pkg/schedule/schedulers/balance_key_range.go

@@ -0,0 +1,151 @@
+package schedulers
+
+import (
+	"net/http"
+	"time"
+
+	"github.com/gorilla/mux"
+	"github.com/unrolled/render"
+
+	"github.com/pingcap/log"
+
+	"github.com/tikv/pd/pkg/core"
+	"github.com/tikv/pd/pkg/core/constant"
+	"github.com/tikv/pd/pkg/errs"
+	sche "github.com/tikv/pd/pkg/schedule/core"
+	"github.com/tikv/pd/pkg/schedule/filter"
+	"github.com/tikv/pd/pkg/schedule/operator"
+	"github.com/tikv/pd/pkg/schedule/plan"
+	"github.com/tikv/pd/pkg/schedule/types"
+	"github.com/tikv/pd/pkg/utils/syncutil"
+)
+
+type balanceKeyRangeSchedulerHandler struct {
+	rd     *render.Render
+	config *balanceKeyRangeSchedulerConfig
+}
+
+func newBalanceKeyRangeHandler(conf *balanceKeyRangeSchedulerConfig) http.Handler {
+	handler := &balanceKeyRangeSchedulerHandler{
+		config: conf,
+		rd:     render.New(render.Options{IndentJSON: true}),
+	}
+	router := mux.NewRouter()
+	router.HandleFunc("/config", handler.updateConfig).Methods(http.MethodPost)
+	router.HandleFunc("/list", handler.listConfig).Methods(http.MethodGet)
+	return router
+}
+
+func (handler *balanceKeyRangeSchedulerHandler) updateConfig(w http.ResponseWriter, _ *http.Request) {
+	handler.rd.JSON(w, http.StatusBadRequest, "update config is not supported")
+}
+
+func (handler *balanceKeyRangeSchedulerHandler) listConfig(w http.ResponseWriter, _ *http.Request) {
+	conf := handler.config.clone()
+	if err := handler.rd.JSON(w, http.StatusOK, conf); err != nil {
+		log.Error("failed to marshal balance key range scheduler config", errs.ZapError(err))
+	}
+}
+
+type balanceKeyRangeSchedulerConfig struct {
+	syncutil.RWMutex
+	schedulerConfig
+	balanceKeyRangeSchedulerParam
+}
+
+type balanceKeyRangeSchedulerParam struct {
+	Role    string          `json:"role"`
+	Engine  string          `json:"engine"`
+	Timeout time.Duration   `json:"timeout"`
+	Ranges  []core.KeyRange `json:"ranges"`
+}
+
+func (conf *balanceKeyRangeSchedulerConfig) encodeConfig() ([]byte, error) {
+	conf.RLock()
+	defer conf.RUnlock()
+	return EncodeConfig(conf)
+}
+
+func (conf *balanceKeyRangeSchedulerConfig) clone() *balanceKeyRangeSchedulerParam {
+	conf.RLock()
+	defer conf.RUnlock()
+	ranges := make([]core.KeyRange, len(conf.Ranges))
+	copy(ranges, conf.Ranges)
+	return &balanceKeyRangeSchedulerParam{
+		Ranges:  ranges,
+		Role:    conf.Role,
+		Engine:  conf.Engine,
+		Timeout: conf.Timeout,
+	}
+}
+
+// EncodeConfig serializes the config.
+func (s *balanceKeyRangeScheduler) EncodeConfig() ([]byte, error) {
+	return s.conf.encodeConfig()
+}
+
+// ReloadConfig reloads the config.
+func (s *balanceKeyRangeScheduler) ReloadConfig() error {
+	s.conf.Lock()
+	defer s.conf.Unlock()
+
+	newCfg := &balanceKeyRangeSchedulerConfig{}
+	if err := s.conf.load(newCfg); err != nil {
+		return err
+	}
+	s.conf.Ranges = newCfg.Ranges
+	s.conf.Timeout = newCfg.Timeout
+	s.conf.Role = newCfg.Role
+	s.conf.Engine = newCfg.Engine
+	return nil
+}
+
+type balanceKeyRangeScheduler struct {
+	*BaseScheduler
+	conf          *balanceKeyRangeSchedulerConfig
+	handler       http.Handler
+	filters       []filter.Filter
+	filterCounter *filter.Counter
+}
+
+// ServeHTTP implements the http.Handler interface.
+func (s *balanceKeyRangeScheduler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
+	s.handler.ServeHTTP(w, r)
+}
+
+// Schedule schedules the balance key range operator.
+func (*balanceKeyRangeScheduler) Schedule(_cluster sche.SchedulerCluster, _dryRun bool) ([]*operator.Operator, []plan.Plan) {
+	log.Debug("balance key range scheduler is scheduling, need to implement")
+	return nil, nil
+}
+
+// IsScheduleAllowed checks if the scheduler is allowed to schedule new operators.
+func (s *balanceKeyRangeScheduler) IsScheduleAllowed(cluster sche.SchedulerCluster) bool {
+	allowed := s.OpController.OperatorCount(operator.OpKeyRange) < cluster.GetSchedulerConfig().GetRegionScheduleLimit()
+	if !allowed {
+		operator.IncOperatorLimitCounter(s.GetType(), operator.OpKeyRange)
+	}
+	return allowed
+}
+
+// BalanceKeyRangeCreateOption is used to create a scheduler with an option.
+type BalanceKeyRangeCreateOption func(s *balanceKeyRangeScheduler)
+
+// newBalanceKeyRangeScheduler creates a scheduler that tends to keep given peer role on
+// special store balanced.
+func newBalanceKeyRangeScheduler(opController *operator.Controller, conf *balanceKeyRangeSchedulerConfig, options ...BalanceKeyRangeCreateOption) Scheduler {
+	s := &balanceKeyRangeScheduler{
+		BaseScheduler: NewBaseScheduler(opController, types.BalanceKeyRangeScheduler, conf),
+		conf:          conf,
+		handler:       newBalanceKeyRangeHandler(conf),
+	}
+	for _, option := range options {
+		option(s)
+	}
+	s.filters = []filter.Filter{
+		&filter.StoreStateFilter{ActionScope: s.GetName(), TransferLeader: true, OperatorLevel: constant.Medium},
+		filter.NewSpecialUseFilter(s.GetName()),
+	}
+	s.filterCounter = filter.NewCounter(s.GetName())
+	return s
+}

pkg/schedule/schedulers/init.go

@@ -15,9 +15,11 @@
 package schedulers
 
 import (
+	"net/url"
 	"strconv"
 	"strings"
 	"sync"
+	"time"
 
 	"github.com/tikv/pd/pkg/core"
 	"github.com/tikv/pd/pkg/errs"
@@ -545,4 +547,56 @@
 		conf.init(sche.GetName(), storage, conf)
 		return sche, nil
 	})
+
+	// balance key range scheduler
+	// args: [role, engine, timeout, range1, range2, ...]
+	RegisterSliceDecoderBuilder(types.BalanceKeyRangeScheduler, func(args []string) ConfigDecoder {
+		return func(v any) error {
+			conf, ok := v.(*balanceKeyRangeSchedulerConfig)
+			if !ok {
+				return errs.ErrScheduleConfigNotExist.FastGenByArgs()
+			}
+			if len(args) < 4 {
+				return errs.ErrSchedulerConfig.FastGenByArgs("args length must be greater than 3")
+			}
+			role, err := url.QueryUnescape(args[0])
+			if err != nil {
+				return errs.ErrQueryUnescape.Wrap(err)
+			}
+			engine, err := url.QueryUnescape(args[1])
+			if err != nil {
+				return errs.ErrQueryUnescape.Wrap(err)
+			}
+			timeout, err := url.QueryUnescape(args[2])
+			if err != nil {
+				return errs.ErrQueryUnescape.Wrap(err)
+			}
+			duration, err := time.ParseDuration(timeout)
+			if err != nil {
+				return errs.ErrURLParse.Wrap(err)
+			}
+			ranges, err := getKeyRanges(args[3:])
+			if err != nil {
+				return err
+			}
+			conf.Ranges = ranges
+			conf.Engine = engine
+			conf.Role = role
+			conf.Timeout = duration
+			return nil
+		}
+	})
+
+	RegisterScheduler(types.BalanceKeyRangeScheduler, func(opController *operator.Controller,
+		storage endpoint.ConfigStorage, decoder ConfigDecoder, _ ...func(string) error) (Scheduler, error) {
+		conf := &balanceKeyRangeSchedulerConfig{
+			schedulerConfig: newBaseDefaultSchedulerConfig(),
+		}
+		if err := decoder(conf); err != nil {
+			return nil, err
+		}
+		sche := newBalanceKeyRangeScheduler(opController, conf)
+		conf.init(sche.GetName(), storage, conf)
+		return sche, nil
+	})
 }

pkg/schedule/types/type.go

@@ -70,6 +70,8 @@ const (
 	TransferWitnessLeaderScheduler CheckerSchedulerType = "transfer-witness-leader-scheduler"
 	// LabelScheduler is label scheduler name.
 	LabelScheduler CheckerSchedulerType = "label-scheduler"
+	// BalanceKeyRangeScheduler is balance key range scheduler name.
+	BalanceKeyRangeScheduler CheckerSchedulerType = "balance-key-range-scheduler"
 )
 
 // TODO: SchedulerTypeCompatibleMap and ConvertOldStrToType should be removed after
@@ -97,6 +99,7 @@ var (
 		SplitBucketScheduler:           "split-bucket",
 		TransferWitnessLeaderScheduler: "transfer-witness-leader",
 		LabelScheduler:                 "label",
+		BalanceKeyRangeScheduler:       "balance-key-range",
 	}
 
 	// ConvertOldStrToType exists for compatibility.
@@ -120,6 +123,7 @@ var (
 		"split-bucket":            SplitBucketScheduler,
 		"transfer-witness-leader": TransferWitnessLeaderScheduler,
 		"label":                   LabelScheduler,
+		"balance-key-range":       BalanceKeyRangeScheduler,
 	}
 
 	// StringToSchedulerType is a map to convert the scheduler string to the CheckerSchedulerType.
@@ -143,6 +147,7 @@ var (
 		"split-bucket-scheduler":            SplitBucketScheduler,
 		"transfer-witness-leader-scheduler": TransferWitnessLeaderScheduler,
 		"label-scheduler":                   LabelScheduler,
+		"balance-key-range-scheduler":       BalanceKeyRangeScheduler,
 	}
 
 	// DefaultSchedulers is the default scheduler types.

server/api/scheduler.go

@@ -99,6 +99,39 @@
 	}
 
 	switch tp {
+	case types.BalanceKeyRangeScheduler:
+		exist, _ := h.IsSchedulerExisted(name)
+		if exist {
+			h.r.JSON(w, http.StatusBadRequest, "The scheduler already exists, pls remove the exist scheduler first.")
+			return
+		}
+		if err := apiutil.CollectStringOption("role", input, collector); err != nil {
+			h.r.JSON(w, http.StatusInternalServerError, err.Error())
+			return
+		}
+		if err := apiutil.CollectStringOption("engine", input, collector); err != nil {
+			h.r.JSON(w, http.StatusInternalServerError, err.Error())
+			return
+		}
+		defaultTimeout := "1h"
+		if err := apiutil.CollectStringOption("timeout", input, collector); err != nil {
+			if errors.ErrorEqual(err, errs.ErrOptionNotExist) {
+				collector(defaultTimeout)
+			} else {
+				h.r.JSON(w, http.StatusInternalServerError, err.Error())
+				return
+			}
+		}
+
+		if err := apiutil.CollectEscapeStringOption("start_key", input, collector); err != nil {
+			h.r.JSON(w, http.StatusInternalServerError, err.Error())
+			return
+		}
+
+		if err := apiutil.CollectEscapeStringOption("end_key", input, collector); err != nil {
+			h.r.JSON(w, http.StatusInternalServerError, err.Error())
+			return
+		}
 	case types.ScatterRangeScheduler:
 		if err := apiutil.CollectEscapeStringOption("start_key", input, collector); err != nil {
 			h.r.JSON(w, http.StatusInternalServerError, err.Error())

server/cluster/cluster_test.go

@@ -3209,6 +3209,20 @@ func TestAddScheduler(t *testing.T) {
 	re.NoError(err)
 	re.NoError(controller.AddScheduler(gls))
 
+	_, err = schedulers.CreateScheduler(types.BalanceKeyRangeScheduler, oc, storage.NewStorageWithMemoryBackend(), schedulers.ConfigSliceDecoder(types.BalanceKeyRangeScheduler, []string{}), controller.RemoveScheduler)
+	re.Error(err)
+
+	gls, err = schedulers.CreateScheduler(types.BalanceKeyRangeScheduler, oc, storage.NewStorageWithMemoryBackend(), schedulers.ConfigSliceDecoder(types.BalanceKeyRangeScheduler, []string{"learner", "tiflash", "1h", "100", "200"}), controller.RemoveScheduler)
+	re.NoError(err)
+	re.NoError(controller.AddScheduler(gls))
+	conf, err = gls.EncodeConfig()
+	re.NoError(err)
+	data = make(map[string]any)
+	re.NoError(json.Unmarshal(conf, &data))
+	re.Equal("learner", data["role"])
+	re.Equal("tiflash", data["engine"])
+	re.Equal(float64(time.Hour.Nanoseconds()), data["timeout"])
+
 	hb, err := schedulers.CreateScheduler(types.BalanceHotRegionScheduler, oc, storage.NewStorageWithMemoryBackend(), schedulers.ConfigJSONDecoder([]byte("{}")))
 	re.NoError(err)
 	conf, err = hb.EncodeConfig()