From de9f13d809e86f8bf24292aac0af16d97ef91195 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Juli=C3=A1n=20Montes=20de=20Oca?= Date: Mon, 20 Apr 2026 14:49:30 +0200 Subject: [PATCH] Add Serve::with_executor and Executor trait for custom task spawning (#3704) ## Motivation axum::serve hardcodes TokioExecutor for spawning connection tasks and hyper's internal HTTP/2 tasks. This is the one thing that cannot be customized by wrapping axum's API from the outside, users are forced to reimplement the entire serve loop (~150 lines) just to swap the executor. This is needed for use cases like runtime telemetry (e.g. dial9-tokio-telemetry) where we want to wrap task spawning, so that we can capture things like wake events. The goal of this feature is to allow hooking into the current tokio spawns (by being able to set an executor that provides them, rather than hardcoding them) to attach instrumentation. The goal is not the make serve fully runtime agnostic. ## Solution This PR adds Serve::with_executor() and WithGracefulShutdown::with_executor() builder methods that take an Executor, and an axum::serve::Executor trait for defining the Executor interface. The default remains a new TokioExecutor (that spawns tokio spawns just like before), so we maintain backward compatibility. Small design note: this defines a new axum::serve::Executor trait rather than reuse of hyper::rt::Executor directly, because the latter is generic at the trait level and I couldn't find a way to bound E to cover all the internal future types hyper needs. So this new trait uses a generic method instead, and then a HyperExecutor adapter bridges to hyper. --- axum/CHANGELOG.md | 3 + axum/src/serve/mod.rs | 352 +++++++++++++++++++++++++++++++++++++++--- 2 files changed, 332 insertions(+), 23 deletions(-) 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;