moirai_core/pool/thread_local.rs
1use crate::platform::*;
2
3/// Thread-local task pool for zero-allocation task execution
4pub struct ThreadLocalPool<T> {
5 /// Stack of available objects
6 pool: UnsafeCell<Vec<T>>,
7 /// Maximum pool size
8 max_size: usize,
9 /// Marker to ensure !Send and !Sync
10 _marker: PhantomData<*const T>,
11}
12
13impl<T> ThreadLocalPool<T> {
14 /// Create a new thread-local pool
15 #[must_use]
16 pub fn new(max_size: usize) -> Self {
17 Self {
18 pool: UnsafeCell::new(Vec::with_capacity(max_size)),
19 max_size,
20 _marker: PhantomData,
21 }
22 }
23
24 /// Get an object from the pool or create a new one
25 pub fn get_or_create<F>(&self, create: F) -> T
26 where
27 F: FnOnce() -> T,
28 {
29 // SAFETY: `ThreadLocalPool` is !Sync (and !Send) via its raw-pointer
30 // phantom marker, so `&self` implies same-thread access; the cell is
31 // therefore not aliased and unique `&mut` is sound.
32 unsafe {
33 let pool = &mut *self.pool.get();
34 pool.pop().unwrap_or_else(create)
35 }
36 }
37
38 /// Return an object to the pool
39 pub fn put(&self, obj: T) {
40 // SAFETY: same-thread uniqueness as in `get_or_create`; the cell is
41 // never aliased on a !Sync type.
42 unsafe {
43 let pool = &mut *self.pool.get();
44 if pool.len() < self.max_size {
45 pool.push(obj);
46 }
47 }
48 }
49}