feat: add max_connection_age config to grpc server (#7031)

* feat: add `max_connection_age` config to grpc server

Signed-off-by: luofucong <luofc@foxmail.com>

* Apply suggestions from code review

Co-authored-by: Yingwen <realevenyag@gmail.com>

* fix ci

Signed-off-by: luofucong <luofc@foxmail.com>

---------

Signed-off-by: luofucong <luofc@foxmail.com>
Co-authored-by: Yingwen <realevenyag@gmail.com>

src/cmd/tests/load_config_test.rs

@@ -143,9 +143,11 @@ fn test_load_frontend_example_config() {
                 remote_write: Some(Default::default()),
                 ..Default::default()
             },
-            grpc: GrpcOptions::default()
-                .with_bind_addr("127.0.0.1:4001")
-                .with_server_addr("127.0.0.1:4001"),
+            grpc: GrpcOptions {
+                bind_addr: "127.0.0.1:4001".to_string(),
+                server_addr: "127.0.0.1:4001".to_string(),
+                ..Default::default()
+            },
             internal_grpc: Some(GrpcOptions::internal_default()),
             http: HttpOptions {
                 cors_allowed_origins: vec!["https://example.com".to_string()],

src/flow/src/server.rs

@@ -491,6 +491,7 @@ impl<'a> FlownodeServiceBuilder<'a> {
             max_recv_message_size: opts.grpc.max_recv_message_size.as_bytes() as usize,
             max_send_message_size: opts.grpc.max_send_message_size.as_bytes() as usize,
             tls: opts.grpc.tls.clone(),
+            max_connection_age: opts.grpc.max_connection_age,
         };
         let service = flownode_server.create_flow_service();
         let runtime = common_runtime::global_runtime();

src/frontend/src/server.rs

@@ -68,7 +68,7 @@ where
         }
     }
 
-    pub fn grpc_server_builder(&self, opts: &GrpcOptions) -> Result<GrpcServerBuilder> {
+    fn grpc_server_builder(&self, opts: &GrpcOptions) -> Result<GrpcServerBuilder> {
         let builder = GrpcServerBuilder::new(opts.as_config(), common_runtime::global_runtime())
             .with_tls_config(opts.tls.clone())
             .context(error::InvalidTlsConfigSnafu)?;

src/servers/Cargo.toml

@@ -64,7 +64,6 @@ futures-util.workspace = true
 headers = "0.4"
 hostname = "0.3"
 http.workspace = true
-http-body = "1"
 humantime.workspace = true
 humantime-serde.workspace = true
 hyper = { workspace = true, features = ["full"] }

src/servers/src/grpc.rs

@@ -23,6 +23,7 @@ pub mod prom_query_gateway;
 pub mod region_server;
 
 use std::net::SocketAddr;
+use std::time::Duration;
 
 use api::v1::health_check_server::{HealthCheck, HealthCheckServer};
 use api::v1::{HealthCheckRequest, HealthCheckResponse};
@@ -71,6 +72,11 @@ pub struct GrpcOptions {
     pub runtime_size: usize,
     #[serde(default = "Default::default")]
     pub tls: TlsOption,
+    /// Maximum time that a channel may exist.
+    /// Useful when the server wants to control the reconnection of its clients.
+    /// Default to `None`, means infinite.
+    #[serde(with = "humantime_serde")]
+    pub max_connection_age: Option<Duration>,
 }
 
 impl GrpcOptions {
@@ -111,6 +117,7 @@ impl GrpcOptions {
             max_recv_message_size: self.max_recv_message_size.as_bytes() as usize,
             max_send_message_size: self.max_send_message_size.as_bytes() as usize,
             tls: self.tls.clone(),
+            max_connection_age: self.max_connection_age,
         }
     }
 }
@@ -130,6 +137,7 @@ impl Default for GrpcOptions {
             flight_compression: FlightCompression::ArrowIpc,
             runtime_size: 8,
             tls: TlsOption::default(),
+            max_connection_age: None,
         }
     }
 }
@@ -148,6 +156,7 @@ impl GrpcOptions {
             flight_compression: FlightCompression::ArrowIpc,
             runtime_size: 8,
             tls: TlsOption::default(),
+            max_connection_age: None,
         }
     }
 
@@ -207,6 +216,7 @@ pub struct GrpcServer {
     >,
     bind_addr: Option<SocketAddr>,
     name: Option<String>,
+    config: GrpcServerConfig,
 }
 
 /// Grpc Server configuration
@@ -217,6 +227,10 @@ pub struct GrpcServerConfig {
     // Max gRPC sending(encoding) message size
     pub max_send_message_size: usize,
     pub tls: TlsOption,
+    /// Maximum time that a channel may exist.
+    /// Useful when the server wants to control the reconnection of its clients.
+    /// Default to `None`, means infinite.
+    pub max_connection_age: Option<Duration>,
 }
 
 impl Default for GrpcServerConfig {
@@ -225,6 +239,7 @@ impl Default for GrpcServerConfig {
             max_recv_message_size: DEFAULT_MAX_GRPC_RECV_MESSAGE_SIZE.as_bytes() as usize,
             max_send_message_size: DEFAULT_MAX_GRPC_SEND_MESSAGE_SIZE.as_bytes() as usize,
             tls: TlsOption::default(),
+            max_connection_age: None,
         }
     }
 }
@@ -333,6 +348,10 @@ impl Server for GrpcServer {
             builder = builder.tls_config(tls_config).context(StartGrpcSnafu)?;
         }
 
+        if let Some(max_connection_age) = self.config.max_connection_age {
+            builder = builder.max_connection_age(max_connection_age);
+        }
+
         let mut builder = builder
             .add_routes(routes)
             .add_service(self.create_healthcheck_service())

src/servers/src/grpc/builder.rs

@@ -197,6 +197,7 @@ impl GrpcServerBuilder {
             otel_arrow_service: Mutex::new(self.otel_arrow_service),
             bind_addr: None,
             name: self.name,
+            config: self.config,
         }
     }
 }