schedule: support placement-constrained hot scheduling

pkg/schedule/filter/filters.go

@@ -703,13 +703,17 @@ type ruleLeaderFitFilter struct {
 
 // newRuleLeaderFitFilter creates a filter that ensures after transfer leader with new store,
 // the isolation level will not decrease.
-func newRuleLeaderFitFilter(scope string, cluster *core.BasicCluster, ruleManager *placement.RuleManager, region *core.RegionInfo, srcLeaderStoreID uint64, allowMoveLeader bool) Filter {
+func newRuleLeaderFitFilter(scope string, cluster *core.BasicCluster, ruleManager *placement.RuleManager,
+	region *core.RegionInfo, oldFit *placement.RegionFit, srcLeaderStoreID uint64, allowMoveLeader bool) Filter {
+	if oldFit == nil {
+		oldFit = ruleManager.FitRegion(cluster, region)
+	}
 	return &ruleLeaderFitFilter{
 		scope:            scope,
 		cluster:          cluster,
 		ruleManager:      ruleManager,
 		region:           region,
-		oldFit:           ruleManager.FitRegion(cluster, region),
+		oldFit:           oldFit,
 		srcLeaderStoreID: srcLeaderStoreID,
 		allowMoveLeader:  allowMoveLeader,
 	}
@@ -816,9 +820,16 @@ func NewPlacementSafeguard(scope string, conf config.SharedConfigProvider, clust
 // NewPlacementLeaderSafeguard creates a filter that ensures after transfer a leader with
 // existed peer, the placement restriction will not become worse.
 // Note that it only worked when PlacementRules enabled otherwise it will always permit the sourceStore.
-func NewPlacementLeaderSafeguard(scope string, conf config.SharedConfigProvider, cluster *core.BasicCluster, ruleManager *placement.RuleManager, region *core.RegionInfo, sourceStore *core.StoreInfo, allowMoveLeader bool) Filter {
+func NewPlacementLeaderSafeguard(scope string, conf config.SharedConfigProvider, cluster *core.BasicCluster, ruleManager *placement.RuleManager,
+	region *core.RegionInfo, sourceStore *core.StoreInfo, allowMoveLeader bool) Filter {
+	return NewPlacementLeaderSafeguardWithFit(scope, conf, cluster, ruleManager, region, sourceStore, nil, allowMoveLeader)
+}
+
+// NewPlacementLeaderSafeguardWithFit creates a placement leader safeguard with a reusable region fit.
+func NewPlacementLeaderSafeguardWithFit(scope string, conf config.SharedConfigProvider, cluster *core.BasicCluster, ruleManager *placement.RuleManager,
+	region *core.RegionInfo, sourceStore *core.StoreInfo, oldFit *placement.RegionFit, allowMoveLeader bool) Filter {
 	if conf.IsPlacementRulesEnabled() {
-		return newRuleLeaderFitFilter(scope, cluster, ruleManager, region, sourceStore.GetID(), allowMoveLeader)
+		return newRuleLeaderFitFilter(scope, cluster, ruleManager, region, oldFit, sourceStore.GetID(), allowMoveLeader)
 	}
 	return nil
 }

pkg/schedule/filter/filters_test.go

@@ -144,12 +144,12 @@ func TestRuleFitFilter(t *testing.T) {
 		filter := newRuleFitFilter("", testCluster.GetBasicCluster(), testCluster.GetRuleManager(), region, nil, 1)
 		re.Equal(testCase.sourceRes, filter.Source(testCluster.GetSharedConfig(), testCluster.GetStore(testCase.storeID)).StatusCode)
 		re.Equal(testCase.targetRes, filter.Target(testCluster.GetSharedConfig(), testCluster.GetStore(testCase.storeID)).StatusCode)
-		leaderFilter := newRuleLeaderFitFilter("", testCluster.GetBasicCluster(), testCluster.GetRuleManager(), region, 1, true)
+		leaderFilter := newRuleLeaderFitFilter("", testCluster.GetBasicCluster(), testCluster.GetRuleManager(), region, nil, 1, true)
 		re.Equal(testCase.targetRes, leaderFilter.Target(testCluster.GetSharedConfig(), testCluster.GetStore(testCase.storeID)).StatusCode)
 	}
 
 	// store-6 is not exist in the peers, so it will not allow transferring leader to store 6.
-	leaderFilter := newRuleLeaderFitFilter("", testCluster.GetBasicCluster(), testCluster.GetRuleManager(), region, 1, false)
+	leaderFilter := newRuleLeaderFitFilter("", testCluster.GetBasicCluster(), testCluster.GetRuleManager(), region, nil, 1, false)
 	re.False(leaderFilter.Target(testCluster.GetSharedConfig(), testCluster.GetStore(6)).IsOK())
 }
 
@@ -223,7 +223,7 @@ func TestRuleFitFilterWithPlacementRule(t *testing.T) {
 		{StoreId: 5, Id: 4},
 		{StoreId: 6, Id: 5},
 	}}, &metapb.Peer{StoreId: 1, Id: 1})
-	leaderFilter := newRuleLeaderFitFilter("", testCluster.GetBasicCluster(), testCluster.GetRuleManager(), region, 1, true)
+	leaderFilter := newRuleLeaderFitFilter("", testCluster.GetBasicCluster(), testCluster.GetRuleManager(), region, nil, 1, true)
 	re.Equal(plan.StatusText(plan.StatusStoreNotMatchRule), leaderFilter.Target(testCluster.GetSharedConfig(), testCluster.GetStore(6)).String())
 }
 

pkg/schedule/placement/config.go

@@ -18,12 +18,15 @@ import (
 	"bytes"
 	"encoding/json"
 	"time"
+
+	"github.com/tikv/pd/pkg/core"
 )
 
 // ruleConfig contains rule and rule group configurations.
 type ruleConfig struct {
-	rules  map[[2]string]*Rule   // {group, id} => Rule
-	groups map[string]*RuleGroup // id => RuleGroup
+	rules                map[[2]string]*Rule   // {group, id} => Rule
+	groups               map[string]*RuleGroup // id => RuleGroup
+	mayRestrictStoreLoad [2]bool               // TiKV at index 0, TiFlash at index 1.
 }
 
 func newRuleConfig() *ruleConfig {
@@ -35,6 +38,7 @@ func newRuleConfig() *ruleConfig {
 
 // adjust configs for `buildRuleList` and API use.
 func (c *ruleConfig) adjust() {
+	c.mayRestrictStoreLoad = [2]bool{}
 	// remove all default group configurations.
 	// if there are rules belong to the group, it will be re-add later.
 	for id, g := range c.groups {
@@ -43,6 +47,10 @@ func (c *ruleConfig) adjust() {
 		}
 	}
 	for _, r := range c.rules {
+		if len(r.LabelConstraints) > 0 {
+			c.mayRestrictStoreLoad[0] = c.mayRestrictStoreLoad[0] || ruleMayRestrictEngine(r, "", core.EngineTiKV)
+			c.mayRestrictStoreLoad[1] = c.mayRestrictStoreLoad[1] || ruleMayRestrictEngine(r, core.EngineTiFlash)
+		}
 		g := c.groups[r.GroupID]
 		if g == nil {
 			// create default group configurations.
@@ -54,6 +62,41 @@ func (c *ruleConfig) adjust() {
 	}
 }
 
+// ruleMayRestrictEngine treats an empty value as a store without an engine label.
+// It reports restriction for non-engine constraints or a partial engine match.
+func ruleMayRestrictEngine(rule *Rule, values ...string) bool {
+	hasOtherConstraint := false
+	for i := range rule.LabelConstraints {
+		if rule.LabelConstraints[i].Key != core.EngineKey {
+			hasOtherConstraint = true
+			break
+		}
+	}
+
+	matched := 0
+	for _, value := range values {
+		hasEngineConstraint, matchesValue := false, true
+		for i := range rule.LabelConstraints {
+			constraint := &rule.LabelConstraints[i]
+			if constraint.Key != core.EngineKey {
+				continue
+			}
+			hasEngineConstraint = true
+			if !constraint.matchValue(value) {
+				matchesValue = false
+				break
+			}
+		}
+		if value != "" && !hasEngineConstraint {
+			matchesValue = false
+		}
+		if matchesValue {
+			matched++
+		}
+	}
+	return matched > 0 && (hasOtherConstraint || matched < len(values))
+}
+
 func (c *ruleConfig) getRule(key [2]string) *Rule {
 	return c.rules[key]
 }

pkg/schedule/placement/label_constraint.go

@@ -51,17 +51,19 @@ type LabelConstraint struct {
 
 // MatchStore checks if a store matches the constraint.
 func (c *LabelConstraint) MatchStore(store *core.StoreInfo) bool {
+	return c.matchValue(store.GetLabelValue(c.Key))
+}
+
+func (c *LabelConstraint) matchValue(value string) bool {
 	switch c.Op {
 	case In:
-		label := store.GetLabelValue(c.Key)
-		return label != "" && slice.AnyOf(c.Values, func(i int) bool { return c.Values[i] == label })
+		return value != "" && slice.AnyOf(c.Values, func(i int) bool { return c.Values[i] == value })
 	case NotIn:
-		label := store.GetLabelValue(c.Key)
-		return label == "" || slice.NoneOf(c.Values, func(i int) bool { return c.Values[i] == label })
+		return value == "" || slice.NoneOf(c.Values, func(i int) bool { return c.Values[i] == value })
 	case Exists:
-		return store.GetLabelValue(c.Key) != ""
+		return value != ""
 	case NotExists:
-		return store.GetLabelValue(c.Key) == ""
+		return value == ""
 	}
 	return false
 }

pkg/schedule/placement/rule_manager.go

@@ -368,6 +368,17 @@ func (m *RuleManager) GetRulesCount() int {
 	return len(m.ruleConfig.rules)
 }
 
+// MayRestrictStoreLoad returns whether placement rules may restrict the load
+// statistics population for the given store engine.
+func (m *RuleManager) MayRestrictStoreLoad(isTiKV bool) bool {
+	m.RLock()
+	defer m.RUnlock()
+	if isTiKV {
+		return m.ruleConfig.mayRestrictStoreLoad[0]
+	}
+	return m.ruleConfig.mayRestrictStoreLoad[1]
+}
+
 // GetGroupsCount returns the number of rule groups.
 func (m *RuleManager) GetGroupsCount() int {
 	m.RLock()

pkg/schedule/placement/rule_manager_test.go

@@ -78,6 +78,49 @@ func TestDefault2(t *testing.T) {
 	re.Equal([]string{"zone", "rack", "host"}, rules[1].LocationLabels)
 }
 
+func TestMayRestrictStoreLoad(t *testing.T) {
+	re := require.New(t)
+	_, manager := newTestManager(t, false)
+	re.False(manager.MayRestrictStoreLoad(true))
+	re.False(manager.MayRestrictStoreLoad(false))
+
+	setConstraints := func(constraints ...LabelConstraint) {
+		re.NoError(manager.SetRule(&Rule{
+			GroupID: "group", ID: "rule", Role: Voter, Count: 1,
+			LabelConstraints: constraints,
+		}))
+	}
+
+	setConstraints(LabelConstraint{Key: "zone", Op: In, Values: []string{"z1"}})
+	re.True(manager.MayRestrictStoreLoad(true))
+	re.False(manager.MayRestrictStoreLoad(false))
+
+	setConstraints(LabelConstraint{Key: core.EngineKey, Op: In, Values: []string{core.EngineTiFlash}})
+	re.False(manager.MayRestrictStoreLoad(true))
+	re.False(manager.MayRestrictStoreLoad(false))
+	setConstraints(LabelConstraint{Key: core.EngineKey, Op: NotIn, Values: []string{core.EngineTiFlash}})
+	re.False(manager.MayRestrictStoreLoad(true))
+	re.False(manager.MayRestrictStoreLoad(false))
+
+	setConstraints(
+		LabelConstraint{Key: core.EngineKey, Op: In, Values: []string{core.EngineTiFlash}},
+		LabelConstraint{Key: "zone", Op: In, Values: []string{"z1"}},
+	)
+	re.False(manager.MayRestrictStoreLoad(true))
+	re.True(manager.MayRestrictStoreLoad(false))
+
+	setConstraints()
+	re.False(manager.MayRestrictStoreLoad(true))
+	re.False(manager.MayRestrictStoreLoad(false))
+
+	re.NoError(manager.SetRule(&Rule{
+		GroupID: "group", ID: "restricted", Role: Voter, Count: 1,
+		LabelConstraints: []LabelConstraint{{Key: "zone", Op: In, Values: []string{"z1"}}},
+	}))
+	re.True(manager.MayRestrictStoreLoad(true))
+	re.False(manager.MayRestrictStoreLoad(false))
+}
+
 func TestAdjustRule(t *testing.T) {
 	re := require.New(t)
 	_, manager := newTestManager(t, false)

pkg/schedule/schedulers/hot_region.go

@@ -347,10 +347,21 @@ func (s *hotScheduler) tryAddPendingInfluence(op *operator.Operator, srcStore []
 	return true
 }
 
+func newBalanceReadSolvers(s *hotScheduler, cluster sche.SchedulerCluster) (leaderSolver, peerSolver *balanceSolver) {
+	// The read-peer population contains every TiKV store considered by the
+	// read-leader solver, so both solvers can share the placement precheck.
+	peerSolver = newBalanceSolver(s, cluster, utils.Read, movePeer)
+	placementState := &placementLoadState{
+		enabled:     peerSolver.placementV2Enabled,
+		canRestrict: peerSolver.placementCanRestrict,
+	}
+	leaderSolver = newBalanceSolverWithPlacementState(s, cluster, utils.Read, transferLeader, placementState)
+	return leaderSolver, peerSolver
+}
+
 func (s *hotScheduler) balanceHotReadRegions(cluster sche.SchedulerCluster) []*operator.Operator {
-	leaderSolver := newBalanceSolver(s, cluster, utils.Read, transferLeader)
+	leaderSolver, peerSolver := newBalanceReadSolvers(s, cluster)
 	leaderOps := leaderSolver.solve()
-	peerSolver := newBalanceSolver(s, cluster, utils.Read, movePeer)
 	peerOps := peerSolver.solve()
 	if len(leaderOps) == 0 && len(peerOps) == 0 {
 		return nil

pkg/schedule/schedulers/hot_region_rank_v2_test.go

@@ -20,8 +20,10 @@ import (
 	"github.com/docker/go-units"
 	"github.com/stretchr/testify/require"
 
+	"github.com/tikv/pd/pkg/core"
 	"github.com/tikv/pd/pkg/mock/mockcluster"
 	"github.com/tikv/pd/pkg/schedule/operator"
+	"github.com/tikv/pd/pkg/schedule/placement"
 	"github.com/tikv/pd/pkg/statistics/utils"
 	"github.com/tikv/pd/pkg/storage"
 	"github.com/tikv/pd/pkg/utils/operatorutil"
@@ -141,6 +143,87 @@ func TestHotWriteRegionScheduleWithRevertRegionsDimFirst(t *testing.T) {
 	clearPendingInfluence(hb)
 }
 
+func TestHotWriteRegionScheduleWithRevertRegionsAndPlacementRulesV2(t *testing.T) {
+	for _, testCase := range []struct {
+		name                  string
+		revertTargetMatchesB  bool
+		expectedOperatorCount int
+	}{
+		{name: "valid revert", revertTargetMatchesB: true, expectedOperatorCount: 2},
+		{name: "invalid revert", expectedOperatorCount: 1},
+	} {
+		t.Run(testCase.name, func(t *testing.T) {
+			re := require.New(t)
+			cancel, _, tc, oc := prepareSchedulersTest()
+			defer cancel()
+			tc.SetEnablePlacementRules(true)
+			sche, err := CreateScheduler(writeType, oc, storage.NewStorageWithMemoryBackend(), nil, nil)
+			re.NoError(err)
+			hb := sche.(*hotScheduler)
+			hb.types = []resourceType{writePeer}
+			hb.conf.setDstToleranceRatio(0.0)
+			hb.conf.setSrcToleranceRatio(0.0)
+			hb.conf.setRankFormulaVersion("v2")
+			hb.conf.setHistorySampleDuration(0)
+			hb.conf.WritePeerPriorities = []string{utils.BytePriority, utils.KeyPriority}
+			tc.SetClusterVersion(versioninfo.MinSupportedVersion(versioninfo.Version4_0))
+
+			for id, labels := range map[uint64]map[string]string{
+				1: {"b": "yes"},
+				2: {"a": "yes"},
+				3: {"b": "yes"},
+				4: {"a": "yes"},
+				5: {"a": "yes", "b": "yes"},
+			} {
+				if id == 2 && testCase.revertTargetMatchesB {
+					labels["b"] = "yes"
+				}
+				tc.AddLabelsStore(id, 20, labels)
+			}
+			tc.SetStoreBusy(4, true)
+			re.NoError(tc.SetRule(&placement.Rule{
+				GroupID: placement.DefaultGroupID, ID: placement.DefaultRuleID,
+				Role: placement.Voter, Count: 1,
+				LabelConstraints: []placement.LabelConstraint{{Key: "a", Op: placement.In, Values: []string{"yes"}}},
+			}))
+			re.NoError(tc.SetRule(&placement.Rule{
+				GroupID: placement.DefaultGroupID, ID: "voter-b",
+				Role: placement.Voter, Count: 2,
+				LabelConstraints: []placement.LabelConstraint{{Key: "b", Op: placement.In, Values: []string{"yes"}}},
+			}))
+
+			tc.UpdateStorageWrittenStats(1, 15*units.MiB*utils.StoreHeartBeatReportInterval, 15*units.MiB*utils.StoreHeartBeatReportInterval)
+			tc.UpdateStorageWrittenStats(2, 20*units.MiB*utils.StoreHeartBeatReportInterval, 14*units.MiB*utils.StoreHeartBeatReportInterval)
+			tc.UpdateStorageWrittenStats(3, 15*units.MiB*utils.StoreHeartBeatReportInterval, 15*units.MiB*utils.StoreHeartBeatReportInterval)
+			tc.UpdateStorageWrittenStats(4, 15*units.MiB*utils.StoreHeartBeatReportInterval, 15*units.MiB*utils.StoreHeartBeatReportInterval)
+			tc.UpdateStorageWrittenStats(5, 10*units.MiB*utils.StoreHeartBeatReportInterval, 16*units.MiB*utils.StoreHeartBeatReportInterval)
+			addRegionInfo(tc, utils.Write, []testRegionInfo{
+				{6, []uint64{3, 2, 1}, 3 * units.MiB, 1.8 * units.MiB, 0},
+				{7, []uint64{1, 4, 5}, 0.1 * units.MiB, 2 * units.MiB, 0},
+			})
+			mainRegion := tc.GetRegion(6)
+			sourcePeers := append(mainRegion.GetPeers()[:0:0], mainRegion.GetStorePeer(2))
+			for _, item := range tc.ExpiredWriteItems(mainRegion.Clone(core.SetPeers(sourcePeers))) {
+				tc.Update(item, utils.Write)
+			}
+
+			ops, _ := hb.Schedule(tc, false)
+			re.Len(ops, 1)
+			re.True(hb.searchRevertRegions[writePeer])
+			clearPendingInfluence(hb)
+
+			ops, _ = hb.Schedule(tc, false)
+			re.Len(ops, testCase.expectedOperatorCount)
+			operatorutil.CheckTransferPeer(re, ops[0], operator.OpHotRegion, 2, 5)
+			if testCase.revertTargetMatchesB {
+				operatorutil.CheckTransferPeer(re, ops[1], operator.OpHotRegion, 5, 2)
+			}
+			re.True(hb.searchRevertRegions[writePeer])
+			clearPendingInfluence(hb)
+		})
+	}
+}
+
 func TestHotWriteRegionScheduleWithRevertRegionsDimFirstOnly(t *testing.T) {
 	// This is a test that searchRevertRegions finds a solution of rank 2.
 	re := require.New(t)