🎓 A1Learner
🔍

Rust Task Tracker

A TaskTracker counts the live tasks spawned through it. Its wait() future completes when the count reaches zero.

The two halves, side by side

The token tells. The tracker waits. Neither half alone is a graceful shutdown:

token.cancel();       // tell every task to stop
tracker.close();      // promise: no new tasks are coming
tracker.wait().await; // wait until every task has stopped

Spawn through the tracker

A tracked task counts itself. close and wait then prove it really ran to its end:

let tracker = TaskTracker::new();
tracker.spawn(async { println!("tracked task ran"); });
tracker.close();
tracker.wait().await;
println!("tracker empty = {}", tracker.is_empty());

Track a future you spawn yourself

track_future wraps any future so the tracker counts it, whichever way it is spawned:

tokio::spawn(tracker.track_future(worker()));

Watch the count, then wait for zero

len() tells you how many tracked tasks are still alive:

tracker.spawn(async { sleep(Duration::from_millis(30)).await });
println!("running = {}", tracker.len());
tracker.close();
tracker.wait().await;
println!("running = {}", tracker.len());

Example

// Cargo.toml: tokio = { version = "1", features = ["full"] }, tokio-util = { version = "0.7", features = ["rt"] }
use tokio::time::{sleep, Duration};
use tokio_util::sync::CancellationToken;
use tokio_util::task::TaskTracker;

async fn worker(name: &'static str, token: CancellationToken) {
    token.cancelled().await;
    println!("{name}: stop signal heard, cleaning up");
}

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

    tracker.spawn(worker("worker", token.clone()));

    sleep(Duration::from_millis(20)).await;
    token.cancel(); // tell every task to stop
    tracker.close(); // and promise: no more tasks are coming
    tracker.wait().await; // wait until every task is really done

    println!("all workers stopped");
}

Key points

  • Use TaskTracker::new() from the tokio-util crate, and tracker.Use spawn(task) to start a tracked task.
  • The tracker is cloneable.
  • So many parts of a program can spawn tasks.
  • They spawn into the same tracker
  • Forgetting close() is the classic hang: without it, wait() waits forever — even while zero tasks are running
  • A finished tracked task removes itself at once.
  • A JoinSet works differently.
  • It holds every task's return value.
  • It holds it until you collect it.
  • So a set you never drain grows.
  • For fire-and-forget work at scale, use the tracker.
  • It is the lighter tool
  • Dropping a TaskTracker does NOT stop its tasks — dropping is not a signal.
  • The token is the signal.
  • The tracker only watches
  • A TaskTracker watches live tasks and waits until none are left.

In this example, one worker waits for the cancellation token and then prints its cleanup line. The worker is spawned through the tracker, so the tracker knows about it. After 20 ms, main fires all three shutdown steps. It fires them in order. cancel() tells the worker to stop. close() promises no more tasks are coming. wait() pauses until the tracker is empty. The cleanup line prints before "all workers stopped". That order is guaranteed. wait() cannot return yet. The worker is still running. Leave out close(), and wait() would hang forever, waiting for tasks that might still come.

Note

TaskTracker needs tokio-util with the rt feature: tokio-util = { version = "0.7", features = ["rt"] }.

Try it Yourself

Edit the code, then press Run.

Running is simulated · Rust Playground

Exercise

Every tracked task has finished, but tracker.wait().await never returns. What did you forget?