ntex_util/future/
on_drop.rs

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