diff --git a/server/api/scheduler.go b/server/api/scheduler.go index 1ab8ef1472a..93b35aca97e 100644 --- a/server/api/scheduler.go +++ b/server/api/scheduler.go @@ -86,6 +86,11 @@ func (h *schedulerHandler) Post(w http.ResponseWriter, r *http.Request) { h.r.JSON(w, http.StatusInternalServerError, err.Error()) return } + case "shuffle-region-scheduler": + if err := h.AddShuffleRegionScheduler(); err != nil { + h.r.JSON(w, http.StatusInternalServerError, err.Error()) + return + } } h.r.JSON(w, http.StatusOK, nil) diff --git a/server/handler.go b/server/handler.go index 2fae8dfda2a..8482ef08ca2 100644 --- a/server/handler.go +++ b/server/handler.go @@ -83,3 +83,8 @@ func (h *Handler) AddEvictLeaderScheduler(storeID uint64) error { func (h *Handler) AddShuffleLeaderScheduler() error { return h.AddScheduler(newShuffleLeaderScheduler(h.opt)) } + +// AddShuffleRegionScheduler adds a shuffle-region-scheduler. +func (h *Handler) AddShuffleRegionScheduler() error { + return h.AddScheduler(newShuffleRegionScheduler(h.opt)) +} diff --git a/server/scheduler.go b/server/scheduler.go index c9af537ff00..72533149754 100644 --- a/server/scheduler.go +++ b/server/scheduler.go @@ -189,6 +189,53 @@ func (s *shuffleLeaderScheduler) Schedule(cluster *clusterInfo) Operator { return newTransferLeader(region, region.GetStorePeer(storeID)) } +type shuffleRegionScheduler struct { + opt *scheduleOption + selector Selector +} + +func newShuffleRegionScheduler(opt *scheduleOption) *shuffleRegionScheduler { + var filters []Filter + filters = append(filters, newStateFilter(opt)) + filters = append(filters, newHealthFilter(opt)) + + return &shuffleRegionScheduler{ + opt: opt, + selector: newRandomSelector(filters), + } +} + +func (s *shuffleRegionScheduler) GetName() string { + return "shuffle-region-scheduler" +} + +func (s *shuffleRegionScheduler) GetResourceKind() ResourceKind { + return regionKind +} + +func (s *shuffleRegionScheduler) GetResourceLimit() uint64 { + return s.opt.GetRegionScheduleLimit() +} + +func (s *shuffleRegionScheduler) Prepare(cluster *clusterInfo) error { return nil } + +func (s *shuffleRegionScheduler) Cleanup(cluster *clusterInfo) {} + +func (s *shuffleRegionScheduler) Schedule(cluster *clusterInfo) Operator { + region, oldPeer := scheduleRemovePeer(cluster, s.selector) + if region == nil { + return nil + } + + excludedFilter := newExcludedFilter(nil, region.GetStoreIds()) + newPeer := scheduleAddPeer(cluster, s.selector, excludedFilter) + if newPeer == nil { + return nil + } + + return newTransferPeer(region, oldPeer, newPeer) +} + func newAddPeer(region *regionInfo, peer *metapb.Peer) Operator { addPeer := newAddPeerOperator(region.GetId(), peer) return newRegionOperator(region, addPeer)