🎓 A1Learner
🔍

Rust Tokio Graceful Shutdown

Graceful shutdown is stopping a server in a set order. The server must also know when every task is truly done.

One signal, many listeners

let (notify_shutdown, _) = broadcast::channel::<()>(1);
// Each connection task gets: notify_shutdown.subscribe()

The Shutdown wrapper

struct Shutdown {
    shutdown: bool, // have we been told?
    notify: broadcast::Receiver<()>,
}
// recv() waits until the signal arrives — or the Sender drops.

A handler watches while it works

tokio::select! {
    _ = handler.shutdown.recv() => {
        // Stop once the current command reaches a safe point.
        return;
    }
    frame = connection.read_frame() => { /* one round of work */ }
}

The tracker nobody writes to

struct Handler {
    // Never used directly — when the Handler drops,
    // this sender drops, and the server notices.
    _shutdown_complete: mpsc::Sender<()>,
}

The server waits for the channel to close

drop(notify_shutdown);       // signal every task
drop(shutdown_complete_tx);  // drop the server's own sender
// This loop ends only when every task's sender is gone too:
while shutdown_complete_rx.recv().await.is_some() {}

Example

// Cargo.toml: tokio = { version = "1", features = ["full"] }
use tokio::sync::{broadcast, mpsc};
use tokio::time::{sleep, Duration};

#[tokio::main]
async fn main() {
    let (shutdown_tx, mut shutdown_rx) = broadcast::channel::<()>(1);
    let (tracker_tx, mut tracker_rx) = mpsc::channel::<()>(1);
    let (tick_tx, mut tick_rx) = mpsc::channel::<()>(4);

    let worker_tracker = tracker_tx.clone();
    let worker = tokio::spawn(async move {
        let _tracker = worker_tracker; // dropped when the worker ends
        loop {
            tokio::select! {
                _ = shutdown_rx.recv() => {
                    println!("worker: shutdown signal received - finishing up");
                    break;
                }
                _ = sleep(Duration::from_millis(10)) => {
                    println!("worker: tick");
                    let _ = tick_tx.send(()).await;
                }
            }
        }
    });

    // Wait until the worker has ticked twice, then signal shutdown.
    tick_rx.recv().await;
    tick_rx.recv().await;
    let _ = shutdown_tx.send(());

    // The tracker channel closes only when every sender is dropped:
    // ours, and the one the worker holds.
    drop(tracker_tx);
    while tracker_rx.recv().await.is_some() {}

    worker.await.unwrap();
    println!("shutdown complete: every task finished");
}

Key points

  • The shutdown signal is a broadcast channel.
  • The server holds the only Sender.
  • Every connection task subscribes a Receiver
  • Each task wraps its receiver in a Shutdown helper.
  • The helper remembers "already told" with a flag.
  • Its recv() completes when the signal fires or the Sender is dropped
  • Each task also holds a clone of one mpsc Sender it never uses.
  • The field exists only to be dropped.
  • When the task ends, its sender drops with it
  • The server keeps the matching mpsc Receiver.
  • When every task has ended, every sender is gone.
  • The channel closes, and recv() returns None.
  • "all tasks finished" needs no counting
  • The run order has four steps.
  • Wait for the trigger (ctrl_c).
  • Drop the broadcast Sender.
  • Drop your own mpsc Sender.
  • Wait for the mpsc channel to close
  • Add a Semaphore when a server must cap live connections.
  • The listener takes a permit before accepting.
  • Each handler returns it on drop.
  • Mini-Redis caps at 250
  • Dropping the broadcast Sender is the signal — no final message is ever sent.
  • A graceful shutdown stops new work first and exits only when every task has finished.

In this example, one worker task ticks every 10 milliseconds. It reports each tick to main over a small channel. So the sequence never depends on timing luck. The worker owns the two shutdown pieces from the points. One is a broadcast receiver. The other is a tracker sender it never writes to. After exactly two ticks, main sends on the broadcast channel. The worker's select! takes the shutdown branch, prints its finishing line, and ends — dropping its tracker sender as it goes. Main has already dropped its own tracker sender. So the tracker channel closes. The waiting loop ends. Main prints the completion line. No task was counted anywhere: the channel closing was the whole signal.

Note

A task that subscribes late simply finds the channel already closed.

Try it Yourself

Edit the code, then press Run.

Running is simulated · Rust Playground

Exercise

How does the server know that every connection task has finished?