kernel/
budget.rs

1//! Shared reservation arithmetic. Currently the page arena uses this ledger;
2//! it is not a whole-process RSS cap. Constrained stores additionally bound
3//! WAL/record sizes and allocator bookkeeping counts through ResourceLimits.
4//! Exhaustion is fallible, including arithmetic overflow.
5
6use crate::{Error, Result};
7use std::sync::atomic::{AtomicUsize, Ordering};
8use std::sync::Arc;
9
10#[derive(Debug, Clone, Copy, PartialEq, Eq)]
11pub enum Class { Pool, Wal, Sort, Query }
12
13pub struct MemoryBudget { total: usize, used: AtomicUsize }
14
15impl MemoryBudget {
16    pub fn new(total: usize) -> Self { MemoryBudget { total, used: AtomicUsize::new(0) } }
17    pub fn used(&self) -> usize { self.used.load(Ordering::Relaxed) }
18    pub fn total(&self) -> usize { self.total }
19
20    pub fn reserve(self: &Arc<Self>, class: Class, bytes: usize) -> Result<Reservation> {
21        let mut cur = self.used.load(Ordering::Relaxed);
22        loop {
23            let next = cur.checked_add(bytes).filter(|n| *n <= self.total)
24                .ok_or(Error::OutOfBudget)?;
25            match self.used.compare_exchange_weak(
26                cur, next, Ordering::AcqRel, Ordering::Relaxed) {
27                Ok(_) => return Ok(Reservation { budget: self.clone(), bytes, class }),
28                Err(actual) => cur = actual,
29            }
30        }
31    }
32}
33
34/// `class` is not read today -- `Drop` only needs `bytes` to release the
35/// reservation -- but every `reserve` call already names one, so it is
36/// carried here rather than discarded, for a future per-class usage
37/// breakdown or a same-class-releases-what-it-reserved assertion.
38pub struct Reservation { budget: Arc<MemoryBudget>, bytes: usize, #[allow(dead_code)] class: Class }
39
40impl Drop for Reservation {
41    fn drop(&mut self) { self.budget.used.fetch_sub(self.bytes, Ordering::AcqRel); }
42}
43
44#[cfg(test)]
45mod tests {
46    use super::*;
47    use std::sync::Arc;
48
49    #[test]
50    fn reservations_are_returned_when_dropped() {
51        let b = Arc::new(MemoryBudget::new(1000));
52        { let _r = b.reserve(Class::Sort, 800).unwrap(); assert_eq!(b.used(), 800); }
53        assert_eq!(b.used(), 0);
54    }
55
56    #[test]
57    fn exhaustion_is_an_error_never_a_panic_and_never_a_kill() {
58        let b = Arc::new(MemoryBudget::new(1000));
59        let _a = b.reserve(Class::Sort, 900).unwrap();
60        assert!(matches!(b.reserve(Class::Query, 200), Err(crate::Error::OutOfBudget)));
61    }
62}
kernel/budget.rs

kernel/
budget.rs
