Is there a way to add a cleanup routine to a Future?

laptou 0
Viewed 57

Is there a way to add code to a Future/async fn in Rust such that it will always run when the execution of the Future is terminating, whether this is because the Future is being dropped or because it has finished normally?

async-await rust
1 Answers

I found this implemented by the future-utils crate called Finally but it relies on an older version of the Future trait before it was stabilized in the standard library. I was bamboozled thinking it was the futures-utils crate, which IS up to date, but doesn't have this functionality.

Regardless, here's an implementation that will call a function when a future is dropped:

[dependencies]
pin-project = "1.0.12"

//! on_drop.rs
use std::future::Future;
use std::mem::ManuallyDrop;
use std::pin::Pin;
use std::task::{Poll, Context};

use pin_project::{pin_project, pinned_drop};

#[pin_project(PinnedDrop)]
pub struct OnDrop<Fut, F> where F: FnOnce() {
    #[pin]
    future: Fut,
    drop_fn: ManuallyDrop<F>,
}

impl <Fut, F> OnDrop<Fut, F>
where F: FnOnce()
{
    pub fn new(future: Fut, drop_fn: F) -> Self {
        Self {
            future,
            drop_fn: ManuallyDrop::new(drop_fn),
        }
    }
}

impl<Fut, F> Future for OnDrop<Fut, F>
where 
    F: FnOnce(),
    Fut: Future,
{
    type Output = Fut::Output;

    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        let this = self.project();
        this.future.poll(cx)
    }
}

#[pinned_drop]
impl<Fut, F> PinnedDrop for OnDrop<Fut, F>
where F: FnOnce()
{
    fn drop(self: Pin<&mut Self>) {
        let mut this = self.project();

        // SAFETY: we always construct `drop_fn` with a value and this is the
        // only place we take the value out of it, when the value is being
        // dropped.
        let drop_fn = unsafe { ManuallyDrop::take(&mut this.drop_fn) };
        drop_fn();
    }
}

pub trait FutureOnDropExt {
    fn on_drop<F: FnOnce()>(self, drop_fn: F) -> OnDrop<Self, F> 
    where 
        Self: Sized;
}

impl<Fut> FutureOnDropExt for Fut where Fut: Future {
    fn on_drop<F: FnOnce()>(self, drop_fn: F) -> OnDrop<Self, F> 
    where 
        Self: Sized
    {
        OnDrop::new(self, drop_fn)
    }
}

And can be used like:

//! main.rs
mod on_drop;
use on_drop::FutureOnDropExt;

fn main() {
    async { println!("this doesn't print since we don't poll") }
        .on_drop(|| println!("this does print when dropped"));
}

this does print when dropped

I took a quick look to see if anyone had suggested this for the futures crate but didn't see anything relevant. If you have a genuine use-case for this, I'd suggest proposing such an adapter. Either it could be added in the future, or sparking the idea may find there is a flaw in the technique or a better solution to achieve your goals.

kmdreko0
Related