Skip to main content

ferroui_base/threading/
thread_safe_object_pool.rs

1use std::any::{Any, TypeId};
2use std::collections::HashMap;
3use std::sync::{Arc, Mutex, PoisonError};
4
5/// A pool of reusable objects that can be shared between threads.
6pub struct ThreadSafeObjectPool<T> {
7    stack: Mutex<Vec<T>>,
8}
9
10impl<T: Default> Default for ThreadSafeObjectPool<T> {
11    fn default() -> Self {
12        Self::new()
13    }
14}
15
16impl<T: Default> ThreadSafeObjectPool<T> {
17    pub fn new() -> Self {
18        Self { stack: Mutex::new(Vec::new()) }
19    }
20
21    /// Takes an object out of the pool, creating one when the pool is empty.
22    pub fn get(&self) -> T {
23        let pooled = self.stack.lock().unwrap_or_else(PoisonError::into_inner).pop();
24        pooled.unwrap_or_default()
25    }
26
27    /// Puts the object back into the pool and clears the slot it came from.
28    pub fn return_and_set_null(&self, obj: &mut Option<T>) {
29        let Some(obj) = obj.take() else {
30            return;
31        };
32        self.stack.lock().unwrap_or_else(PoisonError::into_inner).push(obj);
33    }
34}
35
36static DEFAULT_POOLS: Mutex<Option<HashMap<TypeId, Arc<dyn Any + Send + Sync>>>> = Mutex::new(None);
37
38impl<T: Default + Send + 'static> ThreadSafeObjectPool<T> {
39    /// The process-wide pool for `T`.
40    pub fn default_pool() -> Arc<ThreadSafeObjectPool<T>> {
41        let mut pools = DEFAULT_POOLS.lock().unwrap_or_else(PoisonError::into_inner);
42        let pool = pools
43            .get_or_insert_with(HashMap::new)
44            .entry(TypeId::of::<T>())
45            .or_insert_with(|| Arc::new(ThreadSafeObjectPool::<T>::new()))
46            .clone();
47        drop(pools);
48        match pool.downcast::<ThreadSafeObjectPool<T>>() {
49            Ok(pool) => pool,
50            Err(_) => unreachable!("the pool is keyed by its element type"),
51        }
52    }
53}
54
55#[cfg(test)]
56mod tests {
57    use super::*;
58
59    #[derive(Default, Debug, PartialEq)]
60    struct Buffer(Vec<u8>);
61
62    #[test]
63    fn get_creates_and_reuses() {
64        let pool = ThreadSafeObjectPool::<Buffer>::new();
65        let mut first = Some(pool.get());
66        assert_eq!(first, Some(Buffer(Vec::new())));
67        first.as_mut().unwrap().0.push(7);
68
69        pool.return_and_set_null(&mut first);
70        assert!(first.is_none());
71        // Returning an empty slot does nothing.
72        pool.return_and_set_null(&mut first);
73
74        assert_eq!(pool.get(), Buffer(vec![7]));
75        assert_eq!(pool.get(), Buffer(Vec::new()));
76    }
77
78    #[test]
79    fn default_pool_is_shared_per_type_across_threads() {
80        #[derive(Default)]
81        struct Marker(u32);
82
83        let pool = ThreadSafeObjectPool::<Marker>::default_pool();
84        let mut slot = Some(Marker(42));
85        pool.return_and_set_null(&mut slot);
86
87        let value = std::thread::spawn(|| ThreadSafeObjectPool::<Marker>::default_pool().get().0).join().unwrap();
88        assert_eq!(value, 42);
89        assert!(Arc::ptr_eq(&pool, &ThreadSafeObjectPool::<Marker>::default_pool()));
90    }
91}