From 4a55a66ce1ab62cbd9c70c01133f170d4064b7ee Mon Sep 17 00:00:00 2001
From: husharp <jinhao.hu@pingcap.com>
Date: Sun, 4 Feb 2024 18:10:11 +0800
Subject: [PATCH] fix return json

Signed-off-by: husharp <jinhao.hu@pingcap.com>
---
 server/api/scheduler.go                       | 33 ++++++++++---------
 server/handler.go                             | 13 ++++----
 .../pd-ctl/tests/scheduler/scheduler_test.go  | 10 ++++--
 3 files changed, 32 insertions(+), 24 deletions(-)

diff --git a/server/api/scheduler.go b/server/api/scheduler.go
index 8ecec7f721c..b1d3e8c07af 100644
--- a/server/api/scheduler.go
+++ b/server/api/scheduler.go
@@ -142,10 +142,23 @@ func (h *schedulerHandler) CreateScheduler(w http.ResponseWriter, r *http.Reques
 			return
 		}
 
-	case schedulers.GrantLeaderName:
-		h.addEvictOrGrant(w, input, schedulers.GrantLeaderName)
-	case schedulers.EvictLeaderName:
-		h.addEvictOrGrant(w, input, schedulers.EvictLeaderName)
+	case schedulers.GrantLeaderName, schedulers.EvictLeaderName:
+		storeID, ok := input["store_id"].(float64)
+		if !ok {
+			h.r.JSON(w, http.StatusBadRequest, "missing store id")
+			return
+		}
+		exist, err := h.AddEvictOrGrant(storeID, name)
+		if err != nil {
+			h.r.JSON(w, http.StatusInternalServerError, err.Error())
+			return
+		}
+		// we should ensure whether it is the first time to create evict-leader-scheduler
+		// or just update the evict-leader.
+		if exist {
+			h.r.JSON(w, http.StatusOK, "The scheduler has been applied to the store.")
+			return
+		}
 	case schedulers.ShuffleLeaderName:
 		if err := h.AddShuffleLeaderScheduler(); err != nil {
 			h.r.JSON(w, http.StatusInternalServerError, err.Error())
@@ -204,18 +217,6 @@ func (h *schedulerHandler) CreateScheduler(w http.ResponseWriter, r *http.Reques
 	h.r.JSON(w, http.StatusOK, "The scheduler is created.")
 }
 
-func (h *schedulerHandler) addEvictOrGrant(w http.ResponseWriter, input map[string]any, name string) {
-	storeID, ok := input["store_id"].(float64)
-	if !ok {
-		h.r.JSON(w, http.StatusBadRequest, "missing store id")
-		return
-	}
-	err := h.AddEvictOrGrant(storeID, name)
-	if err != nil {
-		h.r.JSON(w, http.StatusInternalServerError, err.Error())
-	}
-}
-
 // @Tags     scheduler
 // @Summary  Delete a scheduler.
 // @Param    name  path  string  true  "The name of the scheduler."
diff --git a/server/handler.go b/server/handler.go
index 7c840eabe60..db7fa3f8952 100644
--- a/server/handler.go
+++ b/server/handler.go
@@ -583,10 +583,10 @@ func (h *Handler) redirectSchedulerUpdate(name string, storeID float64) error {
 }
 
 // AddEvictOrGrant add evict leader scheduler or grant leader scheduler.
-func (h *Handler) AddEvictOrGrant(storeID float64, name string) error {
-	if exist, err := h.IsSchedulerExisted(name); !exist {
+func (h *Handler) AddEvictOrGrant(storeID float64, name string) (exist bool, err error) {
+	if exist, err = h.IsSchedulerExisted(name); !exist {
 		if err != nil && !errors.ErrorEqual(err, errs.ErrSchedulerNotFound.FastGenByArgs()) {
-			return err
+			return exist, err
 		}
 		switch name {
 		case schedulers.EvictLeaderName:
@@ -595,13 +595,14 @@ func (h *Handler) AddEvictOrGrant(storeID float64, name string) error {
 			err = h.AddGrantLeaderScheduler(uint64(storeID))
 		}
 		if err != nil {
-			return err
+			return exist, err
 		}
 	} else {
 		if err := h.redirectSchedulerUpdate(name, storeID); err != nil {
-			return err
+			return exist, err
 		}
 		log.Info("update scheduler", zap.String("scheduler-name", name), zap.Uint64("store-id", uint64(storeID)))
+		return exist, nil
 	}
-	return nil
+	return exist, nil
 }
diff --git a/tools/pd-ctl/tests/scheduler/scheduler_test.go b/tools/pd-ctl/tests/scheduler/scheduler_test.go
index 19467db2282..dcfa23f997f 100644
--- a/tools/pd-ctl/tests/scheduler/scheduler_test.go
+++ b/tools/pd-ctl/tests/scheduler/scheduler_test.go
@@ -94,7 +94,10 @@ func (suite *schedulerTestSuite) TearDownTest() {
 }
 
 func (suite *schedulerTestSuite) TestScheduler() {
-	suite.env.RunTestInTwoModes(suite.checkScheduler)
+	// use a new environment to avoid affecting other tests
+	env := pdTests.NewSchedulingTestEnvironment(suite.T())
+	env.RunTestInTwoModes(suite.checkScheduler)
+	env.Cleanup()
 }
 
 func (suite *schedulerTestSuite) checkScheduler(cluster *pdTests.TestCluster) {
@@ -633,7 +636,10 @@ func (suite *schedulerTestSuite) checkScheduler(cluster *pdTests.TestCluster) {
 }
 
 func (suite *schedulerTestSuite) TestSchedulerDiagnostic() {
-	suite.env.RunTestInTwoModes(suite.checkSchedulerDiagnostic)
+	// use a new environment to avoid affecting other tests
+	env := pdTests.NewSchedulingTestEnvironment(suite.T())
+	env.RunTestInTwoModes(suite.checkSchedulerDiagnostic)
+	env.Cleanup()
 }
 
 func (suite *schedulerTestSuite) checkSchedulerDiagnostic(cluster *pdTests.TestCluster) {