compact-waitgroup 0.3.6

A compact asynchronous WaitGroup synchronization primitive.
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130

use core::pin::Pin;
use core::task::{Context, Poll};

use derive_more::Into;
use pin_project_lite::pin_project;

use crate::group::GroupTokenFactory;
use crate::{GroupToken, MonoGroupToken};

/// Extension trait for futures to automatically release group tokens.
pub trait GroupTokenExt<T>: Sized {
    /// Releases the group token when the future is ready or dropped.
    #[inline]
    fn release_on_ready(self, token: T) -> GroupTokenReleaseOnReady<Self, T> {
        GroupTokenReleaseOnReady {
            inner: self,
            token: Some(token),
        }
    }

    /// Releases the group token when the future is dropped.
    ///
    /// The token is held for the entire lifetime of the future, even if the
    /// future is ready.
    #[inline]
    fn release_on_drop(self, token: T) -> GroupTokenReleaseOnDrop<Self, T> {
        GroupTokenReleaseOnDrop { inner: self, token }
    }
}

/// Extension trait for `FnOnce` to automatically release group tokens.
pub trait GroupTokenFuncExt<T, Output>: Sized {
    /// Releases the group token when the closure returns.
    fn release_on_return(self, token: T) -> impl FnOnce() -> Output + Send;
}

trait GroupTokenType: Sync + Send + 'static {}

impl GroupTokenType for GroupTokenFactory {}
impl GroupTokenType for GroupToken {}
impl GroupTokenType for MonoGroupToken {}
impl<T: GroupTokenType> GroupTokenType for Option<T> {}

impl<T: GroupTokenType, F: Future> GroupTokenExt<T> for F {}

impl<T: GroupTokenType, Output, F: Send + FnOnce() -> Output> GroupTokenFuncExt<T, Output> for F {
    #[inline]
    fn release_on_return(self, token: T) -> impl FnOnce() -> Output + Send {
        move || {
            let res = (self)();
            drop(token);
            res
        }
    }
}

pin_project! {
    /// Wrapper that releases a token when the future is ready or dropped.
    ///
    /// Created by [`GroupTokenExt::release_on_ready`].
    #[derive(Debug, Into)]
    pub struct GroupTokenReleaseOnReady<F, T> {
        #[pin]
        inner: F,
        token: Option<T>,
    }
}

pin_project! {
    /// Wrapper that releases a token when the future is dropped.
    ///
    /// Created by [`GroupTokenExt::release_on_drop`].
    #[derive(Debug, Into)]
    pub struct GroupTokenReleaseOnDrop<F, T> {
        #[pin]
        inner: F,
        token: T,
    }
}

impl<F, T> GroupTokenReleaseOnDrop<F, T> {
    /// Returns a pinned mutable reference to the inner future.
    #[inline]
    pub fn inner_pin(self: Pin<&mut Self>) -> Pin<&mut F> {
        self.project().inner
    }

    /// Returns a reference to the associated token.
    #[inline]
    pub fn group_token(&self) -> &T {
        &self.token
    }
}

impl<F, T> GroupTokenReleaseOnReady<F, T> {
    /// Returns a pinned mutable reference to the inner future.
    #[inline]
    pub fn inner_pin(self: Pin<&mut Self>) -> Pin<&mut F> {
        self.project().inner
    }

    /// Returns a reference to the associated token if not yet released.
    #[inline]
    pub fn group_token(&self) -> Option<&T> {
        self.token.as_ref()
    }
}

impl<F: Future, T> Future for GroupTokenReleaseOnDrop<F, T> {
    type Output = F::Output;

    #[inline]
    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        self.inner_pin().poll(cx)
    }
}

impl<F: Future, T> Future for GroupTokenReleaseOnReady<F, T> {
    type Output = F::Output;

    #[inline]
    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        let this = self.project();
        let res = this.inner.poll(cx);
        if res.is_ready() {
            drop(this.token.take());
        }
        res
    }
}