perf(prof): speed up interrupt_count checks (#4074)

Co-authored-by: Levi Morrison <levi.morrison@datadoghq.com>

Author: realFlowControl

profiling/src/allocation/mod.rs

@@ -199,9 +199,10 @@ pub fn collect_allocation(ptr: *mut c_void, len: size_t) {
         // Check if there's a pending time interrupt that we can handle now
         // instead of waiting for an interrupt handler. This is slightly more
         // accurate and efficient, win-win.
-        let interrupt_count = REQUEST_LOCALS
-            .try_with_borrow(|locals| locals.interrupt_count.swap(0, Ordering::SeqCst))
-            .unwrap_or(0);
+        // SAFETY: allocation samples are collected on an initialized PHP request thread.
+        let globals = unsafe { module_globals::get_profiler_globals() };
+        // SAFETY: the current thread's module globals are valid through GSHUTDOWN.
+        let interrupt_count = unsafe { (*globals).interrupt_count.swap(0, Ordering::Relaxed) };
 
         // SAFETY: execute_data was provided by the engine, and the profiler
         // doesn't mutate it.

profiling/src/capi.rs

@@ -46,7 +46,10 @@ extern "C" fn ddog_php_prof_trigger_time_sample() {
         if locals.system_settings().profiling_enabled {
             // Safety: only vm interrupts are stored there, or possibly null (edges only).
             if let Some(vm_interrupt) = unsafe { locals.vm_interrupt_addr.as_ref() } {
-                locals.interrupt_count.fetch_add(1, Ordering::SeqCst);
+                // SAFETY: this callback runs on an initialized PHP request thread.
+                let globals = unsafe { crate::module_globals::get_profiler_globals() };
+                // SAFETY: the current thread's module globals are valid through GSHUTDOWN.
+                unsafe { (*globals).interrupt_count.fetch_add(1, Ordering::Relaxed) };
                 vm_interrupt.store(true, Ordering::SeqCst);
             }
         }

profiling/src/lib.rs

@@ -425,7 +425,6 @@ pub struct RequestLocals {
     pub system_settings: ptr::NonNull<SystemSettings>,
     pub profiling_experimental_heap_live_enabled: bool,
 
-    pub interrupt_count: AtomicU32,
     pub vm_interrupt_addr: *const AtomicBool,
 }
 
@@ -450,7 +449,6 @@ impl Default for RequestLocals {
             tags: vec![],
             system_settings: SystemSettings::get(),
             profiling_experimental_heap_live_enabled: false,
-            interrupt_count: AtomicU32::new(0),
             vm_interrupt_addr: ptr::null_mut(),
         }
     }
@@ -738,8 +736,11 @@ extern "C" fn rinit(_type: c_int, _module_number: c_int) -> ZendResult {
             }
 
             if let Some(profiler) = Profiler::get() {
+                // SAFETY: PHP module globals are initialized for this request thread.
+                let globals = unsafe { module_globals::get_profiler_globals() };
                 let interrupt = VmInterrupt {
-                    interrupt_count_ptr: &locals.interrupt_count as *const AtomicU32,
+                    // SAFETY: `globals` is valid until this thread's GSHUTDOWN.
+                    interrupt_count_ptr: unsafe { ptr::addr_of!((*globals).interrupt_count) },
                     engine_ptr: locals.vm_interrupt_addr,
                 };
                 profiler.add_interrupt(interrupt);
@@ -795,8 +796,11 @@ extern "C" fn rshutdown(_type: c_int, _module_number: c_int) -> ZendResult {
         // and we don't need to optimize for that.
         if system_settings.profiling_enabled {
             if let Some(profiler) = Profiler::get() {
+                // SAFETY: PHP module globals remain initialized through RSHUTDOWN.
+                let globals = unsafe { module_globals::get_profiler_globals() };
                 let interrupt = VmInterrupt {
-                    interrupt_count_ptr: &locals.interrupt_count,
+                    // SAFETY: `globals` remains valid until this thread's GSHUTDOWN.
+                    interrupt_count_ptr: unsafe { ptr::addr_of!((*globals).interrupt_count) },
                     engine_ptr: locals.vm_interrupt_addr,
                 };
                 profiler.remove_interrupt(interrupt);

profiling/src/module_globals.rs

@@ -2,6 +2,7 @@ use crate::allocation;
 use core::cell::Cell;
 use core::ffi::c_void;
 use core::ptr;
+use core::sync::atomic::AtomicU32;
 
 #[cfg(php_zend_mm_set_custom_handlers_ex)]
 use crate::allocation::allocation_ge84::ZendMMState;
@@ -13,6 +14,12 @@ pub struct ProfilerGlobals {
     /// Wrapped in `Cell` to prevent torn reads/writes when allocation hooks
     /// are called re-entrantly during `rinit()`/`rshutdown()`.
     pub zend_mm_state: Cell<ZendMMState>,
+    /// Number of profiler time interrupts pending for this PHP thread.
+    ///
+    /// The profiler timer thread updates this through a pointer registered by
+    /// the PHP thread, so the value must remain atomic despite living in
+    /// thread-local PHP module globals.
+    pub interrupt_count: AtomicU32,
 }
 
 /// We need TSRM to call into GINIT and GSHUTDOWN to observe spawning and
@@ -29,6 +36,7 @@ pub static mut GLOBALS_ID: i32 = 0;
 #[cfg(not(php_zts))]
 pub static mut GLOBALS: ProfilerGlobals = ProfilerGlobals {
     zend_mm_state: Cell::new(ZendMMState::new()),
+    interrupt_count: AtomicU32::new(0),
 };
 
 #[cfg(php_zts)]
@@ -40,9 +48,13 @@ mod zts {
     }
 
     #[inline]
-    pub unsafe fn tsrmg_bulk(id: i32) -> *mut c_void {
-        let tls = tsrm_get_ls_cache() as *mut *mut *mut c_void;
-        let storage = *tls; // void** storage
+    pub unsafe fn get_ls_cache() -> *mut c_void {
+        tsrm_get_ls_cache()
+    }
+
+    #[inline]
+    pub unsafe fn tsrmg_bulk(ls_cache: *mut c_void, id: i32) -> *mut c_void {
+        let storage = *(ls_cache as *mut *mut *mut c_void); // void** storage
 
         // TSRM_UNSHUFFLE_RSRC_ID(id) is just `id - 1`.
         let idx = (id - 1) as usize;
@@ -51,6 +63,27 @@ mod zts {
     }
 }
 
+#[cfg(php_zts)]
+#[inline]
+pub unsafe fn get_tsrm_ls_cache() -> *mut c_void {
+    zts::get_ls_cache()
+}
+
+#[cfg(php_zts)]
+#[inline]
+pub unsafe fn get_tsrm_resource_from_cache(ls_cache: *mut c_void, id: i32) -> *mut c_void {
+    zts::tsrmg_bulk(ls_cache, id)
+}
+
+#[cfg(php_zts)]
+#[inline]
+pub unsafe fn get_profiler_globals_from_cache(ls_cache: *mut c_void) -> *mut ProfilerGlobals {
+    // SAFETY: As long as this is called during the times documented by
+    // get_profiler_globals(), GLOBALS_ID will be set by PHP.
+    let id = ptr::addr_of!(GLOBALS_ID).read();
+    get_tsrm_resource_from_cache(ls_cache, id).cast()
+}
+
 /// Returns a pointer to the profiler globals for the current thread.
 ///
 /// # Safety
@@ -64,10 +97,7 @@ mod zts {
 pub unsafe fn get_profiler_globals() -> *mut ProfilerGlobals {
     #[cfg(php_zts)]
     {
-        // SAFETY: As long as this is called during the times documented by
-        // our own safety requirements, GLOBALS_ID will be set by PHP.
-        let id = ptr::addr_of!(GLOBALS_ID).read();
-        zts::tsrmg_bulk(id).cast()
+        get_profiler_globals_from_cache(get_tsrm_ls_cache())
     }
 
     #[cfg(not(php_zts))]
@@ -85,12 +115,13 @@ pub unsafe extern "C" fn ginit(_globals_ptr: *mut c_void) {
     #[cfg(php_zts)]
     crate::timeline::timeline_ginit();
 
-    // Initialize ZendMMState in PHP globals for ZTS builds. For NTS builds,
-    // this was already done in its const initializer.
+    // Initialize PHP globals for ZTS builds. For NTS builds, this was already
+    // done in its const initializer.
     #[cfg(php_zts)]
     {
         let globals = _globals_ptr.cast::<ProfilerGlobals>();
         (*globals).zend_mm_state = Cell::new(ZendMMState::new());
+        (*globals).interrupt_count = AtomicU32::new(0);
     }
 
     // SAFETY: this is called in thread ginit as expected, and no other places.
@@ -113,3 +144,19 @@ pub unsafe extern "C" fn gshutdown(_globals_ptr: *mut c_void) {
     // SAFETY: this is called in thread gshutdown as expected, no other places.
     allocation::gshutdown();
 }
+
+// Unit tests are not loaded by PHP, so provide the PHP globals and TSRM symbol
+// needed to link code retained in the test executable.
+#[cfg(test)]
+mod test_symbols {
+    #[cfg(not(php_zts))]
+    #[export_name = "compiler_globals"]
+    static mut TEST_COMPILER_GLOBALS: core::mem::MaybeUninit<crate::zend::zend_compiler_globals> =
+        core::mem::MaybeUninit::zeroed();
+
+    #[cfg(php_zts)]
+    #[no_mangle]
+    unsafe extern "C" fn tsrm_get_ls_cache() -> *mut core::ffi::c_void {
+        core::ptr::null_mut()
+    }
+}

profiling/src/profiling/interrupts.rs

@@ -50,7 +50,7 @@ impl InterruptManager {
             // Reset interrupt counter to prevent sampling during `mshutdown` (PHP 8.0 bug with
             // userland destructors), but leave the interrupt flag unchanged as other extensions
             // may have raised it.
-            (*interrupt.interrupt_count_ptr).store(0, Ordering::SeqCst);
+            (*interrupt.interrupt_count_ptr).store(0, Ordering::Relaxed);
         }
     }
 
@@ -62,7 +62,7 @@ impl InterruptManager {
     pub(super) fn trigger_interrupts(&self) {
         let vm_interrupts = self.vm_interrupts.lock().unwrap();
         vm_interrupts.iter().for_each(|obj| unsafe {
-            (*obj.interrupt_count_ptr).fetch_add(1, Ordering::SeqCst);
+            (*obj.interrupt_count_ptr).fetch_add(1, Ordering::Relaxed);
             (*obj.engine_ptr).store(true, Ordering::SeqCst);
         });
     }

profiling/src/profiling/mod.rs

@@ -1130,6 +1130,7 @@ impl Profiler {
 
     /// Collect a stack sample with elapsed wall time. Collects CPU time if
     /// it's enabled and available.
+    #[export_name = "ddog_php_prof_collect_time"]
     #[cfg_attr(feature = "tracing", tracing::instrument(skip_all, level = "debug"))]
     pub fn collect_time(&self, execute_data: *mut zend_execute_data, interrupt_count: u32) {
         // todo: should probably exclude the wall and CPU time used by collecting the sample.

profiling/src/wall_time.rs

@@ -2,10 +2,10 @@
 //! implementation reasons, it has cpu-time code as well.
 
 use crate::bindings::{zend_execute_data, zend_interrupt_function, VmInterruptFn};
-use crate::{profiling::Profiler, RefCellExt, REQUEST_LOCALS};
+use crate::module_globals;
+use crate::profiling::Profiler;
 use core::ptr;
-use log::debug;
-use std::sync::atomic::Ordering;
+use core::sync::atomic::Ordering;
 
 #[cfg(not(php_frameless))]
 mod execute_internal {
@@ -79,7 +79,7 @@ mod execute_internal {
         unsafe { prev_execute_internal(execute_data, return_value) };
 
         // See safety section of `execute_data_func_is_trampoline` docs for why
-        // the leaf frame is used  instead of the execute_data ptr.
+        // the leaf frame is used instead of the execute_data ptr.
         ddog_php_prof_interrupt_function(leaf_frame);
     }
 
@@ -109,31 +109,29 @@ static mut PREV_INTERRUPT_FUNCTION: Option<VmInterruptFn> = None;
 #[no_mangle]
 #[inline(never)]
 pub extern "C" fn ddog_php_prof_interrupt_function(execute_data: *mut zend_execute_data) {
-    let result = REQUEST_LOCALS.try_with_borrow(|locals| {
-        if !locals.system_settings().profiling_enabled {
-            return;
-        }
-
-        /* Other extensions/modules or the engine itself may trigger an
-         * interrupt, but given how expensive it is to gather a stack trace,
-         * it should only be done if we triggered it ourselves. So
-         * interrupt_count serves dual purposes:
-         *  1. Track how many interrupts there were.
-         *  2. Ensure we don't collect on someone else's interrupt.
-         */
-        let interrupt_count = locals.interrupt_count.swap(0, Ordering::SeqCst);
-        if interrupt_count == 0 {
-            return;
-        }
-
-        if let Some(profiler) = Profiler::get() {
-            // Safety: execute_data was provided by the engine, and the profiler doesn't mutate it.
-            profiler.collect_time(execute_data, interrupt_count);
-        }
-    });
+    // SAFETY: interrupt callbacks run while the current PHP thread's module globals are valid.
+    let atomic_count = unsafe { &(*module_globals::get_profiler_globals()).interrupt_count };
+
+    /* Other extensions/modules or the engine itself may trigger an
+     * interrupt, but given how expensive it is to gather a stack trace,
+     * it should only be done if we triggered it ourselves. So
+     * interrupt_count serves dual purposes:
+     *  1. Track how many interrupts there were.
+     *  2. Ensure we don't collect on someone else's interrupt.
+     */
+    let interrupt_count = atomic_count.swap(0, Ordering::Relaxed);
+    if interrupt_count == 0 {
+        return;
+    }
+    collect_time_if_enabled(execute_data, interrupt_count);
+}
 
-    if let Err(err) = result {
-        debug!("ddog_php_prof_interrupt_function failed to borrow request locals: {err}");
+#[inline(never)]
+#[export_name = "ddog_php_prof_collect_time_if_enabled"]
+extern "C" fn collect_time_if_enabled(execute_data: *mut zend_execute_data, interrupt_count: u32) {
+    if let Some(profiler) = Profiler::get() {
+        // Safety: execute_data was provided by the engine, and the profiler doesn't mutate it.
+        profiler.collect_time(execute_data, interrupt_count);
     }
 }
 
@@ -144,7 +142,9 @@ mod frameless {
         use crate::bindings::{
             zend_flf_functions, zend_flf_handlers, zend_frameless_function_info,
         };
-        use crate::{profiling::Profiler, zend, RefCellExt, REQUEST_LOCALS};
+        use crate::module_globals;
+        use crate::wall_time::collect_time_if_enabled;
+        use crate::zend;
         use dynasmrt::{dynasm, DynasmApi, ExecutableBuffer};
         use log::error;
         use std::ffi::c_void;
@@ -270,24 +270,19 @@ mod frameless {
         #[no_mangle]
         #[inline(never)]
         pub extern "C" fn ddog_php_prof_icall_trampoline_target() {
-            let interrupt_count = REQUEST_LOCALS
-                .try_with_borrow(|locals| {
-                    if !locals.system_settings().profiling_enabled {
-                        return 0;
-                    }
-                    locals.interrupt_count.swap(0, Ordering::SeqCst)
-                })
-                .unwrap_or(0);
+            // SAFETY: frameless handlers run while the current PHP thread's module globals are
+            // valid. Retain the pointer so the authoritative swap reuses the same TSRM lookup.
+            let atomic_count =
+                unsafe { &(*module_globals::get_profiler_globals()).interrupt_count };
 
+            let interrupt_count = atomic_count.swap(0, Ordering::Relaxed);
             if interrupt_count == 0 {
                 return;
             }
 
-            if let Some(profiler) = Profiler::get() {
-                // SAFETY: profiler doesn't mutate execute_data
-                let execute_data = unsafe { zend::ddog_php_prof_get_current_execute_data() };
-                profiler.collect_time(execute_data, interrupt_count);
-            }
+            // Fetching execute data is intentionally delayed until a profiler interrupt is pending.
+            let execute_data = unsafe { zend::ddog_php_prof_get_current_execute_data() };
+            collect_time_if_enabled(execute_data, interrupt_count);
         }
     }