bench(profiling): measure live heap tracking paths

Author: realFlowControl

profiling/Cargo.toml

@@ -61,6 +61,10 @@ perfcnt = "0.8.0"
 name = "stack_walking"
 harness = false
 
+[[bench]]
+name = "heap_live_tracking"
+harness = false
+
 [features]
 default = ["io_profiling"]
 debug_stats = []

profiling/benches/heap_live_tracking.rs

@@ -0,0 +1,115 @@
+use core::cell::{Cell, UnsafeCell};
+use criterion::{black_box, criterion_group, criterion_main, BatchSize, Criterion};
+
+// Compile the production tracker without linking the PHP extension executable.
+#[allow(dead_code)]
+#[path = "../src/profiling/live_heap.rs"]
+mod live_heap;
+use live_heap::{LiveHeapTracker, LocalLiveHeapTracker};
+
+const BASE_ADDRESS: usize = 0x1000_0000;
+const TRACKED_ALLOCATIONS: usize = 2048;
+const ADDRESS_STRIDE: usize = 64;
+
+fn address(index: usize) -> usize {
+    BASE_ADDRESS + index % TRACKED_ALLOCATIONS * ADDRESS_STRIDE
+}
+
+fn untracked_address(index: usize) -> usize {
+    address(index) + 8
+}
+
+struct Tracker {
+    shared: LiveHeapTracker<[usize; 4]>,
+    local: UnsafeCell<LocalLiveHeapTracker>,
+}
+
+impl Tracker {
+    fn new() -> Self {
+        Self {
+            shared: LiveHeapTracker::new(),
+            local: UnsafeCell::new(LocalLiveHeapTracker::new()),
+        }
+    }
+
+    fn track(&self, ptr: usize, sample: [usize; 4]) -> bool {
+        // Criterion invokes setup and measured routines sequentially.
+        unsafe { (&mut *self.local.get()).track(&self.shared, ptr, sample) }
+    }
+
+    fn untrack(&self, ptr: usize) -> Option<[usize; 4]> {
+        // SAFETY: same sequential access as `track`.
+        unsafe { (&mut *self.local.get()).untrack(&self.shared, ptr) }
+    }
+}
+
+fn populated_tracker() -> Tracker {
+    let tracker = Tracker::new();
+    for index in 0..TRACKED_ALLOCATIONS {
+        assert!(tracker.track(address(index), [0; 4]));
+    }
+    tracker
+}
+
+fn benchmark(c: &mut Criterion) {
+    let mut group = c.benchmark_group("heap_live_tracking");
+
+    {
+        let tracker = populated_tracker();
+        let next = Cell::new(0);
+        group.bench_function("allocate_tracked", |b| {
+            b.iter_batched(
+                || {
+                    let index = next.get();
+                    next.set(index + 1);
+                    let ptr = untracked_address(index);
+                    let _ = tracker.untrack(ptr);
+                    (ptr, [0; 4])
+                },
+                |(ptr, sample)| black_box(tracker.track(ptr, sample)),
+                BatchSize::PerIteration,
+            )
+        });
+    }
+
+    {
+        let tracker = populated_tracker();
+        let next = Cell::new(0);
+        group.bench_function("free_tracked", |b| {
+            b.iter_batched(
+                || {
+                    let index = next.get();
+                    next.set(index + 1);
+                    let ptr = untracked_address(index);
+                    assert!(tracker.track(ptr, [0; 4]));
+                    ptr
+                },
+                |ptr| black_box(tracker.untrack(ptr)),
+                BatchSize::PerIteration,
+            )
+        });
+    }
+
+    {
+        let tracker = populated_tracker();
+        let next = Cell::new(0);
+        group.bench_function("free_untracked", |b| {
+            b.iter_batched(
+                || {
+                    let index = next.get();
+                    next.set(index + 1);
+                    let ptr = untracked_address(index);
+                    let _ = tracker.untrack(ptr);
+                    ptr
+                },
+                |ptr| black_box(tracker.untrack(ptr)),
+                BatchSize::PerIteration,
+            )
+        });
+    }
+
+    group.finish();
+}
+
+criterion_group!(benches, benchmark);
+criterion_main!(benches);

profiling/src/allocation/mod.rs

@@ -5,6 +5,7 @@ pub use profiling_stats::*;
 use crate::bindings::{self as zend};
 use crate::config::SystemSettings;
 use crate::module_globals;
+use crate::profiling::live_heap::LocalLiveHeapTracker;
 use crate::profiling::Profiler;
 use crate::{RefCellExt, REQUEST_LOCALS};
 use core::cell::Cell;
@@ -98,6 +99,7 @@ pub struct AllocationProfilingStats {
     rng: ThreadRng,
     #[cfg(not(php_zts))]
     rng: StdRng,
+    live_heap: LocalLiveHeapTracker,
 }
 
 impl AllocationProfilingStats {
@@ -111,6 +113,7 @@ impl AllocationProfilingStats {
             rng: rand::rng(),
             #[cfg(not(php_zts))]
             rng: StdRng::from_os_rng(),
+            live_heap: LocalLiveHeapTracker::new(),
         };
         stats.next_sampling_interval();
         stats

profiling/src/allocation/profiling_stats.rs

@@ -4,6 +4,7 @@
 //! can be used but expose a relatively safe API.
 
 use super::{AllocationProfilingStats, ALLOCATION_PROFILING_INTERVAL};
+use crate::profiling::live_heap::LiveHeapTracker;
 use libc::size_t;
 use std::mem::MaybeUninit;
 
@@ -101,6 +102,28 @@ pub fn allocation_profiling_stats_should_collect(len: size_t) -> bool {
     unsafe { allocation_profiling_stats_mut(f) }
 }
 
+pub(crate) fn live_heap_track<T>(tracker: &LiveHeapTracker<T>, ptr: usize, sample: T) -> bool {
+    // SAFETY: allocation stats are initialized before allocation hooks run,
+    // and the closure does not retain or recursively borrow them.
+    unsafe {
+        allocation_profiling_stats_mut(|stats| {
+            stats
+                .assume_init_mut()
+                .live_heap
+                .track(tracker, ptr, sample)
+        })
+    }
+}
+
+pub(crate) fn live_heap_untrack<T>(tracker: &LiveHeapTracker<T>, ptr: usize) -> Option<T> {
+    // SAFETY: same lifecycle and borrowing guarantees as `live_heap_track`.
+    unsafe {
+        allocation_profiling_stats_mut(|stats| {
+            stats.assume_init_mut().live_heap.untrack(tracker, ptr)
+        })
+    }
+}
+
 /// Initializes the allocation profiler's globals.
 ///
 /// # Safety

profiling/src/profiling/live_heap.rs

@@ -0,0 +1,90 @@
+use dashmap::DashMap;
+use rustc_hash::FxBuildHasher;
+use std::sync::atomic::{AtomicUsize, Ordering};
+
+/// Maximum number of allocations to track for live heap profiling.
+const MAX_SIZE: usize = 4096;
+
+/// Tracks live heap samples by allocation address. FxHasher spreads sequential
+/// ZendMM addresses across DashMap's shards without hashing already-randomized
+/// pointer bytes with SipHash.
+pub(crate) struct LiveHeapTracker<T> {
+    allocations: DashMap<usize, T, FxBuildHasher>,
+    count: AtomicUsize,
+}
+
+impl<T> LiveHeapTracker<T> {
+    pub(crate) fn new() -> Self {
+        Self {
+            allocations: DashMap::with_hasher(FxBuildHasher),
+            count: AtomicUsize::new(0),
+        }
+    }
+
+    pub(crate) fn len(&self) -> usize {
+        self.count.load(Ordering::Relaxed)
+    }
+
+    pub(crate) fn clear(&self) {
+        self.allocations.clear();
+        self.count.store(0, Ordering::Relaxed);
+    }
+
+    fn track(&self, ptr: usize, sample: T) -> bool {
+        // Best-effort cap: in ZTS the count check and insert still race, so
+        // the map can briefly exceed MAX_SIZE.
+        if self.len() >= MAX_SIZE {
+            return false;
+        }
+
+        if self.allocations.insert(ptr, sample).is_none() {
+            self.count.fetch_add(1, Ordering::Relaxed);
+        }
+        true
+    }
+
+    fn untrack(&self, ptr: usize) -> Option<T> {
+        let result = self.allocations.remove(&ptr).map(|(_, sample)| sample);
+        if result.is_some() {
+            self.count.fetch_sub(1, Ordering::Relaxed);
+        }
+        result
+    }
+}
+
+impl<T: Clone> LiveHeapTracker<T> {
+    pub(crate) fn snapshot(&self) -> Vec<T> {
+        self.allocations
+            .iter()
+            .map(|entry| entry.value().clone())
+            .collect()
+    }
+}
+
+impl<T> Default for LiveHeapTracker<T> {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+pub(crate) struct LocalLiveHeapTracker;
+
+impl LocalLiveHeapTracker {
+    pub(crate) const fn new() -> Self {
+        Self
+    }
+
+    pub(crate) fn track<T>(&mut self, tracker: &LiveHeapTracker<T>, ptr: usize, sample: T) -> bool {
+        tracker.track(ptr, sample)
+    }
+
+    pub(crate) fn untrack<T>(&mut self, tracker: &LiveHeapTracker<T>, ptr: usize) -> Option<T> {
+        tracker.untrack(ptr)
+    }
+}
+
+impl Default for LocalLiveHeapTracker {
+    fn default() -> Self {
+        Self::new()
+    }
+}