How to create a FnMut closure that captures a variable and returns a Future, such that it retains the state of the captured variable between subsequent invocations?

Consider the following code, where a function returns a boxed closure. The closure captures a variable and mutates it.
The output shows that the state of the captured variable is retained between subsequent invocations of the closure.

fn make_closure() -> Box<dyn FnMut(u32)> {
    let mut counter = 0;
    Box::new(move |arg: u32| {
        println!("Counter before: {counter}");
        counter += arg;
        println!("Counter after: {counter}");
    })
}

#[test]
fn test2() {
    let mut cl = make_closure();
    // Prints:
    // Counter before: 0
    // Counter after: 1
    cl(1);
    // Prints:
    // Counter before: 1
    // Counter after: 3
    cl(2);
}

I want to update make_closure such that the update of the captured variable occurs in an async block.
In other words, I need to await some async functions before and after updating the captured variable.

This was my initial attempt, but the output is not what I expected:

fn make_closure_v2() -> Box<dyn FnMut(u32) -> Pin<Box<dyn Future<Output = ()>>>> {
    let mut counter = 0;
    Box::new(move |arg: u32| {
        Box::pin(async move {
            println!("Counter before: {counter}");
            counter += arg;
            println!("Counter after: {counter}");
        })
    })
}

#[tokio::test]
async fn test_v2() {
    let mut cl = make_closure_v2();
    // Prints:
    // Counter before: 0
    // Counter after: 1
    cl(1).await;
    // Prints:
    // Counter before: 0
    // Counter after: 2
    cl(2).await;
}

How can I circumvent the issue? Do I have to resort to some usage of unstable async closures?

The problem is that each future gets its own copy of counter. That's what the move keyword does: it moves the value into the closure/future, rather than trying to share it. However, you can't just directly share the value, since that could create multiple mutable references to the same value.

You can use an Arc<AtomicU32> to share the counter between futures, though there might be a more efficient way of doing it. This also uses the most conservative memory ordering (SeqCst) because I am bad at memory ordering :slight_smile:

use std::future::Future;
use std::pin::Pin;

use std::sync::Arc;
use std::sync::atomic::AtomicU32;
use std::sync::atomic::Ordering::SeqCst;

fn make_closure_v3() -> Box<dyn FnMut(u32) -> Pin<Box<dyn Future<Output = ()>>>> {
    let counter = Arc::new(AtomicU32::new(0));
    Box::new(move |arg: u32| {
        let counter = counter.clone();
        Box::pin(async move {
            let counter_value = counter.load(SeqCst);
            println!("Counter before: {counter_value}");
            counter.fetch_add(arg, SeqCst);
            let counter_value = counter.load(SeqCst);
            println!("Counter after: {counter_value}");
        })
    })
}

#[tokio::test]
async fn test_v3() {
    let mut cl = make_closure_v3();
    // Prints:
    // Counter before: 0
    // Counter after: 1
    cl(1).await;
    // Prints:
    // Counter before: 1
    // Counter after: 3
    cl(2).await;
}
2 Likes

In order for the future returned from the closure to be able to modify the closure’s state, it would need to keep a borrow of the closure, however as it stands right now, the Fn… traits do not allow such a thing. I.e. the call(&mut self, args…) method of a closure is not allowed to return anything related to the borrow self. As far as I’m aware, this is a well-known limitation of creating “async closures” that we should try to address in future Rust versions, but last time I checked, there wasn’t a solution for it yet, even when you include unstable features.

In case it isn’t clear what’s going on, the move in the async move is what’s making the difference: It introduces an additional copying of the counter, and this copy happens before anything is mutated, so each closure call ends up starting with the same start value again.

As @DanielKeep already demonstrated, you can circumvent the issue by avoiding the closure call to borrow from the closure itself, by moving to shared ownership as provided by Arc. Note that this also means that the execution order depends on the order of the .await calls, not the closure calls, so if you call the closure twice, but .await the second result first, it’ll be the one observing the initial value of zero.


If you do not need to adhere to any specific FnMut… trait bound, you can of course simply avoid creating any closures and return an object with a async fn(&mut self…)… method:

struct NotAClosure {
    counter: u32,
}
impl NotAClosure {
    async fn call(&mut self, arg: u32) {
        let Self { counter } = self;

        println!("Counter before: {counter}");
        *counter += arg;
        println!("Counter after: {counter}");
    }
}
fn make_closure_v3() -> NotAClosure {
    NotAClosure {
        counter: 0,
    }
}

#[tokio::test]
async fn test_v3() {
    let mut cl = make_closure_v2();
    // Prints:
    // Counter before: 0
    // Counter after: 1
    cl.call(1).await;
    // Prints:
    // Counter before: 1
    // Counter after: 3
    cl.call(2).await;
}
3 Likes

Thank you all, this was my first post and you have been really helpful. Great community!

The API I am designing requires FnMut, and is really ergonomic with it, thus I can't sacrifice it.

@steffahn I see that you mention async closures. Are you referring to the following unstable feature?

If so, can you please help me in determining the signature of the return type? I can't figure it out.

fn make_async_closure() -> ??? {
    let mut counter = 0;
    Box::new(async move |arg: u32| {
        println!("Counter before: {counter}");
        counter += arg;
        println!("Counter after: {counter}");
    })
}

As I mentioned above, as far as I’m aware, this unstable feature unfortunately doesn’t support your use case either.

#![feature(async_closure)]

#[tokio::main]
async fn main() {
    let mut counter = 0;
    let f = async move |arg: u32| {
        println!("Counter before: {counter}");
        counter += arg;
        println!("Counter after: {counter}");
    };
    
    f(1).await;
    f(2).await;
}
Counter before: 0
Counter after: 1
Counter before: 0
Counter after: 2
2 Likes

This topic was automatically closed 90 days after the last reply. We invite you to open a new topic if you have further questions or comments.