Tech Wiki

TOPICSSERIES

[Rust Zero to Production 25] Build a Stoppable Tokio Worker with Channels, select!, and Cancellation

Putting work on a bounded mpsc channel and listening for shutdown with select! does not, by itself, produce a worker that stops cleanly. The contract also needs to say what happens when the queue is full, whether accepted work drains or gets cancelled, which state wins when stop requests race, and who joins the spawned task.

This example targets Rust 2024 with rustc and Cargo 1.98.1 and Tokio 1.53.1. It uses one check worker to pin down bounded backpressure, the monotonic Running < Drain < Cancel transition, cancellation safety inside select!, and tests with paused time. It performs no network I/O.

1. Bound admission with bounded mpsc

Tokio's bounded mpsc buffers at most the configured number of messages. Once the buffer is full, send().await waits until the receiver makes room. Queue capacity therefore puts a visible limit on waiting work and memory growth when producers keep outrunning the consumer.

Callers that must not wait need try_send instead. The example's try_submit returns TrySendError::Full(Check) unchanged when no slot is available. Its test first waits for an event proving that the worker started the first job, fills the capacity-1 queue with a second job, then checks that a third job returns Full. Nothing depends on guessing when the scheduler will poll the worker.

monitor
    .try_submit(Check::new("buffered", Duration::from_secs(60)))
    .unwrap();
let error = monitor
    .try_submit(Check::new("rejected", Duration::from_secs(60)))
    .unwrap_err();
assert!(matches!(error, TrySendError::Full(check) if check.name() == "rejected"));

1.1. Reject invalid capacity with typed errors

In Tokio 1.53.1, mpsc::channel panics when capacity is 0 or larger than its supported maximum. A library boundary can give callers a better failure mode by validating configuration before channel construction. start_monitor checks both limits first.

if queue_capacity == 0 {
    return Err(StartError::ZeroCapacity);
}
if queue_capacity > Semaphore::MAX_PERMITS {
    return Err(StartError::CapacityTooLarge {
        requested: queue_capacity,
        maximum: Semaphore::MAX_PERMITS,
    });
}

let (jobs, receiver) = mpsc::channel(queue_capacity);

The upper-bound test creates its invalid value with Semaphore::MAX_PERMITS.checked_add(1), avoiding integer overflow in the test itself. The result is StartError::CapacityTooLarge { requested, maximum }; capacity 0 produces StartError::ZeroCapacity.

1.2. Return Result when input is already closed

close_input drops the sender held by the monitor and synchronously closes shared admission. After that, submitter() returns a Result instead of unwrapping or panicking.

pub fn submitter(&self) -> Result<Submitter, SubmitterError> {
    self.jobs
        .as_ref()
        .map(|jobs| Submitter {
            jobs: jobs.clone(),
            admission: Arc::clone(&self.admission),
        })
        .ok_or(SubmitterError::InputClosed)
}

Earlier Submitter clones see the same admission state. Once stop or close_input closes it, a new submit fails even if its clone still owns a sender. Only calls admitted before closure remain in the in-flight counter. An RAII guard decrements that counter on success, send failure, or future cancellation.

2. Separate watch notification from monotonic state

A watch channel retains only the latest value sent. That fits shutdown notification, but it is not an ordered event log. If one caller sends Cancel and another sends Drain before the worker observes either value, the latest watched value could become Drain, effectively undoing cancellation.

The example uses watch to wake the worker and stores precedence separately in an AtomicU8. Values 0, 1, and 2 mean Running, Drain, and Cancel. The only permitted direction is Running < Drain < Cancel.

2.1. Drain can only escalate; Cancel is terminal

A clean-shutdown request uses fetch_max(1). It cannot lower a state that has already reached 2. A cancellation request stores 2. Both requests use a short mutex critical section to close admission first, update state, and then notify through watch. The mutex is never held across an await.

let mut admission = self.admission.state.lock().expect("admission mutex poisoned");
if admission.worker_stopped {
    return Err(RequestError::WorkerStopped);
}
admission.open = false;
self.stop_state.store(2, Ordering::Release);
self.stop
    .send(StopSignal::from_state(&self.stop_state))
    .map_err(|_| RequestError::WorkerStopped)

Cancel therefore wins whether requests arrive as cancel-then-drain or drain-then-cancel before observation. One case is easier to miss: a later cancellation must still become terminal after the worker has observed Drain and started emptying the queue.

2.2. A notification is not the state decision

watch::Receiver::changed waits until the watched value changes. The worker calls borrow_and_update(), but it does not use that value alone to choose a stop mode. It rereads the atomic state. This keeps watch's latest-value notification separate from the policy that stop state cannot move backward.

This is not a recommendation to build every state machine from watch + atomic. It works here because there are three states and their merge rule is max. A protocol that must retain the order of every request needs a queue or another ordered log.

3. Put only cancellation-safe operations in select!

tokio::select! waits on several branches, runs the handler for the first completed branch, and cancels the remaining futures. Cancellation here means dropping the losing futures. select! does not make an arbitrary future safe to drop midway; each branch operation needs its own cancellation contract.

3.1. Rely on the documented guarantees for recv and changed

Tokio 1.53.1 documents mpsc::Receiver::recv as cancel safe. If another branch wins, that recv call is guaranteed not to have received a message. watch::Receiver::changed is also cancel safe: a losing call does not mark the value as seen. Those are the specific guarantees that let this worker await both APIs in one loop.

let next = tokio::select! {
    biased;
    changed = stop.changed() => {
        let _ = stop.borrow_and_update();
        if StopSignal::from_state(stop_state) == StopSignal::Cancel {
            cancel_queue(jobs, admission, events, cancelled).await;
            return StopReason::Cancelled;
        }
        if changed.is_err() {
            watch_open = false;
        }
        continue;
    }
    next = jobs.recv() => next,
};

By contrast, if Sender::send loses a select! race, the message is not sent and the value can be dropped. When message loss is unacceptable, reserve capacity first and send through a Permit. This example never places send inside select!. After submit succeeds, the worker queue owns the job.

3.2. biased gives priority and transfers fairness responsibility

By default, select! pseudo-randomly chooses which branch to check first. That provides some fairness, but it does not promise a particular order. All four current select! sites specify biased;, which polls top to bottom. The stop branch appears before job receipt or the timer.

The ordering is deliberate: shutdown latency takes priority. In biased mode, however, the caller is responsible for fairness. A continuously ready early branch can starve later branches while the loop keeps running. Here, observing shutdown either exits the loop or enters drain mode, so the stop branch does not remain ready forever while hiding normal work. A different loop needs its own starvation analysis.

4. Report drain and cancellation differently

A stop request first closes admission synchronously. Cancel also closes the receiver to wake reservations waiting for capacity, then waits until every active admission has succeeded, failed, or been dropped. Only then does it drain the queue. Consequently, every job whose submit returned Ok appears in either completed or cancelled. StopReason::CleanShutdown and StopReason::Cancelled preserve the drain/cancel distinction.

4.1. A clean drain finishes the current timer and accepted queue

When the worker observes Drain during a job, it closes the receiver to new sends and passes the current Check into drain_queue. That function creates and pins a Sleep for the current check, waits for it to finish, then processes jobs already in the queue. The drain request does not cancel completion of the current job.

let timer = sleep(check.duration);
tokio::pin!(timer);
loop {
    if StopSignal::from_state(stop_state) == StopSignal::Cancel {
        let _ = events.send(WorkerEvent::Cancelled(check.name.clone()));
        cancelled.push(check.name);
        cancel_queue(jobs, admission, events, cancelled).await;
        return StopReason::Cancelled;
    }
    if watch_open {
        tokio::select! {
            biased;
            changed = stop.changed() => {
                let _ = stop.borrow_and_update();
                if StopSignal::from_state(stop_state) == StopSignal::Cancel {
                    let _ = events.send(WorkerEvent::Cancelled(check.name.clone()));
                    cancelled.push(check.name);
                    cancel_queue(jobs, admission, events, cancelled).await;
                    return StopReason::Cancelled;
                }
                if changed.is_err() {
                    watch_open = false;
                }
            }
            () = &mut timer => {
                let _ = events.send(WorkerEvent::Completed(check.name.clone()));
                completed.push(check.name);
                break;
            }
        }
    } else {
        (&mut timer).await;
        let _ = events.send(WorkerEvent::Completed(check.name.clone()));
        completed.push(check.name);
        break;
    }
}

The loop keeps polling the same pinned Sleep. Recreating sleep on every trip through the select! loop could keep moving the deadline. If the watch sender closes during draining, the worker sets watch_open to false and awaits that same timer directly.

4.2. Cancel drops the in-flight sleep and classifies the queue

Leaving the branch scope after cancellation drops the in-flight Sleep future. Dropping it cancels the timer without extra cleanup. That statement is limited to this timer future. It does not establish the same behavior for arbitrary futures or external I/O.

The worker records the current check in both WorkerEvent::Cancelled and the report's cancelled list. cancel_queue closes the receiver, calls wait_until_settled for in-flight admissions, and only then uses try_recv to classify every accepted job. This ordering matters because an outstanding permit obtained before Receiver::close can still send afterward. Completed checks remain in completed. Every terminal return also rereads the atomic state while holding the admission mutex and marks worker_stopped. A request_cancel that succeeds before that lock boundary therefore cannot be reported as CleanShutdown.

5. Make shutdown tests deterministic with paused time and events

Tests that sleep in wall-clock time and assume the worker has probably progressed are slow and flaky. Tokio's test runtime can start with start_paused = true, and time::advance moves timer boundaries without waiting in real time. A paused runtime may also auto-advance to the next pending timer when no other work can run, so these tests do not claim a detailed poll order among independently ready tasks.

5.1. Observe first, then advance only the required duration

The clean-shutdown test starts a paused current_thread runtime. It waits for Started("alpha"), proving that the first job is in flight, requests drain, and advances virtual time by 10 seconds. After checking its completion event, the test advances another 5 seconds for the second timer.

monitor.request_clean_shutdown().unwrap();
advance(Duration::from_secs(10)).await;
assert_eq!(
    monitor.next_event().await,
    Some(WorkerEvent::Completed("alpha".into()))
);
assert_eq!(
    monitor.next_event().await,
    Some(WorkerEvent::Started("beta".into()))
);
advance(Duration::from_secs(5)).await;
assert_eq!(
    monitor.next_event().await,
    Some(WorkerEvent::Completed("beta".into()))
);

The backpressure test also waits for Started before filling the queue. Cancellation tests distinguish one completed job, one in-flight job, and one queued job before sending cancel. Separate cases cover cancel-then-drain, drain-then-cancel, and cancellation after drain observation, fixing both monotonic precedence and escalation in the contract.

5.2. Every shutdown path awaits the JoinHandle

Monitor::join(self) drops its job sender and stop sender, then calls handle.await. cancel(self) sends the cancellation request and continues into join. Clean drain, explicit cancellation, and the all-submitters-dropped path do not treat a detached task as cleanup.

WorkerReport::joined_count is 1 in this example. Tests and program output use it to check that the one worker task was actually joined. A failed JoinHandle await returns JoinWorkerError::TaskFailed.

6. Check the result and its limits

Run these commands from the project directory.

cd examples/article-25-tokio-channels-select-cancellation
cargo fmt --all -- --check
cargo check --all-targets --all-features
cargo clippy --all-targets --all-features -- -D warnings
cargo test --all-features
cargo run --quiet

On Rust 1.98.1, Cargo 1.98.1, and Tokio 1.53.1, all five commands should exit with code 0. The suite contains fourteen integration tests and one unit test. Program output is:

stop=clean
completed=alpha,beta
cancelled=0
joined=1

6.1. Contracts established by the tests

The 15 tests cover capacity boundaries, bounded-queue Full, clean drain, monotonic cancel/drain transitions, the terminal race, rejection after stop, concurrent accepted-send classification, a pre-close outstanding permit, submitter cleanup, and exact binary output. Multi-thread terminal stress is supplementary; the permit and admission-close tests provide deterministic regression coverage. The command sequence also runs formatting, compilation checks, and Clippy.

stop=clean says that this binary invocation took the clean-drain path. It does not claim every run ends cleanly. Cancellation tests separately assert StopReason::Cancelled and the names of cancelled jobs.

6.2. Decisions to make before using this in production

Start by defining the acceptance boundary. In this example, the queue owns a job once submit succeeds. Then decide whether drain must finish every accepted job, whether cancel may drop an in-flight operation, and whether shutdown priority outweighs throughput. Determine the cancellation safety of every API used as a branch.

The example covers one timer-backed check worker. A real database write or socket protocol needs a separate contract for external effects that may have happened before its future was dropped. Do not treat select! itself as a safety guarantee; reason from documented operations such as recv and changed. At the final cleanup boundary, await the worker and inspect its report.

Full source code

The complete runnable source for this article is available in the Chapter 25 project on GitHub.

Sources


One response

  1. […] Previous articleBuild a Stoppable Tokio Worker with Channels, select!, and Cancellation […]

Leave a Reply

Your email address will not be published. Required fields are marked *

Tech Wiki

Built with WordPress · Learn in public.