Search code examples
rustasync-awaitborrow-checkerrust-sqlx

Retry an async function that takes mutable ref as argument


I have a function fetch_stuff that needs a mutable ref to a sqlx connection. I would like to retry this function up to N times if it fails. I need this retry behavior for many similar functions so I tried to make a generic retry_up_to_n_times func, that takes an async fnmut and the number of retries allowed as input.

#[derive(Clone)]
pub struct App {
    db_pool: sqlx::PgPool,
}

impl App {
    pub async fn foo(&self) {
        let mut db_conn = self.db_pool.acquire().await.unwrap();
        retry_up_to_n_times(|| fetch_stuff(&mut db_conn), 3).await;
    }
}

pub async fn fetch_stuff(db_conn: &mut sqlx::postgres::PgConnection) -> Result<(), StorageError> {
    //do stuff with db_conn
    Ok(())
}


pub async fn retry_up_to_n_times<F, T, O, E>(mut func: F, max_tries: usize) -> T::Output
where
    F: FnMut() -> T,
    T: std::future::Future<Output = Result<O, E>>,
    E: Error
{
    let mut fail_count = 0;
    loop {
        match func().await {
            Ok(t) => return Ok(t),
            Err(e) => {
                fail_count += 1;
                if fail_count >= max_tries {
                    return Err(e);
                }
            }
        }
    }
}

The compiler gives me this error

error: captured variable cannot escape `FnMut` closure body
  --> src/app/tag.rs:32:32
   |
31 |         let mut db_conn = self.db_pool.acquire().await.unwrap();
   |             ----------- variable defined here
32 |         retry_up_to_n_times(|| fetch_stuff(&mut db_conn), 3).await;
   |                              - ^^^^^^^^^^^^^^^^^-------^
   |                              | |                |
   |                              | |                variable captured here
   |                              | returns an `async` block that contains a reference to a captured variable, which then escapes the closure body
   |                              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

I don't have much experience with rust so I'm not sure I understand it fully. I've read all I could find on the subject and I tried playing with Rc and RefCell, without success. I also tried using tokio_retry, but ended up with the same issue.

Is what I'm trying to do even possible ?


Solution

  • The simplest approach would be to acquire a pool connection in the closure. This is probably what you want anyway; if the connection you're using dies and this causes an error, you should get another connection on the next retry, otherwise it will be doomed to fail repeatedly since the connection is dead.

    pub async fn foo(&self) {
        retry_up_to_n_times(
            || async {
                let mut db_conn = self.db_pool.acquire().await.unwrap();
                fetch_stuff(&mut db_conn).await
            },
            3,
        )
        .await;
    }
    

    (Ideally you'd return the error instead of unwrap()ping it, but I'll leave that as an exercise for you.)

    To explain why this isn't allowed, nothing about FnMut prevents code from calling it a second time before the first Future is dropped. This would cause two futures to effectively have a &mut to the same value, which is a violation of Rust's aliasing rules. So, a FnMut may not return a value that contains an exclusive borrow of a captured value.

    It's a bit obscured by the async stuff in your code, but fetch_stuff(&mut db_conn) returns a future that contains an exclusive borrow of db_conn.

    A simplified example illustrates the problem Rust is trying to prevent:

    use std::fmt::Display;
    
    fn make_two_muts<T: Display>(mut f: impl FnMut() -> T) {
        let m1 = f();
        let m2 = f();
    
        // Can we have two mutable references to the same value coexisting?
        format!("{m1} {m2}");
    }
    
    fn main() {
        let mut x = 1;
    
        // No, because this isn't allowed:
        make_two_muts(move || &mut x);
    }