Skip to main content

cljrs_runtime/env/
gas.rs

1//! Cooperative execution-credit metering shared by every evaluation tier.
2//!
3//! Meter installation is thread-local; compiled async state machines explicitly
4//! capture and reinstall the active meter stack when they are spawned. Nested meters
5//! are charged together so inner evaluations cannot escape an outer budget.
6//! Native code reports exhaustion through per-scope sticky thread-local flags,
7//! allowing the signal to survive callback/JIT bridge boundaries without
8//! contaminating a healthy enclosing or subsequent scope.
9
10use std::cell::RefCell;
11use std::sync::Arc;
12use std::sync::atomic::{AtomicU64, Ordering};
13
14/// A shareable, monotonically-decreasing execution-credit budget.
15#[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    /// Consume `cost` credits, returning false without partially charging when
32    /// the budget cannot cover the whole checkpoint.
33    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/// Installs a meter for the dynamic extent of an evaluation.
48#[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
71/// Charge the active evaluation, or succeed at no cost when unmetered.
72pub 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
101/// Peek at the native-tier exhaustion signal set by a failed charge.
102pub fn is_exhausted() -> bool {
103    EXHAUSTED.with(|exhausted| exhausted.borrow().iter().any(|value| *value))
104}
105
106/// Clone the complete active meter stack for async task propagation.
107pub fn active_meters() -> Vec<Arc<GasMeter>> {
108    ACTIVE.with(|active| active.borrow().clone())
109}
110
111/// Install a captured meter stack in outer-to-inner order.
112pub fn install_meters(meters: &[Arc<GasMeter>]) -> Vec<GasGuard> {
113    meters.iter().cloned().map(GasGuard::install).collect()
114}
115
116/// Take the native-tier exhaustion signal set by a failed charge.
117///
118/// Prefer [`is_exhausted`] at dispatch boundaries; this remains available for
119/// tests and rare code that intentionally owns the current gas scope.
120pub 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}