1
use std::{
2
    future::Future,
3
    pin::Pin,
4
    task::{Context, Poll},
5
};
6

            
7
use crate::{BudgetContext, BudgetContextData, Budgetable, Container};
8

            
9
/// Spends `amount` from the curent budget.
10
///
11
/// This is a future that must be awaited. This future is created by `spend()`.
12
#[derive(Debug)]
13
#[must_use = "budget is not spent until this future is awaited"]
14
pub(crate) struct SpendBudget<'a, Backing, Budget>
15
where
16
    Backing: Container<BudgetContextData<Budget>>,
17
{
18
    pub(crate) context: &'a BudgetContext<Backing, Budget>,
19
    pub(crate) amount: usize,
20
}
21

            
22
impl<'a, Backing, Budget> Future for SpendBudget<'a, Backing, Budget>
23
where
24
    Budget: Budgetable,
25
    Backing: Container<BudgetContextData<Budget>>,
26
{
27
    type Output = ();
28

            
29
464973
    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
30
464993
        self.context.data.map_locked(|data| {
31
464993
            if data.budget.spend(self.amount) {
32
21342
                data.budget.remove_waker(cx.waker());
33
21342
                Poll::Ready(())
34
            } else {
35
                // Not enough budget
36
430282
                match &mut data.paused_future {
37
430282
                    Some(existing_waker) if existing_waker.will_wake(cx.waker()) => {
38
6104
                        *existing_waker = cx.waker().clone();
39
6104
                    }
40
437540
                    waker => {
41
437540
                        *waker = Some(cx.waker().clone());
42
437540
                    }
43
                }
44

            
45
443644
                data.budget.add_waker(cx.waker());
46
443644

            
47
443644
                Poll::Pending
48
            }
49
464986
        })
50
464973
    }
51
}
52

            
53
macro_rules! define_public_interface {
54
    ($modulename:ident, $backing:ident, $moduledocs:literal) => {
55
        #[doc = $moduledocs]
56
        pub mod $modulename {
57
            use std::{
58
                future::Future,
59
                pin::Pin,
60
                task::{Context, Poll},
61
            };
62

            
63
            use crate::{BudgetContextData, Budgetable};
64

            
65
            type Backing<Budget> = crate::$backing<crate::BudgetContextData<Budget>>;
66

            
67
            /// Spends `amount` from the curent budget.
68
            ///
69
            /// This is a future that must be awaited. This future is created by `spend()`.
70
            #[derive(Debug)]
71
            #[must_use = "budget is not spent until this future is awaited"]
72
            pub struct SpendBudget<'a, Budget>(
73
                crate::SpendBudget<'a, crate::$backing<BudgetContextData<Budget>>, Budget>,
74
            )
75
            where
76
                Budget: Budgetable;
77

            
78
            impl<'a, Budget> Future for SpendBudget<'a, Budget>
79
            where
80
                Budget: Budgetable,
81
            {
82
                type Output = ();
83

            
84
464985
                fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
85
464985
                    let inner = Pin::new(&mut self.0);
86
464985
                    inner.poll(cx)
87
464985
                }
88
            }
89

            
90
            impl<'a, Budget> From<super::SpendBudget<'a, Backing<Budget>, Budget>>
91
                for SpendBudget<'a, Budget>
92
            where
93
                Budget: Budgetable,
94
            {
95
21342
                fn from(future: super::SpendBudget<'a, Backing<Budget>, Budget>) -> Self {
96
21342
                    Self(future)
97
21342
                }
98
            }
99
        }
100
    };
101
}
102

            
103
define_public_interface!(
104
    threadsafe,
105
    SyncContainer,
106
    "Threadsafe (`Send + Sync`) budget spending"
107
);
108
define_public_interface!(
109
    singlethreaded,
110
    NotSyncContainer,
111
    "Single-threaded (`!Send + !Sync`) budget spending"
112
);