Support use-tso-server-proxy option and remove use-tso-server-proxy o…

…ption.

Signed-off-by: Bin Shi <binshi.bing@gmail.com>

client/client.go

@@ -244,15 +244,6 @@ func WithForwardingOption(enableForwarding bool) ClientOption {
 	}
 }
 
-// WithTSOServerProxyOption configures the client to use TSO server proxy,
-// i.e., the client will send TSO requests to the API leader (the TSO server
-// proxy) which will forward the requests to the TSO servers.
-func WithTSOServerProxyOption(useTSOServerProxy bool) ClientOption {
-	return func(c *client) {
-		c.option.useTSOServerProxy = useTSOServerProxy
-	}
-}
-
 // WithMaxErrorRetry configures the client max retry times when connect meets error.
 func WithMaxErrorRetry(count int) ClientOption {
 	return func(c *client) {
@@ -596,11 +587,6 @@ func (c *client) setServiceMode(newMode pdpb.ServiceMode) {
 	c.Lock()
 	defer c.Unlock()
 
-	if c.option.useTSOServerProxy && newMode == pdpb.ServiceMode_API_SVC_MODE {
-		// If we are using TSO server proxy, we always use PD_SVC_MODE.
-		newMode = pdpb.ServiceMode_PD_SVC_MODE
-	}
-
 	if newMode == c.serviceMode {
 		return
 	}

client/option.go

@@ -48,13 +48,12 @@ const (
 // It provides the ability to change some PD client's options online from the outside.
 type option struct {
 	// Static options.
-	gRPCDialOptions   []grpc.DialOption
-	timeout           time.Duration
-	maxRetryTimes     int
-	enableForwarding  bool
-	metricsLabels     prometheus.Labels
-	initMetrics       bool
-	useTSOServerProxy bool
+	gRPCDialOptions  []grpc.DialOption
+	timeout          time.Duration
+	maxRetryTimes    int
+	enableForwarding bool
+	metricsLabels    prometheus.Labels
+	initMetrics      bool
 
 	// Dynamic options.
 	dynamicOptions [dynamicOptionCount]atomic.Value

tools/pd-tso-bench/main.go

@@ -55,13 +55,13 @@ var (
 	caPath                         = flag.String("cacert", "", "path of file that contains list of trusted SSL CAs")
 	certPath                       = flag.String("cert", "", "path of file that contains X509 certificate in PEM format")
 	keyPath                        = flag.String("key", "", "path of file that contains X509 key in PEM format")
-	useTSOServerProxy              = flag.Bool("use-tso-server-proxy", false, "whether send tso requests to tso server proxy instead of tso service directly")
 	maxBatchWaitInterval           = flag.Duration("batch-interval", 0, "the max batch wait interval")
 	enableTSOFollowerProxy         = flag.Bool("enable-tso-follower-proxy", false, "whether enable the TSO Follower Proxy")
 	enableFaultInjection           = flag.Bool("enable-fault-injection", false, "whether enable fault injection")
 	faultInjectionRate             = flag.Float64("fault-injection-rate", 0.01, "the failure rate [0.0001, 1]. 0.01 means 1% failure rate")
 	maxTSOSendIntervalMilliseconds = flag.Int("max-send-interval-ms", 0, "max tso send interval in milliseconds, 60s by default")
-	keyspace                       = flag.Uint("keyspace", 0, "the id of the keyspac to access")
+	keyspaceID                     = flag.Uint("keyspace-id", 0, "the id of the keyspace to access")
+	keyspaceName                   = flag.String("keyspace-name", "", "the name of the keyspace to access")
 	wg                             sync.WaitGroup
 )
 
@@ -356,7 +356,7 @@ func reqWorker(ctx context.Context, pdClients []pd.Client, clientIdx int, durCh
 		err                    error
 		maxRetryTime           int           = 120
 		sleepIntervalOnFailure time.Duration = 1000 * time.Millisecond
-		totalSleepBeforeGetTS  time.Duration = 0
+		totalSleepBeforeGetTS  time.Duration
 	)
 	pdCli := pdClients[clientIdx]
 
@@ -422,22 +422,27 @@ func createPDClient(ctx context.Context) (pd.Client, error) {
 	)
 
 	opts := make([]pd.ClientOption, 0)
-	if *useTSOServerProxy {
-		opts = append(opts, pd.WithTSOServerProxyOption(true))
-	}
-
 	opts = append(opts, pd.WithGRPCDialOptions(
 		grpc.WithKeepaliveParams(keepalive.ClientParameters{
 			Time:    keepaliveTime,
 			Timeout: keepaliveTimeout,
 		}),
 	))
 
-	pdCli, err = pd.NewClientWithKeyspace(ctx, uint32(*keyspace), []string{*pdAddrs}, pd.SecurityOption{
-		CAPath:   *caPath,
-		CertPath: *certPath,
-		KeyPath:  *keyPath,
-	}, opts...)
+	if len(*keyspaceName) > 0 {
+		apiCtx := pd.NewAPIContextV2(*keyspaceName)
+		pdCli, err = pd.NewClientWithAPIContext(ctx, apiCtx, []string{*pdAddrs}, pd.SecurityOption{
+			CAPath:   *caPath,
+			CertPath: *certPath,
+			KeyPath:  *keyPath,
+		}, opts...)
+	} else {
+		pdCli, err = pd.NewClientWithKeyspace(ctx, uint32(*keyspaceID), []string{*pdAddrs}, pd.SecurityOption{
+			CAPath:   *caPath,
+			CertPath: *certPath,
+			KeyPath:  *keyPath,
+		}, opts...)
+	}
 	if err != nil {
 		return nil, err
 	}