perf(profiling): remove indirect allocator forwarding

Forward PHP 8.4+ allocations and frees from a single ZendMMState read instead of loading a function pointer, making an indirect Rust call, and reloading the state. Preserve neighboring custom allocator support with a predictable previous-handler check.

Use unchecked heap extraction in the callbacks: rinit stores the heap before they can be invoked, and rshutdown removes them before clearing it. This avoids an Option discriminant check on every allocation and free.

Across six 60-second runs per binary, mean allocation throughput increased from 9,366,444/s to 10,018,262/s (+6.96%) and median throughput increased by 7.56%. A 60-second native sample reduced known forwarding-path self time from 3.599% to 2.438% (-32.3% relative).

Validation: cargo test (22 passed); allocation sampling-distance, memory-peak, and GC PHPTs passed.

https://datadoghq.atlassian.net/browse/PROF-15506

Author: realFlowControl

profiling/src/allocation/allocation_ge84.rs

@@ -34,21 +34,11 @@ pub struct ZendMMState {
     /// The engine's previous custom shutdown function, if there is one.
     prev_custom_mm_shutdown: Option<zend::VmMmCustomShutdownFn>,
     /// Safety: this function pointer is only allowed to point to
-    /// `alloc_prof_prev_alloc()` when at the same time the
-    /// `ZEND_MM_STATE.prev_custom_mm_alloc` is initialised to a valid function
-    /// pointer, otherwise there will be dragons.
-    alloc: unsafe fn(size_t) -> *mut c_void,
-    /// Safety: this function pointer is only allowed to point to
     /// `alloc_prof_prev_realloc()` when at the same time the
     /// `ZEND_MM_STATE.prev_custom_mm_realloc` is initialised to a valid
     /// function pointer, otherwise there will be dragons.
     realloc: unsafe fn(*mut c_void, size_t) -> *mut c_void,
     /// Safety: this function pointer is only allowed to point to
-    /// `alloc_prof_prev_free()` when at the same time the
-    /// `ZEND_MM_STATE.prev_custom_mm_free` is initialised to a valid function
-    /// pointer, otherwise there will be dragons.
-    free: unsafe fn(*mut c_void),
-    /// Safety: this function pointer is only allowed to point to
     /// `alloc_prof_prev_gc()` when at the same time the
     /// `ZEND_MM_STATE.prev_custom_mm_gc` is initialised to a valid function
     /// pointer, otherwise there will be dragons.
@@ -78,9 +68,7 @@ impl ZendMMState {
             prev_custom_mm_free: None,
             prev_custom_mm_gc: None,
             prev_custom_mm_shutdown: None,
-            alloc: super::alloc_prof_panic_alloc,
             realloc: super::alloc_prof_panic_realloc,
-            free: super::alloc_prof_panic_free,
             gc: alloc_prof_panic_gc,
             shutdown: alloc_prof_panic_shutdown,
         }
@@ -127,14 +115,10 @@ pub fn alloc_prof_rinit(heap_live_enabled: bool) {
                     ptr::addr_of_mut!(zend_mm_state.prev_custom_mm_shutdown),
                 );
             }
-            zend_mm_state.alloc = alloc_prof_prev_alloc;
-            zend_mm_state.free = alloc_prof_prev_free;
             zend_mm_state.realloc = alloc_prof_prev_realloc;
             zend_mm_state.gc = alloc_prof_prev_gc;
             zend_mm_state.shutdown = alloc_prof_prev_shutdown;
         } else {
-            zend_mm_state.alloc = alloc_prof_orig_alloc;
-            zend_mm_state.free = alloc_prof_orig_free;
             zend_mm_state.realloc = alloc_prof_orig_realloc;
             zend_mm_state.gc = alloc_prof_orig_gc;
             zend_mm_state.shutdown = alloc_prof_orig_shutdown;
@@ -303,7 +287,7 @@ unsafe fn alloc_prof_malloc_impl(len: size_t) -> *mut c_void {
     #[cfg(feature = "debug_stats")]
     ALLOCATION_PROFILING_SIZE.fetch_add(len as u64, Relaxed);
 
-    let ptr = tls_zend_mm_state_get!(alloc)(len);
+    let ptr = alloc_prof_forward_alloc(len);
 
     // during startup, minit, rinit, ... current_execute_data is null
     // we are only interested in allocations during userland operations
@@ -318,27 +302,19 @@ unsafe fn alloc_prof_malloc_impl(len: size_t) -> *mut c_void {
     ptr
 }
 
-unsafe fn alloc_prof_prev_alloc(len: size_t) -> *mut c_void {
-    // Safety: `ZEND_MM_STATE.prev_custom_mm_alloc` will be initialised in
-    // `alloc_prof_rinit()` and only point to this function when
-    // `prev_custom_mm_alloc` is also initialised.
-    // Note: We use `.unwrap()` instead of `.unwrap_unchecked()` here because a
-    // neighboring extension could misbehave. If that happens, we want a proper
-    // panic with backtrace for debugging rather than undefined behavior.
-    let alloc = tls_zend_mm_state_get!(prev_custom_mm_alloc).unwrap();
-    #[cfg(php_debug)]
-    {
-        alloc(len, ptr::null(), 0, ptr::null(), 0)
+#[inline(always)]
+unsafe fn alloc_prof_forward_alloc(len: size_t) -> *mut c_void {
+    let state = tls_zend_mm_state_copy!();
+    if let Some(alloc) = state.prev_custom_mm_alloc {
+        #[cfg(php_debug)]
+        return alloc(len, ptr::null(), 0, ptr::null(), 0);
+        #[cfg(not(php_debug))]
+        return alloc(len);
     }
-    #[cfg(not(php_debug))]
-    alloc(len)
-}
 
-unsafe fn alloc_prof_orig_alloc(len: size_t) -> *mut c_void {
-    // Safety: `ZEND_MM_STATE.heap` will be initialised in `alloc_prof_rinit()` and custom ZendMM
-    // handlers only point to this function after successful init. Using `unwrap_unchecked()` is
-    // safe here as we have full control over ZendMM with no neighboring extensions.
-    let heap = tls_zend_mm_state_get!(heap).unwrap_unchecked();
+    // SAFETY: this callback is only invoked after rinit stores the heap and
+    // before rshutdown clears it.
+    let heap = state.heap.unwrap_unchecked();
     #[cfg(php_debug)]
     return zend::_zend_mm_alloc(heap, len, ptr::null(), 0, ptr::null(), 0);
     #[cfg(not(php_debug))]
@@ -375,7 +351,7 @@ fn alloc_prof_free_handler(heap_live_enabled: bool) -> zend::VmMmCustomFreeFn {
 
 #[cfg(not(php_debug))]
 unsafe extern "C" fn alloc_prof_free_noop(ptr: *mut c_void) {
-    tls_zend_mm_state_get!(free)(ptr);
+    alloc_prof_forward_free(ptr);
 }
 
 #[cfg(php_debug)]
@@ -386,7 +362,7 @@ unsafe extern "C" fn alloc_prof_free_noop(
     _orig_file: *const c_char,
     _orig_line: c_uint,
 ) {
-    tls_zend_mm_state_get!(free)(ptr);
+    alloc_prof_forward_free(ptr);
 }
 
 #[inline(always)]
@@ -395,30 +371,22 @@ unsafe fn alloc_prof_free_impl(ptr: *mut c_void) {
     if !ptr.is_null() {
         untrack_allocation(ptr);
     }
-    tls_zend_mm_state_get!(free)(ptr);
+    alloc_prof_forward_free(ptr);
 }
 
-unsafe fn alloc_prof_prev_free(ptr: *mut c_void) {
-    // Safety: `ZEND_MM_STATE.prev_custom_mm_free` will be initialised in
-    // `alloc_prof_rinit()` and only point to this function when
-    // `prev_custom_mm_free` is also initialised.
-    // Note: We use `.unwrap()` instead of `.unwrap_unchecked()` here because a
-    // neighboring extension could misbehave. If that happens, we want a proper
-    // panic with backtrace for debugging rather than undefined behavior.
-    let free = tls_zend_mm_state_get!(prev_custom_mm_free).unwrap();
-    #[cfg(php_debug)]
-    {
-        free(ptr, core::ptr::null(), 0, core::ptr::null(), 0)
+#[inline(always)]
+unsafe fn alloc_prof_forward_free(ptr: *mut c_void) {
+    let state = tls_zend_mm_state_copy!();
+    if let Some(free) = state.prev_custom_mm_free {
+        #[cfg(php_debug)]
+        return free(ptr, core::ptr::null(), 0, core::ptr::null(), 0);
+        #[cfg(not(php_debug))]
+        return free(ptr);
     }
-    #[cfg(not(php_debug))]
-    free(ptr)
-}
 
-unsafe fn alloc_prof_orig_free(ptr: *mut c_void) {
-    // Safety: `ZEND_MM_STATE.heap` will be initialised in `alloc_prof_rinit()` and custom ZendMM
-    // handlers only point to this function after successful init. Using `unwrap_unchecked()` is
-    // safe here as we have full control over ZendMM with no neighboring extensions.
-    let heap = tls_zend_mm_state_get!(heap).unwrap_unchecked();
+    // SAFETY: this callback is only invoked after rinit stores the heap and
+    // before rshutdown clears it.
+    let heap = state.heap.unwrap_unchecked();
     #[cfg(php_debug)]
     return zend::_zend_mm_free(heap, ptr, core::ptr::null(), 0, core::ptr::null(), 0);
     #[cfg(not(php_debug))]

profiling/src/allocation/mod.rs

@@ -275,6 +275,7 @@ fn initialization_panic() -> ! {
     panic!("Allocation profiler was not initialized properly. Please fill an issue stating the PHP version and the backtrace from this panic.");
 }
 
+#[cfg(not(php_zend_mm_set_custom_handlers_ex))]
 unsafe fn alloc_prof_panic_alloc(_len: size_t) -> *mut c_void {
     initialization_panic();
 }
@@ -283,6 +284,7 @@ unsafe fn alloc_prof_panic_realloc(_prev_ptr: *mut c_void, _len: size_t) -> *mut
     initialization_panic();
 }
 
+#[cfg(not(php_zend_mm_set_custom_handlers_ex))]
 unsafe fn alloc_prof_panic_free(_ptr: *mut c_void) {
     initialization_panic();
 }