schedule: fix race for region scatter (tikv#3760)

Signed-off-by: yisaer <disxiaofei@163.com>
bufferflies · Jul 9, 2021 · c042cbc · c042cbc
1 parent a993629
commit c042cbc
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 12 deletions.
diff --git a/server/schedule/region_scatterer.go b/server/schedule/region_scatterer.go
@@ -85,15 +85,10 @@ func (s *selectedStores) GetGroupDistribution(group string) (map[uint64]uint64,
 	return s.getDistributionByGroupLocked(group)
 }
 
-// getDistributionByGroupLocked should be called with lock
-func (s *selectedStores) getDistributionByGroupLocked(group string) (map[uint64]uint64, bool) {
-	if result, ok := s.groupDistribution.Get(group); ok {
-		return result.(map[uint64]uint64), true
-	}
-	return nil, false
-}
-
-func (s *selectedStores) totalCountByStore(storeID uint64) uint64 {
+// TotalCountByStore counts the total count by store
+func (s *selectedStores) TotalCountByStore(storeID uint64) uint64 {
+	s.mu.RLock()
+	defer s.mu.RUnlock()
 	groups := s.groupDistribution.GetAllID()
 	totalCount := uint64(0)
 	for _, group := range groups {
@@ -110,6 +105,14 @@ func (s *selectedStores) totalCountByStore(storeID uint64) uint64 {
 	return totalCount
 }
 
+// getDistributionByGroupLocked should be called with lock
+func (s *selectedStores) getDistributionByGroupLocked(group string) (map[uint64]uint64, bool) {
+	if result, ok := s.groupDistribution.Get(group); ok {
+		return result.(map[uint64]uint64), true
+	}
+	return nil, false
+}
+
 // RegionScatterer scatters regions.
 type RegionScatterer struct {
 	ctx            context.Context
@@ -347,7 +350,7 @@ func (r *RegionScatterer) selectCandidates(region *core.RegionInfo, sourceStoreI
 	maxStoreTotalCount := uint64(0)
 	minStoreTotalCount := uint64(math.MaxUint64)
 	for _, store := range r.cluster.GetStores() {
-		count := context.selectedPeer.totalCountByStore(store.GetID())
+		count := context.selectedPeer.TotalCountByStore(store.GetID())
 		if count > maxStoreTotalCount {
 			maxStoreTotalCount = count
 		}
@@ -356,7 +359,7 @@ func (r *RegionScatterer) selectCandidates(region *core.RegionInfo, sourceStoreI
 		}
 	}
 	for _, store := range stores {
-		storeCount := context.selectedPeer.totalCountByStore(store.GetID())
+		storeCount := context.selectedPeer.TotalCountByStore(store.GetID())
 		// If storeCount is equal to the maxStoreTotalCount, we should skip this store as candidate.
 		// If the storeCount are all the same for the whole cluster(maxStoreTotalCount == minStoreTotalCount), any store
 		// could be selected as candidate.

diff --git a/server/schedule/region_scatterer_test.go b/server/schedule/region_scatterer_test.go
@@ -462,7 +462,7 @@ func (s *testScatterRegionSuite) TestRegionFromDifferentGroups(c *C) {
 		max := uint64(0)
 		min := uint64(math.MaxUint64)
 		for i := uint64(1); i <= uint64(storeCount); i++ {
-			count := ss.totalCountByStore(i)
+			count := ss.TotalCountByStore(i)
 			if count > max {
 				max = count
 			}