Fix balance-leader won't work when placement rule enabled (#2725)

pkg/mock/mockcluster/mockcluster.go

@@ -617,3 +617,14 @@ func (mc *Cluster) MockRegionInfo(regionID uint64, leaderStoreID uint64,
 	}
 	return core.NewRegionInfo(region, leader)
 }
+
+// SetStoreLabel set the labels to the target store
+func (mc *Cluster) SetStoreLabel(storeID uint64, labels map[string]string) {
+	store := mc.GetStore(storeID)
+	newLabels := make([]*metapb.StoreLabel, 0, len(labels))
+	for k, v := range labels {
+		newLabels = append(newLabels, &metapb.StoreLabel{Key: k, Value: v})
+	}
+	newStore := store.Clone(core.SetStoreLabels(newLabels))
+	mc.PutStore(newStore)
+}

server/schedule/filter/filters.go

@@ -16,11 +16,13 @@ package filter
 import (
 	"fmt"
 
+	"github.com/pingcap/log"
 	"github.com/pingcap/pd/v4/pkg/slice"
 	"github.com/pingcap/pd/v4/server/core"
 	"github.com/pingcap/pd/v4/server/schedule/opt"
 	"github.com/pingcap/pd/v4/server/schedule/placement"
 	"github.com/pingcap/pd/v4/server/schedule/storelimit"
+	"go.uber.org/zap"
 )
 
 // revive:disable:unused-parameter
@@ -576,9 +578,12 @@ func (f *ruleLeaderFitFilter) Source(opt opt.Options, store *core.StoreInfo) boo
 }
 
 func (f *ruleLeaderFitFilter) Target(opt opt.Options, store *core.StoreInfo) bool {
-	oldLeaderPeerReplica := *f.region.GetLeader()
-	oldLeaderPeerReplica.StoreId = store.GetID()
-	region := f.region.Clone(core.WithReplacePeerStore(f.oldLeaderStoreID, store.GetID()), core.WithLeader(&oldLeaderPeerReplica))
+	targetPeer := f.region.GetStorePeer(store.GetID())
+	if targetPeer == nil {
+		log.Warn("ruleLeaderFitFilter couldn't find peer on target Store", zap.Uint64("target-store", store.GetID()))
+		return false
+	}
+	region := f.region.Clone(core.WithLeader(targetPeer))
 	newFit := f.fitter.FitRegion(region)
 	return placement.CompareRegionFit(f.oldFit, newFit) <= 0
 }

server/schedulers/scheduler_test.go

@@ -545,42 +545,97 @@ func (s *testBalanceLeaderSchedulerWithRuleEnabledSuite) schedule() []*operator.
 }
 
 func (s *testBalanceLeaderSchedulerWithRuleEnabledSuite) TestBalanceLeaderWithConflictRule(c *C) {
-	rule := placement.Rule{
-		GroupID:  "test",
-		ID:       "1",
-		Index:    1,
-		StartKey: []byte(""),
-		EndKey:   []byte(""),
-		Role:     placement.Leader,
-		Count:    1,
-		LabelConstraints: []placement.LabelConstraint{
-			{
-				Key:    "role",
-				Op:     placement.In,
-				Values: []string{"leader"},
-			},
-		},
-		LocationLabels: []string{"host"},
-	}
-	c.Check(s.tc.SetRule(&rule), IsNil)
-	c.Check(s.tc.DeleteRule("pd", "default"), IsNil)
-
-	// Stores:     1    2    3    4
-	// Leaders:    1    0    0    0
-	// Region1:    L    F    F    F
+	// Stores:     1    2    3
+	// Leaders:    1    0    0
+	// Region1:    L    F    F
 	s.tc.AddLeaderStore(1, 1)
 	s.tc.AddLeaderStore(2, 0)
 	s.tc.AddLeaderStore(3, 0)
-	s.tc.AddLeaderStore(4, 0)
-	s.tc.AddLeaderRegion(1, 1, 2, 3, 4)
-	s.tc.AddLabelsStore(1, 1, map[string]string{
-		"role": "leader",
+	s.tc.AddLeaderRegion(1, 1, 2, 3)
+	s.tc.SetStoreLabel(1, map[string]string{
+		"host": "a",
+	})
+	s.tc.SetStoreLabel(2, map[string]string{
+		"host": "b",
+	})
+	s.tc.SetStoreLabel(3, map[string]string{
+		"host": "c",
 	})
-	c.Check(s.schedule(), IsNil)
 
-	// Stores:     1    2    3    4
-	// Leaders:    16   0    0    0
-	// Region1:    L    F    F    F
+	// Stores:     1    2    3
+	// Leaders:    16   0    0
+	// Region1:    L    F    F
 	s.tc.UpdateLeaderCount(1, 16)
-	c.Check(s.schedule(), IsNil)
+	testcases := []struct {
+		name     string
+		rule     placement.Rule
+		schedule bool
+	}{
+		{
+			name: "default Rule",
+			rule: placement.Rule{
+				GroupID:        "pd",
+				ID:             "default",
+				Index:          1,
+				StartKey:       []byte(""),
+				EndKey:         []byte(""),
+				Role:           placement.Voter,
+				Count:          3,
+				LocationLabels: []string{"host"},
+			},
+			schedule: true,
+		},
+		{
+			name: "single store allowed to be placed leader",
+			rule: placement.Rule{
+				GroupID:  "pd",
+				ID:       "default",
+				Index:    1,
+				StartKey: []byte(""),
+				EndKey:   []byte(""),
+				Role:     placement.Leader,
+				Count:    1,
+				LabelConstraints: []placement.LabelConstraint{
+					{
+						Key:    "host",
+						Op:     placement.In,
+						Values: []string{"a"},
+					},
+				},
+				LocationLabels: []string{"host"},
+			},
+			schedule: false,
+		},
+		{
+			name: "2 store allowed to be placed leader",
+			rule: placement.Rule{
+				GroupID:  "pd",
+				ID:       "default",
+				Index:    1,
+				StartKey: []byte(""),
+				EndKey:   []byte(""),
+				Role:     placement.Leader,
+				Count:    1,
+				LabelConstraints: []placement.LabelConstraint{
+					{
+						Key:    "host",
+						Op:     placement.In,
+						Values: []string{"a", "b"},
+					},
+				},
+				LocationLabels: []string{"host"},
+			},
+			schedule: true,
+		},
+	}
+
+	for _, testcase := range testcases {
+		c.Logf(testcase.name)
+		c.Check(s.tc.SetRule(&testcase.rule), IsNil)
+		if testcase.schedule {
+			c.Check(len(s.schedule()), Equals, 1)
+		} else {
+			c.Check(s.schedule(), IsNil)
+		}
+	}
 }