chore(server): Use hyper_util::server::graceful::GracefulShutdown

Author: tottoto

Cargo.toml

@@ -28,3 +28,6 @@ members = [
   "tests/default_stubs",
 ]
 resolver = "2"
+
+[patch.crates-io]
+hyper-util = { git = "https://github.com/hyperium/hyper-util.git", rev = "refs/pull/136/head" }

tonic/Cargo.toml

@@ -37,7 +37,7 @@ server = [
   "dep:async-stream",
   "dep:h2",
   "dep:hyper", "hyper?/server",
-  "dep:hyper-util", "hyper-util?/service", "hyper-util?/server-auto",
+  "dep:hyper-util", "hyper-util?/service", "hyper-util?/server-auto", "hyper-util?/server-graceful",
   "dep:socket2",
   "dep:tokio", "tokio?/macros", "tokio?/net", "tokio?/time",
   "dep:tower", "tower?/util", "tower?/limit",

tonic/src/transport/server/mod.rs

@@ -16,6 +16,7 @@ use crate::service::Routes;
 pub use conn::{Connected, TcpConnectInfo};
 use hyper_util::{
     rt::{TokioExecutor, TokioIo, TokioTimer},
+    server::graceful::GracefulShutdown,
     service::TowerToHyperService,
 };
 #[cfg(feature = "tls")]
@@ -562,10 +563,9 @@ impl<L> Server<L> {
             builder
         };
 
-        let (signal_tx, signal_rx) = tokio::sync::watch::channel(());
-        let signal_tx = Arc::new(signal_tx);
-
-        let graceful = signal.is_some();
+        let graceful = signal
+            .is_some()
+            .then(|| hyper_util::server::graceful::GracefulShutdown::new());
         let mut sig = pin!(Fuse { inner: signal });
         let mut incoming = pin!(incoming);
 
@@ -601,21 +601,13 @@ impl<L> Server<L> {
 
                     let hyper_svc = TowerToHyperService::new(req_svc);
 
-                    serve_connection(io, hyper_svc, server.clone(), graceful.then(|| signal_rx.clone()));
+                    serve_connection(io, hyper_svc, server.clone(), graceful.clone());
                 }
             }
         }
 
-        if graceful {
-            let _ = signal_tx.send(());
-            drop(signal_rx);
-            trace!(
-                "waiting for {} connections to close",
-                signal_tx.receiver_count()
-            );
-
-            // Wait for all connections to close
-            signal_tx.closed().await;
+        if let Some(graceful) = graceful {
+            graceful.shutdown().await;
         }
 
         Ok(())
@@ -628,7 +620,7 @@ fn serve_connection<IO, S>(
     io: ServerIo<IO>,
     hyper_svc: TowerToHyperService<S>,
     builder: ConnectionBuilder,
-    mut watcher: Option<tokio::sync::watch::Receiver<()>>,
+    graceful: Option<GracefulShutdown>,
 ) where
     S: Service<Request<Incoming>, Response = Response<BoxBody>> + Clone + Send + 'static,
     S::Future: Send + 'static,
@@ -638,28 +630,20 @@ fn serve_connection<IO, S>(
 {
     tokio::spawn(async move {
         {
-            let mut sig = pin!(Fuse {
-                inner: watcher.as_mut().map(|w| w.changed()),
-            });
+            let conn = builder.serve_connection(TokioIo::new(io), hyper_svc);
 
-            let mut conn = pin!(builder.serve_connection(TokioIo::new(io), hyper_svc));
+            let result = if let Some(graceful) = graceful {
+                let conn = graceful.watch(conn);
+                conn.await
+            } else {
+                conn.await
+            };
 
-            loop {
-                tokio::select! {
-                    rv = &mut conn => {
-                        if let Err(err) = rv {
-                            debug!("failed serving connection: {:#}", err);
-                        }
-                        break;
-                    },
-                    _ = &mut sig => {
-                        conn.as_mut().graceful_shutdown();
-                    }
-                }
+            if let Err(err) = result {
+                debug!("failed serving connection: {:#}", err);
             }
         }
 
-        drop(watcher);
         trace!("connection closed");
     });
 }