diff --git a/crates/oxc_allocator/README.md b/crates/oxc_allocator/README.md index 07ec7465c5988..4c59af230660f 100644 --- a/crates/oxc_allocator/README.md +++ b/crates/oxc_allocator/README.md @@ -27,3 +27,5 @@ This approach is significantly faster than using the system allocator for AST op - `serialize` - Enables serialization support for `Box` and `Vec` with `serde` - `from_raw_parts` - Adds unsafe `from_raw_parts` method (not recommended for general use) +- `track_allocations` - For internal use only. The APIs provided by this feature are sketchy at best, + and possibly undefined behavior. Do not enable this feature under any circumstances in production code. diff --git a/crates/oxc_allocator/src/alloc.rs b/crates/oxc_allocator/src/alloc.rs index 138f0d1500959..cc0521f620507 100644 --- a/crates/oxc_allocator/src/alloc.rs +++ b/crates/oxc_allocator/src/alloc.rs @@ -95,17 +95,11 @@ impl Alloc for Bump { /// Panics if reserving space for `layout` fails. #[inline(always)] fn alloc(&self, layout: Layout) -> NonNull { - // 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. + // SAFETY: This is UNSOUND (see comment on `get_stats_ref`). But usage is gated behind + // `track_allocations` feature, so should never be compiled in production code. #[cfg(all(feature = "track_allocations", not(feature = "disable_track_allocations")))] unsafe { - use crate::allocator::NUM_ALLOC_FIELD_OFFSET; - use std::sync::atomic::{AtomicUsize, Ordering}; - let num_alloc_ptr = - std::ptr::from_ref(self).byte_offset(NUM_ALLOC_FIELD_OFFSET).cast::(); - let num_alloc = num_alloc_ptr.as_ref().unwrap_unchecked(); - num_alloc.fetch_add(1, Ordering::SeqCst); + crate::tracking::get_stats_ref(self).record_allocation(); } self.alloc_layout(layout) @@ -146,18 +140,11 @@ impl Alloc for Bump { /// Panics / aborts if reserving space for `new_layout` fails. #[inline(always)] unsafe fn grow(&self, ptr: NonNull, old_layout: Layout, new_layout: Layout) -> NonNull { - // 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. + // SAFETY: This is UNSOUND (see comment on `get_stats_ref`). But usage is gated behind + // `track_allocations` feature, so should never be compiled in production code. #[cfg(all(feature = "track_allocations", not(feature = "disable_track_allocations")))] unsafe { - use crate::allocator::NUM_REALLOC_FIELD_OFFSET; - use std::sync::atomic::{AtomicUsize, Ordering}; - let num_realloc_ptr = std::ptr::from_ref(self) - .byte_offset(NUM_REALLOC_FIELD_OFFSET) - .cast::(); - let num_realloc = num_realloc_ptr.as_ref().unwrap_unchecked(); - num_realloc.fetch_add(1, Ordering::SeqCst); + crate::tracking::get_stats_ref(self).record_reallocation(); } // SAFETY: Safety requirements of `Allocator::grow` are the same as for this method diff --git a/crates/oxc_allocator/src/allocator.rs b/crates/oxc_allocator/src/allocator.rs index 50f43ced9de33..ada9c6e586d83 100644 --- a/crates/oxc_allocator/src/allocator.rs +++ b/crates/oxc_allocator/src/allocator.rs @@ -11,6 +11,9 @@ use bumpalo::Bump; use oxc_data_structures::assert_unchecked; +#[cfg(all(feature = "track_allocations", not(feature = "disable_track_allocations")))] +use crate::tracking::AllocationStats; + /// A bump-allocated memory arena. /// /// # Anatomy of an Allocator @@ -218,25 +221,17 @@ use oxc_data_structures::assert_unchecked; #[derive(Default)] pub struct Allocator { 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. + /// Used to track number of allocations made in this allocator when `track_allocations` feature is enabled #[cfg(all(feature = "track_allocations", not(feature = "disable_track_allocations")))] - #[doc(hidden)] - pub num_realloc: std::sync::atomic::AtomicUsize, + pub(crate) stats: AllocationStats, } -// Consts used in `Alloc` trait for allocation tracking -#[cfg(all(feature = "track_allocations", not(feature = "disable_track_allocations")))] -#[expect(clippy::cast_possible_wrap)] -pub const NUM_ALLOC_FIELD_OFFSET: isize = - (offset_of!(Allocator, num_alloc) as isize) - (offset_of!(Allocator, bump) as isize); +/// Offset of `stats` field, relative to `bump` field. +/// Used in `tracking` module for allocation tracking. #[cfg(all(feature = "track_allocations", not(feature = "disable_track_allocations")))] #[expect(clippy::cast_possible_wrap)] -pub const NUM_REALLOC_FIELD_OFFSET: isize = - (offset_of!(Allocator, num_realloc) as isize) - (offset_of!(Allocator, bump) as isize); +pub const STATS_FIELD_OFFSET: isize = + (offset_of!(Allocator, stats) as isize) - (offset_of!(Allocator, bump) as isize); impl Allocator { /// Create a new [`Allocator`] with no initial capacity. @@ -265,9 +260,7 @@ impl Allocator { 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), + stats: AllocationStats::default(), } } @@ -282,9 +275,7 @@ impl Allocator { 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), + stats: AllocationStats::default(), } } @@ -310,7 +301,7 @@ impl Allocator { const { assert!(!std::mem::needs_drop::(), "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.stats.record_allocation(); self.bump.alloc(val) } @@ -334,7 +325,7 @@ impl Allocator { #[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.stats.record_allocation(); self.bump.alloc_str(src) } @@ -357,7 +348,7 @@ impl Allocator { #[inline(always)] pub fn alloc_slice_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.stats.record_allocation(); self.bump.alloc_slice_copy(src) } @@ -372,7 +363,7 @@ impl Allocator { /// Panics if reserving space matching `layout` fails. pub fn alloc_layout(&self, layout: Layout) -> NonNull { #[cfg(all(feature = "track_allocations", not(feature = "disable_track_allocations")))] - self.num_alloc.fetch_add(1, std::sync::atomic::Ordering::SeqCst); + self.stats.record_allocation(); self.bump.alloc_layout(layout) } @@ -425,7 +416,7 @@ impl Allocator { ); #[cfg(all(feature = "track_allocations", not(feature = "disable_track_allocations")))] - self.num_alloc.fetch_add(1, std::sync::atomic::Ordering::SeqCst); + self.stats.record_allocation(); // 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. @@ -520,10 +511,7 @@ impl Allocator { #[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.stats.reset(); self.bump.reset(); } @@ -649,9 +637,7 @@ impl Allocator { 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), + stats: AllocationStats::default(), } } } diff --git a/crates/oxc_allocator/src/lib.rs b/crates/oxc_allocator/src/lib.rs index 2b161fbe838d1..4cadb00f845a2 100644 --- a/crates/oxc_allocator/src/lib.rs +++ b/crates/oxc_allocator/src/lib.rs @@ -42,6 +42,8 @@ mod from_raw_parts; pub mod hash_map; mod string_builder; mod take_in; +#[cfg(all(feature = "track_allocations", not(feature = "disable_track_allocations")))] +mod tracking; mod vec; mod vec2; diff --git a/crates/oxc_allocator/src/tracking.rs b/crates/oxc_allocator/src/tracking.rs new file mode 100644 index 0000000000000..9b71ead4746cb --- /dev/null +++ b/crates/oxc_allocator/src/tracking.rs @@ -0,0 +1,94 @@ +//! Allocation tracking. +//! +//! This module is only loaded when `track_allocations` feature is enabled. +//! This feature is only used in `tasks/track_memory_allocations`. +//! +//! Current implementation is unsound - see comment on [`get_stats_ref`] below. +//! It's OK to use it in our internal `tasks/track_memory_allocations` tool, +//! but we must take great care that this is NEVER enabled in any other circumstances. +//! +//! Even without the unsoundness, we don't want this enabled outside of `tasks/track_memory_allocations`, +//! as it imposes a performance cost on making allocations. +//! +//! The 2nd cargo feature `disable_track_allocations` is to ensure that compiling with `--all-features` +//! will not load this module. +//! +//! As soon as we replace `bumpalo` with our own arena allocator, we'll remove the hack from `get_stats_ref`, +//! and make this sound. + +use std::{ + ptr, + sync::atomic::{AtomicUsize, Ordering::SeqCst}, +}; + +use bumpalo::Bump; + +use crate::{Allocator, allocator::STATS_FIELD_OFFSET}; + +/// Counters of allocations and reallocations made in an [`Allocator`]. +// +// Note: These fields could be `Cell` instead of `AtomicUsize`, because `Allocator` should not +// be `Sync`. But currently it is (which is unsound!) because of other terrible hacks. +#[derive(Default)] +pub struct AllocationStats { + /// Number of allocations + num_alloc: AtomicUsize, + /// Number of reallocations + num_realloc: AtomicUsize, +} + +impl AllocationStats { + /// Record that an allocation was made. + pub(crate) fn record_allocation(&self) { + self.num_alloc.fetch_add(1, SeqCst); + } + + /// Record that a reallocation was made. + pub(crate) fn record_reallocation(&self) { + self.num_realloc.fetch_add(1, SeqCst); + } + + /// Reset allocation counters. + pub(crate) fn reset(&self) { + self.num_alloc.store(0, SeqCst); + self.num_realloc.store(0, SeqCst); + } +} + +impl Allocator { + /// Get number of allocations and reallocations made in this [`Allocator`]. + #[doc(hidden)] + pub fn get_allocation_stats(&self) -> (usize, usize) { + let num_alloc = self.stats.num_alloc.load(SeqCst); + let num_realloc = self.stats.num_realloc.load(SeqCst); + (num_alloc, num_realloc) + } +} + +/// Get reference to [`AllocationStats`] for a [`Bump`]. +/// +/// # SAFETY +/// +/// Caller must guarantee that the `Bump` provided to this function is wrapped in an [`Allocator`]. +/// +/// In Oxc, we never use `Bump` alone, without it being wrapped in an `Allocator`. +/// However, we have no static guarantee of this relationship between `Bump` and `Allocator`, +/// so it's usually impossible for callers to proveably satisfy the safety requirements of this method. +/// +/// Even if the `Bump` *is* wrapped in an `Allocator`, this may still be UB, as we project beyond +/// the bounds of the `&Bump`. Certainly stacked borrows memory model says this is UB, though it's unclear +/// to me (@overlookmotel) whether stacked borrows is unnecessarily strict on this point. +/// +/// +/// This function (and the `track_allocations` feature in general) must only be used for internal tools, +/// and must NEVER be compiled in production code. +pub unsafe fn get_stats_ref(bump: &Bump) -> &AllocationStats { + // We assume the `Bump` is wrapped in an `Allocator`. We can therefore get a pointer to the `stats` + // field of `Allocator` from the memory location of the `Bump`. + // SAFETY: This is UNSOUND. See above. + unsafe { + let stats_ptr = + ptr::from_ref(bump).byte_offset(STATS_FIELD_OFFSET).cast::(); + stats_ptr.as_ref().unwrap_unchecked() + } +} diff --git a/tasks/track_memory_allocations/src/lib.rs b/tasks/track_memory_allocations/src/lib.rs index b85995200bd0f..cbe5deee7e588 100644 --- a/tasks/track_memory_allocations/src/lib.rs +++ b/tasks/track_memory_allocations/src/lib.rs @@ -119,13 +119,9 @@ pub fn run() -> Result<(), io::Error> { let sys_allocs = NUM_ALLOC.load(SeqCst); let sys_reallocs = NUM_REALLOC.load(SeqCst); #[cfg(not(feature = "is_all_features"))] - let arena_allocs = allocator.num_alloc.load(SeqCst); + let (arena_allocs, arena_reallocs) = allocator.get_allocation_stats(); #[cfg(feature = "is_all_features")] - let arena_allocs = 0; - #[cfg(not(feature = "is_all_features"))] - let arena_reallocs = allocator.num_realloc.load(SeqCst); - #[cfg(feature = "is_all_features")] - let arena_reallocs = 0; + let (arena_allocs, arena_reallocs) = (0, 0); let arena_bytes = allocator.used_bytes(); let s = format!(