From ddf810be7ded85eaea9eef20125911bed058afb7 Mon Sep 17 00:00:00 2001 From: JmPotato Date: Mon, 5 Feb 2024 18:08:43 +0800 Subject: [PATCH] pd-ctl: avoid initializing a new PD client for each cmd execution (#7804) ref tikv/pd#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 Co-authored-by: ti-chi-bot[bot] <108142056+ti-chi-bot[bot]@users.noreply.github.com> --- tools/pd-ctl/pdctl/command/cluster_command.go | 17 ++- tools/pd-ctl/pdctl/command/global.go | 119 +++++++++++++++--- tools/pd-ctl/pdctl/command/ping_command.go | 2 + tools/pd-ctl/pdctl/ctl.go | 43 ++----- tools/pd-ctl/tests/helper.go | 3 - 5 files changed, 118 insertions(+), 66 deletions(-) diff --git a/tools/pd-ctl/pdctl/command/cluster_command.go b/tools/pd-ctl/pdctl/command/cluster_command.go index 4fd7b8be61a..0b3cda8a867 100644 --- a/tools/pd-ctl/pdctl/command/cluster_command.go +++ b/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 diff --git a/tools/pd-ctl/pdctl/command/global.go b/tools/pd-ctl/pdctl/command/global.go index d1ea1faa236..4f20b0b35b4 100644 --- a/tools/pd-ctl/pdctl/command/global.go +++ b/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 } diff --git a/tools/pd-ctl/pdctl/command/ping_command.go b/tools/pd-ctl/pdctl/command/ping_command.go index 38e99e0516d..6622b079d47 100644 --- a/tools/pd-ctl/pdctl/command/ping_command.go +++ b/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{ diff --git a/tools/pd-ctl/pdctl/ctl.go b/tools/pd-ctl/pdctl/ctl.go index 773f28fb8bd..d5ead437dbb 100644 --- a/tools/pd-ctl/pdctl/ctl.go +++ b/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 } diff --git a/tools/pd-ctl/tests/helper.go b/tools/pd-ctl/tests/helper.go index 680ba9825d1..5c5ad5fe6e9 100644 --- a/tools/pd-ctl/tests/helper.go +++ b/tools/pd-ctl/tests/helper.go @@ -23,7 +23,6 @@ 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. @@ -31,8 +30,6 @@ func ExecuteCommand(root *cobra.Command, args ...string) (output []byte, err err 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 }