Articles / Rust

923 words

  • rust
  • tokio
  • async
  • concurrency
  • networking

How select! in Tokio Silently Drops Half-Finished Work

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

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 some middlebox doesn't reap the connection.

It compiles, it passes every test you'll think to write, and in a quiet test environment it works fine. That's 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.

Cancelling a future just means dropping it

select! polls every branch's future. The first 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 costs you nothing. But a future can hold state that lives only inside its own fields between polls, invisible to you as the caller. Throw the future away and that state goes with it.

Why read_exact loses bytes

read_exact is the textbook case. Its future remembers how many of the requested bytes it has already copied into your buffer. Here's the sequence:

  1. The header is 4 bytes and 2 have arrived. The future has written them into header[0..2] and is waiting on the socket for the rest.
  2. The ticker branch wins the race at that exact moment, so the read_exact future is dropped.
  3. 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 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 iteration you call stream.read_exact(&mut header) again. It 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.

What "cancellation safety" means in Tokio's docs

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.

  • Not safe: AsyncReadExt::read_exact and AsyncWriteExt::write_all are documented as not cancellation-safe, for exactly this reason.
  • Safe: mpsc::Receiver::recv, oneshot::Receiver, Mutex::lock and Interval::tick. Dropping their futures either does nothing (no message was consumed) or, for a channel, simply means the value stays in the channel for the next call to pick up.

The fix: build the future once and keep polling it

The root problem is that a brand new read_exact future gets constructed on every pass through the loop. 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 matters. &mut read_fut borrows the same future on every inner iteration instead of manufacturing a fresh one.

When the ticker branch wins, read_fut isn't touched at all. It 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. It's worth recognising on sight, because you'll want it any time a non-cancel-safe operation shares a select with a timer or a shutdown signal.

Or let something else own the buffering

The other honest option is to stop hand-rolling the framing. Let a type that already handles partial reads correctly do it for you.

tokio_util::codec::FramedRead with a Decoder keeps its own internal buffer across polls and is designed to tolerate exactly this interleaving. The trick is that 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 a better investment than hand-managing a pinned future.

Two questions before anything goes in a select! branch

Especially when the select! sits inside a loop, ask:

  1. Does this type's documentation say anything about cancellation safety?
  2. 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;
  • hand-written state machines that accumulate data across several .await points inside one async function without storing that state anywhere durable.

Here's the slightly funny part. Channel receivers, timers and lock acquisition are deliberately built to be thrown away mid-flight without consequence. That's exactly why people reach for them 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".

← All articles