From 044aca7f50b01a2c795f5d1c7cc7f622172a0ab8 Mon Sep 17 00:00:00 2001
From: Gyu-Ho Lee <gyuhox@gmail.com>
Date: Fri, 25 Aug 2017 13:32:46 -0700
Subject: [PATCH] integration: configure keepalive parameters for server

Signed-off-by: Gyu-Ho Lee <gyuhox@gmail.com>
---
 integration/cluster.go | 68 +++++++++++++++++++++++++++++-------------
 1 file changed, 47 insertions(+), 21 deletions(-)

diff --git a/integration/cluster.go b/integration/cluster.go
index a164480d0..363fa8309 100644
--- a/integration/cluster.go
+++ b/integration/cluster.go
@@ -53,6 +53,7 @@ import (
 	"github.com/coreos/pkg/capnslog"
 	"github.com/soheilhy/cmux"
 	"google.golang.org/grpc"
+	"google.golang.org/grpc/keepalive"
 )
 
 const (
@@ -90,14 +91,17 @@ var (
 )
 
 type ClusterConfig struct {
-	Size              int
-	PeerTLS           *transport.TLSInfo
-	ClientTLS         *transport.TLSInfo
-	DiscoveryURL      string
-	UseGRPC           bool
-	QuotaBackendBytes int64
-	MaxTxnOps         uint
-	MaxRequestBytes   uint
+	Size                  int
+	PeerTLS               *transport.TLSInfo
+	ClientTLS             *transport.TLSInfo
+	DiscoveryURL          string
+	UseGRPC               bool
+	QuotaBackendBytes     int64
+	MaxTxnOps             uint
+	MaxRequestBytes       uint
+	GRPCKeepAliveMinTime  time.Duration
+	GRPCKeepAliveInterval time.Duration
+	GRPCKeepAliveTimeout  time.Duration
 }
 
 type cluster struct {
@@ -225,12 +229,15 @@ func (c *cluster) HTTPMembers() []client.Member {
 func (c *cluster) mustNewMember(t *testing.T) *member {
 	m := mustNewMember(t,
 		memberConfig{
-			name:              c.name(rand.Int()),
-			peerTLS:           c.cfg.PeerTLS,
-			clientTLS:         c.cfg.ClientTLS,
-			quotaBackendBytes: c.cfg.QuotaBackendBytes,
-			maxTxnOps:         c.cfg.MaxTxnOps,
-			maxRequestBytes:   c.cfg.MaxRequestBytes,
+			name:                  c.name(rand.Int()),
+			peerTLS:               c.cfg.PeerTLS,
+			clientTLS:             c.cfg.ClientTLS,
+			quotaBackendBytes:     c.cfg.QuotaBackendBytes,
+			maxTxnOps:             c.cfg.MaxTxnOps,
+			maxRequestBytes:       c.cfg.MaxRequestBytes,
+			grpcKeepAliveMinTime:  c.cfg.GRPCKeepAliveMinTime,
+			grpcKeepAliveInterval: c.cfg.GRPCKeepAliveInterval,
+			grpcKeepAliveTimeout:  c.cfg.GRPCKeepAliveTimeout,
 		})
 	m.DiscoveryURL = c.cfg.DiscoveryURL
 	if c.cfg.UseGRPC {
@@ -482,6 +489,7 @@ type member struct {
 	s             *etcdserver.EtcdServer
 	serverClosers []func()
 
+	grpcServerOpts []grpc.ServerOption
 	grpcServer     *grpc.Server
 	grpcServerPeer *grpc.Server
 	grpcAddr       string
@@ -496,12 +504,15 @@ type member struct {
 func (m *member) GRPCAddr() string { return m.grpcAddr }
 
 type memberConfig struct {
-	name              string
-	peerTLS           *transport.TLSInfo
-	clientTLS         *transport.TLSInfo
-	quotaBackendBytes int64
-	maxTxnOps         uint
-	maxRequestBytes   uint
+	name                  string
+	peerTLS               *transport.TLSInfo
+	clientTLS             *transport.TLSInfo
+	quotaBackendBytes     int64
+	maxTxnOps             uint
+	maxRequestBytes       uint
+	grpcKeepAliveMinTime  time.Duration
+	grpcKeepAliveInterval time.Duration
+	grpcKeepAliveTimeout  time.Duration
 }
 
 // mustNewMember return an inited member with the given name. If peerTLS is
@@ -558,6 +569,21 @@ func mustNewMember(t *testing.T, mcfg memberConfig) *member {
 		m.MaxRequestBytes = embed.DefaultMaxRequestBytes
 	}
 	m.AuthToken = "simple" // for the purpose of integration testing, simple token is enough
+
+	m.grpcServerOpts = []grpc.ServerOption{}
+	if mcfg.grpcKeepAliveMinTime > time.Duration(0) {
+		m.grpcServerOpts = append(m.grpcServerOpts, grpc.KeepaliveEnforcementPolicy(keepalive.EnforcementPolicy{
+			MinTime:             mcfg.grpcKeepAliveMinTime,
+			PermitWithoutStream: false,
+		}))
+	}
+	if mcfg.grpcKeepAliveInterval > time.Duration(0) &&
+		mcfg.grpcKeepAliveTimeout > time.Duration(0) {
+		m.grpcServerOpts = append(m.grpcServerOpts, grpc.KeepaliveParams(keepalive.ServerParameters{
+			Time:    mcfg.grpcKeepAliveInterval,
+			Timeout: mcfg.grpcKeepAliveTimeout,
+		}))
+	}
 	return m
 }
 
@@ -672,7 +698,7 @@ func (m *member) Launch() error {
 				return err
 			}
 		}
-		m.grpcServer = v3rpc.Server(m.s, tlscfg)
+		m.grpcServer = v3rpc.Server(m.s, tlscfg, m.grpcServerOpts...)
 		m.grpcServerPeer = v3rpc.Server(m.s, peerTLScfg)
 		m.serverClient = v3client.New(m.s)
 		lockpb.RegisterLockServer(m.grpcServer, v3lock.NewLockServer(m.serverClient))