From 330768a4dc6e18d892d1ff1490065966a8e995e7 Mon Sep 17 00:00:00 2001
From: ti-srebot <66930949+ti-srebot@users.noreply.github.com>
Date: Wed, 17 Jun 2020 22:09:18 +0800
Subject: [PATCH] schedule: fix panic when setting store-limit-mode to auto
 (#2544) (#2547)

Signed-off-by: Ryan Leung <rleungx@gmail.com>
---
 pkg/mock/mockoption/mockoption.go           |  8 ++++++++
 server/schedule/operator_controller.go      |  3 ++-
 server/schedule/operator_controller_test.go | 12 ++++++++++++
 3 files changed, 22 insertions(+), 1 deletion(-)

diff --git a/pkg/mock/mockoption/mockoption.go b/pkg/mock/mockoption/mockoption.go
index cbeb7a5323e..feaff33215b 100644
--- 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
index 112dce1a5f4..cf4d581fb67 100755
--- 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
index ac2dfd5ddf1..3e2d3e633c3 100644
--- 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)
+}