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 {}
42
43impl<T: GroupTokenType, F: Future> GroupTokenExt<T> for F {}
44
45impl<T: GroupTokenType, Output, F: Send + FnOnce() -> Output> GroupTokenFuncExt<T, Output> for F {
46    #[inline]
47    fn release_on_return(self, token: T) -> impl FnOnce() -> Output + Send {
48        move || {
49            let res = (self)();
50            drop(token);
51            res
52        }
53    }
54}
55
56pin_project! {
57    /// Wrapper that releases a token when the future is ready or dropped.
58    ///
59    /// Created by [`GroupTokenExt::release_on_ready`].
60    #[derive(Debug, Into)]
61    pub struct GroupTokenReleaseOnReady<F, T> {
62        #[pin]
63        inner: F,
64        token: Option<T>,
65    }
66}
67
68pin_project! {
69    /// Wrapper that releases a token when the future is dropped.
70    ///
71    /// Created by [`GroupTokenExt::release_on_drop`].
72    #[derive(Debug, Into)]
73    pub struct GroupTokenReleaseOnDrop<F, T> {
74        #[pin]
75        inner: F,
76        token: T,
77    }
78}
79
80impl<F, T> GroupTokenReleaseOnDrop<F, T> {
81    /// Returns a pinned mutable reference to the inner future.
82    #[inline]
83    pub fn inner_pin(self: Pin<&mut Self>) -> Pin<&mut F> {
84        self.project().inner
85    }
86
87    /// Returns a reference to the associated token.
88    #[inline]
89    pub fn group_token(&self) -> &T {
90        &self.token
91    }
92}
93
94impl<F, T> GroupTokenReleaseOnReady<F, T> {
95    /// Returns a pinned mutable reference to the inner future.
96    #[inline]
97    pub fn inner_pin(self: Pin<&mut Self>) -> Pin<&mut F> {
98        self.project().inner
99    }
100
101    /// Returns a reference to the associated token if not yet released.
102    #[inline]
103    pub fn group_token(&self) -> Option<&T> {
104        self.token.as_ref()
105    }
106}
107
108impl<F: Future, T> Future for GroupTokenReleaseOnDrop<F, T> {
109    type Output = F::Output;
110
111    #[inline]
112    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
113        self.inner_pin().poll(cx)
114    }
115}
116
117impl<F: Future, T> Future for GroupTokenReleaseOnReady<F, T> {
118    type Output = F::Output;
119
120    #[inline]
121    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
122        let this = self.project();
123        let res = this.inner.poll(cx);
124        if res.is_ready() {
125            drop(this.token.take());
126        }
127        res
128    }
129}
compact_waitgroup/ext.rs

compact_waitgroup/
ext.rs
