scheduler: Introduce BalaceKeyrange scheduler

pkg/mock/mockcluster/mockcluster.go

@@ -368,8 +368,8 @@ func (mc *Cluster) AddRegionStoreWithLeader(storeID uint64, regionCount int, lea
 	}
 }
 
-// AddLabelsStore adds store with specified count of region and labels.
-func (mc *Cluster) AddLabelsStore(storeID uint64, regionCount int, labels map[string]string) {
+// AddLabelsStoreWithLimit adds store with specified count of region and labels and store limit.
+func (mc *Cluster) AddLabelsStoreWithLimit(storeID uint64, regionCount int, labels map[string]string, limit float64) {
 	newLabels := make([]*metapb.StoreLabel, 0, len(labels))
 	for k, v := range labels {
 		newLabels = append(newLabels, &metapb.StoreLabel{Key: k, Value: v})
@@ -388,11 +388,16 @@ func (mc *Cluster) AddLabelsStore(storeID uint64, regionCount int, labels map[st
 		core.SetRegionSize(int64(regionCount)*defaultRegionSize/units.MiB),
 		core.SetLastHeartbeatTS(time.Now()),
 	)
-	mc.SetStoreLimit(storeID, storelimit.AddPeer, 60)
-	mc.SetStoreLimit(storeID, storelimit.RemovePeer, 60)
+	mc.SetStoreLimit(storeID, storelimit.AddPeer, limit)
+	mc.SetStoreLimit(storeID, storelimit.RemovePeer, limit)
 	mc.PutStore(store)
 }
 
+// AddLabelsStore adds store with specified count of region and labels.
+func (mc *Cluster) AddLabelsStore(storeID uint64, regionCount int, labels map[string]string) {
+	mc.AddLabelsStoreWithLimit(storeID, regionCount, labels, 60)
+}
+
 // AddLabersStoreWithLearnerCount adds store with specified count of region, learner and labels.
 func (mc *Cluster) AddLabersStoreWithLearnerCount(storeID uint64, regionCount int, learnerCount int, labels map[string]string) {
 	mc.AddLabelsStore(storeID, regionCount, labels)

pkg/schedule/handler/handler.go

@@ -45,6 +45,7 @@ import (
 	"github.com/tikv/pd/pkg/statistics"
 	"github.com/tikv/pd/pkg/statistics/buckets"
 	"github.com/tikv/pd/pkg/statistics/utils"
+	"github.com/tikv/pd/pkg/storage"
 	"github.com/tikv/pd/pkg/utils/typeutil"
 )
 
@@ -1303,6 +1304,60 @@ func (h *Handler) CheckRegionsReplicated(rawStartKey, rawEndKey string) (string,
 	return state, nil
 }
 
+// BalanceKeyrange create a new balance-keyrange scheduler according to API request.
+func (h *Handler) BalanceKeyrange(data string) (string, error) {
+	sc, err := h.GetSchedulersController()
+	if err != nil {
+		return "", err
+	}
+	exist, _ := sc.IsSchedulerExisted(types.BalanceKeyrangeScheduler.String())
+	if exist {
+		return "Already existed", errs.ErrSchedulerExisted.FastGenByArgs()
+	}
+	oc := h.GetCoordinator().GetOperatorController()
+	controller := h.GetCoordinator().GetSchedulersController()
+	cb := func(s string) error {
+		return controller.RemoveScheduler(s)
+	}
+	s, err := schedulers.CreateScheduler(types.BalanceKeyrangeScheduler, oc, storage.NewStorageWithMemoryBackend(), schedulers.ConfigSliceDecoder(types.BalanceKeyrangeScheduler, []string{data}), cb)
+	if err != nil {
+		return "Created scheduler failed", err
+	}
+	if err = controller.AddScheduler(s); err != nil {
+		return "Add scheduler failed", err
+	}
+	return "Scheduler added successfully", nil
+}
+
+// CheckBalanceKeyrangeStatus returns the status of the balance-keyrange scheduler.
+func (h *Handler) CheckBalanceKeyrangeStatus() (any, error) {
+	sc, err := h.GetSchedulersController()
+	makeJSONResp := func(s string) any {
+		return struct {
+			Scheduling bool   `json:"scheduling"`
+			ErrMsg     string `json:"err_msg"`
+		}{
+			Scheduling: false,
+			ErrMsg:     s,
+		}
+	}
+	if err != nil {
+		return makeJSONResp("get scheduler control error"), err
+	}
+	s := sc.GetScheduler(types.BalanceKeyrangeScheduler.String())
+	if s == nil {
+		return makeJSONResp("no scheduler found"), nil
+	}
+	if s.IsDisable() {
+		return makeJSONResp("scheduler disabled"), nil
+	}
+	st := sc.GetSchedulerStatus(types.BalanceKeyrangeScheduler.String())
+	if st == nil {
+		return makeJSONResp("can't get scheduler status"), nil
+	}
+	return st, nil
+}
+
 // GetRuleManager returns the rule manager.
 func (h *Handler) GetRuleManager() (*placement.RuleManager, error) {
 	c := h.GetCluster()

pkg/schedule/operator/kind.go

@@ -49,6 +49,8 @@ const (
 	OpWitnessLeader
 	// Include witness transfer.
 	OpWitness
+	// Include keyrange balancer
+	OpKeyrange
 	opMax
 )
 

pkg/schedule/operator/operator.go

@@ -24,6 +24,7 @@ import (
 
 	"github.com/prometheus/client_golang/prometheus"
 
+	"github.com/pingcap/failpoint"
 	"github.com/pingcap/kvproto/pkg/metapb"
 
 	"github.com/tikv/pd/pkg/core"
@@ -292,6 +293,9 @@ func (o *Operator) IsEnd() bool {
 
 // CheckSuccess checks if all steps are finished, and update the status.
 func (o *Operator) CheckSuccess() bool {
+	failpoint.Inject("forceSucess", func() {
+		atomic.StoreInt32(&o.currentStep, int32(len(o.steps)))
+	})
 	if atomic.LoadInt32(&o.currentStep) >= int32(len(o.steps)) {
 		return o.status.To(SUCCESS) || o.Status() == SUCCESS
 	}

pkg/schedule/operator/operator_controller.go

@@ -418,6 +418,16 @@ func (oc *Controller) PromoteWaitingOperator() {
 	}
 }
 
+// GetWopCount returns the count in the waiting list of a kind of scheduler.
+func (oc *Controller) GetWopCount(kind string) uint64 {
+	return oc.wopStatus.getCount(kind)
+}
+
+// GetSchedulerMaxWaitingOperator returns the waiting queue side.
+func (oc *Controller) GetSchedulerMaxWaitingOperator() uint64 {
+	return oc.config.GetSchedulerMaxWaitingOperator()
+}
+
 // checkAddOperator checks if the operator can be added.
 // There are several situations that cannot be added:
 // - There is no such region in the cluster
@@ -948,7 +958,7 @@ func (o *records) Put(op *Operator) {
 	o.ttl.Put(id, record)
 }
 
-// ExceedStoreLimit returns true if the store exceeds the cost limit after adding the  Otherwise, returns false.
+// ExceedStoreLimit returns true if the store exceeds the cost limit after adding the opertor. Otherwise, returns false.
 func (oc *Controller) ExceedStoreLimit(ops ...*Operator) bool {
 	// The operator with Urgent priority, like admin operators, should ignore the store limit check.
 	var desc string