From 1d01343bd65b88889c8ed1b0cc2bf9d894851fa6 Mon Sep 17 00:00:00 2001
From: nolouch <nolouch@gmail.com>
Date: Wed, 27 Mar 2019 16:08:40 +0800
Subject: [PATCH] scheduler: fix region scatter may transfer leader to removed
 peer

Signed-off-by: nolouch <nolouch@gmail.com>
---
 server/schedule/operator.go         | 21 +++++++++++++++++++++
 server/schedule/region_scatterer.go |  6 ++++--
 server/schedulers/scheduler_test.go |  7 ++++++-
 3 files changed, 31 insertions(+), 3 deletions(-)

diff --git a/server/schedule/operator.go b/server/schedule/operator.go
index ca78345d2e2..96ad0f1abdc 100644
--- a/server/schedule/operator.go
+++ b/server/schedule/operator.go
@@ -630,3 +630,24 @@ func getIntersectionStores(a []*metapb.Peer, b []*metapb.Peer) map[uint64]struct
 
 	return intersection
 }
+
+// CheckOperatorValied checks if the operator is valid.
+func CheckOperatorValid(op *Operator) bool {
+	removeStores := []uint64{}
+	for _, step := range op.steps {
+		if tr, ok := step.(TransferLeader); ok {
+			for _, store := range removeStores {
+				if store == tr.FromStore {
+					return false
+				}
+				if store == tr.ToStore {
+					return false
+				}
+			}
+		}
+		if rp, ok := step.(RemovePeer); ok {
+			removeStores = append(removeStores, rp.FromStore)
+		}
+	}
+	return true
+}
diff --git a/server/schedule/region_scatterer.go b/server/schedule/region_scatterer.go
index 22bbdb8df88..192f892166d 100644
--- a/server/schedule/region_scatterer.go
+++ b/server/schedule/region_scatterer.go
@@ -96,6 +96,7 @@ func (r *RegionScatterer) scatterRegion(region *core.RegionInfo) *Operator {
 
 	stores := r.collectAvailableStores(region)
 	var kind OperatorKind
+	newRegion := region.Clone()
 	for _, peer := range region.GetPeers() {
 		if len(stores) == 0 {
 			// Reset selected stores if we have no available stores.
@@ -116,13 +117,14 @@ func (r *RegionScatterer) scatterRegion(region *core.RegionInfo) *Operator {
 		delete(stores, newPeer.GetStoreId())
 		r.selected.put(newPeer.GetStoreId())
 
-		op, err := CreateMovePeerOperator("scatter-peer", r.cluster, region, OpAdmin,
+		op, err := CreateMovePeerOperator("scatter-peer", r.cluster, newRegion, OpAdmin,
 			peer.GetStoreId(), newPeer.GetStoreId(), newPeer.GetId())
 		if err != nil {
 			continue
 		}
 		steps = append(steps, op.steps...)
-		steps = append(steps, TransferLeader{ToStore: newPeer.GetStoreId()})
+		newRegion = newRegion.Clone(core.WithRemoveStorePeer(peer.GetStoreId()), core.WithAddPeer(newPeer))
+
 		kind |= op.Kind()
 	}
 
diff --git a/server/schedulers/scheduler_test.go b/server/schedulers/scheduler_test.go
index 3f90cec7243..7c9b393be25 100644
--- a/server/schedulers/scheduler_test.go
+++ b/server/schedulers/scheduler_test.go
@@ -174,6 +174,10 @@ func (s *testScatterRegionSuite) TestFiveStores(c *C) {
 	s.scatter(c, 5, 5)
 }
 
+func (s *testScatterRegionSuite) checkOperator(op *schedule.Operator, c *C) {
+	c.Assert(schedule.CheckOperatorValid(op), IsTrue)
+}
+
 func (s *testScatterRegionSuite) scatter(c *C, numStores, numRegions uint64) {
 	opt := schedule.NewMockSchedulerOptions()
 	tc := schedule.NewMockCluster(opt)
@@ -184,7 +188,7 @@ func (s *testScatterRegionSuite) scatter(c *C, numStores, numRegions uint64) {
 	}
 
 	// Add regions 1~4.
-	seq := newSequencer(numStores)
+	seq := newSequencer(3)
 	// Region 1 has the same distribution with the Region 2, which is used to test selectPeerToReplace.
 	tc.AddLeaderRegion(1, 1, 2, 3)
 	for i := uint64(2); i <= numRegions; i++ {
@@ -196,6 +200,7 @@ func (s *testScatterRegionSuite) scatter(c *C, numStores, numRegions uint64) {
 	for i := uint64(1); i <= numRegions; i++ {
 		region := tc.GetRegion(i)
 		if op := scatterer.Scatter(region); op != nil {
+			s.checkOperator(op, c)
 			tc.ApplyOperator(op)
 		}
 	}