Rust Tokio Async skill for Codex

Done well means: every spawned task has an owner that joins or cancels it, shutdown completes within a bounded time, no blocking call runs on a runtime worker, every select! branch is cancel-safe or deliberately pinned, and timeouts exist on every network wait.

Before changing code

• Read Cargo.toml: Tokio version and enabled features (full, or a list such as rt-multi-thread, macros, sync, time, signal). Missing features cause “cannot find macro select” or “no function sleep” errors.
• Check whether tokio-util (CancellationToken, TaskTracker, codecs) is already a dep.
• Find the runtime entry: #[tokio::main], a manual Builder, or a framework (axum, tonic) that owns it. Never create a second runtime inside async code.

Runtime setup

#[tokio::main]                                   // multi-thread, workers = CPU cores
async fn main() -> anyhow::Result<()> { run().await }

#[tokio::main(flavor = "current_thread")]        // CLIs, small tools, lower overhead
async fn main() -> anyhow::Result<()> { run().await }

Use tokio::runtime::Builder::new_multi_thread().worker_threads(n).enable_all().build()? only when you need explicit sizing or to embed the runtime in a sync program.

Spawning and owning tasks

Need	Use
Fixed set of concurrent futures in one function, no spawn	tokio::join! / try_join!
Dynamic set of tasks, collect results, abort the rest on drop	tokio::task::JoinSet
Fire-and-track background tasks, wait for all at shutdown	tokio_util::task::TaskTracker
Bounded concurrency over many items	Semaphore::acquire_owned + JoinSet, or futures::stream::iter(..).buffer_unordered(n)
CPU-heavy or blocking sync call (std::fs, bcrypt, compression, sync DB driver)	tokio::task::spawn_blocking
Long CPU-bound parallel work	rayon with a oneshot back to async

let mut set = tokio::task::JoinSet::new();
for url in urls {
    let client = client.clone();
    set.spawn(async move { fetch(&client, &url).await });
}
while let Some(res) = set.join_next().await {
    match res {
        Ok(Ok(body)) => handle(body),
        Ok(Err(e)) => tracing::warn!(error = ?e, "fetch failed"),
        Err(join_err) if join_err.is_panic() => std::panic::resume_unwind(join_err.into_panic()),
        Err(_) => {} // cancelled
    }
}

A bare tokio::spawn whose JoinHandle is dropped detaches the task: its panic is swallowed and nobody waits for it at shutdown. Keep the handle or use JoinSet/TaskTracker.

Cancellation and select!

Dropping a future cancels it at its current .await. select! drops every losing branch each time it runs, so a branch future that has consumed data internally loses it.

Operation	Cancel-safe in select! loop?
mpsc::Receiver::recv, broadcast::Receiver::recv, watch::Receiver::changed	Yes
TcpListener::accept, CancellationToken::cancelled, sleep, Interval::tick	Yes
AsyncReadExt::read into a buffer you own	Yes (no data lost)
read_exact, read_to_end, read_line, write_all	No
Mutex::lock, Semaphore::acquire	Yes, but you lose your place in the queue
Your own async fn that awaits more than once with state	Assume no

For a byte stream, wrap the reader in a tokio_util::codec::FramedRead so each branch pulls whole frames from a cancel-safe Stream. For a one-off non-cancel-safe future (an upload, a multi-step handshake), create it once outside the loop and poll it by mutable reference so losing a race does not restart it:

use futures::StreamExt;
use tokio_util::codec::{FramedRead, LinesCodec};

let mut lines = FramedRead::new(reader, LinesCodec::new_with_max_length(64 * 1024));
let upload = upload_snapshot(&client);        // started once, not per iteration
tokio::pin!(upload);
let mut upload_done = false;
loop {
    tokio::select! {
        biased;                                // poll shutdown first
        _ = token.cancelled() => break,
        res = &mut upload, if !upload_done => { upload_done = true; res?; }
        line = lines.next() => match line {
            Some(line) => process(line?).await,
            None => break,                     // peer closed
        },
        Some(cmd) = commands.recv() => apply(cmd),
    }
}

Without biased;, branches are polled in random order. Use biased; when priority matters (shutdown first) and accept that a busy early branch can starve later ones.

Graceful shutdown

use tokio_util::{sync::CancellationToken, task::TaskTracker};

let token = CancellationToken::new();
let tracker = TaskTracker::new();
for i in 0..workers {
    let token = token.child_token();
    tracker.spawn(async move { worker(i, token).await });
}
tokio::signal::ctrl_c().await?;            // add SIGTERM via tokio::signal::unix on Linux
token.cancel();
tracker.close();
if tokio::time::timeout(Duration::from_secs(20), tracker.wait()).await.is_err() {
    tracing::warn!("workers did not stop within 20s");
}

Each worker checks token.cancelled() in its select! and finishes the current unit of work before returning. Kubernetes sends SIGTERM and waits terminationGracePeriodSeconds (30 s by default); keep the shutdown timeout below it.

Timeouts and time

• Wrap every external call: tokio::time::timeout(Duration::from_secs(5), call()).await returns Err(Elapsed) on timeout. Client-level timeouts (reqwest, tonic) are better when available.
• tokio::time::interval fires immediately on the first tick and bursts to catch up after a stall; set interval.set_missed_tick_behavior(MissedTickBehavior::Delay) for pollers.
• Never std::thread::sleep in async code; use tokio::time::sleep.

Shared state and channels

Pattern	Choose
Short critical section, no .await while locked	std::sync::Mutex (or parking_lot)
Must hold the lock across .await	tokio::sync::Mutex (slower; rethink first)
Many readers of config that changes rarely	tokio::sync::watch or arc-swap
Work queue with backpressure	bounded mpsc::channel(n)
Request/response to an actor task	mpsc of commands each carrying a oneshot::Sender
Fan-out events to all subscribers	broadcast (handle RecvError::Lagged)

Prefer owning state in one task (actor pattern) over a widely shared Arc<Mutex<_>>. Unbounded channels hide overload until memory runs out; use bounded ones by default.

“future is not Send” errors

tokio::spawn requires Send + 'static. The usual causes:

// Bad: MutexGuard (not Send) lives across .await
let mut g = state.lock().unwrap();
g.count += 1;
notify().await;                     // error: future is not Send

// Good: scope the guard so it drops before the await
{ let mut g = state.lock().unwrap(); g.count += 1; }
notify().await;

Other sources: Rc, RefCell borrows, &mut to non-Sync data, or a Box<dyn Error> without + Send + Sync held across an await. The compiler note “has type X which is not Send … await occurs here” names the value; drop or replace it before that line.

Async fn in traits (stable since 1.75) does not promise Send futures; for traits used with tokio::spawn, declare fn call(&self) -> impl Future<Output = R> + Send or use the trait-variant crate. Async closures (async || {}) are stable since 1.85.

Gotchas

• Blocking inside async (sync file I/O, reqwest::blocking, heavy JSON, password hashing) stalls every task on that worker. Rule of thumb: more than ~100 microseconds without an .await belongs in spawn_blocking.
• spawn_blocking tasks cannot be cancelled; shutdown waits for them.
• block_on inside a runtime panics (“Cannot start a runtime from within a runtime”).
• JoinHandle::abort cancels at the next .await; a task in a tight CPU loop never sees it.
• Holding a tokio::sync::Mutex while calling code that locks the same mutex deadlocks.
• #[tokio::test] is single-threaded by default; race bugs may need #[tokio::test(flavor = "multi_thread", worker_threads = 2)] to show up.

Verification

#[tokio::test(start_paused = true)]        // needs tokio feature "test-util"
async fn retries_after_backoff() {
    let handle = tokio::spawn(retry_loop());
    tokio::time::advance(Duration::from_secs(30)).await;   // no real waiting
    assert!(handle.await.unwrap().is_ok());
}

• Test shutdown: start the service, cancel the token, assert it exits within the timeout.
• Hangs: run with tokio-console (console-subscriber crate, build with RUSTFLAGS="--cfg tokio_unstable") to see tasks stuck on an await or busy-polling.
• cargo clippy flags await_holding_lock and await_holding_refcell_ref; keep them denied.

Report template

**Change:** worker pool moved from detached tokio::spawn to TaskTracker + CancellationToken
**Cancel safety:** read_frame pinned outside select!; recv branches unchanged
**Blocking calls moved:** argon2 hash -> spawn_blocking
**Shutdown:** SIGTERM to exit 1.8s in local test (limit 20s)
**Checks:** cargo test (incl. paused-time retry test), clippy await_holding_lock clean