Gracefully shutting down or restarting an HTTP server allows existing client requests that have not been disconnected to be properly completed, avoiding inconsistencies in business state caused by abrupt program termination during logic execution. Now, let’s take a look at how to correctly handle graceful shutdown when using the Rust axum HTTP framework. Of course, the ideas presented in this article are not limited to axum; any other framework using the tokio runtime can refer to them.
The project dependencies in the Cargo.toml file are as follows:
[dependencies]
axum = "0.8.4"
tokio = { version = "1.47.1", features = ["full"] }First, let’s see what happens when we do not handle graceful shutdown and our request involves time-consuming logic:
use axum::{Router, routing::get};
#[tokio::main]
async fn main() {
let app = Router::new().route("/demo", get(demo));
println!("PID: {}", std::process::id());
let listener = tokio::net::TcpListener::bind("0.0.0.0:3000").await.unwrap();
axum::serve(listener, app)
.await
.unwrap();
println!("Application quit");
}
async fn demo() -> &'static str {
println!("demo request start");
// Simulate a time-consuming task
tokio::time::sleep(Duration::from_secs(15)).await;
println!("demo request done");
}After executing cargo run to start the server, the console output is as follows:
PID: 12762When we execute curl http://localhost:3000/demo -v to request the /demo endpoint, the console output is as follows:
demo request startIf we execute kill 12762 within 15 seconds or press Ctrl+C in the console, the console output is as follows:
[1] 12762 terminated cargo runAt the same time, the client request is abruptly closed, and the console does not print ‘demo request done’, indicating that our demo endpoint handler was not fully executed, and the server exited.
Now, let’s see how it behaves after enabling graceful shutdown. In axum, we can directly use the graceful shutdown feature provided by the framework:
use axum::{Router, routing::get};
#[tokio::main]
async fn main() {
let app = Router::new().route("/demo", get(demo));
println!("PID: {}", std::process::id());
let listener = tokio::net::TcpListener::bind("0.0.0.0:3000").await.unwrap();
axum::serve(listener, app)
// Register graceful shutdown handler here
.with_graceful_shutdown(shutdown_signal())
.await
.unwrap();
println!("Application quit");
}
// ...
async fn shutdown_signal() {
// Listen for Ctrl+C
let ctrl_c = async {
tokio::signal::ctrl_c()
.await
.expect("failed to install Ctrl+C handler");
println!("Ctrl+C received");
};
// Listen for SIGTERM
#[cfg(unix)]
let terminate = async {
tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate())
.expect("failed to install signal handler")
.recv()
.await;
println!("Terminate signal received");
};
// Cross-platform code compatibility, for example, Windows does not support unix signals, adding this is to facilitate unified select! logic
#[cfg(not(unix))]
let terminate = std::future::pending::<()>();
tokio::select! {
_ = ctrl_c => {},
_ = terminate => {},
}
}Now, we re-execute cargo run, then request the /demo endpoint, and if we kill the process within 15 seconds, we find that the server does not exit immediately but outputs the following content after 15 seconds:
Terminate signal received
demo request done
application quitAnd the /demo endpoint is also responded to normally, which meets our expectations. Additionally, if a new request comes in after the program receives the exit signal, the server will directly refuse to process it, which is also in line with our intentions.
Is the above implementation a one-size-fits-all solution? Of course not. If your program has background tasks, child tasks, or if you generate child tasks in your route handler, these child tasks cannot be gracefully handled by axum’s shutdown; the program will not wait for these child tasks to finish before exiting, but will directly interrupt their execution, which can lead to potential logical exceptions or state inconsistencies. For example:
// ...
let app = Router::new()
.route("/demo", get(demo))
// Add /demo2 endpoint
.route("/demo2", get(demo2));
// ...
async fn demo2() -> &'static str {
println!("demo2 request start");
// Generate a child task here
tokio::spawn(async {
println!("demo2 sub task start");
tokio::time::sleep(Duration::from_secs(15)).await;
println!("demo2 sub task done");
});
println!("demo2 request done");
"Hello, World!"
}When we request /demo2, the console output is as follows:
demo2 request start
demo2 request done
demo2 sub task startAnd the /demo2 request will also end, followed by the console output after 15 seconds:
demo2 sub task doneIf we exit the server within 15 seconds, we will find that ‘demo2 sub task done’ is not printed, and the server exits directly. This shows that axum does not handle whether child tasks have finished, and we need to handle this ourselves.
One solution is to manually wait for the child task to finish before returning the response body:
async fn demo2() -> &'static str {
println!("demo2 request start");
tokio::spawn(async {
println!("demo2 sub task start");
tokio::time::sleep(Duration::from_secs(15)).await;
println!("demo2 sub task done");
})
// Add .await
.await;
println!("demo2 request done");
"Hello, World!"
}But this approach certainly does not meet business requirements, as we create child tasks not to wait for them to finish synchronously.
Let’s look at a more reasonable approach (getting to the point): using tokio’s CancellationToken and TaskTracker to ensure that all child tasks finish executing before the application exits.
CancellationToken is used to allow child tasks or code in need to be aware that the program is about to exit, giving them the opportunity to handle shutdown logic, while TaskTracker is used to wait for all child tasks to truly finish executing. To use these two features, you need to install the tokio_util crate and enable the rt feature (TaskTracker requires this):
cargo add tokio_util -F rt # Current version 0.7.16First, let’s look at how to use CancellationToken:
use std::time::Duration;
use axum::{Router, routing::get};
use tokio_util::sync::CancellationToken;
#[tokio::main]
async fn main() {
// Create a new CancellationToken
let cancel_token = CancellationToken::new();
let cancel_token_clone = cancel_token.clone();
let app = Router::new()
// ...
// Pass the clone to the required places
.route("/demo2", get(|| demo2(cancel_token_clone)));
// ...
axum::serve(listener, app)
// Pass the clone to the required places
.with_graceful_shutdown(shutdown_signal(cancel_token.clone()))
.await
.unwrap();
// ...
}
// ...
// Here we modify the previous demo2 example for better understanding
async fn demo2(cancel_token: CancellationToken) -> &'static str {
println!("demo2 request start");
tokio::spawn(async move {
println!("demo2 sub task start");
// Used to mark the child task sequence number
let mut idx = 0;
// Simulate receiving an external event every 5 seconds, then perform related checks
loop {
idx += 1;
tokio::select! {
// Use .cancelled() to achieve awareness of program termination
_ = cancel_token.cancelled() => {
println!("demo2 sub task cancelled, stopped checking");
break;
}
_ = tokio::time::sleep(Duration::from_secs(5)) => {
// Note: The check task needs to be placed in a new child task, otherwise executing this branch will block the loop, causing select! to not listen to .cancelled()
tokio::spawn(async {
println!("demo2 sub task checking {}", idx);
// some heavy check
tokio::time::sleep(Duration::from_secs(15)).await;
println!("demo2 sub task checking done {}", idx);
});
}
}
}
println!("demo2 sub task done");
});
println!("demo2 request done");
"Hello, World!"
}
async fn shutdown_signal(cancel_token: CancellationToken) {
let ctrl_c = async {
tokio::signal::ctrl_c()
.await
.expect("failed to install Ctrl+C handler");
println!("Ctrl+C received");
// After receiving the exit signal, call cancel() to inform other places that the program is about to exit
cancel_token.cancel();
};
#[cfg(unix)]
let terminate = async {
tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate())
.expect("failed to install signal handler")
.recv()
.await;
println!("Terminate signal received");
// After receiving the exit signal, call cancel() to inform other places that the program is about to exit
cancel_token.cancel();
};
// ...
}Now we run the server and request the /demo2 endpoint, the console output is as follows:
demo2 request start
demo2 request done
demo2 sub task start
// Outputs every 5 seconds
demo2 sub task checking 1
// Outputs every 5 seconds
demo2 sub task checking 2
// Outputs every 5 seconds
demo2 sub task checking 3
// Each child task completes after 15 seconds and outputs
demo2 sub task checking done 1
// Outputs every 5 seconds
demo2 sub task checking 4At this point, if we terminate the program at any time, the console will continue to output:
Ctrl+C received
demo2 sub task cancelled, stopped checking
demo2 sub task doneIndicating that the child task successfully received the exit signal.
However, we have discovered a new problem: the checking task executed every 5 seconds does not guarantee that all tasks will finish executing before the program exits. This is where TaskTracker comes into play; it can ensure that all triggered checking tasks finish executing before the program exits:
use tokio_util::{sync::CancellationToken, task::TaskTracker};
#[tokio::main]
async fn main() {
// ...
// Create a new TaskTracker
let task_tracker = TaskTracker::new();
let task_tracker_clone = task_tracker.clone();
let app = Router::new()
// ...
// Pass the clone to the required places
.route("/demo2", get(|| demo2(cancel_token_clone, task_tracker_clone)),
);
// ...
axum::serve(listener, app)
// Pass the clone to the required places
.with_graceful_shutdown(shutdown_signal(cancel_token.clone(), task_tracker.clone()))
.await
.unwrap();
}
async fn demo2(cancel_token: CancellationToken, task_tracker: TaskTracker) -> &'static str {
println!("demo2 request start");
tokio::spawn(async move {
// ...
loop {
tokio::select! {
// ...
_ = tokio::time::sleep(Duration::from_secs(5)) => {
// ...
// Change tokio::spawn to task_tracker::spawn
task_tracker.spawn(async move {
// ...
});
}
}
}
// ...
});
// ...
}
async fn shutdown_signal(cancel_token: CancellationToken, task_tracker: TaskTracker) {
let ctrl_c = async {
tokio::signal::ctrl_c()
.await
.expect("failed to install Ctrl+C handler");
println!("Ctrl+C received");
cancel_token.cancel();
// Call close() first
task_tracker.close();
// At this point, calling wait().await will wait for all child tasks generated by task_tracker to finish before the await is considered complete
task_tracker.wait().await;
};
#[cfg(unix)]
let terminate = async {
tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate())
.expect("failed to install signal handler")
.recv()
.await;
println!("Terminate signal received");
cancel_token.cancel();
// Same as above
task_tracker.close();
task_tracker.wait().await;
};
#[cfg(not(unix))]
let terminate = std::future::pending::<()>();
tokio::select! {
_ = ctrl_c => {},
_ = terminate => {},
}
}Now we run the program again and request the /demo2 endpoint, the console output is as follows:
demo2 request start
demo2 request done
demo2 sub task start
# Executes every 5 seconds
demo2 sub task checking 1
demo2 sub task checking 2
demo2 sub task checking 3
# Each corresponding task completes after 15 seconds
demo2 sub task checking done 1
demo2 sub task checking 4
demo2 sub task checking done 2
demo2 sub task checking 5
demo2 sub task checking done 3
# ...When we terminate the process while some child tasks have not yet completed, we find that the program does not exit immediately but waits for all existing child tasks to finish executing before exiting, achieving the desired outcome.
So, is this the end?
Sort of, but not quite, because there are still two issues:
The first is that during the waiting period for child tasks to finish, the client continues to send requests, and the server still processes requests instead of directly refusing. Why is that?
This is because axum only stops processing requests when the future passed to the graceful_shutdown parameter becomes ready (in JavaScript terms, this is when the promise resolves). Therefore, we can move the wait() outside:
#[tokio::main]
async fn main() {
// ...
let listener = tokio::net::TcpListener::bind("0.0.0.0:3000").await.unwrap();
axum::serve(listener, app)
.with_graceful_shutdown(shutdown_signal(cancel_token.clone(), task_tracker.clone()))
.await
.unwrap();
// Move wait() here, but close() still needs to be executed in shutdown_signal()
task_tracker.wait().await;
println!("application quit");
}At this point, after terminating the program, even during the waiting period for child tasks to complete, new requests will not be accepted.
The second issue is that after calling task_tracker.close(), it is still possible to continue executing task_tracker.spawn(). It is important to emphasize this point to avoid the subconscious assumption that calling spawn after close will not succeed. Therefore, we need to handle the code logic properly to avoid spawning after close. Of course, in our code, this issue does not exist because all spawns occur during the request handling phase, and after the program ends, it will not accept new requests, so it will not call spawn.
Finally, it is important to ensure that the business side correctly handles the graceful shutdown logic and conducts thorough testing and timeout handling to avoid situations where the program does not exit for a long time when attempting to terminate, which may necessitate using SIGKILL (kill -9 PID) to forcefully terminate the program.