Future on-drop helper (#501)

ntex-util/CHANGES.md

@@ -4,6 +4,8 @@
 
 * Add EitherService/EitherServiceFactory
 
+* Add future on drop handler
+
 ## [2.8.0] - 2024-12-04
 
 * Use updated Service trait

ntex-util/src/future/mod.rs

@@ -7,12 +7,14 @@ pub use futures_sink::Sink;
 mod either;
 mod join;
 mod lazy;
+mod on_drop;
 mod ready;
 mod select;
 
 pub use self::either::Either;
 pub use self::join::{join, join_all};
 pub use self::lazy::{lazy, Lazy};
+pub use self::on_drop::{OnDropFn, OnDropFutureExt};
 pub use self::ready::Ready;
 pub use self::select::select;
 

ntex-util/src/future/on_drop.rs

@@ -0,0 +1,100 @@
+use std::{cell::Cell, fmt, future::Future, pin::Pin, task::Context, task::Poll};
+
+/// Execute fn during drop
+pub struct OnDropFn<F: FnOnce()> {
+    f: Cell<Option<F>>,
+}
+
+impl<F: FnOnce()> OnDropFn<F> {
+    pub fn new(f: F) -> Self {
+        Self {
+            f: Cell::new(Some(f)),
+        }
+    }
+
+    /// Cancel fn execution
+    pub fn cancel(&self) {
+        self.f.take();
+    }
+}
+
+impl<F: FnOnce()> fmt::Debug for OnDropFn<F> {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        f.debug_struct("OnDropFn")
+            .field("f", &std::any::type_name::<F>())
+            .finish()
+    }
+}
+
+impl<F: FnOnce()> Drop for OnDropFn<F> {
+    fn drop(&mut self) {
+        if let Some(f) = self.f.take() {
+            f()
+        }
+    }
+}
+
+/// Trait adds future on_drop support
+pub trait OnDropFutureExt: Future + Sized {
+    fn on_drop<F: FnOnce()>(self, on_drop: F) -> OnDropFuture<Self, F> {
+        OnDropFuture::new(self, on_drop)
+    }
+}
+
+impl<F: Future> OnDropFutureExt for F {}
+
+pin_project_lite::pin_project! {
+    pub struct OnDropFuture<Ft: Future, F: FnOnce()> {
+        #[pin]
+        fut: Ft,
+        on_drop: OnDropFn<F>
+    }
+}
+
+impl<Ft: Future, F: FnOnce()> OnDropFuture<Ft, F> {
+    pub fn new(fut: Ft, on_drop: F) -> Self {
+        Self {
+            fut,
+            on_drop: OnDropFn::new(on_drop),
+        }
+    }
+}
+
+impl<Ft: Future, F: FnOnce()> Future for OnDropFuture<Ft, F> {
+    type Output = Ft::Output;
+
+    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+        let this = self.project();
+        match this.fut.poll(cx) {
+            Poll::Ready(r) => {
+                this.on_drop.cancel();
+                Poll::Ready(r)
+            }
+            Poll::Pending => Poll::Pending,
+        }
+    }
+}
+
+#[cfg(test)]
+mod test {
+    use std::future::{pending, poll_fn};
+
+    use super::*;
+
+    #[ntex_macros::rt_test2]
+    async fn on_drop() {
+        let mut dropped = false;
+
+        let mut f = pending::<()>().on_drop(|| {
+            dropped = true;
+        });
+        poll_fn(|cx| {
+            let _ = Pin::new(&mut f).poll(cx);
+            Poll::Ready(())
+        })
+        .await;
+
+        drop(f);
+        assert!(dropped);
+    }
+}