diff --git a/tonic/src/transport/server/mod.rs b/tonic/src/transport/server/mod.rs index 9112443ba..4a5dab10e 100644 --- a/tonic/src/transport/server/mod.rs +++ b/tonic/src/transport/server/mod.rs @@ -109,6 +109,7 @@ pub struct Server { accept_http1: bool, service_builder: ServiceBuilder, max_connection_age: Option, + max_connection_age_grace: Option, } impl Default for Server { @@ -135,6 +136,7 @@ impl Default for Server { accept_http1: false, service_builder: Default::default(), max_connection_age: None, + max_connection_age_grace: None, } } } @@ -286,6 +288,36 @@ impl Server { } } + /// Sets the maximum duration that a connection may continue to exist + /// **after** the graceful shutdown period (`max_connection_age`) has elapsed. + /// + /// This timeout only takes effect *after* a connection has exceeded its + /// configured `max_connection_age`. Once that happens, the server will begin + /// graceful shutdown for the connection. If the connection does not close + /// gracefully within the `max_connection_age_grace` duration, the server will then + /// forcefully terminate it. + /// + /// If no `max_connection_age` is configured, this forced shutdown timeout will + /// **never trigger**, because the server will not know when to begin the + /// graceful shutdown phase. + /// + /// Default is no limit (`None`). + /// + /// ``` + /// # use tonic::transport::Server; + /// # use tower_service::Service; + /// # use std::time::Duration; + /// # let builder = Server::builder(); + /// builder.max_connection_age_grace(Duration::from_secs(60)); + /// ``` + #[must_use] + pub fn max_connection_age_grace(self, max_connection_age_grace: Duration) -> Self { + Server { + max_connection_age_grace: Some(max_connection_age_grace), + ..self + } + } + /// Set whether HTTP2 Ping frames are enabled on accepted connections. /// /// If `None` is specified, HTTP2 keepalive is disabled, otherwise the duration @@ -573,6 +605,7 @@ impl Server { max_frame_size: self.max_frame_size, accept_http1: self.accept_http1, max_connection_age: self.max_connection_age, + max_connection_age_grace: self.max_connection_age_grace, } } @@ -701,6 +734,7 @@ impl Server { let http2_max_pending_accept_reset_streams = self.http2_max_pending_accept_reset_streams; let http2_max_local_error_reset_streams = self.http2_max_local_error_reset_streams; let max_connection_age = self.max_connection_age; + let max_connection_age_grace = self.max_connection_age_grace; let svc = self.service_builder.service(svc); @@ -780,7 +814,7 @@ impl Server { let hyper_io = TokioIo::new(io); let hyper_svc = TowerToHyperService::new(req_svc.map_request(|req: Request| req.map(Body::new))); - serve_connection(hyper_io, hyper_svc, server.clone(), graceful.then(|| signal_rx.clone()), max_connection_age); + serve_connection(hyper_io, hyper_svc, server.clone(), graceful.then(|| signal_rx.clone()), max_connection_age, max_connection_age_grace); } } } @@ -809,6 +843,7 @@ fn serve_connection( builder: ConnectionBuilder, mut watcher: Option>, max_connection_age: Option, + max_connection_age_grace: Option, ) where B: http_body::Body + Send + 'static, B::Data: Send, @@ -827,7 +862,8 @@ fn serve_connection( let mut conn = pin!(builder.serve_connection(hyper_io, hyper_svc)); - let mut sleep = pin!(sleep_or_pending(max_connection_age)); + let mut graceful_sleep = pin!(sleep_or_pending(max_connection_age)); + let mut forceful_sleep = pin!(sleep_or_pending(None)); loop { tokio::select! { @@ -837,12 +873,23 @@ fn serve_connection( } break; }, - _ = &mut sleep => { + _ = &mut graceful_sleep => { conn.as_mut().graceful_shutdown(); - sleep.set(sleep_or_pending(None)); + graceful_sleep.set(sleep_or_pending(None)); + forceful_sleep.set(sleep_or_pending( + match (max_connection_age, max_connection_age_grace) { + (None, _) => None, + (Some(_), Some(shutdown_after)) => Some(shutdown_after), + (Some(_), None) => None, + } + )); }, _ = &mut sig => { conn.as_mut().graceful_shutdown(); + }, + _ = &mut forceful_sleep => { + debug!("forcefully closed connection"); + break; } } }