compact_waitgroup/
ext.rs

1use core::pin::Pin;
2use core::task::{Context, Poll};
3
4use derive_more::Into;
5use pin_project_lite::pin_project;
6
7use crate::group::GroupTokenFactory;
8use crate::{GroupToken, MonoGroupToken};
9
10/// Extension trait for futures to automatically release group tokens.
11pub trait GroupTokenExt<T>: Sized {
12    /// Releases the group token when the future is ready or dropped.
13    #[inline]
14    fn release_on_ready(self, token: T) -> GroupTokenReleaseOnReady<Self, T> {
15        GroupTokenReleaseOnReady {
16            inner: self,
17            token: Some(token),
18        }
19    }
20
21    /// Releases the group token when the future is dropped.
22    ///
23    /// The token is held for the entire lifetime of the future, even if the
24    /// future is ready.
25    #[inline]
26    fn release_on_drop(self, token: T) -> GroupTokenReleaseOnDrop<Self, T> {
27        GroupTokenReleaseOnDrop { inner: self, token }
28    }
29}
30
31/// Extension trait for `FnOnce` to automatically release group tokens.
32pub trait GroupTokenFuncExt<T, Output>: Sized {
33    /// Releases the group token when the closure returns.
34    fn release_on_return(self, token: T) -> impl FnOnce() -> Output + Send;
35}
36
37trait GroupTokenType: Sync + Send + 'static {}
38
39impl GroupTokenType for GroupTokenFactory {}
40impl GroupTokenType for GroupToken {}
41impl GroupTokenType for MonoGroupToken {}
42impl<T: GroupTokenType> GroupTokenType for Option<T> {}
43
44impl<T: GroupTokenType, F: Future> GroupTokenExt<T> for F {}
45
46impl<T: GroupTokenType, Output, F: Send + FnOnce() -> Output> GroupTokenFuncExt<T, Output> for F {
47    #[inline]
48    fn release_on_return(self, token: T) -> impl FnOnce() -> Output + Send {
49        move || {
50            let res = (self)();
51            drop(token);
52            res
53        }
54    }
55}
56
57pin_project! {
58    /// Wrapper that releases a token when the future is ready or dropped.
59    ///
60    /// Created by [`GroupTokenExt::release_on_ready`].
61    #[derive(Debug, Into)]
62    pub struct GroupTokenReleaseOnReady<F, T> {
63        #[pin]
64        inner: F,
65        token: Option<T>,
66    }
67}
68
69pin_project! {
70    /// Wrapper that releases a token when the future is dropped.
71    ///
72    /// Created by [`GroupTokenExt::release_on_drop`].
73    #[derive(Debug, Into)]
74    pub struct GroupTokenReleaseOnDrop<F, T> {
75        #[pin]
76        inner: F,
77        token: T,
78    }
79}
80
81impl<F, T> GroupTokenReleaseOnDrop<F, T> {
82    /// Returns a pinned mutable reference to the inner future.
83    #[inline]
84    pub fn inner_pin(self: Pin<&mut Self>) -> Pin<&mut F> {
85        self.project().inner
86    }
87
88    /// Returns a reference to the associated token.
89    #[inline]
90    pub fn group_token(&self) -> &T {
91        &self.token
92    }
93}
94
95impl<F, T> GroupTokenReleaseOnReady<F, T> {
96    /// Returns a pinned mutable reference to the inner future.
97    #[inline]
98    pub fn inner_pin(self: Pin<&mut Self>) -> Pin<&mut F> {
99        self.project().inner
100    }
101
102    /// Returns a reference to the associated token if not yet released.
103    #[inline]
104    pub fn group_token(&self) -> Option<&T> {
105        self.token.as_ref()
106    }
107}
108
109impl<F: Future, T> Future for GroupTokenReleaseOnDrop<F, T> {
110    type Output = F::Output;
111
112    #[inline]
113    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
114        self.inner_pin().poll(cx)
115    }
116}
117
118impl<F: Future, T> Future for GroupTokenReleaseOnReady<F, T> {
119    type Output = F::Output;
120
121    #[inline]
122    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
123        let this = self.project();
124        let res = this.inner.poll(cx);
125        if res.is_ready() {
126            drop(this.token.take());
127        }
128        res
129    }
130}
compact_waitgroup/ext.rs

compact_waitgroup/
ext.rs
