Skip to content
rusthacks.dev
Concurrency

Graceful Shutdown in Tokio With Signals and Task Tracking

Shut down a Tokio service cleanly: catch SIGTERM and Ctrl+C, stop accepting work, wait for in flight tasks with TaskTracker, then exit within a deadline.

7 min read
Summarize in AI
Flat illustration for the post “Graceful Shutdown in Tokio With Signals and Task Tracking”, composed around async.

A service that is killed mid-request drops connections, loses queued work and sometimes leaves a half-written file. A service that shuts down gracefully stops accepting new work, finishes what it has started, flushes what it has buffered and exits within a deadline. Deployments, autoscalers and Kubernetes all ask your process to stop with SIGTERM and then, after a grace period, kill it with SIGKILL. This post builds the shutdown sequence in four steps: catching the signal, telling every task, waiting for them with a deadline, and exiting with the right code. Each step runs, including a real SIGTERM delivered to the process. Everything compiled on Rust 1.99, edition 2024, with tokio 1 and tokio-util 0.7, on a Unix system.

The sequence you want

  1. Receive SIGTERM or SIGINT (Ctrl-C).
  2. Stop accepting new work: close the listener, stop pulling from queues.
  3. Tell in-flight tasks to finish (a cancellation token), not to abort.
  4. Wait for them, up to a deadline shorter than the platform's kill timeout.
  5. Flush logs and metrics, then exit with a status that says whether the shutdown was clean.

Each step has a Tokio piece. Step 3 and 4 are the ones people skip, and without them the process exits as soon as main returns, taking the tasks with it: the runtime drops all spawned tasks when it shuts down, with no chance to finish.

Step 1: catch the signal

tokio::signal::ctrl_c() waits for SIGINT on any platform. For SIGTERM, which is what container runtimes send, use the Unix-specific API. A helper that completes on either keeps the main loop simple:

rust
use std::time::Duration;
use tokio::signal::unix::{signal, SignalKind};

async fn shutdown_signal() -> &'static str {
    let mut term = signal(SignalKind::terminate()).expect("install SIGTERM handler");
    tokio::select! {
        _ = tokio::signal::ctrl_c() => "SIGINT",
        _ = term.recv() => "SIGTERM",
    }
}

#[tokio::main(flavor = "current_thread")]
async fn main() {
    let pid = std::process::id().to_string();
    tokio::spawn(async move {
        tokio::time::sleep(Duration::from_millis(100)).await;
        std::process::Command::new("kill").args(["-TERM", &pid]).status().unwrap();
    });

    println!("running, waiting for a signal");
    let which = shutdown_signal().await;
    println!("received {which}");
}
Text
running, waiting for a signal
received SIGTERM

The spawned task sends the process a real SIGTERM with the kill command after 100 ms, so the example needs no human at the keyboard. Install handlers before the work starts: a signal that arrives before signal() is called gets the default action, which is to terminate immediately. Calling signal(SignalKind::terminate()) replaces that default for as long as the returned value lives.

Step 2 and 3: stop accepting, tell the rest

A server loop that selects between the shutdown signal and new connections stops accepting when the signal wins. In-flight handlers hold a clone of a cancellation token and check it at points where stopping is safe, as in cancelling tasks:

rust
use std::time::Duration;
use tokio::net::{TcpListener, TcpStream};
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio_util::sync::CancellationToken;
use tokio_util::task::TaskTracker;

async fn handle(mut sock: TcpStream, token: CancellationToken) {
    let mut buf = [0u8; 64];
    loop {
        tokio::select! {
            _ = token.cancelled() => {
                let _ = sock.write_all(b"server shutting down\n").await;
                return;
            }
            n = sock.read(&mut buf) => match n {
                Ok(0) | Err(_) => return,
                Ok(n) => {
                    tokio::time::sleep(Duration::from_millis(30)).await;
                    if sock.write_all(&buf[..n]).await.is_err() {
                        return;
                    }
                }
            },
        }
    }
}

#[tokio::main]
async fn main() -> std::io::Result<()> {
    let listener = TcpListener::bind("127.0.0.1:0").await?;
    let addr = listener.local_addr()?;
    let token = CancellationToken::new();
    let tracker = TaskTracker::new();

    let server = {
        let token = token.clone();
        let tracker = tracker.clone();
        tokio::spawn(async move {
            loop {
                tokio::select! {
                    _ = token.cancelled() => break,
                    accepted = listener.accept() => {
                        if let Ok((sock, _)) = accepted {
                            tracker.spawn(handle(sock, token.clone()));
                        }
                    }
                }
            }
        })
    };

    let mut client = TcpStream::connect(addr).await?;
    client.write_all(b"ping").await?;
    let mut reply = [0u8; 4];
    client.read_exact(&mut reply).await?;
    println!("echo: {}", String::from_utf8_lossy(&reply));

    token.cancel();
    server.await.unwrap();
    tracker.close();
    tokio::time::timeout(Duration::from_secs(2), tracker.wait()).await.expect("handlers finish in time");

    let mut tail = String::new();
    client.read_to_string(&mut tail).await?;
    println!("client was told: {}", tail.trim());
    Ok(())
}
Text
echo: ping
client was told: server shutting down

The listener stopped accepting as soon as the token fired, because the accept loop is one of the select branches. The connection handler then noticed the token between requests, told its peer why it was going away and returned. The final read on the client side proves it: the client got a goodbye, not a reset. A handler in the middle of a request would, in this design, finish that request first, since the token is only checked when it comes back to the select.

Step 4: wait with a deadline

The wait must be bounded. A kill timeout of 30 seconds from Kubernetes means your drain has to finish in less, with margin for logging. A stuck handler (a client that never finishes sending, a database call that hangs) would otherwise hold you until SIGKILL, which is the exact outcome you were avoiding. The pattern is to race the tracker against the deadline and report which won:

rust
use std::time::Duration;
use tokio::time::{sleep, timeout};
use tokio_util::task::TaskTracker;

#[tokio::main(flavor = "current_thread", start_paused = true)]
async fn main() {
    let tracker = TaskTracker::new();
    tracker.spawn(async { sleep(Duration::from_millis(500)).await });
    tracker.spawn(async { sleep(Duration::from_secs(120)).await });
    tracker.close();

    match timeout(Duration::from_secs(5), tracker.wait()).await {
        Ok(()) => println!("drained cleanly"),
        Err(_) => println!("deadline hit with {} tasks still running", tracker.len()),
    }
}
Text
deadline hit with 1 tasks still running

On the deadline path, log what was abandoned and exit non-zero, so that the orchestrator and your dashboards can tell a clean stop from a forced one. The tasks that did not finish are dropped when the runtime is dropped at the end of main.

Step 5: the exit code

Return from main with ExitCode so a forced stop is visible to whatever started the process: 0 for a clean drain, a non-zero code (conventionally 1, or 128 plus the signal number for a signal death) when the deadline forced the issue. Do not call std::process::exit in the middle of the shutdown, because it skips destructors, including the ones that flush buffered writers and logging layers.

rust
use std::process::ExitCode;
use std::time::Duration;
use tokio::time::{sleep, timeout};
use tokio_util::task::TaskTracker;

#[tokio::main(flavor = "current_thread", start_paused = true)]
async fn main() -> ExitCode {
    let tracker = TaskTracker::new();
    tracker.spawn(async { sleep(Duration::from_secs(1)).await });
    tracker.close();

    let clean = timeout(Duration::from_secs(5), tracker.wait()).await.is_ok();
    println!("clean shutdown: {clean}");
    if clean { ExitCode::SUCCESS } else { ExitCode::from(2) }
}
Text
clean shutdown: true

Using a distinct non-zero code for "forced" lets an alert fire on exactly that case. Because main returns normally, every destructor on the way out still runs, including the one that flushes a buffered log writer.

Details that bite in production

  • Readiness first. In Kubernetes, the pod receives SIGTERM at about the same time as it is removed from the service endpoints, and traffic can keep arriving for a second or two. Start failing the readiness probe on SIGTERM and wait a few seconds before closing the listener, so the load balancer stops sending first. The sequence is in graceful shutdown on Kubernetes.
  • Axum and hyper have a built-in hook: axum::serve(listener, app).with_graceful_shutdown(signal) stops accepting and waits for in-flight requests. Use it, and still bound the wait. See axum graceful shutdown.
  • Background workers that pull from a queue should stop pulling when the token fires, then finish the item in hand; items not acknowledged will be redelivered by the queue.
  • A second Ctrl-C should mean "now". Handle the first signal by starting the drain and the second by exiting immediately; a developer who hits Ctrl-C twice is saying the drain is stuck.
  • Blocking tasks (spawn_blocking) cannot be told to stop. Dropping the runtime waits for them; use shutdown_timeout on a runtime you built yourself (building a runtime) to bound it.

Test the shutdown path the way you test any other path. A test that starts the server, opens a connection, sends the signal (or cancels the token) and asserts that the connection received its goodbye and that the process returned within the deadline takes under a hundred lines, and it is the only way to know the drain works before a deployment finds out. The test above, with a paused clock for the timing parts, is a template. Run it in CI so nobody breaks it by adding a task that is not spawned through the tracker.

When this is more than you need

A batch job or CLI that is safe to rerun does not need a drain; die fast and restart. A service whose requests are all idempotent and short, run behind a load balancer that retries, can often get by with a 5-second delay on SIGTERM and nothing else. The graceful path earns its keep when requests are long (uploads, streaming), when work is acknowledged before it is done, or when a buffered writer would lose data. Even then, shutdown is only half of reliability: the other half is making sure a crash at any moment, including a SIGKILL after the deadline, leaves your data in a state you can recover from.

Also inWeb
Share

Keep reading

Get new posts by email

We send one email when a new post is published. No spam.

Unsubscribe in one click. No spam, ever.