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

Support graceful shutdown on serve (#2398)

Browse Source
This commit is contained in:
David Pedersen
2023-12-29 11:17:58 +00:00
committed by GitHub
parent 56159b0d4e
commit 12e8c6219d
8 changed files with 249 additions and 139 deletions
Download Patch File Download Diff File Expand all files Collapse all files
examples/graceful-shutdown/Cargo.toml
+1 -1
View File
@@ -5,7 +5,7 @@ edition = "2021"
publish = false
[dependencies]
axum = { path = "../../axum" }
axum = { path = "../../axum", features = ["tracing"] }
hyper = { version = "1.0", features = [] }
hyper-util = { version = "0.1", features = ["tokio", "server-auto", "http1"] }
tokio = { version = "1.0", features = ["full"] }
examples/graceful-shutdown/src/main.rs
+10 -103
View File
@@ -10,17 +10,12 @@
use std::time::Duration;
use axum::{extract::Request, routing::get, Router};
use hyper::body::Incoming;
use hyper_util::rt::TokioIo;
use axum::{routing::get, Router};
use tokio::net::TcpListener;
use tokio::signal;
use tokio::sync::watch;
use tokio::time::sleep;
use tower::Service;
use tower_http::timeout::TimeoutLayer;
use tower_http::trace::TraceLayer;
use tracing::debug;
use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt};
#[tokio::main]
@@ -28,10 +23,11 @@ async fn main() {
// Enable tracing.
tracing_subscriber::registry()
.with(
tracing_subscriber::EnvFilter::try_from_default_env()
.unwrap_or_else(|_| "example_graceful_shutdown=debug,tower_http=debug".into()),
tracing_subscriber::EnvFilter::try_from_default_env().unwrap_or_else(|_| {
"example_graceful_shutdown=debug,tower_http=debug,axum=trace".into()
}),
)
.with(tracing_subscriber::fmt::layer())
.with(tracing_subscriber::fmt::layer().without_time())
.init();
// Create a regular axum app.
@@ -48,100 +44,11 @@ async fn main() {
// Create a `TcpListener` using tokio.
let listener = TcpListener::bind("0.0.0.0:3000").await.unwrap();
// Create a watch channel to track tasks that are handling connections and wait for them to
// complete.
let (close_tx, close_rx) = watch::channel(());
// Continuously accept new connections.
loop {
let (socket, remote_addr) = tokio::select! {
// Either accept a new connection...
result = listener.accept() => {
result.unwrap()
}
// ...or wait to receive a shutdown signal and stop the accept loop.
_ = shutdown_signal() => {
debug!("signal received, not accepting new connections");
break;
}
};
debug!("connection {remote_addr} accepted");
// We don't need to call `poll_ready` because `Router` is always ready.
let tower_service = app.clone();
// Clone the watch receiver and move it into the task.
let close_rx = close_rx.clone();
// Spawn a task to handle the connection. That way we can serve multiple connections
// concurrently.
tokio::spawn(async move {
// Hyper has its own `AsyncRead` and `AsyncWrite` traits and doesn't use tokio.
// `TokioIo` converts between them.
let socket = TokioIo::new(socket);
// Hyper also has its own `Service` trait and doesn't use tower. We can use
// `hyper::service::service_fn` to create a hyper `Service` that calls our app through
// `tower::Service::call`.
let hyper_service = hyper::service::service_fn(move |request: Request<Incoming>| {
// We have to clone `tower_service` because hyper's `Service` uses `&self` whereas
// tower's `Service` requires `&mut self`.
//
// We don't need to call `poll_ready` since `Router` is always ready.
tower_service.clone().call(request)
});
// `hyper_util::server::conn::auto::Builder` supports both http1 and http2 but doesn't
// support graceful so we have to use hyper directly and unfortunately pick between
// http1 and http2.
let conn = hyper::server::conn::http1::Builder::new()
.serve_connection(socket, hyper_service)
// `with_upgrades` is required for websockets.
.with_upgrades();
// `graceful_shutdown` requires a pinned connection.
let mut conn = std::pin::pin!(conn);
loop {
tokio::select! {
// Poll the connection. This completes when the client has closed the
// connection, graceful shutdown has completed, or we encounter a TCP error.
result = conn.as_mut() => {
if let Err(err) = result {
debug!("failed to serve connection: {err:#}");
}
break;
}
// Start graceful shutdown when we receive a shutdown signal.
//
// We use a loop to continue polling the connection to allow requests to finish
// after starting graceful shutdown. Our `Router` has `TimeoutLayer` so
// requests will finish after at most 10 seconds.
_ = shutdown_signal() => {
debug!("signal received, starting graceful shutdown");
conn.as_mut().graceful_shutdown();
}
}
}
debug!("connection {remote_addr} closed");
// Drop the watch receiver to signal to `main` that this task is done.
drop(close_rx);
});
}
// We only care about the watch receivers that were moved into the tasks so close the residual
// receiver.
drop(close_rx);
// Close the listener to stop accepting new connections.
drop(listener);
// Wait for all tasks to complete.
debug!("waiting for {} tasks to finish", close_tx.receiver_count());
close_tx.closed().await;
// Run the server with graceful shutdown
axum::serve(listener, app)
.with_graceful_shutdown(shutdown_signal())
.await
.unwrap();
}
async fn shutdown_signal() {