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//!
10//! The same checkpoints also observe *interrupt flags*: an embedder (e.g. the
11//! nREPL server) installs an `Arc<AtomicBool>` with [`InterruptGuard`] for the
12//! extent of an evaluation, and another thread sets it to stop that
13//! evaluation. A set flag makes every subsequent charge fail exactly like an
14//! exhausted budget, so the evaluation unwinds through the same uncatchable
15//! [`EvalError::GasExhausted`](crate::env::error::EvalError::GasExhausted)
16//! path in every tier (tree-walker, IR interpreter, JIT). Use
17//! [`interrupt_requested`] to tell an interrupt from real exhaustion. When no
18//! interrupt is pending the cost is one relaxed atomic load per charge.
19//!
20//! An interrupt is scoped to the evaluation, not to the thread: an async task
21//! is polled under the flags of the evaluation that spawned it
22//! ([`active_interrupts`], [`InterruptScope`]), so interrupting whichever
23//! evaluation is driving the executor leaves other evaluations' tasks alone.
24
25use std::cell::RefCell;
26use std::sync::Arc;
27use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
28
29/// A shareable, monotonically-decreasing execution-credit budget.
30#[derive(Debug)]
31pub struct GasMeter {
32    remaining: AtomicU64,
33}
34
35impl GasMeter {
36    pub fn new(credits: u64) -> Arc<Self> {
37        Arc::new(Self {
38            remaining: AtomicU64::new(credits),
39        })
40    }
41
42    pub fn remaining(&self) -> u64 {
43        self.remaining.load(Ordering::Relaxed)
44    }
45
46    /// Consume `cost` credits, returning false without partially charging when
47    /// the budget cannot cover the whole checkpoint.
48    pub fn charge(&self, cost: u64) -> bool {
49        self.remaining
50            .try_update(Ordering::Relaxed, Ordering::Relaxed, |remaining| {
51                remaining.checked_sub(cost)
52            })
53            .is_ok()
54    }
55}
56
57thread_local! {
58    static ACTIVE: RefCell<Vec<Arc<GasMeter>>> = const { RefCell::new(Vec::new()) };
59    static EXHAUSTED: RefCell<Vec<bool>> = const { RefCell::new(Vec::new()) };
60    static INTERRUPTS: RefCell<Vec<Arc<AtomicBool>>> = const { RefCell::new(Vec::new()) };
61}
62cljrs_gc::mutator_slot!("cljrs-runtime/GAS_ACTIVE", refcell ACTIVE: Vec<Arc<GasMeter>>);
63cljrs_gc::mutator_slot!("cljrs-runtime/GAS_EXHAUSTED", refcell EXHAUSTED: Vec<bool>);
64cljrs_gc::mutator_slot!("cljrs-runtime/GAS_INTERRUPTS", refcell INTERRUPTS: Vec<Arc<AtomicBool>>);
65
66/// Installs an interrupt flag for the dynamic extent of an evaluation on
67/// this thread. While installed, setting the flag (from any thread) makes the
68/// next execution-credit checkpoint fail, unwinding the evaluation as
69/// `GasExhausted`; [`interrupt_requested`] then reports true.
70#[must_use = "dropping InterruptGuard immediately uninstalls the interrupt flag"]
71pub struct InterruptGuard;
72
73impl InterruptGuard {
74    pub fn install(flag: Arc<AtomicBool>) -> Self {
75        INTERRUPTS.with(|flags| flags.borrow_mut().push(flag));
76        Self
77    }
78}
79
80impl Drop for InterruptGuard {
81    fn drop(&mut self) {
82        INTERRUPTS.with(|flags| {
83            flags.borrow_mut().pop();
84        });
85    }
86}
87
88/// True when an interrupt flag installed on this thread has been set.
89pub fn interrupt_requested() -> bool {
90    INTERRUPTS.with(|flags| {
91        flags
92            .borrow()
93            .iter()
94            .any(|flag| flag.load(Ordering::Relaxed))
95    })
96}
97
98/// Clone the interrupt flags installed on this thread for async task
99/// propagation, as [`active_meters`] does for gas meters.
100pub fn active_interrupts() -> Vec<Arc<AtomicBool>> {
101    INTERRUPTS.with(|flags| flags.borrow().clone())
102}
103
104/// Makes `flags` the only interrupt flags this thread observes, until dropped.
105///
106/// An interrupt belongs to one evaluation, but the flag stack is per thread,
107/// and every task on a `LocalSet` is polled on the thread of whichever
108/// evaluation happens to be driving it. A task polled inside an interrupted
109/// evaluation's extent would otherwise fail its next charge, though nobody
110/// interrupted the evaluation that spawned it. So a task is polled under the
111/// flags captured where it was spawned, replacing (not stacking on) the
112/// driver's, which are put back when the poll ends.
113#[must_use = "dropping InterruptScope immediately restores the previous interrupt flags"]
114pub struct InterruptScope {
115    outer: Vec<Arc<AtomicBool>>,
116}
117
118impl InterruptScope {
119    pub fn enter(flags: &[Arc<AtomicBool>]) -> Self {
120        let outer = INTERRUPTS.with(|installed| installed.replace(flags.to_vec()));
121        Self { outer }
122    }
123}
124
125impl Drop for InterruptScope {
126    fn drop(&mut self) {
127        INTERRUPTS.with(|installed| {
128            installed.replace(std::mem::take(&mut self.outer));
129        });
130    }
131}
132
133/// Installs a meter for the dynamic extent of an evaluation.
134#[must_use = "dropping GasGuard immediately uninstalls the active gas meter"]
135pub struct GasGuard;
136
137impl GasGuard {
138    pub fn install(meter: Arc<GasMeter>) -> Self {
139        ACTIVE.with(|active| active.borrow_mut().push(meter));
140        EXHAUSTED.with(|exhausted| exhausted.borrow_mut().push(false));
141        Self
142    }
143}
144
145impl Drop for GasGuard {
146    fn drop(&mut self) {
147        ACTIVE.with(|active| {
148            let mut active = active.borrow_mut();
149            active.pop();
150        });
151        EXHAUSTED.with(|exhausted| {
152            exhausted.borrow_mut().pop();
153        });
154    }
155}
156
157/// Charge the active evaluation, or succeed at no cost when unmetered.
158pub fn charge(cost: u64) -> bool {
159    if interrupt_requested() || is_exhausted() {
160        return false;
161    }
162    let charged = ACTIVE.with(|active| {
163        let active = active.borrow();
164        if active.is_empty() {
165            return true;
166        }
167        if active.iter().any(|meter| meter.remaining() < cost) {
168            return false;
169        }
170        active.iter().all(|meter| meter.charge(cost))
171    });
172    if !charged {
173        ACTIVE.with(|active| {
174            let active = active.borrow();
175            EXHAUSTED.with(|exhausted| {
176                for (index, meter) in active.iter().enumerate() {
177                    if meter.remaining() < cost {
178                        exhausted.borrow_mut()[index] = true;
179                    }
180                }
181            });
182        });
183    }
184    charged
185}
186
187/// Peek at the native-tier exhaustion signal set by a failed charge.
188pub fn is_exhausted() -> bool {
189    interrupt_requested()
190        || EXHAUSTED.with(|exhausted| exhausted.borrow().iter().any(|value| *value))
191}
192
193/// Clone the complete active meter stack for async task propagation.
194pub fn active_meters() -> Vec<Arc<GasMeter>> {
195    ACTIVE.with(|active| active.borrow().clone())
196}
197
198/// Install a captured meter stack in outer-to-inner order.
199pub fn install_meters(meters: &[Arc<GasMeter>]) -> Vec<GasGuard> {
200    meters.iter().cloned().map(GasGuard::install).collect()
201}
202
203/// Take the native-tier exhaustion signal set by a failed charge.
204///
205/// Prefer [`is_exhausted`] at dispatch boundaries; this remains available for
206/// tests and rare code that intentionally owns the current gas scope.
207pub fn take_exhausted() -> bool {
208    EXHAUSTED.with(|exhausted| {
209        let mut exhausted = exhausted.borrow_mut();
210        let was_exhausted = exhausted.iter().any(|value| *value);
211        exhausted.fill(false);
212        was_exhausted
213    })
214}
215
216#[cfg(test)]
217mod tests {
218    use super::*;
219
220    #[test]
221    fn scoped_meter_charges_without_partial_consumption() {
222        let meter = GasMeter::new(3);
223        let _guard = GasGuard::install(meter.clone());
224        assert!(charge(2));
225        assert!(!charge(2));
226        assert_eq!(meter.remaining(), 1);
227    }
228
229    #[test]
230    fn nested_meters_charge_outer_budget() {
231        let outer = GasMeter::new(3);
232        let _outer_guard = GasGuard::install(outer.clone());
233        let inner = GasMeter::new(2);
234        let _inner_guard = GasGuard::install(inner.clone());
235        assert!(charge(2));
236        assert_eq!(outer.remaining(), 1);
237        assert_eq!(inner.remaining(), 0);
238        assert!(!charge(1));
239    }
240
241    #[test]
242    fn inner_exhaustion_does_not_poison_healthy_outer_scope() {
243        let outer = GasMeter::new(10);
244        let _outer_guard = GasGuard::install(outer.clone());
245        {
246            let inner = GasMeter::new(0);
247            let _inner_guard = GasGuard::install(inner);
248            assert!(!charge(1));
249            assert!(is_exhausted());
250        }
251        assert!(!is_exhausted());
252        assert!(charge(1));
253        assert_eq!(outer.remaining(), 9);
254    }
255
256    #[test]
257    fn interrupt_flag_fails_charges_until_uninstalled() {
258        let flag = Arc::new(AtomicBool::new(false));
259        {
260            let _guard = InterruptGuard::install(flag.clone());
261            assert!(charge(1));
262            flag.store(true, Ordering::Relaxed);
263            assert!(!charge(1));
264            assert!(is_exhausted());
265            assert!(interrupt_requested());
266        }
267        assert!(!interrupt_requested());
268        assert!(!is_exhausted());
269        assert!(charge(1));
270    }
271
272    #[test]
273    fn interrupt_scope_replaces_the_installed_flags_and_restores_them() {
274        let driver = Arc::new(AtomicBool::new(true));
275        let _guard = InterruptGuard::install(driver.clone());
276        assert!(!charge(1));
277        {
278            // A task spawned by an evaluation nobody interrupted.
279            let _scope = InterruptScope::enter(&[]);
280            assert!(!interrupt_requested());
281            assert!(charge(1));
282        }
283        {
284            // A task spawned by the interrupted evaluation itself.
285            let _scope = InterruptScope::enter(&[driver]);
286            assert!(!charge(1));
287        }
288        assert!(interrupt_requested());
289        assert_eq!(active_interrupts().len(), 1);
290    }
291
292    proptest::proptest! {
293        /// Inside a scope the thread observes the scope's flags and no
294        /// others; once it is dropped, the flags installed before it,
295        /// unchanged. Holds for any flags, nested scopes included.
296        #[test]
297        fn interrupt_scope_observes_exactly_its_own_flags(
298            installed in proptest::collection::vec(proptest::bool::ANY, 0..4),
299            scoped in proptest::collection::vec(proptest::bool::ANY, 0..4),
300            nested in proptest::collection::vec(proptest::bool::ANY, 0..4),
301        ) {
302            use proptest::prop_assert_eq;
303            let flags = |set: &[bool]| -> Vec<Arc<AtomicBool>> {
304                set.iter().map(|b| Arc::new(AtomicBool::new(*b))).collect()
305            };
306            let any = |set: &[bool]| set.iter().any(|b| *b);
307
308            let _guards: Vec<InterruptGuard> = flags(&installed)
309                .into_iter()
310                .map(InterruptGuard::install)
311                .collect();
312            prop_assert_eq!(interrupt_requested(), any(&installed));
313            {
314                let _scope = InterruptScope::enter(&flags(&scoped));
315                prop_assert_eq!(interrupt_requested(), any(&scoped));
316                prop_assert_eq!(charge(1), !any(&scoped));
317                {
318                    let _nested = InterruptScope::enter(&flags(&nested));
319                    prop_assert_eq!(interrupt_requested(), any(&nested));
320                    prop_assert_eq!(active_interrupts().len(), nested.len());
321                }
322                prop_assert_eq!(interrupt_requested(), any(&scoped));
323                prop_assert_eq!(active_interrupts().len(), scoped.len());
324            }
325            prop_assert_eq!(interrupt_requested(), any(&installed));
326            prop_assert_eq!(active_interrupts().len(), installed.len());
327        }
328    }
329}
cljrs_runtime/env/gas.rs

cljrs_runtime/env/
gas.rs
