tonic: add max connection age grace

tonic/src/transport/server/mod.rs

@@ -109,6 +109,7 @@ pub struct Server<L = Identity> {
     accept_http1: bool,
     service_builder: ServiceBuilder<L>,
     max_connection_age: Option<Duration>,
+    max_connection_age_grace: Option<Duration>,
 }
 
 impl Default for Server<Identity> {
@@ -135,6 +136,7 @@ impl Default for Server<Identity> {
             accept_http1: false,
             service_builder: Default::default(),
             max_connection_age: None,
+            max_connection_age_grace: None,
         }
     }
 }
@@ -286,6 +288,36 @@ impl<L> Server<L> {
         }
     }
 
+    /// 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<L> Server<L> {
             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<L> Server<L> {
         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<L> Server<L> {
                     let hyper_io = TokioIo::new(io);
                     let hyper_svc = TowerToHyperService::new(req_svc.map_request(|req: Request<Incoming>| 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<B, IO, S, E>(
     builder: ConnectionBuilder<E>,
     mut watcher: Option<tokio::sync::watch::Receiver<()>>,
     max_connection_age: Option<Duration>,
+    max_connection_age_grace: Option<Duration>,
 ) where
     B: http_body::Body + Send + 'static,
     B::Data: Send,
@@ -827,7 +862,8 @@ fn serve_connection<B, IO, S, E>(
 
             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<B, IO, S, E>(
                         }
                         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;
                     }
                 }
             }