From 9f653bea5689d535f511b31d6c3638b8fb9d0840 Mon Sep 17 00:00:00 2001
From: disksing <i@disksing.com>
Date: Mon, 3 Jul 2017 17:40:05 +0800
Subject: [PATCH] balance: avoid remove leader peer. (#671)

---
 server/balancer.go      |  8 ++++----
 server/balancer_test.go |  9 ++++++++-
 server/scheduler.go     | 14 ++++----------
 3 files changed, 16 insertions(+), 15 deletions(-)

diff --git a/server/balancer.go b/server/balancer.go
index fc2119152f5..90ed221c2b7 100644
--- a/server/balancer.go
+++ b/server/balancer.go
@@ -206,7 +206,7 @@ func (s *balanceRegionScheduler) transferPeer(cluster *clusterInfo, region *Regi
 	}
 	s.limit = adjustBalanceLimit(cluster, s.GetResourceKind())
 
-	return newTransferPeer(region, oldPeer, newPeer)
+	return newTransferPeer(region, RegionKind, oldPeer, newPeer)
 }
 
 // replicaChecker ensures region has the best replicas.
@@ -372,7 +372,7 @@ func (r *replicaChecker) checkOfflinePeer(region *RegionInfo) Operator {
 		if newPeer == nil {
 			return nil
 		}
-		return newTransferPeer(region, peer, newPeer)
+		return newTransferPeer(region, RegionKind, peer, newPeer)
 	}
 
 	return nil
@@ -391,7 +391,7 @@ func (r *replicaChecker) checkBestReplacement(region *RegionInfo) Operator {
 	if newScore <= oldScore {
 		return nil
 	}
-	return newTransferPeer(region, oldPeer, newPeer)
+	return newTransferPeer(region, RegionKind, oldPeer, newPeer)
 }
 
 // RegionStat records each hot region's statistics
@@ -467,7 +467,7 @@ func (h *balanceHotRegionScheduler) Schedule(cluster *clusterInfo) Operator {
 	// balance by peer
 	srcRegion, srcPeer, destPeer := h.balanceByPeer(cluster)
 	if srcRegion != nil {
-		return newPriorityTransferPeer(srcRegion, srcPeer, destPeer)
+		return newTransferPeer(srcRegion, PriorityKind, srcPeer, destPeer)
 	}
 
 	// balance by leader
diff --git a/server/balancer_test.go b/server/balancer_test.go
index cac20310b1e..26259f6e49f 100644
--- a/server/balancer_test.go
+++ b/server/balancer_test.go
@@ -785,6 +785,13 @@ func checkTransferPeer(c *C, bop Operator, sourceID, targetID uint64) {
 	}
 }
 
+func checkTransferPeerWithLeaderTransfer(c *C, bop Operator, sourceID, targetID uint64) {
+	op := bop.(*regionOperator)
+	c.Assert(op, NotNil)
+	c.Assert(op.Ops, HasLen, 3)
+	checkTransferPeer(c, bop, sourceID, targetID)
+}
+
 func checkTransferLeader(c *C, bop Operator, sourceID, targetID uint64) {
 	var op *transferLeaderOperator
 	switch t := bop.(type) {
@@ -848,7 +855,7 @@ func (s *testBalanceHotRegionSchedulerSuite) TestBalance(c *C) {
 
 	// Will transfer a hot region from store 1 to store 5, because the total count of peers
 	// which is hot for store 1 is more larger than other stores.
-	checkTransferPeer(c, hb.Schedule(cluster), 1, 5)
+	checkTransferPeerWithLeaderTransfer(c, hb.Schedule(cluster), 1, 5)
 
 	// After transfer a hot region from store 1 to store 5
 	//| region_id | leader_sotre | follower_store | follower_store | written_bytes |
diff --git a/server/scheduler.go b/server/scheduler.go
index f89a0dc0dd0..6fa83028241 100644
--- a/server/scheduler.go
+++ b/server/scheduler.go
@@ -234,7 +234,7 @@ func (s *shuffleRegionScheduler) Schedule(cluster *clusterInfo) Operator {
 		return nil
 	}
 
-	return newTransferPeer(region, oldPeer, newPeer)
+	return newTransferPeer(region, RegionKind, oldPeer, newPeer)
 }
 
 func newAddPeer(region *RegionInfo, peer *metapb.Peer) Operator {
@@ -254,7 +254,7 @@ func newRemovePeer(region *RegionInfo, peer *metapb.Peer) Operator {
 	return newRegionOperator(region, RegionKind, removePeer)
 }
 
-func newTransferPeer(region *RegionInfo, oldPeer, newPeer *metapb.Peer) Operator {
+func newTransferPeer(region *RegionInfo, kind ResourceKind, oldPeer, newPeer *metapb.Peer) Operator {
 	addPeer := newAddPeerOperator(region.GetId(), newPeer)
 	removePeer := newRemovePeerOperator(region.GetId(), oldPeer)
 	if region.Leader != nil && region.Leader.GetId() == oldPeer.GetId() {
@@ -263,15 +263,9 @@ func newTransferPeer(region *RegionInfo, oldPeer, newPeer *metapb.Peer) Operator
 			newLeader = follower
 		}
 		transferLeader := newTransferLeaderOperator(region.GetId(), region.Leader, newLeader)
-		return newRegionOperator(region, RegionKind, addPeer, transferLeader, removePeer)
+		return newRegionOperator(region, kind, addPeer, transferLeader, removePeer)
 	}
-	return newRegionOperator(region, RegionKind, addPeer, removePeer)
-}
-
-func newPriorityTransferPeer(region *RegionInfo, oldPeer, newPeer *metapb.Peer) Operator {
-	addPeer := newAddPeerOperator(region.GetId(), newPeer)
-	removePeer := newRemovePeerOperator(region.GetId(), oldPeer)
-	return newRegionOperator(region, PriorityKind, addPeer, removePeer)
+	return newRegionOperator(region, kind, addPeer, removePeer)
 }
 
 func newPriorityTransferLeader(region *RegionInfo, newLeader *metapb.Peer) Operator {