schedule: fix panic when setting store-limit-mode to auto (#2544) (#2547

) Signed-off-by: Ryan Leung <rleungx@gmail.com>
tikv · Jun 17, 2020 · 330768a · 330768a
1 parent 85a49d5
commit 330768a
Show file tree

Hide file tree

Showing 3 changed files with 22 additions and 1 deletion.
diff --git a/pkg/mock/mockoption/mockoption.go b/pkg/mock/mockoption/mockoption.go
@@ -43,6 +43,7 @@ const (
  defaultLeaderSchedulePolicy = "count"
  defaultEnablePlacementRules = false
  defaultKeyType = "table"
+ defaultStoreLimitMode = "manual"
 )
 
 // ScheduleOptions is a mock of ScheduleOptions
@@ -71,6 +72,7 @@ type ScheduleOptions struct {
  TolerantSizeRatio float64
  LowSpaceRatio float64
  HighSpaceRatio float64
+ StoreLimitMode string
  EnableRemoveDownReplica bool
  EnableReplaceOfflineReplica bool
  EnableMakeUpReplica bool
@@ -110,6 +112,7 @@ func NewScheduleOptions() *ScheduleOptions {
  mso.TolerantSizeRatio = defaultTolerantSizeRatio
  mso.LowSpaceRatio = defaultLowSpaceRatio
  mso.HighSpaceRatio = defaultHighSpaceRatio
+ mso.StoreLimitMode = defaultStoreLimitMode
  mso.EnableRemoveDownReplica = true
  mso.EnableReplaceOfflineReplica = true
  mso.EnableMakeUpReplica = true
@@ -279,3 +282,8 @@ func (mso *ScheduleOptions) GetLeaderSchedulePolicy() core.SchedulePolicy {
 func (mso *ScheduleOptions) GetKeyType() core.KeyType {
  return core.StringToKeyType(mso.KeyType)
 }
+
+// GetStoreLimitMode returns the limit mode of store.
+func (mso *ScheduleOptions) GetStoreLimitMode() string {
+ return mso.StoreLimitMode
+}
diff --git a/server/schedule/operator_controller.go b/server/schedule/operator_controller.go
@@ -897,7 +897,8 @@ func (oc *OperatorController) SetAllStoresLimitAuto(rate float64, limitType stor
  for _, s := range stores {
  sid := s.GetID()
  if old, ok := oc.storesLimit[sid]; ok {
- if old[limitType].Mode() == storelimit.Manual {
+ limit, ok1 := old[limitType]
+ if ok1 && limit.Mode() == storelimit.Manual {
  continue
  }
  }

diff --git a/server/schedule/operator_controller_test.go b/server/schedule/operator_controller_test.go
@@ -704,3 +704,15 @@ func (t *testOperatorControllerSuite) TestAddWaitingOperator(c *C) {
  // no space left, new operator can not be added.
  c.Assert(controller.AddWaitingOperator(addPeerOp(0)), Equals, 0)
 }
+
+func (t *testOperatorControllerSuite) TestAutoStoreLimitMode(c *C) {
+ opt := mockoption.NewScheduleOptions()
+ opt.StoreLimitMode = "auto"
+ tc := mockcluster.NewCluster(opt)
+ stream := mockhbstream.NewHeartbeatStreams(tc.ID, true /* no need to run */)
+ oc := NewOperatorController(t.ctx, tc, stream)
+
+ tc.AddLeaderStore(1, 10)
+ oc.SetStoreLimit(1, 10, storelimit.Auto, storelimit.RegionAdd)
+ oc.SetAllStoresLimitAuto(10, storelimit.RegionRemove)
+}