schedule: impl balance range scheduler

pkg/core/basic_cluster.go

@@ -16,6 +16,7 @@ package core
 
 import (
 	"bytes"
+	"encoding/json"
 
 	"github.com/tikv/pd/pkg/core/constant"
 )
@@ -156,6 +157,15 @@ type KeyRange struct {
 	EndKey   []byte `json:"end-key"`
 }
 
+// MarshalJSON marshals to json.
+func (kr KeyRange) MarshalJSON() ([]byte, error) {
+	m := map[string]string{
+		"start-key": HexRegionKeyStr(kr.StartKey),
+		"end-key":   HexRegionKeyStr(kr.EndKey),
+	}
+	return json.Marshal(m)
+}
+
 // NewKeyRange create a KeyRange with the given start key and end key.
 func NewKeyRange(startKey, endKey string) KeyRange {
 	return KeyRange{
@@ -169,6 +179,17 @@ type KeyRanges struct {
 	krs []*KeyRange
 }
 
+// NewKeyRanges creates a KeyRanges.
+func NewKeyRanges(ranges []KeyRange) *KeyRanges {
+	krs := make([]*KeyRange, 0, len(ranges))
+	for _, kr := range ranges {
+		krs = append(krs, &kr)
+	}
+	return &KeyRanges{
+		krs,
+	}
+}
+
 // NewKeyRangesWithSize creates a KeyRanges with the hint size.
 func NewKeyRangesWithSize(size int) *KeyRanges {
 	return &KeyRanges{

pkg/core/constant/kind.go

@@ -66,7 +66,6 @@ const (
 	RegionKind
 	// WitnessKind indicates the witness kind resource
 	WitnessKind
-
 	// ResourceKindLen represents the ResourceKind count
 	ResourceKindLen
 )

pkg/mcs/scheduling/server/cluster.go

@@ -314,7 +314,11 @@ func (c *Cluster) updateScheduler() {
 		)
 		// 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/filter/filters.go

@@ -932,6 +932,8 @@ var (
 	allSpecialEngines = []string{core.EngineTiFlash}
 	// NotSpecialEngines is used to filter the special engine.
 	NotSpecialEngines = placement.LabelConstraint{Key: core.EngineKey, Op: placement.NotIn, Values: allSpecialEngines}
+	// TiFlashEngineConstraint is used to filter the TiFlash engine.
+	TiFlashEngineConstraint = placement.LabelConstraint{Key: core.EngineKey, Op: placement.In, Values: allSpecialEngines}
 )
 
 type isolationFilter struct {

pkg/schedule/operator/operator_controller.go

@@ -15,6 +15,7 @@
 package operator
 
 import (
+	"bytes"
 	"context"
 	"fmt"
 	"strconv"
@@ -828,14 +829,33 @@ func (oc *Controller) GetHistory(start time.Time) []OpHistory {
 	return history
 }
 
+// OpInfluenceOption is used to filter the region.
+// returns true if the region meets the condition, it will ignore this region in the influence calculation.
+// returns false if the region does not meet the condition, it will calculate the influence of this region.
+type OpInfluenceOption func(region *core.RegionInfo) bool
+
+// WithRangeOption returns an OpInfluenceOption that filters the region by the label.
+func WithRangeOption(ranges []core.KeyRange) OpInfluenceOption {
+	return func(region *core.RegionInfo) bool {
+		for _, r := range ranges {
+			// the start key of the region must greater than the given range start key.
+			// the end key of the region must less than the given range end key.
+			if bytes.Compare(region.GetStartKey(), r.StartKey) < 0 || bytes.Compare(r.EndKey, region.GetEndKey()) < 0 {
+				return false
+			}
+		}
+		return true
+	}
+}
+
 // OperatorCount gets the count of operators filtered by kind.
 // kind only has one OpKind.
 func (oc *Controller) OperatorCount(kind OpKind) uint64 {
 	return oc.counts.getCountByKind(kind)
 }
 
 // GetOpInfluence gets OpInfluence.
-func (oc *Controller) GetOpInfluence(cluster *core.BasicCluster) OpInfluence {
+func (oc *Controller) GetOpInfluence(cluster *core.BasicCluster, ops ...OpInfluenceOption) OpInfluence {
 	influence := OpInfluence{
 		StoresInfluence: make(map[uint64]*StoreInfluence),
 	}
@@ -844,6 +864,11 @@ func (oc *Controller) GetOpInfluence(cluster *core.BasicCluster) OpInfluence {
 			op := value.(*Operator)
 			if !op.CheckTimeout() && !op.CheckSuccess() {
 				region := cluster.GetRegion(op.RegionID())
+				for _, opt := range ops {
+					if !opt(region) {
+						return true
+					}
+				}
 				if region != nil {
 					op.UnfinishedInfluence(influence, region)
 				}