schedule: fix location label in placement rule not work well with rul…

pkg/schedule/checker/replica_strategy.go

@@ -93,6 +93,9 @@ func (s *ReplicaStrategy) SelectStoreToAdd(coLocationStores []*core.StoreInfo, e
 // SelectStoreToFix returns a store to replace down/offline old peer. The location
 // placement after scheduling is allowed to be worse than original.
 func (s *ReplicaStrategy) SelectStoreToFix(coLocationStores []*core.StoreInfo, old uint64) (uint64, bool) {
+ if len(coLocationStores) == 0 {
+ return 0, false
+ }
  // trick to avoid creating a slice with `old` removed.
  s.swapStoreToFirst(coLocationStores, old)
  return s.SelectStoreToAdd(coLocationStores[1:])
@@ -101,6 +104,9 @@ func (s *ReplicaStrategy) SelectStoreToFix(coLocationStores []*core.StoreInfo, o
 // SelectStoreToImprove returns a store to replace oldStore. The location
 // placement after scheduling should be better than original.
 func (s *ReplicaStrategy) SelectStoreToImprove(coLocationStores []*core.StoreInfo, old uint64) (uint64, bool) {
+ if len(coLocationStores) == 0 {
+ return 0, false
+ }
  // trick to avoid creating a slice with `old` removed.
  s.swapStoreToFirst(coLocationStores, old)
  oldStore := s.cluster.GetStore(old)

pkg/schedule/checker/rule_checker.go

@@ -391,19 +391,26 @@ func (c *RuleChecker) allowLeader(fit *placement.RegionFit, peer *metapb.Peer) b
 }
 
 func (c *RuleChecker) fixBetterLocation(region *core.RegionInfo, rf *placement.RuleFit) (*operator.Operator, error) {
- if len(rf.Rule.LocationLabels) == 0 || rf.Rule.Count <= 1 {
+ if len(rf.Rule.LocationLabels) == 0 {
  return nil, nil
  }
 
  isWitness := rf.Rule.IsWitness && c.isWitnessEnabled()
  // If the peer to be moved is a witness, since no snapshot is needed, we also reuse the fast failover logic.
  strategy := c.strategy(region, rf.Rule, isWitness)
- ruleStores := c.getRuleFitStores(rf)
- oldStore := strategy.SelectStoreToRemove(ruleStores)
+ regionStores := c.cluster.GetRegionStores(region)
+ oldStore := strategy.SelectStoreToRemove(regionStores)
  if oldStore == 0 {
  return nil, nil
  }
- newStore, filterByTempState := strategy.SelectStoreToImprove(ruleStores, oldStore)
+ var coLocationStores []*core.StoreInfo
+ for _, s := range regionStores {
+ if placement.MatchLabelConstraints(s, rf.Rule.LabelConstraints) {
+ coLocationStores = append(coLocationStores, s)
+ }
+ }
+
+ newStore, filterByTempState := strategy.SelectStoreToImprove(coLocationStores, oldStore)
  if newStore == 0 {
  log.Debug("no replacement store", zap.Uint64("region-id", region.GetID()))
  c.handleFilterState(region, filterByTempState)

pkg/schedule/checker/rule_checker_test.go

@@ -1332,3 +1332,62 @@ func (suite *ruleCheckerTestSuite) TestPendingList() {
  _, exist = suite.rc.pendingList.Get(1)
  suite.False(exist)
 }
+
+func (suite *ruleCheckerTestSuite) TestLocationLabels() {
+ suite.cluster.AddLabelsStore(1, 1, map[string]string{"zone": "z1", "rack": "r1", "host": "h1"})
+ suite.cluster.AddLabelsStore(2, 1, map[string]string{"zone": "z1", "rack": "r1", "host": "h1"})
+ suite.cluster.AddLabelsStore(3, 1, map[string]string{"zone": "z1", "rack": "r2", "host": "h1"})
+ suite.cluster.AddLabelsStore(4, 1, map[string]string{"zone": "z1", "rack": "r2", "host": "h1"})
+ suite.cluster.AddLabelsStore(5, 1, map[string]string{"zone": "z2", "rack": "r3", "host": "h2"})
+ suite.cluster.AddLabelsStore(6, 1, map[string]string{"zone": "z2", "rack": "r3", "host": "h2"})
+ suite.cluster.AddLeaderRegionWithRange(1, "", "", 1, 2, 5)
+ rule1 := &placement.Rule{
+ GroupID: "pd",
+ ID: "test1",
+ Role: placement.Leader,
+ Count: 1,
+ LabelConstraints: []placement.LabelConstraint{
+ {
+ Key: "zone",
+ Op: placement.In,
+ Values: []string{"z1"},
+ },
+ },
+ LocationLabels: []string{"rack"},
+ }
+ rule2 := &placement.Rule{
+ GroupID: "pd",
+ ID: "test2",
+ Role: placement.Voter,
+ Count: 1,
+ LabelConstraints: []placement.LabelConstraint{
+ {
+ Key: "zone",
+ Op: placement.In,
+ Values: []string{"z1"},
+ },
+ },
+ LocationLabels: []string{"rack"},
+ }
+ rule3 := &placement.Rule{
+ GroupID: "pd",
+ ID: "test3",
+ Role: placement.Voter,
+ Count: 1,
+ LabelConstraints: []placement.LabelConstraint{
+ {
+ Key: "zone",
+ Op: placement.In,
+ Values: []string{"z2"},
+ },
+ },
+ LocationLabels: []string{"rack"},
+ }
+ suite.ruleManager.SetRule(rule1)
+ suite.ruleManager.SetRule(rule2)
+ suite.ruleManager.SetRule(rule3)
+ suite.ruleManager.DeleteRule("pd", "default")
+ op := suite.rc.Check(suite.cluster.GetRegion(1))
+ suite.NotNil(op)
+ suite.Equal("move-to-better-location", op.Desc())
+}