Keyboard shortcuts

Press or to navigate between chapters

Press S or / to search in the book

Press ? to show this help

Press Esc to hide this help

Concept: Graceful Shutdown with watch + select!

Kind: Concept (read, do not code).

The scheduler from the last build runs one tick and returns. To make the coordinator a program, we wrap tick in a loop that polls for queued runs forever. But "forever" is a problem: a service has to be able to stop — on Ctrl-C, on a deploy, on a SIGTERM from the orchestrator. And here the course's opening promise comes due. Under worker failure, no eval is lost — but a careless shutdown is itself a kind of failure. Stop the scheduler the wrong way, mid-run, and you drop jobs that were in flight: evals lost not to a crash but to your own exit path.

This chapter is about stopping correctly: a shutdown that drains the in-flight run to completion and only then exits, rather than cancelling it. The two tools are tokio::sync::watch (to carry the signal) and tokio::select! (to react to it without blocking). But the real lesson is about where you check the signal, because that placement is the entire difference between draining and dropping.

Why dropping a future is the danger

Recall from the async arc what a future is: an inert description of work that only advances when it is polled. The runtime polls it under .await. The flip side, the part that bites here: if you stop polling a future — if you drop it — its work simply stops, wherever it was. No unwind, no error, no chance to finish. A half-done eval future that gets dropped never writes its records. The model call already spent is gone; the job is neither done nor recorded.

Now consider the tempting, wrong way to build a stoppable loop. You have a shutdown signal and a run to process, so you race them:

// WRONG: races the run against shutdown
tokio::select! {
    _ = self.tick() => {}          // process the run
    _ = shutdown.changed() => {}   // ...or bail the instant shutdown fires
}

select! polls both futures and, the moment either is ready, runs that branch and drops the other. So if shutdown fires while tick() is halfway through a run, select! drops the tick() future — cancelling the in-flight run, dropping every job still running inside it. That is exactly the eval-loss we swore off. The run was in progress, and we threw it away to exit a few hundred milliseconds sooner.

TRAP The instinct "race the work against the shutdown signal so I can bail immediately" is precisely backwards for anything you must not lose. select! is a cancellation primitive — the losing branch's future is dropped, its work abandoned. Never put work you need to complete on a select! arm against a shutdown signal. Put the shutdown check between units of work, not around one.

The fix: check between ticks, drain within one

The correct shape flips the relationship. The unit of work — one tick, which processes one whole run — runs to completion outside any select!. The shutdown signal is only consulted at the top of the loop (should I claim another run?) and to cut the idle wait short (if I'm sleeping between polls, wake up now). Shutdown never races the run itself.

Concretely: check the flag at the top; if set, stop claiming and break. Otherwise process a whole run. Then, for the idle interval between polls, select! the sleep against the signal — so a shutdown during the idle wait wakes us immediately, but a shutdown during a run is simply observed at the next top-of-loop check, after the run has drained.

tokio::sync::watch is the right channel for this. It is a single-value broadcast latch: one sender sets a value, any number of receivers can cheaply read the latest with borrow() or await a change with changed(). That "many receivers, one latched value" shape is exactly what a shutdown flag is — the scheduler loop and the HTTP server both hold a receiver and both see the one flip. (An mpsc would be consumed by whoever read it first; a watch lets everyone see the same latched true.)

Here is the whole pattern as a runnable toy. process_one stands in for tick — an indivisible unit of work that, once started, runs to completion. Watch where the signal is checked.

use std::time::Duration;
use tokio::sync::watch;
use tokio::time::sleep;

async fn process_one(n: u32) {
    // Stand-in for "process a whole run": once started, it runs to completion.
    println!("  processing unit {n} (takes 100ms)...");
    sleep(Duration::from_millis(100)).await;
    println!("  unit {n} done");
}

async fn poll_loop(mut shutdown: watch::Receiver<bool>, poll: Duration) {
    let mut next = 0u32;
    loop {
        if *shutdown.borrow() {
            println!("shutdown observed at top of loop — stop claiming");
            break;
        }
        // A whole unit of work is drained before we re-check shutdown.
        process_one(next).await;
        next += 1;

        tokio::select! {
            _ = shutdown.changed() => println!("woken early by shutdown signal"),
            _ = sleep(poll) => println!("poll interval elapsed"),
        }
    }
    println!("loop exited cleanly, drained {next} units");
}

#[tokio::main]
async fn main() {
    let (tx, rx) = watch::channel(false);
    let worker = tokio::spawn(poll_loop(rx, Duration::from_millis(30)));

    // Fire shutdown 180ms in — squarely in the middle of unit 1.
    sleep(Duration::from_millis(180)).await;
    println!("main: sending shutdown");
    tx.send(true).unwrap();

    worker.await.unwrap();
}

Predict the output before running — in particular, when shutdown fires at 180ms (unit 1 is mid-flight, from ~130ms to ~230ms), does unit 1 finish or get dropped?

  processing unit 0 (takes 100ms)...
  unit 0 done
poll interval elapsed
  processing unit 1 (takes 100ms)...
main: sending shutdown
  unit 1 done
woken early by shutdown signal
shutdown observed at top of loop — stop claiming
loop exited cleanly, drained 2 units

Read the last four lines closely, because they are the guarantee:

  • main: sending shutdown prints while unit 1 is still running. The signal is now true.
  • unit 1 done prints anyway. The in-flight unit drained — it was not racing the signal, so it ran to completion. That is the eval that would have been lost under the select!-around-the-work design.
  • woken early by shutdown signal — back at the bottom of the loop, changed() had already fired, so the select! returned immediately instead of sleeping out the 30ms. Shutdown is prompt: we don't dawdle in the idle wait.
  • shutdown observed at top of loop — the next iteration's top-of-loop check sees true and breaks. No new unit is claimed.

So the two behaviours we want fall out of one placement decision: drain (the running unit finishes because it is not on a select! arm) and prompt (the loop wakes from its idle sleep the instant the signal flips). Shutdown stops us claiming new work immediately; it never throws away work already in flight.

borrow() does not mark the value seen The top-of-loop check is *shutdown.borrow(), and it is deliberately not borrow_and_update(). changed() resolves when the value has advanced since the last changed()/borrow_and_update() — a plain borrow() reads the value without consuming that "changed" edge. That is why, in the trace, the select!'s changed() still fires at the bottom of the loop even though we read the flag at the top: reading it did not swallow the notification. Reach for borrow_and_update() and you could read true at the top, consume the edge, and then have changed() sit waiting for a *second* change that never comes.
TRAP changed() takes &mut self — it advances the receiver's internal "last seen" version. So the receiver parameter must be mut. Forget it and the select! arm will not compile:
error[E0596]: cannot borrow `shutdown` as mutable, as it is not declared as mutable
 --> src/main.rs:8:17
  |
8 |             _ = shutdown.changed() => {}
  |                 ^^^^^^^^ cannot borrow as mutable
  |
help: consider changing this to be mutable
  |
4 | async fn poll_loop(mut shutdown: watch::Receiver<bool>, poll: Duration) {
  |                    +++

The fix is the compiler's: mut shutdown: watch::Receiver<bool>. It is a small error, but it is the type system reminding you that awaiting a change is a stateful read, not a peek — which is the same distinction the borrow()-vs-changed() split above turns on.

One-for-one: the toy ↔ the real thing

This maps directly onto Scheduler::run_loop, which you build next:

  • poll_loop(mut shutdown: watch::Receiver<bool>, poll: Duration)Scheduler::run_loop(mut shutdown: watch::Receiver<bool>, poll: Duration) — same signature, same watch channel.
  • process_one(n) — an indivisible unit run to completionself.tick() — which processes one whole run (all its jobs, at buffer_unordered concurrency) before returning. Draining "one unit" means draining an entire run's fan-out.
  • if *shutdown.borrow() { break } at the top ↔ the identical top-of-loop check — stop claiming new runs once shutdown is set.
  • select! { changed() / sleep(poll) } at the bottom ↔ the identical idle-wait race — wake promptly on shutdown, otherwise sleep the poll interval. (The real loop also continues immediately when a tick did process a run, so back-to-back queued runs don't wait out a poll interval.)
  • the watch::Sender in main ↔ the coordinator binary's shutdown_tx, fired from tokio::signal::ctrl_c() inside axum's with_graceful_shutdown — one signal that stops the server accepting requests and tells the scheduler to drain and stop, together.

Hold this against the spine one more time. The scheduler drives dyn WorkerHandle; in Part VIII some of those are RemoteWorkers on other machines. When that day comes, a graceful shutdown here will drain an in-flight run whose jobs are executing across the cluster — and because the drain is a property of where we check the signal, not of where the jobs run, it keeps working unchanged. Shutdown drains; it never drops.

Questions to lock

  1. What happens to a future when it is dropped, and why does that make "race the run against shutdown in select!" lose evals?
  2. Where in the loop is the shutdown flag checked, and how does that placement produce both draining (finish the in-flight run) and promptness (don't sleep out the idle interval)?
  3. Why is tokio::sync::watch the right channel for a shutdown signal shared by the scheduler loop and the HTTP server, rather than an mpsc?
  4. Why is the top-of-loop check borrow() and not borrow_and_update(), and why must the receiver be mut?

Next: we build run_loop and wire up the panoptes-control binary — clap flags, the store, the worker pool, the scheduler, and axum::serve with graceful shutdown, all under one signal. The coordinator becomes a program you can run.