Running code during application shutdown

help
1

I am currently rewriting a project from java to rust due to memory usage concerns. It is a music player which requires storing some state on disk. The state is auto-saved periodically but should always be saved when the application shuts down.

In the Java world I used to rely heavily on an event system, where it is very easy to store a Runnable for an indefinite amount of time and call it when needed. I have attempted to do something similar to that in rust using Future but am struggling to unbox it in a way that allows me to await it. I've tried combinations of deref, clone, to_owned, into_future and more but I couldn't figure it out.

At some point I had something that compiled fine, but it was storing Fn() -> () which is not async and I need to be able to await in the handlers I register. Tried with AsyncFn but I could not get it to work, so I looked at what tokio::spawn accepts and tried to use Future instead.

Is what I'm trying to do even possible? More importantly: is this just an inappropriate approach to the problem in this language?

My code:

use std::{sync::LazyLock};

use futures::{FutureExt, future::BoxFuture};
use log::error;
use tokio::{signal, sync::Mutex};

static RAN_ONCE: LazyLock<Mutex<bool>> = LazyLock::new(|| Mutex::new(false));
static HANDLERS: LazyLock<Mutex<Vec<BoxFuture<'static, ()>>>> = LazyLock::new(|| Mutex::new(Vec::new()));

pub async fn register_handler<F>(future: F) where F: Future<Output = ()> + Send + Sync + 'static, {
    HANDLERS.lock().await.push(future.boxed());
}

pub async fn run_handlers() {
    let mut ran_once = RAN_ONCE.lock().await;
    if *ran_once {
        return;
    }
    *ran_once = true;

    drop(ran_once);

    let handlers = HANDLERS.lock().await;

    for i in 0..handlers.len() {
        if let Some(handler) = handlers.get(i) {
            handler.await; // &Pin<Box<dyn futures::Future<Output = ()> + std::marker::Send>> is not a future
        }
    }
}
2

I think you need to await the future by value:

This reverses the order, but any method that removes values from the &mut list should work.
UPDATE: Below look much better :backhand_index_pointing_down:

3

The big thing you're doing that's making your life hard is using indexing, rather than trying to find an iterator over HANDLERS that does what you want. As a result, instead of having exclusive access to the save future so that you can poll it, you have shared access and can't change it (which implies you can't await it, either, since that changes the internal state of the future).

Does

use std::{sync::LazyLock};

use futures::{FutureExt, future::BoxFuture};
use log::error;
use tokio::{signal, sync::Mutex};

static HANDLERS: LazyLock<Mutex<Vec<BoxFuture<'static, ()>>>> = LazyLock::new(|| Mutex::new(Vec::new()));

pub async fn register_handler<F>(future: F) where F: Future<Output = ()> + Send + Sync + 'static, {
    HANDLERS.lock().await.push(future.boxed());
}

pub async fn run_handlers() {
    let mut handlers = HANDLERS.lock().await;

    for handler in handlers.drain(..) {
        handler.await;
    }
}

work for you? This has two simplifications:

  1. We use drain(..) to get an iterator over the locked handlers that lets you take ownership of each handler in turn.
  2. Because each handler is extracted and run as an owned future, we don't need the RAN_ONCE variable - instead, we lean on Rust's type system to guarantee that you get ownership of each future at most once.

There is also an alternative, that I'd prefer, but it's further from your original code:

use std::sync::{LazyLock, Mutex};

use futures::{FutureExt, future::BoxFuture};

static HANDLERS: LazyLock<Mutex<Vec<BoxFuture<'static, ()>>>> =
    LazyLock::new(|| Mutex::new(Vec::new()));

pub fn register_handler<F>(future: F)
where
    F: Future<Output = ()> + Send + Sync + 'static,
{
    HANDLERS.lock().expect("poison").push(future.boxed());
}

pub async fn run_handlers() {
    let handlers: Vec<_> = std::mem::take(&mut HANDLERS.lock().expect("poison"));

    for handler in handlers {
        handler.await;
    }
}

This doesn't use an async lock, and instead removes all of the handlers at once when you call run_handlers, and then runs them one at a time. I prefer this, because if you're using a multi-threaded executor (the default Tokio executor, for example), holding a sync lock over a .await is a compile-time failure (because MutexGuard happens not to be Send), and thus it's easier to see that this can't deadlock (because the lock is held for such a short period of time).

4

Thanks! I thought I was making my life easier by not worrying about which function to use, but not having ownership was the problem.

Do you refer to how a handler may attempt to register another handler? Would this be solvable keeping tokio Mutexes doing something like this?

pub async fn run_handlers() {
    let mut handlers = HANDLERS.lock().await;

    let mut temp_handlers = Vec::new();
    for handler in handlers.drain(..) {
        temp_handlers.push(handler);
    }

    drop(handlers);

    for handler in temp_handlers.drain(..) {
        handler.await;
    }
}

(I am trying to avoid using unwrap or expect in my code, which wouldn't panic here anyway unless something else panicked but better safe than sorry)

5

If you don't care about panics, there's std::sync::nonpoison::Mutex, or parking_lot::Mutex that you can use.

The problem with Tokio's async Mutex is that it can be held over a .await - which means that to reason about whether your program contains deadlocks involves reasoning about all places where the Mutex is held (and making sure you're not confusing yourself). A sync Mutex, like std::sync::Mutex, cannot be held across a .await, which means that you only have to look at the code from .lock() to the first .await.