*: reset config if the input is invalid (tikv#8632) (tikv#8660)

close tikv#8619 Signed-off-by: ti-chi-bot <ti-community-prow-bot@tidb.io> Signed-off-by: Ryan Leung <rleungx@gmail.com> Co-authored-by: Ryan Leung <rleungx@gmail.com>
ti-chi-bot · Nov 7, 2024 · 57d5bdc · 57d5bdc
1 parent f235847
commit 57d5bdc
Showing 5 changed files with 1,162 additions and 7 deletions.
diff --git a/junitfile b/junitfile
diff --git a/pkg/schedule/schedulers/evict_leader.go b/pkg/schedule/schedulers/evict_leader.go
@@ -75,6 +75,9 @@ func (conf *evictLeaderSchedulerConfig) getStores() []uint64 {
 }
 
 func (conf *evictLeaderSchedulerConfig) BuildWithArgs(args []string) error {
+	failpoint.Inject("buildWithArgsErr", func() {
+		failpoint.Return(errors.New("fail to build with args"))
+	})
 	if len(args) < 1 {
 		return errs.ErrSchedulerConfig.FastGenByArgs("id")
 	}
@@ -386,8 +389,15 @@ func (handler *evictLeaderHandler) UpdateConfig(w http.ResponseWriter, r *http.R
 		args = append(args, handler.config.getRanges(id)...)
 	}
 
-	handler.config.BuildWithArgs(args)
-	err := handler.config.Persist()
+	err := handler.config.BuildWithArgs(args)
+	if err != nil {
+		handler.config.Lock()
+		handler.config.cluster.ResumeLeaderTransfer(id)
+		handler.config.Unlock()
+		handler.rd.JSON(w, http.StatusBadRequest, err.Error())
+		return
+	}
+	err = handler.config.Persist()
 	if err != nil {
 		handler.config.removeStore(id)
 		handler.rd.JSON(w, http.StatusInternalServerError, err.Error())

diff --git a/pkg/schedule/schedulers/grant_leader.go b/pkg/schedule/schedulers/grant_leader.go
@@ -296,10 +296,17 @@ func (handler *grantLeaderHandler) UpdateConfig(w http.ResponseWriter, r *http.R
 		args = append(args, handler.config.getRanges(id)...)
 	}
 
-	handler.config.BuildWithArgs(args)
-	err := handler.config.Persist()
+	err := handler.config.BuildWithArgs(args)
 	if err != nil {
-		handler.config.removeStore(id)
+		handler.config.Lock()
+		handler.config.cluster.ResumeLeaderTransfer(id)
+		handler.config.Unlock()
+		handler.rd.JSON(w, http.StatusBadRequest, err.Error())
+		return
+	}
+	err = handler.config.Persist()
+	if err != nil {
+		_, _ = handler.config.removeStore(id)
 		handler.rd.JSON(w, http.StatusInternalServerError, err.Error())
 		return
 	}

diff --git a/plugin/scheduler_example/evict_leader.go b/plugin/scheduler_example/evict_leader.go
@@ -275,9 +275,21 @@ func (handler *evictLeaderHandler) UpdateConfig(w http.ResponseWriter, r *http.R
 		args = append(args, handler.config.getRanges(id)...)
 	}
 
-	handler.config.BuildWithArgs(args)
-	err := handler.config.Persist()
+	err := handler.config.BuildWithArgs(args)
 	if err != nil {
+		handler.config.mu.Lock()
+		handler.config.cluster.ResumeLeaderTransfer(id)
+		handler.config.mu.Unlock()
+		handler.rd.JSON(w, http.StatusBadRequest, err.Error())
+		return
+	}
+
+	err = handler.config.Persist()
+	if err != nil {
+		handler.config.mu.Lock()
+		delete(handler.config.StoreIDWitRanges, id)
+		handler.config.cluster.ResumeLeaderTransfer(id)
+		handler.config.mu.Unlock()
 		handler.rd.JSON(w, http.StatusInternalServerError, err.Error())
 	}
 	handler.rd.JSON(w, http.StatusOK, nil)

diff --git a/tools/pd-ctl/tests/scheduler/scheduler_test.go b/tools/pd-ctl/tests/scheduler/scheduler_test.go
@@ -722,6 +722,64 @@ func (suite *schedulerTestSuite) checkSchedulerDiagnostic(cluster *pdTests.TestC
 	checkSchedulerDescribeCommand("balance-leader-scheduler", "normal", "")
 }
 
+func (suite *schedulerTestSuite) TestEvictLeaderScheduler() {
+	suite.env.RunTestInTwoModes(suite.checkEvictLeaderScheduler)
+}
+
+func (suite *schedulerTestSuite) checkEvictLeaderScheduler(cluster *pdTests.TestCluster) {
+	re := suite.Require()
+	pdAddr := cluster.GetConfig().GetClientURL()
+	cmd := ctl.GetRootCmd()
+
+	stores := []*metapb.Store{
+		{
+			Id:            1,
+			State:         metapb.StoreState_Up,
+			LastHeartbeat: time.Now().UnixNano(),
+		},
+		{
+			Id:            2,
+			State:         metapb.StoreState_Up,
+			LastHeartbeat: time.Now().UnixNano(),
+		},
+		{
+			Id:            3,
+			State:         metapb.StoreState_Up,
+			LastHeartbeat: time.Now().UnixNano(),
+		},
+		{
+			Id:            4,
+			State:         metapb.StoreState_Up,
+			LastHeartbeat: time.Now().UnixNano(),
+		},
+	}
+	for _, store := range stores {
+		pdTests.MustPutStore(re, cluster, store)
+	}
+
+	pdTests.MustPutRegion(re, cluster, 1, 1, []byte("a"), []byte("b"))
+	suite.checkDefaultSchedulers(re, cmd, pdAddr)
+
+	output, err := tests.ExecuteCommand(cmd, []string{"-u", pdAddr, "scheduler", "add", "evict-leader-scheduler", "2"}...)
+	re.NoError(err)
+	re.Contains(string(output), "Success!")
+	failpoint.Enable("github.com/tikv/pd/pkg/schedule/schedulers/buildWithArgsErr", "return(true)")
+	output, err = tests.ExecuteCommand(cmd, []string{"-u", pdAddr, "scheduler", "add", "evict-leader-scheduler", "1"}...)
+	re.NoError(err)
+	re.Contains(string(output), "fail to build with args")
+	failpoint.Disable("github.com/tikv/pd/pkg/schedule/schedulers/buildWithArgsErr")
+	output, err = tests.ExecuteCommand(cmd, []string{"-u", pdAddr, "scheduler", "remove", "evict-leader-scheduler"}...)
+	re.NoError(err)
+	re.Contains(string(output), "Success!")
+	testutil.Eventually(re, func() bool {
+		output, err = tests.ExecuteCommand(cmd, []string{"-u", pdAddr, "scheduler", "add", "evict-leader-scheduler", "1"}...)
+		return err == nil && strings.Contains(string(output), "Success!")
+	})
+	output, err = tests.ExecuteCommand(cmd, []string{"-u", pdAddr, "scheduler", "remove", "evict-leader-scheduler-1"}...)
+	re.NoError(err)
+	re.Contains(string(output), "Success!")
+}
+
 func mustExec(re *require.Assertions, cmd *cobra.Command, args []string, v any) string {
 	output, err := tests.ExecuteCommand(cmd, args...)
 	re.NoError(err)
@@ -740,3 +798,32 @@ func mightExec(re *require.Assertions, cmd *cobra.Command, args []string, v any)
 	}
 	json.Unmarshal(output, v)
 }
+
+func (suite *schedulerTestSuite) checkDefaultSchedulers(re *require.Assertions, cmd *cobra.Command, pdAddr string) {
+	// scheduler show command
+	expected := make(map[string]bool)
+	for _, scheduler := range suite.defaultSchedulers {
+		expected[scheduler] = true
+	}
+	checkSchedulerCommand(re, cmd, pdAddr, nil, expected)
+}
+
+func checkSchedulerCommand(re *require.Assertions, cmd *cobra.Command, pdAddr string, args []string, expected map[string]bool) {
+	if args != nil {
+		echo := mustExec(re, cmd, args, nil)
+		re.Contains(echo, "Success!")
+	}
+	testutil.Eventually(re, func() bool {
+		var schedulers []string
+		mustExec(re, cmd, []string{"-u", pdAddr, "scheduler", "show"}, &schedulers)
+		if len(schedulers) != len(expected) {
+			return false
+		}
+		for _, scheduler := range schedulers {
+			if _, ok := expected[scheduler]; !ok {
+				return false
+			}
+		}
+		return true
+	})
+}