From 4399f48d5f5a43348dbe1ab10dede3e7dd422c86 Mon Sep 17 00:00:00 2001
From: Breezewish <me@breeswish.org>
Date: Wed, 18 Jul 2018 16:57:50 +0800
Subject: [PATCH] add keep alive settings

---
 config/config.go           | 12 ++++++++++--
 config/config.toml.example | 10 +++++++++-
 store/tikv/client.go       | 15 +++++++++++++++
 tidb-server/main.go        |  2 ++
 4 files changed, 36 insertions(+), 3 deletions(-)

diff --git a/config/config.go b/config/config.go
index eab0795b6537c..84e637bf24012 100644
--- a/config/config.go
+++ b/config/config.go
@@ -217,6 +217,12 @@ type TiKVClient struct {
 	// GrpcConnectionCount is the max gRPC connections that will be established
 	// with each tikv-server.
 	GrpcConnectionCount uint `toml:"grpc-connection-count" json:"grpc-connection-count"`
+	// After a duration of this time in seconds if the client doesn't see any activity it pings
+	// the server to see if the transport is still alive.
+	GrpcKeepAliveTime uint `toml:"grpc-keepalive-time" json:"grpc-keepalive-time"`
+	// After having pinged for keepalive check, the client waits for a duration of Timeout in seconds
+	// and if no activity is seen even after that the connection is closed.
+	GrpcKeepAliveTimeout uint `toml:"grpc-keepalive-timeout" json:"grpc-keepalive-timeout"`
 	// CommitTimeout is the max time which command 'commit' will wait.
 	CommitTimeout string `toml:"commit-timeout" json:"commit-timeout"`
 }
@@ -296,8 +302,10 @@ var defaultConf = Config{
 		Reporter: OpenTracingReporter{},
 	},
 	TiKVClient: TiKVClient{
-		GrpcConnectionCount: 16,
-		CommitTimeout:       "41s",
+		GrpcConnectionCount:  16,
+		GrpcKeepAliveTime:    10,
+		GrpcKeepAliveTimeout: 3,
+		CommitTimeout:        "41s",
 	},
 	Binlog: Binlog{
 		WriteTimeout: "15s",
diff --git a/config/config.toml.example b/config/config.toml.example
index ad111b09bde80..727f0c22d43d4 100644
--- a/config/config.toml.example
+++ b/config/config.toml.example
@@ -215,6 +215,14 @@ local-agent-host-port = ""
 # Max gRPC connections that will be established with each tikv-server.
 grpc-connection-count = 16
 
+# After a duration of this time in seconds if the client doesn't see any activity it pings
+# the server to see if the transport is still alive.
+grpc-keepalive-time = 10
+
+# After having pinged for keepalive check, the client waits for a duration of Timeout in seconds
+# and if no activity is seen even after that the connection is closed.
+grpc-keepalive-timeout = 3
+
 # max time for commit command, must be twice bigger than raft election timeout.
 commit-timeout = "41s"
 
@@ -228,4 +236,4 @@ write-timeout = "15s"
 
 # If IgnoreError is true, when writting binlog meets error, TiDB would stop writting binlog,
 # but still provide service.
-ignore-error = false
\ No newline at end of file
+ignore-error = false
diff --git a/store/tikv/client.go b/store/tikv/client.go
index 1da3454cff0e4..6622584d9b021 100644
--- a/store/tikv/client.go
+++ b/store/tikv/client.go
@@ -35,12 +35,22 @@ import (
 	"golang.org/x/net/context"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/credentials"
+	"google.golang.org/grpc/keepalive"
 )
 
 // MaxConnectionCount is the max gRPC connections that will be established with
 // each tikv-server.
 var MaxConnectionCount uint = 16
 
+// GrpcKeepAliveTime is the duration of time after which if the client doesn't see
+// any activity it pings the server to see if the transport is still alive.
+var GrpcKeepAliveTime = time.Duration(10) * time.Second
+
+// GrpcKeepAliveTimeout is the duration of time for which the client waits after having
+// pinged for keepalive check and if no activity is seen even after that the connection
+// is closed.
+var GrpcKeepAliveTimeout = time.Duration(3) * time.Second
+
 // MaxSendMsgSize set max gRPC request message size sent to server. If any request message size is larger than
 // current value, an error will be reported from gRPC.
 var MaxSendMsgSize = 1<<31 - 1
@@ -121,6 +131,11 @@ func (a *connArray) Init(addr string, security config.Security) error {
 			grpc.WithDefaultCallOptions(grpc.MaxCallRecvMsgSize(MaxCallMsgSize)),
 			grpc.WithDefaultCallOptions(grpc.MaxCallSendMsgSize(MaxSendMsgSize)),
 			grpc.WithBackoffMaxDelay(time.Second*3),
+			grpc.WithKeepaliveParams(keepalive.ClientParameters{
+				Time:                GrpcKeepAliveTime,
+				Timeout:             GrpcKeepAliveTimeout,
+				PermitWithoutStream: true,
+			}),
 		)
 		cancel()
 		if err != nil {
diff --git a/tidb-server/main.go b/tidb-server/main.go
index edff6cd4acaec..5cd9ba9607e4d 100644
--- a/tidb-server/main.go
+++ b/tidb-server/main.go
@@ -389,6 +389,8 @@ func setGlobalVars() {
 	if cfg.TiKVClient.GrpcConnectionCount > 0 {
 		tikv.MaxConnectionCount = cfg.TiKVClient.GrpcConnectionCount
 	}
+	tikv.GrpcKeepAliveTime = time.Duration(cfg.TiKVClient.GrpcKeepAliveTime) * time.Second
+	tikv.GrpcKeepAliveTimeout = time.Duration(cfg.TiKVClient.GrpcKeepAliveTimeout) * time.Second
 
 	// set lower_case_table_names
 	variable.SysVars["lower_case_table_names"].Value = strconv.Itoa(cfg.LowerCaseTableNames)