A reference to a Future is a Future (sort of)

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.