api: change store limit by label

server/api/store.go

@@ -457,11 +457,33 @@ func (h *storesHandler) SetAllLimit(w http.ResponseWriter, r *http.Request) {
 		return
 	}
 
-	for _, typ := range typeValues {
-		if err := h.SetAllStoresLimit(ratePerMin, typ); err != nil {
-			h.rd.JSON(w, http.StatusInternalServerError, err.Error())
+	if _, ok := input["labels"]; !ok {
+		for _, typ := range typeValues {
+			if err := h.SetAllStoresLimit(ratePerMin, typ); err != nil {
+				h.rd.JSON(w, http.StatusInternalServerError, err.Error())
+				return
+			}
+		}
+	} else {
+		labelMap := input["labels"].(map[string]interface{})
+		labels := make([]*metapb.StoreLabel, 0, len(input))
+		for k, v := range labelMap {
+			labels = append(labels, &metapb.StoreLabel{
+				Key:   k,
+				Value: v.(string),
+			})
+		}
+
+		if err := config.ValidateLabels(labels); err != nil {
+			apiutil.ErrorResp(h.rd, w, errcode.NewInvalidInputErr(err))
 			return
 		}
+		for _, typ := range typeValues {
+			if err := h.SetLabelStoresLimit(ratePerMin, typ, labels); err != nil {
+				h.rd.JSON(w, http.StatusInternalServerError, err.Error())
+				return
+			}
+		}
 	}
 
 	h.rd.JSON(w, http.StatusOK, "Set store limit successfully.")

server/handler.go

@@ -420,6 +420,24 @@ func (h *Handler) SetAllStoresLimit(ratePerMin float64, limitType storelimit.Typ
 	return nil
 }
 
+// SetLabelStoresLimit is used to set limit of label stores.
+func (h *Handler) SetLabelStoresLimit(ratePerMin float64, limitType storelimit.Type, labels []*metapb.StoreLabel) error {
+	c, err := h.GetRaftCluster()
+	if err != nil {
+		return err
+	}
+	for _, store := range c.GetStores() {
+		for _, label := range labels {
+			for _, sl := range store.GetLabels() {
+				if label.Key == sl.Key && label.Value == sl.Value {
+					c.SetStoreLimit(store.GetID(), limitType, ratePerMin)
+				}
+			}
+		}
+	}
+	return nil
+}
+
 // GetAllStoresLimit is used to get limit of all stores.
 func (h *Handler) GetAllStoresLimit(limitType storelimit.Type) (map[uint64]config.StoreLimitConfig, error) {
 	c, err := h.GetRaftCluster()

tests/pdctl/store/store_test.go

@@ -196,6 +196,14 @@ func (s *storeTestSuite) TestStore(c *C) {
 	limit2 = leaderServer.GetRaftCluster().GetStoreLimitByType(2, storelimit.RemovePeer)
 	c.Assert(limit2, Equals, float64(25))
 
+	// store limit all <key> <value> <rate> <type>
+	args = []string{"-u", pdAddr, "store", "limit", "all", "zone", "uk", "20", "remove-peer"}
+	_, output, err = pdctl.ExecuteCommandC(cmd, args...)
+	c.Log(string(output))
+	c.Assert(err, IsNil)
+	limit1 = leaderServer.GetRaftCluster().GetStoreLimitByType(1, storelimit.RemovePeer)
+	c.Assert(limit1, Equals, float64(20))
+
 	// store limit all 0 is invalid
 	args = []string{"-u", pdAddr, "store", "limit", "all", "0"}
 	_, output, err = pdctl.ExecuteCommandC(cmd, args...)
@@ -214,7 +222,7 @@ func (s *storeTestSuite) TestStore(c *C) {
 	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["1"]["remove-peer"].(float64), Equals, float64(20))
 	c.Assert(allRemovePeerLimit["3"]["remove-peer"].(float64), Equals, float64(25))
 	_, ok = allRemovePeerLimit["2"]["add-peer"]
 	c.Assert(ok, Equals, false)

tools/pd-ctl/pdctl/command/store_command.go

@@ -89,7 +89,7 @@ func NewSetStoreWeightCommand() *cobra.Command {
 // NewStoreLimitCommand returns a limit subcommand of storeCmd.
 func NewStoreLimitCommand() *cobra.Command {
 	c := &cobra.Command{
-		Use:   "limit [<type>]|[<store_id>|<all> <limit> <type>]",
+		Use:   "limit [<type>]|[<store_id>|<all> [<key> <value>]... <limit> <type>]",
 		Short: "show or set a store's rate limit",
 		Long:  "show or set a store's rate limit, <type> can be 'add-peer'(default) or 'remove-peer'",
 		Run:   storeLimitCommandFunc,
@@ -362,8 +362,7 @@ func setStoreWeightCommandFunc(cmd *cobra.Command, args []string) {
 
 func storeLimitCommandFunc(cmd *cobra.Command, args []string) {
 	argsCount := len(args)
-	switch argsCount {
-	case 0, 1:
+	if argsCount <= 1 {
 		prefix := path.Join(storesPrefix, "limit")
 		if argsCount == 1 {
 			prefix += fmt.Sprintf("?type=%s", args[0])
@@ -374,7 +373,7 @@ func storeLimitCommandFunc(cmd *cobra.Command, args []string) {
 			return
 		}
 		cmd.Println(r)
-	case 2, 3:
+	} else if argsCount <= 3 {
 		rate, err := strconv.ParseFloat(args[1], 64)
 		if err != nil || rate <= 0 {
 			cmd.Println("rate should be a number that > 0.")
@@ -394,9 +393,30 @@ func storeLimitCommandFunc(cmd *cobra.Command, args []string) {
 			postInput["type"] = args[2]
 		}
 		postJSON(cmd, prefix, postInput)
-	default:
-		cmd.Usage()
-		return
+	} else {
+		if args[0] != "all" {
+			cmd.Println("Labels are an option of set all stores limit.")
+		} else {
+			postInput := map[string]interface{}{}
+			prefix := path.Join(storesPrefix, "limit")
+			ratePos := argsCount - 1
+			if argsCount%2 == 1 {
+				postInput["type"] = args[argsCount-1]
+				ratePos = argsCount - 2
+			}
+			rate, err := strconv.ParseFloat(args[ratePos], 64)
+			if err != nil || rate <= 0 {
+				cmd.Println("rate should be a number that > 0.")
+				return
+			}
+			postInput["rate"] = rate
+			labels := make(map[string]interface{})
+			for i := 1; i < ratePos; i += 2 {
+				labels[args[i]] = args[i+1]
+			}
+			postInput["labels"] = labels
+			postJSON(cmd, prefix, postInput)
+		}
 	}
 }