harmos_storage/
capacity.rs

1/* ***
2 *
3 * Harmos
4 *
5 * SPDX-License-Identifier: Apache-2.0
6 *
7 * ***
8 */
9
10// Shared Storage Capacity: One Reservation Type And One Wait Primitive
11
12/* === Modules === */
13
14/* === Re-Exports === */
15
16/* === Imports === */
17
18// Standard Library Imports
19use std::collections::BTreeMap;
20use std::sync::atomic::{AtomicUsize, Ordering};
21use std::sync::{Arc, Mutex};
22
23// Third Party Imports
24use tokio::time::Instant;
25
26// Local Imports
27use super::arrow::RetainedMemory;
28use super::manager::{Inner, State};
29use crate::{Error, Result};
30
31/* === Constants === */
32
33/* === Macros === */
34
35/* === Type Aliases === */
36
37/* === Structs === */
38
39#[derive(Clone, Default)]
40pub(super) struct Held(Arc<AtomicUsize>);
41
42pub(super) struct Reservation {
43    inner: Arc<Inner>,
44    pool: Pool,
45    bytes: usize,
46    held: Held,
47}
48
49#[derive(Default)]
50pub(super) struct Growers {
51    next: u64,
52    waiting: BTreeMap<u64, Grower>,
53}
54
55struct Grower {
56    held: Held,
57    need: usize,
58}
59
60struct Waiting {
61    inner: Arc<Inner>,
62    pool: Pool,
63    id: u64,
64    active: bool,
65}
66
67/* === Enums === */
68
69#[derive(Clone, Copy, Debug, PartialEq, Eq)]
70pub(super) enum Pool {
71    Read,
72
73    Write,
74}
75
76#[derive(Debug)]
77pub(super) enum Refusal {
78    Exceeded(Pool),
79
80    Busy,
81
82    Cycle,
83
84    Elapsed,
85
86    Failed(Error),
87}
88
89pub(super) enum Attempt<T> {
90    Ready(T),
91    Wait,
92    Refuse(Refusal),
93}
94
95/* === Traits === */
96
97/* === Implementations === */
98
99impl Held {
100    pub(super) fn bytes(&self) -> usize {
101        self.0.load(Ordering::SeqCst)
102    }
103}
104
105impl Pool {
106    fn resource(self) -> &'static str {
107        match self {
108            Self::Read => "storage read memory",
109            Self::Write => "storage write memory",
110        }
111    }
112
113    // Owner Ceiling: Bytes One Owner Can Ever Hold In This Pool.
114    fn ceiling(self, inner: &Inner) -> usize {
115        match self {
116            Self::Read => inner.settings.budget.read_memory_bytes.get(),
117            Self::Write => inner.settings.request.max_pending_write_bytes.get() / 2,
118        }
119    }
120
121    // Charge Additional Bytes When They Fit Now.
122    fn charge(self, inner: &Inner, state: &mut State, bytes: usize) -> bool {
123        state.collect();
124        let fits = match self {
125            Self::Read => state.status.read_bytes.checked_add(bytes).is_some_and(|total| total <= self.ceiling(inner)),
126            Self::Write => {
127                let capacity = inner.settings.request.max_pending_write_bytes.get();
128                state.scratch_bytes.checked_add(bytes).is_some_and(|total| total <= capacity / 2)
129                    && state.status.pending_bytes.checked_add(bytes).is_some_and(|total| total <= capacity)
130            }
131        };
132        if !fits {
133            return false;
134        }
135        match self {
136            Self::Read => state.read_bytes += bytes,
137            Self::Write => state.scratch_bytes += bytes,
138        }
139        state.collect();
140        true
141    }
142
143    fn index(self) -> usize {
144        match self {
145            Self::Read => 0,
146            Self::Write => 1,
147        }
148    }
149
150    fn release(self, state: &mut State, bytes: usize) {
151        match self {
152            Self::Read => state.read_bytes -= bytes,
153            Self::Write => state.scratch_bytes -= bytes,
154        }
155        state.collect();
156    }
157}
158
159impl Refusal {
160    pub(super) fn into_error(self) -> Error {
161        match self {
162            Self::Exceeded(pool) => Error::ResourceExhausted { resource: pool.resource() },
163            Self::Busy => Error::ResourceExhausted { resource: "storage capacity" },
164            Self::Cycle => Error::ResourceExhausted { resource: "storage capacity cycle" },
165            Self::Elapsed => Error::ResourceExhausted { resource: "storage capacity deadline" },
166            Self::Failed(error) => error,
167        }
168    }
169}
170
171impl From<Refusal> for Error {
172    fn from(refusal: Refusal) -> Self {
173        refusal.into_error()
174    }
175}
176
177impl Reservation {
178    pub(super) fn empty(inner: &Arc<Inner>, pool: Pool, held: &Held) -> Self {
179        Self { inner: Arc::clone(inner), pool, bytes: 0, held: held.clone() }
180    }
181
182    pub(super) fn charged(inner: &Arc<Inner>, pool: Pool, bytes: usize) -> Self {
183        let held = Held::default();
184        held.0.fetch_add(bytes, Ordering::SeqCst);
185        Self { inner: Arc::clone(inner), pool, bytes, held }
186    }
187
188    pub(super) fn lifecycle(inner: &Arc<Inner>, pool: Pool, bytes: usize) -> Result<Self> {
189        let mut state = inner.state.lock().unwrap();
190        if bytes > pool.ceiling(inner) || !pool.charge(inner, &mut state, bytes) {
191            return Err(Error::ResourceExhausted { resource: pool.resource() });
192        }
193        drop(state);
194        let held = Held::default();
195        held.0.fetch_add(bytes, Ordering::SeqCst);
196        Ok(Self { inner: Arc::clone(inner), pool, bytes, held })
197    }
198
199    pub(super) fn try_new(inner: &Arc<Inner>, pool: Pool, held: &Held, bytes: usize) -> std::result::Result<Self, Refusal> {
200        let mut reservation = Self::empty(inner, pool, held);
201        reservation.try_grow(bytes)?;
202        Ok(reservation)
203    }
204
205    pub(super) async fn wait(
206        inner: &Arc<Inner>, pool: Pool, held: &Held, bytes: usize, deadline: Option<Instant>,
207    ) -> std::result::Result<Self, Refusal> {
208        let mut reservation = Self::empty(inner, pool, held);
209        reservation.grow(bytes, deadline).await?;
210        Ok(reservation)
211    }
212
213    pub(super) fn bytes(&self) -> usize {
214        self.bytes
215    }
216
217    pub(super) fn try_grow(&mut self, additional: usize) -> std::result::Result<(), Refusal> {
218        if additional == 0 {
219            return Ok(());
220        }
221        let inner = Arc::clone(&self.inner);
222        let mut state = inner.state.lock().unwrap();
223        match self.attempt(&mut state, additional) {
224            Attempt::Ready(()) => Ok(()),
225            Attempt::Wait => Err(Refusal::Busy),
226            Attempt::Refuse(refusal) => Err(refusal),
227        }
228    }
229
230    pub(super) async fn grow(&mut self, additional: usize, deadline: Option<Instant>) -> std::result::Result<(), Refusal> {
231        if additional == 0 {
232            return Ok(());
233        }
234        let inner = Arc::clone(&self.inner);
235        let mut waiting: Option<Waiting> = None;
236        let result = inner.wait(deadline, |state| self.grow_step(state, additional, &mut waiting)).await;
237        drop(waiting);
238        result
239    }
240
241    pub(super) fn shrink_to(&mut self, bytes: usize) {
242        if bytes >= self.bytes {
243            return;
244        }
245        let released = self.bytes - bytes;
246        self.bytes = bytes;
247        self.held.0.fetch_sub(released, Ordering::SeqCst);
248        let mut state = self.inner.state.lock().unwrap();
249        self.pool.release(&mut state, released);
250        drop(state);
251        self.inner.notify();
252    }
253
254    pub(super) fn disown(&mut self) -> usize {
255        let bytes = std::mem::take(&mut self.bytes);
256        self.held.0.fetch_sub(bytes, Ordering::SeqCst);
257        bytes
258    }
259
260    pub(super) fn split(&mut self, bytes: usize) -> Self {
261        assert!(bytes <= self.bytes, "split exceeds reservation");
262        self.bytes -= bytes;
263        Self { inner: Arc::clone(&self.inner), pool: self.pool, bytes, held: self.held.clone() }
264    }
265
266    pub(super) fn retain(self) -> RetainedMemory {
267        let owner = Arc::new(Mutex::new(self));
268        RetainedMemory::new(owner)
269    }
270
271    fn grow_step(&mut self, state: &mut State, additional: usize, waiting: &mut Option<Waiting>) -> Attempt<()> {
272        match self.attempt(state, additional) {
273            Attempt::Wait if waiting.is_none() && self.held.bytes() != 0 => {
274                match Waiting::join(&self.inner, state, self.pool, &self.held, additional) {
275                    Some(joined) => {
276                        *waiting = Some(joined);
277                        Attempt::Wait
278                    }
279                    None => Attempt::Refuse(Refusal::Cycle),
280                }
281            }
282            attempt => {
283                if !matches!(&attempt, Attempt::Wait)
284                    && let Some(waiting) = waiting
285                {
286                    waiting.leave(state);
287                }
288                attempt
289            }
290        }
291    }
292
293    fn attempt(&mut self, state: &mut State, additional: usize) -> Attempt<()> {
294        if let Err(error) = admitting(state, self.pool) {
295            return Attempt::Refuse(Refusal::Failed(error));
296        }
297        let owned = self.held.bytes();
298        if owned.checked_add(additional).is_none_or(|total| total > self.pool.ceiling(&self.inner)) {
299            return Attempt::Refuse(Refusal::Exceeded(self.pool));
300        }
301        if !self.pool.charge(&self.inner, state, additional) {
302            if self.pool == Pool::Write {
303                self.inner.relieve_pressure(state, additional, 0);
304            }
305            return Attempt::Wait;
306        }
307        self.bytes += additional;
308        self.held.0.fetch_add(additional, Ordering::SeqCst);
309        Attempt::Ready(())
310    }
311}
312
313impl Growers {
314    // Some Grower's Need Fits Once Every Owner That Is Not Growing Released Everything.
315    fn can_progress(&self, ceiling: usize, held: usize, need: usize) -> bool {
316        let waiters = self.waiting.values();
317        let held = waiters.clone().fold(held, |total, grower| total.saturating_add(grower.held.bytes()));
318        let free = ceiling.saturating_sub(held);
319        need <= free || waiters.into_iter().any(|grower| grower.need <= free)
320    }
321}
322
323impl Waiting {
324    fn join(inner: &Arc<Inner>, state: &mut State, pool: Pool, held: &Held, need: usize) -> Option<Self> {
325        let growers = &mut state.growers[pool.index()];
326        if !growers.can_progress(pool.ceiling(inner), held.bytes(), need) {
327            return None;
328        }
329        let id = growers.next;
330        growers.next += 1;
331        growers.waiting.insert(id, Grower { held: held.clone(), need });
332        Some(Self { inner: Arc::clone(inner), pool, id, active: true })
333    }
334
335    // Admission And Departure Share The State Lock: A Running Owner Must Never
336    // Be Counted As A Waiting Holder By Another Admission.
337    fn leave(&mut self, state: &mut State) {
338        state.growers[self.pool.index()].waiting.remove(&self.id);
339        self.active = false;
340    }
341}
342
343impl Drop for Waiting {
344    fn drop(&mut self) {
345        if !self.active {
346            return;
347        }
348        self.inner.state.lock().unwrap().growers[self.pool.index()].waiting.remove(&self.id);
349    }
350}
351
352impl Drop for Reservation {
353    fn drop(&mut self) {
354        if self.bytes == 0 {
355            return;
356        }
357        self.held.0.fetch_sub(self.bytes, Ordering::SeqCst);
358        let mut state = self.inner.state.lock().unwrap();
359        self.pool.release(&mut state, self.bytes);
360        drop(state);
361        self.inner.notify();
362    }
363}
364
365impl Inner {
366    pub(super) async fn wait<T>(
367        &self, deadline: Option<Instant>, mut attempt: impl FnMut(&mut State) -> Attempt<T>,
368    ) -> std::result::Result<T, Refusal> {
369        let mut changed = self.changed.subscribe();
370        loop {
371            changed.borrow_and_update();
372            match attempt(&mut self.state.lock().unwrap()) {
373                Attempt::Ready(value) => return Ok(value),
374                Attempt::Refuse(refusal) => return Err(refusal),
375                Attempt::Wait => {}
376            }
377            let notified = match deadline {
378                Some(deadline) => tokio::time::timeout_at(deadline, changed.changed()).await.map_err(|_| Refusal::Elapsed)?,
379                None => changed.changed().await,
380            };
381            notified.map_err(|_| Refusal::Failed(Error::Closing { domain: "storage" }))?;
382        }
383    }
384}
385
386impl Inner {
387    pub(super) async fn admit_when<T>(&self, mut check: impl FnMut(&mut State) -> Result<Option<T>>) -> Result<T> {
388        self.wait(None, |state| match check(state) {
389            Ok(Some(value)) => Attempt::Ready(value),
390            Ok(None) => Attempt::Wait,
391            Err(error) => Attempt::Refuse(Refusal::Failed(error)),
392        })
393        .await
394        .map_err(Refusal::into_error)
395    }
396}
397
398/* === Public Functions === */
399
400/* === Private Functions === */
401
402pub(super) fn admitting(state: &State, pool: Pool) -> Result<()> {
403    if state.status.closing {
404        return Err(Error::Closing { domain: "storage" });
405    }
406    if !state.ready || !state.status.running {
407        return Err(Error::NotReady { domain: "storage" });
408    }
409    if pool == Pool::Write
410        && let Some(source) = &state.status.failure
411    {
412        return Err(Error::Persistence { source: source.clone() });
413    }
414    Ok(())
415}
416
417/* === Tests === */
418
419#[cfg(test)]
420mod tests {
421    use super::*;
422    use crate::test_backend::{Data, Recovery, history, storage_settings};
423    use std::num::NonZeroUsize;
424
425    #[test]
426    fn cycle_decision_accounting() {
427        let mut settings = storage_settings();
428        settings.budget.read_memory_bytes = NonZeroUsize::new(100).unwrap();
429        settings.response.batch_bytes = NonZeroUsize::new(40).unwrap();
430        let storage = crate::StorageManager::<()>::new(
431            crate::StorageBackends { primary: Arc::new(Data::<()>::new(history())), state: Arc::new(Recovery::new(history())) },
432            settings,
433        )
434        .unwrap();
435        let inner = &storage.inner;
436        {
437            let mut state = inner.state.lock().unwrap();
438            state.ready = true;
439            state.status.running = true;
440        }
441        let mut a = Reservation::try_new(inner, Pool::Read, &Held::default(), 20).unwrap();
442        let mut b = Reservation::try_new(inner, Pool::Read, &Held::default(), 20).unwrap();
443        let external = Reservation::try_new(inner, Pool::Read, &Held::default(), 60).unwrap();
444        let mut waiting_a = None;
445        let mut waiting_b = None;
446        {
447            let mut state = inner.state.lock().unwrap();
448            assert!(matches!(a.grow_step(&mut state, 40, &mut waiting_a), Attempt::Wait));
449            assert!(matches!(b.grow_step(&mut state, 40, &mut waiting_b), Attempt::Wait));
450        }
451        drop(external);
452        let mut waiting_c = None;
453        let mut c = Reservation::empty(inner, Pool::Read, &Held::default());
454        {
455            let mut state = inner.state.lock().unwrap();
456            assert!(matches!(a.grow_step(&mut state, 40, &mut waiting_a), Attempt::Ready(())));
457            // A Has Not Dropped Its Wait Guard. Its Successful Admission Must
458            // Already Have Left The Ledger Before C Can Enter The Same Lock.
459            assert!(matches!(c.attempt(&mut state, 10), Attempt::Ready(())));
460            assert!(matches!(c.grow_step(&mut state, 15, &mut waiting_c), Attempt::Wait));
461        }
462        drop((waiting_a, waiting_b, waiting_c));
463        drop((a, b, c));
464        assert_eq!(storage.status().read_bytes, 0);
465    }
466}
harmos_storage/capacity.rs

harmos_storage/
capacity.rs
