ntex_util/future/
on_drop.rs

1use std::{cell::Cell, fmt, future::Future, pin::Pin, task::Context, task::Poll};
2
3/// Execute fn during drop
4pub struct OnDropFn<F: FnOnce()> {
5    f: Cell<Option<F>>,
6}
7
8impl<F: FnOnce()> OnDropFn<F> {
9    pub fn new(f: F) -> Self {
10        Self {
11            f: Cell::new(Some(f)),
12        }
13    }
14
15    /// Cancel fn execution
16    pub fn cancel(&self) {
17        self.f.take();
18    }
19}
20
21impl<F: FnOnce()> fmt::Debug for OnDropFn<F> {
22    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
23        f.debug_struct("OnDropFn")
24            .field("f", &std::any::type_name::<F>())
25            .finish()
26    }
27}
28
29impl<F: FnOnce()> Drop for OnDropFn<F> {
30    fn drop(&mut self) {
31        if let Some(f) = self.f.take() {
32            f()
33        }
34    }
35}
36
37/// Trait adds future on_drop support
38pub trait OnDropFutureExt: Future + Sized {
39    fn on_drop<F: FnOnce()>(self, on_drop: F) -> OnDropFuture<Self, F> {
40        OnDropFuture::new(self, on_drop)
41    }
42}
43
44impl<F: Future> OnDropFutureExt for F {}
45
46pin_project_lite::pin_project! {
47    pub struct OnDropFuture<Ft: Future, F: FnOnce()> {
48        #[pin]
49        fut: Ft,
50        on_drop: OnDropFn<F>
51    }
52}
53
54impl<Ft: Future, F: FnOnce()> OnDropFuture<Ft, F> {
55    pub fn new(fut: Ft, on_drop: F) -> Self {
56        Self {
57            fut,
58            on_drop: OnDropFn::new(on_drop),
59        }
60    }
61}
62
63impl<Ft: Future, F: FnOnce()> Future for OnDropFuture<Ft, F> {
64    type Output = Ft::Output;
65
66    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
67        let this = self.project();
68        match this.fut.poll(cx) {
69            Poll::Ready(r) => {
70                this.on_drop.cancel();
71                Poll::Ready(r)
72            }
73            Poll::Pending => Poll::Pending,
74        }
75    }
76}
77
78#[cfg(test)]
79mod test {
80    use std::future::{pending, poll_fn};
81
82    use super::*;
83
84    #[ntex_macros::rt_test2]
85    async fn on_drop() {
86        let f = OnDropFn::new(|| ());
87        assert!(format!("{:?}", f).contains("OnDropFn"));
88        f.cancel();
89        assert!(f.f.get().is_none());
90
91        let mut dropped = false;
92        let mut f = pending::<()>().on_drop(|| {
93            dropped = true;
94        });
95        poll_fn(|cx| {
96            let _ = Pin::new(&mut f).poll(cx);
97            Poll::Ready(())
98        })
99        .await;
100
101        drop(f);
102        assert!(dropped);
103    }
104}
ntex_util/future/on_drop.rs

ntex_util/future/
on_drop.rs
