From 495deb85bb5a8a8690604f6be1e39b1b2179d1e2 Mon Sep 17 00:00:00 2001
From: ti-srebot <66930949+ti-srebot@users.noreply.github.com>
Date: Tue, 11 Aug 2020 15:23:46 +0800
Subject: [PATCH] cherry pick #2739 to release-4.0 (#2743)

Signed-off-by: ti-srebot <ti-srebot@pingcap.com>

Co-authored-by: Song Gao <disxiaofei@163.com>
---
 server/api/store.go             | 23 ++++++++++++++++
 server/api/store_test.go        | 49 +++++++++++++++++++++++++++++++++
 tests/pdctl/store/store_test.go |  6 ++--
 3 files changed, 76 insertions(+), 2 deletions(-)

diff --git a/server/api/store.go b/server/api/store.go
index dec122e491a..cd5374561ec 100644
--- a/server/api/store.go
+++ b/server/api/store.go
@@ -464,12 +464,35 @@ func (h *storesHandler) SetAllLimit(w http.ResponseWriter, r *http.Request) {
 // FIXME: details of output json body
 // @Tags store
 // @Summary Get limit of all stores in the cluster.
+// @Param include_tombstone query bool false "include Tombstone" default(false)
 // @Produce json
 // @Success 200 {object} string
 // @Failure 500 {string} string "PD server failed to proceed the request."
 // @Router /stores/limit [get]
 func (h *storesHandler) GetAllLimit(w http.ResponseWriter, r *http.Request) {
 	limits := h.GetScheduleConfig().StoreLimit
+	includeTombstone := false
+	var err error
+	if includeStr := r.URL.Query().Get("include_tombstone"); includeStr != "" {
+		includeTombstone, err = strconv.ParseBool(includeStr)
+		if err != nil {
+			h.rd.JSON(w, http.StatusBadRequest, err.Error())
+			return
+		}
+	}
+	if !includeTombstone {
+		returned := make(map[uint64]config.StoreLimitConfig, len(limits))
+		rc := getCluster(r.Context())
+		for storeID, v := range limits {
+			store := rc.GetStore(storeID)
+			if store == nil || store.IsTombstone() {
+				continue
+			}
+			returned[storeID] = v
+		}
+		h.rd.JSON(w, http.StatusOK, returned)
+		return
+	}
 	h.rd.JSON(w, http.StatusOK, limits)
 }
 
diff --git a/server/api/store_test.go b/server/api/store_test.go
index bf0809c7fde..3468119df2b 100644
--- a/server/api/store_test.go
+++ b/server/api/store_test.go
@@ -346,3 +346,52 @@ func (s *testStoreSuite) TestDownState(c *C) {
 	storeInfo = newStoreInfo(s.svr.GetScheduleConfig(), newStore)
 	c.Assert(storeInfo.Store.StateName, Equals, downStateName)
 }
+
+func (s *testStoreSuite) TestGetAllLimit(c *C) {
+	testcases := []struct {
+		name           string
+		url            string
+		expectedStores map[uint64]struct{}
+	}{
+		{
+			name: "includeTombstone",
+			url:  fmt.Sprintf("%s/stores/limit?include_tombstone=true", s.urlPrefix),
+			expectedStores: map[uint64]struct{}{
+				1: {},
+				4: {},
+				6: {},
+				7: {},
+			},
+		},
+		{
+			name: "excludeTombStone",
+			url:  fmt.Sprintf("%s/stores/limit?include_tombstone=false", s.urlPrefix),
+			expectedStores: map[uint64]struct{}{
+				1: {},
+				4: {},
+				6: {},
+			},
+		},
+		{
+			name: "default",
+			url:  fmt.Sprintf("%s/stores/limit", s.urlPrefix),
+			expectedStores: map[uint64]struct{}{
+				1: {},
+				4: {},
+				6: {},
+			},
+		},
+	}
+
+	for _, testcase := range testcases {
+		c.Logf(testcase.name)
+		info := make(map[uint64]interface{}, 4)
+		err := readJSON(testDialClient, testcase.url, &info)
+		c.Assert(err, IsNil)
+		c.Assert(len(info), Equals, len(testcase.expectedStores))
+		for id := range testcase.expectedStores {
+			_, ok := info[id]
+			c.Assert(ok, Equals, true)
+		}
+	}
+}
diff --git a/tests/pdctl/store/store_test.go b/tests/pdctl/store/store_test.go
index 67aeddd370c..5551873975b 100644
--- a/tests/pdctl/store/store_test.go
+++ b/tests/pdctl/store/store_test.go
@@ -208,14 +208,16 @@ func (s *storeTestSuite) TestStore(c *C) {
 	json.Unmarshal([]byte(echo), &allAddPeerLimit)
 	c.Assert(allAddPeerLimit["1"]["add-peer"].(float64), Equals, float64(20))
 	c.Assert(allAddPeerLimit["3"]["add-peer"].(float64), Equals, float64(20))
-	c.Assert(allAddPeerLimit["2"]["add-peer"].(float64), Equals, float64(20))
+	_, ok := allAddPeerLimit["2"]["add-peer"]
+	c.Assert(ok, Equals, false)
 
 	echo = pdctl.GetEcho([]string{"-u", pdAddr, "store", "limit", "remove-peer"})
 	allRemovePeerLimit := make(map[string]map[string]interface{})
 	json.Unmarshal([]byte(echo), &allRemovePeerLimit)
 	c.Assert(allRemovePeerLimit["1"]["remove-peer"].(float64), Equals, float64(25))
 	c.Assert(allRemovePeerLimit["3"]["remove-peer"].(float64), Equals, float64(25))
-	c.Assert(allRemovePeerLimit["2"]["remove-peer"].(float64), Equals, float64(25))
+	_, ok = allRemovePeerLimit["2"]["add-peer"]
+	c.Assert(ok, Equals, false)
 
 	// store delete <store_id> command
 	c.Assert(storeInfo.Store.State, Equals, metapb.StoreState_Up)