diff --git a/server/cluster.go b/server/cluster.go index 9196cbe78060..01f52610c469 100644 --- a/server/cluster.go +++ b/server/cluster.go @@ -614,18 +614,18 @@ func (c *RaftCluster) GetStoreRegions(storeID uint64) []*core.RegionInfo { } // RandLeaderRegion returns a random region that has leader on the store. -func (c *RaftCluster) RandLeaderRegion(storeID uint64, opts ...core.RegionOption) *core.RegionInfo { - return c.core.RandLeaderRegion(storeID, opts...) +func (c *RaftCluster) RandLeaderRegion(storeID uint64, startKey, endKey []byte, opts ...core.RegionOption) *core.RegionInfo { + return c.core.RandLeaderRegion(storeID, startKey, endKey, opts...) } // RandFollowerRegion returns a random region that has a follower on the store. -func (c *RaftCluster) RandFollowerRegion(storeID uint64, opts ...core.RegionOption) *core.RegionInfo { - return c.core.RandFollowerRegion(storeID, opts...) +func (c *RaftCluster) RandFollowerRegion(storeID uint64, startKey, endKey []byte, opts ...core.RegionOption) *core.RegionInfo { + return c.core.RandFollowerRegion(storeID, startKey, endKey, opts...) } // RandPendingRegion returns a random region that has a pending peer on the store. -func (c *RaftCluster) RandPendingRegion(storeID uint64, opts ...core.RegionOption) *core.RegionInfo { - return c.core.RandPendingRegion(storeID, opts...) +func (c *RaftCluster) RandPendingRegion(storeID uint64, startKey, endKey []byte, opts ...core.RegionOption) *core.RegionInfo { + return c.core.RandPendingRegion(storeID, startKey, endKey, opts...) } // RandHotRegionFromStore randomly picks a hot region in specified store. diff --git a/server/cluster_test.go b/server/cluster_test.go index 931a10183d88..8ccedf00327e 100644 --- a/server/cluster_test.go +++ b/server/cluster_test.go @@ -841,10 +841,10 @@ func (s *testRegionsInfoSuite) Test(c *C) { } for i := uint64(0); i < n; i++ { - region := cache.RandLeaderRegion(i, core.HealthRegion()) + region := cache.RandLeaderRegion(i, []byte(""), []byte(""), core.HealthRegion()) c.Assert(region.GetLeader().GetStoreId(), Equals, i) - region = cache.RandFollowerRegion(i, core.HealthRegion()) + region = cache.RandFollowerRegion(i, []byte(""), []byte(""), core.HealthRegion()) c.Assert(region.GetLeader().GetStoreId(), Not(Equals), i) c.Assert(region.GetStorePeer(i), NotNil) @@ -860,14 +860,14 @@ func (s *testRegionsInfoSuite) Test(c *C) { // All regions will be filtered out if they have pending peers. for i := uint64(0); i < n; i++ { for j := 0; j < cache.GetStoreLeaderCount(i); j++ { - region := cache.RandLeaderRegion(i, core.HealthRegion()) + region := cache.RandLeaderRegion(i, []byte(""), []byte(""), core.HealthRegion()) newRegion := region.Clone(core.WithPendingPeers(region.GetPeers())) cache.SetRegion(newRegion) } - c.Assert(cache.RandLeaderRegion(i, core.HealthRegion()), IsNil) + c.Assert(cache.RandLeaderRegion(i, []byte(""), []byte(""), core.HealthRegion()), IsNil) } for i := uint64(0); i < n; i++ { - c.Assert(cache.RandFollowerRegion(i, core.HealthRegion()), IsNil) + c.Assert(cache.RandFollowerRegion(i, []byte(""), []byte(""), core.HealthRegion()), IsNil) } } diff --git a/server/core/basic_cluster.go b/server/core/basic_cluster.go index 21dde25d3d24..393166361d03 100644 --- a/server/core/basic_cluster.go +++ b/server/core/basic_cluster.go @@ -152,24 +152,24 @@ func (bc *BasicCluster) UpdateStoreStatus(storeID uint64, leaderCount int, regio } // RandFollowerRegion returns a random region that has a follower on the store. -func (bc *BasicCluster) RandFollowerRegion(storeID uint64, opts ...RegionOption) *RegionInfo { +func (bc *BasicCluster) RandFollowerRegion(storeID uint64, startKey, endKey []byte, opts ...RegionOption) *RegionInfo { bc.RLock() defer bc.RUnlock() - return bc.Regions.RandFollowerRegion(storeID, opts...) + return bc.Regions.RandFollowerRegion(storeID, startKey, endKey, opts...) } // RandLeaderRegion returns a random region that has leader on the store. -func (bc *BasicCluster) RandLeaderRegion(storeID uint64, opts ...RegionOption) *RegionInfo { +func (bc *BasicCluster) RandLeaderRegion(storeID uint64, startKey, endKey []byte, opts ...RegionOption) *RegionInfo { bc.RLock() defer bc.RUnlock() - return bc.Regions.RandLeaderRegion(storeID, opts...) + return bc.Regions.RandLeaderRegion(storeID, startKey, endKey, opts...) } // RandPendingRegion returns a random region that has a pending peer on the store. -func (bc *BasicCluster) RandPendingRegion(storeID uint64, opts ...RegionOption) *RegionInfo { +func (bc *BasicCluster) RandPendingRegion(storeID uint64, startKey, endKey []byte, opts ...RegionOption) *RegionInfo { bc.RLock() defer bc.RUnlock() - return bc.Regions.RandPendingRegion(storeID, opts...) + return bc.Regions.RandPendingRegion(storeID, startKey, endKey, opts...) } // GetRegionCount gets the total count of RegionInfo of regionMap. @@ -308,9 +308,9 @@ func (bc *BasicCluster) Length() int { // RegionSetInformer provides access to a shared informer of regions. type RegionSetInformer interface { - RandFollowerRegion(storeID uint64, opts ...RegionOption) *RegionInfo - RandLeaderRegion(storeID uint64, opts ...RegionOption) *RegionInfo - RandPendingRegion(storeID uint64, opts ...RegionOption) *RegionInfo + RandFollowerRegion(storeID uint64, startKey, endKey []byte, opts ...RegionOption) *RegionInfo + RandLeaderRegion(storeID uint64, startKey, endKey []byte, opts ...RegionOption) *RegionInfo + RandPendingRegion(storeID uint64, startKey, endKey []byte, opts ...RegionOption) *RegionInfo GetAverageRegionSize() int64 GetStoreRegionCount(storeID uint64) int GetRegion(id uint64) *RegionInfo diff --git a/server/core/region.go b/server/core/region.go index 83c1c3f1a8a0..1bb7b48e0223 100644 --- a/server/core/region.go +++ b/server/core/region.go @@ -750,23 +750,23 @@ func (r *RegionsInfo) GetStoreLearnerCount(storeID uint64) int { } // RandRegion get a region by random -func (r *RegionsInfo) RandRegion(opts ...RegionOption) *RegionInfo { - return randRegion(r.tree, opts...) +func (r *RegionsInfo) RandRegion(startKey, endKey []byte, opts ...RegionOption) *RegionInfo { + return randRegion(r.tree, startKey, endKey, opts...) } // RandPendingRegion randomly gets a store's region with a pending peer. -func (r *RegionsInfo) RandPendingRegion(storeID uint64, opts ...RegionOption) *RegionInfo { - return randRegion(r.pendingPeers[storeID], opts...) +func (r *RegionsInfo) RandPendingRegion(storeID uint64, startKey, endKey []byte, opts ...RegionOption) *RegionInfo { + return randRegion(r.pendingPeers[storeID], startKey, endKey, opts...) } // RandLeaderRegion randomly gets a store's leader region. -func (r *RegionsInfo) RandLeaderRegion(storeID uint64, opts ...RegionOption) *RegionInfo { - return randRegion(r.leaders[storeID], opts...) +func (r *RegionsInfo) RandLeaderRegion(storeID uint64, startKey, endKey []byte, opts ...RegionOption) *RegionInfo { + return randRegion(r.leaders[storeID], startKey, endKey, opts...) } // RandFollowerRegion randomly gets a store's follower region. -func (r *RegionsInfo) RandFollowerRegion(storeID uint64, opts ...RegionOption) *RegionInfo { - return randRegion(r.followers[storeID], opts...) +func (r *RegionsInfo) RandFollowerRegion(storeID uint64, startKey, endKey []byte, opts ...RegionOption) *RegionInfo { + return randRegion(r.followers[storeID], startKey, endKey, opts...) } // GetLeader return leader RegionInfo by storeID and regionID(now only used in test) @@ -831,9 +831,9 @@ type RegionsContainer interface { RandomRegion(startKey, endKey []byte) *RegionInfo } -func randRegion(regions RegionsContainer, opts ...RegionOption) *RegionInfo { +func randRegion(regions RegionsContainer, startKey, endKey []byte, opts ...RegionOption) *RegionInfo { for i := 0; i < randomRegionMaxRetry; i++ { - region := regions.RandomRegion(nil, nil) + region := regions.RandomRegion(startKey, endKey) if region == nil { return nil } diff --git a/server/core/region_test.go b/server/core/region_test.go index 686731a1b8e9..d94d20cabd40 100644 --- a/server/core/region_test.go +++ b/server/core/region_test.go @@ -148,7 +148,7 @@ func BenchmarkRandomRegion(b *testing.B) { } b.ResetTimer() for i := 0; i < b.N; i++ { - regions.RandRegion() + regions.RandRegion([]byte(""), []byte("")) } } diff --git a/server/namespace_cluster.go b/server/namespace_cluster.go index 3b21816dfef5..2527815259e8 100644 --- a/server/namespace_cluster.go +++ b/server/namespace_cluster.go @@ -62,9 +62,9 @@ func (c *namespaceCluster) checkRegion(region *core.RegionInfo) bool { const randRegionMaxRetry = 10 // RandFollowerRegion returns a random region that has a follower on the store. -func (c *namespaceCluster) RandFollowerRegion(storeID uint64, opts ...core.RegionOption) *core.RegionInfo { +func (c *namespaceCluster) RandFollowerRegion(storeID uint64, startKey, endKey []byte, opts ...core.RegionOption) *core.RegionInfo { for i := 0; i < randRegionMaxRetry; i++ { - r := c.Cluster.RandFollowerRegion(storeID, opts...) + r := c.Cluster.RandFollowerRegion(storeID, startKey, endKey, opts...) if r == nil { return nil } @@ -76,9 +76,9 @@ func (c *namespaceCluster) RandFollowerRegion(storeID uint64, opts ...core.Regio } // RandLeaderRegion returns a random region that has leader on the store. -func (c *namespaceCluster) RandLeaderRegion(storeID uint64, opts ...core.RegionOption) *core.RegionInfo { +func (c *namespaceCluster) RandLeaderRegion(storeID uint64, startKey, endKey []byte, opts ...core.RegionOption) *core.RegionInfo { for i := 0; i < randRegionMaxRetry; i++ { - r := c.Cluster.RandLeaderRegion(storeID, opts...) + r := c.Cluster.RandLeaderRegion(storeID, startKey, endKey, opts...) if r == nil { return nil } diff --git a/server/schedule/range_cluster.go b/server/schedule/range_cluster.go index 30d2256027cd..ba6914265c6e 100644 --- a/server/schedule/range_cluster.go +++ b/server/schedule/range_cluster.go @@ -101,13 +101,13 @@ func (r *RangeCluster) GetTolerantSizeRatio() float64 { } // RandFollowerRegion returns a random region that has a follower on the store. -func (r *RangeCluster) RandFollowerRegion(storeID uint64, opts ...core.RegionOption) *core.RegionInfo { - return r.regions.RandFollowerRegion(storeID, opts...) +func (r *RangeCluster) RandFollowerRegion(storeID uint64, startKey, endKey []byte, opts ...core.RegionOption) *core.RegionInfo { + return r.regions.RandFollowerRegion(storeID, startKey, endKey, opts...) } // RandLeaderRegion returns a random region that has leader on the store. -func (r *RangeCluster) RandLeaderRegion(storeID uint64, opts ...core.RegionOption) *core.RegionInfo { - return r.regions.RandLeaderRegion(storeID, opts...) +func (r *RangeCluster) RandLeaderRegion(storeID uint64, startKey, endKey []byte, opts ...core.RegionOption) *core.RegionInfo { + return r.regions.RandLeaderRegion(storeID, startKey, endKey, opts...) } // GetAverageRegionSize returns the average region approximate size. diff --git a/server/schedulers/balance_leader.go b/server/schedulers/balance_leader.go index 1aa8b9cae1a8..8f303a736fd8 100644 --- a/server/schedulers/balance_leader.go +++ b/server/schedulers/balance_leader.go @@ -29,7 +29,11 @@ import ( func init() { schedule.RegisterScheduler("balance-leader", func(opController *schedule.OperatorController, args []string) (schedule.Scheduler, error) { - return newBalanceLeaderScheduler(opController), nil + startKey, endKey, err := getRangeKeys(args) + if err != nil { + return nil, err + } + return newBalanceLeaderScheduler(opController, []string{startKey, endKey}), nil }) } @@ -39,6 +43,8 @@ const balanceLeaderRetryLimit = 10 type balanceLeaderScheduler struct { *baseScheduler name string + startKey []byte + endKey []byte selector *selector.BalanceSelector taintStores *cache.TTLUint64 opController *schedule.OperatorController @@ -47,12 +53,14 @@ type balanceLeaderScheduler struct { // newBalanceLeaderScheduler creates a scheduler that tends to keep leaders on // each store balanced. -func newBalanceLeaderScheduler(opController *schedule.OperatorController, opts ...BalanceLeaderCreateOption) schedule.Scheduler { +func newBalanceLeaderScheduler(opController *schedule.OperatorController, args []string, opts ...BalanceLeaderCreateOption) schedule.Scheduler { taintStores := newTaintCache() base := newBaseScheduler(opController) s := &balanceLeaderScheduler{ baseScheduler: base, + startKey: []byte(args[0]), + endKey: []byte(args[1]), taintStores: taintStores, opController: opController, counter: balanceLeaderCounter, @@ -160,7 +168,7 @@ func (l *balanceLeaderScheduler) Schedule(cluster schedule.Cluster) []*operator. // the best follower peer and transfers the leader. func (l *balanceLeaderScheduler) transferLeaderOut(cluster schedule.Cluster, source *core.StoreInfo) []*operator.Operator { sourceID := source.GetID() - region := cluster.RandLeaderRegion(sourceID, core.HealthRegion()) + region := cluster.RandLeaderRegion(sourceID, l.startKey, l.endKey, core.HealthRegion()) if region == nil { log.Debug("store has no leader", zap.String("scheduler", l.GetName()), zap.Uint64("store-id", sourceID)) schedulerCounter.WithLabelValues(l.GetName(), "no-leader-region").Inc() @@ -180,7 +188,7 @@ func (l *balanceLeaderScheduler) transferLeaderOut(cluster schedule.Cluster, sou // the worst follower peer and transfers the leader. func (l *balanceLeaderScheduler) transferLeaderIn(cluster schedule.Cluster, target *core.StoreInfo) []*operator.Operator { targetID := target.GetID() - region := cluster.RandFollowerRegion(targetID, core.HealthRegion()) + region := cluster.RandFollowerRegion(targetID, l.startKey, l.endKey, core.HealthRegion()) if region == nil { log.Debug("store has no follower", zap.String("scheduler", l.GetName()), zap.Uint64("store-id", targetID)) schedulerCounter.WithLabelValues(l.GetName(), "no-follower-region").Inc() diff --git a/server/schedulers/balance_region.go b/server/schedulers/balance_region.go index 4fcef5498806..4639d026dc9b 100644 --- a/server/schedulers/balance_region.go +++ b/server/schedulers/balance_region.go @@ -32,7 +32,11 @@ import ( func init() { schedule.RegisterScheduler("balance-region", func(opController *schedule.OperatorController, args []string) (schedule.Scheduler, error) { - return newBalanceRegionScheduler(opController), nil + startKey, endKey, err := getRangeKeys(args) + if err != nil { + return nil, err + } + return newBalanceRegionScheduler(opController, []string{startKey, endKey}), nil }) } @@ -53,6 +57,8 @@ const ( type balanceRegionScheduler struct { *baseScheduler name string + startKey []byte + endKey []byte selector *selector.BalanceSelector opController *schedule.OperatorController hitsCounter *hitsStoreBuilder @@ -61,10 +67,12 @@ type balanceRegionScheduler struct { // newBalanceRegionScheduler creates a scheduler that tends to keep regions on // each store balanced. -func newBalanceRegionScheduler(opController *schedule.OperatorController, opts ...BalanceRegionCreateOption) schedule.Scheduler { +func newBalanceRegionScheduler(opController *schedule.OperatorController, args []string, opts ...BalanceRegionCreateOption) schedule.Scheduler { base := newBaseScheduler(opController) s := &balanceRegionScheduler{ baseScheduler: base, + startKey: []byte(args[0]), + endKey: []byte(args[1]), opController: opController, hitsCounter: newHitsStoreBuilder(hitsStoreTTL, hitsStoreCountThreshold), counter: balanceRegionCounter, @@ -135,14 +143,14 @@ func (s *balanceRegionScheduler) Schedule(cluster schedule.Cluster) []*operator. for i := 0; i < balanceRegionRetryLimit; i++ { // Priority picks the region that has a pending peer. // Pending region may means the disk is overload, remove the pending region firstly. - region := cluster.RandPendingRegion(sourceID, core.HealthRegionAllowPending()) + region := cluster.RandPendingRegion(sourceID, s.startKey, s.endKey, core.HealthRegionAllowPending()) if region == nil { // Then picks the region that has a follower in the source store. - region = cluster.RandFollowerRegion(sourceID, core.HealthRegion()) + region = cluster.RandFollowerRegion(sourceID, s.startKey, s.endKey, core.HealthRegion()) } if region == nil { // Last, picks the region has the leader in the source store. - region = cluster.RandLeaderRegion(sourceID, core.HealthRegion()) + region = cluster.RandLeaderRegion(sourceID, s.startKey, s.endKey, core.HealthRegion()) } if region == nil { schedulerCounter.WithLabelValues(s.GetName(), "no-region").Inc() diff --git a/server/schedulers/evict_leader.go b/server/schedulers/evict_leader.go index e7716161438e..d332fdb1c099 100644 --- a/server/schedulers/evict_leader.go +++ b/server/schedulers/evict_leader.go @@ -27,19 +27,26 @@ import ( func init() { schedule.RegisterScheduler("evict-leader", func(opController *schedule.OperatorController, args []string) (schedule.Scheduler, error) { - if len(args) != 1 { + if len(args) == 0 { return nil, errors.New("evict-leader needs 1 argument") } id, err := strconv.ParseUint(args[0], 10, 64) if err != nil { return nil, errors.WithStack(err) } - return newEvictLeaderScheduler(opController, id), nil + + startKey, endKey, err := getRangeKeys(args[1:]) + if err != nil { + return nil, err + } + return newEvictLeaderScheduler(opController, id, []string{startKey, endKey}), nil }) } type evictLeaderScheduler struct { *baseScheduler + startKey []byte + endKey []byte name string storeID uint64 selector *selector.RandomSelector @@ -47,7 +54,7 @@ type evictLeaderScheduler struct { // newEvictLeaderScheduler creates an admin scheduler that transfers all leaders // out of a store. -func newEvictLeaderScheduler(opController *schedule.OperatorController, storeID uint64) schedule.Scheduler { +func newEvictLeaderScheduler(opController *schedule.OperatorController, storeID uint64, args []string) schedule.Scheduler { name := fmt.Sprintf("evict-leader-scheduler-%d", storeID) filters := []filter.Filter{ filter.StoreStateFilter{ActionScope: name, TransferLeader: true}, @@ -55,6 +62,8 @@ func newEvictLeaderScheduler(opController *schedule.OperatorController, storeID base := newBaseScheduler(opController) return &evictLeaderScheduler{ baseScheduler: base, + startKey: []byte(args[0]), + endKey: []byte(args[1]), name: name, storeID: storeID, selector: selector.NewRandomSelector(filters), @@ -83,7 +92,7 @@ func (s *evictLeaderScheduler) IsScheduleAllowed(cluster schedule.Cluster) bool func (s *evictLeaderScheduler) Schedule(cluster schedule.Cluster) []*operator.Operator { schedulerCounter.WithLabelValues(s.GetName(), "schedule").Inc() - region := cluster.RandLeaderRegion(s.storeID, core.HealthRegion()) + region := cluster.RandLeaderRegion(s.storeID, s.startKey, s.endKey, core.HealthRegion()) if region == nil { schedulerCounter.WithLabelValues(s.GetName(), "no-leader").Inc() return nil diff --git a/server/schedulers/grant_leader.go b/server/schedulers/grant_leader.go index 44fa55baa782..fb65452c1450 100644 --- a/server/schedulers/grant_leader.go +++ b/server/schedulers/grant_leader.go @@ -25,30 +25,39 @@ import ( func init() { schedule.RegisterScheduler("grant-leader", func(opController *schedule.OperatorController, args []string) (schedule.Scheduler, error) { - if len(args) != 1 { + if len(args) == 0 { return nil, errors.New("grant-leader needs 1 argument") } id, err := strconv.ParseUint(args[0], 10, 64) if err != nil { return nil, errors.WithStack(err) } - return newGrantLeaderScheduler(opController, id), nil + + startKey, endKey, err := getRangeKeys(args[1:]) + if err != nil { + return nil, err + } + return newGrantLeaderScheduler(opController, id, []string{startKey, endKey}), nil }) } // grantLeaderScheduler transfers all leaders to peers in the store. type grantLeaderScheduler struct { *baseScheduler - name string - storeID uint64 + startKey []byte + endKey []byte + name string + storeID uint64 } // newGrantLeaderScheduler creates an admin scheduler that transfers all leaders // to a store. -func newGrantLeaderScheduler(opController *schedule.OperatorController, storeID uint64) schedule.Scheduler { +func newGrantLeaderScheduler(opController *schedule.OperatorController, storeID uint64, args []string) schedule.Scheduler { base := newBaseScheduler(opController) return &grantLeaderScheduler{ baseScheduler: base, + startKey: []byte(args[0]), + endKey: []byte(args[1]), name: fmt.Sprintf("grant-leader-scheduler-%d", storeID), storeID: storeID, } @@ -75,7 +84,7 @@ func (s *grantLeaderScheduler) IsScheduleAllowed(cluster schedule.Cluster) bool func (s *grantLeaderScheduler) Schedule(cluster schedule.Cluster) []*operator.Operator { schedulerCounter.WithLabelValues(s.GetName(), "schedule").Inc() - region := cluster.RandFollowerRegion(s.storeID, core.HealthRegion()) + region := cluster.RandFollowerRegion(s.storeID, s.startKey, s.endKey, core.HealthRegion()) if region == nil { schedulerCounter.WithLabelValues(s.GetName(), "no-follower").Inc() return nil diff --git a/server/schedulers/label.go b/server/schedulers/label.go index be1cc0d547db..7d3485bba4ac 100644 --- a/server/schedulers/label.go +++ b/server/schedulers/label.go @@ -26,14 +26,20 @@ import ( func init() { schedule.RegisterScheduler("label", func(opController *schedule.OperatorController, args []string) (schedule.Scheduler, error) { - return newLabelScheduler(opController), nil + startKey, endKey, err := getRangeKeys(args) + if err != nil { + return nil, err + } + return newLabelScheduler(opController, []string{startKey, endKey}), nil }) } const labelSchedulerName = "label-scheduler" type labelScheduler struct { - name string + name string + startKey []byte + endKey []byte *baseScheduler selector *selector.BalanceSelector } @@ -41,11 +47,13 @@ type labelScheduler struct { // LabelScheduler is mainly based on the store's label information for scheduling. // Now only used for reject leader schedule, that will move the leader out of // the store with the specific label. -func newLabelScheduler(opController *schedule.OperatorController) schedule.Scheduler { +func newLabelScheduler(opController *schedule.OperatorController, args []string) schedule.Scheduler { filters := []filter.Filter{ filter.StoreStateFilter{ActionScope: labelSchedulerName, TransferLeader: true}, } return &labelScheduler{ + startKey: []byte(args[0]), + endKey: []byte(args[1]), name: labelSchedulerName, baseScheduler: newBaseScheduler(opController), selector: selector.NewBalanceSelector(core.LeaderKind, filters), @@ -79,7 +87,7 @@ func (s *labelScheduler) Schedule(cluster schedule.Cluster) []*operator.Operator } log.Debug("label scheduler reject leader store list", zap.Reflect("stores", rejectLeaderStores)) for id := range rejectLeaderStores { - if region := cluster.RandLeaderRegion(id); region != nil { + if region := cluster.RandLeaderRegion(id, s.startKey, s.endKey); region != nil { log.Debug("label scheduler selects region to transfer leader", zap.Uint64("region-id", region.GetID())) excludeStores := make(map[uint64]struct{}) for _, p := range region.GetDownPeers() { diff --git a/server/schedulers/random_merge.go b/server/schedulers/random_merge.go index 7c44fb723435..15afa69af718 100644 --- a/server/schedulers/random_merge.go +++ b/server/schedulers/random_merge.go @@ -25,7 +25,11 @@ import ( func init() { schedule.RegisterScheduler("random-merge", func(opController *schedule.OperatorController, args []string) (schedule.Scheduler, error) { - return newRandomMergeScheduler(opController), nil + startKey, endKey, err := getRangeKeys(args) + if err != nil { + return nil, err + } + return newRandomMergeScheduler(opController, []string{startKey, endKey}), nil }) } @@ -34,12 +38,14 @@ const randomMergeName = "random-merge-scheduler" type randomMergeScheduler struct { name string *baseScheduler + startKey []byte + endKey []byte selector *selector.RandomSelector } // newRandomMergeScheduler creates an admin scheduler that randomly picks two adjacent regions // then merges them. -func newRandomMergeScheduler(opController *schedule.OperatorController) schedule.Scheduler { +func newRandomMergeScheduler(opController *schedule.OperatorController, args []string) schedule.Scheduler { filters := []filter.Filter{ filter.StoreStateFilter{ActionScope: randomMergeName, MoveRegion: true}, } @@ -47,6 +53,8 @@ func newRandomMergeScheduler(opController *schedule.OperatorController) schedule return &randomMergeScheduler{ name: randomMergeName, baseScheduler: base, + startKey: []byte(args[0]), + endKey: []byte(args[1]), selector: selector.NewRandomSelector(filters), } } @@ -72,7 +80,7 @@ func (s *randomMergeScheduler) Schedule(cluster schedule.Cluster) []*operator.Op schedulerCounter.WithLabelValues(s.GetName(), "no-source-store").Inc() return nil } - region := cluster.RandLeaderRegion(store.GetID(), core.HealthRegion()) + region := cluster.RandLeaderRegion(store.GetID(), s.startKey, s.endKey, core.HealthRegion()) if region == nil { schedulerCounter.WithLabelValues(s.GetName(), "no-region").Inc() return nil diff --git a/server/schedulers/scatter_range.go b/server/schedulers/scatter_range.go index 559e42a05e28..e8fb23936394 100644 --- a/server/schedulers/scatter_range.go +++ b/server/schedulers/scatter_range.go @@ -59,11 +59,13 @@ func newScatterRangeScheduler(opController *schedule.OperatorController, args [] rangeName: args[2], balanceLeader: newBalanceLeaderScheduler( opController, + []string{"", ""}, WithBalanceLeaderName("scatter-range-leader"), WithBalanceLeaderCounter(scatterRangeLeaderCounter), ), balanceRegion: newBalanceRegionScheduler( opController, + []string{"", ""}, WithBalanceRegionName("scatter-range-region"), WithBalanceRegionCounter(scatterRangeRegionCounter), ), diff --git a/server/schedulers/shuffle_leader.go b/server/schedulers/shuffle_leader.go index 94e05bb6ab5f..4457bd9b330a 100644 --- a/server/schedulers/shuffle_leader.go +++ b/server/schedulers/shuffle_leader.go @@ -23,7 +23,11 @@ import ( func init() { schedule.RegisterScheduler("shuffle-leader", func(opController *schedule.OperatorController, args []string) (schedule.Scheduler, error) { - return newShuffleLeaderScheduler(opController), nil + startKey, endKey, err := getRangeKeys(args) + if err != nil { + return nil, err + } + return newShuffleLeaderScheduler(opController, []string{startKey, endKey}), nil }) } @@ -32,12 +36,14 @@ const shuffleLeaderName = "shuffle-leader-scheduler" type shuffleLeaderScheduler struct { name string *baseScheduler + startKey []byte + endKey []byte selector *selector.RandomSelector } // newShuffleLeaderScheduler creates an admin scheduler that shuffles leaders // between stores. -func newShuffleLeaderScheduler(opController *schedule.OperatorController) schedule.Scheduler { +func newShuffleLeaderScheduler(opController *schedule.OperatorController, args []string) schedule.Scheduler { filters := []filter.Filter{ filter.StoreStateFilter{ActionScope: shuffleLeaderName, TransferLeader: true}, } @@ -45,6 +51,8 @@ func newShuffleLeaderScheduler(opController *schedule.OperatorController) schedu return &shuffleLeaderScheduler{ name: shuffleLeaderName, baseScheduler: base, + startKey: []byte(args[0]), + endKey: []byte(args[1]), selector: selector.NewRandomSelector(filters), } } @@ -72,7 +80,7 @@ func (s *shuffleLeaderScheduler) Schedule(cluster schedule.Cluster) []*operator. schedulerCounter.WithLabelValues(s.GetName(), "no-target-store").Inc() return nil } - region := cluster.RandFollowerRegion(targetStore.GetID(), core.HealthRegion()) + region := cluster.RandFollowerRegion(targetStore.GetID(), s.startKey, s.endKey, core.HealthRegion()) if region == nil { schedulerCounter.WithLabelValues(s.GetName(), "no-follower").Inc() return nil diff --git a/server/schedulers/shuffle_region.go b/server/schedulers/shuffle_region.go index ca57a7edf691..7c4952b7d857 100644 --- a/server/schedulers/shuffle_region.go +++ b/server/schedulers/shuffle_region.go @@ -26,7 +26,11 @@ import ( func init() { schedule.RegisterScheduler("shuffle-region", func(opController *schedule.OperatorController, args []string) (schedule.Scheduler, error) { - return newShuffleRegionScheduler(opController), nil + startKey, endKey, err := getRangeKeys(args) + if err != nil { + return nil, err + } + return newShuffleRegionScheduler(opController, []string{startKey, endKey}), nil }) } @@ -35,12 +39,14 @@ const shuffleRegionName = "shuffle-region-scheduler" type shuffleRegionScheduler struct { name string *baseScheduler + startKey []byte + endKey []byte selector *selector.RandomSelector } // newShuffleRegionScheduler creates an admin scheduler that shuffles regions // between stores. -func newShuffleRegionScheduler(opController *schedule.OperatorController) schedule.Scheduler { +func newShuffleRegionScheduler(opController *schedule.OperatorController, args []string) schedule.Scheduler { filters := []filter.Filter{ filter.StoreStateFilter{ActionScope: shuffleRegionName, MoveRegion: true}, } @@ -48,6 +54,8 @@ func newShuffleRegionScheduler(opController *schedule.OperatorController) schedu return &shuffleRegionScheduler{ name: shuffleRegionName, baseScheduler: base, + startKey: []byte(args[0]), + endKey: []byte(args[1]), selector: selector.NewRandomSelector(filters), } } @@ -98,9 +106,9 @@ func (s *shuffleRegionScheduler) scheduleRemovePeer(cluster schedule.Cluster) (* return nil, nil } - region := cluster.RandFollowerRegion(source.GetID(), core.HealthRegion()) + region := cluster.RandFollowerRegion(source.GetID(), s.startKey, s.endKey, core.HealthRegion()) if region == nil { - region = cluster.RandLeaderRegion(source.GetID(), core.HealthRegion()) + region = cluster.RandLeaderRegion(source.GetID(), s.startKey, s.endKey, core.HealthRegion()) } if region == nil { schedulerCounter.WithLabelValues(s.GetName(), "no-region").Inc() diff --git a/server/schedulers/utils.go b/server/schedulers/utils.go index 0824d197ce50..6bfd11e2feec 100644 --- a/server/schedulers/utils.go +++ b/server/schedulers/utils.go @@ -14,6 +14,7 @@ package schedulers import ( + "net/url" "time" "github.com/montanaflynn/stats" @@ -21,6 +22,7 @@ import ( "github.com/pingcap/pd/server/core" "github.com/pingcap/pd/server/schedule" "github.com/pingcap/pd/server/schedule/operator" + "github.com/pkg/errors" ) const ( @@ -113,3 +115,24 @@ const ( func newTaintCache() *cache.TTLUint64 { return cache.NewIDTTL(taintCacheGCInterval, taintCacheTTL) } + +func getRangeKeys(args []string) (startKey, endKey string, err error) { + if len(args) > 2 { + return "", "", errors.New("the number of arguments should be not more than 2") + } + startKey, endKey = "", "" + if len(args) > 0 { + startKey, err = url.QueryUnescape(args[0]) + if err != nil { + return "", "", err + } + } + + if len(args) > 1 { + endKey, err = url.QueryUnescape(args[1]) + if err != nil { + return "", "", err + } + } + return startKey, endKey, nil +} diff --git a/tools/pd-simulator/simulator/drive.go b/tools/pd-simulator/simulator/drive.go index 0a2fa86acc33..b912e2fa1ad9 100644 --- a/tools/pd-simulator/simulator/drive.go +++ b/tools/pd-simulator/simulator/drive.go @@ -160,7 +160,7 @@ func (d *Driver) TickCount() int64 { // GetBootstrapInfo returns a valid bootstrap store and region. func (d *Driver) GetBootstrapInfo(r *RaftEngine) (*metapb.Store, *metapb.Region, error) { - origin := r.RandRegion() + origin := r.RandRegion([]byte(""), []byte("")) if origin == nil { return nil, nil, errors.New("no region found for bootstrap") } diff --git a/tools/pd-simulator/simulator/raft.go b/tools/pd-simulator/simulator/raft.go index b227b36065c1..c707ff81cb4a 100644 --- a/tools/pd-simulator/simulator/raft.go +++ b/tools/pd-simulator/simulator/raft.go @@ -287,10 +287,10 @@ func (r *RaftEngine) SearchRegion(regionKey []byte) *core.RegionInfo { } // RandRegion gets a region by random -func (r *RaftEngine) RandRegion() *core.RegionInfo { +func (r *RaftEngine) RandRegion(startKey, endKey []byte) *core.RegionInfo { r.RLock() defer r.RUnlock() - return r.regionsInfo.RandRegion() + return r.regionsInfo.RandRegion(startKey, endKey) } func (r *RaftEngine) allocID(storeID uint64) (uint64, error) {