support setting float for qps

Signed-off-by: Ryan Leung <[email protected]>
tikv · Nov 8, 2024 · d53b4b5 · d53b4b5
1 parent 051907b
commit d53b4b5
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 13 deletions.
diff --git a/tools/pd-ctl/pdctl/command/config_command.go b/tools/pd-ctl/pdctl/command/config_command.go
@@ -471,14 +471,20 @@ func setServiceMiddlewareCommandFunc(cmd *cobra.Command, args []string) {
 	input := map[string]any{
 		"label": args[1],
 	}
-	value, err := strconv.ParseUint(args[3], 10, 64)
-	if err != nil {
-		cmd.Println(err)
-		return
-	}
+
 	if strings.ToLower(args[2]) == "qps" {
+		value, err := strconv.ParseFloat(args[3], 64)
+		if err != nil {
+			cmd.Println(err)
+			return
+		}
 		input["qps"] = value
 	} else if strings.ToLower(args[2]) == "concurrency" {
+		value, err := strconv.ParseUint(args[3], 10, 64)
+		if err != nil {
+			cmd.Println(err)
+			return
+		}
 		input["concurrency"] = value
 	} else {
 		cmd.Println("Input is invalid, should be qps or concurrency")

diff --git a/tools/pd-ctl/tests/config/config_test.go b/tools/pd-ctl/tests/config/config_test.go
@@ -1005,21 +1005,21 @@ func TestServiceMiddlewareConfig(t *testing.T) {
 	re.NoError(err)
 	conf.AuditConfig.EnableAudit = false
 	check()
-	_, err = tests.ExecuteCommand(cmd, "-u", pdAddr, "config", "set", "service-middleware", "rate-limit", "GetRegion", "qps", "100")
+	_, err = tests.ExecuteCommand(cmd, "-u", pdAddr, "config", "set", "service-middleware", "rate-limit", "GetRegion", "qps", "100.1")
 	re.NoError(err)
-	conf.RateLimitConfig.LimiterConfig["GetRegion"] = ratelimit.DimensionConfig{QPS: 100, QPSBurst: 100}
+	conf.RateLimitConfig.LimiterConfig["GetRegion"] = ratelimit.DimensionConfig{QPS: 100.1, QPSBurst: 100}
 	check()
-	_, err = tests.ExecuteCommand(cmd, "-u", pdAddr, "config", "set", "service-middleware", "grpc-rate-limit", "GetRegion", "qps", "101")
+	_, err = tests.ExecuteCommand(cmd, "-u", pdAddr, "config", "set", "service-middleware", "grpc-rate-limit", "GetRegion", "qps", "101.1")
 	re.NoError(err)
-	conf.GRPCRateLimitConfig.LimiterConfig["GetRegion"] = ratelimit.DimensionConfig{QPS: 101, QPSBurst: 101}
+	conf.GRPCRateLimitConfig.LimiterConfig["GetRegion"] = ratelimit.DimensionConfig{QPS: 101.1, QPSBurst: 101}
 	check()
 	_, err = tests.ExecuteCommand(cmd, "-u", pdAddr, "config", "set", "service-middleware", "rate-limit", "GetRegion", "concurrency", "10")
 	re.NoError(err)
-	conf.RateLimitConfig.LimiterConfig["GetRegion"] = ratelimit.DimensionConfig{QPS: 100, QPSBurst: 100, ConcurrencyLimit: 10}
+	conf.RateLimitConfig.LimiterConfig["GetRegion"] = ratelimit.DimensionConfig{QPS: 100.1, QPSBurst: 100, ConcurrencyLimit: 10}
 	check()
 	_, err = tests.ExecuteCommand(cmd, "-u", pdAddr, "config", "set", "service-middleware", "grpc-rate-limit", "GetRegion", "concurrency", "11")
 	re.NoError(err)
-	conf.GRPCRateLimitConfig.LimiterConfig["GetRegion"] = ratelimit.DimensionConfig{QPS: 101, QPSBurst: 101, ConcurrencyLimit: 11}
+	conf.GRPCRateLimitConfig.LimiterConfig["GetRegion"] = ratelimit.DimensionConfig{QPS: 101.1, QPSBurst: 101, ConcurrencyLimit: 11}
 	check()
 	output, err := tests.ExecuteCommand(cmd, "-u", pdAddr, "config", "set", "service-middleware", "xxx", "GetRegion", "qps", "1000")
 	re.NoError(err)
@@ -1032,14 +1032,14 @@ func TestServiceMiddlewareConfig(t *testing.T) {
 	re.Contains(string(output), "Input is invalid")
 	output, err = tests.ExecuteCommand(cmd, "-u", pdAddr, "config", "set", "service-middleware", "grpc-rate-limit", "GetRegion", "qps", "xxx")
 	re.NoError(err)
-	re.Contains(string(output), "strconv.ParseUint")
+	re.Contains(string(output), "strconv.ParseFloat")
 	_, err = tests.ExecuteCommand(cmd, "-u", pdAddr, "config", "set", "service-middleware", "grpc-rate-limit", "enable-grpc-rate-limit", "false")
 	re.NoError(err)
 	conf.GRPCRateLimitConfig.EnableRateLimit = false
 	check()
 	_, err = tests.ExecuteCommand(cmd, "-u", pdAddr, "config", "set", "service-middleware", "rate-limit", "GetRegion", "concurrency", "0")
 	re.NoError(err)
-	conf.RateLimitConfig.LimiterConfig["GetRegion"] = ratelimit.DimensionConfig{QPS: 100, QPSBurst: 100}
+	conf.RateLimitConfig.LimiterConfig["GetRegion"] = ratelimit.DimensionConfig{QPS: 100.1, QPSBurst: 100}
 	check()
 	_, err = tests.ExecuteCommand(cmd, "-u", pdAddr, "config", "set", "service-middleware", "rate-limit", "GetRegion", "qps", "0")
 	re.NoError(err)