scheduler: support range for schedulers

server/cluster.go

@@ -634,8 +634,8 @@ func (c *RaftCluster) RandPendingRegion(storeID uint64, ranges []core.KeyRange,
 }
 
 // RandLearnerRegion returns a random region that has a learner peer on the store.
-func (c *RaftCluster) RandLearnerRegion(storeID uint64, opts ...core.RegionOption) *core.RegionInfo {
-	return c.core.RandLearnerRegion(storeID, opts...)
+func (c *RaftCluster) RandLearnerRegion(storeID uint64, ranges []core.KeyRange, opts ...core.RegionOption) *core.RegionInfo {
+	return c.core.RandLearnerRegion(storeID, ranges, opts...)
 }
 
 // RandHotRegionFromStore randomly picks a hot region in specified store.

server/cluster_test.go

@@ -833,10 +833,10 @@ func (s *testRegionsInfoSuite) Test(c *C) {
 	}
 
 	for i := uint64(0); i < n; i++ {
-		region := cache.RandLeaderRegion(i, []core.KeyRange{core.NewKeyRange("", "")}, core.HealthRegion())
+		region := cache.RandLeaderRegion(i, []core.KeyRange{core.NewKeyRange("", "")})
 		c.Assert(region.GetLeader().GetStoreId(), Equals, i)
 
-		region = cache.RandFollowerRegion(i, []core.KeyRange{core.NewKeyRange("", "")}, core.HealthRegion())
+		region = cache.RandFollowerRegion(i, []core.KeyRange{core.NewKeyRange("", "")})
 		c.Assert(region.GetLeader().GetStoreId(), Not(Equals), i)
 
 		c.Assert(region.GetStorePeer(i), NotNil)

server/core/basic_cluster.go

@@ -188,6 +188,9 @@ func (bc *BasicCluster) RandLearnerRegion(storeID uint64, ranges []KeyRange, opt
 
 func (bc *BasicCluster) selectRegion(regions []*RegionInfo, opts ...RegionOption) *RegionInfo {
 	for _, r := range regions {
+		if r == nil {
+			break
+		}
 		if slice.AllOf(opts, func(i int) bool { return opts[i](r) }) {
 			return r
 		}

server/core/region.go

@@ -505,9 +505,14 @@ func (rst *regionSubTree) RandomRegions(n int, startKey, endKey []byte) []*Regio
 	if rst.length() == 0 {
 		return nil
 	}
+
 	regions := make([]*RegionInfo, 0, n)
 	for i := 0; i < n; i++ {
-		regions = append(regions, rst.regionTree.RandomRegion(startKey, endKey))
+		region := rst.regionTree.RandomRegion(startKey, endKey)
+		if region == nil || !isInvolved(region, startKey, endKey) {
+			continue
+		}
+		regions = append(regions, region)
 	}
 	return regions
 }
@@ -741,42 +746,56 @@ func (r *RegionsInfo) GetStoreLearnerCount(storeID uint64) int {
 
 // RandPendingRegion randomly gets a store's region with a pending peer.
 func (r *RegionsInfo) RandPendingRegion(storeID uint64, ranges []KeyRange) *RegionInfo {
-	return r.pendingPeers[storeID].RandomRegion(ranges)
+	startKey, endKey := r.GetKeys(ranges)
+	return r.pendingPeers[storeID].RandomRegion(startKey, endKey)
 }
 
 // RandPendingRegions randomly gets a store's n regions with a pending peer.
 func (r *RegionsInfo) RandPendingRegions(storeID uint64, ranges []KeyRange, n int) []*RegionInfo {
-	return r.pendingPeers[storeID].RandomRegions(ranges, n)
+	startKey, endKey := r.GetKeys(ranges)
+	return r.pendingPeers[storeID].RandomRegions(n, startKey, endKey)
 }
 
 // RandLeaderRegion randomly gets a store's leader region.
 func (r *RegionsInfo) RandLeaderRegion(storeID uint64, ranges []KeyRange) *RegionInfo {
-	return r.leaders[storeID].RandomRegion(ranges)
+	startKey, endKey := r.GetKeys(ranges)
+	return r.leaders[storeID].RandomRegion(startKey, endKey)
 }
 
 // RandLeaderRegions randomly gets a store's n leader regions.
 func (r *RegionsInfo) RandLeaderRegions(storeID uint64, ranges []KeyRange, n int) []*RegionInfo {
-	return r.leaders[storeID].RandomRegions(ranges, n)
+	startKey, endKey := r.GetKeys(ranges)
+	return r.leaders[storeID].RandomRegions(n, startKey, endKey)
 }
 
 // RandFollowerRegion randomly gets a store's follower region.
 func (r *RegionsInfo) RandFollowerRegion(storeID uint64, ranges []KeyRange) *RegionInfo {
-	return r.followers[storeID].RandomRegion(ranges)
+	startKey, endKey := r.GetKeys(ranges)
+	return r.followers[storeID].RandomRegion(startKey, endKey)
 }
 
 // RandFollowerRegions randomly gets a store's n follower regions.
 func (r *RegionsInfo) RandFollowerRegions(storeID uint64, ranges []KeyRange, n int) []*RegionInfo {
-	return r.followers[storeID].RandomRegions(ranges, n)
+	startKey, endKey := r.GetKeys(ranges)
+	return r.followers[storeID].RandomRegions(n, startKey, endKey)
 }
 
 // RandLearnerRegion randomly gets a store's learner region.
 func (r *RegionsInfo) RandLearnerRegion(storeID uint64, ranges []KeyRange) *RegionInfo {
-	return r.learners[storeID].RandomRegion(ranges)
+	startKey, endKey := r.GetKeys(ranges)
+	return r.learners[storeID].RandomRegion(startKey, endKey)
 }
 
 // RandLearnerRegions randomly gets a store's n learner regions.
 func (r *RegionsInfo) RandLearnerRegions(storeID uint64, ranges []KeyRange, n int) []*RegionInfo {
-	return r.learners[storeID].RandomRegions(ranges, n)
+	startKey, endKey := r.GetKeys(ranges)
+	return r.learners[storeID].RandomRegions(n, startKey, endKey)
+}
+
+// GetKeys gets the start key and end key from random key range.
+func (r *RegionsInfo) GetKeys(ranges []KeyRange) ([]byte, []byte) {
+	idx := rand.Intn(len(ranges))
+	return ranges[idx].StartKey, ranges[idx].EndKey
 }
 
 // GetLeader return leader RegionInfo by storeID and regionID(now only used in test)
@@ -881,6 +900,10 @@ func DiffRegionKeyInfo(origin *RegionInfo, other *RegionInfo) string {
 	return strings.Join(ret, ", ")
 }
 
+func isInvolved(region *RegionInfo, startKey, endKey []byte) bool {
+	return bytes.Compare(region.GetStartKey(), startKey) >= 0 && (len(endKey) == 0 || (len(region.GetEndKey()) > 0 && bytes.Compare(region.GetEndKey(), endKey) <= 0))
+}
+
 // HexRegionKey converts region key to hex format. Used for formating region in
 // logs.
 func HexRegionKey(key []byte) []byte {

server/core/region_test.go

@@ -145,7 +145,7 @@ func BenchmarkRandomRegion(b *testing.B) {
 	}
 	b.ResetTimer()
 	for i := 0; i < b.N; i++ {
-		regions.RandLeaderRegion([]KeyRange{NewKeyRange("","")}, 1)
+		regions.RandLeaderRegion(1, []KeyRange{NewKeyRange("", "")})
 	}
 }