From 45358d4551be7366329766291f45f232f64d3312 Mon Sep 17 00:00:00 2001
From: Evgenii Stratonikov <e.stratonikov@yadro.com>
Date: Fri, 3 Feb 2023 16:43:09 +0300
Subject: [PATCH] [#2] rpc/client: Allow to override low-level gRPC options

Signed-off-by: Evgenii Stratonikov <e.stratonikov@yadro.com>
---
 rpc/client/client.go  |  9 +++++++++
 rpc/client/connect.go | 15 +--------------
 rpc/client/options.go | 24 +++++++++++++++++++++---
 3 files changed, 31 insertions(+), 17 deletions(-)

diff --git a/rpc/client/client.go b/rpc/client/client.go
index e184d33..7e914db 100644
--- a/rpc/client/client.go
+++ b/rpc/client/client.go
@@ -1,5 +1,10 @@
 package client
 
+import (
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/credentials"
+)
+
 // Client represents client for exchanging messages
 // with a remote server using Protobuf RPC.
 type Client struct {
@@ -15,5 +20,9 @@ func New(opts ...Option) *Client {
 		opt(&c.cfg)
 	}
 
+	if c.tlsCfg != nil {
+		c.grpcDialOpts = append(c.grpcDialOpts, grpc.WithTransportCredentials(credentials.NewTLS(c.tlsCfg)))
+	}
+
 	return &c
 }
diff --git a/rpc/client/connect.go b/rpc/client/connect.go
index 2f62130..29f4189 100644
--- a/rpc/client/connect.go
+++ b/rpc/client/connect.go
@@ -8,8 +8,6 @@ import (
 	"net/url"
 
 	grpcstd "google.golang.org/grpc"
-	"google.golang.org/grpc/credentials"
-	"google.golang.org/grpc/credentials/insecure"
 )
 
 var errInvalidEndpoint = errors.New("invalid endpoint options")
@@ -23,21 +21,10 @@ func (c *Client) openGRPCConn(ctx context.Context) error {
 		return errInvalidEndpoint
 	}
 
-	var creds credentials.TransportCredentials
-
-	if c.tlsCfg != nil {
-		creds = credentials.NewTLS(c.tlsCfg)
-	} else {
-		creds = insecure.NewCredentials()
-	}
-
 	dialCtx, cancel := context.WithTimeout(ctx, c.dialTimeout)
 	var err error
 
-	c.conn, err = grpcstd.DialContext(dialCtx, c.addr,
-		grpcstd.WithTransportCredentials(creds),
-		grpcstd.WithBlock(),
-	)
+	c.conn, err = grpcstd.DialContext(dialCtx, c.addr, c.grpcDialOpts...)
 
 	cancel()
 
diff --git a/rpc/client/options.go b/rpc/client/options.go
index 0e2455c..005594c 100644
--- a/rpc/client/options.go
+++ b/rpc/client/options.go
@@ -5,6 +5,8 @@ import (
 	"time"
 
 	"google.golang.org/grpc"
+	"google.golang.org/grpc/credentials/insecure"
+	"google.golang.org/grpc/keepalive"
 )
 
 const (
@@ -21,19 +23,28 @@ type cfg struct {
 	dialTimeout time.Duration
 	rwTimeout   time.Duration
 
-	tlsCfg *tls.Config
+	tlsCfg       *tls.Config
+	grpcDialOpts []grpc.DialOption
 
 	conn *grpc.ClientConn
 }
 
 const (
-	defaultDialTimeout = 5 * time.Second
-	defaultRWTimeout   = 1 * time.Minute
+	defaultDialTimeout      = 5 * time.Second
+	defaultKeepAliveTimeout = 5 * time.Second
+	defaultRWTimeout        = 1 * time.Minute
 )
 
 func (c *cfg) initDefault() {
 	c.dialTimeout = defaultDialTimeout
 	c.rwTimeout = defaultRWTimeout
+	c.grpcDialOpts = []grpc.DialOption{
+		grpc.WithBlock(),
+		grpc.WithTransportCredentials(insecure.NewCredentials()),
+		grpc.WithKeepaliveParams(keepalive.ClientParameters{
+			Timeout: defaultKeepAliveTimeout,
+		}),
+	}
 }
 
 // WithNetworkAddress returns option to specify
@@ -115,3 +126,10 @@ func WithGRPCConn(v *grpc.ClientConn) Option {
 		}
 	}
 }
+
+// WithGRPCDialOptions returns an option to specify grpc.DialOption.
+func WithGRPCDialOptions(opts []grpc.DialOption) Option {
+	return func(c *cfg) {
+		c.grpcDialOpts = append(c.grpcDialOpts, opts...)
+	}
+}