Merge bc33d82 into c774c99

Makefile

@@ -62,7 +62,7 @@ endif
 update:
 	which glide >/dev/null || curl https://glide.sh/get | sh
 	which glide-vc || go get -v -u github.com/sgotti/glide-vc
-	rm -r vendor && mv _vendor/vendor vendor || true
+	rm -rf vendor && mv _vendor/vendor vendor || true
 	rm -rf _vendor
 ifdef PKG
 	glide get --strip-vendor --skip-test ${PKG}

glide.yaml

@@ -92,7 +92,7 @@ import:
   - jsonpb
   - proto
 - package: github.com/google/btree
-  version: 7d79101e329e5a3adf994758c578dab82b90c017
+  version: 925471ac9e2131377a91e1595defec898166fe49
 - package: github.com/gorilla/context
   version: ~1.1.0
 - package: github.com/gorilla/mux

server/balancer_test.go

@@ -41,8 +41,7 @@ func (s *testBalancerSuite) SetUpSuite(c *C) {
 }
 
 func (s *testBalancerSuite) newClusterInfo(c *C) *clusterInfo {
-	clusterInfo := newClusterInfo(s.getRootPath())
-	clusterInfo.idAlloc = newMockIDAllocator()
+	clusterInfo := newClusterInfo(newMockIDAllocator())
 
 	// Set cluster info.
 	meta := &metapb.Cluster{
@@ -64,7 +63,7 @@ func (s *testBalancerSuite) newClusterInfo(c *C) *clusterInfo {
 
 		addr := fmt.Sprintf("127.0.0.1:%d", i)
 		store := s.newStore(c, id, addr)
-		clusterInfo.addStore(store)
+		clusterInfo.setStore(newStoreInfo(store))
 	}
 
 	// Add 1 peer, id will be 5.
@@ -77,10 +76,7 @@ func (s *testBalancerSuite) newClusterInfo(c *C) *clusterInfo {
 	c.Assert(err, IsNil)
 
 	region := s.newRegion(c, id, []byte{}, []byte{}, []*metapb.Peer{peer}, nil)
-	clusterInfo.regions.addRegion(region)
-
-	// Set leader store region.
-	clusterInfo.regions.leaders.update(region.GetId(), peer.GetStoreId())
+	clusterInfo.addRegion(newRegionInfo(region, peer))
 
 	stores := clusterInfo.getStores()
 	c.Assert(stores, HasLen, 4)
@@ -98,16 +94,14 @@ func (s *testBalancerSuite) updateStore(c *C, clusterInfo *clusterInfo, storeID
 		ReceivingSnapCount: receivingSnapCount,
 	}
 
-	ok := clusterInfo.updateStoreStatus(stats)
-	c.Assert(ok, IsTrue)
+	err := clusterInfo.handleStoreHeartbeat(stats)
+	c.Assert(err, IsNil)
 }
 
 func (s *testBalancerSuite) updateStoreState(c *C, clusterInfo *clusterInfo, storeID uint64, state metapb.StoreState) {
 	store := clusterInfo.getStore(storeID)
 	store.State = state
-	clusterInfo.addStore(store.Store)
-	ok := clusterInfo.updateStoreStatus(store.stats.StoreStats)
-	c.Assert(ok, IsTrue)
+	clusterInfo.setStore(store)
 }
 
 func (s *testBalancerSuite) addRegionPeer(c *C, clusterInfo *clusterInfo, storeID uint64, region *regionInfo) {