xgrammar/support/
recursion_guard.rs

1//! Recursion-depth guard.
2//!
3//! A process-global maximum depth (configurable, env-overridable) bounds recursion in
4//! the parser and grammar functors; a per-thread counter tracks the live depth via an
5//! RAII [`RecursionGuard`]. Recursion is always within a single thread, so the counter is
6//! thread-local (each rayon worker gets its own).
7
8use std::{
9    cell::Cell,
10    sync::atomic::{AtomicI32, Ordering},
11};
12
13/// Default maximum recursion depth when unset and no env override is present.
14pub const DEFAULT_MAX_RECURSION_DEPTH: i32 = 10_000;
15
16/// Hard upper bound on any configured maximum.
17pub const MAX_REASONABLE_RECURSION_DEPTH: i32 = 1_000_000;
18
19const ENV_VAR: &str = "XGRAMMAR_MAX_RECURSION_DEPTH";
20
21// 0 is a sentinel meaning "not yet initialized".
22static MAX_RECURSION_DEPTH: AtomicI32 = AtomicI32::new(0);
23
24thread_local! {
25    static CURRENT_DEPTH: Cell<i32> = const { Cell::new(0) };
26}
27
28/// Error returned when an invalid maximum recursion depth is requested, or when the live
29/// depth exceeds the configured maximum.
30#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
31pub enum RecursionError {
32    /// A requested maximum was non-positive or exceeded [`MAX_REASONABLE_RECURSION_DEPTH`].
33    #[error(
34        "RecursionGuard: Maximum recursion depth must be positive and less than {MAX_REASONABLE_RECURSION_DEPTH}, \
35         got: {0}"
36    )]
37    InvalidMax(i32),
38    /// The live recursion depth exceeded the configured maximum.
39    #[error("RecursionGuard: Maximum recursion depth exceeded. Current depth: {current}, Max allowed: {max}")]
40    DepthExceeded {
41        /// The depth reached when the limit was hit.
42        current: i32,
43        /// The configured maximum.
44        max: i32,
45    },
46}
47
48fn load_max_from_env() -> i32 {
49    match std::env::var(ENV_VAR).ok().and_then(|v| v.parse::<i32>().ok()) {
50        Some(v) if v > 0 && v <= MAX_REASONABLE_RECURSION_DEPTH => v,
51        _ => DEFAULT_MAX_RECURSION_DEPTH,
52    }
53}
54
55/// Returns the configured maximum recursion depth, initializing it from the environment
56/// (or the default) on first access.
57#[must_use]
58pub fn get_max_recursion_depth() -> i32 {
59    let current = MAX_RECURSION_DEPTH.load(Ordering::Relaxed);
60    if current != 0 {
61        return current;
62    }
63    let initial = load_max_from_env();
64    MAX_RECURSION_DEPTH.store(initial, Ordering::Relaxed);
65    initial
66}
67
68/// Sets the maximum recursion depth.
69///
70/// # Errors
71/// Returns [`RecursionError::InvalidMax`] if `max_depth` is non-positive or exceeds
72/// [`MAX_REASONABLE_RECURSION_DEPTH`].
73pub fn set_max_recursion_depth(max_depth: i32) -> Result<(), RecursionError> {
74    if max_depth <= 0 || max_depth > MAX_REASONABLE_RECURSION_DEPTH {
75        return Err(RecursionError::InvalidMax(max_depth));
76    }
77    MAX_RECURSION_DEPTH.store(max_depth, Ordering::Relaxed);
78    Ok(())
79}
80
81/// Resets the calling thread's live recursion depth to zero.
82pub fn reset_recursion_depth() {
83    CURRENT_DEPTH.with(|d| d.set(0));
84}
85
86/// RAII guard that increments the calling thread's recursion depth on creation and
87/// decrements it on drop.
88#[derive(Debug)]
89#[must_use = "the guard must be held for the duration of the recursive call"]
90pub struct RecursionGuard {
91    // Not constructible outside `enter`; not `Send` (thread-local bound).
92    _private: (),
93}
94
95impl RecursionGuard {
96    /// Enters one level of recursion.
97    ///
98    /// # Errors
99    /// Returns [`RecursionError::DepthExceeded`] if entering would exceed the configured
100    /// maximum; the depth counter is left unchanged in that case.
101    pub fn enter() -> Result<Self, RecursionError> {
102        let max = get_max_recursion_depth();
103        let current = CURRENT_DEPTH.with(|d| {
104            let next = d.get() + 1;
105            d.set(next);
106            next
107        });
108        if current > max {
109            CURRENT_DEPTH.with(|d| d.set(d.get() - 1));
110            return Err(RecursionError::DepthExceeded {
111                current,
112                max,
113            });
114        }
115        Ok(Self {
116            _private: (),
117        })
118    }
119}
120
121impl Drop for RecursionGuard {
122    fn drop(&mut self) {
123        CURRENT_DEPTH.with(|d| d.set(d.get() - 1));
124    }
125}
xgrammar/support/recursion_guard.rs

xgrammar/support/
recursion_guard.rs
