ntex_util/future/
on_drop.rs1#![allow(clippy::unused_unit)]
2use std::{cell::Cell, fmt, future::Future, pin::Pin, task::Context, task::Poll};
3
4pub 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 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
38pub 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}