filters: add rule filters (#1917)

Signed-off-by: disksing <i@disksing.com>

server/core/region_option.go

@@ -246,3 +246,14 @@ func WithPromoteLearner(peerID uint64) RegionCreateOption {
 		}
 	}
 }
+
+// WithReplacePeerStore replaces a peer's storeID with another ID.
+func WithReplacePeerStore(peerID, newStoreID uint64) RegionCreateOption {
+	return func(region *RegionInfo) {
+		for _, p := range region.GetPeers() {
+			if p.GetId() == peerID {
+				p.StoreId = newStoreID
+			}
+		}
+	}
+}

server/schedule/filter/filters.go

@@ -20,6 +20,7 @@ import (
 	"github.com/pingcap/pd/pkg/slice"
 	"github.com/pingcap/pd/server/core"
 	"github.com/pingcap/pd/server/schedule/opt"
+	"github.com/pingcap/pd/server/schedule/placement"
 )
 
 //revive:disable:unused-parameter
@@ -492,3 +493,78 @@ func (f *BlacklistStoreFilter) filter(store *core.StoreInfo) bool {
 	_, ok := f.blacklist[store.GetID()]
 	return ok
 }
+
+// labelConstraintFilter is a filter that selects stores satisfy the constraints.
+type labelConstraintFilter struct {
+	scope       string
+	constraints []placement.LabelConstraint
+}
+
+// NewLabelConstaintFilter creates a filter that selects stores satisfy the constraints.
+func NewLabelConstaintFilter(scope string, constraints []placement.LabelConstraint) Filter {
+	return labelConstraintFilter{scope: scope, constraints: constraints}
+}
+
+// Scope returns the scheduler or the checker which the filter acts on.
+func (f labelConstraintFilter) Scope() string {
+	return f.scope
+}
+
+// Type returns the name of the filter.
+func (f labelConstraintFilter) Type() string {
+	return "label-constraint-filter"
+}
+
+// Source filters stores when select them as schedule source.
+func (f labelConstraintFilter) Source(opt opt.Options, store *core.StoreInfo) bool {
+	return !placement.MatchLabelConstraints(store, f.constraints)
+}
+
+// Target filters stores when select them as schedule target.
+func (f labelConstraintFilter) Target(opt opt.Options, store *core.StoreInfo) bool {
+	return !placement.MatchLabelConstraints(store, f.constraints)
+}
+
+// RegionFitter is the interface that can fit a region against placement rules.
+type RegionFitter interface {
+	FitRegion(*core.RegionInfo) *placement.RegionFit
+}
+
+type ruleFitFilter struct {
+	scope   string
+	fitter  RegionFitter
+	region  *core.RegionInfo
+	oldFit  *placement.RegionFit
+	oldPeer uint64
+}
+
+// NewRuleFitFilter creates a filter that ensures after replace a peer with new
+// one, the isolation level will not decrease. Its function is the same as
+// distinctScoreFilter but used when placement rules is enabled.
+func NewRuleFitFilter(scope string, fitter RegionFitter, region *core.RegionInfo, oldPeerID uint64) Filter {
+	return &ruleFitFilter{
+		scope:   scope,
+		fitter:  fitter,
+		region:  region,
+		oldFit:  fitter.FitRegion(region),
+		oldPeer: oldPeerID,
+	}
+}
+
+func (f *ruleFitFilter) Scope() string {
+	return f.scope
+}
+
+func (f *ruleFitFilter) Type() string {
+	return "rule-fit-filter"
+}
+
+func (f *ruleFitFilter) Source(opt opt.Options, store *core.StoreInfo) bool {
+	return false
+}
+
+func (f *ruleFitFilter) Target(opt opt.Options, store *core.StoreInfo) bool {
+	region := f.region.Clone(core.WithReplacePeerStore(f.oldPeer, store.GetID()))
+	newFit := f.fitter.FitRegion(region)
+	return placement.CompareRegionFit(f.oldFit, newFit) > 0
+}

server/schedule/filter/filters_test.go

@@ -20,6 +20,7 @@ import (
 	"github.com/pingcap/pd/pkg/mock/mockcluster"
 	"github.com/pingcap/pd/pkg/mock/mockoption"
 	"github.com/pingcap/pd/server/core"
+	"github.com/pingcap/pd/server/schedule/placement"
 )
 
 func Test(t *testing.T) {
@@ -44,3 +45,38 @@ func (s *testFiltersSuite) TestPendingPeerFilter(c *C) {
 	c.Assert(filter.Source(tc, newStore), IsFalse)
 	c.Assert(filter.Target(tc, newStore), IsFalse)
 }
+
+func (s *testFiltersSuite) TestLabelConstraintsFilter(c *C) {
+	opt := mockoption.NewScheduleOptions()
+	tc := mockcluster.NewCluster(opt)
+	store1 := core.NewStoreInfoWithLabel(1, 1, map[string]string{"id": "1"})
+	store2 := core.NewStoreInfoWithLabel(1, 1, map[string]string{"id": "2"})
+	filter1 := NewLabelConstaintFilter("", []placement.LabelConstraint{{Key: "id", Op: "in", Values: []string{"1"}}})
+	filter2 := NewLabelConstaintFilter("", []placement.LabelConstraint{{Key: "id", Op: "in", Values: []string{"2"}}})
+	c.Assert(filter1.Source(tc, store1), IsFalse)
+	c.Assert(filter1.Target(tc, store2), IsTrue)
+	c.Assert(filter2.Source(tc, store1), IsTrue)
+	c.Assert(filter2.Target(tc, store2), IsFalse)
+}
+
+func (s *testFiltersSuite) TestRuleFitFilter(c *C) {
+	opt := mockoption.NewScheduleOptions()
+	opt.EnablePlacementRules = true
+	opt.LocationLabels = []string{"zone"}
+	tc := mockcluster.NewCluster(opt)
+	tc.AddLabelsStore(1, 1, map[string]string{"zone": "z1"})
+	tc.AddLabelsStore(2, 1, map[string]string{"zone": "z1"})
+	tc.AddLabelsStore(3, 1, map[string]string{"zone": "z2"})
+	tc.AddLabelsStore(4, 1, map[string]string{"zone": "z2"})
+	tc.AddLabelsStore(5, 1, map[string]string{"zone": "z3"})
+	region := core.NewRegionInfo(&metapb.Region{Peers: []*metapb.Peer{
+		{StoreId: 1, Id: 1},
+		{StoreId: 3, Id: 3},
+		{StoreId: 5, Id: 5},
+	}}, &metapb.Peer{StoreId: 1, Id: 1})
+
+	filter := NewRuleFitFilter("", tc, region, 1)
+	c.Assert(filter.Target(tc, tc.GetStore(2)), IsFalse)
+	c.Assert(filter.Target(tc, tc.GetStore(4)), IsTrue)
+	c.Assert(filter.Source(tc, tc.GetStore(4)), IsFalse)
+}