From bc75501d1200662ba4f6c4205ce796204aedf534 Mon Sep 17 00:00:00 2001 From: Scott Gerring Date: Thu, 23 Jul 2026 11:21:10 +0200 Subject: [PATCH] feat(heap-profiling): move semaphore check into fast path --- .../src/allocator.rs | 15 +++---- libdd-profiling-heap-allocator/src/lib.rs | 15 ++++++- .../datadog/heap/allocation_requested.h | 9 +++++ .../include/datadog/heap/probes.h | 11 +++++ .../src/generated/bindings.rs | 3 ++ libdd-profiling-heap-sampler/src/lib.rs | 40 ++++++++++++++----- libdd-profiling-heap-sampler/src/probes.c | 7 +++- 7 files changed, 79 insertions(+), 21 deletions(-) diff --git a/libdd-profiling-heap-allocator/src/allocator.rs b/libdd-profiling-heap-allocator/src/allocator.rs index 17e7434c49..dbb1d01775 100644 --- a/libdd-profiling-heap-allocator/src/allocator.rs +++ b/libdd-profiling-heap-allocator/src/allocator.rs @@ -151,23 +151,18 @@ mod tests { // Touches sampler TLS internals, which only exist on Linux. #[cfg(target_os = "linux")] #[test] - fn lazy_init_populates_tls_on_first_alloc() { - // Spin a fresh thread so we start with uninitialized sampler TLS. + fn alloc_skips_tls_when_no_profiler_attached() { + // With the USDT semaphore inactive, alloc should not touch TLS. std::thread::spawn(|| unsafe { - assert!( - dd_tl_state_get().is_null(), - "fresh thread should have NULL sampler TLS" - ); + assert!(dd_tl_state_get().is_null()); let sampled = SampledAllocator::::DEFAULT; let layout = Layout::from_size_align(64, 8).unwrap(); let p = sampled.alloc(layout); assert!(!p.is_null()); - assert!( - !dd_tl_state_get().is_null(), - "TLS should be populated after the first alloc" - ); + // Semaphore inactive: TLS never initialized. + assert!(dd_tl_state_get().is_null()); sampled.dealloc(p, layout); }) diff --git a/libdd-profiling-heap-allocator/src/lib.rs b/libdd-profiling-heap-allocator/src/lib.rs index d17a498650..494351ada1 100644 --- a/libdd-profiling-heap-allocator/src/lib.rs +++ b/libdd-profiling-heap-allocator/src/lib.rs @@ -21,7 +21,9 @@ //! # Example //! //! ```no_run -//! use libdd_profiling_heap_allocator::{set_default_sampling_distance, SampledAllocator}; +//! use libdd_profiling_heap_allocator::{ +//! is_profiler_attached, set_default_sampling_distance, SampledAllocator, +//! }; //! use std::alloc::System; //! //! #[global_allocator] @@ -41,6 +43,17 @@ mod allocator; pub use allocator::SampledAllocator; +/// Returns true when an external profiler is currently attached to the +/// allocator's `ddheap:alloc` USDT. +/// +/// This is a point-in-time read of the USDT semaphore. A profiler can attach +/// or detach immediately after this returns. Use it as a diagnostic/readiness +/// signal, not as a synchronization primitive. +#[inline] +pub fn is_profiler_attached() -> bool { + libdd_profiling_heap_sampler::is_profiler_attached() +} + /// Set the default mean sample distance (bytes between samples) for the /// heap sampler. /// diff --git a/libdd-profiling-heap-sampler/include/datadog/heap/allocation_requested.h b/libdd-profiling-heap-sampler/include/datadog/heap/allocation_requested.h index fb6a6e85a0..40f0c151db 100644 --- a/libdd-profiling-heap-sampler/include/datadog/heap/allocation_requested.h +++ b/libdd-profiling-heap-sampler/include/datadog/heap/allocation_requested.h @@ -34,6 +34,12 @@ #include +#ifdef __linux__ +# include +#else +# define USDT_IS_ACTIVE(group, name) (1) +#endif + /* * Return type for dd_allocation_requested, paired with the `req` * argument of dd_allocation_created. @@ -107,6 +113,9 @@ static inline __attribute__((always_inline, warn_unused_result)) dd_alloc_req_t dd_allocation_requested(size_t size, size_t alignment) { dd_alloc_req_t out = { size, size, alignment, 0 }; + // If no tracer is attached to ddheap:alloc, skip all sampling logic. + if (__builtin_expect(!USDT_IS_ACTIVE(ddheap, alloc), 1)) return out; + // If we don't have TLS yet (this is defensive and should never happen), or // the reentry guard is set (meaning a sampled allocation is already in // flight on this thread and something in its slow path triggered another diff --git a/libdd-profiling-heap-sampler/include/datadog/heap/probes.h b/libdd-profiling-heap-sampler/include/datadog/heap/probes.h index 28913194bc..c3d4e0a67d 100644 --- a/libdd-profiling-heap-sampler/include/datadog/heap/probes.h +++ b/libdd-profiling-heap-sampler/include/datadog/heap/probes.h @@ -12,6 +12,7 @@ #ifndef DD_SAMPLERS_PROBES_H #define DD_SAMPLERS_PROBES_H +#include #include #ifdef __linux__ @@ -42,4 +43,14 @@ void dd_probe_alloc(void *user, uint64_t size, uint64_t weight); */ void dd_probe_free(void *ptr); +/* + * Returns true when an external profiler is currently attached to the + * ddheap:alloc USDT in this object file. + * + * This is a point-in-time read of the alloc probe's USDT semaphore. A profiler + * can attach or detach immediately after this returns. Use it as a + * diagnostic/readiness signal, not as a synchronization primitive. + */ +bool dd_heap_profiler_attached(void); + #endif diff --git a/libdd-profiling-heap-sampler/src/generated/bindings.rs b/libdd-profiling-heap-sampler/src/generated/bindings.rs index 9ba8e6fbcc..6ee2eaaf85 100644 --- a/libdd-profiling-heap-sampler/src/generated/bindings.rs +++ b/libdd-profiling-heap-sampler/src/generated/bindings.rs @@ -201,3 +201,6 @@ unsafe extern "C" { unsafe extern "C" { pub fn dd_probe_free(ptr: *mut ::std::os::raw::c_void); } +unsafe extern "C" { + pub fn dd_heap_profiler_attached() -> bool; +} diff --git a/libdd-profiling-heap-sampler/src/lib.rs b/libdd-profiling-heap-sampler/src/lib.rs index df2fa05cb8..9b47b10b51 100644 --- a/libdd-profiling-heap-sampler/src/lib.rs +++ b/libdd-profiling-heap-sampler/src/lib.rs @@ -52,6 +52,33 @@ pub fn set_default_sampling_distance(distance_bytes: u64) { unsafe { sys::dd_set_default_sampling_interval(distance_bytes) } } +#[cfg(target_os = "linux")] +extern "C" { + #[link_name = "dd_heap_profiler_attached"] + fn dd_heap_profiler_attached_ffi() -> bool; +} + +/// Returns true when an external profiler is currently attached to the +/// `ddheap:alloc` USDT in this object file. +/// +/// This is a point-in-time read of the alloc probe's USDT semaphore. A profiler +/// can attach or detach immediately after this returns. Use it as a +/// diagnostic/readiness signal, not as a synchronization primitive. +#[cfg(target_os = "linux")] +#[inline] +pub fn is_profiler_attached() -> bool { + // SAFETY: dd_heap_profiler_attached performs a single semaphore read and + // has no preconditions. + unsafe { dd_heap_profiler_attached_ffi() } +} + +/// Non-Linux builds do not expose USDT heap probes. +#[cfg(not(target_os = "linux"))] +#[inline] +pub fn is_profiler_attached() -> bool { + false +} + /// Whether heap sampling is enabled for this process. Users of this /// library can use this to check if they should setup allocation tracking /// or not. @@ -177,19 +204,14 @@ mod tests { } #[test] - fn requested_initializes_tls_on_first_use() { + fn requested_skips_tls_init_when_no_profiler_attached() { std::thread::spawn(|| unsafe { - assert!( - dd_tl_state_get().is_null(), - "fresh thread should start without sampler TLS" - ); + assert!(dd_tl_state_get().is_null()); let req = dd_allocation_requested(1, 1); assert_eq!(req.size, 1); assert_eq!(req.weight, 0); - assert!( - !dd_tl_state_get().is_null(), - "request should initialize sampler TLS" - ); + // Semaphore is inactive (no profiler), so TLS is never touched. + assert!(dd_tl_state_get().is_null()); }) .join() .unwrap(); diff --git a/libdd-profiling-heap-sampler/src/probes.c b/libdd-profiling-heap-sampler/src/probes.c index b82930fe4b..1f90bd9a09 100644 --- a/libdd-profiling-heap-sampler/src/probes.c +++ b/libdd-profiling-heap-sampler/src/probes.c @@ -21,11 +21,12 @@ #include #include +#include /* Save / restore errno: an attached USDT consumer may perturb it. */ void dd_probe_alloc(void *user, uint64_t size, uint64_t weight) { int saved_errno = errno; - USDT(ddheap, alloc, user, size, weight); + USDT_WITH_SEMA(ddheap, alloc, user, size, weight); errno = saved_errno; } @@ -38,3 +39,7 @@ void dd_probe_free(void *ptr) { (void)ptr; #endif } + +bool dd_heap_profiler_attached(void) { + return USDT_IS_ACTIVE(ddheap, alloc); +}