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)