From 8e775ae5116e7090c7c69016a6aff4c4a2d9e200 Mon Sep 17 00:00:00 2001 From: lhy1024 Date: Wed, 28 Jun 2023 13:47:34 +0800 Subject: [PATCH] tools: add keepalive for pd-tso-bench (#6699) close tikv/pd#6681 Signed-off-by: lhy1024 Co-authored-by: ti-chi-bot[bot] <108142056+ti-chi-bot[bot]@users.noreply.github.com> --- client/grpcutil/grpcutil.go | 1 + tools/pd-tso-bench/main.go | 16 +++++++++++++++- 2 files changed, 16 insertions(+), 1 deletion(-) diff --git a/client/grpcutil/grpcutil.go b/client/grpcutil/grpcutil.go index 59b7224f29e..125f1125721 100644 --- 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 index a0041b9ef3e..e82a9dcabba 100644 --- 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)