Cancellation Safety in Tokio: Why select! Can Silently Drop Your Work

Here's a bug that will not panic, will not deadlock, and will not show up in your logs. It just quietly corrupts a byte stream, and only under load, and only sometimes. If you have written a connection handler with tokio::select! in a loop, there's a reasonable chance you have already shipped this.

use std::io;
use std::time::Duration;
use tokio::io::AsyncReadExt;
use tokio::net::TcpStream;

async fn handle_connection(mut stream: TcpStream) -> io::Result<()> {
    let mut header = [0u8; 4];
    let mut ticker = tokio::time::interval(Duration::from_secs(5));

    loop {
        tokio::select! {
            result = stream.read_exact(&mut header) => {
                result?;
                dispatch(&header);
            }
            _ = ticker.tick() => {
                send_keepalive(&mut stream).await?;
            }
        }
    }
}

This looks entirely reasonable. Read a fixed-size frame header, and every five seconds send a keepalive so the connection doesn't get reaped by some middlebox. It compiles, it passes every test you'll think to write, and in a quiet test environment it works fine, because the keepalive branch almost never wins a race against a read that's already in flight.

Put it under real network conditions, where a header arrives in two TCP segments a few milliseconds apart, and occasionally the ticker fires in that gap. Now the frame is silently destroyed.

What select! actually does when it "cancels" something

select! polls every branch's future. The first one to return Poll::Ready wins, its result is used, and every other branch's future is dropped. That's the entire mechanism: cancellation in async Rust isn't a signal or an interrupt, it's just Drop running on a future that hadn't finished polling.

For most futures that's harmless. Dropping a Sleep that hasn't fired yet costs you nothing. But a future can hold state that only lives inside its own fields between polls, invisible to you as the caller, and if you throw the future away, that state goes with it.

read_exact is the textbook case. Its future remembers how many of the requested bytes it has already copied into your buffer. Say the header is 4 bytes and 2 have arrived: the future has written those 2 bytes into header[0..2] and is waiting on the socket for the rest. If the ticker branch wins the race at that exact moment, the read_exact future is dropped. Those 2 bytes are still sitting in header, but the fact that there were 2 of them is gone, because that count lived only inside the now-dropped future. Worse: those bytes have already been pulled out of the kernel's socket buffer. They are not still waiting to be read again. On the next loop iteration, you call stream.read_exact(&mut header) again, which happily overwrites header from byte 0 with whatever arrives next, which is actually the tail end of the frame you just half-lost. Your header parsing is now permanently offset by two bytes, for the rest of that connection's life, with no error anywhere to tell you.

This is what Tokio's docs mean by "cancellation safety". It has nothing to do with panics or resource leaks. A method is cancellation-safe if dropping its future partway through a select! leaves no invisible, unrecoverable progress behind. AsyncReadExt::read_exact and AsyncWriteExt::write_all are documented as not cancellation-safe for exactly this reason. mpsc::Receiver::recv, oneshot::Receiver, Mutex::lock, and Interval::tick are documented as safe: dropping their futures either does nothing (no message was consumed) or, in the case of a channel, simply means the value stays in the channel for the next call to pick up.

The fix: stop recreating the future every iteration

The root problem in the example above is that a brand new read_exact future gets constructed on every pass through the outer loop, so any partial progress lives exactly one iteration and then evaporates. The fix is to construct the future once, pin it, and keep polling the same instance across multiple select! calls until it actually completes.

use std::io;
use std::pin::pin;
use std::time::Duration;
use tokio::io::AsyncReadExt;
use tokio::net::TcpStream;

async fn handle_connection(mut stream: TcpStream) -> io::Result<()> {
    let mut ticker = tokio::time::interval(Duration::from_secs(5));

    loop {
        let mut header = [0u8; 4];
        let mut read_fut = pin!(stream.read_exact(&mut header));

        loop {
            tokio::select! {
                result = &mut read_fut => {
                    result?;
                    break;
                }
                _ = ticker.tick() => {
                    send_keepalive(&mut stream).await?;
                }
            }
        }

        dispatch(&header);
    }
}

The change is small but the effect is important: &mut read_fut borrows the same future on every inner loop iteration instead of manufacturing a fresh one. When the ticker branch wins, read_fut is not touched at all, it just sits there, still holding its partial-read state, ready to be polled again next time round the inner loop. Nothing is lost because nothing was dropped. This "outer future survives the losing race" pattern is the standard answer to cancellation-safety problems in Tokio and it's worth recognising on sight, because you'll want it any time a non-cancel-safe operation needs to coexist with a timer or a shutdown signal in the same select.

Or let something else own the buffering

The other honest option is to stop hand-rolling the framing and let a type that already handles partial reads correctly do it for you. tokio_util::codec::FramedRead combined with a Decoder keeps its own internal buffer across polls and is designed to tolerate exactly this kind of interleaving, because the buffer lives on the Framed value, not inside a transient per-call future. If you're parsing a real wire protocol rather than one fixed-size header, this is usually the better investment than hand-managing a pinned future.

A quick mental checklist

Before putting an async call inside a select! branch, especially one sitting inside a loop, it's worth asking two questions: does this type's documentation say anything about cancellation safety, and if I drop this future half-finished, is there state I can't get back? In practice the operations that bite people are the "exact" and "all" family on I/O (read_exact, write_all, copy) and hand-written state machines that accumulate data across multiple .await points inside a single async function without storing that state anywhere durable. Channel receivers, timers, and lock acquisition are, by contrast, deliberately built to be thrown away mid-flight without consequence, which is exactly why they're the branches people reach for without a second thought, and exactly why the other branch in the select is the one that needs scrutiny.

None of this is a flaw in select! itself. It does exactly what it says: poll everything, take the winner, drop the rest. The bug lives in the assumption that "drop" means "nothing happened".