chore(heap-profiling): more review feedback cleanup

Author: scottgerring

Cargo.lock

@@ -33,3 +33,6 @@ libdd-profiling-heap-gotter = { path = "../libdd-profiling-heap-gotter" }
 
 [target.'cfg(target_os = "linux")'.dev-dependencies]
 libc = "0.2"
+# For the USDT-note ELF inspection test (tests/usdt_notes.rs); mirrors how
+# libdd-otel-thread-ctx-ffi pulls in its dep's `sanity-check` feature.
+libdd-profiling-heap-sampler = { path = "../libdd-profiling-heap-sampler", features = ["sanity-check"] }

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

@@ -0,0 +1,33 @@
+// Copyright 2026-Present Datadog, Inc. https://www.datadoghq.com/
+// SPDX-License-Identifier: Apache-2.0
+
+//! Verify the ddheap USDT probes appear exactly once in the built shared
+//! library's `.note.stapsdt`. A duplicate entry (e.g. from `static inline` +
+//! bindgen's `wrap_static_fns`, or LTO inlining across TUs) would make an
+//! attached consumer fire twice, so this guards the "one note per probe"
+//! invariant documented in the sampler's `probes.c`.
+//!
+//! Inspecting the linked cdylib (rather than running the check in-process)
+//! validates the real shipped artifact. Mirrors
+//! libdd-otel-thread-ctx-ffi/tests/elf_properties.rs.
+//!
+//! The cdylib path is derived at runtime from the test executable location.
+//! Both the test binary and the cdylib live in `target/<[triple/]profile>/deps/`.
+
+#![cfg(target_os = "linux")]
+
+use std::path::PathBuf;
+
+fn cdylib_path() -> PathBuf {
+    let exe = std::env::current_exe().expect("failed to read current executable path");
+    exe.parent()
+        .expect("unexpected test executable path structure")
+        .join("liblibdd_profiling_heap_gotter_ffi.so")
+}
+
+#[test]
+#[cfg_attr(miri, ignore)]
+fn ddheap_probes_have_one_note_each() {
+    let path = cdylib_path();
+    libdd_profiling_heap_sampler::usdt_check::check_usdt_probes_in(&path).unwrap();
+}

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

@@ -61,6 +61,8 @@ mod elf;
 #[cfg(all(target_os = "linux", target_pointer_width = "64"))]
 mod hooks;
 
+#[cfg(all(target_os = "linux", target_pointer_width = "64"))]
+use std::sync::atomic::{AtomicBool, Ordering};
 #[cfg(all(target_os = "linux", target_pointer_width = "64"))]
 use std::sync::{Mutex, MutexGuard, TryLockError};
 
@@ -73,10 +75,29 @@ use elf::SymbolOverrides;
 #[cfg(all(target_os = "linux", target_pointer_width = "64"))]
 static GLOBAL_OVERRIDES: Mutex<Option<SymbolOverrides>> = Mutex::new(None);
 
+/// Set when a `dlopen` hook loses the `try_lock` race in
+/// `update_heap_overrides`: its newly-loaded library wasn't scanned, so the
+/// thread holding the lock must re-scan before releasing. Prevents a
+/// concurrent `dlopen` from being missed until the next unrelated scan.
+#[cfg(all(target_os = "linux", target_pointer_width = "64"))]
+static RESCAN_PENDING: AtomicBool = AtomicBool::new(false);
+
+/// Run `update_overrides` while the lock is held until no rescan is pending.
+/// Called by the lock holder after its own scan so a `dlopen` that raced the
+/// lock still gets picked up. `update_overrides` is cheap when `dlpi_adds`
+/// hasn't changed, so a spurious drain is nearly free.
+#[cfg(all(target_os = "linux", target_pointer_width = "64"))]
+fn drain_pending_rescans(so: &mut SymbolOverrides) {
+    while RESCAN_PENDING.swap(false, Ordering::AcqRel) {
+        so.update_overrides();
+    }
+}
+
 #[cfg(all(target_os = "linux", target_pointer_width = "64"))]
 fn lock_global_overrides() -> MutexGuard<'static, Option<SymbolOverrides>> {
     GLOBAL_OVERRIDES
         .lock()
+        // Recover from poison: hot path never takes this lock, registry re-applies idempotently.
         .unwrap_or_else(|poisoned| poisoned.into_inner())
 }
 
@@ -111,6 +132,7 @@ pub fn install_heap_overrides() -> bool {
     }
     let so = guard.as_mut().unwrap();
     so.apply_overrides();
+    drain_pending_rescans(so);
     // Heuristic: at least one ORIG slot resolved.
     any_orig_resolved()
 }
@@ -131,11 +153,18 @@ pub fn update_heap_overrides() {
     // outer apply_overrides, which already walks every library.
     let mut guard = match GLOBAL_OVERRIDES.try_lock() {
         Ok(guard) => guard,
+        // Recover from poison: same rationale as lock_global_overrides.
         Err(TryLockError::Poisoned(poisoned)) => poisoned.into_inner(),
-        Err(TryLockError::WouldBlock) => return,
+        // Lock held by another install/update: flag a rescan so that thread
+        // picks up the library we just loaded before it releases the lock.
+        Err(TryLockError::WouldBlock) => {
+            RESCAN_PENDING.store(true, Ordering::Release);
+            return;
+        }
     };
     if let Some(so) = guard.as_mut() {
         so.update_overrides();
+        drain_pending_rescans(so);
     }
 }
 
@@ -179,55 +208,31 @@ pub fn test_hook_hits() -> u64 {
 #[cfg(all(target_os = "linux", target_pointer_width = "64"))]
 fn register_all(so: &mut SymbolOverrides) {
     use hooks::*;
-    use std::sync::atomic::AtomicUsize;
 
     // Register one entry per supported symbol. The install path stores
     // via `store(Release)` and hooks read via `load(Acquire)`; both go
     // through the typed atomic to avoid racing plain writes against
     // atomic loads.
-    fn reg(so: &mut SymbolOverrides, name: &str, hook_addr: usize, slot: &'static AtomicUsize) {
-        so.register(name, hook_addr, slot);
-    }
-
-    reg(
-        so,
-        "malloc",
-        gotter_malloc as *const () as usize,
-        &ORIG_MALLOC,
-    );
-    reg(so, "free", gotter_free as *const () as usize, &ORIG_FREE);
-    reg(
-        so,
-        "calloc",
-        gotter_calloc as *const () as usize,
-        &ORIG_CALLOC,
-    );
-    reg(
-        so,
+    so.register("malloc", gotter_malloc as *const () as usize, &ORIG_MALLOC);
+    so.register("free", gotter_free as *const () as usize, &ORIG_FREE);
+    so.register("calloc", gotter_calloc as *const () as usize, &ORIG_CALLOC);
+    so.register(
         "realloc",
         gotter_realloc as *const () as usize,
         &ORIG_REALLOC,
     );
-    reg(
-        so,
+    so.register(
         "posix_memalign",
         gotter_posix_memalign as *const () as usize,
         &ORIG_POSIX_MEMALIGN,
     );
-    reg(
-        so,
+    so.register(
         "aligned_alloc",
         gotter_aligned_alloc as *const () as usize,
         &ORIG_ALIGNED_ALLOC,
     );
-    reg(
-        so,
-        "dlopen",
-        gotter_dlopen as *const () as usize,
-        &ORIG_DLOPEN,
-    );
-    reg(
-        so,
+    so.register("dlopen", gotter_dlopen as *const () as usize, &ORIG_DLOPEN);
+    so.register(
         "pthread_create",
         gotter_pthread_create as *const () as usize,
         &ORIG_PTHREAD_CREATE,

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

@@ -24,11 +24,30 @@ use std::ffi::c_void;
 use libdd_profiling_heap_sampler::{dd_sample_flag_peek, dd_tl_state_get_or_init};
 use serial_test::serial;
 
+/// Warn (once) when these tests aren't running under nextest. The GOT install
+/// is permanent, so isolation depends on each test getting its own process;
+/// nextest sets `NEXTEST=1`, plain `cargo test` shares one process and leaks
+/// a prior test's install into later ones.
+fn warn_if_not_isolated() {
+    use std::sync::Once;
+    static ONCE: Once = Once::new();
+    ONCE.call_once(|| {
+        if std::env::var_os("NEXTEST").is_none() {
+            eprintln!(
+                "warning: heap-gotter install tests need per-test process isolation; \
+                 run with `cargo nextest run`. Under `cargo test` a prior test's \
+                 permanent GOT install leaks into later ones."
+            );
+        }
+    });
+}
+
 /// After install the heap should still be functional and no recursive
 /// crash should occur when malloc/free go through the patched GOT.
 #[test]
 #[serial]
 fn install_keeps_heap_functional() {
+    warn_if_not_isolated();
     extern "C" {
         fn malloc(size: usize) -> *mut c_void;
     }
@@ -40,9 +59,17 @@ fn install_keeps_heap_functional() {
     );
 
     unsafe {
-        let p = malloc(64);
+        let p = malloc(64) as *mut u8;
         assert!(!p.is_null(), "malloc returned NULL post-install");
-        libc::free(p);
+        // Write then read back the whole buffer to prove it's real, usable
+        // memory, not just a non-NULL (possibly mis-offset) pointer.
+        for i in 0..64 {
+            p.add(i).write((i as u8) ^ 0xA5);
+        }
+        for i in 0..64 {
+            assert_eq!(p.add(i).read(), (i as u8) ^ 0xA5, "byte {i} corrupted");
+        }
+        libc::free(p as *mut c_void);
     }
 }
 
@@ -54,6 +81,7 @@ fn install_keeps_heap_functional() {
 #[test]
 #[serial]
 fn install_produces_sampled_allocations() {
+    warn_if_not_isolated();
     let installed = libdd_profiling_heap_gotter::install_heap_overrides();
     assert!(installed);
 
@@ -97,6 +125,7 @@ fn install_produces_sampled_allocations() {
 #[test]
 #[serial]
 fn realloc_null_produces_sampled_allocation() {
+    warn_if_not_isolated();
     let installed = libdd_profiling_heap_gotter::install_heap_overrides();
     assert!(installed);
 
@@ -130,6 +159,7 @@ fn realloc_null_produces_sampled_allocation() {
 #[test]
 #[serial]
 fn page_aligned_allocations_are_unsampled() {
+    warn_if_not_isolated();
     let installed = libdd_profiling_heap_gotter::install_heap_overrides();
     assert!(installed);
 
@@ -162,6 +192,7 @@ fn page_aligned_allocations_are_unsampled() {
 #[test]
 #[serial]
 fn realloc_of_sampled_allocation_preserves_data() {
+    warn_if_not_isolated();
     let installed = libdd_profiling_heap_gotter::install_heap_overrides();
     assert!(installed);
 
@@ -246,6 +277,7 @@ unsafe fn alloc_aligned(align: usize, size: usize) -> *mut c_void {
 #[test]
 #[serial]
 fn realloc_stress_across_alignments_preserves_data() {
+    warn_if_not_isolated();
     // Mirrors the demo's menu, plus 2048 to bracket the 1024 cap on both
     // sides. Small alignments sample; those above the cap pass through.
     const ALIGNMENTS: &[usize] = &[1, 8, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096, 8192];
@@ -336,3 +368,32 @@ fn realloc_stress_across_alignments_preserves_data() {
     #[cfg(not(feature = "live-heap"))]
     let _ = saw_sampled;
 }
+
+/// Invalid alignments must surface the real allocator's error unchanged:
+/// aligned_alloc returns NULL + errno EINVAL, posix_memalign returns EINVAL.
+/// The errno check also guards that our post-call sampler bookkeeping doesn't
+/// clobber the real allocator's errno.
+#[test]
+#[serial]
+fn invalid_alignment_passes_through_error() {
+    warn_if_not_isolated();
+    let installed = libdd_profiling_heap_gotter::install_heap_overrides();
+    assert!(installed);
+
+    unsafe {
+        // aligned_alloc: non-power-of-two alignment -> NULL + errno EINVAL.
+        *libc::__errno_location() = 0;
+        let p = libc::aligned_alloc(3, 9);
+        assert!(p.is_null(), "invalid aligned_alloc should fail");
+        assert_eq!(
+            std::io::Error::last_os_error().raw_os_error(),
+            Some(libc::EINVAL),
+            "errno must survive our post-call bookkeeping"
+        );
+
+        // posix_memalign: invalid alignment -> EINVAL return value.
+        let mut out: *mut c_void = std::ptr::null_mut();
+        let rc = libc::posix_memalign(&mut out, 3, 64);
+        assert_eq!(rc, libc::EINVAL, "posix_memalign should return EINVAL");
+    }
+}

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

@@ -22,6 +22,14 @@ default = []
 # against allocs for retained/live-heap profiling, and sample
 # frees to emit ddheap:free.
 live-heap = []
+# Runtime ELF self-inspection of the USDT notes (see src/usdt_check.rs). Opt-in
+# so the crate keeps zero runtime deps by default; only pulls in elf/anyhow.
+sanity-check = ["dep:elf", "dep:anyhow"]
+
+# Only used by the `sanity-check` feature; the crate has no runtime deps otherwise.
+[dependencies]
+anyhow = { version = "1.0", optional = true }
+elf = { version = "0.7", optional = true }
 
 [build-dependencies]
 cc = "1.1.31"

libdd-profiling-heap-sampler/Cargo.toml

@@ -95,6 +95,14 @@ mod linux {
         // Translate the `live-heap` cargo feature into a C define.
         let live_heap = env::var_os("CARGO_FEATURE_LIVE_HEAP").is_some();
 
+        // Define NDEBUG for optimized builds so C `assert()`s are stripped in
+        // release and live in dev/test (incl. `cargo test`). cc doesn't set
+        // this itself; `.debug()` only controls `-g`. Key off OPT_LEVEL, not
+        // DEBUG: this workspace's release profile keeps line-table debuginfo,
+        // so cargo reports DEBUG=true there too; OPT_LEVEL is "0" only for the
+        // unoptimized dev/test profile.
+        let optimized = env::var("OPT_LEVEL").as_deref() != Ok("0");
+
         let mut build = cc::Build::new();
         build
             .files(SOURCES)
@@ -112,6 +120,9 @@ mod linux {
             // loads it works on both glibc and musl without allocation
             // concerns (see tl_state.h for the full analysis).
             .flag_if_supported("-mtls-dialect=gnu2");
+        if optimized {
+            build.define("NDEBUG", None);
+        }
         build.compile("dd_heap_sampler");
 
         // `allocation_requested.c` calls `log()`; glibc keeps it in libm,

libdd-profiling-heap-sampler/build.rs

@@ -29,7 +29,7 @@ this case.
 Although x86-64 typically supports pointer tagging and it was (briefly) enabled
 in the kernel, it was pulled back out around Spectre due to security concerns. 
 
-Instead we steal, bytes from the allocation itself - when we decide to sample, we ask the
+Instead we steal bytes from the allocation itself - when we decide to sample, we ask the
 underlying allocator for more memory than the caller requested, then place
 the user-visible pointer some way into that block, leaving room for a
 16-byte header just before it.
@@ -69,15 +69,16 @@ an unsampled allocation, and if that allocation's contents happen to look
 like our magic at the right offset, we'd wrongly treat it as sampled. Zeroing
 prevents that.
 
-**Keeping the three sites in sync**
+**Keeping the sizing in one place**
 
-The formula for the bumped allocation size has to be computed identically in
-three places: when we decide how much extra to ask for (`bumped_alloc_size`
-in `allocation_requested.c`), when we place the header and user pointer
-(`x86_apply` in `sample_flag.h`), and when we work out how big the original
-allocation was so we can free the right amount (`dd_allocation_freed_slow`
-in `allocation_freed.c`). If any of these disagree we corrupt memory. Touch
-one, check the other two.
+The bumped allocation size is computed by a single shared inline,
+`x86_bumped_size` in `sample_flag.h`. Both the alloc side (`bumped_alloc_size`
+in `allocation_requested.c`, deciding how much extra to ask for) and the free
+side (`dd_allocation_freed_slow` in `allocation_freed.c`, recovering the size
+for sized-free) call it, so the formula can't drift between them. The header
+placement (`x86_apply`, also in `sample_flag.h`) shares the same
+`x86_base_offset` helper for the reserve. All the inlines are
+`always_inline`, so this centralisation adds no call overhead.
 
 **Alignment cap**