tokio-rs/axum
1
0
Fork 0
mirror of https://github.com/tokio-rs/axum.git synced 2026-08-07 00:00:15 +02:00
Code Issues Packages Projects Releases Wiki Activity

Add Serve::with_executor and Executor trait for custom task spawning (#3704)

Browse Source 
## 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<Fut> 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<E> adapter bridges to hyper.
This commit is contained in:
Julián Montes de Oca
2026-04-20 14:49:30 +02:00
committed by GitHub
parent 309d1bd953
commit de9f13d809
2 changed files with 332 additions and 23 deletions
Download Patch File Download Diff File Expand all files Collapse all files
axum/CHANGELOG.md
+3
View File
@@ -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<String>` ([#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
axum/src/serve/mod.rs
+329 -23
View File
@@ -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<L, M, S, B>(listener: L, make_service: M) -> Serve<L, M, S, B>
pub fn serve<L, M, S, B>(listener: L, make_service: M) -> Serve<L, M, S, B, TokioExecutor>
where
L: Listener,
M: for<'a> Service<IncomingStream<'a, L>, 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<Fut>(&self, fut: Fut) -> JoinHandle<Fut::Output>
/// 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<T>` where `T: Executor`.
#[cfg(all(feature = "tokio", any(feature = "http1", feature = "http2")))]
pub trait Executor: Clone + Send + Sync + 'static {
/// Execute a task.
fn execute<Fut>(&self, fut: Fut) -> JoinHandle<Fut::Output>
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<Fut>(&self, fut: Fut) -> JoinHandle<Fut::Output>
where
Fut: Future + Send + 'static,
Fut::Output: Send + 'static,
{
tokio::spawn(fut)
}
}
#[cfg(all(feature = "tokio", any(feature = "http1", feature = "http2")))]
impl<T> Executor for Arc<T>
where
T: Executor,
{
fn execute<Fut>(&self, fut: Fut) -> JoinHandle<Fut::Output>
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<L, M, S, B> {
pub struct Serve<L, M, S, B, E = TokioExecutor> {
listener: L,
make_service: M,
executor: E,
_marker: PhantomData<fn(B) -> S>,
}
#[cfg(all(feature = "tokio", any(feature = "http1", feature = "http2")))]
impl<L, M, S, B> Serve<L, M, S, B>
impl<L, M, S, B, E> Serve<L, M, S, B, E>
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<F>(self, signal: F) -> WithGracefulShutdown<L, M, S, F, B>
pub fn with_graceful_shutdown<F>(self, signal: F) -> WithGracefulShutdown<L, M, S, F, B, E>
where
F: Future<Output = ()> + 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<L::Addr> {
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<Fut>(&self, fut: Fut) -> JoinHandle<Fut::Output>
/// # 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<E2>(self, executor: E2) -> Serve<L, M, S, B, E2>
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<L, M, S, B> Serve<L, M, S, B>
impl<L, M, S, B, E> Serve<L, M, S, B, E>
where
L: Listener,
L::Addr: Debug,
@@ -184,11 +316,13 @@ where
B: HttpBody + Send + 'static,
B::Data: Send,
B::Error: Into<Box<dyn StdError + Send + Sync>>,
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<L, M, S, B> Debug for Serve<L, M, S, B>
impl<L, M, S, B, E> Debug for Serve<L, M, S, B, E>
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<L, M, S, B> IntoFuture for Serve<L, M, S, B>
impl<L, M, S, B, E> IntoFuture for Serve<L, M, S, B, E>
where
L: Listener,
L::Addr: Debug,
@@ -235,6 +380,7 @@ where
B: HttpBody + Send + 'static,
B::Data: Send,
B::Error: Into<Box<dyn StdError + Send + Sync>>,
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<L, M, S, F, B> {
pub struct WithGracefulShutdown<L, M, S, F, B, E = TokioExecutor> {
listener: L,
make_service: M,
executor: E,
signal: F,
_marker: PhantomData<fn(B) -> S>,
}
#[cfg(all(feature = "tokio", any(feature = "http1", feature = "http2")))]
impl<L, M, S, F, B> WithGracefulShutdown<L, M, S, F, B>
impl<L, M, S, F, B, E> WithGracefulShutdown<L, M, S, F, B, E>
where
L: Listener,
{
@@ -263,10 +410,27 @@ where
pub fn local_addr(&self) -> io::Result<L::Addr> {
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<E2>(self, executor: E2) -> WithGracefulShutdown<L, M, S, F, B, E2>
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<L, M, S, F, B> WithGracefulShutdown<L, M, S, F, B>
impl<L, M, S, F, B, E> WithGracefulShutdown<L, M, S, F, B, E>
where
L: Listener,
L::Addr: Debug,
@@ -278,17 +442,19 @@ where
B: HttpBody + Send + 'static,
B::Data: Send,
B::Error: Into<Box<dyn StdError + Send + Sync>>,
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<L, M, S, F, B> Debug for WithGracefulShutdown<L, M, S, F, B>
impl<L, M, S, F, B, E> Debug for WithGracefulShutdown<L, M, S, F, B, E>
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<L, M, S, F, B> IntoFuture for WithGracefulShutdown<L, M, S, F, B>
impl<L, M, S, F, B, E> IntoFuture for WithGracefulShutdown<L, M, S, F, B, E>
where
L: Listener,
L::Addr: Debug,
@@ -356,6 +532,7 @@ where
B: HttpBody + Send + 'static,
B::Data: Send,
B::Error: Into<Box<dyn StdError + Send + Sync>>,
E: Executor,
{
type Output = ();
type IntoFuture = private::ServeFuture<()>;
@@ -365,12 +542,27 @@ where
}
}
async fn handle_connection<L, M, S, B>(
/// Adapts axum's [`Executor`] to hyper's [`hyper::rt::Executor`].
#[derive(Clone)]
struct HyperExecutor<E>(E);
impl<E, Fut> hyper::rt::Executor<Fut> for HyperExecutor<E>
where
E: Executor,
Fut: Future<Output = ()> + Send + 'static,
{
fn execute(&self, fut: Fut) {
drop(self.0.execute(fut));
}
}
async fn handle_connection<L, M, S, B, E>(
make_service: &mut M,
signal_tx: &watch::Sender<()>,
close_rx: &watch::Receiver<()>,
io: <L as Listener>::Io,
remote_addr: <L as Listener>::Addr,
executor: &E,
) where
L: Listener,
L::Addr: Debug,
@@ -381,6 +573,7 @@ async fn handle_connection<L, M, S, B>(
B: HttpBody + Send + 'static,
B::Data: Send,
B::Error: Into<Box<dyn StdError + Send + Sync>>,
E: Executor,
{
let io = TokioIo::new(io);
@@ -404,9 +597,10 @@ async fn handle_connection<L, M, S, B>(
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::<UdsConnectInfo>(),
);
// 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<std::sync::atomic::AtomicUsize>);
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<Fut>(&self, fut: Fut) -> JoinHandle<Fut::Output>
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;