From dc5b8dd7162d8d304c0800d6f5386bdd8779945a Mon Sep 17 00:00:00 2001
From: Kevin Wan <wanjunfeng@gmail.com>
Date: Fri, 3 Mar 2023 23:40:17 +0800
Subject: [PATCH] feat: support grpc client keepalive config (#2950)

---
 zrpc/client.go      | 14 ++++++++++++++
 zrpc/client_test.go |  9 +++++----
 zrpc/config.go      | 27 +++++++++++++++------------
 3 files changed, 34 insertions(+), 16 deletions(-)

diff --git a/zrpc/client.go b/zrpc/client.go
index 9966d038..db00d047 100644
--- a/zrpc/client.go
+++ b/zrpc/client.go
@@ -8,8 +8,11 @@ import (
 	"github.com/zeromicro/go-zero/zrpc/internal/auth"
 	"github.com/zeromicro/go-zero/zrpc/internal/clientinterceptors"
 	"google.golang.org/grpc"
+	"google.golang.org/grpc/keepalive"
 )
 
+const defaultClientKeepaliveTime = 20 * time.Second
+
 var (
 	// WithDialOption is an alias of internal.WithDialOption.
 	WithDialOption = internal.WithDialOption
@@ -62,6 +65,11 @@ func NewClient(c RpcClientConf, options ...ClientOption) (Client, error) {
 	if c.Timeout > 0 {
 		opts = append(opts, WithTimeout(time.Duration(c.Timeout)*time.Millisecond))
 	}
+	if c.KeepaliveTime > 0 {
+		opts = append(opts, WithDialOption(grpc.WithKeepaliveParams(keepalive.ClientParameters{
+			Time: c.KeepaliveTime,
+		})))
+	}
 
 	opts = append(opts, options...)
 
@@ -90,6 +98,12 @@ func NewClientWithTarget(target string, opts ...ClientOption) (Client, error) {
 		Timeout:    true,
 	}
 
+	opts = append([]ClientOption{
+		WithDialOption(grpc.WithKeepaliveParams(keepalive.ClientParameters{
+			Time: defaultClientKeepaliveTime,
+		})),
+	}, opts...)
+
 	return internal.NewClient(target, middlewares, opts...)
 }
 
diff --git a/zrpc/client_test.go b/zrpc/client_test.go
index 368ca2ee..70da01b2 100644
--- a/zrpc/client_test.go
+++ b/zrpc/client_test.go
@@ -113,10 +113,11 @@ func TestDepositServer_Deposit(t *testing.T) {
 	)
 	tarConfClient := MustNewClient(
 		RpcClientConf{
-			Target:  "foo",
-			App:     "foo",
-			Token:   "bar",
-			Timeout: 1000,
+			Target:        "foo",
+			App:           "foo",
+			Token:         "bar",
+			Timeout:       1000,
+			KeepaliveTime: time.Second * 15,
 			Middlewares: ClientMiddlewaresConf{
 				Trace:      true,
 				Duration:   true,
diff --git a/zrpc/config.go b/zrpc/config.go
index 7f772807..646382fd 100644
--- a/zrpc/config.go
+++ b/zrpc/config.go
@@ -1,6 +1,8 @@
 package zrpc
 
 import (
+	"time"
+
 	"github.com/zeromicro/go-zero/core/discov"
 	"github.com/zeromicro/go-zero/core/service"
 	"github.com/zeromicro/go-zero/core/stores/redis"
@@ -14,6 +16,19 @@ type (
 	// ServerMiddlewaresConf defines whether to use server middlewares.
 	ServerMiddlewaresConf = internal.ServerMiddlewaresConf
 
+	// A RpcClientConf is a rpc client config.
+	RpcClientConf struct {
+		Etcd          discov.EtcdConf `json:",optional,inherit"`
+		Endpoints     []string        `json:",optional"`
+		Target        string          `json:",optional"`
+		App           string          `json:",optional"`
+		Token         string          `json:",optional"`
+		NonBlock      bool            `json:",optional"`
+		Timeout       int64           `json:",default=2000"`
+		KeepaliveTime time.Duration   `json:",default=20s"`
+		Middlewares   ClientMiddlewaresConf
+	}
+
 	// A RpcServerConf is a rpc server config.
 	RpcServerConf struct {
 		service.ServiceConf
@@ -29,18 +44,6 @@ type (
 		Health      bool `json:",default=true"`
 		Middlewares ServerMiddlewaresConf
 	}
-
-	// A RpcClientConf is a rpc client config.
-	RpcClientConf struct {
-		Etcd        discov.EtcdConf `json:",optional,inherit"`
-		Endpoints   []string        `json:",optional"`
-		Target      string          `json:",optional"`
-		App         string          `json:",optional"`
-		Token       string          `json:",optional"`
-		NonBlock    bool            `json:",optional"`
-		Timeout     int64           `json:",default=2000"`
-		Middlewares ClientMiddlewaresConf
-	}
 )
 
 // NewDirectClientConf returns a RpcClientConf.