refactor: object pool

Author: SyMind

benches/bench_complex_replace_source.rs

@@ -9,7 +9,8 @@ pub use criterion::*;
 pub use codspeed_criterion_compat::*;
 
 use rspack_sources::{
-  BoxSource, MapOptions, OriginalSource, ReplaceSource, SourceExt,
+  using_object_pool, BoxSource, MapOptions, OriginalSource, ReplaceSource,
+  SourceExt,
 };
 
 static LARGE_REPLACE_SOURCE: LazyLock<BoxSource> = LazyLock::new(|| {
@@ -36723,8 +36724,10 @@ static LARGE_REPLACE_SOURCE: LazyLock<BoxSource> = LazyLock::new(|| {
 pub fn benchmark_complex_replace_source_map(b: &mut Bencher) {
   let source = LARGE_REPLACE_SOURCE.clone();
 
-  b.iter(|| {
-    black_box(source.map(&MapOptions::default()));
+  using_object_pool(|| {
+    b.iter(|| {
+      black_box(source.map(&MapOptions::default()));
+    });
   });
 }
 

benches/benchmark_repetitive_react_components.rs

@@ -9,9 +9,9 @@ pub use criterion::*;
 pub use codspeed_criterion_compat::*;
 
 use rspack_sources::{
-  BoxSource, ConcatSource, MapOptions, OriginalSource, RawStringSource,
-  ReplaceSource, ReplacementEnforce, Source, SourceExt, SourceMap,
-  SourceMapSource, SourceMapSourceOptions,
+  using_object_pool, BoxSource, ConcatSource, MapOptions, OriginalSource,
+  RawStringSource, ReplaceSource, ReplacementEnforce, Source, SourceExt,
+  SourceMap, SourceMapSource, SourceMapSourceOptions,
 };
 
 static REPETITIVE_1K_REACT_COMPONENTS_SOURCE: LazyLock<BoxSource> =
@@ -3504,8 +3504,10 @@ static REPETITIVE_1K_REACT_COMPONENTS_SOURCE: LazyLock<BoxSource> =
 pub fn benchmark_repetitive_react_components_map(b: &mut Bencher) {
   let source = REPETITIVE_1K_REACT_COMPONENTS_SOURCE.clone();
 
-  b.iter(|| {
-    black_box(source.map(&MapOptions::default()));
+  using_object_pool(|| {
+    b.iter(|| {
+      black_box(source.map(&MapOptions::default()));
+    });
   });
 }
 

src/helpers.rs

@@ -833,7 +833,7 @@ where
                     match inner_source_contents.get(&inner_source_index) {
                       Some(Some(source_content)) => Some(
                         split_into_lines(source_content)
-                          .map(|line| WithIndices::new(line))
+                          .map(WithIndices::new)
                           .collect(),
                       ),
                       _ => None,
@@ -933,7 +933,7 @@ where
                     match inner_source_contents.get(&inner_source_index) {
                       Some(Some(source_content)) => Some(
                         split_into_lines(source_content)
-                          .map(|line| WithIndices::new(line))
+                          .map(WithIndices::new)
                           .collect(),
                       ),
                       _ => None,

src/lib.rs

@@ -7,7 +7,7 @@ mod encoder;
 mod error;
 mod helpers;
 mod linear_map;
-mod memory_pool;
+mod object_pool;
 mod original_source;
 mod raw_source;
 mod replace_source;
@@ -40,3 +40,5 @@ pub mod stream_chunks {
 }
 
 pub use helpers::{decode_mappings, encode_mappings};
+
+pub use object_pool::using_object_pool;

src/memory_pool.rs renamed to src/object_pool.rs

@@ -1,6 +1,4 @@
-use std::{
-  cell::RefCell, collections::BTreeMap, rc::Rc, sync::atomic::AtomicBool,
-};
+use std::{cell::RefCell, collections::BTreeMap, rc::Rc};
 
 // Vector pooling minimum capacity threshold
 // Recommended threshold: 64
@@ -11,7 +9,7 @@ use std::{
 // 4. Empirical value: 64 is a proven balance point in real projects
 const MIN_POOL_CAPACITY: usize = 64;
 
-trait Poolable {
+pub trait Poolable {
   fn with_capacity(capacity: usize) -> Self;
   fn capacity(&self) -> usize;
   fn clear(&mut self);
@@ -34,50 +32,64 @@ impl<T> Poolable for Vec<T> {
 /// A memory pool for reusing `T` allocations to reduce memory allocation overhead.
 #[derive(Default, Debug)]
 pub struct ObjectPool<T> {
-  usize_vec_pool: RefCell<BTreeMap<usize, Vec<T>>>,
+  objects: Rc<RefCell<BTreeMap<usize, Vec<T>>>>,
+}
+
+impl<T> Clone for ObjectPool<T> {
+  fn clone(&self) -> Self {
+    Self {
+      objects: self.objects.clone(),
+    }
+  }
 }
 
 impl<T: Poolable> ObjectPool<T> {
   /// Retrieves a reusable `T` from the pool with at least the requested capacity.
   pub fn pull(&self, requested_capacity: usize) -> T {
     if requested_capacity < MIN_POOL_CAPACITY
-      || self.usize_vec_pool.borrow().is_empty()
+      || self.objects.borrow().is_empty()
     {
       return T::with_capacity(requested_capacity);
     }
-    let mut usize_vec_pool = self.usize_vec_pool.borrow_mut();
-    if let Some((_, bucket)) =
-      usize_vec_pool.range_mut(requested_capacity..).next()
-    {
-      if let Some(mut v) = bucket.pop() {
-        v.clear();
-        return v;
+    let mut objects = self.objects.borrow_mut();
+    if let Some((_, bucket)) = objects.range_mut(requested_capacity..).next() {
+      if let Some(mut object) = bucket.pop() {
+        object.clear();
+        return object;
       }
     }
     T::with_capacity(requested_capacity)
   }
 
   /// Returns a `T` to the pool for future reuse.
-  pub fn ret(&self, object: T) {
+  pub fn return_to_pool(&self, object: T) {
     if object.capacity() < MIN_POOL_CAPACITY {
       return;
     }
-    let mut usize_vec_pool = self.usize_vec_pool.borrow_mut();
+    let mut objects = self.objects.borrow_mut();
     let cap = object.capacity();
-    let bucket = usize_vec_pool.entry(cap).or_default();
+    let bucket = objects.entry(cap).or_default();
     bucket.push(object);
   }
 }
 
+/// A smart pointer that holds a pooled object and automatically returns it to the pool when dropped.
+///
+/// `Pooled<T>` implements RAII (Resource Acquisition Is Initialization) pattern to manage
+/// pooled objects lifecycle. When the `Pooled` instance is dropped, the contained object
+/// is automatically returned to its associated pool for future reuse.
 #[derive(Debug)]
 pub struct Pooled<T: Poolable> {
   object: Option<T>,
-  pool: Rc<ObjectPool<T>>,
+  pool: Option<ObjectPool<T>>,
 }
 
 impl<T: Poolable> Pooled<T> {
-  fn new(pool: Rc<ObjectPool<T>>, requested_capacity: usize) -> Self {
-    let object = pool.pull(requested_capacity);
+  pub fn new(pool: Option<ObjectPool<T>>, requested_capacity: usize) -> Self {
+    let object = match &pool {
+      Some(pool) => pool.pull(requested_capacity),
+      None => T::with_capacity(requested_capacity),
+    };
     Self {
       object: Some(object),
       pool,
@@ -96,7 +108,9 @@ impl<T: Poolable> Pooled<T> {
 impl<T: Poolable> Drop for Pooled<T> {
   fn drop(&mut self) {
     if let Some(object) = self.object.take() {
-      self.pool.ret(object);
+      if let Some(pool) = &self.pool {
+        pool.return_to_pool(object);
+      }
     }
   }
 }
@@ -115,14 +129,33 @@ impl<T: Poolable> std::ops::DerefMut for Pooled<T> {
   }
 }
 
-pub(crate) const USING_OBJECT_POOL: AtomicBool = AtomicBool::new(false);
+thread_local! {
+  pub static USIZE_VEC_POOL: RefCell<Option<ObjectPool<Vec<usize>>>> = RefCell::default();
+}
 
+/// Executes a function with object pooling enabled for the current thread.
+///
+/// This function temporarily enables a thread-local object pool for `Vec<usize>` allocations,
+/// executes the provided closure, and then cleans up the pool to prevent memory leaks.
 pub fn using_object_pool<F, R>(f: F) -> R
 where
   F: FnOnce() -> R,
 {
-  USING_OBJECT_POOL.store(true, std::sync::atomic::Ordering::SeqCst);
+  // Initialize the thread-local pool if needed
+  USIZE_VEC_POOL.with(|pool| {
+    let mut pool_ref = pool.borrow_mut();
+    if pool_ref.is_none() {
+      *pool_ref = Some(ObjectPool::default());
+    }
+  });
+
   let result = f();
-  USING_OBJECT_POOL.store(false, std::sync::atomic::Ordering::SeqCst);
+
+  // Clean up the pool to prevent memory retention
+  // This ensures no memory is held between different pooling sessions
+  USIZE_VEC_POOL.with(|pool| {
+    pool.borrow_mut().take();
+  });
+
   result
 }

src/replace_source.rs

@@ -391,7 +391,7 @@ impl StreamChunks for ReplaceSource {
           match source_content {
             SourceContent::Raw(source) => {
               let lines = split_into_lines(source)
-                .map(|line| WithIndices::new(line))
+                .map(WithIndices::new)
                 .collect::<Vec<_>>();
               let matched =
                 check_content_at_position(&lines, line, column, expected_chunk);

src/with_indices.rs

@@ -1,6 +1,9 @@
 use std::{cell::OnceCell, marker::PhantomData};
 
-use crate::helpers::SourceText;
+use crate::{
+  helpers::SourceText,
+  object_pool::{Pooled, USIZE_VEC_POOL},
+};
 
 #[derive(Debug)]
 pub struct WithIndices<'text, S>
@@ -10,7 +13,7 @@ where
   /// line is a string reference
   pub line: S,
   /// the byte position of each `char` in `line` string slice .
-  pub indices_indexes: OnceCell<PooledUsizeVec<'context>>,
+  pub char_byte_indices: OnceCell<Pooled<Vec<usize>>>,
   data: PhantomData<&'text S>,
 }
 
@@ -20,7 +23,7 @@ where
 {
   pub fn new(line: S) -> Self {
     Self {
-      indices_indexes: OnceCell::new(),
+      char_byte_indices: OnceCell::new(),
       line,
       data: PhantomData,
     }
@@ -32,15 +35,17 @@ where
       return S::default();
     }
 
-    let indices_indexes = self.indices_indexes.get_or_init(|| {
-      let mut vec = PooledUsizeVec::new(self.line.len());
+    let char_byte_indices = self.char_byte_indices.get_or_init(|| {
+      let mut vec = USIZE_VEC_POOL.with(|pool| {
+        Pooled::new(pool.borrow().as_ref().cloned(), self.line.len())
+      });
       vec.extend(self.line.char_indices().map(|(i, _)| i));
       vec
     });
 
     let str_len = self.line.len();
-    let start = *indices_indexes.get(start_index).unwrap_or(&str_len);
-    let end = *indices_indexes.get(end_index).unwrap_or(&str_len);
+    let start = *char_byte_indices.get(start_index).unwrap_or(&str_len);
+    let end = *char_byte_indices.get(end_index).unwrap_or(&str_len);
 
     #[allow(unsafe_code)]
     unsafe {

tests/compat_source.rs

@@ -6,8 +6,8 @@ use rspack_sources::stream_chunks::{
   stream_chunks_default, GeneratedInfo, OnChunk, OnName, OnSource, StreamChunks,
 };
 use rspack_sources::{
-  ConcatSource, MapOptions, MemoryPool, RawStringSource, Rope, Source,
-  SourceExt, SourceMap, SourceValue,
+  ConcatSource, MapOptions, RawStringSource, Rope, Source, SourceExt,
+  SourceMap, SourceValue,
 };
 
 #[derive(Debug, Eq)]
@@ -42,14 +42,12 @@ impl Source for CompatSource {
 impl StreamChunks for CompatSource {
   fn stream_chunks<'a>(
     &'a self,
-    memory_pool: &'a MemoryPool,
     options: &MapOptions,
     on_chunk: OnChunk<'_, 'a>,
     on_source: OnSource<'_, 'a>,
     on_name: OnName<'_, 'a>,
   ) -> GeneratedInfo {
     stream_chunks_default(
-      memory_pool,
       self.0,
       self.1.as_ref(),
       options,