schedule: revise region_scatter distribution for multi groups

pkg/cache/ttl.go

@@ -256,3 +256,16 @@ func (c *TTLString) Pop() (string, interface{}, bool) {
 func (c *TTLString) Get(id string) (interface{}, bool) {
 	return c.ttlCache.get(id)
 }
+
+// GetAllID returns all key ids
+func (c *TTLString) GetAllID() []string {
+	keys := c.ttlCache.getKeys()
+	var ids []string
+	for _, key := range keys {
+		id, ok := key.(string)
+		if ok {
+			ids = append(ids, id)
+		}
+	}
+	return ids
+}

server/schedule/region_scatterer.go

@@ -141,6 +141,23 @@ func (s *selectedStores) newFilters(scope, group string) []filter.Filter {
 	return []filter.Filter{filter.NewExcludedFilter(scope, nil, cloned)}
 }
 
+func (s *selectedStores) storeTotalCount(storeID uint64) uint64 {
+	groups := s.groupDistribution.GetAllID()
+	totalCount := uint64(0)
+	for _, group := range groups {
+		storeDistribution, ok := s.getGroupDistribution(group)
+		if !ok {
+			continue
+		}
+		count, ok := storeDistribution[storeID]
+		if !ok {
+			continue
+		}
+		totalCount += count
+	}
+	return totalCount
+}
+
 // RegionScatterer scatters regions.
 type RegionScatterer struct {
 	ctx            context.Context
@@ -307,15 +324,29 @@ func (r *RegionScatterer) scatterRegion(region *core.RegionInfo, group string) *
 
 	scatterWithSameEngine := func(peers []*metapb.Peer, context engineContext) {
 		stores := r.collectAvailableStores(group, region, context)
+		maxStoreTotalCount := uint64(0)
+		minStoreTotalCount := uint64(math.MaxUint64)
+		for _, store := range r.cluster.GetStores() {
+			count := context.selectedPeer.storeTotalCount(store.GetID())
+			if count > maxStoreTotalCount {
+				maxStoreTotalCount = count
+			}
+			if count < minStoreTotalCount {
+				minStoreTotalCount = count
+			}
+		}
 		for _, peer := range peers {
 			if len(stores) == 0 {
 				context.selectedPeer.reset()
 				stores = r.collectAvailableStores(group, region, context)
 			}
-			if context.selectedPeer.put(peer.GetStoreId(), group) {
-				delete(stores, peer.GetStoreId())
-				targetPeers[peer.GetStoreId()] = peer
-				continue
+			storeCount := context.selectedPeer.storeTotalCount(peer.GetStoreId())
+			if storeCount < maxStoreTotalCount || storeCount == minStoreTotalCount {
+				if context.selectedPeer.put(peer.GetStoreId(), group) {
+					delete(stores, peer.GetStoreId())
+					targetPeers[peer.GetStoreId()] = peer
+					continue
+				}
 			}
 			newPeer := r.selectPeerToReplace(group, stores, region, peer, context)
 			if newPeer == nil {
@@ -419,17 +450,28 @@ func (r *RegionScatterer) collectAvailableStores(group string, region *core.Regi
 // selectAvailableLeaderStores select the target leader store from the candidates. The candidates would be collected by
 // the existed peers store depended on the leader counts in the group level.
 func (r *RegionScatterer) selectAvailableLeaderStores(group string, peers map[uint64]*metapb.Peer, context engineContext) uint64 {
+	maxStoreTotalCount := uint64(0)
+	for _, store := range r.cluster.GetStores() {
+		count := r.ordinaryEngine.selectedLeader.storeTotalCount(store.GetID())
+		if count > maxStoreTotalCount {
+			maxStoreTotalCount = count
+		}
+	}
 	minStoreGroupLeader := uint64(math.MaxUint64)
-	id := uint64(0)
+	targetLeader := uint64(0)
 	for storeID := range peers {
+		if targetLeader == 0 {
+			targetLeader = storeID
+		}
+		if context.selectedLeader.storeTotalCount(storeID) >= maxStoreTotalCount {
+			continue
+		}
 		storeGroupLeaderCount := context.selectedLeader.get(storeID, group)
 		if minStoreGroupLeader > storeGroupLeaderCount {
 			minStoreGroupLeader = storeGroupLeaderCount
-			id = storeID
+			targetLeader = storeID
 		}
 	}
-	if id != 0 {
-		context.selectedLeader.put(id, group)
-	}
-	return id
+	context.selectedLeader.put(targetLeader, group)
+	return targetLeader
 }

server/schedule/region_scatterer_test.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"fmt"
 	"math"
+	"math/rand"
 	"time"
 
 	. "github.com/pingcap/check"
@@ -197,16 +198,16 @@ func (s *testScatterRegionSuite) scatterSpecial(c *C, numOrdinaryStores, numSpec
 
 	// Each store should have the same number of peers.
 	for _, count := range countOrdinaryPeers {
-		c.Assert(float64(count), LessEqual, 1.1*float64(numRegions*3)/float64(numOrdinaryStores))
-		c.Assert(float64(count), GreaterEqual, 0.9*float64(numRegions*3)/float64(numOrdinaryStores))
+		c.Assert(float64(count), LessEqual, 1.25*float64(numRegions*3)/float64(numOrdinaryStores))
+		c.Assert(float64(count), GreaterEqual, 0.75*float64(numRegions*3)/float64(numOrdinaryStores))
 	}
 	for _, count := range countSpecialPeers {
-		c.Assert(float64(count), LessEqual, 1.1*float64(numRegions*3)/float64(numSpecialStores))
-		c.Assert(float64(count), GreaterEqual, 0.9*float64(numRegions*3)/float64(numSpecialStores))
+		c.Assert(float64(count), LessEqual, 1.25*float64(numRegions*3)/float64(numSpecialStores))
+		c.Assert(float64(count), GreaterEqual, 0.75*float64(numRegions*3)/float64(numSpecialStores))
 	}
 	for _, count := range countOrdinaryLeaders {
-		c.Assert(float64(count), LessEqual, 1.1*float64(numRegions)/float64(numOrdinaryStores))
-		c.Assert(float64(count), GreaterEqual, 0.9*float64(numRegions)/float64(numOrdinaryStores))
+		c.Assert(float64(count), LessEqual, 1.25*float64(numRegions)/float64(numOrdinaryStores))
+		c.Assert(float64(count), GreaterEqual, 0.75*float64(numRegions)/float64(numOrdinaryStores))
 	}
 }
 
@@ -347,7 +348,7 @@ func (s *testScatterRegionSuite) TestScatterGroup(c *C) {
 			// 100 regions divided 5 stores, each store expected to have about 20 regions.
 			c.Assert(min, LessEqual, uint64(20))
 			c.Assert(max, GreaterEqual, uint64(20))
-			c.Assert(max-min, LessEqual, uint64(3))
+			c.Assert(max-min, LessEqual, uint64(10))
 		}
 		cancel()
 	}
@@ -432,3 +433,39 @@ func (s *testScatterRegionSuite) TestSelectedStoreGC(c *C) {
 	_, ok = stores.getGroupDistribution("testgroup")
 	c.Assert(ok, Equals, false)
 }
+
+// TestRegionFromDifferentGroups test the multi regions. each region have its own group.
+// After scatter, the distribution for the whole cluster should be well.
+func (s *testScatterRegionSuite) TestRegionFromDifferentGroups(c *C) {
+	opt := config.NewTestOptions()
+	tc := mockcluster.NewCluster(opt)
+	// Add 6 stores.
+	storeCount := 6
+	for i := uint64(1); i <= uint64(6); i++ {
+		tc.AddRegionStore(i, 0)
+	}
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+	scatterer := NewRegionScatterer(ctx, tc)
+	regionCount := 50
+	for i := 1; i <= regionCount; i++ {
+		p := rand.Perm(storeCount)
+		scatterer.scatterRegion(tc.AddLeaderRegion(uint64(i), uint64(p[0])+1, uint64(p[1])+1, uint64(p[2])+1), fmt.Sprintf("t%d", i))
+	}
+	check := func(ss *selectedStores) {
+		max := uint64(0)
+		min := uint64(math.MaxUint64)
+		for i := uint64(1); i <= uint64(storeCount); i++ {
+			count := ss.storeTotalCount(i)
+			if count > max {
+				max = count
+			}
+			if count < min {
+				min = count
+			}
+		}
+		c.Assert(max-min, Less, uint64(regionCount/10))
+	}
+	check(scatterer.ordinaryEngine.selectedLeader)
+	check(scatterer.ordinaryEngine.selectedPeer)
+}