1
1
#![doc = include_str!(".crate-docs.md")]
2
#![warn(
3
    clippy::cargo,
4
    missing_docs,
5
    clippy::pedantic,
6
    future_incompatible,
7
    rust_2018_idioms
8
)]
9
#![allow(
10
    clippy::option_if_let_else,
11
    clippy::module_name_repetitions,
12
    clippy::missing_errors_doc
13
)]
14

            
15
use std::{
16
    cell::RefCell,
17
    marker::PhantomData,
18
    rc::Rc,
19
    sync::{Arc, Mutex},
20
    task::Waker,
21
};
22

            
23
use crate::{sealed::BudgetableSealed, spend::SpendBudget};
24

            
25
/// A budget implementation compatible with any async executor.
26
pub mod asynchronous;
27
/// A standalone implementation does not require another async executor and
28
/// blocks the current thread while executing.
29
pub mod blocking;
30
mod replenishable;
31
/// Shared implementation of budget spending.
32
pub mod spend;
33

            
34
pub use replenishable::ReplenishableBudget;
35

            
36
#[derive(Debug)]
37
struct BudgetContext<Backing, Budget>
38
where
39
    Backing: Container<BudgetContextData<Budget>>,
40
{
41
    data: Backing,
42
    _budget: PhantomData<Budget>,
43
}
44

            
45
impl<Backing, Budget> Clone for BudgetContext<Backing, Budget>
46
where
47
    Backing: Container<BudgetContextData<Budget>>,
48
{
49
321
    fn clone(&self) -> Self {
50
321
        Self {
51
321
            data: self.data.cloned(),
52
321
            _budget: PhantomData,
53
321
        }
54
321
    }
55
}
56

            
57
#[derive(Debug)]
58
struct BudgetContextData<Budget> {
59
    budget: Budget,
60
    paused_future: Option<Waker>,
61
}
62

            
63
/// The result of a completed future.
64
#[derive(Debug)]
65
pub struct BudgetResult<T, Budget> {
66
    /// The output of the future.
67
    pub output: T,
68
    /// The budget after completing the future.
69
    pub balance: Budget,
70
}
71

            
72
impl<Backing, Budget> BudgetContext<Backing, Budget>
73
where
74
    Budget: Budgetable,
75
    Backing: Container<BudgetContextData<Budget>>,
76
{
77
    #[must_use]
78
    fn budget(&self) -> usize {
79
        self.data.map_locked(|data| data.budget.get())
80
    }
81

            
82
21342
    fn spend(&self, amount: usize) -> SpendBudget<'_, Backing, Budget> {
83
21342
        SpendBudget {
84
21342
            context: self,
85
21342
            amount,
86
21342
        }
87
21342
    }
88
}
89

            
90
/// A type that can be used as a budget.
91
///
92
/// Current implementors are:
93
///
94
/// - [`usize`]
95
/// - [`ReplenishableBudget`]
96
pub trait Budgetable: sealed::BudgetableSealed {}
97

            
98
mod sealed {
99
    pub trait BudgetableSealed: Clone + std::fmt::Debug + Unpin + 'static {
100
        fn get(&self) -> usize;
101
        fn spend(&mut self, amount: usize) -> bool;
102
        fn replenish(&mut self, amount: usize);
103
        fn add_waker(&self, waker: &std::task::Waker);
104
        fn remove_waker(&self, waker: &std::task::Waker);
105
        fn park_for_budget(&self);
106
    }
107
}
108

            
109
impl Budgetable for usize {}
110

            
111
impl BudgetableSealed for usize {
112
    fn spend(&mut self, amount: usize) -> bool {
113
134
        if let Some(remaining) = self.checked_sub(amount) {
114
58
            *self = remaining;
115
58
            true
116
        } else {
117
76
            false
118
        }
119
134
    }
120

            
121
    fn get(&self) -> usize {
122
        *self
123
    }
124

            
125
76
    fn replenish(&mut self, amount: usize) {
126
76
        *self = self.saturating_add(amount);
127
76
    }
128

            
129
91
    fn add_waker(&self, _waker: &std::task::Waker) {}
130

            
131
117
    fn remove_waker(&self, _waker: &std::task::Waker) {}
132

            
133
15
    fn park_for_budget(&self) {
134
15
        std::thread::park();
135
15
    }
136
}
137

            
138
trait Container<T>: Unpin + 'static {
139
    fn new(contained: T) -> Self;
140
    fn cloned(&self) -> Self;
141
    fn map_locked<R, F: FnOnce(&mut T) -> R>(&self, map: F) -> R;
142
}
143

            
144
#[derive(Debug)]
145
struct SyncContainer<T>(Arc<Mutex<T>>);
146

            
147
impl<T> Clone for SyncContainer<T> {
148
    fn clone(&self) -> Self {
149
        Self(self.0.clone())
150
    }
151
}
152

            
153
impl<T> Container<T> for SyncContainer<T>
154
where
155
    T: 'static,
156
{
157
1365368
    fn map_locked<R, F: FnOnce(&mut T) -> R>(&self, map: F) -> R {
158
1365368
        let mut locked = self.0.lock().unwrap();
159
1365368
        map(&mut locked)
160
1365368
    }
161

            
162
1059
    fn cloned(&self) -> Self {
163
1059
        Self(self.0.clone())
164
1059
    }
165

            
166
241
    fn new(contained: T) -> Self {
167
241
        Self(Arc::new(Mutex::new(contained)))
168
241
    }
169
}
170

            
171
#[derive(Debug)]
172
struct NotSyncContainer<T>(Rc<RefCell<T>>);
173

            
174
impl<T> Clone for NotSyncContainer<T> {
175
    fn clone(&self) -> Self {
176
        Self(self.0.clone())
177
    }
178
}
179

            
180
impl<T> Container<T> for NotSyncContainer<T>
181
where
182
    T: 'static,
183
{
184
112634
    fn map_locked<R, F: FnOnce(&mut T) -> R>(&self, map: F) -> R {
185
112634
        let mut locked = self.0.borrow_mut();
186
112634
        map(&mut *locked)
187
112634
    }
188

            
189
116
    fn cloned(&self) -> Self {
190
116
        Self(self.0.clone())
191
116
    }
192

            
193
13
    fn new(contained: T) -> Self {
194
13
        Self(Rc::new(RefCell::new(contained)))
195
13
    }
196
}
197

            
198
#[cfg(test)]
199
mod tests;