400 - Dealing with cancel safety in async Rust | RFD | Oxide

Example: A corrected spawn_send_task

The select! loops section had an example spawn_send_task function which was identified as buggy. By using Sender::reserve, this code can be made correct.

This version of the code is a bit more complex. It reads as:

fn spawn_send_task(sender: tokio::sync::mpsc::Sender<String>) {
    let strings: Vec<String> = vec![
        "foo".to_owned(),
        "bar".to_owned(),
        "baz".to_owned(),
    ];
    tokio::task::spawn(async move {
        // interval.tick() completes execution at 0s, 1s, 2s...
        let mut interval = tokio::time::interval(Duration::from_secs(1));
        let mut strings = strings.into_iter();
        // std::vec::IntoIter implements ExactSizeIterator, which has a len()
        // method which tells us if the iterator is non-empty. If your
        // iterator doesn't implement ExactSizeIterator, you can instead
        // use iterator.peekable(), and peek() to see if any values remain.
        while strings.len() > 0 {
            tokio::select! {
                permit = sender.reserve() => {
                    match permit {
                        Ok(permit) => {
                            let value = strings.next().unwrap();
                            // send() is synchronous because the async part
                            // (waiting for a slot to be available) has
                            // already been completed by the time the permit is
                            // acquired.
                            permit.send(value);
                        }
                        Err(_) => {
                            println!("receiver closed");
                            break;
                        }
                    }
                }
                _ = interval.tick() => {
                    println!("interval tick");
                }
            }
        }
    });
}

There is also a way to solve this issue that doesn’t make you lose your place in line. This is covered in Example: spawn_send_task redux.

Reserve and related patterns

mpsc::Sender::reserve is an example of a common pattern, where the general sequence of async operations is:

Some convenience methods like Sender::send combine both steps. These APIs can be easier to use sometimes. But providing these kinds of convenience APIs can lead to users accidentally introducing bugs, as seen in select! loops.

If you’re writing such APIs, consider not providing convenience methods like Sender::send, and instead mentioning cancel safety in the documentation. While it can be a bit surprising at first (a sender that can’t just send values?), not providing such convenience methods makes it easier to train users into writing correct code. Whether this is a good idea or not depends on the situation.

Now let’s look at a case study which shows how to split up complex operations in practice.

1. A StreamExt::next future.

2. A future returned from a method called InstanceSerialConsoleHelper::recv.

3. An instance of SinkExt::send.

4. InstanceSerialConsoleHelper::send, which was a wrapper around SinkExt::send.

• Future 2 was not cancel-safe because it read a value off of a stream, then processed it using further awaits in the middle. If this future was cancelled in the middle, the value could have been read off the stream but not processed, and lost.

• Future 3 was not cancel-safe because SinkExt::send is not cancel-safe, for the same reason that mpsc::sender::Send is not cancel-safe.

• Future 4 was not cancel-safe because it was a wrapper over SinkExt::send.

Note

Making InstanceSerialConsoleHelper::recv cancel-safe

let mut helper: InstanceSerialConsoleHelper = /* ... */;
tokio::select! {
    res = helper.recv() => {
        let message = res?.await?;
        // operate on message
    }
    // ... other select branches
}

Making SinkExt::send cancel-safe

use cancel_safe_futures::prelude::*;
let mut sink = /* ... */;
tokio::select! {
    res = sink.reserve() => {
        let permit = res?;
        res.send(value).await?;
        // operate on message
    }
    // ... other select branches
}

Making InstanceSerialConsoleHelper::send cancel-safe

InstanceSerialConsoleHelper::send was a wrapper around SinkExt::send that didn’t do anything else. To address cancel-safety issues with this method, we removed it and instead made InstanceSerialConsoleHelper implement Sink. Then, the reserve method implemented above can be used for this select branch as well.

• internally in a field, or

• externally via a &mut parameter.

1. fn remaining(&self) → usize: Reports the number of remaining bytes. (There’s also a has_remaining(&self) → bool method.)

2. fn chunk(&self) → &[u8]: Returns the next chunk of bytes.

3. fn advance(&mut self, cnt: usize): Advances the buffer by a particular count of bytes.

async fn write(writer: &mut W, data: &[u8]) -> std::io::Result<()>
where
    W: AsyncWrite + Unpin,
{
    // Cursor<&[u8]> implements the bytes::Buf trait.
    let mut cursor = Cursor::new(data);
    while cursor.has_remaining() {
        tokio::select! {
            res = writer.write_all_buf(&mut cursor) => {
                res?;
            }
            // ... some other branch
        }
    }
    Ok(())
}

• Almost all uses of mutexes are to temporarily violate code invariants within a critical section, restoring them by the end of the operation.

• Mutexes are almost always shared between futures (otherwise why would a mutex be used at all), which means that the "cancellation blast radius" of a future holding a mutex extends to other futures.

• std::sync::Mutex has a notion of lock poisoning in case a panic occurs within the critical section. Tokio’s mutexes don’t, either for cancellations or for panics.

Alternatives to Tokio mutexes

The recommended alternative to a Tokio mutex is a message-passing design, also known as the actor model. Message-passing designs is described in this Tokio tutorial. Rather than having futures access some common shared state, this scheme has a single manager task that has full (non-shared) ownership of this state. This task receives requests in serial order, via an MPSC channel, and sends responses over (typically) oneshot channels.

Another option is to switch to std::sync::Mutex, and not hold locks across await points. Be aware, however, that the worker thread is fully occupied while waiting for the lock. This is not a problem if the lock is uncontended and only held for short periods of time.

If you really must use a Tokio mutex

It isn’t impossible to write correct code that manages Tokio mutexes, just very difficult. Some approaches you could take:

All of these options require careful analysis.

struct MyHandler {
    inner: /* some stream of data */,
}
impl MyHandler {
    // This is cancel-unsafe because message will be dropped if
    // the future is cancelled in the middle of the operation.
    async fn next(&mut self) -> MyMessage {
        let message = self.inner.next().await;
        process_message(&message).await;
        message
    }
}

1. There’s extra ceremony involved with carefully setting up the join handle. This is easy to get wrong. (Is there an abstraction we can write to make this easier?)

2. If cancellation is actually desired in some situations, then the task has to be manually cancelled (either via aborts or by using an explicit cancellation channel).

3. tokio::task::spawn requires a 'static bound, so it can’t borrow data from self or elsewhere. Whether this is okay, an inconvenience, or a deal-breaker, depends on the situation^[11].

1. try_send on bounded channels: A bounded channel’s Sender provides a synchronous try_send method, which returns an error if the channel is full. This method can be used to externalize backpressure to clients, e.g. by returning an HTTP 429 Too Many Requests error if the channel is full.

2. Watch channels: If you only care about the last value sent to a channel, consider using a Tokio watch channel. A watch channel is a single-producer, multi-consumer channel which only stores one value and discards the current value on receiving a new one, eliminating backpressure concerns.

3. Unbounded channels: A final option that is not recommended, but still available, is to use unbounded MPSC channels. Unbounded channels do not have a capacity limit, and their send method is synchronous. Using unbounded channels has a significant downside, namely the lack of backpressure. Among other things, this can result in memory consumption blowing up if the unbounded channel isn’t emptied. How this concern trades off against cancel-safety issues is a case-by-case decision.

• The code that sent reports had to be asynchronous.

• Since this code used a mutex that was held across an await point, that mutex had to be a Tokio mutex rather than std::sync::Mutex.

• Other code that used the mutex also had to be asynchronous.

The solution

In Omicron PR #3950 we switched to using a watch channel. We used the fact that progress reports were cumulative; only the last progress was interesting to the tracker.

This had ripple effects, all positive:

An alternative would have been to rewrite this code into a message-passing style, but this was a more expedient way to achieve cancel safety.

use std::future::Future;
struct TransactionContext {
    // ...
}
struct ConnectionPool {
    // ...
}
impl ConnectionPool {
    fn execute_transaction<F, Fut>(future_fn: F) -> Result<T, E>
    where
        F: FnOnce(TransactionContext) -> Fut,
        Fut: Future<Output = Result<R, E>>,
    {
        let cx: TransactionContext = /* ... */;
        // ... begin transaction
        match future_fn(cx).await {
            Ok(value) => {
                // ... commit transaction
                Ok(value)
            }
            Err(error) => {
                // ... rollback transaction
                Err(error)
            }
        }
    }
}

• Don’t name cancel-unsafe methods next() or recv(). It is natural to use these sorts of methods in select! loops.

• Don’t name methods reserve() or acquire() unless they’re mostly cancel-safe. This pattern-matches against the Tokio methods that are similar.

• Use names that indicate an irrevocable action that shouldn’t be repeated. For example, a method named self.http_post_data(), by referring to an HTTP POST, clearly indicates an action that shouldn’t be recreated in a select! loop.

• Use the type system if possible. Consider a method like fn execute(self), where self isn’t cloneable. It is not going to be possible to use it in a select! loop. Instead, users must create the future returned by the method outside of the select! loop.

  • Note that a method like execute(self) can’t be used in a select! loop, but it can can still be called once and cancelled (e.g. with a timeout). Often this is okay since the entire operation is aborted in that case.