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>

Author: pront

changelog.d/fix_allocation_tracing_dealloc_panic.fix.md

@@ -0,0 +1 @@
+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.

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,135 @@ 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 sentinel_does_not_collide_with_registered_ids() {
+        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 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_ne!(
+            raw_id, UNTRACKED,
+            "header must not be UNTRACKED when tracking is on"
+        );
+
+        unsafe { allocator.dealloc(ptr, layout) };
+    }
+
+    /// Verify the UNTRACED sentinel dealloc path: manually write UNTRACED
+    /// into the header and confirm dealloc skips trace_deallocation.
+    #[test]
+    #[serial]
+    fn untraced_sentinel_dealloc_skips_tracing() {
+        let allocator = GroupedTraceableAllocator::new(System, CountingTracer::new());
+        let layout = Layout::from_size_align(64, 8).unwrap();
+
+        let _guard = TrackingGuard::enable();
+        let (wrapped_layout, offset) = get_wrapped_layout(layout);
+        let ptr = unsafe { System.alloc(wrapped_layout) };
+        assert!(!ptr.is_null());
+
+        unsafe { ptr.add(offset).cast::<u8>().write(UNTRACED) };
+
+        unsafe { allocator.dealloc(ptr, layout) };
+
+        assert_eq!(allocator.tracer.deallocs.load(Ordering::Relaxed), 0);
+    }
+}