tools: add keepalive for pd-tso-bench (tikv#6699)

close tikv#6681 Signed-off-by: lhy1024 <[email protected]> Co-authored-by: ti-chi-bot[bot] <108142056+ti-chi-bot[bot]@users.noreply.github.com>
rleungx · Aug 2, 2023 · 0935c73 · 0935c73
1 parent 54f755b
commit 0935c73
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 1 deletion.
diff --git a/client/grpcutil/grpcutil.go b/client/grpcutil/grpcutil.go
@@ -78,6 +78,7 @@ func BuildForwardContext(ctx context.Context, addr string) context.Context {
 func GetOrCreateGRPCConn(ctx context.Context, clientConns *sync.Map, addr string, tlsCfg *tlsutil.TLSConfig, opt ...grpc.DialOption) (*grpc.ClientConn, error) {
 	conn, ok := clientConns.Load(addr)
 	if ok {
+		// TODO: check the connection state.
 		return conn.(*grpc.ClientConn), nil
 	}
 	tlsConfig, err := tlsCfg.ToTLSConfig()

diff --git a/tools/pd-tso-bench/main.go b/tools/pd-tso-bench/main.go
@@ -33,6 +33,13 @@ import (
 	"github.com/prometheus/client_golang/prometheus/promhttp"
 	pd "github.com/tikv/pd/client"
 	"go.uber.org/zap"
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/keepalive"
+)
+
+const (
+	keepaliveTime    = 10 * time.Second
+	keepaliveTimeout = 3 * time.Second
 )
 
 var (
@@ -95,11 +102,18 @@ func bench(mainCtx context.Context) {
 			err   error
 		)
 
+		opt := pd.WithGRPCDialOptions(
+			grpc.WithKeepaliveParams(keepalive.ClientParameters{
+				Time:    keepaliveTime,
+				Timeout: keepaliveTimeout,
+			}),
+		)
+
 		pdCli, err = pd.NewClientWithContext(mainCtx, []string{*pdAddrs}, pd.SecurityOption{
 			CAPath:   *caPath,
 			CertPath: *certPath,
 			KeyPath:  *keyPath,
-		})
+		}, opt)
 
 		pdCli.UpdateOption(pd.MaxTSOBatchWaitInterval, *maxBatchWaitInterval)
 		pdCli.UpdateOption(pd.EnableTSOFollowerProxy, *enableTSOFollowerProxy)