refactor: using_object_pool

Author: SyMind

benches/bench.rs

@@ -30,7 +30,14 @@ use benchmark_repetitive_react_components::{
   benchmark_repetitive_react_components_source,
 };
 
+<<<<<<< Updated upstream
 use crate::bench_complex_replace_source::benchmark_complex_replace_source_size;
+=======
+use crate::{
+  bench_complex_replace_source::benchmark_complex_replace_source_map_in_using_object_pool,
+  benchmark_repetitive_react_components::benchmark_repetitive_react_components_map_in_using_object_pool,
+};
+>>>>>>> Stashed changes
 
 const HELLOWORLD_JS: &str = include_str!(concat!(
   env!("CARGO_MANIFEST_DIR"),
@@ -161,6 +168,11 @@ fn bench_rspack_sources(criterion: &mut Criterion) {
     benchmark_complex_replace_source_map,
   );
 
+  group.bench_function(
+    "complex_replace_source_map_in_using_object_pool",
+    benchmark_complex_replace_source_map_in_using_object_pool,
+  );
+
   group.bench_function(
     "complex_replace_source_source",
     benchmark_complex_replace_source_source,
@@ -188,6 +200,11 @@ fn bench_rspack_sources(criterion: &mut Criterion) {
     benchmark_repetitive_react_components_map,
   );
 
+  group.bench_function(
+    "repetitive_react_components_map_in_using_object_pool",
+    benchmark_repetitive_react_components_map_in_using_object_pool,
+  );
+
   group.bench_function(
     "repetitive_react_components_source",
     benchmark_repetitive_react_components_source,

benches/bench_complex_replace_source.rs

@@ -36724,6 +36724,16 @@ 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()));
+  });
+}
+
+pub fn benchmark_complex_replace_source_map_in_using_object_pool(
+  b: &mut Bencher,
+) {
+  let source = LARGE_REPLACE_SOURCE.clone();
+
   using_object_pool(|| {
     b.iter(|| {
       black_box(source.map(&MapOptions::default()));

benches/benchmark_repetitive_react_components.rs

@@ -3504,6 +3504,16 @@ 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()));
+  });
+}
+
+pub fn benchmark_repetitive_react_components_map_in_using_object_pool(
+  b: &mut Bencher,
+) {
+  let source = REPETITIVE_1K_REACT_COMPONENTS_SOURCE.clone();
+
   using_object_pool(|| {
     b.iter(|| {
       black_box(source.map(&MapOptions::default()));

src/helpers.rs

@@ -11,6 +11,7 @@ use crate::{
   decoder::MappingsDecoder,
   encoder::create_encoder,
   linear_map::LinearMap,
+  object_pool::cleanup_idle_object_pool,
   source::{Mapping, OriginalLocation},
   with_indices::WithIndices,
   MapOptions, Rope, SourceMap,
@@ -62,8 +63,12 @@ pub fn get_map<'a, S: StreamChunks>(
     },
   );
   let mappings = mappings_encoder.drain();
-  (!mappings.is_empty())
-    .then(|| SourceMap::new(mappings, sources, sources_content, names))
+  let source_map = (!mappings.is_empty())
+    .then(|| SourceMap::new(mappings, sources, sources_content, names));
+
+  cleanup_idle_object_pool();
+
+  source_map
 }
 
 /// [StreamChunks] abstraction, see [webpack-sources source.streamChunks](https://github.com/webpack/webpack-sources/blob/9f98066311d53a153fdc7c633422a1d086528027/lib/helpers/streamChunks.js#L13).

src/object_pool.rs

@@ -1,4 +1,9 @@
-use std::{cell::RefCell, collections::BTreeMap, rc::Rc};
+use std::{
+  cell::{OnceCell, RefCell},
+  collections::BTreeMap,
+  rc::Rc,
+  sync::atomic::AtomicBool,
+};
 
 // Vector pooling minimum capacity threshold
 // Recommended threshold: 64
@@ -71,6 +76,10 @@ impl<T: Poolable> ObjectPool<T> {
     let bucket = objects.entry(cap).or_default();
     bucket.push(object);
   }
+
+  pub fn clear(&self) {
+    self.objects.borrow_mut().clear();
+  }
 }
 
 /// A smart pointer that holds a pooled object and automatically returns it to the pool when dropped.
@@ -130,9 +139,11 @@ impl<T: Poolable> std::ops::DerefMut for Pooled<T> {
 }
 
 thread_local! {
-  pub static USIZE_VEC_POOL: RefCell<Option<ObjectPool<Vec<usize>>>> = RefCell::default();
+  pub static USIZE_VEC_POOL: OnceCell<ObjectPool<Vec<usize>>> = OnceCell::default();
 }
 
+pub static IN_USING_OBJECT_POOL: AtomicBool = AtomicBool::new(false);
+
 /// Executes a function with object pooling enabled for the current thread.
 ///
 /// This function temporarily enables a thread-local object pool for `Vec<usize>` allocations,
@@ -141,21 +152,35 @@ pub fn using_object_pool<F, R>(f: F) -> R
 where
   F: FnOnce() -> R,
 {
+  IN_USING_OBJECT_POOL.store(true, std::sync::atomic::Ordering::Relaxed);
   // 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());
-    }
+  USIZE_VEC_POOL.with(|once_cell| {
+    once_cell.get_or_init(ObjectPool::default);
   });
 
   let result = f();
 
-  // 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();
+  IN_USING_OBJECT_POOL.store(false, std::sync::atomic::Ordering::Relaxed);
+  USIZE_VEC_POOL.with(|once_cell| {
+    if let Some(pool) = once_cell.get() {
+      pool.clear();
+    }
   });
 
   result
 }
+
+/// Cleans up the object pool when not in pooling mode to prevent memory retention.
+///
+/// This function is called automatically after map operations complete to ensure
+/// that memory is not retained unnecessarily outside of pooling contexts.
+pub fn cleanup_idle_object_pool() {
+  // Only clear if we're not in an explicit pooling context
+  if !IN_USING_OBJECT_POOL.load(std::sync::atomic::Ordering::Relaxed) {
+    USIZE_VEC_POOL.with(|once_cell| {
+      if let Some(pool) = once_cell.get() {
+        pool.clear();
+      }
+    });
+  }
+}

src/with_indices.rs

@@ -2,7 +2,7 @@ use std::{cell::OnceCell, marker::PhantomData};
 
 use crate::{
   helpers::SourceText,
-  object_pool::{Pooled, USIZE_VEC_POOL},
+  object_pool::{ObjectPool, Pooled, USIZE_VEC_POOL},
 };
 
 #[derive(Debug)]
@@ -36,8 +36,9 @@ where
     }
 
     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())
+      let mut vec = USIZE_VEC_POOL.with(|once_cell| {
+        let pool = once_cell.get_or_init(ObjectPool::default);
+        Pooled::new(Some(pool.clone()), self.line.len())
       });
       vec.extend(self.line.char_indices().map(|(i, _)| i));
       vec