Repository: pingcap/pd
Pull request #2051 (2019-12-24 18:51:38)

schedule: make region scheduler and checkers aware of placement rules
### What problem does this PR solve? <!--add the issue link with summary if it exists-->
Continue merge placement rules.

### What is changed and how it works?
- skip strict label match if placement rules enabled
- disallow merge regions that across different rules
- replace learnerChecker and replicaChecker with ruleChecker if placement rules enabled
- replace distinctScoreFilter with RuleFitFilter if placement rules enabled
- add tests

### Check List <!--REMOVE the items that are not applicable-->

Tests <!-- At least one of them must be included. -->
 - Unit test

codecov-io: # [Codecov](https://codecov.io/gh/pingcap/pd/pull/2051?src=pr&el=h1) Report
> Merging [#2051](https://codecov.io/gh/pingcap/pd/pull/2051?src=pr&el=desc) into [master](https://codecov.io/gh/pingcap/pd/commit/4336d595e1e1100451152be2590d2d56de51ddaf?src=pr&el=desc) will **increase** coverage by `0.18%`.
> The diff coverage is `79.68%`.

[![Impacted file tree graph](https://codecov.io/gh/pingcap/pd/pull/2051/graphs/tree.svg?width=650&token=J1yRL471OI&height=150&src=pr)](https://codecov.io/gh/pingcap/pd/pull/2051?src=pr&el=tree)

```diff
@@            Coverage Diff             @@
##           master    #2051      +/-   ##
==========================================
+ Coverage   76.65%   76.84%   +0.18%     
==========================================
  Files         185      185              
  Lines       18829    18865      +36     
==========================================
+ Hits        14433    14496      +63     
+ Misses       3285     3257      -28     
- Partials     1111     1112       +1
```


| [Impacted Files](https://codecov.io/gh/pingcap/pd/pull/2051?src=pr&el=tree) | Coverage Δ | |
|---|---|---|
| [server/config/config.go](https://codecov.io/gh/pingcap/pd/pull/2051/diff?src=pr&el=tree#diff-c2VydmVyL2NvbmZpZy9jb25maWcuZ28=) | `83.47% <ø> (ø)` | :arrow_up: |
| [server/cluster/coordinator.go](https://codecov.io/gh/pingcap/pd/pull/2051/diff?src=pr&el=tree#diff-c2VydmVyL2NsdXN0ZXIvY29vcmRpbmF0b3IuZ28=) | `78.39% <100%> (+10.8%)` | :arrow_up: |
| [server/schedule/checker\_controller.go](https://codecov.io/gh/pingcap/pd/pull/2051/diff?src=pr&el=tree#diff-c2VydmVyL3NjaGVkdWxlL2NoZWNrZXJfY29udHJvbGxlci5nbw==) | `76.66% <66.66%> (-14.64%)` | :arrow_down: |
| [server/schedulers/balance\_region.go](https://codecov.io/gh/pingcap/pd/pull/2051/diff?src=pr&el=tree#diff-c2VydmVyL3NjaGVkdWxlcnMvYmFsYW5jZV9yZWdpb24uZ28=) | `84.29% <81.81%> (+2.88%)` | :arrow_up: |
| [server/cluster/cluster.go](https://codecov.io/gh/pingcap/pd/pull/2051/diff?src=pr&el=tree#diff-c2VydmVyL2NsdXN0ZXIvY2x1c3Rlci5nbw==) | `81.05% <83.33%> (+0.24%)` | :arrow_up: |
| [server/schedule/checker/merge\_checker.go](https://codecov.io/gh/pingcap/pd/pull/2051/diff?src=pr&el=tree#diff-c2VydmVyL3NjaGVkdWxlL2NoZWNrZXIvbWVyZ2VfY2hlY2tlci5nbw==) | `70.73% <85.71%> (+1.28%)` | :arrow_up: |
| [server/schedulers/base\_scheduler.go](https://codecov.io/gh/pingcap/pd/pull/2051/diff?src=pr&el=tree#diff-c2VydmVyL3NjaGVkdWxlcnMvYmFzZV9zY2hlZHVsZXIuZ28=) | `60.86% <0%> (-8.7%)` | :arrow_down: |
| [server/schedulers/shuffle\_hot\_region.go](https://codecov.io/gh/pingcap/pd/pull/2051/diff?src=pr&el=tree#diff-c2VydmVyL3NjaGVkdWxlcnMvc2h1ZmZsZV9ob3RfcmVnaW9uLmdv) | `62.62% <0%> (-5.06%)` | :arrow_down: |
| [server/schedulers/random\_merge.go](https://codecov.io/gh/pingcap/pd/pull/2051/diff?src=pr&el=tree#diff-c2VydmVyL3NjaGVkdWxlcnMvcmFuZG9tX21lcmdlLmdv) | `60.71% <0%> (-3.58%)` | :arrow_down: |
| [server/kv/etcd\_kv.go](https://codecov.io/gh/pingcap/pd/pull/2051/diff?src=pr&el=tree#diff-c2VydmVyL2t2L2V0Y2Rfa3YuZ28=) | `81.81% <0%> (-2.6%)` | :arrow_down: |
| ... and [16 more](https://codecov.io/gh/pingcap/pd/pull/2051/diff?src=pr&el=tree-more) | |

------

[Continue to review full report at Codecov](https://codecov.io/gh/pingcap/pd/pull/2051?src=pr&el=continue).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/pingcap/pd/pull/2051?src=pr&el=footer). Last update [4336d59...5a0fd92](https://codecov.io/gh/pingcap/pd/pull/2051?src=pr&el=lastupdated). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments).

disksing: Please take a look @rleungx 
disksing: /merge
sre-bot: /run-all-tests
sre-bot: @disksing merge failed.
disksing: /merge
sre-bot: /run-all-tests
sre-bot: @disksing merge failed.
nolouch: /merge
sre-bot: /run-all-tests

Patch:
diff --git a/pkg/mock/mockcluster/mockcluster.go b/pkg/mock/mockcluster/mockcluster.go
index 348421e8d..d16c5832b 100644
--- a/pkg/mock/mockcluster/mockcluster.go
+++ b/pkg/mock/mockcluster/mockcluster.go
@@ -43,7 +43,8 @@ type Cluster struct {
 
 // NewCluster creates a new Cluster
 func NewCluster(opt *mockoption.ScheduleOptions) *Cluster {
-	ruleManager, _ := placement.NewRuleManager(core.NewStorage(kv.NewMemoryKV()), opt.MaxReplicas, opt.GetLocationLabels())
+	ruleManager := placement.NewRuleManager(core.NewStorage(kv.NewMemoryKV()))
+	ruleManager.Initialize(opt.MaxReplicas, opt.GetLocationLabels())
 	return &Cluster{
 		BasicCluster:    core.NewBasicCluster(),
 		IDAllocator:     mockid.NewIDAllocator(),
diff --git a/server/cluster/cluster.go b/server/cluster/cluster.go
index e2440f2e2..6c205781f 100644
--- a/server/cluster/cluster.go
+++ b/server/cluster/cluster.go
@@ -200,6 +200,14 @@ func (c *RaftCluster) Start(s Server) error {
 		return nil
 	}
 
+	c.ruleManager = placement.NewRuleManager(c.storage)
+	if c.IsPlacementRulesEnabled() {
+		err = c.ruleManager.Initialize(c.opt.GetMaxReplicas(), c.opt.GetLocationLabels())
+		if err != nil {
+			return err
+		}
+	}
+
 	c.coordinator = newCoordinator(c.ctx, cluster, s.GetHBStreams())
 	c.regionStats = statistics.NewRegionStatistics(c.opt)
 	c.quit = make(chan struct{})
@@ -797,7 +805,16 @@ func (c *RaftCluster) PutStore(store *metapb.Store) error {
 			core.SetStoreLabels(labels),
 		)
 	}
-	// Check location labels.
+	if err = c.checkStoreLabels(s); err != nil {
+		return err
+	}
+	return c.putStoreLocked(s)
+}
+
+func (c *RaftCluster) checkStoreLabels(s *core.StoreInfo) error {
+	if c.opt.IsPlacementRulesEnabled() {
+		return nil
+	}
 	keysSet := make(map[string]struct{})
 	for _, k := range c.GetLocationLabels() {
 		keysSet[k] = struct{}{}
@@ -821,7 +838,7 @@ func (c *RaftCluster) PutStore(store *metapb.Store) error {
 			}
 		}
 	}
-	return c.putStoreLocked(s)
+	return nil
 }
 
 // RemoveStore marks a store as offline in cluster.
diff --git a/server/cluster/coordinator.go b/server/cluster/coordinator.go
index b7e7819d5..f9f082146 100644
--- a/server/cluster/coordinator.go
+++ b/server/cluster/coordinator.go
@@ -78,7 +78,7 @@ func newCoordinator(ctx context.Context, cluster *RaftCluster, hbStreams opt.Hea
 		ctx:             ctx,
 		cancel:          cancel,
 		cluster:         cluster,
-		checkers:        schedule.NewCheckerController(ctx, cluster, opController),
+		checkers:        schedule.NewCheckerController(ctx, cluster, cluster.ruleManager, opController),
 		regionScatterer: schedule.NewRegionScatterer(cluster),
 		schedulers:      make(map[string]*scheduleController),
 		opController:    opController,
diff --git a/server/config/config.go b/server/config/config.go
index 842735a8b..dc9e86e16 100644
--- a/server/config/config.go
+++ b/server/config/config.go
@@ -819,7 +819,7 @@ type ReplicationConfig struct {
 	StrictlyMatchLabel bool `toml:"strictly-match-label,omitempty" json:"strictly-match-label,string"`
 
 	// When PlacementRules feature is enabled. MaxReplicas and LocationLabels are not uesd any more.
-	EnablePlacementRules bool // Keep it false before full feature get merged. `toml:"enable-placement-rules" json:"enable-placement-rules,string"`
+	EnablePlacementRules bool `toml:"enable-placement-rules" json:"enable-placement-rules,string"`
 }
 
 func (c *ReplicationConfig) clone() *ReplicationConfig {
diff --git a/server/schedule/checker/merge_checker.go b/server/schedule/checker/merge_checker.go
index c11314b6e..d2e68db0c 100644
--- a/server/schedule/checker/merge_checker.go
+++ b/server/schedule/checker/merge_checker.go
@@ -14,6 +14,7 @@
 package checker
 
 import (
+	"bytes"
 	"context"
 	"time"
 
@@ -23,23 +24,26 @@ import (
 	"github.com/pingcap/pd/server/core"
 	"github.com/pingcap/pd/server/schedule/operator"
 	"github.com/pingcap/pd/server/schedule/opt"
+	"github.com/pingcap/pd/server/schedule/placement"
 	"go.uber.org/zap"
 )
 
 // MergeChecker ensures region to merge with adjacent region when size is small
 type MergeChecker struct {
-	cluster    opt.Cluster
-	splitCache *cache.TTLUint64
-	startTime  time.Time // it's used to judge whether server recently start.
+	cluster     opt.Cluster
+	splitCache  *cache.TTLUint64
+	ruleManager *placement.RuleManager
+	startTime   time.Time // it's used to judge whether server recently start.
 }
 
 // NewMergeChecker creates a merge checker.
-func NewMergeChecker(ctx context.Context, cluster opt.Cluster) *MergeChecker {
+func NewMergeChecker(ctx context.Context, cluster opt.Cluster, ruleManager *placement.RuleManager) *MergeChecker {
 	splitCache := cache.NewIDTTL(ctx, time.Minute, cluster.GetSplitMergeInterval())
 	return &MergeChecker{
-		cluster:    cluster,
-		splitCache: splitCache,
-		startTime:  time.Now(),
+		cluster:     cluster,
+		splitCache:  splitCache,
+		ruleManager: ruleManager,
+		startTime:   time.Now(),
 	}
 }
 
@@ -137,6 +141,17 @@ func (m *MergeChecker) checkTarget(region, adjacent *core.RegionInfo) bool {
 
 // allowMerge returns true if two regions can be merged according to the key type.
 func (m *MergeChecker) allowMerge(region *core.RegionInfo, adjacent *core.RegionInfo) bool {
+	var start, end []byte
+	if bytes.Equal(region.GetEndKey(), adjacent.GetStartKey()) && len(region.GetEndKey()) != 0 {
+		start, end = region.GetStartKey(), adjacent.GetEndKey()
+	} else if bytes.Equal(adjacent.GetEndKey(), region.GetStartKey()) && len(adjacent.GetEndKey()) != 0 {
+		start, end = adjacent.GetStartKey(), region.GetEndKey()
+	} else {
+		return false
+	}
+	if m.cluster.IsPlacementRulesEnabled() && len(m.ruleManager.GetSplitKeys(start, end)) > 0 {
+		return false
+	}
 	strategy := m.cluster.GetKeyType()
 	switch strategy {
 	case core.Table:
diff --git a/server/schedule/checker_controller.go b/server/schedule/checker_controller.go
index 77d0d179d..d0edb36a3 100644
--- a/server/schedule/checker_controller.go
+++ b/server/schedule/checker_controller.go
@@ -20,6 +20,7 @@ import (
 	"github.com/pingcap/pd/server/schedule/checker"
 	"github.com/pingcap/pd/server/schedule/operator"
 	"github.com/pingcap/pd/server/schedule/opt"
+	"github.com/pingcap/pd/server/schedule/placement"
 )
 
 // CheckerController is used to manage all checkers.
@@ -28,18 +29,20 @@ type CheckerController struct {
 	opController   *OperatorController
 	learnerChecker *checker.LearnerChecker
 	replicaChecker *checker.ReplicaChecker
+	ruleChecker    *checker.RuleChecker
 	mergeChecker   *checker.MergeChecker
 }
 
 // NewCheckerController create a new CheckerController.
 // TODO: isSupportMerge should be removed.
-func NewCheckerController(ctx context.Context, cluster opt.Cluster, opController *OperatorController) *CheckerController {
+func NewCheckerController(ctx context.Context, cluster opt.Cluster, ruleManager *placement.RuleManager, opController *OperatorController) *CheckerController {
 	return &CheckerController{
 		cluster:        cluster,
 		opController:   opController,
 		learnerChecker: checker.NewLearnerChecker(cluster),
 		replicaChecker: checker.NewReplicaChecker(cluster),
-		mergeChecker:   checker.NewMergeChecker(ctx, cluster),
+		ruleChecker:    checker.NewRuleChecker(cluster, ruleManager),
+		mergeChecker:   checker.NewMergeChecker(ctx, cluster, ruleManager),
 	}
 }
 
@@ -49,15 +52,25 @@ func (c *CheckerController) CheckRegion(region *core.RegionInfo) (bool, []*opera
 	// Don't check isRaftLearnerEnabled cause it maybe disable learner feature but there are still some learners to promote.
 	opController := c.opController
 	checkerIsBusy := true
-	if op := c.learnerChecker.Check(region); op != nil {
-		return false, []*operator.Operator{op}
-	}
-	if opController.OperatorCount(operator.OpReplica) < c.cluster.GetReplicaScheduleLimit() {
-		checkerIsBusy = false
-		if op := c.replicaChecker.Check(region); op != nil {
-			return checkerIsBusy, []*operator.Operator{op}
+	if c.cluster.IsPlacementRulesEnabled() {
+		if opController.OperatorCount(operator.OpReplica) < c.cluster.GetReplicaScheduleLimit() {
+			checkerIsBusy = false
+			if op := c.ruleChecker.Check(region); op != nil {
+				return checkerIsBusy, []*operator.Operator{op}
+			}
+		}
+	} else {
+		if op := c.learnerChecker.Check(region); op != nil {
+			return false, []*operator.Operator{op}
+		}
+		if opController.OperatorCount(operator.OpReplica) < c.cluster.GetReplicaScheduleLimit() {
+			checkerIsBusy = false
+			if op := c.replicaChecker.Check(region); op != nil {
+				return checkerIsBusy, []*operator.Operator{op}
+			}
 		}
 	}
+
 	if c.mergeChecker != nil && opController.OperatorCount(operator.OpMerge) < c.cluster.GetMergeScheduleLimit() {
 		checkerIsBusy = false
 		if ops := c.mergeChecker.Check(region); ops != nil {
diff --git a/server/schedule/placement/rule_manager.go b/server/schedule/placement/rule_manager.go
index 5203572dc..82c33a9f0 100644
--- a/server/schedule/placement/rule_manager.go
+++ b/server/schedule/placement/rule_manager.go
@@ -28,23 +28,28 @@ import (
 type RuleManager struct {
 	store *core.Storage
 	sync.RWMutex
-	rules     map[[2]string]*Rule
-	splitKeys [][]byte
+	initialized bool
+	rules       map[[2]string]*Rule
+	splitKeys   [][]byte
 }
 
 // NewRuleManager creates a RuleManager instance.
-func NewRuleManager(store *core.Storage, maxReplica int, locationLabels []string) (*RuleManager, error) {
-	m := &RuleManager{
+func NewRuleManager(store *core.Storage) *RuleManager {
+	return &RuleManager{
 		store: store,
 		rules: make(map[[2]string]*Rule),
 	}
-	if err := m.initialize(maxReplica, locationLabels); err != nil {
-		return nil, err
-	}
-	return m, nil
 }
 
-func (m *RuleManager) initialize(maxReplica int, locationLabels []string) error {
+// Initialize loads rules from storage. If Placement Rules feature is never enabled, it creates default rule that is
+// compatible with previous configuration.
+func (m *RuleManager) Initialize(maxReplica int, locationLabels []string) error {
+	m.Lock()
+	defer m.Unlock()
+	if m.initialized {
+		return nil
+	}
+
 	var rules []*Rule
 	ok, err := m.store.LoadRules(&rules)
 	if err != nil {
@@ -66,6 +71,7 @@ func (m *RuleManager) initialize(maxReplica int, locationLabels []string) error
 		m.rules[r.Key()] = r
 	}
 	m.updateSplitKeys()
+	m.initialized = true
 	return nil
 }
 
diff --git a/server/schedulers/balance_region.go b/server/schedulers/balance_region.go
index 252deaf73..0f0a77d2c 100644
--- a/server/schedulers/balance_region.go
+++ b/server/schedulers/balance_region.go
@@ -138,17 +138,21 @@ func (s *balanceRegionScheduler) Schedule(cluster opt.Cluster) []*operator.Opera
 		sourceID := source.GetID()
 
 		for i := 0; i < balanceRegionRetryLimit; i++ {
-			// Priority picks the region that has a pending peer.
+			// Priority pick the region that has a pending peer.
 			// Pending region may means the disk is overload, remove the pending region firstly.
 			region := cluster.RandPendingRegion(sourceID, s.conf.Ranges, opt.HealthAllowPending(cluster), opt.ReplicatedRegion(cluster))
 			if region == nil {
-				// Then picks the region that has a follower in the source store.
+				// Then pick the region that has a follower in the source store.
 				region = cluster.RandFollowerRegion(sourceID, s.conf.Ranges, opt.HealthRegion(cluster), opt.ReplicatedRegion(cluster))
 			}
 			if region == nil {
-				// Last, picks the region has the leader in the source store.
+				// Then pick the region has the leader in the source store.
 				region = cluster.RandLeaderRegion(sourceID, s.conf.Ranges, opt.HealthRegion(cluster), opt.ReplicatedRegion(cluster))
 			}
+			if region == nil {
+				// Finally pick learner.
+				region = cluster.RandLearnerRegion(sourceID, s.conf.Ranges, opt.HealthRegion(cluster), opt.ReplicatedRegion(cluster))
+			}
 			if region == nil {
 				schedulerCounter.WithLabelValues(s.GetName(), "no-region").Inc()
 				continue
@@ -180,24 +184,35 @@ func (s *balanceRegionScheduler) transferPeer(cluster opt.Cluster, region *core.
 	source := cluster.GetStore(sourceStoreID)
 	if source == nil {
 		log.Error("failed to get the source store", zap.Uint64("store-id", sourceStoreID))
+		return nil
 	}
-	scoreGuard := filter.NewDistinctScoreFilter(s.GetName(), cluster.GetLocationLabels(), stores, source)
-	replicaChecker := checker.NewReplicaChecker(cluster, s.GetName())
 	exclude := make(map[uint64]struct{})
 	excludeFilter := filter.NewExcludedFilter(s.GetName(), nil, exclude)
 	for {
-		storeID, _ := replicaChecker.SelectBestReplacementStore(region, oldPeer, scoreGuard, excludeFilter)
-		if storeID == 0 {
+		var target *core.StoreInfo
+		if cluster.IsPlacementRulesEnabled() {
+			scoreGuard := filter.NewRuleFitFilter(s.GetName(), cluster, region, sourceStoreID)
+			fit := cluster.FitRegion(region)
+			rf := fit.GetRuleFit(oldPeer.GetId())
+			if rf == nil {
+				schedulerCounter.WithLabelValues(s.GetName(), "skip-orphan-peer").Inc()
+				return nil
+			}
+			target = checker.SelectStoreToReplacePeerByRule(s.GetName(), cluster, region, fit, rf, oldPeer, scoreGuard, excludeFilter)
+		} else {
+			scoreGuard := filter.NewDistinctScoreFilter(s.GetName(), cluster.GetLocationLabels(), stores, source)
+			replicaChecker := checker.NewReplicaChecker(cluster, s.GetName())
+			storeID, _ := replicaChecker.SelectBestReplacementStore(region, oldPeer, scoreGuard, excludeFilter)
+			if storeID != 0 {
+				target = cluster.GetStore(storeID)
+			}
+		}
+		if target == nil {
 			schedulerCounter.WithLabelValues(s.GetName(), "no-replacement").Inc()
 			return nil
 		}
-		exclude[storeID] = struct{}{} // exclude next round.
+		exclude[target.GetID()] = struct{}{} // exclude next round.
 
-		target := cluster.GetStore(storeID)
-		if target == nil {
-			log.Error("failed to get the target store", zap.Uint64("store-id", storeID))
-			continue
-		}
 		regionID := region.GetID()
 		sourceID := source.GetID()
 		targetID := target.GetID()
@@ -210,12 +225,8 @@ func (s *balanceRegionScheduler) transferPeer(cluster opt.Cluster, region *core.
 			continue
 		}
 
-		newPeer, err := cluster.AllocPeer(storeID)
-		if err != nil {
-			schedulerCounter.WithLabelValues(s.GetName(), "no-peer").Inc()
-			return nil
-		}
-		op, err := operator.CreateMovePeerOperator(BalanceRegionType, cluster, region, operator.OpBalance, oldPeer.GetStoreId(), newPeer)
+		newPeer := &metapb.Peer{StoreId: target.GetID(), IsLearner: oldPeer.IsLearner}
+		op, err := operator.CreateMovePeerOperator("balance-region", cluster, region, operator.OpBalance, oldPeer.GetStoreId(), newPeer)
 		if err != nil {
 			schedulerCounter.WithLabelValues(s.GetName(), "create-operator-fail").Inc()
 			return nil
diff --git a/server/server.go b/server/server.go
index ca825a18c..1cef6d3f4 100644
--- a/server/server.go
+++ b/server/server.go
@@ -669,6 +669,12 @@ func (s *Server) SetReplicationConfig(cfg config.ReplicationConfig) error {
 	if err := cfg.Validate(); err != nil {
 		return err
 	}
+	if cfg.EnablePlacementRules {
+		// initialize rule manager.
+		if err := s.GetRaftCluster().GetRuleManager().Initialize(int(cfg.MaxReplicas), cfg.LocationLabels); err != nil {
+			return err
+		}
+	}
 	old := s.scheduleOpt.GetReplication().Load()
 	s.scheduleOpt.GetReplication().Store(&cfg)
 	if err := s.scheduleOpt.Persist(s.storage); err != nil {


Test patch:
diff --git a/server/api/label_test.go b/server/api/label_test.go
index b2e461e73..203f5c8f0 100644
--- a/server/api/label_test.go
+++ b/server/api/label_test.go
@@ -274,6 +274,22 @@ func (s *testStrictlyLabelsStoreSuite) TestStoreMatch(c *C) {
 			c.Assert(strings.Contains(err.Error(), t.expectError), IsTrue)
 		}
 	}
+
+	// enable placement rules. Report no error any more.
+	c.Assert(postJSON(fmt.Sprintf("%s/config", s.urlPrefix), []byte(`{"enable-placement-rules":"true"}`)), IsNil)
+	for _, t := range cases {
+		_, err := s.svr.PutStore(context.Background(), &pdpb.PutStoreRequest{
+			Header: &pdpb.RequestHeader{ClusterId: s.svr.ClusterID()},
+			Store: &metapb.Store{
+				Id:      t.store.Id,
+				Address: fmt.Sprintf("tikv%d", t.store.Id),
+				State:   t.store.State,
+				Labels:  t.store.Labels,
+				Version: t.store.Version,
+			},
+		})
+		c.Assert(err, IsNil)
+	}
 }
 
 func (s *testStrictlyLabelsStoreSuite) TearDownSuite(c *C) {
diff --git a/server/schedule/checker/merge_checker_test.go b/server/schedule/checker/merge_checker_test.go
index 7bc8b0ff3..6b8c3bdb4 100644
--- a/server/schedule/checker/merge_checker_test.go
+++ b/server/schedule/checker/merge_checker_test.go
@@ -15,6 +15,7 @@ package checker
 
 import (
 	"context"
+	"encoding/hex"
 	"testing"
 	"time"
 
@@ -26,6 +27,7 @@ import (
 	"github.com/pingcap/pd/server/core"
 	"github.com/pingcap/pd/server/schedule/operator"
 	"github.com/pingcap/pd/server/schedule/opt"
+	"github.com/pingcap/pd/server/schedule/placement"
 	"go.uber.org/goleak"
 )
 
@@ -127,7 +129,7 @@ func (s *testMergeCheckerSuite) SetUpTest(c *C) {
 		s.cluster.PutRegion(region)
 	}
 	s.ctx, s.cancel = context.WithCancel(context.Background())
-	s.mc = NewMergeChecker(s.ctx, s.cluster)
+	s.mc = NewMergeChecker(s.ctx, s.cluster, s.cluster.RuleManager)
 }
 
 func (s *testMergeCheckerSuite) TearDownTest(c *C) {
@@ -164,6 +166,25 @@ func (s *testMergeCheckerSuite) TestBasic(c *C) {
 	c.Assert(ops[0].RegionID(), Equals, s.regions[2].GetID())
 	c.Assert(ops[1].RegionID(), Equals, s.regions[3].GetID())
 
+	// merge cannot across rule key.
+	s.cluster.EnablePlacementRules = true
+	s.mc.ruleManager.SetRule(&placement.Rule{
+		GroupID:     "pd",
+		ID:          "test",
+		Index:       1,
+		Override:    true,
+		StartKeyHex: hex.EncodeToString([]byte("x")),
+		EndKeyHex:   hex.EncodeToString([]byte("z")),
+		Role:        placement.Voter,
+		Count:       3,
+	})
+	// region 2 can only merge with previous region now.
+	ops = s.mc.Check(s.regions[2])
+	c.Assert(ops, NotNil)
+	c.Assert(ops[0].RegionID(), Equals, s.regions[2].GetID())
+	c.Assert(ops[1].RegionID(), Equals, s.regions[1].GetID())
+	s.mc.ruleManager.DeleteRule("test", "test")
+
 	// Skip recently split regions.
 	s.cluster.ScheduleOptions.SplitMergeInterval = time.Hour
 	s.mc.RecordRegionSplit([]uint64{s.regions[2].GetID()})
@@ -442,7 +463,7 @@ func (s *testMergeCheckerSuite) TestCache(c *C) {
 		s.cluster.PutRegion(region)
 	}
 
-	s.mc = NewMergeChecker(s.ctx, s.cluster)
+	s.mc = NewMergeChecker(s.ctx, s.cluster, s.cluster.RuleManager)
 
 	ops := s.mc.Check(s.regions[1])
 	c.Assert(ops, IsNil)
diff --git a/server/schedule/operator_controller_test.go b/server/schedule/operator_controller_test.go
index 83f3e7bd9..b1e543f64 100644
--- a/server/schedule/operator_controller_test.go
+++ b/server/schedule/operator_controller_test.go
@@ -522,7 +522,7 @@ func (t *testOperatorControllerSuite) TestStoreLimitWithMerge(c *C) {
 		tc.PutRegion(region)
 	}
 
-	mc := checker.NewMergeChecker(t.ctx, tc)
+	mc := checker.NewMergeChecker(t.ctx, tc, tc.RuleManager)
 	oc := NewOperatorController(t.ctx, tc, mockhbstream.NewHeartbeatStream())
 
 	cfg.StoreBalanceRate = 60
diff --git a/server/schedule/placement/rule_manager_test.go b/server/schedule/placement/rule_manager_test.go
index 2c3dfc627..7875ab52d 100644
--- a/server/schedule/placement/rule_manager_test.go
+++ b/server/schedule/placement/rule_manager_test.go
@@ -32,7 +32,8 @@ type testManagerSuite struct {
 func (s *testManagerSuite) SetUpTest(c *C) {
 	s.store = core.NewStorage(kv.NewMemoryKV())
 	var err error
-	s.manager, err = NewRuleManager(s.store, 3, []string{"zone", "rack", "host"})
+	s.manager = NewRuleManager(s.store)
+	err = s.manager.Initialize(3, []string{"zone", "rack", "host"})
 	c.Assert(err, IsNil)
 }
 
@@ -77,7 +78,8 @@ func (s *testManagerSuite) TestSaveLoad(c *C) {
 	for _, r := range rules {
 		s.manager.SetRule(r)
 	}
-	m2, err := NewRuleManager(s.store, 3, []string{"no", "labels"})
+	m2 := NewRuleManager(s.store)
+	err := m2.Initialize(3, []string{"no", "labels"})
 	c.Assert(err, IsNil)
 	c.Assert(m2.GetAllRules(), HasLen, 2)
 	c.Assert(m2.GetRule("pd", "default"), DeepEquals, rules[0])
diff --git a/server/schedulers/balance_test.go b/server/schedulers/balance_test.go
index 7e4e4bc76..cdc9bfc82 100644
--- a/server/schedulers/balance_test.go
+++ b/server/schedulers/balance_test.go
@@ -557,14 +557,20 @@ func (s *testBalanceRegionSchedulerSuite) TestBalance(c *C) {
 
 func (s *testBalanceRegionSchedulerSuite) TestReplicas3(c *C) {
 	opt := mockoption.NewScheduleOptions()
+	newTestReplication(opt, 3, "zone", "rack", "host")
+
 	tc := mockcluster.NewCluster(opt)
 	oc := schedule.NewOperatorController(s.ctx, nil, nil)
 
-	newTestReplication(opt, 3, "zone", "rack", "host")
-
 	sb, err := schedule.CreateScheduler(BalanceRegionType, oc, core.NewStorage(kv.NewMemoryKV()), schedule.ConfigSliceDecoder(BalanceRegionType, []string{"", ""}))
 	c.Assert(err, IsNil)
 
+	s.checkReplica3(c, tc, opt, sb)
+	opt.EnablePlacementRules = true
+	s.checkReplica3(c, tc, opt, sb)
+}
+
+func (s *testBalanceRegionSchedulerSuite) checkReplica3(c *C, tc *mockcluster.Cluster, opt *mockoption.ScheduleOptions, sb schedule.Scheduler) {
 	// Store 1 has the largest region score, so the balancer try to replace peer in store 1.
 	tc.AddLabelsStore(1, 16, map[string]string{"zone": "z1", "rack": "r1", "host": "h1"})
 	tc.AddLabelsStore(2, 15, map[string]string{"zone": "z1", "rack": "r2", "host": "h1"})
@@ -612,14 +618,20 @@ func (s *testBalanceRegionSchedulerSuite) TestReplicas3(c *C) {
 
 func (s *testBalanceRegionSchedulerSuite) TestReplicas5(c *C) {
 	opt := mockoption.NewScheduleOptions()
+	newTestReplication(opt, 5, "zone", "rack", "host")
+
 	tc := mockcluster.NewCluster(opt)
 	oc := schedule.NewOperatorController(s.ctx, nil, nil)
 
-	newTestReplication(opt, 5, "zone", "rack", "host")
-
 	sb, err := schedule.CreateScheduler(BalanceRegionType, oc, core.NewStorage(kv.NewMemoryKV()), schedule.ConfigSliceDecoder(BalanceRegionType, []string{"", ""}))
 	c.Assert(err, IsNil)
 
+	s.checkReplica5(c, tc, opt, sb)
+	opt.EnablePlacementRules = true
+	s.checkReplica5(c, tc, opt, sb)
+}
+
+func (s *testBalanceRegionSchedulerSuite) checkReplica5(c *C, tc *mockcluster.Cluster, opt *mockoption.ScheduleOptions, sb schedule.Scheduler) {
 	tc.AddLabelsStore(1, 4, map[string]string{"zone": "z1", "rack": "r1", "host": "h1"})
 	tc.AddLabelsStore(2, 5, map[string]string{"zone": "z2", "rack": "r1", "host": "h1"})
 	tc.AddLabelsStore(3, 6, map[string]string{"zone": "z3", "rack": "r1", "host": "h1"})
@@ -737,6 +749,12 @@ func (s *testBalanceRegionSchedulerSuite) TestReplacePendingRegion(c *C) {
 	sb, err := schedule.CreateScheduler(BalanceRegionType, oc, core.NewStorage(kv.NewMemoryKV()), schedule.ConfigSliceDecoder(BalanceRegionType, []string{"", ""}))
 	c.Assert(err, IsNil)
 
+	s.checkReplacePendingRegion(c, tc, opt, sb)
+	opt.EnablePlacementRules = true
+	s.checkReplacePendingRegion(c, tc, opt, sb)
+}
+
+func (s *testBalanceRegionSchedulerSuite) checkReplacePendingRegion(c *C, tc *mockcluster.Cluster, opt *mockoption.ScheduleOptions, sb schedule.Scheduler) {
 	// Store 1 has the largest region score, so the balance scheduler try to replace peer in store 1.
 	tc.AddLabelsStore(1, 16, map[string]string{"zone": "z1", "rack": "r1", "host": "h1"})
 	tc.AddLabelsStore(2, 7, map[string]string{"zone": "z1", "rack": "r2", "host": "h1"})