diff --git a/pkg/schedule/checker/replica_strategy.go b/pkg/schedule/checker/replica_strategy.go index 27e6301c3dc..797d8bb2853 100644 --- a/pkg/schedule/checker/replica_strategy.go +++ b/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) diff --git a/pkg/schedule/checker/rule_checker.go b/pkg/schedule/checker/rule_checker.go index aad87ec54b3..d405941300e 100644 --- a/pkg/schedule/checker/rule_checker.go +++ b/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) diff --git a/pkg/schedule/checker/rule_checker_test.go b/pkg/schedule/checker/rule_checker_test.go index c204345faa3..ebe1d1aadaf 100644 --- a/pkg/schedule/checker/rule_checker_test.go +++ b/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()) +}