diff --git a/axum/CHANGELOG.md b/axum/CHANGELOG.md index 471df7e9..7bffad57 100644 --- a/axum/CHANGELOG.md +++ b/axum/CHANGELOG.md @@ -16,6 +16,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 (because it was already never terminating if that method wasn't used) ([#3601]) - **added:** New `ListenerExt::limit_connections` allows limiting concurrent `axum::serve` connections ([#3489]) - **added:** `MethodRouter::method_filter` ([#3586]) +- **added:** `serve::Executor` trait and `Serve::with_executor` for customizing how connection + tasks are spawned, enabling use cases like tracing and telemetry instrumentation ([#3704]) - **changed:** `serve` has an additional generic argument and can now work with any response body type, not just `axum::body::Body` ([#3205]) - **changed:** `Redirect` constructors now accept any `impl Into` ([#3635]) @@ -43,6 +45,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 [#3620]: https://github.com/tokio-rs/axum/pull/3620 [#3656]: https://github.com/tokio-rs/axum/pull/3656 [#3611]: https://github.com/tokio-rs/axum/pull/3611 +[#3704]: https://github.com/tokio-rs/axum/pull/3704 # 0.8.8 diff --git a/axum/src/serve/mod.rs b/axum/src/serve/mod.rs index 05751d2f..3698934c 100644 --- a/axum/src/serve/mod.rs +++ b/axum/src/serve/mod.rs @@ -8,16 +8,17 @@ use std::{ io, marker::PhantomData, pin::pin, + sync::Arc, }; use axum_core::{body::Body, extract::Request, response::Response}; use futures_util::FutureExt; use http_body::Body as HttpBody; use hyper::body::Incoming; -use hyper_util::rt::{TokioExecutor, TokioIo, TokioTimer}; +use hyper_util::rt::{TokioIo, TokioTimer}; #[cfg(any(feature = "http1", feature = "http2"))] use hyper_util::{server::conn::auto::Builder, service::TowerToHyperService}; -use tokio::sync::watch; +use tokio::{sync::watch, task::JoinHandle}; use tower::ServiceExt as _; use tower_service::Service; @@ -27,9 +28,10 @@ pub use self::listener::{ConnLimiter, ConnLimiterIo, Listener, ListenerExt, TapI /// Serve the service with the supplied listener. /// -/// This method of running a service is intentionally simple and doesn't support any configuration. +/// This method of running a service is intentionally simple and doesn't support much configuration. /// hyper's default configuration applies (including [timeouts]); use hyper or hyper-util if you -/// need configuration. +/// need more control. You can supply a custom [`Executor`] via [`Serve::with_executor`] to +/// control how connection tasks are spawned. /// /// It supports both HTTP/1 as well as HTTP/2. /// @@ -98,7 +100,7 @@ pub use self::listener::{ConnLimiter, ConnLimiterIo, Listener, ListenerExt, TapI /// [`HandlerWithoutStateExt::into_make_service_with_connect_info`]: crate::handler::HandlerWithoutStateExt::into_make_service_with_connect_info /// [`HandlerService::into_make_service_with_connect_info`]: crate::handler::HandlerService::into_make_service_with_connect_info #[cfg(all(feature = "tokio", any(feature = "http1", feature = "http2")))] -pub fn serve(listener: L, make_service: M) -> Serve +pub fn serve(listener: L, make_service: M) -> Serve where L: Listener, M: for<'a> Service, Error = Infallible, Response = S>, @@ -111,21 +113,99 @@ where Serve { listener, make_service, + executor: TokioExecutor, _marker: PhantomData, } } +/// A Tokio executor used by [`serve`] to spawn connection tasks, graceful shutdown +/// tasks, and hyper's internal tasks (e.g. HTTP/2 connection management). +/// +/// The default executor is [`TokioExecutor`], which simply calls to +/// [`tokio::spawn`]. A custom implementation can be provided to wrap +/// spawned tasks, e.g. to add tracing or telemetry. +/// +/// Spawned futures rely on Tokio primitives internally, so the executor +/// must run them within a Tokio runtime context (e.g. via [`tokio::spawn`]). +/// +/// # Example +/// +/// An executor that wraps every spawned task in a [`tracing`] span. +/// +/// ``` +/// use std::future::Future; +/// use axum::serve::Executor; +/// use tokio::task::JoinHandle; +/// use tracing::Instrument; +/// +/// #[derive(Clone)] +/// struct InstrumentedExecutor; +/// +/// impl Executor for InstrumentedExecutor { +/// fn execute(&self, fut: Fut) -> JoinHandle +/// where +/// Fut: Future + Send + 'static, +/// Fut::Output: Send + 'static, +/// { +/// let span = tracing::info_span!("axum.serve.task"); +/// tokio::spawn(fut.instrument(span)) +/// } +/// } +/// ``` +/// +/// If your executor is expensive to clone, wrap it in an `Arc`. +/// A blanket implementation is provided for `Arc` where `T: Executor`. +#[cfg(all(feature = "tokio", any(feature = "http1", feature = "http2")))] +pub trait Executor: Clone + Send + Sync + 'static { + /// Execute a task. + fn execute(&self, fut: Fut) -> JoinHandle + where + Fut: Future + Send + 'static, + Fut::Output: Send + 'static; +} + +/// The default executor, which uses [`tokio::spawn`]. +#[cfg(all(feature = "tokio", any(feature = "http1", feature = "http2")))] +#[derive(Clone, Debug)] +pub struct TokioExecutor; + +#[cfg(all(feature = "tokio", any(feature = "http1", feature = "http2")))] +impl Executor for TokioExecutor { + fn execute(&self, fut: Fut) -> JoinHandle + where + Fut: Future + Send + 'static, + Fut::Output: Send + 'static, + { + tokio::spawn(fut) + } +} + +#[cfg(all(feature = "tokio", any(feature = "http1", feature = "http2")))] +impl Executor for Arc +where + T: Executor, +{ + fn execute(&self, fut: Fut) -> JoinHandle + where + Fut: Future + Send + 'static, + Fut::Output: Send + 'static, + { + self.as_ref().execute(fut) + } +} + /// Future returned by [`serve`]. #[cfg(all(feature = "tokio", any(feature = "http1", feature = "http2")))] #[must_use = "futures must be awaited or polled"] -pub struct Serve { +pub struct Serve { listener: L, make_service: M, + executor: E, _marker: PhantomData S>, } #[cfg(all(feature = "tokio", any(feature = "http1", feature = "http2")))] -impl Serve +impl Serve where L: Listener, { @@ -154,13 +234,14 @@ where /// /// Similarly to [`serve`], although this future resolves to `io::Result<()>`, it will never /// error. It returns `Ok(())` only after the `signal` future completes. - pub fn with_graceful_shutdown(self, signal: F) -> WithGracefulShutdown + pub fn with_graceful_shutdown(self, signal: F) -> WithGracefulShutdown where F: Future + Send + 'static, { WithGracefulShutdown { listener: self.listener, make_service: self.make_service, + executor: self.executor, signal, _marker: PhantomData, } @@ -170,10 +251,61 @@ where pub fn local_addr(&self) -> io::Result { self.listener.local_addr() } + + /// Provide a custom [`Executor`] to use for spawning connection tasks and + /// hyper's internal tasks (e.g. HTTP/2). + /// + /// The default is [`TokioExecutor`]. See the [`Executor`] docs for how to + /// implement a custom one. + /// + /// This method can be called before or after [`with_graceful_shutdown`]. + /// + /// # Example + /// + /// ``` + /// use axum::{Router, routing::get, serve::Executor}; + /// # use std::future::Future; + /// # use tokio::task::JoinHandle; + /// # + /// # #[derive(Clone)] + /// # struct MyExecutor; + /// # + /// # impl Executor for MyExecutor { + /// # fn execute(&self, fut: Fut) -> JoinHandle + /// # where + /// # Fut: Future + Send + 'static, + /// # Fut::Output: Send + 'static, + /// # { + /// # tokio::spawn(fut) + /// # } + /// # } + /// # + /// # async { + /// let router = Router::new().route("/", get(|| async { "Hello, World!" })); + /// let listener = tokio::net::TcpListener::bind("0.0.0.0:3000").await.unwrap(); + /// + /// axum::serve(listener, router) + /// .with_executor(MyExecutor) + /// .await; + /// # }; + /// ``` + /// + /// [`with_graceful_shutdown`]: Serve::with_graceful_shutdown + pub fn with_executor(self, executor: E2) -> Serve + where + E2: Executor, + { + Serve { + listener: self.listener, + make_service: self.make_service, + executor, + _marker: PhantomData, + } + } } #[cfg(all(feature = "tokio", any(feature = "http1", feature = "http2")))] -impl Serve +impl Serve where L: Listener, L::Addr: Debug, @@ -184,11 +316,13 @@ where B: HttpBody + Send + 'static, B::Data: Send, B::Error: Into>, + E: Executor, { async fn run(self) -> ! { let Self { mut listener, mut make_service, + executor, _marker, } = self; @@ -197,34 +331,45 @@ where loop { let (io, remote_addr) = listener.accept().await; - handle_connection(&mut make_service, &signal_tx, &close_rx, io, remote_addr).await; + handle_connection( + &mut make_service, + &signal_tx, + &close_rx, + io, + remote_addr, + &executor, + ) + .await; } } } #[cfg(all(feature = "tokio", any(feature = "http1", feature = "http2")))] -impl Debug for Serve +impl Debug for Serve where L: Debug + 'static, M: Debug, + E: Debug, { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { let Self { listener, make_service, + executor, _marker: _, } = self; let mut s = f.debug_struct("Serve"); s.field("listener", listener) - .field("make_service", make_service); + .field("make_service", make_service) + .field("executor", executor); s.finish() } } #[cfg(all(feature = "tokio", any(feature = "http1", feature = "http2")))] -impl IntoFuture for Serve +impl IntoFuture for Serve where L: Listener, L::Addr: Debug, @@ -235,6 +380,7 @@ where B: HttpBody + Send + 'static, B::Data: Send, B::Error: Into>, + E: Executor, { type Output = Infallible; type IntoFuture = private::ServeFuture; @@ -247,15 +393,16 @@ where /// Serve future with graceful shutdown enabled. #[cfg(all(feature = "tokio", any(feature = "http1", feature = "http2")))] #[must_use = "futures must be awaited or polled"] -pub struct WithGracefulShutdown { +pub struct WithGracefulShutdown { listener: L, make_service: M, + executor: E, signal: F, _marker: PhantomData S>, } #[cfg(all(feature = "tokio", any(feature = "http1", feature = "http2")))] -impl WithGracefulShutdown +impl WithGracefulShutdown where L: Listener, { @@ -263,10 +410,27 @@ where pub fn local_addr(&self) -> io::Result { self.listener.local_addr() } + + /// Provide a custom [`Executor`] to use for spawning connection tasks and + /// hyper's internal tasks (e.g. HTTP/2). + /// + /// See [`Serve::with_executor`] for details. + pub fn with_executor(self, executor: E2) -> WithGracefulShutdown + where + E2: Executor, + { + WithGracefulShutdown { + listener: self.listener, + make_service: self.make_service, + executor, + signal: self.signal, + _marker: PhantomData, + } + } } #[cfg(all(feature = "tokio", any(feature = "http1", feature = "http2")))] -impl WithGracefulShutdown +impl WithGracefulShutdown where L: Listener, L::Addr: Debug, @@ -278,17 +442,19 @@ where B: HttpBody + Send + 'static, B::Data: Send, B::Error: Into>, + E: Executor, { async fn run(self) { let Self { mut listener, mut make_service, + executor, signal, _marker, } = self; let (signal_tx, signal_rx) = watch::channel(()); - tokio::spawn(async move { + executor.execute(async move { signal.await; trace!("received graceful shutdown signal. Telling tasks to shutdown"); drop(signal_rx); @@ -305,7 +471,15 @@ where } }; - handle_connection(&mut make_service, &signal_tx, &close_rx, io, remote_addr).await; + handle_connection( + &mut make_service, + &signal_tx, + &close_rx, + io, + remote_addr, + &executor, + ) + .await; } drop(close_rx); @@ -320,17 +494,19 @@ where } #[cfg(all(feature = "tokio", any(feature = "http1", feature = "http2")))] -impl Debug for WithGracefulShutdown +impl Debug for WithGracefulShutdown where L: Debug + 'static, M: Debug, S: Debug, F: Debug, + E: Debug, { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { let Self { listener, make_service, + executor: _, signal, _marker: _, } = self; @@ -344,7 +520,7 @@ where } #[cfg(all(feature = "tokio", any(feature = "http1", feature = "http2")))] -impl IntoFuture for WithGracefulShutdown +impl IntoFuture for WithGracefulShutdown where L: Listener, L::Addr: Debug, @@ -356,6 +532,7 @@ where B: HttpBody + Send + 'static, B::Data: Send, B::Error: Into>, + E: Executor, { type Output = (); type IntoFuture = private::ServeFuture<()>; @@ -365,12 +542,27 @@ where } } -async fn handle_connection( +/// Adapts axum's [`Executor`] to hyper's [`hyper::rt::Executor`]. +#[derive(Clone)] +struct HyperExecutor(E); + +impl hyper::rt::Executor for HyperExecutor +where + E: Executor, + Fut: Future + Send + 'static, +{ + fn execute(&self, fut: Fut) { + drop(self.0.execute(fut)); + } +} + +async fn handle_connection( make_service: &mut M, signal_tx: &watch::Sender<()>, close_rx: &watch::Receiver<()>, io: ::Io, remote_addr: ::Addr, + executor: &E, ) where L: Listener, L::Addr: Debug, @@ -381,6 +573,7 @@ async fn handle_connection( B: HttpBody + Send + 'static, B::Data: Send, B::Error: Into>, + E: Executor, { let io = TokioIo::new(io); @@ -404,9 +597,10 @@ async fn handle_connection( let signal_tx = signal_tx.clone(); let close_rx = close_rx.clone(); - tokio::spawn(async move { + let hyper_executor = HyperExecutor(executor.clone()); + executor.execute(async move { #[allow(unused_mut)] - let mut builder = Builder::new(TokioExecutor::new()); + let mut builder = Builder::new(hyper_executor); // Enable Hyper's default HTTP/1 request header timeout. #[cfg(feature = "http1")] @@ -509,6 +703,7 @@ mod tests { use tokio::{ io::{self, AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt}, net::TcpListener, + task::JoinHandle, }; use tower::ServiceBuilder; @@ -680,10 +875,51 @@ mod tests { UnixListener::bind("").unwrap(), handler.into_make_service_with_connect_info::(), ); + + // with_executor + let router: Router = Router::new(); + let exec = TestExecutor::new(); + serve(TcpListener::bind(addr).await.unwrap(), router.clone()).with_executor(exec.clone()); + serve(TcpListener::bind(addr).await.unwrap(), router.clone()) + .with_executor(exec.clone()) + .with_graceful_shutdown(std::future::pending()); + serve(TcpListener::bind(addr).await.unwrap(), router.clone()) + .with_graceful_shutdown(std::future::pending()) + .with_executor(exec.clone()); + serve(TcpListener::bind(addr).await.unwrap(), get(handler)).with_executor(exec.clone()); + serve( + TcpListener::bind(addr).await.unwrap(), + handler.into_make_service(), + ) + .with_executor(exec); } async fn handler() {} + #[derive(Clone)] + struct TestExecutor(std::sync::Arc); + + impl TestExecutor { + fn new() -> Self { + Self(std::sync::Arc::new(std::sync::atomic::AtomicUsize::new(0))) + } + + fn count(&self) -> usize { + self.0.load(std::sync::atomic::Ordering::SeqCst) + } + } + + impl super::Executor for TestExecutor { + fn execute(&self, fut: Fut) -> JoinHandle + where + Fut: std::future::Future + Send + 'static, + Fut::Output: Send + 'static, + { + self.0.fetch_add(1, std::sync::atomic::Ordering::SeqCst); + tokio::spawn(fut) + } + } + #[crate::test] async fn test_serve_local_addr() { let router: Router = Router::new(); @@ -807,6 +1043,76 @@ mod tests { assert_eq!(body, "Hello, World!"); } + #[crate::test] + async fn serving_with_custom_executor() { + let (client, server) = io::duplex(1024); + let listener = ReadyListener(Some(server)); + + let app = Router::new().route("/", get(|| async { "Hello, World!" })); + + let executor = TestExecutor::new(); + tokio::spawn( + serve(listener, app) + .with_executor(executor.clone()) + .into_future(), + ); + + let stream = TokioIo::new(client); + let (mut sender, conn) = hyper::client::conn::http1::handshake(stream).await.unwrap(); + tokio::spawn(conn); + + let request = Request::builder().body(Body::empty()).unwrap(); + + let response = sender.send_request(request).await.unwrap(); + assert_eq!(response.status(), StatusCode::OK); + + let body = Body::new(response.into_body()); + let body = to_bytes(body, usize::MAX).await.unwrap(); + let body = String::from_utf8(body.to_vec()).unwrap(); + assert_eq!(body, "Hello, World!"); + + // One task per connection for HTTP/1. + assert_eq!(executor.count(), 1); + } + + #[crate::test] + #[cfg(feature = "http2")] + async fn serving_with_custom_executor_http2() { + use hyper_util::rt::TokioExecutor; + + let (client, server) = io::duplex(1024); + let listener = ReadyListener(Some(server)); + + let app = Router::new().route("/", get(|| async { "Hello, World!" })); + + let executor = TestExecutor::new(); + tokio::spawn( + serve(listener, app) + .with_executor(executor.clone()) + .into_future(), + ); + + let io = TokioIo::new(client); + let (mut sender, conn) = hyper::client::conn::http2::Builder::new(TokioExecutor::new()) + .handshake(io) + .await + .unwrap(); + tokio::spawn(conn); + + let request = Request::builder().body(Body::empty()).unwrap(); + + let response = sender.send_request(request).await.unwrap(); + assert_eq!(response.status(), StatusCode::OK); + + let body = Body::new(response.into_body()); + let body = to_bytes(body, usize::MAX).await.unwrap(); + let body = String::from_utf8(body.to_vec()).unwrap(); + assert_eq!(body, "Hello, World!"); + + // Two tasks: axum's connection, and hyper's internal HTTP/2 task. + assert_eq!(executor.count(), 2); + } + #[crate::test] async fn serving_with_custom_body_type() { struct CustomBody;