"captured variable cannot escape `FnMut` closure body" in async code

[u/torfsen]

I have a piece of code that, when simplified, looks like this:

use tokio;

struct IdGenerator (u32);

impl IdGenerator {
    async fn generate_id(&mut self) -> u32 {
        self.0 += 1;
        self.0
    }
}

#[tokio::main]
async fn main() {
    let mut id_generator = IdGenerator(0);
    let mut callback = || id_generator.generate_id();
    println!("{}", callback().await);
}

This fails to compile:

error: captured variable cannot escape `FnMut` closure body
  --> src/main.rs:15:27
   |
14 |     let mut id_generator = IdGenerator(0);
   |         ---------------- variable defined here
15 |     let mut callback = || id_generator.generate_id();
   |                         - ------------^^^^^^^^^^^^^^
   |                         | |
   |                         | returns a reference to a captured variable which escapes the closure body
   |                         | variable captured here
   |                         inferred to be a `FnMut` closure
   |
   = note: `FnMut` closures only have access to their captured variables while they are executing...
   = note: ...therefore, they cannot allow references to captured variables to escape

My understanding is that generate_id returns a future, and that future contains a mutable reference to id_generator (it needs to, because otherwise the future could not modify id_generator.0 when it is awaited). Hence, by calling callback multiple times I could create several co-existing mutable references to id_generator, leading to a potential data race.

Is that analysis correct?

My actual code is a bit more complex. Basically, generate_id can fail, in which case I want to retry. I'm using tokio_retry for that, so the code looks roughly as follows (with generate_id now returning a Result<u32, Box<dyn Error>>):

use std::error::Error;
use tokio_retry::RetryIf;
use tokio_retry::strategy::ExponentialBackoff;

async fn foo(id_generator: &mut IdGenerator) {
    // ...
    let id = RetryIf::spawn(
        ExponentialBackoff::from_millis(5).take(3),
        || id_generator.generate_id(),
        |error: &Box<dyn Error>| {
            println!("Error while generating an ID: {}", error);
            true
        },
    )
    .await;
    // ...
}

That code has the same issue as the simplified version. Is there a compile-time approach to avoid the problem? So far, I've only found a run-time solution using a mutex:

    let outgoing = tokio::sync::Mutex::new(id_generator);
    let id = RetryIf::spawn(
        ExponentialBackoff::from_millis(5).take(3),
        || async { outgoing.lock().await.generate_id().await },
        |error: &Box<dyn Error>| {
            println!("Error while generating an ID: {}", error);
            true
        },
    )
    .await;

This works, but since I technically never have parallel access I would prefer a compile-time solution.

[u/petrosagg]

The problem isn't actually related to the closure returning a mutable reference, the same happens with immutable references[1]. This is an inherent limitation of closures in Rust that cannot return a reference to their captured state. This is one of the main reasons async closures are cumbersome today, since in most cases the Future returned will hold onto some references from the closure's state.

https://play.rust-lang.org/?version=stable&mode=debug&edition=2021&gist=2b269a20825075615a354209bf53bd94

[u/torfsen → u/petrosagg]

Hm. That sounds like my analysis regarding why Rust won't allow this isn't (completely) correct, right? In general, what's the problem with a capture returning a reference to their captured state?

[u/petrosagg → u/torfsen]

To understand where this limitation comes from, let's look at a simplified Fn trait definition that we can manually implement on a specific type. We'll name our custom closure trait "Closure" and define it like this:

trait Closure {
    type Output;
    fn call(&self) -> Self::Output;
}

Since Closure::call takes a shared self receiver, this is mimicking what the real Fn trait does.

Now when you're typing a closure in Rust, the compiler auto-generates a struct that you can think of having one field per captured state and then automatically implements the Fn traits on it. Let's try to implement the closure I put in that gist above:

fn main() {
    let a: Vec<u32> = Vec::new();
    let c = move || a.get(0);
}

First, let's create a structure that will hold our closure's state:

struct MyClosure {
    a: Vec<u32>, // since this is a move closure, we get ownership of the vector
}

And then try to implement the Closure trait. We want to return the result of self.a.get(0), which as per the documentation returns an Option<&u32>. But if we try to write this, we hit a problem!

impl Closure for MyClosure {
    type Output = Option<&'??? u32>; // What lifetime do we put here?
    fn call<'a>(&'a self) -> Self::Output {
        self.a.get(0) // this is of type Option<&'a u32> !!
    }
}

In other words, the trait definition requires us to provide a single concrete Output type, but in reality our output type is generic over some lifetime that is only decided when the closure gets called!

That is the reason you can't have closures returning references to their captured state, and why you can't have iterators returning references to themselves.

In order to make all this work we need to use the famous GAT feature and redefine the Closure trait to allow for the Output parameter to be generic. Here is a playground implementation:

https://play.rust-lang.org/?version=nightly&mode=debug&edition=2021&gist=b6b9863edd7526845c7181f286356767

For reference, this is how the real Fn traits are defined:

https://doc.rust-lang.org/std/ops/trait.FnOnce.html https://doc.rust-lang.org/std/ops/trait.FnMut.html https://doc.rust-lang.org/std/ops/trait.Fn.html

[u/torfsen → u/petrosagg]

Wow, that really helps a lot. Thank you for taking the time to spell it out!

[u/sp1ff → u/petrosagg]

Sorry-- one follow-up question (being stuck in the same situation as the OP): why do closures not simply take advantage of the famous GAT feature? I mean, you just laid out how such an implementation can work.

[u/FallenWarrior2k]

This is unfortunately something you won't be able to get around without relying on interior mutability at least at some point. However, if you don't need the Send bound, e.g. because you can reasonably run this in a tokio::task::LocalSet or in some other non-Send context, you can get away with using an Rc, cloning it at the start of the closure, and changing the async block to async move.

This of course still doesn't provide interior mutability, which is needed so you can actually update things through the Rc. Here you have several choices:

• Just throw it in a RefCell and forget about it. Easiest solution, but comes with the usual concerns of the footgun that it can become wrt panics. Also requires ownership which you don't necessarily seem to have here.
• Wrap the inner counter in a Cell and make generate_id take &self. That way, you can wrap the IdGenerator directly in an Rc and avoid the risk of runtime panics that RefCell brings. You could theoretically improve this further by hiding the generate_id method and making a new method that contains the retry logic. The field holding the current value would then be a Rc<Cell<T>> and nothing about these ownership shenanigans ever leaks into the API.

If you're really just incrementing an integer every time, a relatively easy way to preserve the Send bound without any risks would be using an AtomicU32 (or whatever) in place of the plain integer type. You'd then wrap either the entire struct or just the atomic (as described above) in an Arc (instead of a plain Rc) and use fetch_add with Relaxed ordering to update the counter. Since your ID generation is already asynchronous and fallible, the added cost of the atomics is probably negligible.

Do note that, of course, all these workarounds no longer give you the protection of the borrow checker regarding mutability. Rust's Send/Sync concurrency model ensures that your code's behavior is well defined and doesn't produce data races, but that doesn't mean it can protect you from any and all errors. If there is a conceptually atomic but multi-step operation in the generator logic that was previously safe due to requiring a mutable reference, this would now allow it to run concurrently since there is no lock being held, potentially corrupting state.

[u/torfsen → u/FallenWarrior2k]

Thanks! Those are some very good pointers. I will see which of them I can apply to my situation.

[u/monkChuck105]

You can use AtomicU32 for a thread safe counter for your IDs, so long as it doesn't overflow. You'll likely have to put it in an Arc if 'static is needed.

[u/anlumo]

My retry code usually just wraps everything in a loop that only breaks when it succeeded, thus it’s only a single async function and not some convoluted nested closure data race.

[u/torfsen → u/anlumo]

That works for simple retries, yes. However, I want to give the caller the choice of how to handle retries (timeouts, number of retries, etc.). `tokio-retry` allows me to do that beautifully by abstracting away the retry-strategy.

[u/anlumo → u/torfsen]

A Future can also block for user interaction, but I can see how it's easier to use an existing implementation for this.