Skip to content

Commit eca063f

Browse files
refactor: encapsulate ALLOCATION_PROFILING_STATS for safety (#3519)
Co-authored-by: Florian Engelhardt <[email protected]>
1 parent 7209999 commit eca063f

5 files changed

Lines changed: 139 additions & 74 deletions

File tree

profiling/src/allocation/allocation_ge84.rs

Lines changed: 5 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,6 @@
1-
use crate::allocation::{collect_allocation, ALLOCATION_PROFILING_STATS};
1+
use crate::allocation::{
2+
allocation_profiling_stats_mut, allocation_profiling_stats_should_collect, collect_allocation,
3+
};
24
use crate::bindings::{self as zend};
35
use crate::{RefCellExt, PROFILER_NAME};
46
use core::{cell::Cell, ptr};
@@ -366,14 +368,7 @@ unsafe extern "C" fn alloc_prof_malloc(len: size_t) -> *mut c_void {
366368
return ptr;
367369
}
368370

369-
let should_collect = ALLOCATION_PROFILING_STATS
370-
.try_with(|stats| {
371-
// Safety: ALLOCATION_PROFILING_STATS is thread-local and initialized in GINIT.
372-
unsafe { (*(*stats.get()).as_mut_ptr()).should_collect_allocation(len) }
373-
})
374-
.unwrap_unchecked();
375-
376-
if should_collect {
371+
if allocation_profiling_stats_should_collect(len) {
377372
collect_allocation(len);
378373
}
379374

@@ -442,14 +437,7 @@ unsafe extern "C" fn alloc_prof_realloc(prev_ptr: *mut c_void, len: size_t) -> *
442437
return ptr;
443438
}
444439

445-
let should_collect = ALLOCATION_PROFILING_STATS
446-
.try_with(|stats| {
447-
// Safety: ALLOCATION_PROFILING_STATS is thread-local and initialized in GINIT.
448-
unsafe { (*(*stats.get()).as_mut_ptr()).should_collect_allocation(len) }
449-
})
450-
.unwrap_unchecked();
451-
452-
if should_collect {
440+
if allocation_profiling_stats_should_collect(len) {
453441
collect_allocation(len);
454442
}
455443

profiling/src/allocation/allocation_le83.rs

Lines changed: 3 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
use crate::allocation::{collect_allocation, ALLOCATION_PROFILING_STATS};
1+
use crate::allocation::{allocation_profiling_stats_should_collect, collect_allocation};
22
use crate::bindings::{
33
self as zend, datadog_php_install_handler, datadog_php_zif_handler,
44
ddog_php_prof_copy_long_into_zval,
@@ -357,14 +357,7 @@ unsafe extern "C" fn alloc_prof_malloc(len: size_t) -> *mut c_void {
357357
return ptr;
358358
}
359359

360-
let should_collect = ALLOCATION_PROFILING_STATS
361-
.try_with(|stats| {
362-
// Safety: ALLOCATION_PROFILING_STATS is thread-local and initialized in GINIT.
363-
unsafe { (*(*stats.get()).as_mut_ptr()).should_collect_allocation(len) }
364-
})
365-
.unwrap_unchecked();
366-
367-
if should_collect {
360+
if allocation_profiling_stats_should_collect(len) {
368361
collect_allocation(len);
369362
}
370363

@@ -423,14 +416,7 @@ unsafe extern "C" fn alloc_prof_realloc(prev_ptr: *mut c_void, len: size_t) -> *
423416
return ptr;
424417
}
425418

426-
let should_collect = ALLOCATION_PROFILING_STATS
427-
.try_with(|stats| {
428-
// Safety: ALLOCATION_PROFILING_STATS is thread-local and initialized in GINIT.
429-
unsafe { (*(*stats.get()).as_mut_ptr()).should_collect_allocation(len) }
430-
})
431-
.unwrap_unchecked();
432-
433-
if should_collect {
419+
if allocation_profiling_stats_should_collect(len) {
434420
collect_allocation(len);
435421
}
436422

profiling/src/allocation/mod.rs

Lines changed: 9 additions & 38 deletions
Original file line numberDiff line numberDiff line change
@@ -1,20 +1,22 @@
1+
mod tls_allocation_profiling_stats;
2+
3+
#[cfg(php_zend_mm_set_custom_handlers_ex)]
4+
pub mod allocation_ge84;
5+
#[cfg(not(php_zend_mm_set_custom_handlers_ex))]
6+
pub mod allocation_le83;
7+
8+
pub use tls_allocation_profiling_stats::*;
9+
110
use crate::bindings::{self as zend};
211
use crate::profiling::Profiler;
312
use crate::{RefCellExt, REQUEST_LOCALS};
413
use libc::size_t;
514
use log::{debug, error, trace};
615
use rand::rngs::ThreadRng;
716
use rand_distr::{Distribution, Poisson};
8-
use std::cell::UnsafeCell;
917
use std::ffi::c_void;
10-
use std::mem::MaybeUninit;
1118
use std::sync::atomic::{AtomicU64, Ordering};
1219

13-
#[cfg(php_zend_mm_set_custom_handlers_ex)]
14-
pub mod allocation_ge84;
15-
#[cfg(not(php_zend_mm_set_custom_handlers_ex))]
16-
pub mod allocation_le83;
17-
1820
/// Default sampling interval in bytes (4MB)
1921
pub const DEFAULT_ALLOCATION_SAMPLING_INTERVAL: u64 = 1024 * 4096;
2022

@@ -86,37 +88,6 @@ pub fn collect_allocation(len: size_t) {
8688
}
8789
}
8890

89-
thread_local! {
90-
pub(crate) static ALLOCATION_PROFILING_STATS: UnsafeCell<MaybeUninit<AllocationProfilingStats>> =
91-
const { UnsafeCell::new(MaybeUninit::uninit()) };
92-
}
93-
94-
pub fn alloc_prof_ginit() {
95-
unsafe {
96-
// Eagerly initialize the allocation profiling stats before handling first request
97-
ALLOCATION_PROFILING_STATS
98-
.try_with(|slot| {
99-
// SAFETY:
100-
// - `thread_local!` guarantees `slot` is per-thread, and `UnsafeCell::get()`
101-
// gives us a mutable pointer to the inner `MaybeUninit`.
102-
// - We must not call this twice on the same thread.
103-
let slot_ptr: *mut MaybeUninit<AllocationProfilingStats> = slot.get();
104-
(*slot_ptr).write(AllocationProfilingStats::new());
105-
})
106-
.unwrap_unchecked();
107-
}
108-
109-
#[cfg(not(php_zend_mm_set_custom_handlers_ex))]
110-
allocation_le83::alloc_prof_ginit();
111-
#[cfg(php_zend_mm_set_custom_handlers_ex)]
112-
allocation_ge84::alloc_prof_ginit();
113-
}
114-
115-
pub fn alloc_prof_gshutdown() {
116-
#[cfg(php_zend_mm_set_custom_handlers_ex)]
117-
allocation_ge84::alloc_prof_gshutdown();
118-
}
119-
12091
#[cfg(not(php_zend_mm_set_custom_handlers_ex))]
12192
pub fn alloc_prof_startup() {
12293
allocation_le83::alloc_prof_startup();
Lines changed: 118 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,118 @@
1+
//! The thread-local allocation profiling stats are held in this module.
2+
//! The stats are used on the hot-path of allocation, so this code is
3+
//! performance sensitive. It is encapsulated so that some unsafe techniques
4+
//! can be used but expose a relatively safe API.
5+
6+
use super::AllocationProfilingStats;
7+
use libc::size_t;
8+
use std::cell::UnsafeCell;
9+
use std::mem::MaybeUninit;
10+
11+
#[cfg(php_zend_mm_set_custom_handlers_ex)]
12+
use super::allocation_ge84;
13+
#[cfg(not(php_zend_mm_set_custom_handlers_ex))]
14+
use super::allocation_le83;
15+
16+
thread_local! {
17+
/// This is initialized in ginit, before any memory allocator hooks are
18+
/// installed. During a request, all accesses will be initialized.
19+
///
20+
/// This is not pub so that unsafe code can be contained to this module.
21+
static ALLOCATION_PROFILING_STATS: UnsafeCell<MaybeUninit<AllocationProfilingStats>> =
22+
const { UnsafeCell::new(MaybeUninit::uninit()) };
23+
}
24+
25+
/// Accesses the thread-local [`AllocationProfilingStats`], passing a mutable
26+
/// reference to the contained `MaybeUninit` to `F`.
27+
///
28+
/// # Safety
29+
///
30+
/// 1. There should not be any active borrows to the thread-local variable
31+
/// [`AllocationProfilingStats`] when this function is called.
32+
/// 2. Function `F` should not do anything which causes a new borrow on
33+
/// [`AllocationProfilingStats`].
34+
/// 3. Do not call this function in ALLOCATION_PROFILING_STATS's destructor,
35+
/// as it assumes that [`std::thread::LocalKey::try_with`] cannot fail.
36+
///
37+
/// This is not pub to limit caller's ability to violate these conditions.
38+
unsafe fn allocation_profiling_stats_mut<F, R>(f: F) -> R
39+
where
40+
F: FnOnce(&mut MaybeUninit<AllocationProfilingStats>) -> R,
41+
{
42+
let result = ALLOCATION_PROFILING_STATS.try_with(|cell| {
43+
let ptr: *mut MaybeUninit<AllocationProfilingStats> = cell.get();
44+
// SAFETY: the cell is statically initialized to [`MaybeUninit::uninit`] so the
45+
// _cell_ is valid and initialized memory. As required by this own
46+
// function's safety requirements, there should not be any active borrows
47+
// to [`ALLOCATION_PROFILING_STATS`], so this mutable dereference is sound.
48+
let uninit = unsafe { &mut *ptr };
49+
f(uninit)
50+
});
51+
unsafe {
52+
// SAFETY: this function is not called in a destructor, therefore it
53+
// cannot return an AccessError:
54+
// > If the key has been destroyed (which may happen if this is called
55+
// > in a destructor), this function will return an AccessError.
56+
result.unwrap_unchecked()
57+
}
58+
}
59+
60+
/// Given the provided allocation length `len`, return whether the allocation
61+
/// should be collected. This is a mutable operation, as the thread-local
62+
/// variable will be modified to reduce the distance until the next sample.
63+
pub fn allocation_profiling_stats_should_collect(len: size_t) -> bool {
64+
let f = |maybe_uninit: &mut MaybeUninit<AllocationProfilingStats>| {
65+
// SAFETY: ALLOCATION_PROFILING_STATS was initialized in GINIT.
66+
let stats = unsafe { maybe_uninit.assume_init_mut() };
67+
stats.should_collect_allocation(len)
68+
};
69+
70+
// SAFETY:
71+
// 1. This function doesn't expose any way for the caller to keep a
72+
// borrow alive, nor do the other public functions, so there cannot be
73+
// any existing borrows alive.
74+
// 2. This closure will not cause any new borrows.
75+
// 3. This function isn't called during ALLOCATION_PROFILING_STATS's dtor,
76+
// as MaybeUninit's destructor does nothing, you have to specifically drop
77+
// it. Even if the destructor were called, AllocationProfilingStats's dtor
78+
// doesn't access the TLS variable (it can't, it doesn't have access).
79+
unsafe { allocation_profiling_stats_mut(f) }
80+
}
81+
82+
/// Initializes the allocation profiler's globals.
83+
///
84+
/// # Safety
85+
///
86+
/// Must be called once per PHP thread ginit.
87+
pub unsafe fn ginit() {
88+
// SAFETY:
89+
// 1. During ginit, there will not be any other borrows.
90+
// 2. This closure will not make new borrows.
91+
// 3. This is not during the thread-local destructor.
92+
unsafe {
93+
allocation_profiling_stats_mut(|uninit| {
94+
uninit.write(AllocationProfilingStats::new());
95+
})
96+
};
97+
98+
#[cfg(not(php_zend_mm_set_custom_handlers_ex))]
99+
allocation_le83::alloc_prof_ginit();
100+
#[cfg(php_zend_mm_set_custom_handlers_ex)]
101+
allocation_ge84::alloc_prof_ginit();
102+
}
103+
104+
/// Shuts down the allocation profiler's globals.
105+
///
106+
/// # Safety
107+
///
108+
/// Must be called once per PHP thread gshutdown.
109+
pub unsafe fn gshutdown() {
110+
#[cfg(php_zend_mm_set_custom_handlers_ex)]
111+
allocation_ge84::alloc_prof_gshutdown();
112+
113+
// SAFETY:
114+
// 1. During gshutdown, there will not be any other borrows.
115+
// 2. This closure will not make new borrows.
116+
// 3. This is not during the thread-local destructor.
117+
unsafe { allocation_profiling_stats_mut(|maybe_uninit| maybe_uninit.assume_init_drop()) }
118+
}

profiling/src/lib.rs

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -214,14 +214,16 @@ unsafe extern "C" fn ginit(_globals_ptr: *mut c_void) {
214214
#[cfg(php_zts)]
215215
timeline::timeline_ginit();
216216

217-
allocation::alloc_prof_ginit();
217+
// SAFETY: this is called in thread ginit as expected, and no other places.
218+
allocation::ginit();
218219
}
219220

220221
unsafe extern "C" fn gshutdown(_globals_ptr: *mut c_void) {
221222
#[cfg(php_zts)]
222223
timeline::timeline_gshutdown();
223224

224-
allocation::alloc_prof_gshutdown();
225+
// SAFETY: this is called in thread gshutdown as expected, no other places.
226+
allocation::gshutdown();
225227
}
226228

227229
// Important note on the PHP lifecycle:

0 commit comments

Comments
 (0)