oxc_allocator/
pool.rs

1use std::{iter, mem::ManuallyDrop, ops::Deref, sync::Mutex};
2
3use crate::Allocator;
4
5/// A thread-safe pool for reusing [`Allocator`] instances to reduce allocation overhead.
6///
7/// Internally uses a `Vec` protected by a `Mutex` to store available allocators.
8pub struct AllocatorPool {
9    allocators: Mutex<Vec<Allocator>>,
10}
11
12impl AllocatorPool {
13    /// Creates a new [`AllocatorPool`] for use across the specified number of threads.
14    pub fn new(thread_count: usize) -> AllocatorPool {
15        let allocators = iter::repeat_with(Allocator::new).take(thread_count).collect();
16        AllocatorPool { allocators: Mutex::new(allocators) }
17    }
18
19    /// Retrieves an [`Allocator`] from the pool, or creates a new one if the pool is empty.
20    ///
21    /// Returns an [`AllocatorGuard`] that gives access to the allocator.
22    ///
23    /// # Panics
24    ///
25    /// Panics if the underlying mutex is poisoned.
26    pub fn get(&self) -> AllocatorGuard<'_> {
27        let allocator = {
28            let mut allocators = self.allocators.lock().unwrap();
29            allocators.pop()
30        };
31        let allocator = allocator.unwrap_or_else(Allocator::new);
32
33        AllocatorGuard { allocator: ManuallyDrop::new(allocator), pool: self }
34    }
35
36    /// Add an [`Allocator`] to the pool.
37    ///
38    /// The `Allocator` should be empty, ready to be re-used.
39    ///
40    /// # Panics
41    ///
42    /// Panics if the underlying mutex is poisoned.
43    fn add(&self, allocator: Allocator) {
44        let mut allocators = self.allocators.lock().unwrap();
45        allocators.push(allocator);
46    }
47}
48
49/// A guard object representing exclusive access to an [`Allocator`] from the pool.
50///
51/// On drop, the `Allocator` is reset and returned to the pool.
52pub struct AllocatorGuard<'alloc_pool> {
53    allocator: ManuallyDrop<Allocator>,
54    pool: &'alloc_pool AllocatorPool,
55}
56
57impl Deref for AllocatorGuard<'_> {
58    type Target = Allocator;
59
60    fn deref(&self) -> &Self::Target {
61        &self.allocator
62    }
63}
64
65impl Drop for AllocatorGuard<'_> {
66    /// Return [`Allocator`] back to the pool.
67    fn drop(&mut self) {
68        // SAFETY: After taking ownership of the `Allocator`, we do not touch the `ManuallyDrop` again
69        let mut allocator = unsafe { ManuallyDrop::take(&mut self.allocator) };
70        allocator.reset();
71        self.pool.add(allocator);
72    }
73}
oxc_allocator/pool.rs

oxc_allocator/
pool.rs
