From 4072fdc87b75ddee468ff1bc1821f3618a7441fc Mon Sep 17 00:00:00 2001 From: ZenoTan Date: Mon, 7 Sep 2020 18:45:50 +0800 Subject: [PATCH 1/8] Add basic change Signed-off-by: ZenoTan --- server/cluster/cluster.go | 20 ++++++++++++++++++-- 1 file changed, 18 insertions(+), 2 deletions(-) diff --git a/server/cluster/cluster.go b/server/cluster/cluster.go index 09260ffd627..7936fe2706c 100644 --- a/server/cluster/cluster.go +++ b/server/cluster/cluster.go @@ -1779,8 +1779,24 @@ func (c *RaftCluster) SetStoreLimit(storeID uint64, typ storelimit.Type, ratePer } // SetAllStoresLimit sets all store limit for a given type and rate. -func (c *RaftCluster) SetAllStoresLimit(typ storelimit.Type, ratePerMin float64) { - c.opt.SetAllStoresLimit(typ, ratePerMin) +func (c *RaftCluster) SetAllStoresLimit(typ storelimit.Type, ratePerMin float64, labels... config.StoreLabel) { + if len(labels) == 0 { + c.opt.SetAllStoresLimit(typ, ratePerMin) + } else { + c.setLabelStoresLimit(typ, ratePerMin, labels...) + } +} + +func (c *RaftCluster) setLabelStoresLimit(typ storelimit.Type, ratePerMin float64, labels... config.StoreLabel) { + for _, store := range c.GetStores() { + for _, sl := range store.GetLabels() { + for _, l := range labels { + if sl.Key == l.Key && sl.Value == l.Value { + c.SetStoreLimit(store.GetID(), typ, ratePerMin) + } + } + } + } } // GetClusterVersion returns the current cluster version. From a473ec3bf6ca3add281d14b2d6ade4be2e85dbe4 Mon Sep 17 00:00:00 2001 From: ZenoTan Date: Tue, 8 Sep 2020 11:55:55 +0800 Subject: [PATCH 2/8] Change api Signed-off-by: ZenoTan --- server/api/store.go | 21 +++++++++++++++++---- server/cluster/cluster.go | 20 ++------------------ server/handler.go | 18 ++++++++++++++++++ 3 files changed, 37 insertions(+), 22 deletions(-) diff --git a/server/api/store.go b/server/api/store.go index 330be679eef..c454544f3dd 100644 --- a/server/api/store.go +++ b/server/api/store.go @@ -457,10 +457,23 @@ 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()) - return + if labels, 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 { + labelValues, ok := labels.([]config.StoreLabel) + if !ok || labelValues == nil { + h.rd.JSON(w, http.StatusBadRequest, "invalid store labels") + } + for _, typ := range typeValues { + if err := h.SetLabelStoresLimit(ratePerMin, typ, labelValues); err != nil { + h.rd.JSON(w, http.StatusInternalServerError, err.Error()) + return + } } } diff --git a/server/cluster/cluster.go b/server/cluster/cluster.go index 7936fe2706c..09260ffd627 100644 --- a/server/cluster/cluster.go +++ b/server/cluster/cluster.go @@ -1779,24 +1779,8 @@ func (c *RaftCluster) SetStoreLimit(storeID uint64, typ storelimit.Type, ratePer } // SetAllStoresLimit sets all store limit for a given type and rate. -func (c *RaftCluster) SetAllStoresLimit(typ storelimit.Type, ratePerMin float64, labels... config.StoreLabel) { - if len(labels) == 0 { - c.opt.SetAllStoresLimit(typ, ratePerMin) - } else { - c.setLabelStoresLimit(typ, ratePerMin, labels...) - } -} - -func (c *RaftCluster) setLabelStoresLimit(typ storelimit.Type, ratePerMin float64, labels... config.StoreLabel) { - for _, store := range c.GetStores() { - for _, sl := range store.GetLabels() { - for _, l := range labels { - if sl.Key == l.Key && sl.Value == l.Value { - c.SetStoreLimit(store.GetID(), typ, ratePerMin) - } - } - } - } +func (c *RaftCluster) SetAllStoresLimit(typ storelimit.Type, ratePerMin float64) { + c.opt.SetAllStoresLimit(typ, ratePerMin) } // GetClusterVersion returns the current cluster version. diff --git a/server/handler.go b/server/handler.go index 74528acae43..f291f5bf4ae 100644 --- a/server/handler.go +++ b/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 []config.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() From a46f6b9d9ce7ae1940f5913bc492c94b610a86aa Mon Sep 17 00:00:00 2001 From: ZenoTan Date: Tue, 8 Sep 2020 14:11:46 +0800 Subject: [PATCH 3/8] pd ctl support Signed-off-by: ZenoTan --- server/api/store.go | 19 ++++++++--- server/handler.go | 2 +- tests/pdctl/store/store_test.go | 10 +++++- tools/pd-ctl/pdctl/command/store_command.go | 35 ++++++++++++++++----- 4 files changed, 52 insertions(+), 14 deletions(-) diff --git a/server/api/store.go b/server/api/store.go index c454544f3dd..af13c5a1766 100644 --- a/server/api/store.go +++ b/server/api/store.go @@ -457,7 +457,7 @@ func (h *storesHandler) SetAllLimit(w http.ResponseWriter, r *http.Request) { return } - if labels, ok := input["labels"]; !ok { + 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()) @@ -465,12 +465,21 @@ func (h *storesHandler) SetAllLimit(w http.ResponseWriter, r *http.Request) { } } } else { - labelValues, ok := labels.([]config.StoreLabel) - if !ok || labelValues == nil { - h.rd.JSON(w, http.StatusBadRequest, "invalid store labels") + 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, labelValues); err != nil { + if err := h.SetLabelStoresLimit(ratePerMin, typ, labels); err != nil { h.rd.JSON(w, http.StatusInternalServerError, err.Error()) return } diff --git a/server/handler.go b/server/handler.go index f291f5bf4ae..87fe3d443ab 100644 --- a/server/handler.go +++ b/server/handler.go @@ -421,7 +421,7 @@ func (h *Handler) SetAllStoresLimit(ratePerMin float64, limitType storelimit.Typ } // SetLabelStoresLimit is used to set limit of label stores. -func (h *Handler) SetLabelStoresLimit(ratePerMin float64, limitType storelimit.Type, labels []config.StoreLabel) error { +func (h *Handler) SetLabelStoresLimit(ratePerMin float64, limitType storelimit.Type, labels []*metapb.StoreLabel) error { c, err := h.GetRaftCluster() if err != nil { return err diff --git a/tests/pdctl/store/store_test.go b/tests/pdctl/store/store_test.go index e136e909864..30ac64b4266 100644 --- a/tests/pdctl/store/store_test.go +++ b/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 + 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) diff --git a/tools/pd-ctl/pdctl/command/store_command.go b/tools/pd-ctl/pdctl/command/store_command.go index 3b3724778bd..62ffcda5630 100644 --- a/tools/pd-ctl/pdctl/command/store_command.go +++ b/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 []|[| ]", + Use: "limit []|[| [ ]... ]", Short: "show or set a store's rate limit", Long: "show or set a store's rate limit, 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,31 @@ 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.") + return + } 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) + } } } From b5d19ddc4e9d9307c6b7572e003e235a158e3aee Mon Sep 17 00:00:00 2001 From: ZenoTan Date: Tue, 8 Sep 2020 14:28:59 +0800 Subject: [PATCH 4/8] Fix Signed-off-by: ZenoTan --- tools/pd-ctl/pdctl/command/store_command.go | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tools/pd-ctl/pdctl/command/store_command.go b/tools/pd-ctl/pdctl/command/store_command.go index 62ffcda5630..533cb0a89cd 100644 --- a/tools/pd-ctl/pdctl/command/store_command.go +++ b/tools/pd-ctl/pdctl/command/store_command.go @@ -401,8 +401,8 @@ func storeLimitCommandFunc(cmd *cobra.Command, args []string) { postInput := map[string]interface{}{} prefix := path.Join(storesPrefix, "limit") ratePos := argsCount - 1 - if argsCount % 2 == 1 { - postInput["type"] = args[argsCount - 1] + if argsCount%2 == 1 { + postInput["type"] = args[argsCount-1] ratePos = argsCount - 2 } rate, err := strconv.ParseFloat(args[ratePos], 64) From 79b144e63c1d1d3b4a8ed93d1a01d0a28ef81c7d Mon Sep 17 00:00:00 2001 From: ZenoTan Date: Tue, 8 Sep 2020 14:55:02 +0800 Subject: [PATCH 5/8] Fix Signed-off-by: ZenoTan --- tools/pd-ctl/pdctl/command/store_command.go | 1 - 1 file changed, 1 deletion(-) diff --git a/tools/pd-ctl/pdctl/command/store_command.go b/tools/pd-ctl/pdctl/command/store_command.go index 533cb0a89cd..447a9036cb7 100644 --- a/tools/pd-ctl/pdctl/command/store_command.go +++ b/tools/pd-ctl/pdctl/command/store_command.go @@ -396,7 +396,6 @@ func storeLimitCommandFunc(cmd *cobra.Command, args []string) { } else { if args[0] != "all" { cmd.Println("Labels are an option of set all stores limit.") - return } else { postInput := map[string]interface{}{} prefix := path.Join(storesPrefix, "limit") From ef7cb44d70ddc4ec192aa701b2630d6124e7463b Mon Sep 17 00:00:00 2001 From: ZenoTan Date: Thu, 10 Sep 2020 10:35:01 +0800 Subject: [PATCH 6/8] Address comment Signed-off-by: ZenoTan --- tests/pdctl/store/store_test.go | 1 - tools/pd-ctl/pdctl/command/store_command.go | 11 ++++++----- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/tests/pdctl/store/store_test.go b/tests/pdctl/store/store_test.go index 30ac64b4266..173d716a55d 100644 --- a/tests/pdctl/store/store_test.go +++ b/tests/pdctl/store/store_test.go @@ -199,7 +199,6 @@ func (s *storeTestSuite) TestStore(c *C) { // store limit all 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)) diff --git a/tools/pd-ctl/pdctl/command/store_command.go b/tools/pd-ctl/pdctl/command/store_command.go index 447a9036cb7..4f92bc18a18 100644 --- a/tools/pd-ctl/pdctl/command/store_command.go +++ b/tools/pd-ctl/pdctl/command/store_command.go @@ -25,6 +25,7 @@ import ( var ( storesPrefix = "pd/api/v1/stores" + storesLimitPrefix = "pd/api/v1/stores/limit" storePrefix = "pd/api/v1/store/%v" ) @@ -363,7 +364,7 @@ func setStoreWeightCommandFunc(cmd *cobra.Command, args []string) { func storeLimitCommandFunc(cmd *cobra.Command, args []string) { argsCount := len(args) if argsCount <= 1 { - prefix := path.Join(storesPrefix, "limit") + prefix := storesLimitPrefix if argsCount == 1 { prefix += fmt.Sprintf("?type=%s", args[0]) } @@ -382,7 +383,7 @@ func storeLimitCommandFunc(cmd *cobra.Command, args []string) { // if the storeid is "all", set limits for all stores var prefix string if args[0] == "all" { - prefix = path.Join(storesPrefix, "limit") + prefix = storesLimitPrefix } else { prefix = fmt.Sprintf(path.Join(storePrefix, "limit"), args[0]) } @@ -398,7 +399,7 @@ func storeLimitCommandFunc(cmd *cobra.Command, args []string) { cmd.Println("Labels are an option of set all stores limit.") } else { postInput := map[string]interface{}{} - prefix := path.Join(storesPrefix, "limit") + prefix := storesLimitPrefix ratePos := argsCount - 1 if argsCount%2 == 1 { postInput["type"] = args[argsCount-1] @@ -439,7 +440,7 @@ func showAllStoresLimitCommandFunc(cmd *cobra.Command, args []string) { cmd.Usage() return } - prefix := path.Join(storesPrefix, "limit") + prefix := storesLimitPrefix if len(args) == 1 { prefix += fmt.Sprintf("?type=%s", args[0]) } @@ -472,7 +473,7 @@ func setAllLimitCommandFunc(cmd *cobra.Command, args []string) { cmd.Println("rate should be a number that > 0.") return } - prefix := path.Join(storesPrefix, "limit") + prefix := storesLimitPrefix input := map[string]interface{}{ "rate": rate, } From b35c55325ab7989e23a48439545cca85ace90f7a Mon Sep 17 00:00:00 2001 From: ZenoTan Date: Thu, 10 Sep 2020 10:45:53 +0800 Subject: [PATCH 7/8] Reformat Signed-off-by: ZenoTan --- tools/pd-ctl/pdctl/command/store_command.go | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tools/pd-ctl/pdctl/command/store_command.go b/tools/pd-ctl/pdctl/command/store_command.go index 4f92bc18a18..cc00cfecb83 100644 --- a/tools/pd-ctl/pdctl/command/store_command.go +++ b/tools/pd-ctl/pdctl/command/store_command.go @@ -24,9 +24,9 @@ import ( ) var ( - storesPrefix = "pd/api/v1/stores" + storesPrefix = "pd/api/v1/stores" storesLimitPrefix = "pd/api/v1/stores/limit" - storePrefix = "pd/api/v1/store/%v" + storePrefix = "pd/api/v1/store/%v" ) // NewStoreCommand return a stores subcommand of rootCmd From 291c46dbdaa663839a34340f9237418bb4293350 Mon Sep 17 00:00:00 2001 From: ZenoTan Date: Thu, 10 Sep 2020 11:02:18 +0800 Subject: [PATCH 8/8] Remove Signed-off-by: ZenoTan --- tests/pdctl/store/store_test.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/pdctl/store/store_test.go b/tests/pdctl/store/store_test.go index 173d716a55d..450f33f591a 100644 --- a/tests/pdctl/store/store_test.go +++ b/tests/pdctl/store/store_test.go @@ -198,7 +198,7 @@ func (s *storeTestSuite) TestStore(c *C) { // store limit all args = []string{"-u", pdAddr, "store", "limit", "all", "zone", "uk", "20", "remove-peer"} - _, output, err = pdctl.ExecuteCommandC(cmd, args...) + _, _, err = pdctl.ExecuteCommandC(cmd, args...) c.Assert(err, IsNil) limit1 = leaderServer.GetRaftCluster().GetStoreLimitByType(1, storelimit.RemovePeer) c.Assert(limit1, Equals, float64(20))