checker: fix the conflict between tiflash learner and location labels (…

…#6660) close #6662 Signed-off-by: Ryan Leung <[email protected]>
tikv · Jun 25, 2023 · f4d774a · f4d774a
1 parent ac31f87
commit f4d774a
Show file tree

Hide file tree

Showing 2 changed files with 35 additions and 2 deletions.
diff --git a/pkg/schedule/checker/rule_checker.go b/pkg/schedule/checker/rule_checker.go
@@ -398,12 +398,13 @@ func (c *RuleChecker) fixBetterLocation(region *core.RegionInfo, rf *placement.R
 	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)
-	regionStores := c.cluster.GetRegionStores(region)
-	oldStore := strategy.SelectStoreToRemove(regionStores)
+	ruleStores := c.getRuleFitStores(rf)
+	oldStore := strategy.SelectStoreToRemove(ruleStores)
 	if oldStore == 0 {
 		return nil, nil
 	}
 	var coLocationStores []*core.StoreInfo
+	regionStores := c.cluster.GetRegionStores(region)
 	for _, s := range regionStores {
 		if placement.MatchLabelConstraints(s, rf.Rule.LabelConstraints) {
 			coLocationStores = append(coLocationStores, s)

diff --git a/pkg/schedule/checker/rule_checker_test.go b/pkg/schedule/checker/rule_checker_test.go
@@ -1391,3 +1391,35 @@ func (suite *ruleCheckerTestSuite) TestLocationLabels() {
 	suite.NotNil(op)
 	suite.Equal("move-to-better-location", op.Desc())
 }
+
+func (suite *ruleCheckerTestSuite) TestTiFlashLocationLabels() {
+	suite.cluster.SetEnableUseJointConsensus(true)
+	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.AddLabelsStore(7, 1, map[string]string{"engine": "tiflash"})
+	suite.cluster.AddRegionWithLearner(1, 1, []uint64{3, 5}, []uint64{7})
+
+	rule1 := &placement.Rule{
+		GroupID: "tiflash",
+		ID:      "test1",
+		Role:    placement.Learner,
+		Count:   1,
+		LabelConstraints: []placement.LabelConstraint{
+			{
+				Key:    "engine",
+				Op:     placement.In,
+				Values: []string{"tiflash"},
+			},
+		},
+	}
+	suite.ruleManager.SetRule(rule1)
+	rule := suite.ruleManager.GetRule("pd", "default")
+	rule.LocationLabels = []string{"zone", "rack", "host"}
+	suite.ruleManager.SetRule(rule)
+	op := suite.rc.Check(suite.cluster.GetRegion(1))
+	suite.Nil(op)
+}