From 76dc560ca60f40fa5ce0369d2b739334be300afc Mon Sep 17 00:00:00 2001
From: Ryan Leung <rleungx@gmail.com>
Date: Thu, 26 Sep 2024 10:51:01 +0800
Subject: [PATCH] tools: support changing service-middleware config through
 pd-ctl (#8645)

close tikv/pd#4373

Signed-off-by: Ryan Leung <rleungx@gmail.com>

Co-authored-by: ti-chi-bot[bot] <108142056+ti-chi-bot[bot]@users.noreply.github.com>
---
 tools/pd-ctl/pdctl/command/config_command.go | 100 ++++++++++++++++---
 tools/pd-ctl/tests/config/config_test.go     |  85 ++++++++++++++++
 2 files changed, 172 insertions(+), 13 deletions(-)

diff --git a/tools/pd-ctl/pdctl/command/config_command.go b/tools/pd-ctl/pdctl/command/config_command.go
index 4efaae3fa13..a8e5af73638 100644
--- a/tools/pd-ctl/pdctl/command/config_command.go
+++ b/tools/pd-ctl/pdctl/command/config_command.go
@@ -17,6 +17,7 @@ package command
 import (
 	"bytes"
 	"encoding/json"
+	"fmt"
 	"net/http"
 	"net/url"
 	"os"
@@ -33,19 +34,20 @@ import (
 )
 
 const (
-	configPrefix          = "pd/api/v1/config"
-	schedulePrefix        = "pd/api/v1/config/schedule"
-	replicatePrefix       = "pd/api/v1/config/replicate"
-	labelPropertyPrefix   = "pd/api/v1/config/label-property"
-	clusterVersionPrefix  = "pd/api/v1/config/cluster-version"
-	rulesPrefix           = "pd/api/v1/config/rules"
-	rulesBatchPrefix      = "pd/api/v1/config/rules/batch"
-	rulePrefix            = "pd/api/v1/config/rule"
-	ruleGroupPrefix       = "pd/api/v1/config/rule_group"
-	ruleGroupsPrefix      = "pd/api/v1/config/rule_groups"
-	replicationModePrefix = "pd/api/v1/config/replication-mode"
-	ruleBundlePrefix      = "pd/api/v1/config/placement-rule"
-	pdServerPrefix        = "pd/api/v1/config/pd-server"
+	configPrefix                  = "pd/api/v1/config"
+	schedulePrefix                = "pd/api/v1/config/schedule"
+	replicatePrefix               = "pd/api/v1/config/replicate"
+	labelPropertyPrefix           = "pd/api/v1/config/label-property"
+	clusterVersionPrefix          = "pd/api/v1/config/cluster-version"
+	rulesPrefix                   = "pd/api/v1/config/rules"
+	rulesBatchPrefix              = "pd/api/v1/config/rules/batch"
+	rulePrefix                    = "pd/api/v1/config/rule"
+	ruleGroupPrefix               = "pd/api/v1/config/rule_group"
+	ruleGroupsPrefix              = "pd/api/v1/config/rule_groups"
+	replicationModePrefix         = "pd/api/v1/config/replication-mode"
+	ruleBundlePrefix              = "pd/api/v1/config/placement-rule"
+	pdServerPrefix                = "pd/api/v1/config/pd-server"
+	serviceMiddlewareConfigPrefix = "pd/api/v1/service-middleware/config"
 	// flagFromAPIServer has no influence for pd mode, but it is useful for us to debug in api mode.
 	flagFromAPIServer = "from_api_server"
 )
@@ -77,6 +79,7 @@ func NewShowConfigCommand() *cobra.Command {
 	sc.AddCommand(NewShowClusterVersionCommand())
 	sc.AddCommand(newShowReplicationModeCommand())
 	sc.AddCommand(NewShowServerConfigCommand())
+	sc.AddCommand(NewShowServiceMiddlewareConfigCommand())
 	sc.Flags().Bool(flagFromAPIServer, false, "read data from api server rather than micro service")
 	return sc
 }
@@ -151,6 +154,16 @@ func NewShowServerConfigCommand() *cobra.Command {
 	}
 }
 
+// NewShowReplicationConfigCommand return a show all subcommand of show subcommand
+func NewShowServiceMiddlewareConfigCommand() *cobra.Command {
+	sc := &cobra.Command{
+		Use:   "service-middleware",
+		Short: "show service middleware config of PD",
+		Run:   showServiceMiddlewareConfigCommandFunc,
+	}
+	return sc
+}
+
 // NewSetConfigCommand return a set subcommand of configCmd
 func NewSetConfigCommand() *cobra.Command {
 	sc := &cobra.Command{
@@ -161,6 +174,7 @@ func NewSetConfigCommand() *cobra.Command {
 	sc.AddCommand(NewSetLabelPropertyCommand())
 	sc.AddCommand(NewSetClusterVersionCommand())
 	sc.AddCommand(newSetReplicationModeCommand())
+	sc.AddCommand(newSetServiceMiddlewareCommand())
 	return sc
 }
 
@@ -192,6 +206,14 @@ func newSetReplicationModeCommand() *cobra.Command {
 	}
 }
 
+func newSetServiceMiddlewareCommand() *cobra.Command {
+	return &cobra.Command{
+		Use:   "service-middleware <type> [<key> <value> | <label> <qps|concurrency> <value>]",
+		Short: "set service middleware config",
+		Run:   setServiceMiddlewareCommandFunc,
+	}
+}
+
 // NewDeleteConfigCommand a set subcommand of cfgCmd
 func NewDeleteConfigCommand() *cobra.Command {
 	sc := &cobra.Command{
@@ -333,6 +355,16 @@ func showServerCommandFunc(cmd *cobra.Command, _ []string) {
 	cmd.Println(r)
 }
 
+func showServiceMiddlewareConfigCommandFunc(cmd *cobra.Command, _ []string) {
+	header := buildHeader(cmd)
+	r, err := doRequest(cmd, serviceMiddlewareConfigPrefix, http.MethodGet, header)
+	if err != nil {
+		cmd.Printf("Failed to get config: %s\n", err)
+		return
+	}
+	cmd.Println(r)
+}
+
 func postConfigDataWithPath(cmd *cobra.Command, key, value, path string) error {
 	var val any
 	data := make(map[string]any)
@@ -422,6 +454,48 @@ func setReplicationModeCommandFunc(cmd *cobra.Command, args []string) {
 	}
 }
 
+func setServiceMiddlewareCommandFunc(cmd *cobra.Command, args []string) {
+	if len(args) != 3 && len(args) != 4 {
+		cmd.Println(cmd.UsageString())
+		return
+	}
+
+	if len(args) == 3 {
+		cfg := map[string]any{
+			fmt.Sprintf("%s.%s", args[0], args[1]): args[2],
+		}
+		postJSON(cmd, serviceMiddlewareConfigPrefix, cfg)
+		return
+	}
+
+	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" {
+		input["qps"] = value
+	} else if strings.ToLower(args[2]) == "concurrency" {
+		input["concurrency"] = value
+	} else {
+		cmd.Println("Input is invalid, should be qps or concurrency")
+		return
+	}
+
+	if args[0] == "rate-limit" {
+		input["type"] = "label"
+		postJSON(cmd, serviceMiddlewareConfigPrefix+"/rate-limit", input)
+		return
+	} else if args[0] == "grpc-rate-limit" {
+		postJSON(cmd, serviceMiddlewareConfigPrefix+"/grpc-rate-limit", input)
+		return
+	}
+	cmd.Printf("Failed to get correct type: %s\n", args[0])
+}
+
 // NewPlacementRulesCommand placement rules subcommand
 func NewPlacementRulesCommand() *cobra.Command {
 	c := &cobra.Command{
diff --git a/tools/pd-ctl/tests/config/config_test.go b/tools/pd-ctl/tests/config/config_test.go
index 17fd0b4ec63..aa9c624a303 100644
--- a/tools/pd-ctl/tests/config/config_test.go
+++ b/tools/pd-ctl/tests/config/config_test.go
@@ -30,6 +30,7 @@ import (
 	"github.com/pingcap/kvproto/pkg/metapb"
 	"github.com/stretchr/testify/require"
 	"github.com/stretchr/testify/suite"
+	"github.com/tikv/pd/pkg/ratelimit"
 	sc "github.com/tikv/pd/pkg/schedule/config"
 	"github.com/tikv/pd/pkg/schedule/placement"
 	"github.com/tikv/pd/pkg/utils/testutil"
@@ -937,6 +938,90 @@ func TestReplicationMode(t *testing.T) {
 	check()
 }
 
+func TestServiceMiddlewareConfig(t *testing.T) {
+	re := require.New(t)
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+	cluster, err := pdTests.NewTestCluster(ctx, 1)
+	re.NoError(err)
+	defer cluster.Destroy()
+	err = cluster.RunInitialServers()
+	re.NoError(err)
+	re.NotEmpty(cluster.WaitLeader())
+	pdAddr := cluster.GetConfig().GetClientURL()
+	cmd := ctl.GetRootCmd()
+
+	store := &metapb.Store{
+		Id:            1,
+		State:         metapb.StoreState_Up,
+		LastHeartbeat: time.Now().UnixNano(),
+	}
+	leaderServer := cluster.GetLeaderServer()
+	re.NoError(leaderServer.BootstrapCluster())
+	pdTests.MustPutStore(re, cluster, store)
+
+	conf := config.ServiceMiddlewareConfig{
+		AuditConfig: config.AuditConfig{
+			EnableAudit: true,
+		},
+		RateLimitConfig: config.RateLimitConfig{
+			EnableRateLimit: true,
+			LimiterConfig:   make(map[string]ratelimit.DimensionConfig),
+		},
+		GRPCRateLimitConfig: config.GRPCRateLimitConfig{
+			EnableRateLimit: true,
+			LimiterConfig:   make(map[string]ratelimit.DimensionConfig),
+		},
+	}
+
+	check := func() {
+		output, err := tests.ExecuteCommand(cmd, "-u", pdAddr, "config", "show", "service-middleware")
+		re.NoError(err)
+		var conf2 config.ServiceMiddlewareConfig
+		re.NoError(json.Unmarshal(output, &conf2))
+		re.Equal(conf, conf2)
+	}
+
+	check()
+
+	_, err = tests.ExecuteCommand(cmd, "-u", pdAddr, "config", "set", "service-middleware", "audit", "enable-audit", "false")
+	re.NoError(err)
+	conf.AuditConfig.EnableAudit = false
+	check()
+	_, err = tests.ExecuteCommand(cmd, "-u", pdAddr, "config", "set", "service-middleware", "rate-limit", "GetRegion", "qps", "100")
+	re.NoError(err)
+	conf.RateLimitConfig.LimiterConfig["GetRegion"] = ratelimit.DimensionConfig{QPS: 100, QPSBurst: 100}
+	check()
+	_, err = tests.ExecuteCommand(cmd, "-u", pdAddr, "config", "set", "service-middleware", "grpc-rate-limit", "GetRegion", "qps", "101")
+	re.NoError(err)
+	conf.GRPCRateLimitConfig.LimiterConfig["GetRegion"] = ratelimit.DimensionConfig{QPS: 101, 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}
+	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}
+	check()
+	output, err := tests.ExecuteCommand(cmd, "-u", pdAddr, "config", "set", "service-middleware", "xxx", "GetRegion", "qps", "1000")
+	re.NoError(err)
+	re.Contains(string(output), "correct type")
+	output, err = tests.ExecuteCommand(cmd, "-u", pdAddr, "config", "set", "service-middleware", "grpc-rate-limit", "xxx", "qps", "1000")
+	re.NoError(err)
+	re.Contains(string(output), "There is no label matched.")
+	output, err = tests.ExecuteCommand(cmd, "-u", pdAddr, "config", "set", "service-middleware", "grpc-rate-limit", "GetRegion", "xxx", "1000")
+	re.NoError(err)
+	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")
+	_, 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()
+}
+
 func (suite *configTestSuite) TestUpdateDefaultReplicaConfig() {
 	suite.env.RunTestBasedOnMode(suite.checkUpdateDefaultReplicaConfig)
 }