Graceful Shutdown

Goal of This Episode

Assemble the tools from earlier episodes into a complete graceful shutdown flow.

When a server needs to stop, the crudest approach is to just kill it — but then work in progress is cut off mid-flight, possibly leaving corrupted data and unanswered requests. graceful shutdown is the politer way: on receiving a stop request, don’t hard-cut — instead, “tell everyone to wrap up → wait for work in hand to finish; on timeout, request cancellation of the remaining work.”

Break it into three ingredients:

Signal source: how to know “it’s time to stop.”
Broadcasting shutdown: how to tell every worker “we’re wrapping up.”
Waiting for the drain: how to wait for all workers to finish up.

We’ll match each one to a tool we’ve already learned.

Assembling the Three Ingredients

The signal source is tokio::signal::ctrl_c() — it’s a Future; .awaiting it waits until the user presses Ctrl-C (in practice you’d also listen for SIGTERM).
Broadcasting shutdown uses Episode 26’s watch as a shutdown flag: one-to-many, and workers who subscribe late can still read the current state.
Waiting for the drain uses Episode 31’s JoinSet — join_next() until it’s empty.

Each worker internally uses select! to wait for “the next job” and “the shutdown signal” at the same time. If a job arrives first, leave the select! and process it; if shutdown arrives first, stop taking new jobs:

extern crate tokio;

use std::time::Duration;
use tokio::sync::watch;
use tokio::task::JoinSet;
use tokio::time::{sleep, timeout};

async fn wait_next_job(id: u32, next_job: &mut u32) -> u32 {
    // here sleep stands in for "waiting for the next job to come in".
    // this wait may be cancelled by shutdown; actual processing happens outside the select!.
    sleep(Duration::from_millis(500)).await;
    let job = *next_job;
    *next_job += 1;
    println!("worker {} got job {}", id, job);
    job
}

async fn process_job(id: u32, job: u32) {
    // this stands in for "actually processing the job".
    // it's deliberately outside the select!, so shutdown can't cancel it midway.
    sleep(Duration::from_millis(300)).await;
    println!("worker {} finished job {}", id, job);
}

async fn worker(id: u32, mut shutdown: watch::Receiver<bool>) {
    let mut next_job = 0;

    loop {
        let job = tokio::select! {
            // waiting for the next job: this can be cancelled by shutdown
            job = wait_next_job(id, &mut next_job) => job,
            // the wrap-up signal
            _ = shutdown.changed() => {
                println!("worker {} got the shutdown signal, exiting", id);
                break;
            }
        };

        // process outside select!, so shutdown cannot drop it midway
        process_job(id, job).await;
    }
}

async fn drain_workers(workers: &mut JoinSet<()>) {
    while let Some(result) = workers.join_next().await {
        if let Err(error) = result {
            eprintln!("worker ended abnormally: {}", error);
        }
    }
}

#[tokio::main]
async fn main() {
    // the watch flag used to broadcast shutdown
    let (shutdown_tx, shutdown_rx) = watch::channel(false);

    // manage all workers with a JoinSet
    let mut workers = JoinSet::new();
    for id in 0..3 {
        workers.spawn(worker(id, shutdown_rx.clone()));
    }

    // 1. wait for the signal
    tokio::signal::ctrl_c().await.expect("failed to listen for Ctrl-C");
    println!("got Ctrl-C, starting graceful shutdown");

    // 2. broadcast the wrap-up
    shutdown_tx.send(true).expect("no worker is listening");

    // 3. wait for all workers to drain, with a 5-second deadline
    match timeout(Duration::from_secs(5), drain_workers(&mut workers)).await {
        Ok(()) => println!("all workers have ended"),
        Err(_) => {
            println!("drain timed out; requesting cancellation of the remaining workers");
            workers.abort_all();
        }
    }
}

timeout(Duration::from_secs(5), future) sets a five-second deadline for waiting on this future.

It is itself a Future. When the wait completes, .await yields Ok(the inner output); if the wait times out, it yields Err(_). In this example, the inner future is:

drain_workers(&mut workers)

That is, “check each worker’s completion result until the JoinSet is empty.” drain_workers prints any JoinError it receives and keeps waiting for the other workers, so Ok(()) only means the wait completed, not that every worker ended normally. On timeout, we call abort_all() to request cancellation of the remaining Tasks.

The Cancellation Safety Design Point

Here’s a key design choice echoing Episodes 27 and 28: place the select! deliberately. In the worker above, the select! waits on “the next job” and “shutdown”; once a job is actually in hand, we leave the select! and only then call process_job. So when shutdown wins, what gets dropped (cancelled) is the resumable wait for the next job — not work that has already started.

If instead you put the real processing inside a branch that can lose to shutdown, operations that aren’t safely cancellable — like read_exact — could be cut off midway, and the data lost with them. This is the cancellation safety we emphasized earlier, applied concretely to shutdown.

Set a Deadline for Wrapping Up

These five seconds are the waiting period for workers to finish the work in hand on their own. On timeout, we request cancellation of the remaining Tasks; once cancellation is requested, in-flight work can be interrupted even though process_job is outside the select!.

This is not a guaranteed time limit for the whole program to exit. Tokio needs Tasks to yield control so it can handle scheduling and cancellation; a loop or blocking call that never yields cannot be interrupted immediately by timeout or abort_all(). A spawn_blocking Task that has already started cannot be stopped by abort_all() either. Returning from abort_all() only means cancellation has been requested.

A Better-fitting Tool: `CancellationToken`

Using watch as a shutdown flag works, but it feels a bit like “borrowing” a state-broadcast tool to serve as a switch. tokio-util provides a tool designed for “cancellation” from the ground up — CancellationToken, with semantics that fit better. tokio-util isn’t part of Tokio proper, so add the dependency first:

[dependencies]
tokio-util = "0.7"

(As with tokio-stream in Episode 30, the - in the crate name becomes _ in code: use tokio_util::....)

Swapping it in for the watch above:

extern crate tokio;
extern crate tokio_util;

use std::time::Duration;
use tokio::task::JoinSet;
use tokio::time::{sleep, timeout};
use tokio_util::sync::CancellationToken;

async fn wait_next_job(id: u32, next_job: &mut u32) -> u32 {
    sleep(Duration::from_millis(500)).await;
    let job = *next_job;
    *next_job += 1;
    println!("worker {} got job {}", id, job);
    job
}

async fn process_job(id: u32, job: u32) {
    sleep(Duration::from_millis(300)).await;
    println!("worker {} finished job {}", id, job);
}

async fn worker(id: u32, token: CancellationToken) {
    let mut next_job = 0;

    loop {
        let job = tokio::select! {
            job = wait_next_job(id, &mut next_job) => job,
            _ = token.cancelled() => { // wait directly on "being cancelled"
                println!("worker {} got cancelled, exiting", id);
                break;
            }
        };

        process_job(id, job).await;
    }
}

async fn drain_workers(workers: &mut JoinSet<()>) {
    while let Some(result) = workers.join_next().await {
        if let Err(error) = result {
            eprintln!("worker ended abnormally: {}", error);
        }
    }
}

#[tokio::main]
async fn main() {
    let token = CancellationToken::new();

    let mut workers = JoinSet::new();
    for id in 0..3 {
        workers.spawn(worker(id, token.clone())); // each worker gets a clone
    }

    tokio::signal::ctrl_c().await.expect("failed to listen for Ctrl-C");
    token.cancel(); // one command, everyone cancelled

    match timeout(Duration::from_secs(5), drain_workers(&mut workers)).await {
        Ok(()) => println!("all workers have ended"),
        Err(_) => {
            println!("drain timed out; requesting cancellation of the remaining workers");
            workers.abort_all();
        }
    }
}

token.cancelled() is a Future that waits for “being cancelled”; one call to token.cancel() wakes every worker holding a clone. It reads as what it is — cancellation — and fits the need better than borrowing watch as a switch.

Recap

Graceful shutdown: signal the wrap-up and wait for work in hand to finish; on timeout, request cancellation of the remaining work.
Three ingredients: signal source (tokio::signal::ctrl_c()), broadcasting shutdown (a watch flag), waiting for the drain (check each result from join_next() until the JoinSet is empty).
select! is a good fit for waiting on “the next job” and “shutdown” at once; if the real work can’t be safely cancelled, use select! only to obtain the job, then leave the select! to process it, so shutdown can’t drop in-flight work midway (cancellation safety).
Use timeout to set the deadline for workers to wrap up on their own; on timeout, request cancellation with abort_all().
The better-fitting tool is tokio_util’s CancellationToken: token.cancel() gives the order and every token.cancelled() wakes up — semantically a better match than borrowing watch.

Learning Rust from Zero