From ca88b69d24a6944191cf807eddee58431dd8a448 Mon Sep 17 00:00:00 2001
From: Kevin Wan <wanjunfeng@gmail.com>
Date: Sat, 21 May 2022 19:38:27 +0800
Subject: [PATCH] feat: set default connection idle time for grpc servers
 (#1922)

* feat: set default connection idle time for grpc servers

* feat: add grpc health check
---
 zrpc/internal/rpcserver.go |  7 +++++++
 zrpc/internal/server.go    | 11 +++++++++++
 2 files changed, 18 insertions(+)

diff --git a/zrpc/internal/rpcserver.go b/zrpc/internal/rpcserver.go
index be378deb..98ef126f 100644
--- a/zrpc/internal/rpcserver.go
+++ b/zrpc/internal/rpcserver.go
@@ -7,6 +7,7 @@ import (
 	"github.com/zeromicro/go-zero/core/stat"
 	"github.com/zeromicro/go-zero/zrpc/internal/serverinterceptors"
 	"google.golang.org/grpc"
+	"google.golang.org/grpc/health/grpc_health_v1"
 )
 
 type (
@@ -71,9 +72,15 @@ func (s *rpcServer) Start(register RegisterFn) error {
 		WithStreamServerInterceptors(streamInterceptors...))
 	server := grpc.NewServer(options...)
 	register(server)
+
+	// register the health check service
+	grpc_health_v1.RegisterHealthServer(server, s.health)
+	s.health.Resume()
+
 	// we need to make sure all others are wrapped up,
 	// so we do graceful stop at shutdown phase instead of wrap up phase
 	waitForCalled := proc.AddWrapUpListener(func() {
+		s.health.Shutdown()
 		server.GracefulStop()
 	})
 	defer waitForCalled()
diff --git a/zrpc/internal/server.go b/zrpc/internal/server.go
index ecdf86dd..5a5d5655 100644
--- a/zrpc/internal/server.go
+++ b/zrpc/internal/server.go
@@ -1,10 +1,16 @@
 package internal
 
 import (
+	"time"
+
 	"github.com/zeromicro/go-zero/core/stat"
 	"google.golang.org/grpc"
+	"google.golang.org/grpc/health"
+	"google.golang.org/grpc/keepalive"
 )
 
+const defaultConnectionIdleDuration = time.Minute * 5
+
 type (
 	// RegisterFn defines the method to register a server.
 	RegisterFn func(*grpc.Server)
@@ -20,6 +26,7 @@ type (
 
 	baseRpcServer struct {
 		address            string
+		health             *health.Server
 		metrics            *stat.Metrics
 		options            []grpc.ServerOption
 		streamInterceptors []grpc.StreamServerInterceptor
@@ -30,7 +37,11 @@ type (
 func newBaseRpcServer(address string, rpcServerOpts *rpcServerOptions) *baseRpcServer {
 	return &baseRpcServer{
 		address: address,
+		health:  health.NewServer(),
 		metrics: rpcServerOpts.metrics,
+		options: []grpc.ServerOption{grpc.KeepaliveParams(keepalive.ServerParameters{
+			MaxConnectionIdle: defaultConnectionIdleDuration,
+		})},
 	}
 }