Skip to main content

kaish_vfs/
budget.rs

1//! Shared byte budget for memory-resident filesystems.
2//!
3//! Counting and limiting are different concerns (see
4//! `docs/kaish-overlayfs.md`, "Byte accounting"): every memory-resident
5//! filesystem counts its own bytes unconditionally via
6//! [`Filesystem::resident_bytes`](crate::Filesystem::resident_bytes); a
7//! `ByteBudget` is the optional, shared *limit*. One `Arc<ByteBudget>` handed
8//! to several filesystems (a kernel's `/` scratch, a workspace overlay's
9//! upper and bases) gives one number — and one loud failure mode — per
10//! consumer.
11
12use std::io;
13use std::sync::atomic::{AtomicU64, Ordering};
14
15/// A shared cap on memory-resident bytes.
16///
17/// Cloned via `Arc` into every filesystem that should draw from the same
18/// pool. Exceeding the limit fails the write loudly, ENOSPC-style — an
19/// in-band error a model reads and adapts to; fail loud over quietly eating
20/// RAM.
21#[derive(Debug)]
22pub struct ByteBudget {
23    used: AtomicU64,
24    limit: u64,
25    /// Names the budget in errors so the failure points at the knob.
26    label: String,
27}
28
29impl ByteBudget {
30    /// A budget of `limit` bytes labeled "memory" in error messages.
31    pub fn new(limit: u64) -> Self {
32        Self::labeled(limit, "memory")
33    }
34
35    /// A budget whose errors name `label` (e.g. a config knob like
36    /// `vfs-budget`), so the failure tells the reader what to raise.
37    pub fn labeled(limit: u64, label: impl Into<String>) -> Self {
38        Self {
39            used: AtomicU64::new(0),
40            limit,
41            label: label.into(),
42        }
43    }
44
45    /// The configured cap in bytes.
46    pub fn limit(&self) -> u64 {
47        self.limit
48    }
49
50    /// Bytes currently charged against the budget.
51    pub fn used(&self) -> u64 {
52        self.used.load(Ordering::Acquire)
53    }
54
55    /// The label used in error messages and introspection.
56    pub fn label(&self) -> &str {
57        &self.label
58    }
59
60    /// Bytes still available.
61    pub fn remaining(&self) -> u64 {
62        self.limit.saturating_sub(self.used())
63    }
64
65    /// Reserve `bytes` against the budget, or fail loudly without charging.
66    pub fn try_charge(&self, bytes: u64) -> io::Result<()> {
67        let exhausted = |current: u64| {
68            io::Error::new(
69                io::ErrorKind::StorageFull,
70                format!(
71                    "{} budget exhausted: {} bytes used + {} requested exceeds the {} byte limit",
72                    self.label, current, bytes, self.limit
73                ),
74            )
75        };
76        let mut current = self.used.load(Ordering::Relaxed);
77        loop {
78            // An arithmetic overflow is past any limit by definition.
79            let next = current.checked_add(bytes).ok_or_else(|| exhausted(current))?;
80            if next > self.limit {
81                return Err(exhausted(current));
82            }
83            match self
84                .used
85                .compare_exchange_weak(current, next, Ordering::AcqRel, Ordering::Relaxed)
86            {
87                Ok(_) => return Ok(()),
88                Err(actual) => current = actual,
89            }
90        }
91    }
92
93    /// Return `bytes` to the budget.
94    ///
95    /// Panics if it would underflow: crediting more than was charged means
96    /// the caller's accounting is wrong, and a limiter that under-reports is
97    /// worse than no limiter — crash over drift.
98    pub fn credit(&self, bytes: u64) {
99        let previous = self.used.fetch_sub(bytes, Ordering::AcqRel);
100        assert!(
101            previous >= bytes,
102            "ByteBudget '{}' underflow: credited {} bytes with only {} charged — accounting bug",
103            self.label,
104            bytes,
105            previous
106        );
107    }
108}
109
110#[cfg(test)]
111mod tests {
112    use super::*;
113
114    #[test]
115    fn test_charge_and_credit() {
116        let budget = ByteBudget::new(100);
117        budget.try_charge(60).unwrap();
118        assert_eq!(budget.used(), 60);
119        assert_eq!(budget.remaining(), 40);
120        budget.credit(20);
121        assert_eq!(budget.used(), 40);
122    }
123
124    #[test]
125    fn test_exceeding_fails_loudly_without_charging() {
126        let budget = ByteBudget::labeled(100, "vfs-budget");
127        budget.try_charge(90).unwrap();
128        let error = budget.try_charge(11).unwrap_err();
129        assert_eq!(error.kind(), io::ErrorKind::StorageFull);
130        assert!(error.to_string().contains("vfs-budget"));
131        assert!(error.to_string().contains("100"));
132        // The failed charge left the budget untouched.
133        assert_eq!(budget.used(), 90);
134        budget.try_charge(10).unwrap();
135    }
136
137    #[test]
138    fn test_overflow_treated_as_full() {
139        let budget = ByteBudget::new(u64::MAX);
140        budget.try_charge(u64::MAX - 1).unwrap();
141        let error = budget.try_charge(u64::MAX).unwrap_err();
142        assert_eq!(error.kind(), io::ErrorKind::StorageFull);
143    }
144
145    #[test]
146    #[should_panic(expected = "underflow")]
147    fn test_credit_underflow_panics() {
148        let budget = ByteBudget::new(100);
149        budget.try_charge(10).unwrap();
150        budget.credit(11);
151    }
152}