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

cljrs_runtime/env/
gas.rs
