Apr 10, 2024
If T
implements Future
, then &mut T
also implements Future
. You can use this fact to add cancellation safety to a non-cancel-safe future.
All of this comes down to select!
as a tool for running concurrent jobs. From its documentation:
When using select! in a loop to receive messages from multiple sources, you should make sure that the receive call is cancellation safe to avoid losing messages.
[…]
To determine whether your own methods are cancellation safe, look for the location of uses of
.await
. This is because when an asynchronous method is cancelled, that always happens at an.await
. If your function behaves correctly even if it is restarted while waiting at an.await
, then it is cancellation safe.
The docs go into more depth about why this is, but the bottom line is that, when select!
ing against N futures, N - 1 of them will be cancelled (and dropped), abandoning any work they were doing up to the last await point. In practice, it’s hard to make a future that is entirely cancellation safe: it has to do its work atomically, or in ways that can be restarted or unwound after the fact. Cancellation safety, unfortunately, does not compose - a sequence of two cancellation-safe operations is not necessarily a cancellation-safe operation.
What I didn’t know when I last used this, and what the docs unfortunately do not say, is that stored futures can be made cancellation-safe by taking references to them. A reference to a future is in most ways a transparent The operations they perform will update the stored future, but only the reference, and not the underlying stored future, will be cancelled.
Thus:
async fn run_tasks() {
loop {
tokio::select! {
a = task_a() => { completed(a) },
b = task_b() => { completed(b) },
}
}
}
requires that both task_a
and task_b
return cancellation-safe futures, as the loop will start both tasks from the beginning each time it enters the select!
. On the other hand:
async fn run_tasks() {
let fut_a = task_a(); // no .await
let fut_b = task_b(); // ditto
tokio::pin!(fut_a);
tokio::pin!(fut_b);
loop {
tokio::select! {
a = &mut fut_a => { completed(a) },
b = &mut fut_b => { completed(b) },
}
}
}
does not require that either task_a
or task_b
be cancellation-safe, as the loop will resume both future at their respective last await points each time it enters the select!
.
Note that both futures have to remain awaitable, though - if either of them terminates, this will panic on the next select!
. This makes this technique more useful with futures::select_all
, where a completed task can be taken out of the argument list for future selects.