compact_waitgroup/
ext.rs

1use core::{
2    pin::Pin,
3    task::{Context, Poll},
4};
5
6use derive_more::Into;
7use pin_project_lite::pin_project;
8
9use crate::{GroupToken, MonoGroupToken, group::GroupTokenFactory};
10
11/// Extension trait for futures to automatically release group tokens.
12pub trait GroupTokenExt<T>: Sized {
13    /// Releases the group token when the future is ready or dropped.
14    #[inline]
15    fn release_on_ready(self, token: T) -> GroupTokenReleaseOnReady<Self, T> {
16        GroupTokenReleaseOnReady {
17            inner: self,
18            token: Some(token),
19        }
20    }
21
22    /// Releases the group token when the future is dropped.
23    ///
24    /// The token is held for the entire lifetime of the future, even if the
25    /// future is ready.
26    #[inline]
27    fn release_on_drop(self, token: T) -> GroupTokenReleaseOnDrop<Self, T> {
28        GroupTokenReleaseOnDrop { inner: self, token }
29    }
30}
31
32/// Extension trait for `FnOnce` to automatically release group tokens.
33pub trait GroupTokenFuncExt<T, Output>: Sized {
34    /// Releases the group token when the closure returns.
35    fn release_on_return(self, token: T) -> impl FnOnce() -> Output + Send;
36}
37
38trait GroupTokenType: Sync + Send + 'static {}
39
40impl GroupTokenType for GroupTokenFactory {}
41impl GroupTokenType for GroupToken {}
42impl GroupTokenType for MonoGroupToken {}
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
