pd-ctl: avoid initializing a new PD client for each cmd execution (#7804

)

ref #7300

Avoid initializing a new PD client for each cmd execution. Now we will only initialize a new one when the cluster changes.

Signed-off-by: JmPotato <ghzpotato@gmail.com>

Co-authored-by: ti-chi-bot[bot] <108142056+ti-chi-bot[bot]@users.noreply.github.com>

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

@@ -14,18 +14,15 @@
 
 package command
 
-import (
-	"context"
-
-	"github.com/spf13/cobra"
-)
+import "github.com/spf13/cobra"
 
 // NewClusterCommand return a cluster subcommand of rootCmd
 func NewClusterCommand() *cobra.Command {
 	cmd := &cobra.Command{
-		Use:   "cluster",
-		Short: "show the cluster information",
-		Run:   showClusterCommandFunc,
+		Use:               "cluster",
+		Short:             "show the cluster information",
+		PersistentPreRunE: requirePDClient,
+		Run:               showClusterCommandFunc,
 	}
 	cmd.AddCommand(NewClusterStatusCommand())
 	return cmd
@@ -42,7 +39,7 @@ func NewClusterStatusCommand() *cobra.Command {
 }
 
 func showClusterCommandFunc(cmd *cobra.Command, _ []string) {
-	info, err := PDCli.GetCluster(context.Background())
+	info, err := PDCli.GetCluster(cmd.Context())
 	if err != nil {
 		cmd.Printf("Failed to get the cluster information: %s\n", err)
 		return
@@ -51,7 +48,7 @@ func showClusterCommandFunc(cmd *cobra.Command, _ []string) {
 }
 
 func showClusterStatusCommandFunc(cmd *cobra.Command, _ []string) {
-	status, err := PDCli.GetClusterStatus(context.Background())
+	status, err := PDCli.GetClusterStatus(cmd.Context())
 	if err != nil {
 		cmd.Printf("Failed to get the cluster status: %s\n", err)
 		return

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

@@ -16,6 +16,7 @@ package command
 
 import (
 	"bytes"
+	"crypto/tls"
 	"encoding/json"
 	"io"
 	"net/http"
@@ -24,6 +25,7 @@ import (
 	"strings"
 
 	"github.com/pingcap/errors"
+	"github.com/pingcap/kvproto/pkg/metapb"
 	"github.com/spf13/cobra"
 	pd "github.com/tikv/pd/client/http"
 	"github.com/tikv/pd/pkg/utils/apiutil"
@@ -32,44 +34,127 @@ import (
 
 const (
 	pdControlCallerID = "pd-ctl"
-	pingPrefix        = "pd/api/v1/ping"
+	clusterPrefix     = "pd/api/v1/cluster"
 )
 
+func initTLSConfig(caPath, certPath, keyPath string) (*tls.Config, error) {
+	tlsInfo := transport.TLSInfo{
+		CertFile:      certPath,
+		KeyFile:       keyPath,
+		TrustedCAFile: caPath,
+	}
+	tlsConfig, err := tlsInfo.ClientConfig()
+	if err != nil {
+		return nil, errors.WithStack(err)
+	}
+	return tlsConfig, nil
+}
+
 // PDCli is a pd HTTP client
 var PDCli pd.Client
 
-// SetNewPDClient creates a PD HTTP client with the given PD addresses and options.
-func SetNewPDClient(addrs []string, opts ...pd.ClientOption) {
+func requirePDClient(cmd *cobra.Command, _ []string) error {
+	var (
+		caPath string
+		err    error
+	)
+	caPath, err = cmd.Flags().GetString("cacert")
+	if err == nil && len(caPath) != 0 {
+		var certPath, keyPath string
+		certPath, err = cmd.Flags().GetString("cert")
+		if err != nil {
+			return err
+		}
+		keyPath, err = cmd.Flags().GetString("key")
+		if err != nil {
+			return err
+		}
+		return initNewPDClientWithTLS(cmd, caPath, certPath, keyPath)
+	}
+	return initNewPDClient(cmd)
+}
+
+// shouldInitPDClient checks whether we should create a new PD client according to the cluster information.
+func shouldInitPDClient(cmd *cobra.Command) (bool, error) {
+	// Get the cluster information the current command assigned to.
+	newClusterInfoJSON, err := doRequest(cmd, clusterPrefix, http.MethodGet, http.Header{})
+	if err != nil {
+		return false, err
+	}
+	newClusterInfo := &metapb.Cluster{}
+	err = json.Unmarshal([]byte(newClusterInfoJSON), newClusterInfo)
+	if err != nil {
+		return false, err
+	}
+	// If the PD client is nil and we get the cluster information successfully,
+	// we should initialize the PD client directly.
+	if PDCli == nil {
+		return true, nil
+	}
+	// Get current cluster information that the PD client connects to.
+	currentClusterInfo, err := PDCli.GetCluster(cmd.Context())
+	if err != nil {
+		return true, nil
+	}
+	// Compare the cluster ID to determine whether we should re-initialize the PD client.
+	return currentClusterInfo.GetId() == 0 || newClusterInfo.GetId() != currentClusterInfo.GetId(), nil
+}
+
+func initNewPDClient(cmd *cobra.Command, opts ...pd.ClientOption) error {
+	if should, err := shouldInitPDClient(cmd); !should || err != nil {
+		return err
+	}
 	if PDCli != nil {
 		PDCli.Close()
 	}
-	PDCli = pd.NewClient(pdControlCallerID, addrs, opts...)
+	PDCli = pd.NewClient(pdControlCallerID, getEndpoints(cmd), opts...)
+	return nil
+}
+
+func initNewPDClientWithTLS(cmd *cobra.Command, caPath, certPath, keyPath string) error {
+	tlsConfig, err := initTLSConfig(caPath, certPath, keyPath)
+	if err != nil {
+		return err
+	}
+	initNewPDClient(cmd, pd.WithTLSConfig(tlsConfig))
+	return nil
 }
 
-// TODO: replace dialClient with PDCli
+// TODO: replace dialClient with the PD HTTP client completely.
 var dialClient = &http.Client{
 	Transport: apiutil.NewCallerIDRoundTripper(http.DefaultTransport, pdControlCallerID),
 }
 
-// InitHTTPSClient creates https client with ca file
-func InitHTTPSClient(pdAddrs, caPath, certPath, keyPath string) error {
-	tlsInfo := transport.TLSInfo{
-		CertFile:      certPath,
-		KeyFile:       keyPath,
-		TrustedCAFile: caPath,
+// RequireHTTPSClient creates a HTTPS client if the related flags are set
+func RequireHTTPSClient(cmd *cobra.Command, args []string) error {
+	caPath, err := cmd.Flags().GetString("cacert")
+	if err == nil && len(caPath) != 0 {
+		certPath, err := cmd.Flags().GetString("cert")
+		if err != nil {
+			return err
+		}
+		keyPath, err := cmd.Flags().GetString("key")
+		if err != nil {
+			return err
+		}
+		err = initHTTPSClient(caPath, certPath, keyPath)
+		if err != nil {
+			cmd.Println(err)
+			return err
+		}
 	}
-	tlsConfig, err := tlsInfo.ClientConfig()
+	return nil
+}
+
+func initHTTPSClient(caPath, certPath, keyPath string) error {
+	tlsConfig, err := initTLSConfig(caPath, certPath, keyPath)
 	if err != nil {
-		return errors.WithStack(err)
+		return err
 	}
-
 	dialClient = &http.Client{
 		Transport: apiutil.NewCallerIDRoundTripper(
 			&http.Transport{TLSClientConfig: tlsConfig}, pdControlCallerID),
 	}
-
-	SetNewPDClient(strings.Split(pdAddrs, ","), pd.WithTLSConfig(tlsConfig.Clone()))
-
 	return nil
 }
 

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

@@ -21,6 +21,8 @@ import (
 	"github.com/spf13/cobra"
 )
 
+const pingPrefix = "pd/api/v1/ping"
+
 // NewPingCommand return a ping subcommand of rootCmd
 func NewPingCommand() *cobra.Command {
 	m := &cobra.Command{

tools/pd-ctl/pdctl/ctl.go

@@ -35,15 +35,19 @@ func init() {
 // GetRootCmd is exposed for integration tests. But it can be embedded into another suite, too.
 func GetRootCmd() *cobra.Command {
 	rootCmd := &cobra.Command{
-		Use:   "pd-ctl",
-		Short: "Placement Driver control",
+		Use:               "pd-ctl",
+		Short:             "Placement Driver control",
+		PersistentPreRunE: command.RequireHTTPSClient,
+		SilenceErrors:     true,
 	}
 
-	rootCmd.PersistentFlags().StringP("pd", "u", "http://127.0.0.1:2379", "address of pd")
+	rootCmd.PersistentFlags().StringP("pd", "u", "http://127.0.0.1:2379", "address of PD")
 	rootCmd.PersistentFlags().String("cacert", "", "path of file that contains list of trusted SSL CAs")
 	rootCmd.PersistentFlags().String("cert", "", "path of file that contains X509 certificate in PEM format")
 	rootCmd.PersistentFlags().String("key", "", "path of file that contains X509 key in PEM format")
 
+	rootCmd.Flags().ParseErrorsWhitelist.UnknownFlags = true
+
 	rootCmd.AddCommand(
 		command.NewConfigCommand(),
 		command.NewRegionCommand(),
@@ -70,39 +74,6 @@ func GetRootCmd() *cobra.Command {
 		command.NewResourceManagerCommand(),
 	)
 
-	rootCmd.Flags().ParseErrorsWhitelist.UnknownFlags = true
-	rootCmd.SilenceErrors = true
-
-	rootCmd.PersistentPreRunE = func(cmd *cobra.Command, args []string) error {
-		addrs, err := cmd.Flags().GetString("pd")
-		if err != nil {
-			return err
-		}
-
-		// TODO: refine code after replace dialClient with PDCli
-		CAPath, err := cmd.Flags().GetString("cacert")
-		if err == nil && len(CAPath) != 0 {
-			certPath, err := cmd.Flags().GetString("cert")
-			if err != nil {
-				return err
-			}
-
-			keyPath, err := cmd.Flags().GetString("key")
-			if err != nil {
-				return err
-			}
-
-			if err := command.InitHTTPSClient(addrs, CAPath, certPath, keyPath); err != nil {
-				rootCmd.Println(err)
-				return err
-			}
-		} else {
-			command.SetNewPDClient(strings.Split(addrs, ","))
-		}
-
-		return nil
-	}
-
 	return rootCmd
 }
 

tools/pd-ctl/tests/helper.go

@@ -23,16 +23,13 @@ import (
 	"github.com/tikv/pd/pkg/core"
 	"github.com/tikv/pd/pkg/response"
 	"github.com/tikv/pd/pkg/utils/typeutil"
-	"github.com/tikv/pd/tools/pd-ctl/pdctl/command"
 )
 
 // ExecuteCommand is used for test purpose.
 func ExecuteCommand(root *cobra.Command, args ...string) (output []byte, err error) {
 	buf := new(bytes.Buffer)
 	root.SetOut(buf)
 	root.SetArgs(args)
-	command.SetNewPDClient([]string{args[1]})
-	defer command.PDCli.Close()
 	err = root.Execute()
 	return buf.Bytes(), err
 }