1use std::cell::RefCell;
11use std::sync::Arc;
12use std::sync::atomic::{AtomicU64, Ordering};
13
14#[derive(Debug)]
16pub struct GasMeter {
17 remaining: AtomicU64,
18}
19
20impl GasMeter {
21 pub fn new(credits: u64) -> Arc<Self> {
22 Arc::new(Self {
23 remaining: AtomicU64::new(credits),
24 })
25 }
26
27 pub fn remaining(&self) -> u64 {
28 self.remaining.load(Ordering::Relaxed)
29 }
30
31 pub fn charge(&self, cost: u64) -> bool {
34 self.remaining
35 .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |remaining| {
36 remaining.checked_sub(cost)
37 })
38 .is_ok()
39 }
40}
41
42thread_local! {
43 static ACTIVE: RefCell<Vec<Arc<GasMeter>>> = const { RefCell::new(Vec::new()) };
44 static EXHAUSTED: RefCell<Vec<bool>> = const { RefCell::new(Vec::new()) };
45}
46
47#[must_use = "dropping GasGuard immediately uninstalls the active gas meter"]
49pub struct GasGuard;
50
51impl GasGuard {
52 pub fn install(meter: Arc<GasMeter>) -> Self {
53 ACTIVE.with(|active| active.borrow_mut().push(meter));
54 EXHAUSTED.with(|exhausted| exhausted.borrow_mut().push(false));
55 Self
56 }
57}
58
59impl Drop for GasGuard {
60 fn drop(&mut self) {
61 ACTIVE.with(|active| {
62 let mut active = active.borrow_mut();
63 active.pop();
64 });
65 EXHAUSTED.with(|exhausted| {
66 exhausted.borrow_mut().pop();
67 });
68 }
69}
70
71pub fn charge(cost: u64) -> bool {
73 if is_exhausted() {
74 return false;
75 }
76 let charged = ACTIVE.with(|active| {
77 let active = active.borrow();
78 if active.is_empty() {
79 return true;
80 }
81 if active.iter().any(|meter| meter.remaining() < cost) {
82 return false;
83 }
84 active.iter().all(|meter| meter.charge(cost))
85 });
86 if !charged {
87 ACTIVE.with(|active| {
88 let active = active.borrow();
89 EXHAUSTED.with(|exhausted| {
90 for (index, meter) in active.iter().enumerate() {
91 if meter.remaining() < cost {
92 exhausted.borrow_mut()[index] = true;
93 }
94 }
95 });
96 });
97 }
98 charged
99}
100
101pub fn is_exhausted() -> bool {
103 EXHAUSTED.with(|exhausted| exhausted.borrow().iter().any(|value| *value))
104}
105
106pub fn active_meters() -> Vec<Arc<GasMeter>> {
108 ACTIVE.with(|active| active.borrow().clone())
109}
110
111pub fn install_meters(meters: &[Arc<GasMeter>]) -> Vec<GasGuard> {
113 meters.iter().cloned().map(GasGuard::install).collect()
114}
115
116pub fn take_exhausted() -> bool {
121 EXHAUSTED.with(|exhausted| {
122 let mut exhausted = exhausted.borrow_mut();
123 let was_exhausted = exhausted.iter().any(|value| *value);
124 exhausted.fill(false);
125 was_exhausted
126 })
127}
128
129#[cfg(test)]
130mod tests {
131 use super::*;
132
133 #[test]
134 fn scoped_meter_charges_without_partial_consumption() {
135 let meter = GasMeter::new(3);
136 let _guard = GasGuard::install(meter.clone());
137 assert!(charge(2));
138 assert!(!charge(2));
139 assert_eq!(meter.remaining(), 1);
140 }
141
142 #[test]
143 fn nested_meters_charge_outer_budget() {
144 let outer = GasMeter::new(3);
145 let _outer_guard = GasGuard::install(outer.clone());
146 let inner = GasMeter::new(2);
147 let _inner_guard = GasGuard::install(inner.clone());
148 assert!(charge(2));
149 assert_eq!(outer.remaining(), 1);
150 assert_eq!(inner.remaining(), 0);
151 assert!(!charge(1));
152 }
153
154 #[test]
155 fn inner_exhaustion_does_not_poison_healthy_outer_scope() {
156 let outer = GasMeter::new(10);
157 let _outer_guard = GasGuard::install(outer.clone());
158 {
159 let inner = GasMeter::new(0);
160 let _inner_guard = GasGuard::install(inner);
161 assert!(!charge(1));
162 assert!(is_exhausted());
163 }
164 assert!(!is_exhausted());
165 assert!(charge(1));
166 assert_eq!(outer.remaining(), 9);
167 }
168}