fix(profiling): revert unsafe optimization (#3541)

Author: morrisonlevi

profiling/src/allocation/allocation_ge84.rs

@@ -1,6 +1,4 @@
-use crate::allocation::{
-    allocation_profiling_stats_mut, allocation_profiling_stats_should_collect, collect_allocation,
-};
+use crate::allocation::{collect_allocation, ALLOCATION_PROFILING_STATS};
 use crate::bindings::{self as zend};
 use crate::{RefCellExt, PROFILER_NAME};
 use core::{cell::Cell, ptr};
@@ -368,7 +366,9 @@ unsafe extern "C" fn alloc_prof_malloc(len: size_t) -> *mut c_void {
         return ptr;
     }
 
-    if allocation_profiling_stats_should_collect(len) {
+    if ALLOCATION_PROFILING_STATS
+        .borrow_mut_or_false(|allocations| allocations.should_collect_allocation(len))
+    {
         collect_allocation(len);
     }
 
@@ -437,7 +437,9 @@ unsafe extern "C" fn alloc_prof_realloc(prev_ptr: *mut c_void, len: size_t) -> *
         return ptr;
     }
 
-    if allocation_profiling_stats_should_collect(len) {
+    if ALLOCATION_PROFILING_STATS
+        .borrow_mut_or_false(|allocations| allocations.should_collect_allocation(len))
+    {
         collect_allocation(len);
     }
 

profiling/src/allocation/allocation_le83.rs

@@ -1,4 +1,4 @@
-use crate::allocation::{allocation_profiling_stats_should_collect, collect_allocation};
+use crate::allocation::{collect_allocation, ALLOCATION_PROFILING_STATS};
 use crate::bindings::{
     self as zend, datadog_php_install_handler, datadog_php_zif_handler,
     ddog_php_prof_copy_long_into_zval,
@@ -357,7 +357,9 @@ unsafe extern "C" fn alloc_prof_malloc(len: size_t) -> *mut c_void {
         return ptr;
     }
 
-    if allocation_profiling_stats_should_collect(len) {
+    if ALLOCATION_PROFILING_STATS
+        .borrow_mut_or_false(|allocations| allocations.should_collect_allocation(len))
+    {
         collect_allocation(len);
     }
 
@@ -416,7 +418,9 @@ unsafe extern "C" fn alloc_prof_realloc(prev_ptr: *mut c_void, len: size_t) -> *
         return ptr;
     }
 
-    if allocation_profiling_stats_should_collect(len) {
+    if ALLOCATION_PROFILING_STATS
+        .borrow_mut_or_false(|allocations| allocations.should_collect_allocation(len))
+    {
         collect_allocation(len);
     }
 

profiling/src/allocation/mod.rs

@@ -1,19 +1,16 @@
-mod tls_allocation_profiling_stats;
-
 #[cfg(php_zend_mm_set_custom_handlers_ex)]
 pub mod allocation_ge84;
 #[cfg(not(php_zend_mm_set_custom_handlers_ex))]
 pub mod allocation_le83;
 
-pub use tls_allocation_profiling_stats::*;
-
 use crate::bindings::{self as zend};
 use crate::profiling::Profiler;
 use crate::{RefCellExt, REQUEST_LOCALS};
 use libc::size_t;
 use log::{debug, error, trace};
 use rand::rngs::ThreadRng;
 use rand_distr::{Distribution, Poisson};
+use std::cell::RefCell;
 use std::ffi::c_void;
 use std::sync::atomic::{AtomicU64, Ordering};
 
@@ -88,6 +85,23 @@ pub fn collect_allocation(len: size_t) {
     }
 }
 
+thread_local! {
+    static ALLOCATION_PROFILING_STATS: RefCell<AllocationProfilingStats> =
+        RefCell::new(AllocationProfilingStats::new());
+}
+
+pub fn alloc_prof_ginit() {
+    #[cfg(not(php_zend_mm_set_custom_handlers_ex))]
+    allocation_le83::alloc_prof_ginit();
+    #[cfg(php_zend_mm_set_custom_handlers_ex)]
+    allocation_ge84::alloc_prof_ginit();
+}
+
+pub fn alloc_prof_gshutdown() {
+    #[cfg(php_zend_mm_set_custom_handlers_ex)]
+    allocation_ge84::alloc_prof_gshutdown();
+}
+
 #[cfg(not(php_zend_mm_set_custom_handlers_ex))]
 pub fn alloc_prof_startup() {
     allocation_le83::alloc_prof_startup();

profiling/src/allocation/tls_allocation_profiling_stats.rs

@@ -215,15 +215,15 @@ unsafe extern "C" fn ginit(_globals_ptr: *mut c_void) {
     timeline::timeline_ginit();
 
     // SAFETY: this is called in thread ginit as expected, and no other places.
-    allocation::ginit();
+    allocation::alloc_prof_ginit();
 }
 
 unsafe extern "C" fn gshutdown(_globals_ptr: *mut c_void) {
     #[cfg(php_zts)]
     timeline::timeline_gshutdown();
 
     // SAFETY: this is called in thread gshutdown as expected, no other places.
-    allocation::gshutdown();
+    allocation::alloc_prof_gshutdown();
 }
 
 // Important note on the PHP lifecycle: