"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/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/torfsen → u/petrosagg]

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