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
11pub trait GroupTokenExt<T>: Sized {
12    #[inline]
13    fn release_on_drop(self, token: T) -> GroupTokenReleaseOnDrop<Self, T> {
14        GroupTokenReleaseOnDrop { inner: self, token }
15    }
16
17    #[inline]
18    fn release_on_ready(self, token: T) -> GroupTokenReleaseOnReady<Self, T> {
19        GroupTokenReleaseOnReady {
20            inner: self,
21            token: Some(token),
22        }
23    }
24}
25
26trait GroupTokenType {}
27
28impl GroupTokenType for GroupTokenFactory {}
29impl GroupTokenType for GroupToken {}
30impl GroupTokenType for MonoGroupToken {}
31
32impl<T: GroupTokenType, F: Future> GroupTokenExt<T> for F {}
33
34pin_project! {
35    #[derive(Debug, Into)]
36    pub struct GroupTokenReleaseOnDrop<F, T> {
37        #[pin]
38        inner: F,
39        token: T,
40    }
41}
42
43pin_project! {
44    #[derive(Debug, Into)]
45    pub struct GroupTokenReleaseOnReady<F, T> {
46        #[pin]
47        inner: F,
48        token: Option<T>,
49    }
50}
51
52impl<F, T> GroupTokenReleaseOnDrop<F, T> {
53    #[inline]
54    pub fn inner_pin(self: Pin<&mut Self>) -> Pin<&mut F> {
55        self.project().inner
56    }
57
58    #[inline]
59    pub fn group_token(&self) -> &T {
60        &self.token
61    }
62}
63
64impl<F, T> GroupTokenReleaseOnReady<F, T> {
65    #[inline]
66    pub fn inner_pin(self: Pin<&mut Self>) -> Pin<&mut F> {
67        self.project().inner
68    }
69
70    #[inline]
71    pub fn group_token(&self) -> Option<&T> {
72        self.token.as_ref()
73    }
74}
75
76impl<F: Future, T> Future for GroupTokenReleaseOnDrop<F, T> {
77    type Output = F::Output;
78
79    #[inline]
80    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
81        self.inner_pin().poll(cx)
82    }
83}
84
85impl<F: Future, T> Future for GroupTokenReleaseOnReady<F, T> {
86    type Output = F::Output;
87
88    #[inline]
89    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
90        let this = self.project();
91        let res = this.inner.poll(cx);
92        if res.is_ready() {
93            drop(this.token.take());
94        }
95        res
96    }
97}
compact_waitgroup/ext.rs

compact_waitgroup/
ext.rs
