serve: Extract serve_connection out of handle_connection

axum/src/serve/mod.rs

@@ -15,7 +15,10 @@ use hyper::body::Incoming;
 use hyper_util::rt::{TokioExecutor, TokioIo};
 #[cfg(any(feature = "http1", feature = "http2"))]
 use hyper_util::{server::conn::auto::Builder, service::TowerToHyperService};
-use tokio::sync::watch;
+use tokio::{
+    io::{AsyncRead, AsyncWrite},
+    sync::watch,
+};
 use tower::ServiceExt as _;
 use tower_service::Service;
 
@@ -370,46 +373,61 @@ async fn handle_connection<L, M, S>(
         .await
         .unwrap_or_else(|err| match err {});
 
-    let tower_service = make_service
+    let conn_service = make_service
         .call(IncomingStream {
             io: &io,
             remote_addr,
         })
         .await
-        .unwrap_or_else(|err| match err {})
-        .map_request(|req: Request<Incoming>| req.map(Body::new));
+        .unwrap_or_else(|err| match err {});
 
-    let hyper_service = TowerToHyperService::new(tower_service);
-    let signal_tx = signal_tx.clone();
-    let close_rx = close_rx.clone();
+    tokio::spawn(serve_connection(
+        io,
+        conn_service,
+        signal_tx.clone(),
+        close_rx.clone(),
+    ));
+}
 
-    tokio::spawn(async move {
-        #[allow(unused_mut)]
-        let mut builder = Builder::new(TokioExecutor::new());
-        // CONNECT protocol needed for HTTP/2 websockets
-        #[cfg(feature = "http2")]
-        builder.http2().enable_connect_protocol();
+async fn serve_connection<I, S>(
+    io: TokioIo<I>,
+    conn_service: S,
+    signal_tx: watch::Sender<()>,
+    close_rx: watch::Receiver<()>,
+) where
+    I: AsyncRead + AsyncWrite + Unpin + Send + 'static,
+    S: Service<Request, Response = Response, Error = Infallible> + Clone + Send + 'static,
+    S::Future: Send,
+{
+    let hyper_service = TowerToHyperService::new(
+        conn_service.map_request(|req: Request<Incoming>| req.map(Body::new)),
+    );
 
-        let mut conn = pin!(builder.serve_connection_with_upgrades(io, hyper_service));
-        let mut signal_closed = pin!(signal_tx.closed().fuse());
+    #[allow(unused_mut)]
+    let mut builder = Builder::new(TokioExecutor::new());
+    // CONNECT protocol needed for HTTP/2 websockets
+    #[cfg(feature = "http2")]
+    builder.http2().enable_connect_protocol();
 
-        loop {
-            tokio::select! {
-                result = conn.as_mut() => {
-                    if let Err(_err) = result {
-                        trace!("failed to serve connection: {_err:#}");
-                    }
-                    break;
-                }
-                _ = &mut signal_closed => {
-                    trace!("signal received in task, starting graceful shutdown");
-                    conn.as_mut().graceful_shutdown();
+    let mut conn = pin!(builder.serve_connection_with_upgrades(io, hyper_service));
+    let mut signal_closed = pin!(signal_tx.closed().fuse());
+
+    loop {
+        tokio::select! {
+            result = conn.as_mut() => {
+                if let Err(_err) = result {
+                    trace!("failed to serve connection: {_err:#}");
                 }
+                break;
+            }
+            _ = &mut signal_closed => {
+                trace!("signal received in task, starting graceful shutdown");
+                conn.as_mut().graceful_shutdown();
             }
         }
+    }
 
-        drop(close_rx);
-    });
+    drop(close_rx);
 }
 
 /// An incoming stream.