From 4e66dcc33b74c8e3c0bdacfe0c24596b35caa239 Mon Sep 17 00:00:00 2001
From: Minoru Osuka <minoru.osuka@gmail.com>
Date: Tue, 24 Mar 2020 14:11:12 +0900
Subject: [PATCH] Add keepalive options

---
 client/grpc_client.go | 14 +++++++++++++-
 server/grpc_server.go | 13 +++++++++++++
 2 files changed, 26 insertions(+), 1 deletion(-)

diff --git a/client/grpc_client.go b/client/grpc_client.go
index 95d1176..10fdd3b 100644
--- a/client/grpc_client.go
+++ b/client/grpc_client.go
@@ -18,12 +18,14 @@ import (
 	"context"
 	"log"
 	"math"
+	"time"
 
 	"github.com/golang/protobuf/ptypes/empty"
 	"github.com/mosuka/cete/errors"
 	"github.com/mosuka/cete/protobuf"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/keepalive"
 	"google.golang.org/grpc/status"
 )
 
@@ -37,7 +39,10 @@ type GRPCClient struct {
 }
 
 func NewGRPCClient(address string) (*GRPCClient, error) {
-	baseCtx := context.TODO()
+	return NewGRPCClientWithContext(address, context.TODO())
+}
+
+func NewGRPCClientWithContext(address string, baseCtx context.Context) (*GRPCClient, error) {
 	ctx, cancel := context.WithCancel(baseCtx)
 
 	dialOpts := []grpc.DialOption{
@@ -46,6 +51,13 @@ func NewGRPCClient(address string) (*GRPCClient, error) {
 			grpc.MaxCallSendMsgSize(math.MaxInt64),
 			grpc.MaxCallRecvMsgSize(math.MaxInt64),
 		),
+		grpc.WithKeepaliveParams(
+			keepalive.ClientParameters{
+				Time:                1 * time.Second,
+				Timeout:             5 * time.Second,
+				PermitWithoutStream: true,
+			},
+		),
 	}
 
 	conn, err := grpc.DialContext(ctx, address, dialOpts...)
diff --git a/server/grpc_server.go b/server/grpc_server.go
index 5cb0699..77d0a63 100644
--- a/server/grpc_server.go
+++ b/server/grpc_server.go
@@ -17,13 +17,16 @@ package server
 import (
 	"math"
 	"net"
+	"time"
 
 	grpcmiddleware "github.com/grpc-ecosystem/go-grpc-middleware"
 	grpczap "github.com/grpc-ecosystem/go-grpc-middleware/logging/zap"
+	grpc_prometheus "github.com/grpc-ecosystem/go-grpc-prometheus"
 	"github.com/mosuka/cete/metric"
 	"github.com/mosuka/cete/protobuf"
 	"go.uber.org/zap"
 	"google.golang.org/grpc"
+	"google.golang.org/grpc/keepalive"
 )
 
 type GRPCServer struct {
@@ -53,6 +56,15 @@ func NewGRPCServer(address string, raftServer *RaftServer, logger *zap.Logger) (
 				grpczap.UnaryServerInterceptor(grpcLogger),
 			),
 		),
+		grpc.KeepaliveParams(
+			keepalive.ServerParameters{
+				//MaxConnectionIdle:     0,
+				//MaxConnectionAge:      0,
+				//MaxConnectionAgeGrace: 0,
+				Time:    5 * time.Second,
+				Timeout: 5 * time.Second,
+			},
+		),
 	)
 
 	service, err := NewGRPCService(raftServer, logger)
@@ -65,6 +77,7 @@ func NewGRPCServer(address string, raftServer *RaftServer, logger *zap.Logger) (
 
 	// Initialize all metrics.
 	metric.GrpcMetrics.InitializeMetrics(server)
+	grpc_prometheus.Register(server)
 
 	listener, err := net.Listen("tcp", address)
 	if err != nil {