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/// Runs a callback when this guard is dropped.
5///
6/// Call [`cancel`](Self::cancel) to drop the guard without running it.
7pub struct OnDropFn<F: FnOnce()> {
8    f: Cell<Option<F>>,
9}
10
11impl<F: FnOnce()> OnDropFn<F> {
12    /// Creates a guard that runs `f` when dropped.
13    pub fn new(f: F) -> Self {
14        Self {
15            f: Cell::new(Some(f)),
16        }
17    }
18
19    /// Discards the callback so it will not run on drop.
20    pub fn cancel(&self) {
21        self.f.take();
22    }
23}
24
25impl<F: FnOnce()> fmt::Debug for OnDropFn<F> {
26    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
27        f.debug_struct("OnDropFn")
28            .field("f", &std::any::type_name::<F>())
29            .finish()
30    }
31}
32
33impl<F: FnOnce()> Drop for OnDropFn<F> {
34    fn drop(&mut self) {
35        if let Some(f) = self.f.take() {
36            f();
37        }
38    }
39}
40
41/// Extension trait for running a callback when a pending future is dropped.
42pub trait OnDropFutureExt: Future + Sized {
43    /// Wraps this future so that `on_drop` is called if the future is dropped
44    /// before it completes. The callback is cancelled once the future
45    /// completes, whatever its output.
46    fn on_drop<F: FnOnce()>(self, on_drop: F) -> OnDropFuture<Self, F> {
47        OnDropFuture::new(self, on_drop)
48    }
49}
50
51impl<F: Future> OnDropFutureExt for F {}
52
53pin_project_lite::pin_project! {
54    /// Future wrapper that runs a callback if dropped before completion.
55    pub struct OnDropFuture<Ft: Future, F: FnOnce()> {
56        #[pin]
57        fut: Ft,
58        on_drop: OnDropFn<F>
59    }
60}
61
62impl<Ft: Future, F: FnOnce()> OnDropFuture<Ft, F> {
63    /// Creates a new `OnDropFuture` that calls `on_drop` if `fut` is dropped
64    /// before it completes.
65    pub fn new(fut: Ft, on_drop: F) -> Self {
66        Self {
67            fut,
68            on_drop: OnDropFn::new(on_drop),
69        }
70    }
71}
72
73impl<Ft: Future, F: FnOnce()> Future for OnDropFuture<Ft, F> {
74    type Output = Ft::Output;
75
76    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
77        let this = self.project();
78        match this.fut.poll(cx) {
79            Poll::Ready(r) => {
80                this.on_drop.cancel();
81                Poll::Ready(r)
82            }
83            Poll::Pending => Poll::Pending,
84        }
85    }
86}
87
88#[cfg(test)]
89mod test {
90    use std::future::{pending, poll_fn};
91
92    use super::*;
93
94    #[ntex::test]
95    async fn on_drop() {
96        let f = OnDropFn::new(|| ());
97        assert!(format!("{f:?}").contains("OnDropFn"));
98        f.cancel();
99        assert!(f.f.get().is_none());
100
101        let mut dropped = false;
102        let mut f = pending::<()>().on_drop(|| {
103            dropped = true;
104        });
105        poll_fn(|cx| {
106            let _ = Pin::new(&mut f).poll(cx);
107            Poll::Ready(())
108        })
109        .await;
110
111        drop(f);
112        assert!(dropped);
113    }
114
115    #[ntex::test]
116    async fn on_drop_completed() {
117        let mut dropped = false;
118        let f = std::future::ready(1).on_drop(|| dropped = true);
119        assert_eq!(f.await, 1);
120        assert!(!dropped);
121    }
122}
ntex_util/future/on_drop.rs

ntex_util/future/
on_drop.rs
