diff --git a/server/schedulers/evict_leader.go b/server/schedulers/evict_leader.go index 1862c337de8..d1f6d392b49 100644 --- a/server/schedulers/evict_leader.go +++ b/server/schedulers/evict_leader.go @@ -85,6 +85,16 @@ type evictLeaderSchedulerConfig struct { cluster opt.Cluster } +func (conf *evictLeaderSchedulerConfig) getStores() []uint64 { + conf.mu.RLock() + defer conf.mu.RUnlock() + stores := make([]uint64, 0, len(conf.StoreIDWithRanges)) + for storeID := range conf.StoreIDWithRanges { + stores = append(stores, storeID) + } + return stores +} + func (conf *evictLeaderSchedulerConfig) BuildWithArgs(args []string) error { if len(args) != 1 { return errs.ErrSchedulerConfig.FastGenByArgs("id") @@ -189,6 +199,11 @@ func newEvictLeaderScheduler(opController *schedule.OperatorController, conf *ev } } +// EvictStoreIDs returns the IDs of the evict-stores. +func (s *evictLeaderScheduler) EvictStoreIDs() []uint64 { + return s.conf.getStores() +} + func (s *evictLeaderScheduler) ServeHTTP(w http.ResponseWriter, r *http.Request) { s.handler.ServeHTTP(w, r) } diff --git a/tests/server/cluster/cluster_test.go b/tests/server/cluster/cluster_test.go index 1b623331670..2ad1afc4ec0 100644 --- a/tests/server/cluster/cluster_test.go +++ b/tests/server/cluster/cluster_test.go @@ -1229,7 +1229,7 @@ func (s *clusterTestSuite) TestTransferLeaderForScheduler(c *C) { "store_id": 2, }) // Check scheduler updated. - c.Assert(len(rc.GetSchedulers()), Equals, 5) + c.Assert(len(rc.GetSchedulers()), Equals, 4) checkEvictLeaderSchedulerExist(c, rc, true) checkEvictLeaderStoreIDs(c, rc, []uint64{1, 2}) @@ -1248,7 +1248,7 @@ func (s *clusterTestSuite) TestTransferLeaderForScheduler(c *C) { time.Sleep(time.Second) c.Assert(leaderServer.GetRaftCluster().IsPrepared(), IsTrue) // Check scheduler updated. - c.Assert(len(rc.GetSchedulers()), Equals, 5) + c.Assert(len(rc.GetSchedulers()), Equals, 4) checkEvictLeaderSchedulerExist(c, rc, true) checkEvictLeaderStoreIDs(c, rc, []uint64{1, 2}) @@ -1266,7 +1266,7 @@ func (s *clusterTestSuite) TestTransferLeaderForScheduler(c *C) { time.Sleep(time.Second) c.Assert(leaderServer.GetRaftCluster().IsPrepared(), IsTrue) // Check scheduler updated - c.Assert(len(rc.GetSchedulers()), Equals, 5) + c.Assert(len(rc.GetSchedulers()), Equals, 4) checkEvictLeaderSchedulerExist(c, rc, true) checkEvictLeaderStoreIDs(c, rc, []uint64{1, 2})