chore(heap-profiling): address review feedback and general cleanup

Author: scottgerring

libdd-profiling-heap-allocator/benches/sampler_overhead.rs

@@ -38,6 +38,10 @@ mod linux_bench {
             // Return a stable aligned pointer with mapped bytes before it. The sampler's free path
             // may inspect header-sized bytes immediately before the user pointer when
             // checking for sampled allocations.
+            //
+            // Always returns the same fixed pointer. This allocator isn't tracking real
+            // capacity or state; it exists purely to eliminate the real allocator's cost
+            // from the benchmark so it measures the sampler's own overhead.
             unsafe { ptr::addr_of_mut!(NOOP_BUFFER.0).cast::<u8>().add(4096) }
         }
 
@@ -48,10 +52,21 @@ mod linux_bench {
         unsafe { ptr::addr_of_mut!(NOOP_BUFFER.0).cast::<u8>().add(4096) }
     }
 
+    /// # Safety
+    ///
+    /// Must be called on a thread that isn't concurrently tearing down its
+    /// TLS (i.e. not from a destructor); otherwise identical to
+    /// `dd_tl_state_get_or_init`'s own safety contract.
     unsafe fn sampler_tl_state() -> *mut libdd_profiling_heap_sampler::dd_tl_state_t {
         unsafe { dd_tl_state_get_or_init() }
     }
 
+    /// Pins this thread's sampler state onto the fast (unsampled) path for
+    /// the rest of the benchmark. `remaining_bytes` starts at a huge
+    /// negative value and `sampling_interval` at a huge positive one, so
+    /// benchmark-sized allocations can never drive `remaining_bytes`
+    /// non-negative and trigger the slow/sampled path. Dividing by 4 keeps
+    /// headroom against overflow while summing allocation sizes.
     unsafe fn pin_sampler_to_fast_path() {
         let tl = unsafe { sampler_tl_state() };
         if !tl.is_null() {
@@ -64,6 +79,10 @@ mod linux_bench {
         }
     }
 
+    /// Forces the next allocation on this thread onto the slow/sampled
+    /// path. `512 * 1024` matches `DD_SAMPLING_INTERVAL_DEFAULT` (see
+    /// tl_state.h), so this benchmarks the sampled path against the same
+    /// interval used in production rather than an arbitrary value.
     unsafe fn force_next_allocation_to_sample() {
         let tl = unsafe { sampler_tl_state() };
         if !tl.is_null() {

libdd-profiling-heap-allocator/src/allocator.rs

@@ -71,6 +71,13 @@ unsafe impl<A: GlobalAlloc> GlobalAlloc for SampledAllocator<A> {
             return self.inner.dealloc(ptr, layout);
         }
         let freed = dd_allocation_freed(ptr.cast(), layout.size(), layout.align());
+        // `layout.align()` is reused here rather than anything derived from
+        // `freed`: alignment never changes between the original `alloc` and
+        // this `dealloc`, so pairing `freed.size` with the caller's own
+        // alignment is exactly what `GlobalAlloc::dealloc`'s safety contract
+        // already requires (the layout passed here must match the one used
+        // for the original allocation). This isn't a guarantee `SampledAllocator`
+        // adds — it's just satisfying the contract our caller is on the hook for.
         let inner_layout = Layout::from_size_align_unchecked(freed.size, layout.align());
         self.inner.dealloc(freed.ptr.cast(), inner_layout);
     }

libdd-profiling-heap-allocator/src/lib.rs

@@ -14,8 +14,9 @@
 //!
 //! # Features
 //!
-//! * `live-heap` (off by default) — enables live-heap tracking: allocations are flagged and frees
-//!   are sampled, so a profiler can balance allocs against frees. Off = allocation profiling only.
+//! * `live-heap` (off by default) - enables live-heap tracking: sampled allocations are flagged at
+//!   alloc time, and that flag is detected again on free, so a profiler can pair each free back to
+//!   its sample and balance allocs against frees. Off = allocation profiling only.
 //!
 //! # Example
 //!

libdd-profiling-heap-gotter-ffi/Cargo.toml

@@ -16,6 +16,8 @@ bench = false
 [features]
 default = ["cbindgen"]
 cbindgen = ["build_common/cbindgen", "libdd-common-ffi/cbindgen"]
+# See libdd-profiling-heap-gotter's test-support feature. Not for production use.
+test-support = ["libdd-profiling-heap-gotter/test-support"]
 
 [build-dependencies]
 build_common = { path = "../build-common" }

libdd-profiling-heap-gotter-ffi/examples/cdylib_demo.rs

@@ -17,6 +17,13 @@
 //! The gotter-ffi crate is Linux-only; on other targets the example
 //! compiles to a no-op `main` so clippy/test on non-Linux don't fail
 //! with "configured out".
+//!
+//! This demo deliberately loads the library via `dlopen`/`dlsym` even
+//! though a real Rust binary would just link the crate directly. That's
+//! contrived for Rust specifically, but it's the pattern every other
+//! language's SDK actually has to use to consume this FFI surface, so
+//! exercising it here catches issues (symbol visibility, ABI mismatches)
+//! that a direct `extern crate` link would hide.
 
 #[cfg(not(target_os = "linux"))]
 fn main() -> Result<(), String> {
@@ -54,6 +61,11 @@ mod linux {
             Ok(Self(handle))
         }
 
+        /// # Safety
+        ///
+        /// `T` must be exactly the function pointer type of the symbol named
+        /// `name` in the loaded library; a mismatch is transmuted silently
+        /// and will call into the wrong signature.
         unsafe fn symbol<T>(&self, name: &CStr) -> Result<T, String>
         where
             T: Copy,
@@ -85,6 +97,11 @@ mod linux {
         }
     }
 
+    /// Infers the built cdylib's path from `cargo run --example`'s layout:
+    /// the example binary lands at `target/<profile>/examples/<name>`, and
+    /// the cdylib is a sibling of `examples/` at `target/<profile>/<lib>`.
+    /// `DDOG_HEAP_GOTTER_FFI_CDYLIB` overrides this for anyone driving the
+    /// build differently.
     fn cdylib_path() -> Result<PathBuf, String> {
         if let Some(path) = std::env::var_os("DDOG_HEAP_GOTTER_FFI_CDYLIB") {
             return Ok(PathBuf::from(path));

libdd-profiling-heap-gotter-ffi/src/lib.rs

@@ -62,3 +62,14 @@ pub extern "C" fn ddog_heap_gotter_update() -> VoidResult {
 pub extern "C" fn ddog_heap_gotter_is_installed() -> bool {
     libdd_profiling_heap_gotter::heap_overrides_are_installed()
 }
+
+/// Test-only: number of times a patched hook (`malloc`/`free`) has run in
+/// this process. Lets integration tests prove the patched GOT was actually
+/// exercised, not just that nothing crashed. Not part of the production API
+/// surface; only compiled in with the `test-support` feature.
+#[cfg(feature = "test-support")]
+#[no_mangle]
+#[must_use]
+pub extern "C" fn ddog_heap_gotter_test_hook_hits() -> u64 {
+    libdd_profiling_heap_gotter::test_hook_hits()
+}

libdd-profiling-heap-gotter-ffi/tests/install.rs

@@ -14,6 +14,8 @@
 #![cfg(all(target_os = "linux", target_pointer_width = "64", not(miri)))]
 
 use libdd_common_ffi::VoidResult;
+#[cfg(feature = "test-support")]
+use libdd_profiling_heap_gotter_ffi::ddog_heap_gotter_test_hook_hits;
 use libdd_profiling_heap_gotter_ffi::{ddog_heap_gotter_install, ddog_heap_gotter_is_installed};
 
 #[track_caller]
@@ -37,9 +39,19 @@ fn install_patches_the_got() {
         "is_installed should be true after install"
     );
 
-    // Touch the heap while installed so the patched GOT actually gets
-    // used. We just need the process to still be alive after this.
+    // Touch the heap while installed so the patched GOT actually gets used.
+    #[cfg(feature = "test-support")]
+    let hits_before = ddog_heap_gotter_test_hook_hits();
+
     let v: Vec<u8> = vec![0; 128];
     assert_eq!(v.len(), 128);
     drop(v);
+
+    // With test-support, prove the hooks actually ran rather than just that
+    // the process survived.
+    #[cfg(feature = "test-support")]
+    assert!(
+        ddog_heap_gotter_test_hook_hits() > hits_before,
+        "expected the malloc/free hooks to run"
+    );
 }

libdd-profiling-heap-gotter/Cargo.toml

@@ -19,6 +19,10 @@ bench = false
 [features]
 # Forward live-heap (retained-heap) tracking down to the sampler
 live-heap = ["libdd-profiling-heap-sampler/live-heap"]
+# Exposes a hook-hit counter so integration tests in other crates (e.g.
+# libdd-profiling-heap-gotter-ffi) can prove the patched GOT was actually
+# used, not just that nothing crashed. Not for production use.
+test-support = []
 
 [target.'cfg(target_os = "linux")'.dependencies]
 libc = "0.2"

libdd-profiling-heap-gotter/src/elf.rs

@@ -531,6 +531,13 @@ struct PatchedLibrary {
     /// Identifies the library at this base address, so we can detect
     /// base-address reuse after a `dlclose` + `dlopen` places a
     /// different library at the same load address.
+    ///
+    /// Known limitation: detection keys on this name, not on library
+    /// contents. A different version of the same library reloaded at the
+    /// same base (identical path, changed contents, same base despite
+    /// ASLR) would slip through and we would restore stale GOT values.
+    /// This is judged unlikely enough in practice to document rather than
+    /// guard against with additional fingerprinting.
     library_name: String,
     /// Set each pass in which this library was seen; used to drop entries
     /// for libraries that have since been unloaded.

libdd-profiling-heap-gotter/src/hooks.rs

@@ -17,8 +17,9 @@ use core::ffi::{c_char, c_int, c_void};
 use std::sync::atomic::{AtomicUsize, Ordering};
 
 use libdd_profiling_heap_sampler::{
-    dd_alloc_req_t, dd_allocation_created, dd_allocation_freed, dd_allocation_realloc_commit,
-    dd_allocation_realloc_prepare, dd_allocation_requested, dd_tl_state_init,
+    dd_alloc_req_is_sampled, dd_alloc_req_t, dd_allocation_created, dd_allocation_freed,
+    dd_allocation_realloc_commit, dd_allocation_realloc_prepare, dd_allocation_requested,
+    dd_tl_state_init,
 };
 
 /// Resolved address of the real `malloc`; filled by `install_heap_overrides`.
@@ -39,9 +40,27 @@ pub(crate) static ORIG_DLOPEN: AtomicUsize = AtomicUsize::new(0);
 /// Resolved address of the real `pthread_create`.
 pub(crate) static ORIG_PTHREAD_CREATE: AtomicUsize = AtomicUsize::new(0);
 
+/// Counts hook invocations so integration tests outside this crate can
+/// prove the patched GOT was actually exercised, not just that nothing
+/// crashed. Only `malloc`/`free` increment it: that's enough to prove the
+/// hooks ran without adding bookkeeping to every symbol.
+#[cfg(feature = "test-support")]
+pub(crate) static HOOK_HITS: AtomicUsize = AtomicUsize::new(0);
+
 /// Load a resolved function pointer from one of the `ORIG_*` slots.
+///
+/// # Safety
+///
+/// `T` must be exactly the `extern "C" fn(...)` pointer type that
+/// `apply_overrides` writes into `slot` (see the `ORIG_*` docs above); if
+/// `slot` is still `0`, this returns `None` rather than transmuting.
 #[inline]
 unsafe fn load_fn<T>(slot: &AtomicUsize) -> Option<T> {
+    // `Acquire` pairs with the `store(Release)` in elf.rs's
+    // `apply_overrides`, which runs before the GOT is ever patched to
+    // route calls into these hooks. That gives a real happens-before
+    // edge: once this observes a non-zero slot, it's guaranteed to see
+    // the fully-published address, not just "no torn read".
     let v = slot.load(Ordering::Acquire);
     if v == 0 {
         None
@@ -95,17 +114,21 @@ type PthreadCreateFn = unsafe extern "C" fn(
 
 #[no_mangle]
 pub unsafe extern "C" fn gotter_malloc(size: usize) -> *mut c_void {
+    #[cfg(feature = "test-support")]
+    HOOK_HITS.fetch_add(1, Ordering::Relaxed);
     let Some(real): Option<MallocFn> = load_fn(&ORIG_MALLOC) else {
         return std::ptr::null_mut();
     };
     // Default alignment for malloc on glibc is 2*sizeof(void*) == 16.
-    let req = dd_allocation_requested(size, core::mem::align_of::<u64>() * 2);
+    let req = dd_allocation_requested(size, core::mem::align_of::<*mut c_void>() * 2);
     let raw = real(req.size);
     dd_allocation_created(raw, req)
 }
 
 #[no_mangle]
 pub unsafe extern "C" fn gotter_free(ptr: *mut c_void) {
+    #[cfg(feature = "test-support")]
+    HOOK_HITS.fetch_add(1, Ordering::Relaxed);
     let Some(real): Option<FreeFn> = load_fn(&ORIG_FREE) else {
         return;
     };
@@ -125,14 +148,14 @@ pub unsafe extern "C" fn gotter_calloc(nmemb: usize, size: usize) -> *mut c_void
     let Some(total) = nmemb.checked_mul(size) else {
         return real(nmemb, size);
     };
-    let req = dd_allocation_requested(total, core::mem::align_of::<u64>() * 2);
+    let req = dd_allocation_requested(total, core::mem::align_of::<*mut c_void>() * 2);
     // calloc takes (nmemb, size); when the sampler bumps `req.size` we
     // funnel the extra bytes into the size argument (nmemb stays 1's
     // worth conceptually). The simplest robust path is to switch to a
     // single (1, req.size) allocation when sampling kicks in, so the
     // underlying allocator zeroes everything we hand back. Unsampled
     // path keeps the user's (nmemb, size) verbatim.
-    let raw = if req.weight == 0 {
+    let raw = if !dd_alloc_req_is_sampled(req) {
         real(nmemb, size)
     } else {
         real(1, req.size)