diff --git a/axum/src/serve/mod.rs b/axum/src/serve/mod.rs index de1cd6f9..ca6c39af 100644 --- a/axum/src/serve/mod.rs +++ b/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( .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| 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( + io: TokioIo, + conn_service: S, + signal_tx: watch::Sender<()>, + close_rx: watch::Receiver<()>, +) where + I: AsyncRead + AsyncWrite + Unpin + Send + 'static, + S: Service + Clone + Send + 'static, + S::Future: Send, +{ + let hyper_service = TowerToHyperService::new( + conn_service.map_request(|req: Request| 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.