From 8c8ee3b7a5e7889129475a77a7e3ebf00f2157ef Mon Sep 17 00:00:00 2001
From: Karim Radhouani <medkarimrdi@gmail.com>
Date: Tue, 19 Mar 2024 14:41:50 -0700
Subject: [PATCH] allow enabling and configuring gRPC keepalive for gNMI
 clients

---
 docs/user_guide/targets/targets.md | 15 +++++++++++++++
 pkg/api/types/target.go            | 16 ++++++++++++++++
 2 files changed, 31 insertions(+)

diff --git a/docs/user_guide/targets/targets.md b/docs/user_guide/targets/targets.md
index cdc19e8e..35ad1636 100644
--- a/docs/user_guide/targets/targets.md
+++ b/docs/user_guide/targets/targets.md
@@ -294,6 +294,21 @@ targets:
     # When empty or set to 0s, the Golang default (15s) applies.
     # Disabled if set to a negative value.
     tcp-keepalive: 0s
+    # sets gRPC keepalive parameters according to: 
+    # https://github.com/grpc/proposal/blob/master/A8-client-side-keepalive.md
+    grpc-keepalive:
+      # After a duration of this time if the client doesn't see any activity 
+      # it pings the server to see if the transport is still alive. 
+      # If set below 10s, a minimum value of 10s will be used instead.
+      time:
+      # After having pinged for keepalive check, the client waits 
+      # for a duration of Timeout and if no activity is seen even 
+      # after that the connection is closed.
+      timeout:
+      # If true, client sends keepalive pings even with no active RPCs. 
+      # If false, when there are no active RPCs, 
+      # Time and Timeout will be ignored and no keepalive pings will be sent.
+      permit-without-stream: false
 ```
 
 ### Example
diff --git a/pkg/api/types/target.go b/pkg/api/types/target.go
index c93cead8..7605c220 100644
--- a/pkg/api/types/target.go
+++ b/pkg/api/types/target.go
@@ -22,6 +22,7 @@ import (
 	"google.golang.org/grpc/credentials/insecure"
 	"google.golang.org/grpc/credentials/oauth"
 	"google.golang.org/grpc/encoding/gzip"
+	"google.golang.org/grpc/keepalive"
 
 	"github.com/openconfig/gnmic/pkg/api/utils"
 )
@@ -150,6 +151,13 @@ type TargetConfig struct {
 	Metadata         map[string]string `mapstructure:"metadata,omitempty" yaml:"metadata,omitempty" json:"metadata,omitempty"`
 	CipherSuites     []string          `mapstructure:"cipher-suites,omitempty" yaml:"cipher-suites,omitempty" json:"cipher-suites,omitempty"`
 	TCPKeepalive     time.Duration     `mapstructure:"tcp-keepalive,omitempty" yaml:"tcp-keepalive,omitempty" json:"tcp-keepalive,omitempty"`
+	GRPCKeepalive    *clientKeepalive  `mapstructure:"grpc-keepalive,omitempty" yaml:"grpc-keepalive,omitempty" json:"grpc-keepalive,omitempty"`
+}
+
+type clientKeepalive struct {
+	Time                time.Duration `mapstructure:"time,omitempty"`
+	Timeout             time.Duration `mapstructure:"timeout,omitempty"`
+	PermitWithoutStream bool          `mapstructure:"permit-without-stream,omitempty"`
 }
 
 func (tc TargetConfig) String() string {
@@ -227,6 +235,14 @@ func (tc *TargetConfig) GrpcDialOptions() ([]grpc.DialOption, error) {
 	if tc.Gzip != nil && *tc.Gzip {
 		tOpts = append(tOpts, grpc.WithDefaultCallOptions(grpc.UseCompressor(gzip.Name)))
 	}
+	// gRPC keepalive
+	if tc.GRPCKeepalive != nil {
+		tOpts = append(tOpts, grpc.WithKeepaliveParams(keepalive.ClientParameters{
+			Time:                tc.GRPCKeepalive.Time,
+			Timeout:             tc.GRPCKeepalive.Timeout,
+			PermitWithoutStream: tc.GRPCKeepalive.PermitWithoutStream,
+		}))
+	}
 	// insecure
 	if tc.Insecure != nil && *tc.Insecure {
 		tOpts = append(tOpts,