feat(allocator): optimize memory usage with dual allocator pools for import plugin

apps/oxlint/src/js_plugins/external_linter.rs

@@ -77,12 +77,11 @@ fn wrap_lint_file(cb: JsLintFileCb) -> ExternalLinterLintFileCb {
             let (tx, rx) = channel();
 
             // SAFETY: This function is only called when an `ExternalLinter` exists.
-            // When that is the case, the `AllocatorPool` used to create `Allocator`s is created with
-            // `AllocatorPool::new_fixed_size`, so all `Allocator`s are created via `FixedSizeAllocator`.
-            // This is somewhat sketchy, as we don't have a type-level guarantee of this invariant,
-            // but it does hold at present.
-            // When we replace `bumpalo` with a custom allocator, we can close this soundness hole.
-            // TODO: Do that.
+            // When that is the case, the `AllocatorPool` is created with `AllocatorPool::new_dual`,
+            // and `Allocator`s for files that need linting are created via `FixedSizeAllocator`
+            // using `get_maybe_fixed_size(true)`.
+            // This is validated with debug assertions in the linter runtime before calling this function.
+            // The `AllocatorGuard` tracks whether an allocator is fixed-size via the `is_fixed_size` field.
             let (buffer_id, buffer) = unsafe { get_buffer(allocator) };
 
             // Send data to JS

apps/oxlint/src/js_plugins/raw_fs.rs

@@ -6,16 +6,17 @@ use std::{
     ptr, slice,
 };
 
-use oxc_allocator::Allocator;
+use oxc_allocator::{Allocator, AllocatorGuard};
 use oxc_linter::RuntimeFileSystem;
 
 /// File system used when JS plugins are in use.
 ///
 /// Identical to `OsFileSystem`, except that `read_to_arena_str` reads the file's contents into
 /// start of the allocator, instead of the end. This conforms to what raw transfer needs.
 ///
-/// Must only be used in conjunction with `AllocatorPool` created with `new_fixed_size`,
-/// which wraps `Allocator`s with a custom `Drop` impl, which makes `read_to_arena_str` safe.
+/// Must only be used with fixed-size `Allocator`s (created via `AllocatorPool::new_fixed_size`
+/// or `AllocatorPool::new_dual` with `get_maybe_fixed_size(true)`),
+/// which have a custom `Drop` impl that makes `read_to_arena_str` safe.
 ///
 /// This is a temporary solution. When we replace `bumpalo` with our own allocator, all strings
 /// will be written at start of the arena, so then `OsFileSystem` will work fine, and we can
@@ -33,10 +34,15 @@ impl RuntimeFileSystem for RawTransferFileSystem {
     fn read_to_arena_str<'a>(
         &self,
         path: &Path,
-        allocator: &'a Allocator,
+        allocator_guard: &'a AllocatorGuard,
     ) -> Result<&'a str, std::io::Error> {
-        // SAFETY: Caller promises not to allow `allocator` to be dropped
-        unsafe { read_to_arena_str(path, allocator) }
+        #[cfg(all(feature = "napi", target_pointer_width = "64", target_endian = "little"))]
+        if allocator_guard.is_fixed_size() {
+            // SAFETY: Caller promises not to allow `allocator` to be dropped
+            unsafe { return read_to_arena_str(path, allocator_guard) }
+        }
+
+        oxc_linter::read_to_arena_str(path, allocator_guard)
     }
 
     fn write_file(&self, path: &Path, content: &str) -> Result<(), std::io::Error> {

crates/oxc_allocator/src/pool/mod.rs

@@ -18,39 +18,70 @@ pub use fixed_size::{FixedSizeAllocatorMetadata, free_fixed_size_allocator};
 
 /// A thread-safe pool for reusing [`Allocator`] instances to reduce allocation overhead.
 ///
-/// Uses either:
-/// 1. Standard allocators - suitable for general use.
-/// 2. Fixed-size allocators - compatible with raw transfer.
+/// Supports three configurations:
+/// 1. Standard allocators only - suitable for general use ([`AllocatorPool::new`]).
+/// 2. Fixed-size allocators only - compatible with raw transfer ([`AllocatorPool::new_fixed_size`]).
+/// 3. Dual mode with both types - allows choosing allocator type at runtime ([`AllocatorPool::new_dual`]).
 ///
-/// Standard allocator pool is created by [`AllocatorPool::new`].
-/// Fixed-size allocator pool is created by [`AllocatorPool::new_fixed_size`].
+/// The dual mode is useful when some allocations need raw transfer (e.g., files to be linted with JS plugins)
+/// while others don't (e.g., dependency files only parsed for module resolution).
 ///
 /// Fixed-size allocators are only supported on 64-bit little-endian platforms at present,
 /// and require the `fixed_size` Cargo feature to be enabled.
-#[repr(transparent)]
-pub struct AllocatorPool(AllocatorPoolInner);
-
-/// Inner type of [`AllocatorPool`], holding either a standard or fixed-size allocator pool.
-enum AllocatorPoolInner {
-    Standard(StandardAllocatorPool),
+pub struct AllocatorPool {
+    standard: StandardAllocatorPool,
     #[cfg(all(feature = "fixed_size", target_pointer_width = "64", target_endian = "little"))]
-    FixedSize(FixedSizeAllocatorPool),
+    fixed_size: FixedSizeAllocatorPool,
 }
 
 impl AllocatorPool {
     /// Create a new [`AllocatorPool`] for use across the specified number of threads,
-    /// which uses standard allocators.
+    /// which uses standard allocators only.
     pub fn new(thread_count: usize) -> AllocatorPool {
-        Self(AllocatorPoolInner::Standard(StandardAllocatorPool::new(thread_count)))
+        Self {
+            standard: StandardAllocatorPool::new(thread_count),
+            #[cfg(all(
+                feature = "fixed_size",
+                target_pointer_width = "64",
+                target_endian = "little"
+            ))]
+            fixed_size: FixedSizeAllocatorPool::new(0),
+        }
     }
 
     /// Create a new [`AllocatorPool`] for use across the specified number of threads,
-    /// which uses fixed-size allocators (suitable for raw transfer).
+    /// which uses fixed-size allocators only (suitable for raw transfer).
     #[cfg(feature = "fixed_size")]
     pub fn new_fixed_size(thread_count: usize) -> AllocatorPool {
         #[cfg(all(target_pointer_width = "64", target_endian = "little"))]
         {
-            Self(AllocatorPoolInner::FixedSize(FixedSizeAllocatorPool::new(thread_count)))
+            Self {
+                standard: StandardAllocatorPool::new(0), // Not used, but need to initialize
+                fixed_size: FixedSizeAllocatorPool::new(thread_count),
+            }
+        }
+
+        #[cfg(not(all(target_pointer_width = "64", target_endian = "little")))]
+        {
+            let _thread_count = thread_count; // Avoid unused vars lint warning
+            panic!("Fixed size allocators are only supported on 64-bit little-endian platforms");
+        }
+    }
+
+    /// Create a new [`AllocatorPool`] for use across the specified number of threads,
+    /// which contains both standard and fixed-size allocators.
+    ///
+    /// This allows choosing the allocator type at runtime via [`AllocatorPool::get_maybe_fixed_size`].
+    /// Useful when some allocations need raw transfer (e.g., files to be linted with JS plugins)
+    /// while others don't (e.g., dependency files only parsed for module resolution).
+    #[cfg(feature = "fixed_size")]
+    pub fn new_dual(thread_count: usize) -> AllocatorPool {
+        #[cfg(all(target_pointer_width = "64", target_endian = "little"))]
+        {
+            Self {
+                standard: StandardAllocatorPool::new(thread_count),
+                fixed_size: FixedSizeAllocatorPool::new(thread_count),
+            }
         }
 
         #[cfg(not(all(target_pointer_width = "64", target_endian = "little")))]
@@ -64,21 +95,80 @@ impl AllocatorPool {
     ///
     /// Returns an [`AllocatorGuard`] that gives access to the allocator.
     ///
+    /// If the pool was created with [`AllocatorPool::new_fixed_size`], returns a fixed-size allocator.
+    /// Otherwise, returns a standard allocator.
+    ///
     /// # Panics
     ///
     /// Panics if the underlying mutex is poisoned.
     pub fn get(&self) -> AllocatorGuard<'_> {
-        let allocator = match &self.0 {
-            AllocatorPoolInner::Standard(pool) => pool.get(),
+        if cfg!(all(feature = "fixed_size", target_pointer_width = "64", target_endian = "little"))
+        {
+            let allocator = self.fixed_size.get();
+            return AllocatorGuard {
+                allocator: ManuallyDrop::new(allocator),
+                is_fixed_size: true,
+                pool: self,
+            };
+        }
+
+        let allocator = self.standard.get();
+        AllocatorGuard {
+            allocator: ManuallyDrop::new(allocator),
             #[cfg(all(
                 feature = "fixed_size",
                 target_pointer_width = "64",
                 target_endian = "little"
             ))]
-            AllocatorPoolInner::FixedSize(pool) => pool.get(),
-        };
+            is_fixed_size: false,
+            pool: self,
+        }
+    }
 
-        AllocatorGuard { allocator: ManuallyDrop::new(allocator), pool: self }
+    /// Retrieve an [`Allocator`] from the pool, choosing between standard and fixed-size allocators.
+    ///
+    /// Returns an [`AllocatorGuard`] that gives access to the allocator.
+    ///
+    /// # Parameters
+    ///
+    /// * `fixed_size` - If `true`, returns a fixed-size allocator (suitable for raw transfer).
+    ///   If `false`, returns a standard allocator.
+    ///
+    /// # Panics
+    ///
+    /// Panics if:
+    /// - The underlying mutex is poisoned.
+    /// - `fixed_size` is `true` but the pool was not created with [`AllocatorPool::new_dual`] or [`AllocatorPool::new_fixed_size`].
+    #[cfg(feature = "fixed_size")]
+    pub fn get_maybe_fixed_size(&self, fixed_size: bool) -> AllocatorGuard<'_> {
+        #[cfg(all(target_pointer_width = "64", target_endian = "little"))]
+        {
+            if fixed_size {
+                let allocator = self.fixed_size.get();
+                return AllocatorGuard {
+                    allocator: ManuallyDrop::new(allocator),
+                    is_fixed_size: true,
+                    pool: self,
+                };
+            }
+        }
+
+        #[cfg(not(all(target_pointer_width = "64", target_endian = "little")))]
+        if fixed_size {
+            panic!("Fixed size allocators are only supported on 64-bit little-endian platforms");
+        }
+
+        let allocator = self.standard.get();
+        AllocatorGuard {
+            allocator: ManuallyDrop::new(allocator),
+            #[cfg(all(
+                feature = "fixed_size",
+                target_pointer_width = "64",
+                target_endian = "little"
+            ))]
+            is_fixed_size: false,
+            pool: self,
+        }
     }
 
     /// Add an [`Allocator`] to the pool.
@@ -88,20 +178,32 @@ impl AllocatorPool {
     /// # Panics
     ///
     /// Panics if the underlying mutex is poisoned.
+    #[cfg(all(feature = "fixed_size", target_pointer_width = "64", target_endian = "little"))]
+    fn add(&self, allocator: Allocator, is_fixed_size: bool) {
+        // SAFETY: This method is only called from `AllocatorGuard::drop`.
+        // `AllocatorGuard`s are only created by `AllocatorPool::get` or `AllocatorPool::get_maybe_fixed_size`,
+        // so the `Allocator` must have been created by this pool. Therefore, it is the correct type for the pool.
+        unsafe {
+            if is_fixed_size {
+                self.fixed_size.add(allocator);
+            } else {
+                self.standard.add(allocator);
+            }
+        }
+    }
+
+    /// Add an [`Allocator`] to the pool (for platforms without fixed-size support).
+    #[cfg(not(all(
+        feature = "fixed_size",
+        target_pointer_width = "64",
+        target_endian = "little"
+    )))]
     fn add(&self, allocator: Allocator) {
         // SAFETY: This method is only called from `AllocatorGuard::drop`.
         // `AllocatorGuard`s are only created by `AllocatorPool::get`, so the `Allocator` must have
         // been created by this pool. Therefore, it is the correct type for the pool.
         unsafe {
-            match &self.0 {
-                AllocatorPoolInner::Standard(pool) => pool.add(allocator),
-                #[cfg(all(
-                    feature = "fixed_size",
-                    target_pointer_width = "64",
-                    target_endian = "little"
-                ))]
-                AllocatorPoolInner::FixedSize(pool) => pool.add(allocator),
-            }
+            self.standard.add(allocator);
         }
     }
 }
@@ -111,9 +213,19 @@ impl AllocatorPool {
 /// On drop, the `Allocator` is reset and returned to the pool.
 pub struct AllocatorGuard<'alloc_pool> {
     allocator: ManuallyDrop<Allocator>,
+    #[cfg(all(feature = "fixed_size", target_pointer_width = "64", target_endian = "little"))]
+    is_fixed_size: bool,
     pool: &'alloc_pool AllocatorPool,
 }
 
+impl AllocatorGuard<'_> {
+    /// Returns `true` if this allocator is a fixed-size allocator (suitable for raw transfer).
+    #[cfg(all(feature = "fixed_size", target_pointer_width = "64", target_endian = "little"))]
+    pub fn is_fixed_size(&self) -> bool {
+        self.is_fixed_size
+    }
+}
+
 impl Deref for AllocatorGuard<'_> {
     type Target = Allocator;
 
@@ -127,6 +239,15 @@ impl Drop for AllocatorGuard<'_> {
     fn drop(&mut self) {
         // SAFETY: After taking ownership of the `Allocator`, we do not touch the `ManuallyDrop` again
         let allocator = unsafe { ManuallyDrop::take(&mut self.allocator) };
+
+        #[cfg(all(feature = "fixed_size", target_pointer_width = "64", target_endian = "little"))]
+        self.pool.add(allocator, self.is_fixed_size);
+
+        #[cfg(not(all(
+            feature = "fixed_size",
+            target_pointer_width = "64",
+            target_endian = "little"
+        )))]
         self.pool.add(allocator);
     }
 }

crates/oxc_language_server/src/linter/isolated_lint_handler.rs

@@ -8,7 +8,6 @@ use oxc_data_structures::rope::Rope;
 use rustc_hash::FxHashSet;
 use tower_lsp_server::{UriExt, lsp_types::Uri};
 
-use oxc_allocator::Allocator;
 use oxc_linter::{
     AllowWarnDeny, ConfigStore, DisableDirectives, Fix, FixKind, LINTABLE_EXTENSIONS, LintOptions,
     LintRunner, LintRunnerBuilder, LintServiceOptions, Linter, Message, PossibleFixes,
@@ -49,13 +48,13 @@ impl RuntimeFileSystem for IsolatedLintHandlerFileSystem {
     fn read_to_arena_str<'a>(
         &'a self,
         path: &Path,
-        allocator: &'a Allocator,
+        allocator_guard: &'a oxc_allocator::AllocatorGuard,
     ) -> Result<&'a str, std::io::Error> {
         if path == self.path_to_lint {
             return Ok(&self.source_text);
         }
 
-        read_to_arena_str(path, allocator)
+        read_to_arena_str(path, allocator_guard)
     }
 
     fn write_file(&self, _path: &Path, _content: &str) -> Result<(), std::io::Error> {