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    /// Wraps this future so that `on_drop` is called if the future is dropped
41    /// before it completes. The callback is cancelled if the future resolves
42    /// successfully.
43    fn on_drop<F: FnOnce()>(self, on_drop: F) -> OnDropFuture<Self, F> {
44        OnDropFuture::new(self, on_drop)
45    }
46}
47
48impl<F: Future> OnDropFutureExt for F {}
49
50pin_project_lite::pin_project! {
51    pub struct OnDropFuture<Ft: Future, F: FnOnce()> {
52        #[pin]
53        fut: Ft,
54        on_drop: OnDropFn<F>
55    }
56}
57
58impl<Ft: Future, F: FnOnce()> OnDropFuture<Ft, F> {
59    /// Creates a new `OnDropFuture` that calls `on_drop` if `fut` is dropped
60    /// before it completes.
61    pub fn new(fut: Ft, on_drop: F) -> Self {
62        Self {
63            fut,
64            on_drop: OnDropFn::new(on_drop),
65        }
66    }
67}
68
69impl<Ft: Future, F: FnOnce()> Future for OnDropFuture<Ft, F> {
70    type Output = Ft::Output;
71
72    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
73        let this = self.project();
74        match this.fut.poll(cx) {
75            Poll::Ready(r) => {
76                this.on_drop.cancel();
77                Poll::Ready(r)
78            }
79            Poll::Pending => Poll::Pending,
80        }
81    }
82}
83
84#[cfg(test)]
85mod test {
86    use std::future::{pending, poll_fn};
87
88    use super::*;
89
90    #[ntex::test]
91    async fn on_drop() {
92        let f = OnDropFn::new(|| ());
93        assert!(format!("{f:?}").contains("OnDropFn"));
94        f.cancel();
95        assert!(f.f.get().is_none());
96
97        let mut dropped = false;
98        let mut f = pending::<()>().on_drop(|| {
99            dropped = true;
100        });
101        poll_fn(|cx| {
102            let _ = Pin::new(&mut f).poll(cx);
103            Poll::Ready(())
104        })
105        .await;
106
107        drop(f);
108        assert!(dropped);
109    }
110}
ntex_util/future/on_drop.rs

ntex_util/future/
on_drop.rs
