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

harmos_storage/
capacity.rs
