1use std::io;
13use std::sync::atomic::{AtomicU64, Ordering};
14
15#[derive(Debug)]
22pub struct ByteBudget {
23 used: AtomicU64,
24 limit: u64,
25 label: String,
27}
28
29impl ByteBudget {
30 pub fn new(limit: u64) -> Self {
32 Self::labeled(limit, "memory")
33 }
34
35 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 pub fn limit(&self) -> u64 {
47 self.limit
48 }
49
50 pub fn used(&self) -> u64 {
52 self.used.load(Ordering::Acquire)
53 }
54
55 pub fn label(&self) -> &str {
57 &self.label
58 }
59
60 pub fn remaining(&self) -> u64 {
62 self.limit.saturating_sub(self.used())
63 }
64
65 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 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 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 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}