test(parser): track number of allocations (#12555)

- related: oxc-project/backlog#5

This PR adds a reliable method of tracking the number of allocations made while running the parser. This includes both arena allocations (via a new feature), as well as system allocations (by creating a new global allocator). Reductions in these numbers should correlate to real-world performance improvements that aren't quantified in CodSpeed currently. This will also help prevent regressions where we accidentally allocate lots more memory than we expect.

Originally I had included total system bytes allocated, but it was tricky to get this number to match exactly between platforms. I opted not try and make this perfect but instead use the total number of allocations which is a good proxy for bytes anyway.

Author: camchenry

.cargo/config.toml

@@ -6,6 +6,7 @@ codecov = "llvm-cov --workspace --ignore-filename-regex tasks"
 coverage = "run -p oxc_coverage --profile coverage --"
 benchmark = "bench -p oxc_benchmark"
 minsize = "run -p oxc_minsize --profile coverage --"
+allocs = "run -p oxc_track_memory_allocations --profile coverage --"
 rule = "run -p rulegen"
 
 # Build oxlint in release mode

.github/workflows/ci.yml

@@ -282,6 +282,36 @@ jobs:
           cargo minsize
           git diff --exit-code
 
+  allocs:
+    name: Allocations
+    runs-on: ubuntu-latest
+    steps:
+      - uses: taiki-e/checkout-action@b13d20b7cda4e2f325ef19895128f7ff735c0b3d # v1.3.1
+      - uses: dorny/paths-filter@de90cc6fb38fc0963ad72b210f1f284cd68cea36 # v3.0.2
+        id: filter
+        with:
+          filters: |
+            src:
+              - '.github/workflows/ci.yml'
+              - 'crates/oxc_ast/**'
+              - 'crates/oxc_data_structures/**'
+              - 'crates/oxc_parser/**'
+              - 'crates/oxc_allocator/**'
+              - 'tasks/allocs/**'
+
+      - uses: oxc-project/setup-rust@cd82e1efec7fef815e2c23d296756f31c7cdc03d # v1.0.0
+        if: steps.filter.outputs.src == 'true'
+        with:
+          cache-key: allocs
+          save-cache: ${{ github.ref_name == 'main' }}
+
+      - name: Check allocations
+        if: steps.filter.outputs.src == 'true'
+        run: |
+          cargo allocs
+          git diff --exit-code ||
+          (echo 'Allocations have changed. Run the `cargo allocs` command to update the allocation snapshot, otherwise please fix the regression.' && exit 1)
+
   ast_changes:
     name: AST Changes
     runs-on: ubuntu-latest

Cargo.lock

@@ -40,3 +40,5 @@ fixed_size = ["from_raw_parts", "dep:oxc_ast_macros"]
 disable_fixed_size = []
 from_raw_parts = []
 serialize = ["dep:serde", "oxc_estree/serialize"]
+track_allocations = []
+disable_track_allocations = []

crates/oxc_allocator/Cargo.toml

@@ -95,6 +95,25 @@ impl Alloc for Bump {
     /// Panics if reserving space for `layout` fails.
     #[inline(always)]
     fn alloc(&self, layout: Layout) -> NonNull<u8> {
+        // SAFETY: We only use `Bump` inside of `Allocator` in oxc, so the `self` reference should
+        // also be pointing to a valid `Allocator` struct, which we can use for finding the stats fields.
+        // This will go away when we add a custom allocator to oxc.
+        #[cfg(all(feature = "track_allocations", not(feature = "disable_track_allocations")))]
+        unsafe {
+            use crate::Allocator;
+            use std::{
+                mem::offset_of,
+                ptr,
+                sync::atomic::{AtomicUsize, Ordering},
+            };
+            #[expect(clippy::cast_possible_wrap)]
+            const OFFSET: isize = (offset_of!(Allocator, num_alloc) as isize)
+                - (offset_of!(Allocator, bump) as isize);
+            let num_alloc_ptr = ptr::from_ref(self).byte_offset(OFFSET).cast::<AtomicUsize>();
+            let num_alloc = num_alloc_ptr.as_ref().unwrap_unchecked();
+            num_alloc.fetch_add(1, Ordering::SeqCst);
+        }
+
         self.alloc_layout(layout)
     }
 
@@ -133,6 +152,25 @@ impl Alloc for Bump {
     /// Panics / aborts if reserving space for `new_layout` fails.
     #[inline(always)]
     unsafe fn grow(&self, ptr: NonNull<u8>, old_layout: Layout, new_layout: Layout) -> NonNull<u8> {
+        // SAFETY: We only use `Bump` inside of `Allocator` in oxc, so the `self` reference should
+        // also be pointing to a valid `Allocator` struct, which we can use for finding the stats fields.
+        // This will go away when we add a custom allocator to oxc.
+        #[cfg(all(feature = "track_allocations", not(feature = "disable_track_allocations")))]
+        unsafe {
+            use crate::Allocator;
+            use std::{
+                mem::offset_of,
+                ptr,
+                sync::atomic::{AtomicUsize, Ordering},
+            };
+            #[expect(clippy::cast_possible_wrap)]
+            const OFFSET: isize = (offset_of!(Allocator, num_realloc) as isize)
+                - (offset_of!(Allocator, bump) as isize);
+            let num_realloc_ptr = ptr::from_ref(self).byte_offset(OFFSET).cast::<AtomicUsize>();
+            let num_realloc = num_realloc_ptr.as_ref().unwrap_unchecked();
+            num_realloc.fetch_add(1, Ordering::SeqCst);
+        }
+
         // SAFETY: Safety requirements of `Allocator::grow` are the same as for this method
         let res = unsafe { Allocator::grow(&self, ptr, old_layout, new_layout) };
         match res {

crates/oxc_allocator/src/alloc.rs

@@ -214,7 +214,20 @@ use oxc_data_structures::assert_unchecked;
 /// [`HashMap::new_in`]: crate::HashMap::new_in
 #[derive(Default)]
 pub struct Allocator {
+    #[cfg(not(all(feature = "track_allocations", not(feature = "disable_track_allocations"))))]
     bump: Bump,
+    // NOTE: We need to expose `bump` publicly here for calculating its field offset in memory.
+    #[cfg(all(feature = "track_allocations", not(feature = "disable_track_allocations")))]
+    #[doc(hidden)]
+    pub bump: Bump,
+    /// Used to track the total number of allocations made in this allocator when the `track_allocations` feature is enabled.
+    #[cfg(all(feature = "track_allocations", not(feature = "disable_track_allocations")))]
+    #[doc(hidden)]
+    pub num_alloc: std::sync::atomic::AtomicUsize,
+    /// Used to track the total number of re-allocations made in this allocator when the `track_allocations` feature is enabled.
+    #[cfg(all(feature = "track_allocations", not(feature = "disable_track_allocations")))]
+    #[doc(hidden)]
+    pub num_realloc: std::sync::atomic::AtomicUsize,
 }
 
 impl Allocator {
@@ -241,7 +254,13 @@ impl Allocator {
     #[expect(clippy::inline_always)]
     #[inline(always)]
     pub fn new() -> Self {
-        Self { bump: Bump::new() }
+        Self {
+            bump: Bump::new(),
+            #[cfg(all(feature = "track_allocations", not(feature = "disable_track_allocations")))]
+            num_alloc: std::sync::atomic::AtomicUsize::new(0),
+            #[cfg(all(feature = "track_allocations", not(feature = "disable_track_allocations")))]
+            num_realloc: std::sync::atomic::AtomicUsize::new(0),
+        }
     }
 
     /// Create a new [`Allocator`] with specified capacity.
@@ -252,7 +271,13 @@ impl Allocator {
     #[expect(clippy::inline_always)]
     #[inline(always)]
     pub fn with_capacity(capacity: usize) -> Self {
-        Self { bump: Bump::with_capacity(capacity) }
+        Self {
+            bump: Bump::with_capacity(capacity),
+            #[cfg(all(feature = "track_allocations", not(feature = "disable_track_allocations")))]
+            num_alloc: std::sync::atomic::AtomicUsize::new(0),
+            #[cfg(all(feature = "track_allocations", not(feature = "disable_track_allocations")))]
+            num_realloc: std::sync::atomic::AtomicUsize::new(0),
+        }
     }
 
     /// Allocate an object in this [`Allocator`] and return an exclusive reference to it.
@@ -276,6 +301,9 @@ impl Allocator {
     pub fn alloc<T>(&self, val: T) -> &mut T {
         const { assert!(!std::mem::needs_drop::<T>(), "Cannot allocate Drop type in arena") };
 
+        #[cfg(all(feature = "track_allocations", not(feature = "disable_track_allocations")))]
+        self.num_alloc.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
+
         self.bump.alloc(val)
     }
 
@@ -297,6 +325,9 @@ impl Allocator {
     #[expect(clippy::inline_always)]
     #[inline(always)]
     pub fn alloc_str<'alloc>(&'alloc self, src: &str) -> &'alloc str {
+        #[cfg(all(feature = "track_allocations", not(feature = "disable_track_allocations")))]
+        self.num_alloc.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
+
         self.bump.alloc_str(src)
     }
 
@@ -317,6 +348,9 @@ impl Allocator {
     #[expect(clippy::inline_always)]
     #[inline(always)]
     pub fn alloc_slice_copy<T: Copy>(&self, src: &[T]) -> &mut [T] {
+        #[cfg(all(feature = "track_allocations", not(feature = "disable_track_allocations")))]
+        self.num_alloc.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
+
         self.bump.alloc_slice_copy(src)
     }
 
@@ -329,6 +363,9 @@ impl Allocator {
     ///
     /// Panics if reserving space matching `layout` fails.
     pub fn alloc_layout(&self, layout: Layout) -> NonNull<u8> {
+        #[cfg(all(feature = "track_allocations", not(feature = "disable_track_allocations")))]
+        self.num_alloc.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
+
         self.bump.alloc_layout(layout)
     }
 
@@ -379,6 +416,9 @@ impl Allocator {
             "attempted to create a string longer than `isize::MAX` bytes"
         );
 
+        #[cfg(all(feature = "track_allocations", not(feature = "disable_track_allocations")))]
+        self.num_alloc.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
+
         // Create actual `&str` in a separate function, to ensure that `alloc_concat_strs_array`
         // is inlined, so that compiler has knowledge to remove the overflow checks above.
         // When some of `strings` are static, this function is usually only a few instructions.
@@ -471,6 +511,12 @@ impl Allocator {
     #[expect(clippy::inline_always)]
     #[inline(always)]
     pub fn reset(&mut self) {
+        #[cfg(all(feature = "track_allocations", not(feature = "disable_track_allocations")))]
+        {
+            self.num_alloc.store(0, std::sync::atomic::Ordering::SeqCst);
+            self.num_realloc.store(0, std::sync::atomic::Ordering::SeqCst);
+        }
+
         self.bump.reset();
     }
 
@@ -592,7 +638,13 @@ impl Allocator {
     #[expect(clippy::inline_always)]
     #[inline(always)]
     pub(crate) fn from_bump(bump: Bump) -> Self {
-        Self { bump }
+        Self {
+            bump,
+            #[cfg(all(feature = "track_allocations", not(feature = "disable_track_allocations")))]
+            num_alloc: std::sync::atomic::AtomicUsize::new(0),
+            #[cfg(all(feature = "track_allocations", not(feature = "disable_track_allocations")))]
+            num_realloc: std::sync::atomic::AtomicUsize::new(0),
+        }
     }
 }
 

crates/oxc_allocator/src/allocator.rs

@@ -132,6 +132,10 @@ benchmark:
 benchmark-one *args:
   cargo benchmark --bench {{args}} --no-default-features --features {{args}}
 
+# Update memory allocation snapshots.
+allocs:
+  cargo allocs
+
 # Automatically DRY up Cargo.toml manifests in a workspace.
 autoinherit:
   cargo binstall cargo-autoinherit

justfile

@@ -0,0 +1,30 @@
+[package]
+name = "oxc_track_memory_allocations"
+version = "0.0.0"
+edition.workspace = true
+license.workspace = true
+publish = false
+
+[lints]
+workspace = true
+
+[lib]
+test = false
+doctest = false
+
+[[bin]]
+name = "oxc_track_memory_allocations"
+test = false
+doctest = false
+
+[dependencies]
+oxc_allocator = { workspace = true, features = ["track_allocations"] }
+oxc_parser = { workspace = true }
+oxc_tasks_common = { workspace = true }
+
+humansize = { workspace = true }
+mimalloc-safe = { workspace = true }
+
+[features]
+# Sentinel flag that should only get enabled if we're running under `--all-features`
+is_all_features = []

tasks/track_memory_allocations/Cargo.toml

@@ -0,0 +1,3 @@
+# Oxc allocations stats
+
+This task keeps track of the number of system allocations as well as arena allocations and the total number of bytes allocated. This is used for monitoring possible regressions in allocations and improvements in memory usage.

tasks/track_memory_allocations/README.md

@@ -0,0 +1,14 @@
+File                       | File size      || Sys allocs     | Sys reallocs   || Arena allocs   | Arena reallocs | Arena bytes    
+-------------------------------------------------------------------------------------------------------------------------------------------
+checker.ts                 | 2.92 MB        ||          10161 |             21 ||         268665 |          23341 | 12.422 MB     
+
+cal.com.tsx                | 1.06 MB        ||           2209 |             54 ||         138188 |          13712 | 7.768 MB      
+
+RadixUIAdoptionSection.jsx | 2.52 kB        ||              4 |              0 ||            365 |             66 | 19.104 kB     
+
+pdf.mjs                    | 567.30 kB      ||            688 |             71 ||          90678 |           8148 | 4.308 MB      
+
+antd.js                    | 4.12 MB        ||           6860 |            282 ||         509752 |          55282 | 25.574 MB     
+
+binder.ts                  | 193.08 kB      ||            540 |              7 ||          16807 |           1475 | 873.704 kB    
+