fix(observability): fix panic in allocation tracing dealloc path (#25136)

* fix(observability): fix panic in allocation tracing when deallocating pre-tracking memory

When `--allocation-tracing` is enabled at runtime, the custom allocator
wraps every allocation with an extra byte to store the allocation group
ID. Previously, allocations made before tracking was enabled used the
original (unwrapped) layout. When those were later freed, `dealloc`
read an out-of-bounds byte as the group ID, hitting
`NonZeroU8::new_unchecked(0)` -- undefined behavior that recent Rust
toolchains (>= ~1.78) turn into an abort in debug builds.

Additionally, reentrant allocations (wrapped layout but tracing closure
skipped) left the group ID header uninitialized, causing misattributed
deallocations and skewed per-group memory accounting.

Fix: always allocate with the wrapped layout, regardless of whether
tracking is currently enabled. The group ID header byte is set to:
- UNTRACKED (0): tracking was off at allocation time
- UNTRACED (u8::MAX): tracking was on but the tracing closure was
  skipped due to reentrancy
- A real group ID (1..254): normal traced allocation

On deallocation, all paths free with the wrapped layout (which is
always correct now). UNTRACKED and UNTRACED skip `trace_deallocation`
to keep per-group accounting balanced.

This eliminates:
- The original panic/UB from `NonZero::new_unchecked(0)`
- Layout mismatches for pre-tracking allocations (including realloc)
- Accounting skew from uninitialized headers on reentrant allocations

This bug has been latent since #15221 (Nov 2022) which introduced the
runtime toggle.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

* add authors

* fix(observability): strengthen allocation tracing tests

Harden test coverage based on code review feedback:
- Assert exact ROOT group ID and trace counters in tracked alloc test
- Add end-to-end reentrant allocation test that exercises the real
  reentrancy path via thread-local borrow
- Remove redundant manual-sentinel test now covered by reentrant test
- Rename sentinel collision test to reflect what it actually checks

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

---------

Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Author: pront

changelog.d/fix_allocation_tracing_dealloc_panic.fix.md

@@ -0,0 +1,3 @@
+Fixed a panic (abort) when running with `--allocation-tracing` in debug builds, caused by deallocating memory that was allocated before tracking was enabled. Also fixed per-group memory accounting skew for reentrant allocations whose tracing closure was skipped, which left the group ID header uninitialized and caused deallocations to be attributed to wrong groups.
+
+authors: pront

src/internal_telemetry/allocations/allocator/tracing_allocator.rs

@@ -9,8 +9,24 @@ use super::{
 };
 use crate::internal_telemetry::allocations::TRACK_ALLOCATIONS;
 
+/// Header value for allocations made while tracking is disabled.
+/// On deallocation we free with the wrapped layout but skip tracing.
+const UNTRACKED: u8 = 0;
+
+/// Header value for allocations whose tracing closure was skipped due to
+/// reentrancy. `register()` never hands out `u8::MAX`, so this cannot
+/// collide with a real group ID. On deallocation we free with the wrapped
+/// layout but skip `trace_deallocation` to keep accounting balanced.
+const UNTRACED: u8 = u8::MAX;
+
 /// A tracing allocator that groups allocation events by groups.
 ///
+/// Every allocation made through this allocator uses the "wrapped" layout:
+/// the requested layout extended by one byte to store an allocation group
+/// ID. This byte is always present regardless of whether tracking is
+/// enabled, which guarantees that `dealloc` can always read a valid header
+/// and free with the correct (wrapped) layout.
+///
 /// This allocator can only be used when specified via `#[global_allocator]`.
 pub struct GroupedTraceableAllocator<A, T> {
     allocator: A,
@@ -29,11 +45,6 @@ unsafe impl<A: GlobalAlloc, T: Tracer> GlobalAlloc for GroupedTraceableAllocator
     #[inline]
     unsafe fn alloc(&self, object_layout: Layout) -> *mut u8 {
         unsafe {
-            if !TRACK_ALLOCATIONS.load(Ordering::Relaxed) {
-                return self.allocator.alloc(object_layout);
-            }
-
-            // Allocate our wrapped layout and make sure the allocation succeeded.
             let (actual_layout, offset_to_group_id) = get_wrapped_layout(object_layout);
             let actual_ptr = self.allocator.alloc(actual_layout);
             if actual_ptr.is_null() {
@@ -42,6 +53,16 @@ unsafe impl<A: GlobalAlloc, T: Tracer> GlobalAlloc for GroupedTraceableAllocator
 
             let group_id_ptr = actual_ptr.add(offset_to_group_id).cast::<u8>();
 
+            if !TRACK_ALLOCATIONS.load(Ordering::Relaxed) {
+                group_id_ptr.write(UNTRACKED);
+                return actual_ptr;
+            }
+
+            // Write the untraced sentinel so that `dealloc` always finds a
+            // valid header, even when the tracing closure below is skipped
+            // due to reentrancy.
+            group_id_ptr.write(UNTRACED);
+
             let object_size = object_layout.size();
 
             try_with_suspended_allocation_group(
@@ -58,19 +79,19 @@ unsafe impl<A: GlobalAlloc, T: Tracer> GlobalAlloc for GroupedTraceableAllocator
     #[inline]
     unsafe fn dealloc(&self, object_ptr: *mut u8, object_layout: Layout) {
         unsafe {
-            if !TRACK_ALLOCATIONS.load(Ordering::Relaxed) {
-                self.allocator.dealloc(object_ptr, object_layout);
-                return;
-            }
-            // Regenerate the wrapped layout so we know where we have to look, as the pointer we've given relates to the
-            // requested layout, not the wrapped layout that was actually allocated.
             let (wrapped_layout, offset_to_group_id) = get_wrapped_layout(object_layout);
-
             let raw_group_id = object_ptr.add(offset_to_group_id).cast::<u8>().read();
 
-            // Deallocate before tracking, just to make sure we're reclaiming memory as soon as possible.
+            // Always free with the wrapped layout since all allocations
+            // (tracked or not) use it.
             self.allocator.dealloc(object_ptr, wrapped_layout);
 
+            // Skip tracing for untracked (tracking was off) and untraced
+            // (reentrant, closure skipped) allocations.
+            if raw_group_id == UNTRACKED || raw_group_id == UNTRACED {
+                return;
+            }
+
             let object_size = object_layout.size();
             let source_group_id = AllocationGroupId::from_raw(raw_group_id);
 
@@ -96,3 +117,158 @@ fn get_wrapped_layout(object_layout: Layout) -> (Layout, usize) {
 
     (actual_layout.pad_to_align(), offset_to_group_id)
 }
+
+#[cfg(test)]
+mod tests {
+    use std::{
+        alloc::{GlobalAlloc, Layout, System},
+        sync::atomic::{AtomicUsize, Ordering},
+    };
+
+    use serial_test::serial;
+
+    use super::*;
+    use crate::internal_telemetry::allocations::allocator::{
+        token::AllocationGroupId, tracer::Tracer,
+    };
+
+    /// RAII guard that enables `TRACK_ALLOCATIONS` on creation and
+    /// restores it to `false` on drop, ensuring cleanup even if the
+    /// test panics.
+    struct TrackingGuard;
+
+    impl TrackingGuard {
+        fn enable() -> Self {
+            TRACK_ALLOCATIONS.store(true, Ordering::Release);
+            Self
+        }
+    }
+
+    impl Drop for TrackingGuard {
+        fn drop(&mut self) {
+            TRACK_ALLOCATIONS.store(false, Ordering::Release);
+        }
+    }
+
+    struct CountingTracer {
+        allocs: AtomicUsize,
+        deallocs: AtomicUsize,
+    }
+
+    impl CountingTracer {
+        const fn new() -> Self {
+            Self {
+                allocs: AtomicUsize::new(0),
+                deallocs: AtomicUsize::new(0),
+            }
+        }
+    }
+
+    impl Tracer for CountingTracer {
+        fn trace_allocation(&self, _size: usize, _group_id: AllocationGroupId) {
+            self.allocs.fetch_add(1, Ordering::Relaxed);
+        }
+
+        fn trace_deallocation(&self, _size: usize, _source_group_id: AllocationGroupId) {
+            self.deallocs.fetch_add(1, Ordering::Relaxed);
+        }
+    }
+
+    #[test]
+    fn sentinels_do_not_collide_with_root_id() {
+        assert_eq!(UNTRACED, u8::MAX);
+        assert_ne!(UNTRACED, AllocationGroupId::ROOT.as_raw());
+        assert_ne!(UNTRACKED, AllocationGroupId::ROOT.as_raw());
+    }
+
+    /// Allocations made while tracking is disabled get UNTRACKED (0) in the
+    /// header. Deallocating them (whether tracking is on or off) must use
+    /// the wrapped layout and skip tracing.
+    #[test]
+    #[serial]
+    fn untracked_alloc_dealloc_skips_tracing() {
+        let allocator = GroupedTraceableAllocator::new(System, CountingTracer::new());
+        let layout = Layout::from_size_align(64, 8).unwrap();
+
+        // Tracking starts off (default state). Allocate while disabled.
+        let ptr = unsafe { allocator.alloc(layout) };
+        assert!(!ptr.is_null());
+
+        // Header must be UNTRACKED.
+        let (_, offset) = get_wrapped_layout(layout);
+        let raw_id = unsafe { ptr.add(offset).cast::<u8>().read() };
+        assert_eq!(raw_id, UNTRACKED);
+
+        // Enable tracking, then dealloc -- must not panic, no trace events.
+        let _guard = TrackingGuard::enable();
+        unsafe { allocator.dealloc(ptr, layout) };
+
+        assert_eq!(allocator.tracer.allocs.load(Ordering::Relaxed), 0);
+        assert_eq!(allocator.tracer.deallocs.load(Ordering::Relaxed), 0);
+    }
+
+    /// Tracked allocation: header is a valid non-zero, non-sentinel group
+    /// ID, tracing fires for both alloc and dealloc, and dealloc completes
+    /// without panic.
+    #[test]
+    #[serial]
+    fn tracked_alloc_dealloc_does_not_panic() {
+        let allocator = GroupedTraceableAllocator::new(System, CountingTracer::new());
+        let layout = Layout::from_size_align(64, 8).unwrap();
+
+        let _guard = TrackingGuard::enable();
+        let ptr = unsafe { allocator.alloc(layout) };
+        assert!(!ptr.is_null());
+
+        let (_, offset) = get_wrapped_layout(layout);
+        let raw_id = unsafe { ptr.add(offset).cast::<u8>().read() };
+        assert_eq!(
+            raw_id,
+            AllocationGroupId::ROOT.as_raw(),
+            "header must be the ROOT group ID"
+        );
+        assert_eq!(allocator.tracer.allocs.load(Ordering::Relaxed), 1);
+
+        unsafe { allocator.dealloc(ptr, layout) };
+
+        assert_eq!(allocator.tracer.deallocs.load(Ordering::Relaxed), 1);
+    }
+
+    /// End-to-end reentrant allocation: hold a mutable borrow on the
+    /// thread-local group stack so `try_with_suspended_allocation_group`
+    /// skips the tracing closure. The header must be UNTRACED and both
+    /// trace counters must stay at zero through alloc + dealloc.
+    #[test]
+    #[serial]
+    fn reentrant_alloc_writes_untraced_and_skips_tracing() {
+        use crate::internal_telemetry::allocations::allocator::token::LOCAL_ALLOCATION_GROUP_STACK;
+
+        let allocator = GroupedTraceableAllocator::new(System, CountingTracer::new());
+        let layout = Layout::from_size_align(64, 8).unwrap();
+
+        let _guard = TrackingGuard::enable();
+
+        // Hold a mutable borrow to simulate reentrancy -- any nested
+        // `try_borrow_mut` inside `try_with_suspended_allocation_group`
+        // will fail, causing the tracing closure to be skipped.
+        LOCAL_ALLOCATION_GROUP_STACK.with(|group_stack| {
+            let _borrow = group_stack.borrow_mut();
+
+            let ptr = unsafe { allocator.alloc(layout) };
+            assert!(!ptr.is_null());
+
+            let (_, offset) = get_wrapped_layout(layout);
+            let raw_id = unsafe { ptr.add(offset).cast::<u8>().read() };
+            assert_eq!(
+                raw_id, UNTRACED,
+                "reentrant alloc must write UNTRACED sentinel"
+            );
+
+            assert_eq!(allocator.tracer.allocs.load(Ordering::Relaxed), 0);
+
+            unsafe { allocator.dealloc(ptr, layout) };
+
+            assert_eq!(allocator.tracer.deallocs.load(Ordering::Relaxed), 0);
+        });
+    }
+}