From 5d931cd78f1890bf42ac18425c392ba85a77e7ed Mon Sep 17 00:00:00 2001
From: Gaurav-Kumbhat <Gaurav.Kumbhat@ibm.com>
Date: Sun, 11 May 2025 09:40:44 -0500
Subject: [PATCH] :sparkles: Add tonic endpoint settings

Signed-off-by: Gaurav-Kumbhat <Gaurav.Kumbhat@ibm.com>
---
 ginepro/src/balanced_channel.rs | 42 +++++++++++++++++++++++++++++++++
 ginepro/src/service_probe.rs    | 22 +++++++++++++++++
 2 files changed, 64 insertions(+)

diff --git a/ginepro/src/balanced_channel.rs b/ginepro/src/balanced_channel.rs
index eea84c3..b0603a4 100644
--- a/ginepro/src/balanced_channel.rs
+++ b/ginepro/src/balanced_channel.rs
@@ -103,6 +103,9 @@ pub struct LoadBalancedChannelBuilder<T, S> {
     connect_timeout: Option<Duration>,
     tls_config: Option<ClientTlsConfig>,
     lookup_service: Option<T>,
+    keep_alive_timeout: Option<Duration>,
+    http2_keep_alive_interval: Option<Duration>,
+    keep_alive_while_idle: bool,
 }
 
 impl<S> LoadBalancedChannelBuilder<DnsResolver, S>
@@ -125,6 +128,9 @@ where
             tls_config: None,
             lookup_service: None,
             resolution_strategy: ResolutionStrategy::Lazy,
+            keep_alive_timeout: None,
+            http2_keep_alive_interval: None,
+            keep_alive_while_idle: false,
         }
     }
 
@@ -141,6 +147,9 @@ where
             timeout: self.timeout,
             connect_timeout: self.connect_timeout,
             resolution_strategy: self.resolution_strategy,
+            keep_alive_timeout: self.keep_alive_timeout,
+            http2_keep_alive_interval: self.http2_keep_alive_interval,
+            keep_alive_while_idle: self.keep_alive_while_idle,
         }
     }
 }
@@ -177,6 +186,36 @@ where
         }
     }
 
+    /// Set the maximum time client should wait for idle keep-alive connections
+    ///
+    /// Uses hyper’s default otherwise.
+    pub fn keep_alive_timeout(self, keep_alive_timeout: Duration) -> LoadBalancedChannelBuilder<T, S> {
+        Self {
+            keep_alive_timeout: Some(keep_alive_timeout),
+            ..self
+        }
+    }
+
+    /// Set http2 KEEP_ALIVE_INTERVAL
+    ///
+    /// Uses hyper’s default otherwise.
+    pub fn http2_keep_alive_interval(self, http2_keep_alive_interval: Duration) -> LoadBalancedChannelBuilder<T, S> {
+        Self {
+            http2_keep_alive_interval: Some(http2_keep_alive_interval),
+            ..self
+        }
+    }
+
+    /// Set http2 KEEP_ALIVE_WHILE_IDLE
+    ///
+    /// Uses hyper’s default otherwise.
+    pub fn keep_alive_while_idle(self, keep_alive_while_idle: bool) -> LoadBalancedChannelBuilder<T, S> {
+        Self {
+            keep_alive_while_idle: keep_alive_while_idle,
+            ..self
+        }
+    }
+
     /// Set the [`ResolutionStrategy`].
     ///
     /// Default set to [`ResolutionStrategy::Lazy`].
@@ -234,6 +273,9 @@ where
             probe_interval: self
                 .probe_interval
                 .unwrap_or_else(|| Duration::from_secs(10)),
+            keep_alive_timeout: self.keep_alive_timeout,
+            http2_keep_alive_interval: self.http2_keep_alive_interval,
+            keep_alive_while_idle: self.keep_alive_while_idle,
         };
 
         let tls_config = self.tls_config.map(|mut tls_config| {
diff --git a/ginepro/src/service_probe.rs b/ginepro/src/service_probe.rs
index 535d2fa..d20f654 100644
--- a/ginepro/src/service_probe.rs
+++ b/ginepro/src/service_probe.rs
@@ -2,6 +2,7 @@ use crate::{LookupService, ServiceDefinition};
 use std::collections::HashSet;
 use std::net::SocketAddr;
 use tokio::sync::mpsc::Sender;
+use tokio::time::Duration;
 use tonic::transport::{channel::Endpoint, ClientTlsConfig};
 use tower::discover::Change;
 
@@ -41,6 +42,9 @@ where
     endpoints: HashSet<SocketAddr>,
     endpoint_reporter: Sender<Change<SocketAddr, Endpoint>>,
     tls_config: Option<ClientTlsConfig>,
+    keep_alive_timeout: Option<Duration>,
+    http2_keep_alive_interval: Option<Duration>,
+    keep_alive_while_idle: bool,
 }
 
 /// Config parameters to customize the behavior of `GrpcServiceProbe`.
@@ -59,6 +63,12 @@ where
     pub endpoint_timeout: Option<tokio::time::Duration>,
     /// A connection timeout that will be applied to every endpoint.
     pub endpoint_connect_timeout: Option<tokio::time::Duration>,
+    /// Keep alive timeout for the underlying connection.
+    pub keep_alive_timeout: Option<Duration>,
+    /// HTTP2 keep alive interval
+    pub http2_keep_alive_interval: Option<Duration>,
+    /// Enable keep alive while idle for the underlying connection.
+    pub keep_alive_while_idle: bool,
 }
 
 impl<Lookup: LookupService> GrpcServiceProbe<Lookup> {
@@ -78,6 +88,9 @@ impl<Lookup: LookupService> GrpcServiceProbe<Lookup> {
             endpoint_reporter,
             scheme: http::uri::Scheme::HTTP,
             tls_config: None,
+            keep_alive_timeout: config.keep_alive_timeout,
+            http2_keep_alive_interval: config.http2_keep_alive_interval,
+            keep_alive_while_idle: config.keep_alive_while_idle,
         }
     }
 
@@ -231,6 +244,15 @@ impl<Lookup: LookupService> GrpcServiceProbe<Lookup> {
         if let Some(ref connect_timeout) = self.endpoint_connect_timeout {
             endpoint = endpoint.connect_timeout(*connect_timeout)
         }
+        if let Some(ref timeout) = self.keep_alive_timeout{
+            endpoint = endpoint.keep_alive_timeout(*timeout);
+        }
+        if let Some(ref inteval) = self.http2_keep_alive_interval{
+            endpoint = endpoint.http2_keep_alive_interval(*inteval);
+        }
+        if self.keep_alive_while_idle{
+            endpoint = endpoint.keep_alive_while_idle(true);
+        }
 
         Some(endpoint)
     }