Avoid shared shutdown waiter contention (#3867)

Signed-off-by: jthomson04 <[email protected]>

axum/CHANGELOG.md

@@ -23,6 +23,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - **added:** Add `RawPathParams::from_request_extensions` ([#3757])
 - **changed:** `serve` has an additional generic argument and can now work with any response body
   type, not just `axum::body::Body` ([#3205])
+- **changed:** Reduced contention in `axum::serve` shutdown notification with many
+  active connections ([#3867])
 - **changed:** `Redirect` constructors now accept any `impl Into<String>` ([#3635])
 - **changed:** Updated `matchit` allowing for routes with captures and static prefixes and suffixes ([#3702])
 - **fixed:** Responses to `HEAD` will not accidentally reply with `content-length: 0` anymore ([#3742])
@@ -44,6 +46,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 [#3836]: https://github.com/tokio-rs/axum/pull/3836
 [#3757]: https://github.com/tokio-rs/axum/pull/3757
 [#3801]: https://github.com/tokio-rs/axum/pull/3801
+[#3867]: https://github.com/tokio-rs/axum/pull/3867
 
 # 0.8.9
 

axum/src/serve/mod.rs

@@ -329,14 +329,14 @@ where
             _marker,
         } = self;
 
-        let (signal_tx, _signal_rx) = watch::channel(());
+        let (_signal_tx, signal_rx) = watch::channel(());
         let (_close_tx, close_rx) = watch::channel(());
 
         loop {
             let (io, remote_addr) = listener.accept().await;
             handle_connection(
                 &mut make_service,
-                &signal_tx,
+                &signal_rx,
                 &close_rx,
                 io,
                 remote_addr,
@@ -456,28 +456,31 @@ where
             _marker,
         } = self;
 
-        let (signal_tx, signal_rx) = watch::channel(());
+        let (signal_tx, mut signal_rx) = watch::channel(());
         executor.execute(async move {
             signal.await;
             trace!("received graceful shutdown signal. Telling tasks to shutdown");
-            drop(signal_rx);
+            drop(signal_tx);
         });
 
         let (close_tx, close_rx) = watch::channel(());
 
         loop {
             let (io, remote_addr) =
-                match select(pin!(listener.accept()), pin!(signal_tx.closed())).await {
+                match select(pin!(listener.accept()), pin!(signal_rx.changed())).await {
                     Either::Left((conn, _)) => conn,
-                    Either::Right(_) => {
+                    Either::Right((Err(_), _)) => {
                         trace!("signal received, not accepting new connections");
                         break;
                     }
+                    Either::Right((Ok(()), _)) => {
+                        unreachable!("shutdown channel never sends values")
+                    }
                 };
 
             handle_connection(
                 &mut make_service,
-                &signal_tx,
+                &signal_rx,
                 &close_rx,
                 io,
                 remote_addr,
@@ -562,7 +565,7 @@ where
 
 async fn handle_connection<L, M, S, B, E>(
     make_service: &mut M,
-    signal_tx: &watch::Sender<()>,
+    signal_rx: &watch::Receiver<()>,
     close_rx: &watch::Receiver<()>,
     io: <L as Listener>::Io,
     remote_addr: <L as Listener>::Addr,
@@ -579,6 +582,7 @@ async fn handle_connection<L, M, S, B, E>(
     B::Error: Into<Box<dyn StdError + Send + Sync>>,
     E: Executor,
 {
+    let mut signal_rx = signal_rx.clone();
     let io = TokioIo::new(io);
 
     trace!("connection {remote_addr:?} accepted");
@@ -598,7 +602,6 @@ async fn handle_connection<L, M, S, B, E>(
         .map_request(|req: Request<Incoming>| req.map(Body::new));
 
     let hyper_service = TowerToHyperService::new(tower_service);
-    let signal_tx = signal_tx.clone();
     let close_rx = close_rx.clone();
 
     let hyper_executor = HyperExecutor(executor.clone());
@@ -615,7 +618,7 @@ async fn handle_connection<L, M, S, B, E>(
         builder.http2().enable_connect_protocol();
 
         let mut conn = pin!(builder.serve_connection_with_upgrades(io, hyper_service));
-        let mut signal_closed = pin!(signal_tx.closed().fuse());
+        let mut signal_closed = pin!(signal_rx.changed().fuse());
 
         loop {
             match select(conn.as_mut(), &mut signal_closed).await {
@@ -625,10 +628,13 @@ async fn handle_connection<L, M, S, B, E>(
                     }
                     break;
                 }
-                Either::Right(_) => {
+                Either::Right((Err(_), _)) => {
                     trace!("signal received in task, starting graceful shutdown");
                     conn.as_mut().graceful_shutdown();
                 }
+                Either::Right((Ok(()), _)) => {
+                    unreachable!("shutdown channel never sends values")
+                }
             }
         }