Skip to main content

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}