Skip to main content

seq_core/
arena.rs

1//! Arena Allocator - Thread-local bump allocation for Values
2//!
3//! Uses bumpalo for fast bump allocation of Strings and Variants.
4//! Each OS thread has an arena that's used by strands executing on it.
5//!
6//! # Design
7//! - Thread-local Bump allocator
8//! - Fast allocation (~5ns vs ~100ns for malloc)
9//! - Periodic reset to prevent unbounded growth
10//! - Manual reset when strand completes
11//!
12//! # ⚠️ IMPORTANT: Thread-Local, Not Strand-Local
13//!
14//! The arena is **thread-local**, not strand-local. This has implications
15//! if May's scheduler migrates a strand to a different thread:
16//!
17//! **What happens:**
18//! - Strand starts on Thread A, allocates strings from Arena A
19//! - May migrates strand to Thread B (rare, but possible)
20//! - Strand now allocates from Arena B
21//! - When strand exits, Arena B is reset (not Arena A)
22//! - Arena A still contains the strings from earlier allocation
23//!
24//! **Why this is safe:**
25//! - Arena reset only happens on strand exit (see `scheduler.rs:512-525`)
26//! - A migrated strand continues executing, doesn't trigger reset of Arena A
27//! - Arena A will be reset when the *next* strand on Thread A exits
28//! - Channel sends clone to global allocator (see `seqstring.rs:123-132`)
29//!
30//! **Performance impact:**
31//! - Minimal in practice - May rarely migrates strands
32//! - If migration occurs, some memory stays in old arena until next reset
33//! - Auto-reset at 10MB threshold prevents unbounded growth
34//!
35//! **Alternative considered:**
36//! Strand-local arenas would require passing arena pointer with every
37//! strand migration. This adds complexity and overhead for a rare case.
38//! Thread-local is simpler and faster for the common case.
39//!
40//! See: `docs/ARENA_ALLOCATION_DESIGN.md` for full design rationale.
41
42use crate::memory_stats::{get_or_register_slot, update_arena_stats};
43use bumpalo::Bump;
44use std::cell::RefCell;
45
46/// Configuration for the arena
47const ARENA_RESET_THRESHOLD: usize = 10 * 1024 * 1024; // 10MB - reset when we exceed this
48
49// Thread-local arena for value allocations
50thread_local! {
51    static ARENA: RefCell<Bump> = {
52        // Register thread with memory stats registry once during initialization
53        get_or_register_slot();
54        RefCell::new(Bump::new())
55    };
56    static ARENA_BYTES_ALLOCATED: RefCell<usize> = const { RefCell::new(0) };
57}
58
59/// Execute a function with access to the thread-local arena.
60///
61/// The `&Bump` is live only for the duration of the closure, so the closure
62/// must consume any borrowed `&str` it allocates. To retain an allocated
63/// string past the call, wrap it with `SeqString` via
64/// `seqstring::arena_string` (the canonical user-facing entry point).
65///
66/// # Performance
67/// ~5ns vs ~100ns for global allocator (20x faster).
68pub fn with_arena<F, R>(f: F) -> R
69where
70    F: FnOnce(&Bump) -> R,
71{
72    // Thread registration happens once during ARENA initialization,
73    // not on every arena access (keeping the fast path fast).
74    ARENA.with(|arena| {
75        let bump = arena.borrow();
76        let result = f(&bump);
77
78        // Track allocation size
79        let allocated = bump.allocated_bytes();
80        drop(bump); // Drop borrow before accessing ARENA_BYTES_ALLOCATED
81
82        ARENA_BYTES_ALLOCATED.with(|bytes| {
83            *bytes.borrow_mut() = allocated;
84        });
85
86        // Update cross-thread memory stats registry
87        update_arena_stats(allocated);
88
89        // Auto-reset if threshold exceeded
90        if should_reset() {
91            arena_reset();
92        }
93
94        result
95    })
96}
97
98/// Reset the thread-local arena
99///
100/// Call this when a strand completes to free memory.
101/// Also called automatically when arena exceeds threshold.
102pub fn arena_reset() {
103    ARENA.with(|arena| {
104        arena.borrow_mut().reset();
105    });
106    ARENA_BYTES_ALLOCATED.with(|bytes| {
107        *bytes.borrow_mut() = 0;
108    });
109    // Update cross-thread memory stats registry
110    update_arena_stats(0);
111}
112
113/// Check if arena should be reset (exceeded threshold)
114fn should_reset() -> bool {
115    ARENA_BYTES_ALLOCATED.with(|bytes| *bytes.borrow() > ARENA_RESET_THRESHOLD)
116}
117
118/// Get current arena statistics
119pub fn arena_stats() -> ArenaStats {
120    // Read from our tracked bytes instead of Bump's internal state
121    // This ensures consistency with arena_reset() which sets ARENA_BYTES_ALLOCATED to 0
122    let allocated = ARENA_BYTES_ALLOCATED.with(|bytes| *bytes.borrow());
123    ArenaStats {
124        allocated_bytes: allocated,
125    }
126}
127
128/// Arena statistics for debugging/monitoring
129#[derive(Debug, Clone, Copy)]
130pub struct ArenaStats {
131    pub allocated_bytes: usize,
132}
133
134#[cfg(test)]
135mod tests {
136    use super::*;
137
138    #[test]
139    fn test_arena_reset() {
140        arena_reset(); // Start fresh
141
142        // Allocate some strings via with_arena
143        with_arena(|arena| {
144            let _s1 = arena.alloc_str("Hello");
145            let _s2 = arena.alloc_str("World");
146        });
147
148        let stats_before = arena_stats();
149        assert!(stats_before.allocated_bytes > 0);
150
151        // Reset arena
152        arena_reset();
153
154        let stats_after = arena_stats();
155        // After reset, allocated bytes should be much less than before
156        // (Bump might keep some internal overhead, so we don't assert == 0)
157        assert!(
158            stats_after.allocated_bytes < stats_before.allocated_bytes,
159            "Arena should have less memory after reset (before: {}, after: {})",
160            stats_before.allocated_bytes,
161            stats_after.allocated_bytes
162        );
163    }
164
165    #[test]
166    fn test_with_arena() {
167        arena_reset(); // Start fresh
168
169        // We can't return the &str from the closure (lifetime issue)
170        // Instead, test that allocation works and stats update
171        let len = with_arena(|arena| {
172            let s = arena.alloc_str("Test string");
173            assert_eq!(s, "Test string");
174            s.len()
175        });
176
177        assert_eq!(len, 11);
178
179        let stats = arena_stats();
180        assert!(stats.allocated_bytes > 0);
181    }
182
183    #[test]
184    fn test_auto_reset_threshold() {
185        arena_reset(); // Start fresh
186
187        // Allocate just under threshold
188        let big_str = "x".repeat(ARENA_RESET_THRESHOLD / 2);
189        with_arena(|arena| {
190            let _s = arena.alloc_str(&big_str);
191        });
192
193        let stats1 = arena_stats();
194        let initial_bytes = stats1.allocated_bytes;
195        assert!(initial_bytes > 0);
196
197        // Allocate more to exceed threshold - this should trigger auto-reset
198        let big_str2 = "y".repeat(ARENA_RESET_THRESHOLD / 2 + 1000);
199        with_arena(|arena| {
200            let _s = arena.alloc_str(&big_str2);
201        });
202
203        // Arena should have been reset and re-allocated with just the second string
204        let stats2 = arena_stats();
205        // After reset, we should only have the second allocation
206        // (which is slightly larger than ARENA_RESET_THRESHOLD / 2)
207        assert!(
208            stats2.allocated_bytes < initial_bytes + (ARENA_RESET_THRESHOLD / 2 + 2000),
209            "Arena should have reset: stats2={}, initial={}, threshold={}",
210            stats2.allocated_bytes,
211            initial_bytes,
212            ARENA_RESET_THRESHOLD
213        );
214    }
215}