From c58b1c4abb8f6c1377611ed0593737d732092381 Mon Sep 17 00:00:00 2001 From: Cameron Clark Date: Wed, 17 Dec 2025 17:28:35 +0000 Subject: [PATCH] refactor(allocator): extract CapacityLimit struct for pool capacity limiting Extract Condvar and max_count logic into a separate CapacityLimit struct. This ensures the Condvar is only allocated when max_count is Some, and makes the capacity limiting logic reusable and generic. --- crates/oxc_allocator/src/pool/fixed_size.rs | 108 ++++++++++++++++---- crates/oxc_allocator/src/pool/mod.rs | 12 ++- crates/oxc_linter/src/service/runtime.rs | 2 +- 3 files changed, 100 insertions(+), 22 deletions(-) diff --git a/crates/oxc_allocator/src/pool/fixed_size.rs b/crates/oxc_allocator/src/pool/fixed_size.rs index 0d6be9f9cc773..715a6de90e396 100644 --- a/crates/oxc_allocator/src/pool/fixed_size.rs +++ b/crates/oxc_allocator/src/pool/fixed_size.rs @@ -3,7 +3,7 @@ use std::{ mem::{self, ManuallyDrop}, ptr::NonNull, sync::{ - Mutex, + Condvar, Mutex, atomic::{AtomicBool, AtomicU32, Ordering}, }, }; @@ -18,6 +18,43 @@ use crate::{ const TWO_GIB: usize = 1 << 31; const FOUR_GIB: usize = 1 << 32; +/// Capacity limiter for a pool that blocks when the maximum capacity is reached. +struct CapacityLimit { + /// Maximum number of items that can be created + max_count: u32, + /// Condition variable to signal when an item is returned to the pool + available: Condvar, +} + +impl CapacityLimit { + /// Create a new [`CapacityLimit`] with the given maximum count. + fn new(max_count: u32) -> Self { + Self { max_count, available: Condvar::new() } + } + + /// Check if the pool is at capacity. + fn is_at_capacity(&self, current_count: u32) -> bool { + current_count >= self.max_count + } + + /// Wait for an item to become available and return it. + /// + /// This method blocks until an item is returned to the pool. + fn wait_for_item(&self, mut items: std::sync::MutexGuard<'_, Vec>) -> T { + loop { + items = self.available.wait(items).unwrap(); + if let Some(item) = items.pop() { + return item; + } + } + } + + /// Notify one waiting thread that an item is available. + fn notify_available(&self) { + self.available.notify_one(); + } +} + /// A thread-safe pool for reusing [`Allocator`] instances, that uses fixed-size allocators, /// suitable for use with raw transfer. pub struct FixedSizeAllocatorPool { @@ -25,36 +62,36 @@ pub struct FixedSizeAllocatorPool { allocators: Mutex>, /// ID to assign to next `Allocator` that's created next_id: AtomicU32, + /// Capacity limiter. `None` means no limit (default). + capacity_limit: Option, } impl FixedSizeAllocatorPool { /// Create a new [`FixedSizeAllocatorPool`] for use across the specified number of threads. - pub fn new(thread_count: usize) -> FixedSizeAllocatorPool { + /// + /// If `max_count` is `Some`, the pool will block when trying to get an allocator + /// if the maximum number of allocators has been reached, waiting until one is returned. + /// If `max_count` is `None`, there is no limit on the number of allocators. + pub fn new(thread_count: usize, max_count: Option) -> FixedSizeAllocatorPool { // Each allocator consumes a large block of memory, so create them on demand instead of upfront, // in case not all threads end up being used (e.g. language server without `import` plugin) let allocators = Vec::with_capacity(thread_count); - FixedSizeAllocatorPool { allocators: Mutex::new(allocators), next_id: AtomicU32::new(0) } + FixedSizeAllocatorPool { + allocators: Mutex::new(allocators), + next_id: AtomicU32::new(0), + capacity_limit: max_count.map(CapacityLimit::new), + } } /// Retrieve an [`Allocator`] from the pool, or create a new one if the pool is empty. /// + /// If `max_count` was set and the maximum number of allocators has been reached, + /// this method will block until an allocator is returned to the pool. + /// /// # Panics /// Panics if the underlying mutex is poisoned. pub fn get(&self) -> Allocator { - let fixed_size_allocator = { - let mut allocators = self.allocators.lock().unwrap(); - allocators.pop() - }; - - let fixed_size_allocator = fixed_size_allocator.unwrap_or_else(|| { - // Each allocator needs to have a unique ID, but the order those IDs are assigned in - // doesn't matter, so `Ordering::Relaxed` is fine - let id = self.next_id.fetch_add(1, Ordering::Relaxed); - // Protect against IDs wrapping around. - // TODO: Does this work? Do we need it anyway? - assert!(id < u32::MAX, "Created too many allocators"); - FixedSizeAllocator::new(id) - }); + let fixed_size_allocator = self.get_impl(); // Unwrap `FixedSizeAllocator`. // `add` method will wrap it again, before returning it to pool, ensuring it gets dropped properly. @@ -66,6 +103,32 @@ impl FixedSizeAllocatorPool { ManuallyDrop::into_inner(allocator) } + fn get_impl(&self) -> FixedSizeAllocator { + let mut allocators = self.allocators.lock().unwrap(); + + // Try to get an allocator from the pool + if let Some(alloc) = allocators.pop() { + return alloc; + } + + // Check if we're at maximum capacity + if let Some(capacity_limit) = &self.capacity_limit { + let current_count = self.next_id.load(Ordering::Relaxed); + if capacity_limit.is_at_capacity(current_count) { + // At maximum capacity - wait for an item to be returned + return capacity_limit.wait_for_item(allocators); + } + } + + // Create a new allocator. + // Each allocator needs to have a unique ID, but the order those IDs are assigned in + // doesn't matter, so `Ordering::Relaxed` is fine. + let id = self.next_id.fetch_add(1, Ordering::Relaxed); + // Protect against IDs wrapping around. + assert!(id < u32::MAX, "Created too many allocators"); + FixedSizeAllocator::new(id) + } + /// Add an [`Allocator`] to the pool. /// /// The `Allocator` is reset by this method, so it's ready to be re-used. @@ -80,8 +143,15 @@ impl FixedSizeAllocatorPool { FixedSizeAllocator { allocator: ManuallyDrop::new(allocator) }; fixed_size_allocator.reset(); - let mut allocators = self.allocators.lock().unwrap(); - allocators.push(fixed_size_allocator); + { + let mut allocators = self.allocators.lock().unwrap(); + allocators.push(fixed_size_allocator); + } + + // Notify one waiting thread that an allocator is available (if capacity is limited) + if let Some(capacity_limit) = &self.capacity_limit { + capacity_limit.notify_available(); + } } } diff --git a/crates/oxc_allocator/src/pool/mod.rs b/crates/oxc_allocator/src/pool/mod.rs index 5c558da03de39..8b353a89623a5 100644 --- a/crates/oxc_allocator/src/pool/mod.rs +++ b/crates/oxc_allocator/src/pool/mod.rs @@ -46,16 +46,24 @@ impl AllocatorPool { /// Create a new [`AllocatorPool`] for use across the specified number of threads, /// which uses fixed-size allocators (suitable for raw transfer). + /// + /// If `max_count` is `Some`, the pool will block when trying to get an allocator + /// if the maximum number of allocators has been reached, waiting until one is returned. + /// If `max_count` is `None`, there is no limit on the number of allocators. #[cfg(feature = "fixed_size")] - pub fn new_fixed_size(thread_count: usize) -> AllocatorPool { + pub fn new_fixed_size(thread_count: usize, max_count: Option) -> AllocatorPool { #[cfg(all(target_pointer_width = "64", target_endian = "little"))] { - Self(AllocatorPoolInner::FixedSize(FixedSizeAllocatorPool::new(thread_count))) + Self(AllocatorPoolInner::FixedSize(FixedSizeAllocatorPool::new( + thread_count, + max_count, + ))) } #[cfg(not(all(target_pointer_width = "64", target_endian = "little")))] { let _thread_count = thread_count; // Avoid unused vars lint warning + let _max_count = max_count; // Avoid unused vars lint warning panic!("Fixed size allocators are only supported on 64-bit little-endian platforms"); } } diff --git a/crates/oxc_linter/src/service/runtime.rs b/crates/oxc_linter/src/service/runtime.rs index 8946ff9065035..ddd52481a9233 100644 --- a/crates/oxc_linter/src/service/runtime.rs +++ b/crates/oxc_linter/src/service/runtime.rs @@ -212,7 +212,7 @@ impl Runtime { // If an external linter is used (JS plugins), we must use fixed-size allocators, // for compatibility with raw transfer let allocator_pool = if linter.has_external_linter() { - AllocatorPool::new_fixed_size(thread_count) + AllocatorPool::new_fixed_size(thread_count, None) } else { AllocatorPool::new(thread_count) };