From be77a1cb935065cfbfc9a2412193b6c659dae5f4 Mon Sep 17 00:00:00 2001 From: Scott Gerring Date: Thu, 25 Jun 2026 09:44:03 +0200 Subject: [PATCH 01/16] feat: add libdd-heap- allocation tracking --- .github/CODEOWNERS | 1 + .../verify-heap-sampler-bindings.yml | 49 ++ Cargo.lock | 57 +- Cargo.toml | 4 + NOTICE | 8 + libdd-heap-allocator/Cargo.toml | 19 + libdd-heap-allocator/README.md | 25 + libdd-heap-allocator/examples/usdt_demo.rs | 61 ++ libdd-heap-allocator/src/allocator.rs | 158 +++++ libdd-heap-allocator/src/lib.rs | 27 + libdd-heap-gotter-ffi/Cargo.toml | 33 + libdd-heap-gotter-ffi/README.md | 44 ++ libdd-heap-gotter-ffi/build.rs | 16 + libdd-heap-gotter-ffi/cbindgen.toml | 37 ++ libdd-heap-gotter-ffi/examples/cdylib_demo.rs | 148 +++++ libdd-heap-gotter-ffi/src/lib.rs | 74 +++ libdd-heap-gotter-ffi/tests/install.rs | 48 ++ libdd-heap-gotter/Cargo.toml | 24 + libdd-heap-gotter/README.md | 12 + .../examples/gotter_usdt_demo.rs | 99 +++ libdd-heap-gotter/src/elf.rs | 594 ++++++++++++++++++ libdd-heap-gotter/src/hooks.rs | 311 +++++++++ libdd-heap-gotter/src/lib.rs | 247 ++++++++ libdd-heap-gotter/tests/install.rs | 54 ++ libdd-heap-sampler/Cargo.toml | 28 + libdd-heap-sampler/README.md | 131 ++++ libdd-heap-sampler/build.rs | 201 ++++++ .../include/datadog/heap/allocation_created.h | 50 ++ .../include/datadog/heap/allocation_freed.h | 72 +++ .../datadog/heap/allocation_requested.h | 97 +++ .../include/datadog/heap/probes.h | 40 ++ .../include/datadog/heap/sample_flag.h | 175 ++++++ .../include/datadog/heap/tl_state.h | 149 +++++ libdd-heap-sampler/src/allocation_created.c | 46 ++ libdd-heap-sampler/src/allocation_freed.c | 36 ++ libdd-heap-sampler/src/allocation_requested.c | 115 ++++ libdd-heap-sampler/src/generated/bindings.rs | 149 +++++ .../dd_heap_sampler_static_wrappers.c | 20 + libdd-heap-sampler/src/lib.rs | 111 ++++ libdd-heap-sampler/src/probes.c | 28 + libdd-heap-sampler/src/sample_flag.c | 94 +++ libdd-heap-sampler/src/tl_state.c | 72 +++ libdd-heap-sampler/vendor/README.md | 32 + libdd-heap-sampler/vendor/usdt.h | 545 ++++++++++++++++ tools/docker/Dockerfile.build | 5 + 45 files changed, 4345 insertions(+), 1 deletion(-) create mode 100644 .github/workflows/verify-heap-sampler-bindings.yml create mode 100644 libdd-heap-allocator/Cargo.toml create mode 100644 libdd-heap-allocator/README.md create mode 100644 libdd-heap-allocator/examples/usdt_demo.rs create mode 100644 libdd-heap-allocator/src/allocator.rs create mode 100644 libdd-heap-allocator/src/lib.rs create mode 100644 libdd-heap-gotter-ffi/Cargo.toml create mode 100644 libdd-heap-gotter-ffi/README.md create mode 100644 libdd-heap-gotter-ffi/build.rs create mode 100644 libdd-heap-gotter-ffi/cbindgen.toml create mode 100644 libdd-heap-gotter-ffi/examples/cdylib_demo.rs create mode 100644 libdd-heap-gotter-ffi/src/lib.rs create mode 100644 libdd-heap-gotter-ffi/tests/install.rs create mode 100644 libdd-heap-gotter/Cargo.toml create mode 100644 libdd-heap-gotter/README.md create mode 100644 libdd-heap-gotter/examples/gotter_usdt_demo.rs create mode 100644 libdd-heap-gotter/src/elf.rs create mode 100644 libdd-heap-gotter/src/hooks.rs create mode 100644 libdd-heap-gotter/src/lib.rs create mode 100644 libdd-heap-gotter/tests/install.rs create mode 100644 libdd-heap-sampler/Cargo.toml create mode 100644 libdd-heap-sampler/README.md create mode 100644 libdd-heap-sampler/build.rs create mode 100644 libdd-heap-sampler/include/datadog/heap/allocation_created.h create mode 100644 libdd-heap-sampler/include/datadog/heap/allocation_freed.h create mode 100644 libdd-heap-sampler/include/datadog/heap/allocation_requested.h create mode 100644 libdd-heap-sampler/include/datadog/heap/probes.h create mode 100644 libdd-heap-sampler/include/datadog/heap/sample_flag.h create mode 100644 libdd-heap-sampler/include/datadog/heap/tl_state.h create mode 100644 libdd-heap-sampler/src/allocation_created.c create mode 100644 libdd-heap-sampler/src/allocation_freed.c create mode 100644 libdd-heap-sampler/src/allocation_requested.c create mode 100644 libdd-heap-sampler/src/generated/bindings.rs create mode 100644 libdd-heap-sampler/src/generated/dd_heap_sampler_static_wrappers.c create mode 100644 libdd-heap-sampler/src/lib.rs create mode 100644 libdd-heap-sampler/src/probes.c create mode 100644 libdd-heap-sampler/src/sample_flag.c create mode 100644 libdd-heap-sampler/src/tl_state.c create mode 100644 libdd-heap-sampler/vendor/README.md create mode 100644 libdd-heap-sampler/vendor/usdt.h diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS index 4836f91915..0ccd3849bc 100644 --- a/.github/CODEOWNERS +++ b/.github/CODEOWNERS @@ -49,6 +49,7 @@ libdd-crashtracker*/ @DataDog/libdatadog-profiling libdd-data-pipeline*/ @DataDog/libdatadog-apm libdd-ddsketch*/ @DataDog/libdatadog-apm @DataDog/apm-common-components-core libdd-dogstatsd-client @DataDog/apm-common-components-core +libdd-heap-*/ @DataDog/libdatadog-profiling libdd-http-client @DataDog/apm-common-components-core libdd-agent-client @DataDog/apm-common-components-core libdd-library-config*/ @DataDog/apm-sdk-capabilities-rust diff --git a/.github/workflows/verify-heap-sampler-bindings.yml b/.github/workflows/verify-heap-sampler-bindings.yml new file mode 100644 index 0000000000..b6e06a2655 --- /dev/null +++ b/.github/workflows/verify-heap-sampler-bindings.yml @@ -0,0 +1,49 @@ +name: 'Verify libdd-heap-sampler bindings' +on: + pull_request: + types: [opened, synchronize, reopened] + paths: + - 'libdd-heap-sampler/**' + - '.github/workflows/verify-heap-sampler-bindings.yml' +env: + CARGO_TERM_COLOR: always + CARGO_INCREMENTAL: 0 +jobs: + verify-bindings: + name: "Verify libdd-heap-sampler generated bindings are in sync" + runs-on: ubuntu-latest + steps: + - name: Checkout sources + uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # 4.2.2 + - name: Install libclang-dev + # Only this verification job needs libclang; the normal build + # path for libdd-heap-sampler consumes the checked-in bindings + # under `src/generated/` and does not depend on bindgen at all. + run: | + sudo apt-get update + sudo apt-get install -y libclang-dev + - name: Read Rust version from rust-toolchain.toml + id: rust-version + run: echo "version=$(grep -Po '^channel = "\K[^"]+' rust-toolchain.toml)" >> $GITHUB_OUTPUT + - name: Install ${{ steps.rust-version.outputs.version }} toolchain + run: rustup set profile minimal && rustup install ${{ steps.rust-version.outputs.version }} && rustup default ${{ steps.rust-version.outputs.version }} + - name: Cache [rust] + uses: Swatinem/rust-cache@f13886b937689c021905a6b90929199931d60db1 # 2.8.1 + with: + cache-targets: true + cache-bin: true + - name: Regenerate bindings + # `LIBDD_HEAP_SAMPLER_REGEN=1` rewrites `src/generated/*` + # in-place via bindgen. If the committed files were stale the + # git diff below fails and instructs the author how to refresh + # them. We use an env var (rather than a cargo feature) so that + # unrelated `--all-features` CI jobs cannot accidentally invoke + # bindgen — see libdd-heap-sampler/Cargo.toml for the rationale. + run: LIBDD_HEAP_SAMPLER_REGEN=1 cargo build -p libdd-heap-sampler + - name: Verify committed bindings match regenerated output + run: | + if ! git diff --exit-code libdd-heap-sampler/src/generated/; then + echo "::error::libdd-heap-sampler/src/generated/ is out of date." + echo "Run \`LIBDD_HEAP_SAMPLER_REGEN=1 cargo build -p libdd-heap-sampler\` locally and commit the result." + exit 1 + fi diff --git a/Cargo.lock b/Cargo.lock index 3df371d202..ccc9d77d4d 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -357,7 +357,7 @@ version = "0.13.13" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f8bce4948d2520386c6d92a6ea2d472300257702242e5a1d01d6add52bd2e7c1" dependencies = [ - "bindgen", + "bindgen 0.72.1", "cc", "cmake", "dunce", @@ -488,6 +488,26 @@ dependencies = [ "serde", ] +[[package]] +name = "bindgen" +version = "0.71.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5f58bf3d7db68cfbac37cfc485a8d711e87e064c3d0fe0435b92f7a407f9d6b3" +dependencies = [ + "bitflags", + "cexpr", + "clang-sys", + "itertools", + "log", + "prettyplease", + "proc-macro2", + "quote", + "regex", + "rustc-hash", + "shlex", + "syn 2.0.87", +] + [[package]] name = "bindgen" version = "0.72.1" @@ -3078,6 +3098,41 @@ dependencies = [ "tracing", ] +[[package]] +name = "libdd-heap-allocator" +version = "0.1.0" +dependencies = [ + "libdd-heap-sampler", +] + +[[package]] +name = "libdd-heap-gotter" +version = "0.1.0" +dependencies = [ + "libc", + "libdd-heap-sampler", +] + +[[package]] +name = "libdd-heap-gotter-ffi" +version = "37.0.0" +dependencies = [ + "anyhow", + "build_common", + "function_name", + "libc", + "libdd-common-ffi", + "libdd-heap-gotter", +] + +[[package]] +name = "libdd-heap-sampler" +version = "0.1.0" +dependencies = [ + "bindgen 0.71.1", + "cc", +] + [[package]] name = "libdd-http-client" version = "37.0.0" diff --git a/Cargo.toml b/Cargo.toml index a4a09f13a8..f0474f32e2 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -5,6 +5,10 @@ members = [ "builder", "libdd-alloc", + "libdd-heap-sampler", + "libdd-heap-allocator", + "libdd-heap-gotter", + "libdd-heap-gotter-ffi", "libdd-crashtracker", "libdd-crashtracker-ffi", "datadog-ffe", diff --git a/NOTICE b/NOTICE index ab7dccadf5..07440e227e 100644 --- a/NOTICE +++ b/NOTICE @@ -2,3 +2,11 @@ Datadog libdatadog Copyright 2021-2022 Datadog, Inc. This product includes software developed at Datadog (). + +-- + +This product bundles a copy of the libbpf/usdt single-header USDT +library in `libdd-heap-sampler/vendor/usdt.h`. That file is licensed +under the BSD 2-Clause License, Copyright (c) 2024 Meta Platforms, Inc. +and affiliates. The SPDX identifier and copyright notice are retained +verbatim in the file header. Upstream: . diff --git a/libdd-heap-allocator/Cargo.toml b/libdd-heap-allocator/Cargo.toml new file mode 100644 index 0000000000..56f9d5ff19 --- /dev/null +++ b/libdd-heap-allocator/Cargo.toml @@ -0,0 +1,19 @@ +# Copyright 2025-Present Datadog, Inc. https://www.datadoghq.com/ +# SPDX-License-Identifier: Apache-2.0 + +[package] +name = "libdd-heap-allocator" +version = "0.1.0" +description = "Rust GlobalAlloc wrapper that drives libdd-heap-sampler." +homepage = "https://github.com/DataDog/libdatadog/tree/main/libdd-heap-allocator" +repository = "https://github.com/DataDog/libdatadog/tree/main/libdd-heap-allocator" +edition.workspace = true +rust-version.workspace = true +license.workspace = true +publish = false + +[lib] +bench = false + +[dependencies] +libdd-heap-sampler = { path = "../libdd-heap-sampler" } diff --git a/libdd-heap-allocator/README.md b/libdd-heap-allocator/README.md new file mode 100644 index 0000000000..8e07d5838d --- /dev/null +++ b/libdd-heap-allocator/README.md @@ -0,0 +1,25 @@ +# libdd-heap-allocator + +Rust `GlobalAlloc` wrapper with USDT-based heap profiling, effectively implementing [libdd-heap-sampler](../libdd-heap-sampler) for Rust apps at compile time. This lets Rust users quickly setup sampled heap profiling within their application regardless of the particular allocator they are using. + +For this to work _well_, you should make sure everything passes through the global allocator! + +Usage: + +```rust +use libdd_heap_allocator::SampledAllocator; +use std::alloc::System; + +// Wrap the default system allocator +#[global_allocator] +static ALLOC: SampledAllocator = SampledAllocator::::DEFAULT; +``` + +To wrap a custom allocator instead; note that this is kind of ill-advised; we want to see _all_ allocations for the process: + +```rust +#[global_allocator] +static ALLOC: SampledAllocator = SampledAllocator::new(MyAllocator::new()); +``` + +See [`examples/usdt_demo.rs`](examples/usdt_demo.rs) for a runnable demo that fires USDT probes in a loop for `bpftrace` to observe. \ No newline at end of file diff --git a/libdd-heap-allocator/examples/usdt_demo.rs b/libdd-heap-allocator/examples/usdt_demo.rs new file mode 100644 index 0000000000..ebd4829dcd --- /dev/null +++ b/libdd-heap-allocator/examples/usdt_demo.rs @@ -0,0 +1,61 @@ +// Copyright 2025-Present Datadog, Inc. https://www.datadoghq.com/ +// SPDX-License-Identifier: Apache-2.0 + +//! Sample app exercising `libdd-heap-allocator` as the global allocator. +//! +//! Install `SampledAllocator` globally, then loop producing +//! allocations (strings joined into a single buffer) so a tracer attached +//! to the `ddheap:alloc` USDT probe sees samples fire periodically. +//! +//! Run (Linux, inside the crate's Lima VM): +//! ``` +//! cargo run --example usdt_demo -p libdd-heap-allocator +//! ``` +//! and in another shell, attach a tracer, e.g. +//! ``` +//! sudo bpftrace -p -e 'usdt:*:ddheap:alloc { printf("alloc %p %d %d\n", arg0, arg1, arg2); }' +//! ``` +//! +//! `SampledAllocator` is Linux-only; on other targets the example +//! compiles to an empty `main` so clippy/test on non-Linux don't fail +//! with "configured out". + +#[cfg(not(target_os = "linux"))] +fn main() {} + +#[cfg(target_os = "linux")] +fn main() { + linux::main(); +} + +#[cfg(target_os = "linux")] +mod linux { + use libdd_heap_allocator::SampledAllocator; + use std::alloc::System; + use std::thread::sleep; + use std::time::Duration; + + #[global_allocator] + static ALLOC: SampledAllocator = SampledAllocator::::DEFAULT; + + pub fn main() { + println!( + "pid={}; attach a tracer on 'usdt:*:ddheap:alloc'", + std::process::id() + ); + + let mut i: u64 = 0; + loop { + // ~1000 small allocations + one larger join: plenty of alloc + // pressure to cross the default 512 KiB sampling interval over + // a handful of iterations. + let parts: Vec = (0..1000) + .map(|j| format!("chunk-{i}-{j}-with-some-padding-to-make-it-meaningful")) + .collect(); + let joined = parts.join(", "); + println!("[{i}] joined {} bytes", joined.len()); + i = i.wrapping_add(1); + sleep(Duration::from_secs(1)); + } + } +} // mod linux diff --git a/libdd-heap-allocator/src/allocator.rs b/libdd-heap-allocator/src/allocator.rs new file mode 100644 index 0000000000..ac996e0f5d --- /dev/null +++ b/libdd-heap-allocator/src/allocator.rs @@ -0,0 +1,158 @@ +// Copyright 2025-Present Datadog, Inc. https://www.datadoghq.com/ +// SPDX-License-Identifier: Apache-2.0 + +use core::alloc::{GlobalAlloc, Layout}; +use std::alloc::System; + +#[cfg(target_os = "linux")] +use libdd_heap_sampler::{dd_allocation_created, dd_allocation_freed, dd_allocation_requested}; + +/// `GlobalAlloc` wrapper that routes each alloc/dealloc through +/// `libdd-heap-sampler` before forwarding to the inner allocator `A`. +/// +/// The default `realloc` / `alloc_zeroed` impls from [`GlobalAlloc`] are +/// inherited; they dispatch back to `alloc` / `dealloc`, so sampling +/// still fires for those paths. +/// +/// On non-Linux targets the sampler integration is a no-op and this is +/// just a transparent forwarder to `inner`; callers can use the same +/// `#[global_allocator]` setup unconditionally. +pub struct SampledAllocator { + inner: A, +} + +impl SampledAllocator { + /// Wrap an allocator. `const` so it's usable directly in a + /// `#[global_allocator]` static. + pub const fn new(inner: A) -> Self { + Self { inner } + } +} + +impl SampledAllocator { + /// Default wrap of the system allocator, usable directly in a + /// `#[global_allocator]` static. + pub const DEFAULT: Self = Self { inner: System }; +} + +unsafe impl GlobalAlloc for SampledAllocator { + #[cfg(target_os = "linux")] + unsafe fn alloc(&self, layout: Layout) -> *mut u8 { + let req = dd_allocation_requested(layout.size(), layout.align()); + // Sampled paths may bump the size for inline flag storage; + // forward the returned size to the inner allocator verbatim. + let inner_layout = Layout::from_size_align_unchecked(req.size, layout.align()); + let raw = self.inner.alloc(inner_layout); + dd_allocation_created(raw.cast(), req).cast() + } + + #[cfg(not(target_os = "linux"))] + unsafe fn alloc(&self, layout: Layout) -> *mut u8 { + self.inner.alloc(layout) + } + + #[cfg(target_os = "linux")] + unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) { + let freed = dd_allocation_freed(ptr.cast(), layout.size(), layout.align()); + let inner_layout = Layout::from_size_align_unchecked(freed.size, layout.align()); + self.inner.dealloc(freed.ptr.cast(), inner_layout); + } + + #[cfg(not(target_os = "linux"))] + unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) { + self.inner.dealloc(ptr, layout); + } +} + +// Tests dispatch through the wrapped allocator and (on Linux) into the +// sampler's C-side TLS primitives via FFI; miri can't execute those. +#[cfg(all(test, not(miri)))] +mod tests { + use super::*; + use core::sync::atomic::{AtomicUsize, Ordering}; + #[cfg(target_os = "linux")] + use libdd_heap_sampler::dd_tl_state_get; + + /// Minimal `GlobalAlloc` that forwards to `System` while recording + /// counters so tests can assert the sampled wrapper passed the right + /// size/align through. + struct CountingSystem { + alloc_count: AtomicUsize, + dealloc_count: AtomicUsize, + last_alloc_size: AtomicUsize, + } + + impl CountingSystem { + const fn new() -> Self { + Self { + alloc_count: AtomicUsize::new(0), + dealloc_count: AtomicUsize::new(0), + last_alloc_size: AtomicUsize::new(0), + } + } + } + + unsafe impl GlobalAlloc for CountingSystem { + unsafe fn alloc(&self, layout: Layout) -> *mut u8 { + self.alloc_count.fetch_add(1, Ordering::Relaxed); + self.last_alloc_size.store(layout.size(), Ordering::Relaxed); + System.alloc(layout) + } + unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) { + self.dealloc_count.fetch_add(1, Ordering::Relaxed); + System.dealloc(ptr, layout); + } + } + + #[test] + fn alloc_dealloc_forwards_to_inner() { + let sampled = SampledAllocator::new(CountingSystem::new()); + let layout = Layout::from_size_align(128, 16).unwrap(); + + unsafe { + let p = sampled.alloc(layout); + assert!(!p.is_null()); + + assert_eq!(sampled.inner.alloc_count.load(Ordering::Relaxed), 1); + // Sampler may bump size for flag storage once that's wired up; + // today it returns the requested size verbatim. + assert!(sampled.inner.last_alloc_size.load(Ordering::Relaxed) >= 128); + + sampled.dealloc(p, layout); + assert_eq!(sampled.inner.dealloc_count.load(Ordering::Relaxed), 1); + } + } + + // 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. + std::thread::spawn(|| unsafe { + assert!( + dd_tl_state_get().is_null(), + "fresh thread should have NULL sampler TLS" + ); + + 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" + ); + + sampled.dealloc(p, layout); + }) + .join() + .unwrap(); + } + + #[test] + fn default_matches_new_system() { + let _ = SampledAllocator::::DEFAULT; + let _ = SampledAllocator::new(System); + } +} diff --git a/libdd-heap-allocator/src/lib.rs b/libdd-heap-allocator/src/lib.rs new file mode 100644 index 0000000000..7d53c61249 --- /dev/null +++ b/libdd-heap-allocator/src/lib.rs @@ -0,0 +1,27 @@ +// Copyright 2025-Present Datadog, Inc. https://www.datadoghq.com/ +// SPDX-License-Identifier: Apache-2.0 + +//! Rust `GlobalAlloc` wrapper that drives `libdd-heap-sampler` around each +//! allocation. Wrap any underlying allocator with [`SampledAllocator`]; on +//! each alloc/dealloc the sampler's decision/flag/USDT machinery runs +//! around the inner call. +//! +//! `SampledAllocator` is portable across targets, so callers can use it in +//! a single `#[global_allocator]` static without cfg-gating. The sampling +//! integration (USDT probes via `libdd-heap-sampler`) is Linux-only; on +//! every other target the wrapper compiles to a transparent pass-through +//! to the inner allocator. +//! +//! # Example +//! +//! ```no_run +//! use libdd_heap_allocator::SampledAllocator; +//! use std::alloc::System; +//! +//! #[global_allocator] +//! static ALLOC: SampledAllocator = SampledAllocator::::DEFAULT; +//! ``` + +mod allocator; + +pub use allocator::SampledAllocator; diff --git a/libdd-heap-gotter-ffi/Cargo.toml b/libdd-heap-gotter-ffi/Cargo.toml new file mode 100644 index 0000000000..d754ddea11 --- /dev/null +++ b/libdd-heap-gotter-ffi/Cargo.toml @@ -0,0 +1,33 @@ +# Copyright 2025-Present Datadog, Inc. https://www.datadoghq.com/ +# SPDX-License-Identifier: Apache-2.0 + +[package] +name = "libdd-heap-gotter-ffi" +edition.workspace = true +version.workspace = true +rust-version.workspace = true +license.workspace = true +publish = false + +[lib] +crate-type = ["lib", "staticlib", "cdylib"] +bench = false + +[features] +default = ["cbindgen"] +cbindgen = ["build_common/cbindgen", "libdd-common-ffi/cbindgen"] + +[build-dependencies] +build_common = { path = "../build-common" } + +[dependencies] +anyhow = "1.0" +function_name = "0.3.0" +libdd-common-ffi = { path = "../libdd-common-ffi", default-features = false } +# libdd-heap-gotter exposes the same public surface on every target +# (no-ops on non-Linux), so this dependency is unconditional and the +# FFI layer is portable too. +libdd-heap-gotter = { path = "../libdd-heap-gotter" } + +[target.'cfg(target_os = "linux")'.dev-dependencies] +libc = "0.2" diff --git a/libdd-heap-gotter-ffi/README.md b/libdd-heap-gotter-ffi/README.md new file mode 100644 index 0000000000..8b1e5e90d5 --- /dev/null +++ b/libdd-heap-gotter-ffi/README.md @@ -0,0 +1,44 @@ +# libdd-heap-gotter-ffi + +C FFI bindings for `libdd-heap-gotter`. + +## Overview + +`libdd-heap-gotter-ffi` exposes a small C ABI for installing heap-profiling GOT table interposition from language runtimes such as Python and Ruby. + +The API installs hooks for supported allocator symbols, updates hooks after new libraries are loaded, and restores patched GOT entries when profiling is disabled. + +## API + +- `ddog_heap_gotter_install()` - install heap GOT overrides in the current process. +- `ddog_heap_gotter_update()` - re-scan loaded libraries and patch newly introduced GOT entries. +- `ddog_heap_gotter_restore()` - restore every GOT entry patched by install/update. +- `ddog_heap_gotter_is_installed()` - return whether overrides are currently installed. + +## Important lifetime note + +The shared library containing these hooks must remain loaded while overrides are installed. Patched GOT entries point at functions in this library, so unloading it before calling `ddog_heap_gotter_restore()` can leave dangling function pointers and crash the process. + +## Building + +This crate follows the standard libdatadog FFI layout: it builds `staticlib` and `cdylib` artifacts and generates a C header with cbindgen. + +```bash +cargo check -p libdd-heap-gotter-ffi +``` + +## Dynamic-loading demo + +The `cdylib_demo` example loads the generated shared library with `dlopen`, resolves the C ABI symbols with `dlsym`, installs the GOT hooks, and produces allocation pressure. + +```bash +cargo build -p libdd-heap-gotter-ffi +cargo run -p libdd-heap-gotter-ffi --example cdylib_demo +``` + +If the cdylib is somewhere else, set: + +```bash +DDOG_HEAP_GOTTER_FFI_CDYLIB=/path/to/liblibdd_heap_gotter_ffi.so \ + cargo run -p libdd-heap-gotter-ffi --example cdylib_demo +``` diff --git a/libdd-heap-gotter-ffi/build.rs b/libdd-heap-gotter-ffi/build.rs new file mode 100644 index 0000000000..b61ab698a8 --- /dev/null +++ b/libdd-heap-gotter-ffi/build.rs @@ -0,0 +1,16 @@ +// Copyright 2025-Present Datadog, Inc. https://www.datadoghq.com/ +// SPDX-License-Identifier: Apache-2.0 + +//! Build script for `libdd-heap-gotter-ffi`. Generates the C header for +//! the FFI surface via cbindgen. + +extern crate build_common; + +use build_common::generate_and_configure_header; + +fn main() { + println!("cargo:rerun-if-changed=src/*"); + println!("cargo:rerun-if-changed=cbindgen.toml"); + println!("cargo:rerun-if-changed=build.rs"); + generate_and_configure_header("heap_gotter.h"); +} diff --git a/libdd-heap-gotter-ffi/cbindgen.toml b/libdd-heap-gotter-ffi/cbindgen.toml new file mode 100644 index 0000000000..ff29d2d549 --- /dev/null +++ b/libdd-heap-gotter-ffi/cbindgen.toml @@ -0,0 +1,37 @@ +# Copyright 2025-Present Datadog, Inc. https://www.datadoghq.com/ +# SPDX-License-Identifier: Apache-2.0 + +language = "C" +cpp_compat = true +tab_width = 2 +header = """// Copyright 2025-Present Datadog, Inc. https://www.datadoghq.com/ +// SPDX-License-Identifier: Apache-2.0 +""" +include_guard = "DDOG_HEAP_GOTTER_H" +style = "both" +pragma_once = true +no_includes = true +sys_includes = ["stdbool.h", "stddef.h", "stdint.h"] +includes = ["common.h"] + +[parse] +parse_deps = true +include = ["libdd-common-ffi", "libdd-heap-gotter"] + +[export] +prefix = "ddog_" +renaming_overrides_prefixing = true + +[export.mangle] +rename_types = "PascalCase" + +[export.rename] +"VoidResult" = "ddog_VoidResult" +"Error" = "ddog_Error" + +[enum] +prefix_with_name = true +rename_variants = "ScreamingSnakeCase" + +[fn] +must_use = "DDOG_CHECK_RETURN" diff --git a/libdd-heap-gotter-ffi/examples/cdylib_demo.rs b/libdd-heap-gotter-ffi/examples/cdylib_demo.rs new file mode 100644 index 0000000000..4558008829 --- /dev/null +++ b/libdd-heap-gotter-ffi/examples/cdylib_demo.rs @@ -0,0 +1,148 @@ +// Copyright 2025-Present Datadog, Inc. https://www.datadoghq.com/ +// SPDX-License-Identifier: Apache-2.0 + +//! Demonstrates using `libdd-heap-gotter-ffi` as an actual dynamically-loaded C ABI library. +//! +//! Build the cdylib first: +//! ```sh +//! cargo build -p libdd-heap-gotter-ffi +//! ``` +//! +//! Then run this demo: +//! ```sh +//! cargo run -p libdd-heap-gotter-ffi --example cdylib_demo +//! ``` +//! +//! 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". + +#[cfg(not(target_os = "linux"))] +fn main() -> Result<(), String> { + Ok(()) +} + +#[cfg(target_os = "linux")] +fn main() -> Result<(), String> { + linux::main() +} + +#[cfg(target_os = "linux")] +mod linux { + use libdd_common_ffi::VoidResult; + use std::ffi::{CStr, CString}; + use std::path::{Path, PathBuf}; + use std::thread::sleep; + use std::time::Duration; + + const LIB_NAME: &str = "liblibdd_heap_gotter_ffi.so"; + + type InstallFn = unsafe extern "C" fn() -> VoidResult; + type RestoreFn = unsafe extern "C" fn() -> VoidResult; + type IsInstalledFn = unsafe extern "C" fn() -> bool; + + struct DlopenHandle(*mut libc::c_void); + + impl DlopenHandle { + fn open(path: &Path) -> Result { + let path = + CString::new(path.to_string_lossy().as_bytes()).map_err(|e| e.to_string())?; + let handle = unsafe { libc::dlopen(path.as_ptr(), libc::RTLD_NOW | libc::RTLD_LOCAL) }; + if handle.is_null() { + return Err(dlerror()); + } + Ok(Self(handle)) + } + + unsafe fn symbol(&self, name: &CStr) -> Result + where + T: Copy, + { + let ptr = unsafe { libc::dlsym(self.0, name.as_ptr()) }; + if ptr.is_null() { + return Err(dlerror()); + } + Ok(unsafe { std::mem::transmute_copy(&ptr) }) + } + } + + impl Drop for DlopenHandle { + fn drop(&mut self) { + unsafe { + libc::dlclose(self.0); + } + } + } + + fn dlerror() -> String { + let err = unsafe { libc::dlerror() }; + if err.is_null() { + "unknown dlerror".to_string() + } else { + unsafe { CStr::from_ptr(err) } + .to_string_lossy() + .into_owned() + } + } + + fn cdylib_path() -> Result { + if let Some(path) = std::env::var_os("DDOG_HEAP_GOTTER_FFI_CDYLIB") { + return Ok(PathBuf::from(path)); + } + + let exe = std::env::current_exe().map_err(|e| e.to_string())?; + let profile_dir = exe + .parent() + .and_then(|p| p.parent()) + .ok_or_else(|| format!("could not infer target profile dir from {}", exe.display()))?; + Ok(profile_dir.join(LIB_NAME)) + } + + fn check(result: VoidResult, operation: &str) -> Result<(), String> { + match result { + VoidResult::Ok => Ok(()), + VoidResult::Err(err) => Err(format!("{operation} failed: {err}")), + } + } + + pub fn main() -> Result<(), String> { + let lib_path = cdylib_path()?; + if !lib_path.exists() { + return Err(format!( + "{} does not exist; run `cargo build -p libdd-heap-gotter-ffi` first", + lib_path.display() + )); + } + + println!("pid={}", std::process::id()); + println!("loading {}", lib_path.display()); + let lib = DlopenHandle::open(&lib_path)?; + + let install: InstallFn = unsafe { lib.symbol(c"ddog_heap_gotter_install")? }; + let restore: RestoreFn = unsafe { lib.symbol(c"ddog_heap_gotter_restore")? }; + let is_installed: IsInstalledFn = unsafe { lib.symbol(c"ddog_heap_gotter_is_installed")? }; + + println!("pre-install is_installed={}", unsafe { is_installed() }); + check(unsafe { install() }, "ddog_heap_gotter_install")?; + println!("post-install is_installed={}", unsafe { is_installed() }); + println!("attach a tracer on `usdt:*:ddheap:*`; producing allocation pressure..."); + + for i in 0..30_u64 { + let parts: Vec = (0..1000) + .map(|j| format!("chunk-{i}-{j}-with-some-padding-to-make-it-meaningful")) + .collect(); + let joined = parts.join(", "); + println!("[{i}] joined {} bytes", joined.len()); + sleep(Duration::from_secs(1)); + } + + check(unsafe { restore() }, "ddog_heap_gotter_restore")?; + println!("post-restore is_installed={}", unsafe { is_installed() }); + + // Keep `lib` alive until after restore. The GOT entries patched by install point at + // functions in this cdylib; dropping/dlclosing it while installed can leave + // dangling function pointers. + drop(lib); + Ok(()) + } +} // mod linux diff --git a/libdd-heap-gotter-ffi/src/lib.rs b/libdd-heap-gotter-ffi/src/lib.rs new file mode 100644 index 0000000000..0f02e5489c --- /dev/null +++ b/libdd-heap-gotter-ffi/src/lib.rs @@ -0,0 +1,74 @@ +// Copyright 2025-Present Datadog, Inc. https://www.datadoghq.com/ +// SPDX-License-Identifier: Apache-2.0 + +//! C FFI bindings for [`libdd_heap_gotter`]. Exposes install / update / +//! restore / is-installed entry points as `extern "C"` functions so +//! language runtimes (Python, Ruby, …) can drive GOT-based heap +//! profiling from their own native extension code. + +#![cfg_attr(not(test), deny(clippy::panic))] +#![cfg_attr(not(test), deny(clippy::unwrap_used))] +#![cfg_attr(not(test), deny(clippy::expect_used))] +#![cfg_attr(not(test), deny(clippy::todo))] +#![cfg_attr(not(test), deny(clippy::unimplemented))] + +use function_name::named; +use libdd_common_ffi::{wrap_with_void_ffi_result, VoidResult}; + +// `libdd_heap_gotter` exposes the same public surface on every target. +// On non-Linux the underlying functions are no-ops, so callers that +// invoke these FFI entry points outside Linux observe a clean error +// from `ddog_heap_gotter_install` (nothing was overridden) without +// having to `#ifdef` their integration code. + +/// Install GOT overrides for supported heap-allocation symbols in the current process. +/// +/// The library containing these hooks must remain loaded until +/// `ddog_heap_gotter_restore` has been called. GOT entries are patched to point at +/// functions in this library, so unloading it while installed can leave dangling function pointers. +/// +/// On non-Linux targets this returns an error indicating that nothing +/// could be installed; the rest of the API can still be called safely. +#[no_mangle] +#[must_use] +#[named] +pub extern "C" fn ddog_heap_gotter_install() -> VoidResult { + wrap_with_void_ffi_result!({ + let installed = libdd_heap_gotter::install_heap_overrides(); + anyhow::ensure!(installed, "no heap GOT overrides could be installed"); + }) +} + +/// Re-scan loaded libraries and patch newly-introduced GOT entries. +/// +/// This is normally called automatically by the installed `dlopen` hook, but language runtimes may +/// call it explicitly after unusual native-extension loading flows. No-op on non-Linux targets. +#[no_mangle] +#[must_use] +#[named] +pub extern "C" fn ddog_heap_gotter_update() -> VoidResult { + wrap_with_void_ffi_result!({ + libdd_heap_gotter::update_heap_overrides(); + }) +} + +/// Restore every GOT entry patched by `ddog_heap_gotter_install`. +/// +/// Call this before unloading the shared library that provides the gotter hooks. No-op on +/// non-Linux targets. +#[no_mangle] +#[must_use] +#[named] +pub extern "C" fn ddog_heap_gotter_restore() -> VoidResult { + wrap_with_void_ffi_result!({ + libdd_heap_gotter::restore_heap_overrides(); + }) +} + +/// Return whether heap GOT overrides are currently installed in this process. Always `false` on +/// non-Linux targets. +#[no_mangle] +#[must_use] +pub extern "C" fn ddog_heap_gotter_is_installed() -> bool { + libdd_heap_gotter::heap_overrides_are_installed() +} diff --git a/libdd-heap-gotter-ffi/tests/install.rs b/libdd-heap-gotter-ffi/tests/install.rs new file mode 100644 index 0000000000..ac0b6a7a39 --- /dev/null +++ b/libdd-heap-gotter-ffi/tests/install.rs @@ -0,0 +1,48 @@ +// Copyright 2025-Present Datadog, Inc. https://www.datadoghq.com/ +// SPDX-License-Identifier: Apache-2.0 + +//! Smoke test for the FFI install/restore roundtrip. +//! +//! Installing GOT overrides mutates global process state, so this lives +//! as an integration test in its own binary rather than a unit test. +//! Miri can't execute the underlying dl_iterate_phdr/mprotect calls. + +#![cfg(all(target_os = "linux", not(miri)))] + +use libdd_common_ffi::VoidResult; +use libdd_heap_gotter_ffi::{ + ddog_heap_gotter_install, ddog_heap_gotter_is_installed, ddog_heap_gotter_restore, +}; + +fn assert_ok(result: VoidResult, what: &str) { + match result { + VoidResult::Ok => {} + VoidResult::Err(err) => panic!("{what} failed: {err}"), + } +} + +#[test] +fn install_restore_roundtrip() { + assert!( + !ddog_heap_gotter_is_installed(), + "expected clean process state before install" + ); + + assert_ok(ddog_heap_gotter_install(), "ddog_heap_gotter_install"); + assert!( + ddog_heap_gotter_is_installed(), + "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. + let v: Vec = vec![0; 128]; + assert_eq!(v.len(), 128); + drop(v); + + assert_ok(ddog_heap_gotter_restore(), "ddog_heap_gotter_restore"); + assert!( + !ddog_heap_gotter_is_installed(), + "is_installed should be false after restore" + ); +} diff --git a/libdd-heap-gotter/Cargo.toml b/libdd-heap-gotter/Cargo.toml new file mode 100644 index 0000000000..c15a334220 --- /dev/null +++ b/libdd-heap-gotter/Cargo.toml @@ -0,0 +1,24 @@ +# Copyright 2025-Present Datadog, Inc. https://www.datadoghq.com/ +# SPDX-License-Identifier: Apache-2.0 + +[package] +name = "libdd-heap-gotter" +version = "0.1.0" +description = "GOT-table interposition that routes a running process's allocator through libdd-heap-sampler." +homepage = "https://github.com/DataDog/libdatadog/tree/main/libdd-heap-gotter" +repository = "https://github.com/DataDog/libdatadog/tree/main/libdd-heap-gotter" +edition.workspace = true +rust-version.workspace = true +license.workspace = true +publish = false + +[lib] +crate-type = ["lib", "cdylib"] +bench = false + +[target.'cfg(target_os = "linux")'.dependencies] +libc = "0.2" +libdd-heap-sampler = { path = "../libdd-heap-sampler" } + +[target.'cfg(target_os = "linux")'.dev-dependencies] +libc = "0.2" diff --git a/libdd-heap-gotter/README.md b/libdd-heap-gotter/README.md new file mode 100644 index 0000000000..b1a297ee17 --- /dev/null +++ b/libdd-heap-gotter/README.md @@ -0,0 +1,12 @@ +# libdd-heap-gotter + +GOTter implements our GOT-patching mechanism to wrap (dynamically!) linked allocators in a running process. +This follows the same approach as `ddprof`, and may prove useful to inject via our tracing libraries into running processes such as python. + +It contains: + +* A set of functions such as `gotter_malloc` that will be used to override the _originals_ of these functions +* Overrides for _other bits_ we need for this to work robustly in a running process - including e.g. `fork` +* A function to install the overrides in a running process `install_heap_overrides()` + +This will be used in places such as the python profiler to install the heap profiler at runtime to capture native allocations. diff --git a/libdd-heap-gotter/examples/gotter_usdt_demo.rs b/libdd-heap-gotter/examples/gotter_usdt_demo.rs new file mode 100644 index 0000000000..df703e1014 --- /dev/null +++ b/libdd-heap-gotter/examples/gotter_usdt_demo.rs @@ -0,0 +1,99 @@ +// Copyright 2025-Present Datadog, Inc. https://www.datadoghq.com/ +// SPDX-License-Identifier: Apache-2.0 + +//! Sample app exercising `libdd-heap-gotter`: install the GOT overrides +//! at startup, then loop producing allocations through the *unmodified* +//! Rust system allocator. Once `install_heap_overrides()` has run, those +//! `malloc`/`free` calls resolve through the patched GOT entries and +//! flow into `libdd-heap-sampler`, firing `ddheap:alloc` / `ddheap:free` +//! USDTs. +//! +//! Run (Linux, inside the crate's Lima VM, or natively under a container): +//! ``` +//! cargo run --example gotter_usdt_demo -p libdd-heap-gotter +//! ``` +//! Add `-- --stress` to keep a CPU core busy enough for CPU profiles: +//! ``` +//! cargo run --example gotter_usdt_demo -p libdd-heap-gotter -- --stress +//! ``` +//! and in another shell, attach a tracer, e.g. +//! ``` +//! sudo bpftrace -p -e 'usdt:*:ddheap:alloc { printf("alloc %p %d %d\n", arg0, arg1, arg2); }' +//! ``` +//! +//! The gotter crate is Linux-only; on other targets the example +//! compiles to an empty `main` so clippy/test on non-Linux don't fail +//! with "configured out". + +#[cfg(not(target_os = "linux"))] +fn main() {} + +#[cfg(target_os = "linux")] +fn main() { + linux::main(); +} + +#[cfg(target_os = "linux")] +mod linux { + use std::hint::black_box; + use std::thread::sleep; + use std::time::{Duration, Instant}; + + fn burn_cpu_for(duration: Duration, mut state: u64) -> u64 { + let deadline = Instant::now() + duration; + while Instant::now() < deadline { + // A tiny integer workload that is deterministic, dependency-chained, + // and opaque to the optimizer. This keeps one CPU core busy without + // changing the allocation profile this example is meant to exercise. + state = + state.wrapping_mul(0x9e37_79b9_7f4a_7c15).rotate_left(17) ^ 0xbf58_476d_1ce4_e5b9; + black_box(state); + } + state + } + + pub fn main() { + let stress = std::env::args().any(|arg| arg == "--stress"); + + println!( + "pid={}; pre-install. Attach a tracer on 'usdt:*:ddheap:*'. stress={stress}", + std::process::id() + ); + + // Make a few allocations before install to demonstrate the pre-patch + // baseline — no USDTs should fire for these. + { + let warmup: Vec = (0..16).map(|i| format!("warmup-{i}")).collect(); + println!("pre-install warmup: {} entries", warmup.len()); + } + + sleep(Duration::from_secs(2)); + + // Install GOT overrides. After this, malloc/free/calloc/realloc + // calls anywhere in the process (including those issued by the Rust + // System allocator backing `Vec`, `String`, etc.) route through + // `libdd-heap-sampler`. + let ok = libdd_heap_gotter::install_heap_overrides(); + println!("install_heap_overrides() -> {ok}"); + + let mut i: u64 = 0; + let mut cpu_state = 0x1234_5678_9abc_def0; + loop { + // ~1000 small allocations + one larger join: plenty of alloc + // pressure to cross the default 512 KiB sampling interval over + // a handful of iterations. + let parts: Vec = (0..1000) + .map(|j| format!("chunk-{i}-{j}-with-some-padding-to-make-it-meaningful")) + .collect(); + let joined = parts.join(", "); + println!("[{i}] joined {} bytes", joined.len()); + i = i.wrapping_add(1); + + if stress { + cpu_state = burn_cpu_for(Duration::from_secs(1), cpu_state ^ i); + } else { + sleep(Duration::from_secs(1)); + } + } + } +} // mod linux diff --git a/libdd-heap-gotter/src/elf.rs b/libdd-heap-gotter/src/elf.rs new file mode 100644 index 0000000000..b07020de53 --- /dev/null +++ b/libdd-heap-gotter/src/elf.rs @@ -0,0 +1,594 @@ +// Copyright 2025-Present Datadog, Inc. https://www.datadoghq.com/ +// SPDX-License-Identifier: Apache-2.0 + +//! GOT-table interposition primitives. +//! +//! Port of ddprof's `src/lib/elfutils.cc` + the parts of +//! `symbol_overrides.cc` that drive it. Walks each loaded ELF object via +//! `dl_iterate_phdr`, parses its `PT_DYNAMIC` for the symbol/string/hash +//! tables and the relocation arrays, and rewrites GOT entries whose +//! symbol name is in the override map. Records the previous values so +//! the overrides can be reverted. +//! +//! Scope: +//! * 64-bit ELF only (`Elf64_*`). Other targets are gated out at compile time via `#[cfg]` on the +//! parent module. +//! * GNU hash tables only - `DT_HASH` is skipped, matching ddprof (sometimes points into kernel +//! space on old glibcs). +//! * REL / RELA / JMPREL relocation arrays. + +use core::ffi::{c_char, c_int, c_void}; +use std::collections::HashMap; +use std::ffi::CStr; + +use libc::{ + dl_iterate_phdr, dl_phdr_info, mprotect, sysconf, Elf64_Phdr, Elf64_Rel, Elf64_Rela, Elf64_Sym, + PROT_READ, PROT_WRITE, PT_DYNAMIC, PT_LOAD, _SC_PAGESIZE, +}; + +// ELF dynamic-section tags and friends. The `libc` crate doesn't export +// these (they're processor-independent ELF spec constants), so we name +// them locally. Values come from ``. +#[allow(non_camel_case_types)] +#[repr(C)] +struct Elf64_Dyn { + d_tag: i64, + d_un: u64, // d_val / d_ptr union; we only ever read it as u64 +} +const DT_NULL: i64 = 0; +const DT_STRTAB: i64 = 5; +const DT_SYMTAB: i64 = 6; +const DT_RELA: i64 = 7; +const DT_RELASZ: i64 = 8; +const DT_STRSZ: i64 = 10; +const DT_REL: i64 = 17; +const DT_RELSZ: i64 = 18; +const DT_PLTREL: i64 = 20; +const DT_JMPREL: i64 = 23; +const DT_PLTRELSZ: i64 = 2; +const DT_GNU_HASH: i64 = 0x6fff_fef5; +const STN_UNDEF: u32 = 0; + +/// The subset of an ELF object's `PT_DYNAMIC` entries needed to find and rewrite GOT entries. +struct DynamicInfo { + strtab: *const c_char, + strtab_size: usize, + symtab: *const Elf64_Sym, + sym_count: u32, + rels: *const Elf64_Rel, + rels_count: usize, + relas: *const Elf64_Rela, + relas_count: usize, + jmprels: *const Elf64_Rela, + jmprels_count: usize, + gnu_hash: *const u32, + base_address: usize, +} + +impl DynamicInfo { + /// Read DT_* entries out of a PT_DYNAMIC array. ddprof's + /// `retrieve_dynamic_info` handles the glibc-vs-musl quirk where + /// glibc stores absolute addresses in DT entries while musl stores + /// load-relative offsets; we use the same `addr > base ? addr : base + addr` + /// heuristic. + unsafe fn from_phdr(info: &dl_phdr_info) -> Option { + let phdrs = std::slice::from_raw_parts(info.dlpi_phdr, info.dlpi_phnum as usize); + let dyn_phdr = phdrs.iter().find(|p| p.p_type == PT_DYNAMIC)?; + let dyn_begin = (info.dlpi_addr as usize + dyn_phdr.p_vaddr as usize) as *const Elf64_Dyn; + let base = info.dlpi_addr as usize; + let correct = |a: u64| -> usize { + let a = a as usize; + if a > base { + a + } else { + base + a + } + }; + + let mut strtab: *const c_char = std::ptr::null(); + let mut strtab_size: usize = 0; + let mut symtab: *const Elf64_Sym = std::ptr::null(); + let mut rels: *const Elf64_Rel = std::ptr::null(); + let mut rels_size: usize = 0; + let mut relas: *const Elf64_Rela = std::ptr::null(); + let mut relas_size: usize = 0; + let mut jmprels: *const Elf64_Rela = std::ptr::null(); + let mut jmprels_size: usize = 0; + let mut gnu_hash: *const u32 = std::ptr::null(); + let mut pltrel_type: i64 = 0; + + let mut it = dyn_begin; + loop { + let d = &*it; + if d.d_tag == DT_NULL { + break; + } + let v = d.d_un; + match d.d_tag { + DT_STRTAB => strtab = correct(v) as *const c_char, + DT_STRSZ => strtab_size = v as usize, + DT_SYMTAB => symtab = correct(v) as *const Elf64_Sym, + DT_GNU_HASH => gnu_hash = correct(v) as *const u32, + DT_REL => rels = correct(v) as *const Elf64_Rel, + DT_RELA => relas = correct(v) as *const Elf64_Rela, + DT_JMPREL => jmprels = correct(v) as *const Elf64_Rela, + DT_RELSZ => rels_size = v as usize, + DT_RELASZ => relas_size = v as usize, + DT_PLTRELSZ => jmprels_size = v as usize, + DT_PLTREL => pltrel_type = v as i64, + _ => {} + } + it = it.add(1); + } + + // JMPREL entries are RELA only if DT_PLTREL says so. + if pltrel_type != DT_RELA { + jmprels = std::ptr::null(); + jmprels_size = 0; + } + + if strtab.is_null() || symtab.is_null() || gnu_hash.is_null() { + return None; + } + + let sym_count = gnu_hash_symbol_count(gnu_hash); + + Some(Self { + strtab, + strtab_size, + symtab, + sym_count, + rels, + rels_count: rels_size / core::mem::size_of::(), + relas, + relas_count: relas_size / core::mem::size_of::(), + jmprels, + jmprels_count: jmprels_size / core::mem::size_of::(), + gnu_hash, + base_address: base, + }) + } + + unsafe fn sym_name(&self, idx: u32) -> Option<&CStr> { + if (idx as usize) >= self.sym_count as usize { + return None; + } + let sym = &*self.symtab.add(idx as usize); + let off = sym.st_name as usize; + if off >= self.strtab_size { + return None; + } + Some(CStr::from_ptr(self.strtab.add(off))) + } +} + +/// Compute the GNU symbol hash used by `DT_GNU_HASH` tables. +/// See . +fn gnu_hash(name: &[u8]) -> u32 { + let mut h: u32 = 5381; + for c in name { + h = h.wrapping_shl(5).wrapping_add(h).wrapping_add(*c as u32); + } + h +} + +unsafe fn gnu_hash_symbol_count(hashtab: *const u32) -> u32 { + let nbuckets = *hashtab; + let symbias = *hashtab.add(1); + let bloom_size = *hashtab.add(2); + // 4 header words + bloom (one Elf64_Addr per entry == 2 u32s) + let mut p = hashtab.add(4 + 2 * bloom_size as usize); + let buckets = std::slice::from_raw_parts(p, nbuckets as usize); + p = p.add(nbuckets as usize); + let chain_zero = p.offset(-(symbias as isize)); + + if nbuckets == 0 { + return 0; + } + let mut idx = *buckets.iter().max().unwrap(); + while *chain_zero.add(idx as usize) & 1 == 0 { + idx += 1; + } + idx + 1 +} + +unsafe fn gnu_hash_lookup(info: &DynamicInfo, name: &[u8]) -> Option { + let hashtab = info.gnu_hash; + let nbuckets = *hashtab; + let symbias = *hashtab.add(1); + let bloom_size = *hashtab.add(2); + let bloom_shift = *hashtab.add(3); + let bloom = hashtab.add(4) as *const u64; + let mut p = hashtab.add(4 + 2 * bloom_size as usize); + let buckets = p; + p = p.add(nbuckets as usize); + let chain_zero = p.offset(-(symbias as isize)); + + if nbuckets == 0 { + return None; + } + + let h = gnu_hash(name); + let word = *bloom.add(((h / 64) & (bloom_size - 1)) as usize); + let bit1 = h & 63; + let bit2 = (h >> bloom_shift) & 63; + if ((word >> bit1) & (word >> bit2) & 1) == 0 { + return None; + } + + let mut symidx = *buckets.add((h % nbuckets) as usize); + if symidx == STN_UNDEF { + return None; + } + + loop { + let chain_h = *chain_zero.add(symidx as usize); + if ((chain_h ^ h) >> 1) == 0 { + if let Some(sname) = info.sym_name(symidx) { + if sname.to_bytes() == name && check_sym(&*info.symtab.add(symidx as usize)) { + return Some(*info.symtab.add(symidx as usize)); + } + } + } + if chain_h & 1 != 0 { + break; + } + symidx += 1; + } + None +} + +/// Mirror of ddprof's `check`: defining symbol, function/object/notype. +fn check_sym(sym: &Elf64_Sym) -> bool { + const SHN_ABS: u16 = 0xfff1; + let stt = sym.st_info & 0xf; + if sym.st_value == 0 && sym.st_shndx != SHN_ABS { + return false; + } + // STT_NOTYPE(0), STT_OBJECT(1), STT_FUNC(2), STT_GNU_IFUNC(10) + matches!(stt, 0 | 1 | 2 | 10) +} + +/// Visit each loaded ELF object once. `is_exe` is true only on the +/// first callback (the main executable). The callback returns `true` to +/// stop iteration. +fn iterate_libraries bool>(mut cb: F) { + struct Ctx<'a> { + cb: &'a mut dyn FnMut(&dl_phdr_info, bool) -> bool, + is_first: bool, + } + let mut ctx = Ctx { + cb: &mut cb, + is_first: true, + }; + + unsafe extern "C" fn trampoline( + info: *mut dl_phdr_info, + _size: libc::size_t, + data: *mut c_void, + ) -> c_int { + let ctx = &mut *(data as *mut Ctx); + let is_exe = ctx.is_first; + ctx.is_first = false; + if (ctx.cb)(&*info, is_exe) { + 1 + } else { + 0 + } + } + + unsafe { + dl_iterate_phdr(Some(trampoline), &mut ctx as *mut _ as *mut c_void); + } +} + +/// Temporarily make the containing page writable and replace one GOT entry. +unsafe fn override_entry(addr: usize, new_value: usize) -> bool { + let page = sysconf(_SC_PAGESIZE) as usize; + let aligned = (addr & !(page - 1)) as *mut c_void; + if mprotect(aligned, page, PROT_READ | PROT_WRITE) != 0 { + return false; + } + std::ptr::write_unaligned(addr as *mut usize, new_value); + true +} + +/// Read one GOT entry without assuming pointer alignment. +unsafe fn read_entry(addr: usize) -> usize { + std::ptr::read_unaligned(addr as *const usize) +} + +#[derive(Clone, Copy)] +pub struct LookupResult { + pub address: usize, + #[allow(dead_code)] // exposed for diagnostics / future filtering + pub size: u64, +} + +/// Look up a symbol across loaded objects, returning the first +/// non-zero-sized definition whose address is not `not_this_symbol`. +/// Mirrors ddprof's `lookup_symbol(..., accept_null_sized_symbol=false)`. +pub fn lookup_symbol(name: &str, not_this_symbol: usize) -> Option { + let needle = name.as_bytes(); + let mut found: Option = None; + iterate_libraries(|info, _is_exe| unsafe { + let lib_name = if info.dlpi_name.is_null() { + "" + } else { + CStr::from_ptr(info.dlpi_name).to_str().unwrap_or("") + }; + if lib_name.contains("linux-vdso") || lib_name.contains("/ld-linux") { + return false; + } + let Some(dyn_info) = DynamicInfo::from_phdr(info) else { + return false; + }; + if let Some(sym) = gnu_hash_lookup(&dyn_info, needle) { + if sym.st_size > 0 { + let addr = sym.st_value as usize + dyn_info.base_address; + if addr != not_this_symbol { + found = Some(LookupResult { + address: addr, + size: sym.st_size, + }); + return true; // stop + } + } + } + false + }); + found +} + +/// Per-library revert info: GOT addr -> original value at that addr. +#[derive(Default)] +struct LibraryRevertInfo { + /// Kept for diagnostics / future logging; we identify libraries by + /// base address elsewhere. + #[allow(dead_code)] + library_name: String, + old_value_per_address: HashMap, + processed: bool, +} + +/// One registered override entry. +struct OverrideInfo { + /// Output slot the install path fills with the resolved real symbol + /// address (so hooks can call through it). `AtomicUsize` is what the + /// hook side sees; we just store its address here. + ref_slot: *mut usize, + /// Address of our hook function, written into matching GOT entries. + new_symbol: usize, + /// If a GOT entry's address equals this, leave it alone. Used to + /// avoid clobbering our own ref slot's relocation in this library + /// (otherwise applying our override would replace the resolved real + /// symbol with our hook, causing infinite recursion when the hook + /// calls back through `ref_slot`). + do_not_override_this_symbol: usize, +} + +unsafe impl Send for OverrideInfo {} +unsafe impl Sync for OverrideInfo {} + +/// Mirror of ddprof's `SymbolOverrides`. Holds the override table and +/// the per-library revert info needed to undo writes. +#[derive(Default)] +pub struct SymbolOverrides { + overrides: HashMap, + revert_info_per_library: HashMap, + last_seen_nb_libs: i32, +} + +impl SymbolOverrides { + pub fn new() -> Self { + Self { + overrides: HashMap::new(), + revert_info_per_library: HashMap::new(), + last_seen_nb_libs: -1, + } + } + + /// Register an override. `ref_slot` is filled in by `apply_overrides` + /// with the resolved address of the real symbol so the hook can call + /// through it. + pub fn register(&mut self, name: &str, hook: usize, ref_slot: *mut usize) { + self.overrides.insert( + name.to_string(), + OverrideInfo { + ref_slot, + new_symbol: hook, + // Filled in by apply_overrides: we set it to the address + // of our own `ref_slot` once we know it. For a static + // Rust we can pass 0 - see note in apply_overrides. + do_not_override_this_symbol: 0, + }, + ); + } + + /// Resolve real-symbol addresses, then walk every loaded library and + /// patch GOT entries. + pub fn apply_overrides(&mut self) { + // 1. Resolve each override's underlying real symbol so hooks can forward through it. + // Excluding our own hook function address avoids picking up a self-reference (when the + // gotter library itself exports the same name - it won't in our case, but cheap + // insurance). + let resolved: Vec<(String, usize)> = self + .overrides + .iter() + .filter_map(|(name, ov)| { + lookup_symbol(name, ov.new_symbol).map(|r| (name.clone(), r.address)) + }) + .collect(); + for (name, addr) in resolved { + if let Some(ov) = self.overrides.get_mut(&name) { + unsafe { *ov.ref_slot = addr }; + } + } + self.update_overrides(); + } + + /// Process any newly-loaded libraries (e.g. after `dlopen`). + /// No-op if the loaded-library count hasn't changed. + pub fn update_overrides(&mut self) { + // `dl_phdr_info::dlpi_adds` is incremented on every dlopen. + // ddprof uses it as a cheap "did anything change?" probe. + let mut nb_loaded: i32 = -1; + iterate_libraries(|info, _| { + nb_loaded = info.dlpi_adds as i32; + true + }); + if nb_loaded == self.last_seen_nb_libs { + return; + } + self.last_seen_nb_libs = nb_loaded; + + for v in self.revert_info_per_library.values_mut() { + v.processed = false; + } + + // SAFETY: closure runs synchronously inside dl_iterate_phdr. + let self_ptr = self as *mut Self as usize; + iterate_libraries(move |info, _is_exe| unsafe { + let this = &mut *(self_ptr as *mut Self); + let lib_name = if info.dlpi_name.is_null() { + String::new() + } else { + CStr::from_ptr(info.dlpi_name) + .to_string_lossy() + .into_owned() + }; + if lib_name.contains("linux-vdso") || lib_name.contains("/ld-linux") { + return false; + } + if let Some(dyn_info) = DynamicInfo::from_phdr(info) { + this.apply_to_library(&dyn_info, lib_name); + } + false + }); + + // Drop any tracked libraries that have been unloaded. + self.revert_info_per_library.retain(|_, v| v.processed); + } + + /// Restore every GOT entry we touched. + pub fn restore_overrides(&mut self) { + let info_per_lib = std::mem::take(&mut self.revert_info_per_library); + for (_base, revert) in info_per_lib { + unsafe { + for (addr, old) in revert.old_value_per_address { + override_entry(addr, old); + } + } + } + self.last_seen_nb_libs = -1; + } + + unsafe fn apply_to_library(&mut self, dyn_info: &DynamicInfo, library_name: String) { + let (entry_is_new, _) = match self.revert_info_per_library.entry(dyn_info.base_address) { + std::collections::hash_map::Entry::Vacant(e) => { + e.insert(LibraryRevertInfo { + library_name: library_name.clone(), + processed: true, + ..Default::default() + }); + (true, ()) + } + std::collections::hash_map::Entry::Occupied(mut e) => { + e.get_mut().processed = true; + (false, ()) + } + }; + if !entry_is_new { + return; + } + + // Hand-managed split borrow so we can pass &overrides + &mut revert. + let revert = self + .revert_info_per_library + .get_mut(&dyn_info.base_address) + .unwrap(); + + if !dyn_info.rels.is_null() { + let relocs = std::slice::from_raw_parts(dyn_info.rels, dyn_info.rels_count); + for reloc in relocs { + Self::process_relocation( + &self.overrides, + dyn_info, + elf64_r_sym(reloc.r_info) as u32, + reloc.r_offset as usize, + revert, + ); + } + } + for slice_ptr_and_len in [ + (dyn_info.relas, dyn_info.relas_count), + (dyn_info.jmprels, dyn_info.jmprels_count), + ] { + if slice_ptr_and_len.0.is_null() { + continue; + } + let relocs = std::slice::from_raw_parts(slice_ptr_and_len.0, slice_ptr_and_len.1); + for reloc in relocs { + Self::process_relocation( + &self.overrides, + dyn_info, + elf64_r_sym(reloc.r_info) as u32, + reloc.r_offset as usize, + revert, + ); + } + } + } + + unsafe fn process_relocation( + overrides: &HashMap, + dyn_info: &DynamicInfo, + sym_index: u32, + r_offset: usize, + revert: &mut LibraryRevertInfo, + ) { + // st_name -> string in strtab. Walk lazily: we look up the + // name in the override map; if it's not there, skip. + let sym = &*dyn_info.symtab.add(sym_index as usize); + let name_off = sym.st_name as usize; + if name_off == 0 || name_off >= dyn_info.strtab_size { + return; + } + let cstr = CStr::from_ptr(dyn_info.strtab.add(name_off)); + let Ok(name) = cstr.to_str() else { return }; + + let Some(ov) = overrides.get(name) else { + return; + }; + // `ref_slot==0` means we never resolved the real symbol, so the + // hook would call a NULL pointer. Skip. + let real = unsafe { *ov.ref_slot }; + if real == 0 { + return; + } + + let addr = r_offset + dyn_info.base_address; + if addr == ov.do_not_override_this_symbol { + return; + } + if revert.old_value_per_address.contains_key(&addr) { + return; + } + revert.old_value_per_address.insert(addr, read_entry(addr)); + override_entry(addr, ov.new_symbol); + } +} + +fn elf64_r_sym(info: u64) -> u64 { + info >> 32 +} + +// Phdr/PT_LOAD ranges are unused for now; kept for future "skip self +// library" logic. Silence unused warnings. +#[allow(dead_code)] +const _: () = { + let _ = PT_LOAD; +}; + +#[allow(dead_code)] +fn _phdr_marker(_: Elf64_Phdr) {} diff --git a/libdd-heap-gotter/src/hooks.rs b/libdd-heap-gotter/src/hooks.rs new file mode 100644 index 0000000000..68b0d267d9 --- /dev/null +++ b/libdd-heap-gotter/src/hooks.rs @@ -0,0 +1,311 @@ +// Copyright 2025-Present Datadog, Inc. https://www.datadoghq.com/ +// SPDX-License-Identifier: Apache-2.0 + +//! GOT hook functions and their per-symbol "real" function pointer slots. +//! +//! Each `gotter_*` function: +//! 1. Calls `dd_allocation_requested` (or skips, for free-side hooks). +//! 2. Forwards to the real symbol via its `ORIG_*` slot, which the install path fills in by +//! symbol-table lookup. +//! 3. Calls `dd_allocation_created` / `dd_allocation_freed` to fire the USDT and close the reentry +//! guard. +//! +//! Modeled on ddprof `src/lib/symbol_overrides.cc`, minus the C++ allocator +//! family (operator new/delete) and the mmap/munmap pair which aren't +//! supported by the sampler yet. + +use core::cell::Cell; +use core::ffi::{c_char, c_int, c_void}; +use std::sync::atomic::{AtomicUsize, Ordering}; + +use libdd_heap_sampler::{ + dd_alloc_req_t, dd_allocation_created, dd_allocation_freed, dd_allocation_requested, + dd_tl_state_get, dd_tl_state_init, +}; + +// Per-thread reentry guard for the gotter shims themselves. Distinct +// from the sampler's `dd_tl_state_t::reentry_guard` because that one +// lives inside a struct we have to *allocate* (via `calloc`) on first +// touch - and on a freshly-installed gotter that first `calloc` lands +// right back in `gotter_calloc`. So before we look at the sampler TLS +// at all, we set this flag; any reentry while it's set forwards +// straight through to the real allocator with no sampling. +// +// `const { Cell::new(false) }` keeps the TLS slot lazy-init-free - +// macOS and glibc both initialise it without an allocation. +std::thread_local! { + static IN_HOOK: Cell = const { Cell::new(false) }; +} + +struct GotterReentry(bool); + +impl GotterReentry { + fn enter() -> Self { + let was = IN_HOOK.with(|c| { + let prev = c.get(); + c.set(true); + prev + }); + GotterReentry(was) + } + fn reentered(&self) -> bool { + self.0 + } +} + +impl Drop for GotterReentry { + fn drop(&mut self) { + IN_HOOK.with(|c| c.set(self.0)); + } +} + +/// Resolved address of the real `malloc`; filled by `install_heap_overrides`. +/// The bare `usize` payload is a function pointer. +pub(crate) static ORIG_MALLOC: AtomicUsize = AtomicUsize::new(0); +/// Resolved address of the real `free`. +pub(crate) static ORIG_FREE: AtomicUsize = AtomicUsize::new(0); +/// Resolved address of the real `calloc`. +pub(crate) static ORIG_CALLOC: AtomicUsize = AtomicUsize::new(0); +/// Resolved address of the real `realloc`. +pub(crate) static ORIG_REALLOC: AtomicUsize = AtomicUsize::new(0); +/// Resolved address of the real `posix_memalign`. +pub(crate) static ORIG_POSIX_MEMALIGN: AtomicUsize = AtomicUsize::new(0); +/// Resolved address of the real `aligned_alloc`. +pub(crate) static ORIG_ALIGNED_ALLOC: AtomicUsize = AtomicUsize::new(0); +/// Resolved address of the real `dlopen`. +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); + +/// Ensure the sampler's per-thread state exists before recording an allocation. +#[inline] +fn ensure_tls() { + unsafe { + if dd_tl_state_get().is_null() { + // dd_tl_state_init calls calloc internally; the gotter + // reentry guard around the caller stops us from re-entering + // this path through our own gotter_calloc hook. + dd_tl_state_init(); + } + } +} + +/// Load a resolved function pointer from one of the `ORIG_*` slots. +#[inline] +unsafe fn load_fn(slot: &AtomicUsize) -> Option { + let v = slot.load(Ordering::Acquire); + if v == 0 { + None + } else { + // SAFETY: caller guarantees T is the right `extern "C" fn(...)` + // type and that `v` was written by `apply_overrides` with a + // function of that signature. + Some(core::mem::transmute_copy::(&v)) + } +} + +/// Signature of the real `malloc`. +type MallocFn = unsafe extern "C" fn(usize) -> *mut c_void; +/// Signature of the real `free`. +type FreeFn = unsafe extern "C" fn(*mut c_void); +/// Signature of the real `calloc`. +type CallocFn = unsafe extern "C" fn(usize, usize) -> *mut c_void; +/// Signature of the real `realloc`. +type ReallocFn = unsafe extern "C" fn(*mut c_void, usize) -> *mut c_void; +/// Signature of the real `posix_memalign`. +type PosixMemalignFn = unsafe extern "C" fn(*mut *mut c_void, usize, usize) -> c_int; +/// Signature of the real `aligned_alloc`. +type AlignedAllocFn = unsafe extern "C" fn(usize, usize) -> *mut c_void; +/// Signature of the real `dlopen`. +type DlopenFn = unsafe extern "C" fn(*const c_char, c_int) -> *mut c_void; +/// Signature of a `pthread_create` start routine. +type StartRoutine = unsafe extern "C" fn(*mut c_void) -> *mut c_void; +/// Signature of the real `pthread_create`. +type PthreadCreateFn = unsafe extern "C" fn( + *mut libc::pthread_t, + *const libc::pthread_attr_t, + StartRoutine, + *mut c_void, +) -> c_int; + +#[no_mangle] +pub unsafe extern "C" fn gotter_malloc(size: usize) -> *mut c_void { + let Some(real): Option = load_fn(&ORIG_MALLOC) else { + return std::ptr::null_mut(); + }; + let guard = GotterReentry::enter(); + if guard.reentered() { + return real(size); + } + ensure_tls(); + // Default alignment for malloc on glibc is 2*sizeof(void*) == 16. + let req = dd_allocation_requested(size, core::mem::align_of::() * 2); + let raw = real(req.size); + dd_allocation_created(raw, req) +} + +#[no_mangle] +pub unsafe extern "C" fn gotter_free(ptr: *mut c_void) { + let Some(real): Option = load_fn(&ORIG_FREE) else { + return; + }; + if ptr.is_null() { + return; + } + let guard = GotterReentry::enter(); + if guard.reentered() { + real(ptr); + return; + } + let freed = dd_allocation_freed(ptr, 0, 0); + real(freed.ptr); +} + +#[no_mangle] +pub unsafe extern "C" fn gotter_calloc(nmemb: usize, size: usize) -> *mut c_void { + let Some(real): Option = load_fn(&ORIG_CALLOC) else { + return std::ptr::null_mut(); + }; + let guard = GotterReentry::enter(); + if guard.reentered() { + return real(nmemb, size); + } + ensure_tls(); + let total = nmemb.saturating_mul(size); + let req = dd_allocation_requested(total, core::mem::align_of::() * 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 { + real(nmemb, size) + } else { + real(1, req.size) + }; + dd_allocation_created(raw, req) +} + +#[no_mangle] +pub unsafe extern "C" fn gotter_realloc(ptr: *mut c_void, size: usize) -> *mut c_void { + let Some(real): Option = load_fn(&ORIG_REALLOC) else { + return std::ptr::null_mut(); + }; + let guard = GotterReentry::enter(); + if guard.reentered() { + return real(ptr, size); + } + ensure_tls(); + // realloc decomposes into a free + an alloc from the sampler's + // perspective. If `ptr` was sampled we have to clear its flag so the + // underlying allocator sees the raw pointer; size==0 is a free. + let freed_ptr = if ptr.is_null() { + std::ptr::null_mut() + } else { + dd_allocation_freed(ptr, 0, 0).ptr + }; + + if size == 0 { + // realloc(p, 0) == free(p). Just forward. + return real(freed_ptr, 0); + } + + let req = dd_allocation_requested(size, core::mem::align_of::() * 2); + let raw = real(freed_ptr, req.size); + dd_allocation_created(raw, req) +} + +#[no_mangle] +pub unsafe extern "C" fn gotter_posix_memalign( + memptr: *mut *mut c_void, + alignment: usize, + size: usize, +) -> c_int { + let Some(real): Option = load_fn(&ORIG_POSIX_MEMALIGN) else { + return libc::ENOMEM; + }; + let guard = GotterReentry::enter(); + if guard.reentered() { + return real(memptr, alignment, size); + } + ensure_tls(); + let req = dd_allocation_requested(size, alignment); + let ret = real(memptr, alignment, req.size); + if ret == 0 && !memptr.is_null() { + let raw = *memptr; + *memptr = dd_allocation_created(raw, req); + } + ret +} + +#[no_mangle] +pub unsafe extern "C" fn gotter_aligned_alloc(alignment: usize, size: usize) -> *mut c_void { + let Some(real): Option = load_fn(&ORIG_ALIGNED_ALLOC) else { + return std::ptr::null_mut(); + }; + let guard = GotterReentry::enter(); + if guard.reentered() { + return real(alignment, size); + } + ensure_tls(); + let req = dd_allocation_requested(size, alignment); + let raw = real(alignment, req.size); + dd_allocation_created(raw, req) +} + +/// Forward `dlopen` and then patch any GOT entries introduced by the newly-loaded library. +#[no_mangle] +pub unsafe extern "C" fn gotter_dlopen(filename: *const c_char, flags: c_int) -> *mut c_void { + let Some(real): Option = load_fn(&ORIG_DLOPEN) else { + // Hooks not yet wired up; calling real() would NPE - punt to libc. + return libc::dlopen(filename, flags); + }; + let handle = real(filename, flags); + // New library may have introduced new GOT entries that need patching. + crate::update_heap_overrides(); + handle +} + +/// Args we package up so the wrapped start routine sees its original +/// arg through our trampoline. Matches ddprof's `Args = tuple<...>`. +struct PthreadCreateArgs { + start: StartRoutine, + arg: *mut c_void, +} + +unsafe extern "C" fn pthread_start_trampoline(arg: *mut c_void) -> *mut c_void { + let boxed: Box = Box::from_raw(arg as *mut PthreadCreateArgs); + // Materialise per-thread sampler state up front so the first + // tracked alloc on this thread doesn't have to. + dd_tl_state_init(); + (boxed.start)(boxed.arg) +} + +#[no_mangle] +pub unsafe extern "C" fn gotter_pthread_create( + thread: *mut libc::pthread_t, + attr: *const libc::pthread_attr_t, + start: StartRoutine, + arg: *mut c_void, +) -> c_int { + let Some(real): Option = load_fn(&ORIG_PTHREAD_CREATE) else { + return libc::EAGAIN; + }; + let boxed = Box::new(PthreadCreateArgs { start, arg }); + let raw = Box::into_raw(boxed); + let rc = real(thread, attr, pthread_start_trampoline, raw as *mut c_void); + if rc != 0 { + // Reclaim the box; trampoline won't run. + drop(Box::from_raw(raw)); + } + rc +} + +// Touch the sampler-side reentry guard helpers indirectly; silences +// "unused import" warnings about `dd_alloc_req_t` once cfg-gated paths +// extend hooks later. +#[allow(dead_code)] +fn _types_anchor() -> dd_alloc_req_t { + dd_alloc_req_t { size: 0, weight: 0 } +} diff --git a/libdd-heap-gotter/src/lib.rs b/libdd-heap-gotter/src/lib.rs new file mode 100644 index 0000000000..e330107f32 --- /dev/null +++ b/libdd-heap-gotter/src/lib.rs @@ -0,0 +1,247 @@ +// Copyright 2025-Present Datadog, Inc. https://www.datadoghq.com/ +// SPDX-License-Identifier: Apache-2.0 + +//! GOT-table interposition for heap profiling. +//! +//! This crate installs hook functions over a running process's dynamic +//! symbol relocations (the GOT / PLT-resolved entries) so that calls +//! such as `malloc` and `free` are routed through +//! [`libdd-heap-sampler`] without recompiling or relinking the target +//! application. The approach mirrors ddprof's `src/lib/symbol_overrides.cc` +//! + `src/lib/elfutils.cc`. +//! +//! The public API is available on every platform so downstream code +//! never has to `#[cfg]`-guard its callers. The GOT-patching machinery +//! itself only exists on Linux (where `dl_iterate_phdr` + ELF relocs +//! are well-defined); on every other target the entry points compile +//! to no-ops and `heap_overrides_are_installed()` always returns +//! `false`. +//! +//! # Quickstart +//! +//! ```no_run +//! libdd_heap_gotter::install_heap_overrides(); +//! // ... application runs; malloc/free/calloc/realloc/etc. flow through +//! // libdd-heap-sampler and emit ddheap:alloc / ddheap:free USDTs ... +//! libdd_heap_gotter::restore_heap_overrides(); +//! ``` +//! +//! # Status +//! +//! Initial port. Covers: `malloc`, `free`, `calloc`, `realloc`, +//! `posix_memalign`, `aligned_alloc`, plus `dlopen` (to re-scan on new +//! library load) and `pthread_create` (to materialise sampler TLS up +//! front on new threads). +//! +//! Not yet covered: +//! * `operator new` / `operator delete` family +//! * `mmap` / `munmap` (sampler-side API doesn't exist yet) +//! * jemalloc-specific `mallocx`/`dallocx`/etc. +//! * `pthread_atfork` child handler to reset state cleanly across `fork()`. + +#[cfg(target_os = "linux")] +mod elf; +#[cfg(target_os = "linux")] +mod hooks; + +#[cfg(target_os = "linux")] +use std::sync::Mutex; + +#[cfg(target_os = "linux")] +use elf::SymbolOverrides; + +/// Holds the SymbolOverrides registry across calls to `install` / `update` +/// / `restore`. ddprof keeps the equivalent state in +/// `g_symbol_overrides` guarded by `g_mutex` +#[cfg(target_os = "linux")] +static GLOBAL_OVERRIDES: Mutex> = Mutex::new(None); + +/// Install GOT overrides for the supported allocator and helper symbols. +/// Safe to call more than once: the registry is rebuilt and re-applied, +/// which also picks up any libraries loaded since the last call. +/// +/// Returns `true` if at least one symbol was successfully overridden; +/// `false` if nothing could be resolved. This might happen if the +/// target process has already been statically linked against a custom +/// allocator that doesn't appear in the dynamic symbol table. +/// +/// On non-Linux targets this is a no-op that always returns `false` — +/// the GOT-patching path it would otherwise execute has no portable +/// equivalent outside ELF + `dl_iterate_phdr`. +#[cfg(target_os = "linux")] +pub fn install_heap_overrides() -> bool { + let mut guard = GLOBAL_OVERRIDES.lock().expect("gotter mutex poisoned"); + if guard.is_none() { + let mut so = SymbolOverrides::new(); + register_all(&mut so); + *guard = Some(so); + } + let so = guard.as_mut().unwrap(); + so.apply_overrides(); + // Heuristic: at least one ORIG slot resolved. + any_orig_resolved() +} + +/// See the Linux variant above. +#[cfg(not(target_os = "linux"))] +pub fn install_heap_overrides() -> bool { + false +} + +/// Re-scan loaded libraries and patch any newly-introduced GOT entries. +/// Called automatically from the `dlopen` hook; user code typically +/// doesn't need to call this directly. No-op on non-Linux targets. +#[cfg(target_os = "linux")] +pub fn update_heap_overrides() { + // `try_lock` so a dlopen happening on the same thread that owns the + // install lock doesn't deadlock - that thread will finish its + // outer apply_overrides, which already walks every library. + if let Ok(mut guard) = GLOBAL_OVERRIDES.try_lock() { + if let Some(so) = guard.as_mut() { + so.update_overrides(); + } + } +} + +/// See the Linux variant above. +#[cfg(not(target_os = "linux"))] +pub fn update_heap_overrides() {} + +/// Revert every GOT entry we patched. After this call, the process is +/// once again calling the real allocator symbols directly. No-op on +/// non-Linux targets. +#[cfg(target_os = "linux")] +pub fn restore_heap_overrides() { + let mut guard = GLOBAL_OVERRIDES.lock().expect("gotter mutex poisoned"); + if let Some(so) = guard.as_mut() { + so.restore_overrides(); + } + *guard = None; +} + +/// See the Linux variant above. +#[cfg(not(target_os = "linux"))] +pub fn restore_heap_overrides() {} + +/// Return whether heap GOT overrides are currently installed. Always +/// returns `false` on non-Linux targets, since `install_heap_overrides` +/// is a no-op there. +#[cfg(target_os = "linux")] +pub fn heap_overrides_are_installed() -> bool { + GLOBAL_OVERRIDES + .lock() + .map(|guard| guard.is_some()) + .unwrap_or(false) +} + +/// See the Linux variant above. +#[cfg(not(target_os = "linux"))] +pub fn heap_overrides_are_installed() -> bool { + false +} + +/// Register GOT overrides for every symbol this crate currently hooks. +#[cfg(target_os = "linux")] +fn register_all(so: &mut SymbolOverrides) { + use hooks::*; + use std::sync::atomic::AtomicUsize; + + // Register one entry per supported symbol. The `ref_slot` raw + // pointer is to a `'static AtomicUsize`, so it's valid forever. + // AtomicUsize is repr(transparent) over UnsafeCell; we + // intentionally bypass its API for the install-time write because + // we hand the raw `*mut usize` to the ELF GOT scanner. Hooks then + // read it back via `Atomic::load(Acquire)`. + fn reg(so: &mut SymbolOverrides, name: &str, hook_addr: usize, slot: &'static AtomicUsize) { + let slot_ptr = slot as *const AtomicUsize as *mut usize; + so.register(name, hook_addr, slot_ptr); + } + + 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, + "realloc", + gotter_realloc as *const () as usize, + &ORIG_REALLOC, + ); + reg( + so, + "posix_memalign", + gotter_posix_memalign as *const () as usize, + &ORIG_POSIX_MEMALIGN, + ); + reg( + so, + "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, + "pthread_create", + gotter_pthread_create as *const () as usize, + &ORIG_PTHREAD_CREATE, + ); +} + +#[cfg(target_os = "linux")] +fn any_orig_resolved() -> bool { + use hooks::*; + use std::sync::atomic::Ordering; + [ + &ORIG_MALLOC, + &ORIG_FREE, + &ORIG_CALLOC, + &ORIG_REALLOC, + &ORIG_POSIX_MEMALIGN, + &ORIG_ALIGNED_ALLOC, + &ORIG_DLOPEN, + &ORIG_PTHREAD_CREATE, + ] + .iter() + .any(|s| s.load(Ordering::Relaxed) != 0) +} + +// Tests call into the ELF symbol-lookup path (dl_iterate_phdr + +// dynsym parsing of loaded libraries) which miri can't execute, so +// skip the whole module under miri. +#[cfg(all(test, target_os = "linux", not(miri)))] +mod tests { + use super::*; + + /// Smoke test that doesn't actually install (avoids messing with + /// the test binary's allocator) but exercises the symbol-lookup + /// path. `malloc` is always present in a Linux process. + #[test] + fn can_lookup_malloc() { + let r = elf::lookup_symbol("malloc", 0); + assert!(r.is_some(), "expected to find malloc in loaded libraries"); + let r = r.unwrap(); + assert!(r.address != 0); + assert!(r.size > 0); + } + + #[test] + fn unknown_symbol_lookup_returns_none() { + let r = elf::lookup_symbol("definitely_not_a_real_libc_symbol_xyzzy", 0); + assert!(r.is_none()); + } +} diff --git a/libdd-heap-gotter/tests/install.rs b/libdd-heap-gotter/tests/install.rs new file mode 100644 index 0000000000..4cf496aacb --- /dev/null +++ b/libdd-heap-gotter/tests/install.rs @@ -0,0 +1,54 @@ +// Copyright 2025-Present Datadog, Inc. https://www.datadoghq.com/ +// SPDX-License-Identifier: Apache-2.0 + +//! End-to-end smoke test: install the GOT overrides into the live test +//! process, do a libc `malloc`/`free`, and check that the hook ran. +//! +//! Single test only — installing the overrides mutates global process +//! state, so we keep this isolated from the unit-test binary. + +// Integration test invokes the GOT-patching machinery for real +// (dl_iterate_phdr + dlsym + mprotect), which miri can't execute. +#![cfg(all(target_os = "linux", not(miri)))] + +use std::ffi::c_void; + +// We don't have a clean way to instrument the gotter's own hooks from +// outside the crate, so this test goes one level lower: it confirms +// that after install, the heap is still functional and that no +// recursive crash has occurred when malloc/free go through the patched +// GOT. +#[test] +fn install_changes_malloc_dispatch() { + extern "C" { + fn malloc(size: usize) -> *mut c_void; + } + + let installed = libdd_heap_gotter::install_heap_overrides(); + assert!( + installed, + "expected install_heap_overrides to find at least one symbol" + ); + + // Touch the heap to exercise the patched GOT. We can't assert that + // the address differs from the original libc malloc address from + // Rust directly — the `malloc` extern fn item resolves via the same + // GOT we just patched, so reading its address here either gives us + // the PLT entry (still the same) or the post-resolution function + // pointer (now our hook). What we *can* assert is that the heap is + // still functional and that no recursive crash has occurred. + unsafe { + let p = malloc(64); + assert!(!p.is_null(), "malloc returned NULL post-install"); + libc::free(p); + } + + libdd_heap_gotter::restore_heap_overrides(); + + // A second alloc post-restore should also be fine. + unsafe { + let p = malloc(64); + assert!(!p.is_null(), "malloc returned NULL post-restore"); + libc::free(p); + } +} diff --git a/libdd-heap-sampler/Cargo.toml b/libdd-heap-sampler/Cargo.toml new file mode 100644 index 0000000000..7178568517 --- /dev/null +++ b/libdd-heap-sampler/Cargo.toml @@ -0,0 +1,28 @@ +# Copyright 2025-Present Datadog, Inc. https://www.datadoghq.com/ +# SPDX-License-Identifier: Apache-2.0 + +[package] +name = "libdd-heap-sampler" +version = "0.1.0" +description = "Sampling primitives for USDT-based heap profiling." +homepage = "https://github.com/DataDog/libdatadog/tree/main/libdd-heap-sampler" +repository = "https://github.com/DataDog/libdatadog/tree/main/libdd-heap-sampler" +edition.workspace = true +rust-version.workspace = true +license.workspace = true +publish = false + +[lib] +crate-type = ["lib", "staticlib"] +bench = false + +[build-dependencies] +cc = "1.1.31" +# `bindgen` is only *invoked* when the env var LIBDD_HEAP_SAMPLER_REGEN +# is set (see build.rs); it is otherwise pulled in as a build-dep so +# that regen works without a cargo feature. Cargo features would be +# problematic here because CI jobs frequently use `--all-features`, +# which would silently trigger a bindgen regen on hosts that lack +# libclang. An env var side-steps +# that (see README's "Regenerating bindings" section). +bindgen = "0.71" diff --git a/libdd-heap-sampler/README.md b/libdd-heap-sampler/README.md new file mode 100644 index 0000000000..78c82e99fe --- /dev/null +++ b/libdd-heap-sampler/README.md @@ -0,0 +1,131 @@ +# Heap Profiling + +This library forms the foundation of Datadog's application-side heap profiling support. It provides sampling functions that can be used to wrap each of the primary allocation and free functions of an arbitrary allocator. + +For allocations that are sampled as well as the corresponding frees of these allocations, appropriate USDTs are +emitted such that an external process such as the [eBPF full host profiler](https://github.com/open-telemetry/opentelemetry-ebpf-profiler) can collect the samples as +well as the stack trace at the time they are emitted to ultimately emit as a heap profiling event stream. + +## Use Cases + +This profiling infrastructure will initially support these two use cases: + +**Rust compile-time instrumentation** +A crate exposing a `GlobalAlloc` implementation will allow Rust users to, at compile time, opt into heap profiling. We +anticipate this will be shipped as a feature of `dd-trace-rs`. This addresses a pain point both internally with the +increase in Rust adoption in Datadog services, and would adress the same pain point within the broader Rust community. + +**Python Runtime Instrumentation for _native_ library allocation sampling** +Today our python profiler cannot sample allocations occuring behind the FFI; we will extend `ddtrace-py` to load our dynamic +runtime patching mechanism such that, as native libraries are loaded, we intercept their allocators. As many popular python +libraries function largely as API glue around native libraries this will help close the allocation observability gap. + +**Example Apps** +* [libdd-heap-allocator/examples/usdt_demo.rs](../libdd-heap-allocator/examples/usdt_demo.rs) - spins allocating/freeing memory hooked by a `GlobalAllocator` +* [libdd-heap-gotter-ffi/examples/cdylib_demo.rs](../libdd-heap-gotter-ffi/examples/cdylib_demo.rs) - _dynamically loads_ the gotter library, then spins allocating memory hooked by GOT table rewriting + +## Components + +```mermaid +flowchart TB + sampler["libdd-heap-sampler
Sampling decisions and USDT emission"] + allocator["libdd-heap-allocator
Rust GlobalAlloc wrapper"] + hooks["libdd-heap-hooks
Native allocator hooks (jemalloc, etc.)"] + gotter["libdd-heap-gotter
Dynamically inject samplers at runtime"] + + sampler --- allocator + sampler --- hooks + sampler --- gotter +``` + +### Samplers - `libdd-heap-sampler` (you are here) +These are the foundational functions themselves containing the sampling logic and USDTs, and are intended to be used +within higher order constructs that bind them back to concrete allocator callsites. +They are responsible for deciding whether or not to sample, and storing the information required to decide later on, at `free` time, if the given allocation _was_ sampled. We will cover: + +**USDTs** + +The actual USDTs emitted are: + +* `ddheap:alloc(void *user, uint64_t size, uint64_t weight)` — fired on sampled allocations; `user` is the user-visible pointer, `size` in bytes, `weight` is the unbiased size estimator (`nsamples * interval`) +* `ddheap:free(void *ptr)` — fired when a previously-sampled allocation is freed +* `ddheap:mmap` - TODO +* `ddheap:munmap` - TODO + +**Allocations** + +By splitting into `requested` and `created`, these are designed to be generic across different allocation functions +(e.g. `malloc`, `operator new`, `aligned_alloc`, etc.). The job of binding these back to concrete callsites in a +process is left to the other components - e.g. `libddd-heap-gotter`, `libdd-heap-allocator`, etc. + +The allocation-side pair is declared `static inline __attribute__((always_inline))` so the non-sampled fast path inlines +into the wrapper with no function-call overhead. + +Note that the functions on the _allocation_ side will return an _updated_ allocation size. This will generally be the +same as the requested allocation size, but may not always be as the sampling mechanism may choose to increase the +allocation size in order to ease the process of tracking sampling decisions. The caller should pass this returned +value through verbatim to the allocator it is wrapping. + +* _`allocation requested`_ - called _before_ `malloc`, `operator new`, etc. Returns the size to actually allocate plus the sampling decision for this allocation. +* _`allocation created`_ - called _after_ the allocator returns; on sampled allocations applies the flag and emits the USDT. +* _`allocation freed`_ - used by `free`, `operator delete`, etc. + +**Mappings** (TODO!) +* _mapping created_ - used by `mmap` +* _mapping freed_ - used by `munmap` + +### [Rust Allocator `libdd-heap-allocator`](../libdd-heap-allocator) +An implementation of a rust allocator using `libdd-heap-sampler` and wrapping an arbitrary allocator. + +### [Native Allocator Hooks `libdd-heap-hooks`](../libdd-heap-hooks) +These implement the native profiling hooks for the various allocators we support, emitting the same USDTs in the sampling path as `libdd-heap-sampler` does. We will implement this for `jemalloc` first. + +### [GOTter `libdd-heap-gotter`](../libdd-heap-gotter) +GOTter implements our GOT-patching mechanism to wrap (dynamically!) linked allocators in a running process. + +## Regenerating bindings + +The Rust FFI declarations in `src/generated/bindings.rs` and the +static-inline wrapper `src/generated/dd_heap_sampler_static_wrappers.c` +are produced by `bindgen` from the public headers under +`include/datadog/heap/`. **They are checked in so that day-to-day builds +do not depend on `libclang`** — this is the same pattern used by +`rusqlite`, `zstd-sys`, and other bindgen consumers, and it keeps our +internal build images (CentOS, Alpine) working without dragging LLVM in +as a runtime build dep. + +The C implementation is still architecture-specific where needed: +`sample_flag.h` uses x86_64 header-prefixed magic pointers and aarch64 +TBI pointer tagging. The generated Rust ABI intentionally allowlists only +the common public surface, so those internal arch-specific helpers and +constants are not emitted into `bindings.rs` and a single checked-in +binding set works for both supported Linux architectures. + +### Requirements + +- **`libclang`** (`libclang-dev` on Debian/Ubuntu, `libclang-devel` on + RPM distros, `brew install llvm` on macOS). This is the only hard + requirement. + +You do **not** need `bindgen-cli`, a cross libc, a cross Rust toolchain, +or a cross linker — `build.rs` invokes the `bindgen` crate directly and +only emits the arch-independent Rust-facing ABI. + +### Regenerating + +```bash +LIBDD_HEAP_SAMPLER_REGEN=1 cargo build -p libdd-heap-sampler +``` + +This refreshes `src/generated/bindings.rs` and +`src/generated/dd_heap_sampler_static_wrappers.c`. Commit the delta. CI's +`verify-heap-sampler-bindings` workflow runs the same command on every PR +and fails if the checked-in files are stale. + +> **Why an env var, not a `cargo` feature?** Several of libdatadog's +> CI jobs (clippy, coverage, workspace tests) run with +> `--all-features`. Cargo has no way to hide a feature from that flag, +> so a `regen-bindings` feature would be silently enabled everywhere +> and drag `libclang` into every build. An env var is +> invisible to `--all-features`, so `build.rs` only invokes bindgen +> when a human (or the verify-bindings workflow) explicitly asks. diff --git a/libdd-heap-sampler/build.rs b/libdd-heap-sampler/build.rs new file mode 100644 index 0000000000..43ad82d406 --- /dev/null +++ b/libdd-heap-sampler/build.rs @@ -0,0 +1,201 @@ +// Copyright 2025-Present Datadog, Inc. https://www.datadoghq.com/ +// SPDX-License-Identifier: Apache-2.0 + +//! Build script for `libdd-heap-sampler`. +//! +//! On Linux, compiles the C sampler primitives under `src/*.c` and either +//! stages the checked-in bindgen artifacts from `src/generated/` into +//! `OUT_DIR` (default) or regenerates them from the public headers under +//! `include/datadog/heap/` when `LIBDD_HEAP_SAMPLER_REGEN` is set. +//! On every other target this is a no-op and the crate compiles to an +//! empty rlib. + +fn main() { + // USDT/SystemTap (sys/sdt.h) is Linux-only, so the crate compiles to an + // empty rlib on every other target. Build scripts are compiled for the + // host, so use Cargo's target cfg env var rather than #[cfg(target_os)]. + if std::env::var("CARGO_CFG_TARGET_OS").as_deref() == Ok("linux") { + linux::build(); + } +} + +mod linux { + use std::env; + use std::path::{Path, PathBuf}; + + const SOURCES: &[&str] = &[ + "src/allocation_requested.c", + "src/allocation_created.c", + "src/allocation_freed.c", + "src/probes.c", + "src/sample_flag.c", + "src/tl_state.c", + ]; + + // Checked-in bindgen outputs. Regenerated on demand by setting the + // `LIBDD_HEAP_SAMPLER_REGEN` env var (see below); the default build + // path only *reads* these files, so libclang is NOT a build-time + // dependency for downstream consumers. + // + // The C implementation is still architecture-specific where needed + // (`sample_flag.h` selects x86_64 header-magic vs aarch64 TBI at C + // compile time), but the Rust-facing ABI intentionally excludes + // those internal helpers/constants and is arch-independent. + const GENERATED_BINDINGS: &str = "src/generated/bindings.rs"; + const GENERATED_WRAPPER: &str = "src/generated/dd_heap_sampler_static_wrappers.c"; + + pub fn build() { + let out_dir = PathBuf::from(env::var("OUT_DIR").expect("OUT_DIR not set")); + let out_bindings = out_dir.join("bindings.rs"); + let out_wrapper = out_dir.join("dd_heap_sampler_static_wrappers.c"); + + // Regeneration is opt-in via env var, NOT a cargo feature. + // Cargo has no way to hide a feature from `--all-features`, and + // most CI jobs in this workspace run with `--all-features`, + // which would otherwise silently invoke bindgen on runners that + // lack libclang. An env var is invisible to + // `--all-features`, so this stays a deliberate opt-in for the + // regen workflow only. + println!("cargo:rerun-if-env-changed={REGEN_ENV_VAR}"); + if env::var_os(REGEN_ENV_VAR).is_some() { + regen::run(); + } + // Whether we just regenerated or not, we still need the checked-in + // files staged into OUT_DIR so that lib.rs's + // `include!(concat!(env!("OUT_DIR"), "/bindings.rs"))` finds + // them and cc compiles the wrapper. + use_checked_in(&out_bindings, &out_wrapper); + + compile_c(&out_wrapper); + } + + /// Setting this env var to any value triggers a bindgen regen under + /// `src/generated/` on the next build.rs invocation. Requires libclang; + /// see `libdd-heap-sampler/README.md`. + const REGEN_ENV_VAR: &str = "LIBDD_HEAP_SAMPLER_REGEN"; + + fn use_checked_in(out_bindings: &Path, out_wrapper: &Path) { + for src in [GENERATED_BINDINGS, GENERATED_WRAPPER] { + println!("cargo:rerun-if-changed={src}"); + assert!( + Path::new(src).is_file(), + "checked-in bindgen output `{src}` missing. Run \ + `{REGEN_ENV_VAR}=1 cargo build -p libdd-heap-sampler` \ + to regenerate (requires libclang; see README)." + ); + } + std::fs::copy(GENERATED_BINDINGS, out_bindings) + .expect("failed to stage checked-in bindings.rs into OUT_DIR"); + std::fs::copy(GENERATED_WRAPPER, out_wrapper) + .expect("failed to stage checked-in wrapper .c into OUT_DIR"); + } + + fn compile_c(wrap_path: &Path) { + let mut build = cc::Build::new(); + build + .files(SOURCES) + .file(wrap_path) + .include(".") + .include("include") + // See regen::run's matching `-Ivendor` and vendor/README.md. + .include("vendor") + .warnings(true) + .flag_if_supported("-Wextra") + .flag_if_supported("-fcf-protection=none") + // Use TLSDESC for dd_tl_state_storage. For static builds the + // linker relaxes this to local-exec automatically. For dynamic + // 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"); + build.compile("dd_heap_sampler"); + + for f in SOURCES { + println!("cargo:rerun-if-changed={f}"); + } + } + + mod regen { + use super::{GENERATED_BINDINGS, GENERATED_WRAPPER}; + use std::path::Path; + + const HEADERS: &[&str] = &[ + "include/datadog/heap/allocation_requested.h", + "include/datadog/heap/allocation_created.h", + "include/datadog/heap/allocation_freed.h", + "include/datadog/heap/probes.h", + "include/datadog/heap/sample_flag.h", + "include/datadog/heap/tl_state.h", + ]; + + /// Regenerate the checked-in bindings. The allowlist names only the + /// Rust-facing ABI, so bindgen ignores architecture-specific internal + /// constants/helpers from sample_flag.h and the output is reusable on + /// both supported Linux architectures. + pub fn run() { + regen_bindings(); + } + + fn regen_bindings() { + let committed_bindings = GENERATED_BINDINGS; + let committed_wrapper = GENERATED_WRAPPER; + if let Some(parent) = Path::new(committed_bindings).parent() { + std::fs::create_dir_all(parent).expect("failed to create src/generated/"); + } + + let mut builder = bindgen::Builder::default() + .clang_arg("-Iinclude") + // `-Ivendor` makes resolve to the vendored + // libbpf/usdt single-header (BSD-2-Clause, see + // vendor/README.md). We bundle it unconditionally rather + // than depending on the system's `systemtap-sdt-dev` + // package because libbpf/usdt is genuinely standalone + // (no sdt-config.h companion file) and so works + // identically across glibc and musl distros. + .clang_arg("-Ivendor") + .allowlist_function("dd_.*") + .allowlist_type("dd_.*") + .allowlist_var("DD_SAMPLING_INTERVAL_DEFAULT") + // Emit FFI-linkable shims for `static inline` helpers so + // we have a single source of truth for the fast path + // (the C header), reached via one function call from + // Rust. + .wrap_static_fns(true) + .wrap_static_fns_path(committed_wrapper) + .parse_callbacks(Box::new(bindgen::CargoCallbacks::new())); + + for h in HEADERS { + println!("cargo:rerun-if-changed={h}"); + builder = builder.header(*h); + } + + builder + .generate() + .unwrap_or_else(|e| panic!("bindgen failed to generate bindings: {e}")) + .write_to_file(committed_bindings) + .unwrap_or_else(|e| panic!("failed to write {committed_bindings}: {e}")); + + // bindgen doesn't emit any copyright header on its output; + // libdatadog's CI runs `licensecheck` against every `.rs`/`.c` + // file and fails when the Apache-2.0 header is missing. + // Prepend it in the appropriate comment style for each file + // so the checked-in artifacts satisfy the check. + prepend_license_header(Path::new(&committed_bindings), "//"); + prepend_license_header(Path::new(&committed_wrapper), "//"); + } + + fn prepend_license_header(path: &Path, comment: &str) { + let header = format!( + "{c} Copyright 2025-Present Datadog, Inc. https://www.datadoghq.com/\n\ + {c} SPDX-License-Identifier: Apache-2.0\n\ + {c} @generated by libdd-heap-sampler/build.rs via bindgen; do not edit by hand.\n\ + {c} Regenerate with: LIBDD_HEAP_SAMPLER_REGEN=1 cargo build -p libdd-heap-sampler\n\n", + c = comment, + ); + let body = std::fs::read_to_string(path).unwrap_or_else(|e| { + panic!("reading {} to prepend license header: {e}", path.display()) + }); + std::fs::write(path, format!("{header}{body}")) + .unwrap_or_else(|e| panic!("writing {} with license header: {e}", path.display())); + } + } +} diff --git a/libdd-heap-sampler/include/datadog/heap/allocation_created.h b/libdd-heap-sampler/include/datadog/heap/allocation_created.h new file mode 100644 index 0000000000..1d7d86135f --- /dev/null +++ b/libdd-heap-sampler/include/datadog/heap/allocation_created.h @@ -0,0 +1,50 @@ +/** + * @file allocation_created.h + * + * Post-allocation hook for the ddheap sampler. + * + * Call dd_allocation_created() immediately after the underlying allocator + * returns. On the common (unsampled) fast path this is a single branch on + * req.weight and an immediate return of the raw pointer. On the sampled slow + * path it applies the architecture-specific sample flag to the raw pointer, + * emits the ddheap:alloc USDT, and closes the reentry guard that was opened + * by the paired dd_allocation_requested() call. + * + * Always call this even when the allocator returns NULL: the reentry guard + * must be closed regardless of whether the allocation succeeded. + */ +#ifndef DD_SAMPLERS_ALLOCATION_CREATED_H +#define DD_SAMPLERS_ALLOCATION_CREATED_H + +#include /* dd_alloc_req_t */ + +/* + * Slow path. As with allocation_requested, this is fired only when + * we hit a sampled allocation, and is intentionally placed separately + * from dd_allocation_created so that we don't bloat the instruction + * cache for the fast path. + */ +void *dd_allocation_created_slow(void *raw, dd_alloc_req_t req); + +/* + * Post-allocation hook. Pair with dd_allocation_requested. + * + * Wrapper usage: + * dd_alloc_req_t req = dd_allocation_requested(size, alignment); + * void *raw = real_alloc(..., req.size); + * return dd_allocation_created(raw, req); + * + * Fast path (not sampled): returns raw unchanged. + * Slow path (sampled): applies the architecture-specific flag, emits + * the ddheap:alloc USDT, closes the reentry guard opened by the paired + * requested() call, returns the user-visible pointer. + * + * Safe when raw == NULL (allocator failed): no USDT, guard still closed. + */ +static inline __attribute__((always_inline)) +void *dd_allocation_created(void *raw, dd_alloc_req_t req) { + if (__builtin_expect(req.weight == 0, 1)) return raw; + return dd_allocation_created_slow(raw, req); +} + +#endif diff --git a/libdd-heap-sampler/include/datadog/heap/allocation_freed.h b/libdd-heap-sampler/include/datadog/heap/allocation_freed.h new file mode 100644 index 0000000000..0e63bda996 --- /dev/null +++ b/libdd-heap-sampler/include/datadog/heap/allocation_freed.h @@ -0,0 +1,72 @@ +/** + * @file allocation_freed.h + * + * Pre-free hook for the ddheap sampler. + * + * Call dd_allocation_freed() before forwarding to the underlying deallocator. + * On the fast path it checks whether the pointer carries the architecture- + * specific sample flag (a top-byte tag on arm64, a magic header word on + * x86-64); unsampled pointers return immediately with their inputs unchanged. + * On the slow path it emits the ddheap:free USDT and returns the raw pointer + * and adjusted size that the caller must pass to the real free. + * + * The caller must use the ptr and size from the returned dd_alloc_freed_t + * rather than the originals when invoking the underlying deallocator, since + * sampled allocations on x86-64 have a 16-byte header that must be included + * in the free. + */ +#ifndef DD_SAMPLERS_ALLOCATION_FREED_H +#define DD_SAMPLERS_ALLOCATION_FREED_H + +#include + +#include + +/* + * Return type for dd_allocation_freed. Mirrors dd_alloc_req_t on the + * allocation side. + * + * ptr - pointer the caller MUST pass to the underlying deallocator. + * Equals the input for unsampled allocations; equals the raw + * pointer (user - sample_flag_overhead) for sampled ones on + * architectures that use inline flag headers. + * size - size the caller MUST pass to a sized-free variant; equals + * the input for unsampled allocations; may be larger for + * sampled ones that reserved header bytes at alloc time. + */ +typedef struct { + void *ptr; + size_t size; +} dd_alloc_freed_t; + +/* Slow path for sampled frees. */ +dd_alloc_freed_t dd_allocation_freed_slow(void *ptr, void *raw, size_t size, + size_t alignment) + __attribute__((warn_unused_result)); + +/* + * Hook invoked by an allocator wrapper BEFORE performing a free. + * Wraps free, operator delete (sized and unsized), sdallocx, etc. + * + * Checks whether the allocation at `ptr` was previously sampled and + * emits the matching `ddheap:free` USDT if so, then returns the + * (ptr, size) the caller must forward to the underlying deallocator + * verbatim. + * + * Args cover the superset of inputs any free-like call carries: + * ptr - allocation being freed (user pointer returned by alloc) + * size - size the caller knows about, or 0 if unknown (plain free) + * alignment - alignment used at allocation time, or 0 + */ +static inline __attribute__((always_inline)) +dd_alloc_freed_t dd_allocation_freed(void *ptr, size_t size, size_t alignment) { + void *raw; + if (__builtin_expect(dd_sample_flag_check_fast(ptr, &raw), 0)) { + return dd_allocation_freed_slow(ptr, raw, size, alignment); + } + + dd_alloc_freed_t out = { .ptr = ptr, .size = size }; + return out; +} + +#endif diff --git a/libdd-heap-sampler/include/datadog/heap/allocation_requested.h b/libdd-heap-sampler/include/datadog/heap/allocation_requested.h new file mode 100644 index 0000000000..3de2cfc425 --- /dev/null +++ b/libdd-heap-sampler/include/datadog/heap/allocation_requested.h @@ -0,0 +1,97 @@ +/** + * @file allocation_requested.h + * + * Pre-allocation hook for the ddheap sampler. + * + * Call dd_allocation_requested() immediately before invoking the underlying + * allocator. It runs the Poisson sampling decision for this allocation: + * most calls return quickly on the fast path (a counter decrement and a + * branch-not-taken); only allocations that cross the sampling boundary pay + * the cost of drawing a fresh inter-sample interval and setting the reentry + * guard. + * + * The returned dd_alloc_req_t carries two values the caller must forward: + * - size: the number of bytes to request from the real allocator + * (may be larger than the original on architectures that store + * the sample flag in a header word before the user pointer). + * - weight: 0 if not sampled; otherwise the unbiased size estimator + * (nsamples * interval) to attribute to this allocation. + * + * Always pair with dd_allocation_created(), even if the allocator fails. + */ +#ifndef DD_SAMPLERS_ALLOCATION_REQUESTED_H +#define DD_SAMPLERS_ALLOCATION_REQUESTED_H + +#include +#include + +#include + +/* + * Return type for dd_allocation_requested, paired with the `req` + * argument of dd_allocation_created. + * + * size - size the wrapper MUST pass to its underlying allocator. + * Usually the caller's requested size; on sampled + * allocations on architectures that use in-band flagging + * (header magic), this is bumped by the flag overhead. + * weight - 0 if this allocation was not sampled; otherwise the + * unbiased size estimator (nsamples * interval) for + * aggregated reporting. + * + * Kept to 16 bytes so it's returned in registers. + */ +typedef struct { + size_t size; + uint64_t weight; +} dd_alloc_req_t; + +/* Slow path for an allocation request. This is only taken when we think we + * need to sample, and is declared as a separate function to avoid bloating + the instruction cache of the fast path + */ +dd_alloc_req_t dd_allocation_requested_slow(dd_tl_state_t *tl, size_t size, + size_t alignment); + +/* + * Pre-allocation hook. Call BEFORE invoking the underlying allocator. + * This lets us decide if we want to sample the allocation or not, in + * which case in some situations (e.g., x86-64) we may increase the + * allocation size to store sampling metadata. + * + * Something like this - note that this function returns a new allocation size to use + * for the actual alloc request : + * + * dd_alloc_req_t req = dd_allocation_requested(size, alignment); + * void *raw = real_alloc(..., req.size); + * void *user = dd_allocation_created(raw, req); + * return user; + * + * The reentry guard is opened here (on sampled path only) and will be closed + * by the paired dd_allocation_created call. ALWAYS pair them, even on + * allocator failure (pass raw=NULL). + */ +static inline __attribute__((always_inline)) +dd_alloc_req_t dd_allocation_requested(size_t size, size_t alignment) { + dd_alloc_req_t out = { size, 0 }; + + // If we don't have TLS yet, 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 allocation), pass through without sampling. + // Either condition is rare on a hot path, so mark the branch unlikely. + dd_tl_state_t *tl = dd_tl_state_get_or_init_fast(); + if (__builtin_expect(!tl || tl->reentry_guard, 0)) return out; + + // If we haven't crossed the sampling boundary, do nothing and bail. + // Tell the compiler this branch is _likely_ - i.e. we're mostly going to not sample + tl->remaining_bytes += (int64_t)size; + if (__builtin_expect(tl->remaining_bytes < 0, 1)) return out; + + // Sampling path! Jump out to dd_allocation_requested_slow to do the necessary work. + // We're intentionally putting this in a separate function call to avoid bloating + // the i-cache of the fast path. At the point we have decided to sample, we are less + // concerned about the cost of the function call. + return dd_allocation_requested_slow(tl, size, alignment); +} + +#endif diff --git a/libdd-heap-sampler/include/datadog/heap/probes.h b/libdd-heap-sampler/include/datadog/heap/probes.h new file mode 100644 index 0000000000..4e8eabb46c --- /dev/null +++ b/libdd-heap-sampler/include/datadog/heap/probes.h @@ -0,0 +1,40 @@ +/* + * USDT probe emission functions for the ddheap provider. + * + * Defined in probes.c as regular non-inline functions so that each probe + * site has a single, stable address in the final binary. This matters + * because bindgen's wrap_static_fns generates tiny wrapper stubs for any + * static inline function it sees; if the DTRACE_PROBE macros expanded inside + * those stubs the resulting .note.stapsdt entries would carry section-relative + * offsets that bpftrace cannot resolve correctly. + */ + +#ifndef DD_SAMPLERS_PROBES_H +#define DD_SAMPLERS_PROBES_H + +#include + +#ifdef __linux__ + /* libbpf/usdt vendored at libdd-heap-sampler/vendor/usdt.h. Provides + * the variadic USDT() macro that emits the same v3 ELF-note format + * that bpftrace, systemtap, and BPF tracers all consume. */ +# include +#else +# define USDT(provider, name, ...) ((void)0) +#endif + +/* + * Emits the `ddheap:alloc` USDT. + * user - user-visible allocation pointer + * size - allocation size in bytes + * weight - unbiased size estimator (nsamples * interval) + */ +void dd_probe_alloc(void *user, uint64_t size, uint64_t weight); + +/* + * Emits the `ddheap:free` USDT. + * ptr - user-visible pointer being freed + */ +void dd_probe_free(void *ptr); + +#endif diff --git a/libdd-heap-sampler/include/datadog/heap/sample_flag.h b/libdd-heap-sampler/include/datadog/heap/sample_flag.h new file mode 100644 index 0000000000..b7a8db8012 --- /dev/null +++ b/libdd-heap-sampler/include/datadog/heap/sample_flag.h @@ -0,0 +1,175 @@ +/** + * @file sample_flag.h + * + * Architecture-specific "this allocation is sampled" flagging. + + * Two implementations live side-by-side in sample_flag.c under #ifdef: + * amd64 : header-magic (bump size, write magic before user pointer) + * arm64 : TBI pointer tag (flag ignored top byte) + * + * Memory tagging _should_ be the sweet spot as it adds no overhead to read, + * but we will see how it works out in practice. + */ +#ifndef DD_SAMPLERS_SAMPLE_FLAG_H +#define DD_SAMPLERS_SAMPLE_FLAG_H + +#include +#include +#include +#include + +/* + * Per-thread initialisation required by the flagging scheme. + * Must be called once per thread before dd_sample_flag_apply/check are used. + * On arm64 this issues prctl(PR_SET_TAGGED_ADDR_CTRL) so tagged pointers + * survive syscalls; on other architectures this is a no-op. + */ +void dd_sample_flag_thread_init(void); + +/* + * Extra bytes the flag reserves on top of the user's requested size. + * dd_allocation_requested adds this to the size it hands back to the + * wrapper so the underlying allocator reserves room for the flag. 0 on + * architectures that flag via pointer bits instead of memory. + */ +size_t dd_sample_flag_overhead(void); + +/* + * Apply the sampled flag. Takes the raw pointer from the underlying + * allocator (of size `user-requested + dd_sample_flag_overhead()`) and + * returns the user-visible pointer to hand to the application. + */ +void *dd_sample_flag_apply(void *raw); + +/* + * If `user` was previously returned by dd_sample_flag_apply, write the + * raw pointer (to pass to the underlying free) into *raw_out and return + * true. Otherwise leave *raw_out untouched and return false. + */ +bool dd_sample_flag_check(void *user, void **raw_out); + +#if defined(__x86_64__) + +#define DD_HEADER_BYTES 16 +#define DD_PAGE_SIZE 4096 + +/* + * x86-64 marks sampled allocations with a magic word written 16 bytes + * before the user pointer. The allocator hands us a 16-byte aligned raw + * pointer of size user_size + 32; we pick the user pointer inside that + * surplus so reading 16 bytes behind it always stays within a mapped + * page. + * + * Two layouts are used, distinguished by which magic is present: + * DD_MAGIC_A: user = raw + 16 (common case) + * DD_MAGIC_B: user = raw + 32 (only when raw + 16 would land in + * the first 16 bytes of a page; about + * 1/256 of sampled allocations) + * + * With +32 surplus exactly one of {raw+16, raw+32} is always safe: + * both are 16-byte aligned and their page offsets differ by 16, so at + * most one can fall in the unsafe [0, 16) band. The two-magic encoding + * lets us recover raw at free time without any extra metadata. + * This costs us on average 32 bytes on every 512 KiB allocated, or about + * 62 KiB on 1 GiB. + * + * Alternatives considered: + * - +16 surplus with a free+malloc retry when raw+16 is unsafe. + * Saves 16 bytes per sampled allocation in the common case (~99.6%) + * but plumbs malloc/free callbacks into the sampler and risks an + * OOM on the retry path. Not worth it: sampled allocations are + * rare so the absolute byte cost of always +32 is negligible. + * - Drop samples whose user pointer comes back in the unsafe band. + * Simple, but introduces a bias in the sample set + * (allocations whose raw lands at offset 4080 of a page are never + * reported). + */ + +#define DD_MAGIC_A 0xfab1eddec0dedca7ULL /* user = raw + 16 */ +#define DD_MAGIC_B 0xfab1eddec0dedca8ULL /* user = raw + 32 */ + +/* + * Layout helpers. x86_apply and x86_raw_from_user MUST be each other's + * inverse: the magic stamped at apply time is what lets x86_raw_from_user + * recover the same raw at check time. + */ + +/* + * Pick the user pointer inside the [raw, raw + 32) surplus and stamp + * the matching magic into the 16-byte header immediately preceding it. + * The chosen layout is encoded in the magic value itself so + * x86_raw_from_user can recover `raw` at free time without any extra + * metadata. + * + * raw - 16-byte aligned pointer from the underlying allocator, + * backing a region of size user_size + 32. + * + * Returns the user-visible pointer (raw+16 or raw+32), chosen so the + * 16-byte header slot never straddles the start of a page. + */ +static inline __attribute__((always_inline)) +void *x86_apply(void *raw) { + uintptr_t r = (uintptr_t)raw; + uintptr_t u = r + DD_HEADER_BYTES; + uint64_t magic = DD_MAGIC_A; + if ((u & (DD_PAGE_SIZE - 1)) < DD_HEADER_BYTES) { + u = r + 2 * DD_HEADER_BYTES; + magic = DD_MAGIC_B; + } + memcpy((void *)(u - DD_HEADER_BYTES), &magic, sizeof(magic)); + return (void *)u; +} + +/* + * Inverse of x86_apply: given a user pointer previously produced by + * it and the magic read out of the 16-byte header preceding `user` + * (DD_MAGIC_A => user = raw+16, DD_MAGIC_B => raw+32), return the + * original raw pointer to hand back to the underlying free. + */ +static inline __attribute__((always_inline)) +void *x86_raw_from_user(void *user, uint64_t magic) { + uintptr_t u = (uintptr_t)user; + return (void *)(u - (magic == DD_MAGIC_A ? DD_HEADER_BYTES + : 2 * DD_HEADER_BYTES)); +} + +static inline __attribute__((always_inline)) +bool dd_sample_flag_check_fast(void *user, void **raw_out) { + if (((uintptr_t)user & (DD_PAGE_SIZE - 1)) < DD_HEADER_BYTES) { + return false; + } + + void *header = (char *)user - DD_HEADER_BYTES; + uint64_t magic; + memcpy(&magic, header, sizeof(magic)); + if (magic != DD_MAGIC_A && magic != DD_MAGIC_B) { + return false; + } + + const uint64_t zero = 0; + memcpy(header, &zero, sizeof(zero)); + *raw_out = x86_raw_from_user(user, magic); + return true; +} + +#elif defined(__aarch64__) + +#define DD_TBI_TAG 0xDDu +#define DD_TBI_TAG_MASK ((uintptr_t)0xFFu << 56) +#define DD_TBI_TAGGED ((uintptr_t)DD_TBI_TAG << 56) + +static inline __attribute__((always_inline)) +bool dd_sample_flag_check_fast(void *user, void **raw_out) { + uintptr_t addr = (uintptr_t)user; + if ((addr & DD_TBI_TAG_MASK) != DD_TBI_TAGGED) { + return false; + } + *raw_out = (void *)(addr & ~DD_TBI_TAG_MASK); + return true; +} + +#else +# error "dd_sample_flag: unsupported architecture (expected __x86_64__ or __aarch64__)" +#endif + +#endif diff --git a/libdd-heap-sampler/include/datadog/heap/tl_state.h b/libdd-heap-sampler/include/datadog/heap/tl_state.h new file mode 100644 index 0000000000..a99ce5819e --- /dev/null +++ b/libdd-heap-sampler/include/datadog/heap/tl_state.h @@ -0,0 +1,149 @@ +/** + * @file tl_state.h + * + * Per-thread sampler state. How to store this is _nuanced_. We can choose between + * the older, thread-specific data (TSD) APIs (pthread_getspecific, pthread_setspecific, ...), + * and the newer, compiler-side one - __thread. In the latter case we must also + * consider the impact of TLS model, and in all cases, we have to be careful + * to not accidentally infinitely recurse when the access mechanism must allocate. + * + * --- Summary of options ----------------------------------------------------- + * + * A. initial-exec TLS, exclude musl late-dlopen + * Direct TP-relative load on every access. Fast. Works for static builds + * and glibc dynamic builds (our struct fits within glibc's static TLS + * surplus). Hard fails at dlopen time on musl: loader rejects before any + * of our code runs, so no runtime fallback is possible. + * + * B. TLSDESC, single library <-- let's start here + * Works for static builds, glibc dynamic, and musl dynamic. For static + * builds the linker relaxes TLSDESC to local-exec automatically, so + * there is no per-access overhead vs A. For musl dynamic, __tls_get_addr + * is a pure DTV lookup (musl pre-populates at dlopen time): also no + * extra overhead vs what musl can offer. The reentrancy concern on glibc + * is eliminated because the gotter skips ld-linux in its GOT walk, so + * __tls_get_addr's internal malloc never goes through our hook. The only + * case where B costs more than necessary is glibc dynamic: we pay the + * TLSDESC indirect call where initial-exec would give us a direct + * TP-relative load. That is the sole remaining argument for option C. + * (Pre-warming via pthread_create hook was considered and rejected: it + * cannot cover threads that existed before the gotter installed.) + * + * C. Two build variants, caller picks + * Build a -glibc variant (initial-exec) and a -musl variant (TLSDESC). + * Each gets optimal per-access performance for its runtime. The caller + * detects musl/glibc and loads the appropriate one. More complex build + * and deployment story. + * + * Constraints: + * - Must work for both static builds and dynamic builds, including late + * dlopen (we cannot assume we are loaded at startup). + * - Allocation reentrancy during TLS init is a real problem: any allocating + * path inside the sampler re-enters the hook before state exists. + * - The sampler must handle init generically, without relying on the caller + * (gotter, Rust allocator, etc.) doing any injection-specific setup first. + * - Ignoring musl + late dlopen as a supported target may be acceptable, + * and would simplify the dynamic build story considerably. + * + * --- Thread-local mechanism ------------------------------------------------ + * + * Two broad approaches. "Thread-Specific Data" (TSD) is the POSIX runtime + * model; "Thread-Local Storage" (TLS) is the compiler/linker model. + * + * pthread_key_t TSD. A runtime key-value store: pthread_key_create once, + * pthread_getspecific / pthread_setspecific per thread. + * The value is a void*, so you heap-allocate the struct + * and store a pointer. pthread_getspecific is cheap on + * glibc (array lookup in thread memory, no __tls_get_addr). + * ddprof uses this (see allocation_tracker.cc) to avoid + * __tls_get_addr allocating in the Global Dynamic model; + * ddprof builds a universal musl/glibc binary so + * initial-exec _Thread_local is not an option for them. + * Downside: init must heap-allocate the struct, which + * re-enters the allocator hook before state exists. + * + * _Thread_local TLS. C11 standard spelling of GCC's older __thread + * (__thread) extension; both compile identically on GCC and Clang. + * Storage lives in the thread's TLS segment, not the heap. + * No allocation on access, but the TLS model matters (see + * below) and __tls_get_addr can still allocate on first + * access in the Global Dynamic model. +**/ +#ifndef DD_SAMPLERS_TL_STATE_H +#define DD_SAMPLERS_TL_STATE_H + +#include +#include +#include + +/* 512 KiB mean between samples. */ +#define DD_SAMPLING_INTERVAL_DEFAULT (512u * 1024u) + +/* Per-thread state for the Poisson sampler. See file header for the + * rationale behind _Thread_local vs pthread TLS. */ +typedef struct { + uint64_t sampling_interval; /* mean bytes between samples. + This _will probably_ be constant, but if we + drop it in the TLS we afford the eBPF profiler + the opportunity to tune it to adjust overhead + dynamically. Whether or not this turns out to be + a clever idea remains to be seen. */ + int64_t remaining_bytes; /* signed counter; sample when >= 0 */ + bool remaining_bytes_initialized; /* false until first interval drawn */ + bool initialized; /* false until dd_tl_state_init() has run; + field rather than a separate _Thread_local + so only one TLS lookup is needed on the + fast path (avoids a second TLSDESC call). */ + bool reentry_guard; /* Set between dd_allocation_requested_slow and + dd_allocation_created_slow while a sampled + allocation is in flight. + + The slow path can itself trigger allocations: + log() in next_interval() may touch lazy-init'd + libc state on first call; USDT emission and any + attached eBPF consumer can cause incidental + userspace allocation; TLS materialisation on a + fresh thread may call calloc. + + Without this guard, those inner allocations + would re-enter dd_allocation_requested, corrupt + the remaining_bytes counter, and in the worst + case recurse infinitely. While the guard is set + the fast path bails out immediately and the + inner allocation is passed through unsampled. */ + uint32_t rng; /* LCG state */ + +} dd_tl_state_t; + +extern _Thread_local dd_tl_state_t dd_tl_state_storage; + +/* + * Returns the current thread's state, or NULL if not yet initialised. + * Never allocates. Callers must treat NULL as "don't sample". + */ +dd_tl_state_t *dd_tl_state_get(void); + +static inline __attribute__((always_inline)) +dd_tl_state_t *dd_tl_state_get_fast(void) { + if (__builtin_expect(!dd_tl_state_storage.initialized, 0)) return NULL; + return &dd_tl_state_storage; +} + +/* + * Initializes the current thread's tracking state if not already present. + * Returns the new state, or NULL if state already exists on this thread. + */ +dd_tl_state_t *dd_tl_state_init(void); + +/* + * Returns the current thread's tracking state, initializing it on first use. + */ +dd_tl_state_t *dd_tl_state_get_or_init(void); + +static inline __attribute__((always_inline)) +dd_tl_state_t *dd_tl_state_get_or_init_fast(void) { + if (__builtin_expect(dd_tl_state_storage.initialized, 1)) return &dd_tl_state_storage; + return dd_tl_state_init(); +} + +#endif diff --git a/libdd-heap-sampler/src/allocation_created.c b/libdd-heap-sampler/src/allocation_created.c new file mode 100644 index 0000000000..0381a973d7 --- /dev/null +++ b/libdd-heap-sampler/src/allocation_created.c @@ -0,0 +1,46 @@ +// Copyright 2025-Present Datadog, Inc. https://www.datadoghq.com/ +// SPDX-License-Identifier: Apache-2.0 + +#include +#include +#include +#include + +/* + * Slow path for dd_allocation_created. We only arrive here when the paired + * dd_allocation_requested_slow decided to sample (req.weight > 0). + * + * Applies the architecture-specific sample flag to raw (tagging the pointer + * on arm64, writing a header magic word on x86-64) to produce the + * user-visible pointer, then fires the ddheap:alloc USDT so an attached + * profiler can record the sample. Finally closes the reentry guard that + * dd_allocation_requested_slow opened. + * + * raw may be NULL if the underlying allocator failed; in that case we skip + * the flag and USDT but still close the guard. + * + * TODO: on x86-64, consider abandoning the sample when raw falls within + * DD_HEADER_BYTES of a page boundary. The free-side fast path already bails + * in that case (to avoid reading before the page), so those allocations will + * never emit a ddheap:free. Dropping them at alloc time would keep the + * alloc/free pair balanced at the cost of occasionally missing a sample. + * + * (ddprof: AllocTrackerHelper::track / AllocationTracker push_alloc_sample) + */ +void *dd_allocation_created_slow(void *raw, dd_alloc_req_t req) { + void *user = raw; + if (raw != NULL) { + /* Apply the sample flag and fire the USDT. We use the user pointer + * (post-flag) as the USDT argument so the profiler sees the same + * address the application will. */ + user = dd_sample_flag_apply(raw); + dd_probe_alloc(user, (uint64_t)req.size, req.weight); + } + + /* Always close the reentry guard, even on allocation failure (raw == NULL), + * so the thread isn't permanently locked out of sampling. */ + dd_tl_state_t *tl = dd_tl_state_get(); + if (tl) tl->reentry_guard = false; + + return user; +} \ No newline at end of file diff --git a/libdd-heap-sampler/src/allocation_freed.c b/libdd-heap-sampler/src/allocation_freed.c new file mode 100644 index 0000000000..14b5ca4dee --- /dev/null +++ b/libdd-heap-sampler/src/allocation_freed.c @@ -0,0 +1,36 @@ +// Copyright 2025-Present Datadog, Inc. https://www.datadoghq.com/ +// SPDX-License-Identifier: Apache-2.0 + +#include +#include +#include + +/* + * Slow path for dd_allocation_freed. We only arrive here when + * dd_sample_flag_check_fast confirmed that ptr carries the sample flag, + * meaning this allocation was previously sampled. + * + * Fires the ddheap:free USDT with the user-visible pointer, then returns + * the raw pointer and adjusted size that the caller must forward to the + * real deallocator. On x86-64 the size grows by DD_HEADER_BYTES to cover + * the header that was reserved at allocation time; on arm64 the size is + * unchanged since TBI tagging touches only pointer bits. + */ +dd_alloc_freed_t dd_allocation_freed_slow(void *ptr, void *raw, size_t size, + size_t alignment) { + (void)alignment; + + /* Fire with the user-visible pointer, matching what was reported at alloc + * time, so the profiler can correlate the two events by address. */ + dd_probe_free(ptr); + dd_alloc_freed_t out = { + /* Return the raw pointer so the caller passes the real allocation base + * to the deallocator, not the user pointer that may be offset or tagged. */ + .ptr = raw, + /* Recover the full allocation size including any header reserved at + * alloc time. On arm64 dd_sample_flag_overhead() is 0, so this is + * a no-op there. */ + .size = size + dd_sample_flag_overhead(), + }; + return out; +} \ No newline at end of file diff --git a/libdd-heap-sampler/src/allocation_requested.c b/libdd-heap-sampler/src/allocation_requested.c new file mode 100644 index 0000000000..66a9f8f5b6 --- /dev/null +++ b/libdd-heap-sampler/src/allocation_requested.c @@ -0,0 +1,115 @@ +// Copyright 2025-Present Datadog, Inc. https://www.datadoghq.com/ +// SPDX-License-Identifier: Apache-2.0 + +#include +#include + +#include + +/* + * Advances the Park-Miller LCG one step and returns the new 31-bit state. + * Cheap, branch-free PRNG suitable for the sampling hot path. + */ +static uint32_t lcg_next(uint32_t *rng) { + *rng = (uint32_t)(((uint64_t)(*rng) * 48271u) % 2147483647u); + return *rng; +} + +/* + * Draws the next inter-sample gap in bytes from an exponential distribution + * with the given mean. Clamped to [8, 20*mean] to avoid degenerate near-zero + * gaps on one end and unbounded intervals on unlucky draws on the other. + */ +static uint64_t next_interval(uint32_t *rng, uint64_t mean) { + double u = (double)lcg_next(rng) / 2147483647.0; + if (u <= 0.0) u = 1e-10; /* guard against log(0) = -inf */ + double v = -log(u) * (double)mean; + double vmax = 20.0 * (double)mean; + if (v > vmax) v = vmax; /* cap runaway intervals on very lucky draws */ + if (v < 8.0) v = 8.0; /* floor keeps the counter moving forward */ + return (uint64_t)v; +} + +/* + * Called when remaining_bytes has crossed zero, meaning at least one sample + * is owed. Draws fresh intervals until the counter is negative again, counting + * how many samples fired. Returns nsamples * interval as the unbiased weight + * estimator to attribute to this allocation. + * + * On the very first call for a thread, remaining_bytes_initialized is false + * and we draw the initial interval from scratch. If that interval exceeds the + * accumulated byte credit the counter goes back negative and we return 0, + * meaning no sample this time. This is normal and happens at most once per thread. + * + * Note: remaining_bytes has already been incremented by `size` in the inline + * fast path; we arrive here because that increment pushed it to zero or above. + */ +static uint64_t sample(dd_tl_state_t *tl) { + uint64_t interval = tl->sampling_interval; + + if (!tl->remaining_bytes_initialized) { + /* First allocation on this thread: draw the initial interval and + * subtract it from the credit accumulated so far. If we're already + * back in the red, skip the sample; the counter just wasn't large + * enough to cover the first interval. */ + tl->remaining_bytes -= (int64_t)next_interval(&tl->rng, interval); + tl->remaining_bytes_initialized = true; + if (tl->remaining_bytes < 0) return 0; + } + + /* remaining_bytes is >= 0, meaning we've crossed one full interval + * boundary. Integer-divide to find how many full intervals fit in the + * current credit (usually 1, but can be more for very large allocations), + * then keep drawing until we're back in the red. Each iteration is one + * additional sample. */ + size_t nsamples = (size_t)tl->remaining_bytes / interval; + tl->remaining_bytes %= (int64_t)interval; + do { + tl->remaining_bytes -= (int64_t)next_interval(&tl->rng, interval); + ++nsamples; + } while (tl->remaining_bytes >= 0); + + /* Weight is the unbiased estimator: each sample represents `interval` + * bytes on average, so nsamples * interval gives the expected total. */ + return (uint64_t)nsamples * interval; +} + +/* + * Slow path for dd_allocation_requested. We only arrive here when the fast + * path counter has crossed zero. Sets the reentry guard, runs the sampling + * decision, and returns the allocation request the caller should forward to + * the real allocator. + * + * If sample() returns 0 (first-interval miss on a fresh thread) the guard is + * closed here and a no-sample result is returned. Otherwise the guard stays + * open until dd_allocation_created_slow closes it, keeping any allocations + * triggered during the slow path from re-entering the sampler. + * + * (ddprof: AllocationTracker::track_allocation / next_sample_interval) + */ +dd_alloc_req_t dd_allocation_requested_slow(dd_tl_state_t *tl, size_t size, + size_t alignment) { + (void)alignment; + + /* Open the reentry guard before doing anything else. Any allocation that + * happens between here and dd_allocation_created_slow (e.g. inside log() + * or the USDT machinery) will see the guard set and pass through unsampled. */ + tl->reentry_guard = true; + + uint64_t weight = sample(tl); + if (weight == 0) { + /* First-interval miss: no sample this time. Close the guard now since + * dd_allocation_created_slow won't be called on the sampled path. */ + tl->reentry_guard = false; + dd_alloc_req_t out = { size, 0 }; + return out; + } + + dd_alloc_req_t out = { + /* Ask the allocator for extra bytes if the flag scheme needs them + * (16 bytes on x86-64 for the header; 0 on arm64). */ + .size = size + dd_sample_flag_overhead(), + .weight = weight, + }; + return out; +} \ No newline at end of file diff --git a/libdd-heap-sampler/src/generated/bindings.rs b/libdd-heap-sampler/src/generated/bindings.rs new file mode 100644 index 0000000000..5d4078d674 --- /dev/null +++ b/libdd-heap-sampler/src/generated/bindings.rs @@ -0,0 +1,149 @@ +// Copyright 2025-Present Datadog, Inc. https://www.datadoghq.com/ +// SPDX-License-Identifier: Apache-2.0 +// @generated by libdd-heap-sampler/build.rs via bindgen; do not edit by hand. +// Regenerate with: LIBDD_HEAP_SAMPLER_REGEN=1 cargo build -p libdd-heap-sampler + +/* automatically generated by rust-bindgen 0.71.1 */ + +pub const DD_SAMPLING_INTERVAL_DEFAULT: u32 = 524288; +#[repr(C)] +#[derive(Debug, Copy, Clone)] +pub struct dd_tl_state_t { + pub sampling_interval: u64, + pub remaining_bytes: i64, + pub remaining_bytes_initialized: bool, + pub initialized: bool, + pub reentry_guard: bool, + pub rng: u32, +} +#[allow(clippy::unnecessary_operation, clippy::identity_op)] +const _: () = { + ["Size of dd_tl_state_t"][::std::mem::size_of::() - 24usize]; + ["Alignment of dd_tl_state_t"][::std::mem::align_of::() - 8usize]; + ["Offset of field: dd_tl_state_t::sampling_interval"] + [::std::mem::offset_of!(dd_tl_state_t, sampling_interval) - 0usize]; + ["Offset of field: dd_tl_state_t::remaining_bytes"] + [::std::mem::offset_of!(dd_tl_state_t, remaining_bytes) - 8usize]; + ["Offset of field: dd_tl_state_t::remaining_bytes_initialized"] + [::std::mem::offset_of!(dd_tl_state_t, remaining_bytes_initialized) - 16usize]; + ["Offset of field: dd_tl_state_t::initialized"] + [::std::mem::offset_of!(dd_tl_state_t, initialized) - 17usize]; + ["Offset of field: dd_tl_state_t::reentry_guard"] + [::std::mem::offset_of!(dd_tl_state_t, reentry_guard) - 18usize]; + ["Offset of field: dd_tl_state_t::rng"][::std::mem::offset_of!(dd_tl_state_t, rng) - 20usize]; +}; +unsafe extern "C" { + pub fn dd_tl_state_get() -> *mut dd_tl_state_t; +} +unsafe extern "C" { + #[link_name = "dd_tl_state_get_fast__extern"] + pub fn dd_tl_state_get_fast() -> *mut dd_tl_state_t; +} +unsafe extern "C" { + pub fn dd_tl_state_init() -> *mut dd_tl_state_t; +} +unsafe extern "C" { + pub fn dd_tl_state_get_or_init() -> *mut dd_tl_state_t; +} +unsafe extern "C" { + #[link_name = "dd_tl_state_get_or_init_fast__extern"] + pub fn dd_tl_state_get_or_init_fast() -> *mut dd_tl_state_t; +} +#[repr(C)] +#[derive(Debug, Copy, Clone)] +pub struct dd_alloc_req_t { + pub size: usize, + pub weight: u64, +} +#[allow(clippy::unnecessary_operation, clippy::identity_op)] +const _: () = { + ["Size of dd_alloc_req_t"][::std::mem::size_of::() - 16usize]; + ["Alignment of dd_alloc_req_t"][::std::mem::align_of::() - 8usize]; + ["Offset of field: dd_alloc_req_t::size"] + [::std::mem::offset_of!(dd_alloc_req_t, size) - 0usize]; + ["Offset of field: dd_alloc_req_t::weight"] + [::std::mem::offset_of!(dd_alloc_req_t, weight) - 8usize]; +}; +unsafe extern "C" { + pub fn dd_allocation_requested_slow( + tl: *mut dd_tl_state_t, + size: usize, + alignment: usize, + ) -> dd_alloc_req_t; +} +unsafe extern "C" { + #[link_name = "dd_allocation_requested__extern"] + pub fn dd_allocation_requested(size: usize, alignment: usize) -> dd_alloc_req_t; +} +unsafe extern "C" { + pub fn dd_allocation_created_slow( + raw: *mut ::std::os::raw::c_void, + req: dd_alloc_req_t, + ) -> *mut ::std::os::raw::c_void; +} +unsafe extern "C" { + #[link_name = "dd_allocation_created__extern"] + pub fn dd_allocation_created( + raw: *mut ::std::os::raw::c_void, + req: dd_alloc_req_t, + ) -> *mut ::std::os::raw::c_void; +} +unsafe extern "C" { + pub fn dd_sample_flag_thread_init(); +} +unsafe extern "C" { + pub fn dd_sample_flag_overhead() -> usize; +} +unsafe extern "C" { + pub fn dd_sample_flag_apply(raw: *mut ::std::os::raw::c_void) -> *mut ::std::os::raw::c_void; +} +unsafe extern "C" { + pub fn dd_sample_flag_check( + user: *mut ::std::os::raw::c_void, + raw_out: *mut *mut ::std::os::raw::c_void, + ) -> bool; +} +unsafe extern "C" { + #[link_name = "dd_sample_flag_check_fast__extern"] + pub fn dd_sample_flag_check_fast( + user: *mut ::std::os::raw::c_void, + raw_out: *mut *mut ::std::os::raw::c_void, + ) -> bool; +} +#[repr(C)] +#[derive(Debug, Copy, Clone)] +pub struct dd_alloc_freed_t { + pub ptr: *mut ::std::os::raw::c_void, + pub size: usize, +} +#[allow(clippy::unnecessary_operation, clippy::identity_op)] +const _: () = { + ["Size of dd_alloc_freed_t"][::std::mem::size_of::() - 16usize]; + ["Alignment of dd_alloc_freed_t"][::std::mem::align_of::() - 8usize]; + ["Offset of field: dd_alloc_freed_t::ptr"] + [::std::mem::offset_of!(dd_alloc_freed_t, ptr) - 0usize]; + ["Offset of field: dd_alloc_freed_t::size"] + [::std::mem::offset_of!(dd_alloc_freed_t, size) - 8usize]; +}; +unsafe extern "C" { + pub fn dd_allocation_freed_slow( + ptr: *mut ::std::os::raw::c_void, + raw: *mut ::std::os::raw::c_void, + size: usize, + alignment: usize, + ) -> dd_alloc_freed_t; +} +unsafe extern "C" { + #[link_name = "dd_allocation_freed__extern"] + pub fn dd_allocation_freed( + ptr: *mut ::std::os::raw::c_void, + size: usize, + alignment: usize, + ) -> dd_alloc_freed_t; +} +unsafe extern "C" { + pub fn dd_probe_alloc(user: *mut ::std::os::raw::c_void, size: u64, weight: u64); +} +unsafe extern "C" { + pub fn dd_probe_free(ptr: *mut ::std::os::raw::c_void); +} diff --git a/libdd-heap-sampler/src/generated/dd_heap_sampler_static_wrappers.c b/libdd-heap-sampler/src/generated/dd_heap_sampler_static_wrappers.c new file mode 100644 index 0000000000..805de31ced --- /dev/null +++ b/libdd-heap-sampler/src/generated/dd_heap_sampler_static_wrappers.c @@ -0,0 +1,20 @@ +// Copyright 2025-Present Datadog, Inc. https://www.datadoghq.com/ +// SPDX-License-Identifier: Apache-2.0 +// @generated by libdd-heap-sampler/build.rs via bindgen; do not edit by hand. +// Regenerate with: LIBDD_HEAP_SAMPLER_REGEN=1 cargo build -p libdd-heap-sampler + +#include "include/datadog/heap/allocation_requested.h" +#include "include/datadog/heap/allocation_created.h" +#include "include/datadog/heap/allocation_freed.h" +#include "include/datadog/heap/probes.h" +#include "include/datadog/heap/sample_flag.h" +#include "include/datadog/heap/tl_state.h" + +// Static wrappers + +dd_tl_state_t * dd_tl_state_get_fast__extern(void) { return dd_tl_state_get_fast(); } +dd_tl_state_t * dd_tl_state_get_or_init_fast__extern(void) { return dd_tl_state_get_or_init_fast(); } +dd_alloc_req_t dd_allocation_requested__extern(size_t size, size_t alignment) { return dd_allocation_requested(size, alignment); } +void * dd_allocation_created__extern(void *raw, dd_alloc_req_t req) { return dd_allocation_created(raw, req); } +bool dd_sample_flag_check_fast__extern(void *user, void **raw_out) { return dd_sample_flag_check_fast(user, raw_out); } +dd_alloc_freed_t dd_allocation_freed__extern(void *ptr, size_t size, size_t alignment) { return dd_allocation_freed(ptr, size, alignment); } diff --git a/libdd-heap-sampler/src/lib.rs b/libdd-heap-sampler/src/lib.rs new file mode 100644 index 0000000000..1d623a7794 --- /dev/null +++ b/libdd-heap-sampler/src/lib.rs @@ -0,0 +1,111 @@ +// Copyright 2025-Present Datadog, Inc. https://www.datadoghq.com/ +// SPDX-License-Identifier: Apache-2.0 + +//! Rust bindings for the libdd-heap-sampler C library. +//! +//! FFI declarations and `#[repr(C)]` struct mirrors are generated by +//! `bindgen` from the public headers under `include/datadog/heap/` at +//! build time; see `build.rs`. All items are re-exported at the crate +//! root. +//! +//! The `static inline __attribute__((always_inline))` fast-path helpers in +//! the headers have no linker symbol, so bindgen's `wrap_static_fns` +//! emits a C shim for each - the shim's body is the static-inline +//! inlined at its one call site. The shim's link-symbol is +//! `__extern`, but bindgen exposes the Rust fn under the natural C +//! name via `#[link_name]`, so callers write `dd_allocation_requested(…)` +//! the same way they would in C. + +// USDT/SystemTap is Linux-only; on every other target the crate is an empty +// rlib. macOS is `unix` too, so we gate on target_os = "linux" specifically. +#[cfg(target_os = "linux")] +#[allow( + non_camel_case_types, + non_snake_case, + non_upper_case_globals, + dead_code +)] +mod sys { + include!(concat!(env!("OUT_DIR"), "/bindings.rs")); +} + +#[cfg(target_os = "linux")] +pub use sys::*; + +// Tests exercise the C/FFI sampler primitives directly; miri can't +// execute foreign code, so skip the whole module under miri. +#[cfg(all(test, target_os = "linux", not(miri)))] +mod tests { + use super::*; + use core::ffi::c_void; + + #[test] + fn alloc_req_layout_matches_c() { + assert_eq!(core::mem::size_of::(), 16); + assert_eq!(core::mem::align_of::(), 8); + } + + #[test] + fn tl_state_layout_matches_c() { + assert_eq!(core::mem::size_of::(), 24); + assert_eq!(core::mem::align_of::(), 8); + } + + #[test] + fn requested_initializes_tls_on_first_use() { + std::thread::spawn(|| unsafe { + assert!( + dd_tl_state_get().is_null(), + "fresh thread should start without sampler TLS" + ); + 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" + ); + }) + .join() + .unwrap(); + } + + #[test] + fn created_is_passthrough_when_weight_zero() { + unsafe { + let fake = 0xdead_beef_usize as *mut c_void; + let req = dd_alloc_req_t { + size: 64, + weight: 0, + }; + assert_eq!(dd_allocation_created(fake, req), fake); + } + } + + #[test] + fn freed_unsampled_returns_inputs_unchanged() { + // The C side now reads 8 bytes at ptr-16 looking for DD_MAGIC, + // so we need a real buffer underneath. Zero contents won't match, + // so the function returns inputs verbatim. + let mut buf = vec![0u8; 128]; + let ptr = unsafe { buf.as_mut_ptr().add(16) } as *mut c_void; + unsafe { + let freed = dd_allocation_freed(ptr, 64, 8); + assert_eq!(freed.ptr, ptr); + assert_eq!(freed.size, 64); + } + } + + #[test] + fn tl_state_init_then_get_returns_same_non_null() { + std::thread::spawn(|| unsafe { + let initialized = dd_tl_state_init(); + assert!(!initialized.is_null(), "init returned NULL on fresh thread"); + let got = dd_tl_state_get(); + assert_eq!(initialized, got); + assert_eq!((*got).sampling_interval, 512 * 1024); + }) + .join() + .unwrap(); + } +} diff --git a/libdd-heap-sampler/src/probes.c b/libdd-heap-sampler/src/probes.c new file mode 100644 index 0000000000..bdc6362013 --- /dev/null +++ b/libdd-heap-sampler/src/probes.c @@ -0,0 +1,28 @@ +// Copyright 2025-Present Datadog, Inc. https://www.datadoghq.com/ +// SPDX-License-Identifier: Apache-2.0 + +/* + * USDT emission for the ddheap provider. + * + * Kept as non-inline functions in a separate translation unit so that each + * probe has one USDT() expansion and one .note.stapsdt entry. The intent + * is that callers always reach the probe via a call instruction rather + * than having it inlined at multiple sites. + * + * The immediate concern is bindgen's wrap_static_fns: if these were static + * inline, it would generate a wrapper stub for each one containing its own + * USDT() expansion, likely producing duplicate .note.stapsdt entries and + * causing bpftrace to attach twice. LTO could in principle inline these + * across TU boundaries and cause similar problems, though we haven't + * tested that path. + */ + +#include + +void dd_probe_alloc(void *user, uint64_t size, uint64_t weight) { + USDT(ddheap, alloc, user, size, weight); +} + +void dd_probe_free(void *ptr) { + USDT(ddheap, free, ptr); +} \ No newline at end of file diff --git a/libdd-heap-sampler/src/sample_flag.c b/libdd-heap-sampler/src/sample_flag.c new file mode 100644 index 0000000000..9ee2004087 --- /dev/null +++ b/libdd-heap-sampler/src/sample_flag.c @@ -0,0 +1,94 @@ +// Copyright 2025-Present Datadog, Inc. https://www.datadoghq.com/ +// SPDX-License-Identifier: Apache-2.0 + +#include + +#if defined(__x86_64__) + +/* + * x86-64 layout details live in sample_flag.h alongside DD_MAGIC_A/B + * and the x86 helper pair. Summary: allocate user_size + 32, + * let the helper pick the user pointer (raw+16 in the common case, + * raw+32 when raw+16 would land in the unsafe first-16-bytes-of-page + * region), and write the corresponding magic at user-16. On free the + * matched magic tells us which offset to use to recover raw. + */ + +void dd_sample_flag_thread_init(void) { /* nothing needed on x86-64 */ } + +size_t dd_sample_flag_overhead(void) { + // +32 guarantees both raw+16 and raw+32 are available; one of them + // is always safe to use as the user pointer regardless of where raw + // lands relative to a page boundary. + return 2 * DD_HEADER_BYTES; +} + +void *dd_sample_flag_apply(void *raw) { + return x86_apply(raw); +} + +bool dd_sample_flag_check(void *user, void **raw_out) { + return dd_sample_flag_check_fast(user, raw_out); +} + +#elif defined(__aarch64__) + +/* + * arm64: TBI (Top Byte Ignore) pointer tag. + * + * ARMv8 userspace ignores bits 56..63 of a virtual address, so we stash + * DD_TBI_TAG in those bits to mark sampled allocations. The allocation + * itself is untouched (no size bump, no header writes), so overhead is 0. + * + * On apply: OR DD_TBI_TAGGED into the pointer and return it. + * On check: test bits 56..63; if they match DD_TBI_TAG, clear them to + * recover the raw pointer, write it to *raw_out, and return true. + * + * Tagged pointers must survive kernel boundaries (e.g. passing a tagged + * pointer to read/write). dd_sample_flag_thread_init calls + * prctl(PR_SET_TAGGED_ADDR_CTRL) on Linux to enable this. Without it the + * kernel would reject tagged pointers with EFAULT. + * + * TODO: audit interaction with HWASan and MTE, which also use the top byte. + */ + +size_t dd_sample_flag_overhead(void) { + return 0; +} + +void *dd_sample_flag_apply(void *raw) { + return (void *)((uintptr_t)raw | DD_TBI_TAGGED); +} + +bool dd_sample_flag_check(void *user, void **raw_out) { + return dd_sample_flag_check_fast(user, raw_out); +} + +#if defined(__linux__) +#include + +/* PR_SET_TAGGED_ADDR_CTRL / PR_TAGGED_ADDR_ENABLE only reached glibc in + * 2.31 (Feb 2020); some of our internal CI images (notably the + * libddprof-build centos image) pin an older glibc whose + * does not declare them, breaking the build with an + * undeclared-identifier error even though the kernel underneath (any + * Linux ≥ 5.4) accepts the syscall fine. The prctl numbers are stable + * kernel ABI — verified against `linux/prctl.h` — so define them + * locally when the libc header does not. */ +#ifndef PR_SET_TAGGED_ADDR_CTRL +# define PR_SET_TAGGED_ADDR_CTRL 55 +#endif +#ifndef PR_TAGGED_ADDR_ENABLE +# define PR_TAGGED_ADDR_ENABLE (1UL << 0) +#endif + +void dd_sample_flag_thread_init(void) { + prctl(PR_SET_TAGGED_ADDR_CTRL, PR_TAGGED_ADDR_ENABLE, 0, 0, 0); +} +#else +void dd_sample_flag_thread_init(void) { } +#endif + +#else +# error "dd_sample_flag: unsupported architecture (expected __x86_64__ or __aarch64__)" +#endif \ No newline at end of file diff --git a/libdd-heap-sampler/src/tl_state.c b/libdd-heap-sampler/src/tl_state.c new file mode 100644 index 0000000000..34cd8d39c7 --- /dev/null +++ b/libdd-heap-sampler/src/tl_state.c @@ -0,0 +1,72 @@ +// Copyright 2025-Present Datadog, Inc. https://www.datadoghq.com/ +// SPDX-License-Identifier: Apache-2.0 + +#include +#include + +#include +#include + +_Thread_local dd_tl_state_t dd_tl_state_storage; + +/* + * Fills a freshly zeroed dd_tl_state_t with its initial values. + * + * Seeds the LCG from a mix of the TLS storage address and a monotonic + * clock read. Neither source is cryptographic, but the LCG only drives + * allocation *sampling* decisions - all we need is that distinct threads + * start with distinct non-zero seeds so their sample sequences + * decorrelate. The TLS address gives per-thread and (thanks to ASLR) + * per-process variation; the CLOCK_MONOTONIC nanoseconds decorrelate + * threads created back-to-back whose TLS slots differ only in low bits. + * clock_gettime is vDSO-served on Linux and available on every libc we + * ship on, unlike 's getentropy (glibc >= 2.25, musl + * >= 1.1.20). Copies in the default sampling interval so the eBPF + * profiler can tune it per-thread at runtime by writing to the TLS slot + * directly. + */ +static void tl_state_populate(dd_tl_state_t *st) { + *st = (dd_tl_state_t){0}; + + /* Set both flags before doing any work. The subsequent function calls + * (dd_sample_flag_thread_init, clock_gettime) act as compiler barriers, + * so the compiler cannot sink these writes below them. Any allocation + * triggered inside those calls will therefore see both flags set and + * bail out unsampled rather than recursing back into init. + * + */ + st->initialized = true; + st->reentry_guard = true; + + dd_sample_flag_thread_init(); + + struct timespec ts = {0}; + (void)clock_gettime(CLOCK_MONOTONIC, &ts); + uint32_t seed = (uint32_t)((uintptr_t)st + ^ (uintptr_t)ts.tv_nsec + ^ ((uintptr_t)ts.tv_sec << 20)); + st->rng = seed ? seed : 1u; + + st->sampling_interval = DD_SAMPLING_INTERVAL_DEFAULT; + st->reentry_guard = false; +} + +dd_tl_state_t *dd_tl_state_get(void) { + return dd_tl_state_get_fast(); +} + +/* + * Initialises TLS for this thread on first call; returns NULL on subsequent + * calls (state already exists). Callers that need the pointer regardless + * should use dd_tl_state_get_or_init(). + */ +dd_tl_state_t *dd_tl_state_init(void) { + if (dd_tl_state_storage.initialized) return NULL; + + tl_state_populate(&dd_tl_state_storage); + return &dd_tl_state_storage; +} + +dd_tl_state_t *dd_tl_state_get_or_init(void) { + return dd_tl_state_get_or_init_fast(); +} diff --git a/libdd-heap-sampler/vendor/README.md b/libdd-heap-sampler/vendor/README.md new file mode 100644 index 0000000000..dddbcd26df --- /dev/null +++ b/libdd-heap-sampler/vendor/README.md @@ -0,0 +1,32 @@ + +# Vendored headers + +## `usdt.h` + +The libbpf project's single-header USDT (User Statically-Defined +Tracepoint) library. Provides `USDT(group, name, args...)` and friends +that emit the standard v3 ELF-note format consumed by bpftrace, +systemtap, and any BPF-based tracer. + +We use it in preference to systemtap's `` because libbpf/usdt +is genuinely standalone: a single file with no `sdt-config.h` companion +and no other compile-time dependencies on the host system. + +- **Upstream:** +- **Path in upstream:** `usdt.h` (repo root) +- **License:** BSD-2-Clause. The SPDX header is retained verbatim at the + top of the file. Copyright (c) 2024 Meta Platforms, Inc. and affiliates. +- **Vendored from:** upstream `main` as of 2026-06-29. + +To refresh, fetch the latest copy and overwrite the file in place: + +```sh +curl -fsSL https://raw.githubusercontent.com/libbpf/usdt/main/usdt.h \ + > libdd-heap-sampler/vendor/usdt.h +``` + +Verify the BSD-2-Clause SPDX identifier is still on line 1 after +refreshing. diff --git a/libdd-heap-sampler/vendor/usdt.h b/libdd-heap-sampler/vendor/usdt.h new file mode 100644 index 0000000000..549d1f7748 --- /dev/null +++ b/libdd-heap-sampler/vendor/usdt.h @@ -0,0 +1,545 @@ +// SPDX-License-Identifier: BSD-2-Clause +/* + * This single-header library defines a collection of variadic macros for + * defining and triggering USDTs (User Statically-Defined Tracepoints): + * + * - For USDTs without associated semaphore: + * USDT(group, name, args...) + * + * - For USDTs with implicit (transparent to the user) semaphore: + * USDT_WITH_SEMA(group, name, args...) + * USDT_IS_ACTIVE(group, name) + * + * - For USDTs with explicit (user-defined and provided) semaphore: + * USDT_WITH_EXPLICIT_SEMA(sema, group, name, args...) + * USDT_SEMA_IS_ACTIVE(sema) + * + * all of which emit a NOP instruction into the instruction stream, and so + * have *zero* overhead for the surrounding code. USDTs are identified by + * a combination of `group` and `name` identifiers, which is used by external + * tracing tooling (tracers) for identifying exact USDTs of interest. + * + * USDTs can have an associated (2-byte) activity counter (USDT semaphore), + * automatically maintained by Linux kernel whenever any correctly written + * BPF-based tracer is attached to the USDT. This USDT semaphore can be used + * to check whether there is a need to do any extra data collection and + * processing for a given USDT (if necessary), and otherwise avoid extra work + * for a common case of USDT not being traced ("active"). + * + * See documentation for USDT_WITH_SEMA()/USDT_IS_ACTIVE() or + * USDT_WITH_EXPLICIT_SEMA()/USDT_SEMA_IS_ACTIVE() APIs below for details on + * working with USDTs with implicitly or explicitly associated + * USDT semaphores, respectively. + * + * There is also some additional data recorded into an auxiliary note + * section. The data in the note section describes the operands, in terms of + * size and location, used by tracing tooling to know where to find USDT + * arguments. Each location is encoded as an assembler operand string. + * Tracing tools (bpftrace and BPF-based tracers, systemtap, etc) insert + * breakpoints on top of the nop, and decode the location operand-strings, + * like an assembler, to find the values being passed. + * + * The operand strings are selected by the compiler for each operand. + * They are constrained by inline-assembler codes.The default is: + * + * #define USDT_ARG_CONSTRAINT nor + * + * This is a good default if the operands tend to be integral and + * moderate in number (smaller than number of registers). In other + * cases, the compiler may report "'asm' requires impossible reload" or + * similar. In this case, consider simplifying the macro call (fewer + * and simpler operands), reduce optimization, or override the default + * constraints string via: + * + * #define USDT_ARG_CONSTRAINT g + * #include + * + * For some historical description of USDT v3 format (the one used by this + * library and generally recognized and assumed by BPF-based tracing tools) + * see [0]. The more formal specification can be found at [1]. Additional + * argument constraints information can be found at [2]. + * + * Original SystemTap's sys/sdt.h implementation ([3]) was used as a base for + * this USDT library implementation. Current implementation differs *a lot* in + * terms of exposed user API and general usability, which was the main goal + * and focus of the reimplementation work. Nevertheless, underlying recorded + * USDT definitions are fully binary compatible and any USDT-based tooling + * should work equally well with USDTs defined by either SystemTap's or this + * library's USDT implementation. + * + * [0] https://ecos.sourceware.org/ml/systemtap/2010-q3/msg00145.html + * [1] https://sourceware.org/systemtap/wiki/UserSpaceProbeImplementation + * [2] https://gcc.gnu.org/onlinedocs/gcc/Constraints.html + * [3] https://sourceware.org/git/?p=systemtap.git;a=blob;f=includes/sys/sdt.h + */ +#ifndef __USDT_H +#define __USDT_H + +/* + * Changelog: + * + * 0.1.0 + * ----- + * - Initial release + */ +#define USDT_MAJOR_VERSION 0 +#define USDT_MINOR_VERSION 1 +#define USDT_PATCH_VERSION 0 + +/* C++20 and C23 added __VA_OPT__ as a standard replacement for non-standard `##__VA_ARGS__` extension */ +#if (defined(__STDC_VERSION__) && __STDC_VERSION__ > 201710L) || (defined(__cplusplus) && __cplusplus > 201703L) +#define __usdt_va_opt 1 +#define __usdt_va_args(...) __VA_OPT__(,) __VA_ARGS__ +#else +#define __usdt_va_args(...) , ##__VA_ARGS__ +#endif + +/* + * Trigger USDT with `group`:`name` identifier and pass through `args` as its + * arguments. Zero arguments are acceptable as well. No USDT semaphore is + * associated with this USDT. + * + * Such "semaphoreless" USDTs are commonly used when there is no extra data + * collection or processing needed to collect and prepare USDT arguments and + * they are just available in the surrounding code. USDT() macro will just + * record their locations in CPU registers or in memory for tracing tooling to + * be able to access them, if necessary. + */ +#ifdef __usdt_va_opt +#define USDT(group, name, ...) \ + __usdt_probe(group, name, __usdt_sema_none, 0 __VA_OPT__(,) __VA_ARGS__) +#else +#define USDT(group, name, ...) \ + __usdt_probe(group, name, __usdt_sema_none, 0, ##__VA_ARGS__) +#endif + +/* + * Trigger USDT with `group`:`name` identifier and pass through `args` as its + * arguments. Zero arguments are acceptable as well. USDT also get an + * implicitly-defined associated USDT semaphore, which will be "activated" by + * tracing tooling and can be used to check whether USDT is being actively + * observed. + * + * USDTs with semaphore are commonly used when there is a need to perform + * additional data collection and processing to prepare USDT arguments, which + * otherwise might not be necessary for the rest of application logic. In such + * case, USDT semaphore can be used to avoid unnecessary extra work. If USDT + * is not traced (which is presumed to be a common situation), the associated + * USDT semaphore is "inactive", and so there is no need to waste resources to + * prepare USDT arguments. Use USDT_IS_ACTIVE(group, name) to check whether + * USDT is "active". + * + * N.B. There is an inherent (albeit short) gap between checking whether USDT + * is active and triggering corresponding USDT, in which external tracer can + * be attached to an USDT and activate USDT semaphore after the activity check. + * If such a race occurs, tracers might miss one USDT execution. Tracers are + * expected to accommodate such possibility and this is expected to not be + * a problem for applications and tracers. + * + * N.B. Implicit USDT semaphore defined by USDT_WITH_SEMA() is contained + * within a single executable or shared library and is not shared outside + * them. I.e., if you use USDT_WITH_SEMA() with the same USDT group and name + * identifier across executable and shared library, it will work and won't + * conflict, per se, but will define independent USDT semaphores, one for each + * shared library/executable in which USDT_WITH_SEMA(group, name) is used. + * That is, if you attach to this USDT in one shared library (or executable), + * then only USDT semaphore within that shared library (or executable) will be + * updated by the kernel, while other libraries (or executable) will not see + * activated USDT semaphore. In short, it's best to use unique USDT group:name + * identifiers across different shared libraries (and, equivalently, between + * executable and shared library). This is advanced consideration and is + * rarely (if ever) seen in practice, but just to avoid surprises this is + * called out here. (Static libraries become a part of final executable, once + * linked by linker, so the above considerations don't apply to them.) + */ +#ifdef __usdt_va_opt +#define USDT_WITH_SEMA(group, name, ...) \ + __usdt_probe(group, name, \ + __usdt_sema_implicit, __usdt_sema_name(group, name) \ + __VA_OPT__(,) __VA_ARGS__) +#else +#define USDT_WITH_SEMA(group, name, ...) \ + __usdt_probe(group, name, \ + __usdt_sema_implicit, __usdt_sema_name(group, name), \ + ##__VA_ARGS__) +#endif + +struct usdt_sema { volatile unsigned short active; }; + +/* + * Check if USDT with `group`:`name` identifier is "active" (i.e., whether it + * is attached to by external tracing tooling and is actively observed). + * + * This macro can be used to decide whether any additional and potentially + * expensive data collection or processing should be done to pass extra + * information into the given USDT. It is assumed that USDT is triggered with + * USDT_WITH_SEMA() macro which will implicitly define associated USDT + * semaphore. (If one needs more control over USDT semaphore, see + * USDT_DEFINE_SEMA() and USDT_WITH_EXPLICIT_SEMA() macros below.) + * + * N.B. Such checks are necessarily racy and speculative. Between checking + * whether USDT is active and triggering the USDT itself, tracer can be + * detached with no notification. This race should be extremely rare and worst + * case should result in one-time wasted extra data collection and processing. + */ +#define USDT_IS_ACTIVE(group, name) ({ \ + extern struct usdt_sema __usdt_sema_name(group, name) \ + __usdt_asm_name(__usdt_sema_name(group, name)); \ + __usdt_sema_implicit(__usdt_sema_name(group, name)); \ + __usdt_sema_name(group, name).active > 0; \ +}) + +/* + * APIs for working with user-defined explicit USDT semaphores. + * + * This is a less commonly used advanced API for use cases in which user needs + * an explicit control over (potentially shared across multiple USDTs) USDT + * semaphore instance. This can be used when there is a group of logically + * related USDTs that all need extra data collection and processing whenever + * any of a family of related USDTs are "activated" (i.e., traced). In such + * a case, all such related USDTs will be associated with the same shared USDT + * semaphore defined with USDT_DEFINE_SEMA() and the USDTs themselves will be + * triggered with USDT_WITH_EXPLICIT_SEMA() macros, taking an explicit extra + * USDT semaphore identifier as an extra parameter. + */ + +/** + * Underlying C global variable name for user-defined USDT semaphore with + * `sema` identifier. Could be useful for debugging, but normally shouldn't be + * used explicitly. + */ +#define USDT_SEMA(sema) __usdt_sema_##sema + +/* + * Define storage for user-defined USDT semaphore `sema`. + * + * Should be used only once in non-header source file to let compiler allocate + * space for the semaphore variable. Just like with any other global variable. + * + * This macro can be used anywhere where global variable declaration is + * allowed. Just like with global variable definitions, there should be only + * one definition of user-defined USDT semaphore with given `sema` identifier, + * otherwise compiler or linker will complain about duplicate variable + * definition. + * + * For C++, it is allowed to use USDT_DEFINE_SEMA() both in global namespace + * and inside namespaces (including nested namespaces). Just make sure that + * USDT_DECLARE_SEMA() is placed within the namespace where this semaphore is + * referenced, or any of its parent namespaces, so the C++ language-level + * identifier is visible to the code that needs to reference the semaphore. + * At the lowest layer, USDT semaphores have global naming and visibility + * (they have a corresponding `__usdt_sema_` symbol, which can be linked + * against from C or C++ code, if necessary). To keep it simple, putting + * USDT_DECLARE_SEMA() declarations into global namespaces is the simplest + * no-brainer solution. All these aspects are irrelevant for plain C, because + * C doesn't have namespaces and everything is always in the global namespace. + * + * N.B. Due to USDT metadata being recorded in non-allocatable ELF note + * section, it has limitations when it comes to relocations, which, in + * practice, means that it's not possible to correctly share USDT semaphores + * between main executable and shared libraries, or even between multiple + * shared libraries. USDT semaphore has to be contained to individual shared + * library or executable to avoid unpleasant surprises with half-working USDT + * semaphores. We enforce this by marking semaphore ELF symbols as having + * a hidden visibility. This is quite an advanced use case and consideration + * and for most users this should have no consequences whatsoever. + */ +#define USDT_DEFINE_SEMA(sema) \ + struct usdt_sema __usdt_sema_sec USDT_SEMA(sema) \ + __usdt_asm_name(USDT_SEMA(sema)) \ + __attribute__((visibility("hidden"))) = { 0 } + +/* + * Declare extern reference to user-defined USDT semaphore `sema`. + * + * Refers to a variable defined in another compilation unit by + * USDT_DEFINE_SEMA() and allows to use the same USDT semaphore across + * multiple compilation units (i.e., .c and .cpp files). + * + * See USDT_DEFINE_SEMA() notes above for C++ language usage peculiarities. + */ +#define USDT_DECLARE_SEMA(sema) \ + extern struct usdt_sema USDT_SEMA(sema) __usdt_asm_name(USDT_SEMA(sema)) + +/* + * Check if user-defined USDT semaphore `sema` is "active" (i.e., whether it + * is attached to by external tracing tooling and is actively observed). + * + * This macro can be used to decide whether any additional and potentially + * expensive data collection or processing should be done to pass extra + * information into USDT(s) associated with USDT semaphore `sema`. + * + * N.B. Such checks are necessarily racy. Between checking the state of USDT + * semaphore and triggering associated USDT(s), the active tracer might attach + * or detach. This race should be extremely rare and worst case should result + * in one-time missed USDT event or wasted extra data collection and + * processing. USDT-using tracers should be written with this in mind and is + * not a concern of the application defining USDTs with associated semaphore. + */ +#define USDT_SEMA_IS_ACTIVE(sema) (USDT_SEMA(sema).active > 0) + +/* + * Invoke USDT specified by `group` and `name` identifiers and associate + * explicitly user-defined semaphore `sema` with it. Pass through `args` as + * USDT arguments. `args` are optional and zero arguments are acceptable. + * + * Semaphore is defined with the help of USDT_DEFINE_SEMA() macro and can be + * checked whether active with USDT_SEMA_IS_ACTIVE(). + */ +#ifdef __usdt_va_opt +#define USDT_WITH_EXPLICIT_SEMA(sema, group, name, ...) \ + __usdt_probe(group, name, __usdt_sema_explicit, USDT_SEMA(sema), ##__VA_ARGS__) +#else +#define USDT_WITH_EXPLICIT_SEMA(sema, group, name, ...) \ + __usdt_probe(group, name, __usdt_sema_explicit, USDT_SEMA(sema) __VA_OPT__(,) __VA_ARGS__) +#endif + +/* + * Adjustable implementation aspects + */ +#ifndef USDT_ARG_CONSTRAINT +#if defined __powerpc__ +#define USDT_ARG_CONSTRAINT nZr +#elif defined __arm__ +#define USDT_ARG_CONSTRAINT g +#elif defined __loongarch__ +#define USDT_ARG_CONSTRAINT nmr +#else +#define USDT_ARG_CONSTRAINT nor +#endif +#endif /* USDT_ARG_CONSTRAINT */ + +#ifndef USDT_NOP +#if defined(__ia64__) || defined(__s390__) || defined(__s390x__) +#define USDT_NOP nop 0 +#else +#define USDT_NOP nop +#endif +#endif /* USDT_NOP */ + +/* + * Implementation details + */ +/* USDT name for implicitly-defined USDT semaphore, derived from group:name */ +#define __usdt_sema_name(group, name) __usdt_sema_##group##__##name +/* ELF section into which USDT semaphores are put */ +#define __usdt_sema_sec __attribute__((section(".probes"))) + +#define __usdt_concat(a, b) a ## b +#define __usdt_apply(fn, n) __usdt_concat(fn, n) + +#ifndef __usdt_nth +#define __usdt_nth(_, _1, _2, _3, _4, _5, _6, _7, _8, _9, _10, _11, _12, N, ...) N +#endif + +#ifndef __usdt_narg +#ifdef __usdt_va_opt +#define __usdt_narg(...) __usdt_nth(_ __VA_OPT__(,) __VA_ARGS__, 12, 11, 10, 9, 8, 7, 6, 5, 4, 3, 2, 1, 0) +#else +#define __usdt_narg(...) __usdt_nth(_, ##__VA_ARGS__, 12, 11, 10, 9, 8, 7, 6, 5, 4, 3, 2, 1, 0) +#endif +#endif /* __usdt_narg */ + +#define __usdt_hash # +#define __usdt_str_(x) #x +#define __usdt_str(x) __usdt_str_(x) + +#ifndef __usdt_asm_name +#define __usdt_asm_name(name) __asm__(__usdt_str(name)) +#endif + +#define __usdt_asm0() "\n" +#define __usdt_asm1(x) __usdt_str(x) "\n" +#define __usdt_asm2(x, ...) __usdt_str(x) "," __usdt_asm1(__VA_ARGS__) +#define __usdt_asm3(x, ...) __usdt_str(x) "," __usdt_asm2(__VA_ARGS__) +#define __usdt_asm4(x, ...) __usdt_str(x) "," __usdt_asm3(__VA_ARGS__) +#define __usdt_asm5(x, ...) __usdt_str(x) "," __usdt_asm4(__VA_ARGS__) +#define __usdt_asm6(x, ...) __usdt_str(x) "," __usdt_asm5(__VA_ARGS__) +#define __usdt_asm7(x, ...) __usdt_str(x) "," __usdt_asm6(__VA_ARGS__) +#define __usdt_asm8(x, ...) __usdt_str(x) "," __usdt_asm7(__VA_ARGS__) +#define __usdt_asm9(x, ...) __usdt_str(x) "," __usdt_asm8(__VA_ARGS__) +#define __usdt_asm10(x, ...) __usdt_str(x) "," __usdt_asm9(__VA_ARGS__) +#define __usdt_asm11(x, ...) __usdt_str(x) "," __usdt_asm10(__VA_ARGS__) +#define __usdt_asm12(x, ...) __usdt_str(x) "," __usdt_asm11(__VA_ARGS__) +#define __usdt_asm(...) __usdt_apply(__usdt_asm, __usdt_narg(__VA_ARGS__))(__VA_ARGS__) + +#ifdef __LP64__ +#define __usdt_asm_addr .8byte +#else +#define __usdt_asm_addr .4byte +#endif + +#define __usdt_asm_strz_(x) __usdt_asm1(.asciz #x) +#define __usdt_asm_strz(x) __usdt_asm_strz_(x) +#define __usdt_asm_str_(x) __usdt_asm1(.ascii #x) +#define __usdt_asm_str(x) __usdt_asm_str_(x) + +/* "semaphoreless" USDT case */ +#ifndef __usdt_sema_none +#define __usdt_sema_none(sema) +#endif + +/* implicitly defined __usdt_sema__group__name semaphore (using weak symbols) */ +#ifndef __usdt_sema_implicit +#define __usdt_sema_implicit(sema) \ + __asm__ __volatile__ ( \ + __usdt_asm1(.ifndef sema) \ + __usdt_asm3( .pushsection .probes, "aw", "progbits") \ + __usdt_asm1( .weak sema) \ + __usdt_asm1( .hidden sema) \ + __usdt_asm1( .align 2) \ + __usdt_asm1(sema:) \ + __usdt_asm1( .zero 2) \ + __usdt_asm2( .type sema, @object) \ + __usdt_asm2( .size sema, 2) \ + __usdt_asm1( .popsection) \ + __usdt_asm1(.endif) \ + ); +#endif + +/* externally defined semaphore using USDT_DEFINE_SEMA() and passed explicitly by user */ +#ifndef __usdt_sema_explicit +#define __usdt_sema_explicit(sema) \ + __asm__ __volatile__ ("" :: "m" (sema)); +#endif + +/* main USDT definition (nop and .note.stapsdt metadata) */ +#define __usdt_probe(group, name, sema_def, sema, ...) do { \ + sema_def(sema) \ + __asm__ __volatile__ ( \ + __usdt_asm( 990: USDT_NOP) \ + __usdt_asm3( .pushsection .note.stapsdt, "", "note") \ + __usdt_asm1( .balign 4) \ + __usdt_asm3( .4byte 992f-991f,994f-993f,3) \ + __usdt_asm1(991: .asciz "stapsdt") \ + __usdt_asm1(992: .balign 4) \ + __usdt_asm1(993: __usdt_asm_addr 990b) \ + __usdt_asm1( __usdt_asm_addr _.stapsdt.base) \ + __usdt_asm1( __usdt_asm_addr sema) \ + __usdt_asm_strz(group) \ + __usdt_asm_strz(name) \ + __usdt_asm_args(__VA_ARGS__) \ + __usdt_asm1( .ascii "\0") \ + __usdt_asm1(994: .balign 4) \ + __usdt_asm1( .popsection) \ + __usdt_asm1(.ifndef _.stapsdt.base) \ + __usdt_asm5( .pushsection .stapsdt.base,"aG","progbits",.stapsdt.base,comdat)\ + __usdt_asm1( .weak _.stapsdt.base) \ + __usdt_asm1( .hidden _.stapsdt.base) \ + __usdt_asm1(_.stapsdt.base:) \ + __usdt_asm1( .space 1) \ + __usdt_asm2( .size _.stapsdt.base, 1) \ + __usdt_asm1( .popsection) \ + __usdt_asm1(.endif) \ + :: __usdt_asm_ops(__VA_ARGS__) \ + ); \ +} while (0) + +/* + * NB: gdb PR24541 highlighted an unspecified corner of the sdt.h + * operand note format. + * + * The named register may be a longer or shorter (!) alias for the + * storage where the value in question is found. For example, on + * i386, 64-bit value may be put in register pairs, and a register + * name stored would identify just one of them. Previously, gcc was + * asked to emit the %w[id] (16-bit alias of some registers holding + * operands), even when a wider 32-bit value was used. + * + * Bottom line: the byte-width given before the @ sign governs. If + * there is a mismatch between that width and that of the named + * register, then a sys/sdt.h note consumer may need to employ + * architecture-specific heuristics to figure out where the compiler + * has actually put the complete value. + */ +#if defined(__powerpc__) || defined(__powerpc64__) +#define __usdt_argref(id) %I[id]%[id] +#elif defined(__i386__) +#define __usdt_argref(id) %k[id] /* gcc.gnu.org/PR80115 sourceware.org/PR24541 */ +#else +#define __usdt_argref(id) %[id] +#endif + +#define __usdt_asm_arg(n) __usdt_asm_str(%c[__usdt_asz##n]) \ + __usdt_asm1(.ascii "@") \ + __usdt_asm_str(__usdt_argref(__usdt_aval##n)) + +#define __usdt_asm_args0 /* no arguments */ +#define __usdt_asm_args1 __usdt_asm_arg(1) +#define __usdt_asm_args2 __usdt_asm_args1 __usdt_asm1(.ascii " ") __usdt_asm_arg(2) +#define __usdt_asm_args3 __usdt_asm_args2 __usdt_asm1(.ascii " ") __usdt_asm_arg(3) +#define __usdt_asm_args4 __usdt_asm_args3 __usdt_asm1(.ascii " ") __usdt_asm_arg(4) +#define __usdt_asm_args5 __usdt_asm_args4 __usdt_asm1(.ascii " ") __usdt_asm_arg(5) +#define __usdt_asm_args6 __usdt_asm_args5 __usdt_asm1(.ascii " ") __usdt_asm_arg(6) +#define __usdt_asm_args7 __usdt_asm_args6 __usdt_asm1(.ascii " ") __usdt_asm_arg(7) +#define __usdt_asm_args8 __usdt_asm_args7 __usdt_asm1(.ascii " ") __usdt_asm_arg(8) +#define __usdt_asm_args9 __usdt_asm_args8 __usdt_asm1(.ascii " ") __usdt_asm_arg(9) +#define __usdt_asm_args10 __usdt_asm_args9 __usdt_asm1(.ascii " ") __usdt_asm_arg(10) +#define __usdt_asm_args11 __usdt_asm_args10 __usdt_asm1(.ascii " ") __usdt_asm_arg(11) +#define __usdt_asm_args12 __usdt_asm_args11 __usdt_asm1(.ascii " ") __usdt_asm_arg(12) +#define __usdt_asm_args(...) __usdt_apply(__usdt_asm_args, __usdt_narg(__VA_ARGS__)) + +#define __usdt_is_arr(x) (__builtin_classify_type(x) == 14 || __builtin_classify_type(x) == 5) +#define __usdt_arg_size(x) (__usdt_is_arr(x) ? sizeof(void *) : sizeof(x)) + +/* + * We can't use __builtin_choose_expr() in C++, so fall back to table-based + * signedness determination for known types, utilizing templates magic. + */ +#ifdef __cplusplus + +#define __usdt_is_signed(x) (!__usdt_is_arr(x) && __usdt_t<__typeof(x)>::is_signed) + +#include + +template struct __usdt_t { static const bool is_signed = false; }; +template struct __usdt_t : public __usdt_t {}; +template struct __usdt_t : public __usdt_t {}; + +#define __usdt_def_signed(T) \ +template<> struct __usdt_t { static const bool is_signed = true; }; \ +template<> struct __usdt_t { static const bool is_signed = true; }; \ +template<> struct __usdt_t { static const bool is_signed = true; }; \ +template<> struct __usdt_t { static const bool is_signed = true; } +#define __usdt_maybe_signed(T) \ +template<> struct __usdt_t { static const bool is_signed = (T)-1 < (T)1; }; \ +template<> struct __usdt_t { static const bool is_signed = (T)-1 < (T)1; }; \ +template<> struct __usdt_t { static const bool is_signed = (T)-1 < (T)1; }; \ +template<> struct __usdt_t { static const bool is_signed = (T)-1 < (T)1; } + +__usdt_def_signed(signed char); +__usdt_def_signed(short); +__usdt_def_signed(int); +__usdt_def_signed(long); +__usdt_def_signed(long long); +__usdt_maybe_signed(char); +__usdt_maybe_signed(wchar_t); + +#else /* !__cplusplus */ + +#define __usdt_is_inttype(x) (__builtin_classify_type(x) >= 1 && __builtin_classify_type(x) <= 4) +#define __usdt_inttype(x) __typeof(__builtin_choose_expr(__usdt_is_inttype(x), (x), 0U)) +#define __usdt_is_signed(x) ((__usdt_inttype(x))-1 < (__usdt_inttype(x))1) + +#endif /* __cplusplus */ + +#define __usdt_asm_op(n, x) \ + [__usdt_asz##n] "n" ((__usdt_is_signed(x) ? (int)-1 : 1) * (int)__usdt_arg_size(x)), \ + [__usdt_aval##n] __usdt_str(USDT_ARG_CONSTRAINT)(x) + +#define __usdt_asm_ops0() [__usdt_dummy] "g" (0) +#define __usdt_asm_ops1(x) __usdt_asm_op(1, x) +#define __usdt_asm_ops2(a,x) __usdt_asm_ops1(a), __usdt_asm_op(2, x) +#define __usdt_asm_ops3(a,b,x) __usdt_asm_ops2(a,b), __usdt_asm_op(3, x) +#define __usdt_asm_ops4(a,b,c,x) __usdt_asm_ops3(a,b,c), __usdt_asm_op(4, x) +#define __usdt_asm_ops5(a,b,c,d,x) __usdt_asm_ops4(a,b,c,d), __usdt_asm_op(5, x) +#define __usdt_asm_ops6(a,b,c,d,e,x) __usdt_asm_ops5(a,b,c,d,e), __usdt_asm_op(6, x) +#define __usdt_asm_ops7(a,b,c,d,e,f,x) __usdt_asm_ops6(a,b,c,d,e,f), __usdt_asm_op(7, x) +#define __usdt_asm_ops8(a,b,c,d,e,f,g,x) __usdt_asm_ops7(a,b,c,d,e,f,g), __usdt_asm_op(8, x) +#define __usdt_asm_ops9(a,b,c,d,e,f,g,h,x) __usdt_asm_ops8(a,b,c,d,e,f,g,h), __usdt_asm_op(9, x) +#define __usdt_asm_ops10(a,b,c,d,e,f,g,h,i,x) __usdt_asm_ops9(a,b,c,d,e,f,g,h,i), __usdt_asm_op(10, x) +#define __usdt_asm_ops11(a,b,c,d,e,f,g,h,i,j,x) __usdt_asm_ops10(a,b,c,d,e,f,g,h,i,j), __usdt_asm_op(11, x) +#define __usdt_asm_ops12(a,b,c,d,e,f,g,h,i,j,k,x) __usdt_asm_ops11(a,b,c,d,e,f,g,h,i,j,k), __usdt_asm_op(12, x) +#define __usdt_asm_ops(...) __usdt_apply(__usdt_asm_ops, __usdt_narg(__VA_ARGS__))(__VA_ARGS__) + +#endif /* __USDT_H */ diff --git a/tools/docker/Dockerfile.build b/tools/docker/Dockerfile.build index 31fa97614f..8b51661fd5 100644 --- a/tools/docker/Dockerfile.build +++ b/tools/docker/Dockerfile.build @@ -129,6 +129,10 @@ COPY "datadog-ffe-test-suite/Cargo.toml" "datadog-ffe-test-suite/" COPY "datadog-ffe-ffi/Cargo.toml" "datadog-ffe-ffi/" COPY "libdd-otel-thread-ctx/Cargo.toml" "libdd-otel-thread-ctx/" COPY "libdd-otel-thread-ctx-ffi/Cargo.toml" "libdd-otel-thread-ctx-ffi/" +COPY "libdd-heap-sampler/Cargo.toml" "libdd-heap-sampler/" +COPY "libdd-heap-allocator/Cargo.toml" "libdd-heap-allocator/" +COPY "libdd-heap-gotter/Cargo.toml" "libdd-heap-gotter/" +COPY "libdd-heap-gotter-ffi/Cargo.toml" "libdd-heap-gotter-ffi/" RUN find -name "Cargo.toml" | sed -e s#Cargo.toml#src/lib.rs#g | xargs -n 1 sh -c 'mkdir -p $(dirname $1); touch $1; echo $1' create_stubs RUN echo \ bin_tests/src/bin/crashtracker_bin_test.rs \ @@ -160,6 +164,7 @@ RUN echo \ libdd-trace-utils/benches/deserialization.rs \ libdd-trace-utils/benches/main.rs \ libdd-data-pipeline/benches/trace_buffer.rs \ + libdd-heap-allocator/benches/sampler_overhead.rs \ | xargs -n 1 sh -c 'mkdir -p $(dirname $1); touch $1; echo $1' create_stubs # cache dependencies From 13be13c157c49ca160d588a210618cb766ad8ca5 Mon Sep 17 00:00:00 2001 From: Scott Gerring Date: Wed, 1 Jul 2026 11:55:40 +0200 Subject: [PATCH 02/16] address review comments --- libdd-heap-gotter/src/elf.rs | 203 +++++++++++++++--- libdd-heap-gotter/src/hooks.rs | 121 +++++++++-- libdd-heap-gotter/src/lib.rs | 12 +- libdd-heap-gotter/src/realloc_math.rs | 36 ++++ libdd-heap-sampler/build.rs | 4 + .../datadog/heap/allocation_requested.h | 48 +++-- .../include/datadog/heap/sample_flag.h | 166 ++++++++------ libdd-heap-sampler/src/allocation_created.c | 7 +- libdd-heap-sampler/src/allocation_freed.c | 32 ++- libdd-heap-sampler/src/allocation_requested.c | 77 ++++++- libdd-heap-sampler/src/generated/bindings.rs | 25 ++- libdd-heap-sampler/src/lib.rs | 4 +- libdd-heap-sampler/src/sample_flag.c | 75 +++++-- libdd-heap-sampler/src/tl_state.c | 13 +- 14 files changed, 644 insertions(+), 179 deletions(-) create mode 100644 libdd-heap-gotter/src/realloc_math.rs diff --git a/libdd-heap-gotter/src/elf.rs b/libdd-heap-gotter/src/elf.rs index b07020de53..bf30a6c268 100644 --- a/libdd-heap-gotter/src/elf.rs +++ b/libdd-heap-gotter/src/elf.rs @@ -23,8 +23,10 @@ use std::ffi::CStr; use libc::{ dl_iterate_phdr, dl_phdr_info, mprotect, sysconf, Elf64_Phdr, Elf64_Rel, Elf64_Rela, Elf64_Sym, - PROT_READ, PROT_WRITE, PT_DYNAMIC, PT_LOAD, _SC_PAGESIZE, + PROT_EXEC, PROT_READ, PROT_WRITE, PT_DYNAMIC, PT_LOAD, _SC_PAGESIZE, }; +use std::io::{BufRead, BufReader}; +use std::sync::atomic::{AtomicUsize, Ordering}; // ELF dynamic-section tags and friends. The `libc` crate doesn't export // these (they're processor-independent ELF spec constants), so we name @@ -282,15 +284,117 @@ fn iterate_libraries bool>(mut cb: F) { } } -/// Temporarily make the containing page writable and replace one GOT entry. -unsafe fn override_entry(addr: usize, new_value: usize) -> bool { - let page = sysconf(_SC_PAGESIZE) as usize; - let aligned = (addr & !(page - 1)) as *mut c_void; - if mprotect(aligned, page, PROT_READ | PROT_WRITE) != 0 { - return false; +/// A single /proc/self/maps entry: address range + current protection flags. +#[derive(Clone, Copy)] +struct MapEntry { + start: usize, + end: usize, + prot: i32, +} + +/// Parse /proc/self/maps into a sorted list of (range, prot) entries. +/// +/// Used to remember each GOT page's original protection so we can restore +/// it after patching, rather than leaving Full-RELRO pages read-write for +/// the lifetime of the process. +fn read_proc_maps() -> Vec { + let Ok(f) = std::fs::File::open("/proc/self/maps") else { + return Vec::new(); + }; + let mut out = Vec::new(); + for line in BufReader::new(f).lines().map_while(Result::ok) { + let mut parts = line.split_whitespace(); + let Some(range) = parts.next() else { continue }; + let Some(perms) = parts.next() else { continue }; + let Some(dash) = range.find('-') else { + continue; + }; + let Ok(start) = usize::from_str_radix(&range[..dash], 16) else { + continue; + }; + let Ok(end) = usize::from_str_radix(&range[dash + 1..], 16) else { + continue; + }; + let b = perms.as_bytes(); + let mut prot = 0; + if b.first() == Some(&b'r') { + prot |= PROT_READ; + } + if b.get(1) == Some(&b'w') { + prot |= PROT_WRITE; + } + if b.get(2) == Some(&b'x') { + prot |= PROT_EXEC; + } + out.push(MapEntry { start, end, prot }); + } + out +} + +/// Batched GOT-entry patcher that remembers each touched page's +/// original protection and restores it at the end of a patching pass. +/// +/// On Full-RELRO binaries, GOT pages start read-only. The old +/// `override_entry` helper flipped them to read-write and never +/// restored them, weakening RELRO for the process lifetime. This guard +/// mprotects each unique page once (RW), lets the caller write as +/// many entries as it needs, then mprotects each page back to what +/// /proc/self/maps reported at guard-construction time. +struct PageProtGuard { + page_size: usize, + maps: Vec, + // Aligned page base -> original prot flags read from /proc/self/maps. + touched: HashMap, +} + +impl PageProtGuard { + fn new() -> Self { + Self { + page_size: unsafe { sysconf(_SC_PAGESIZE) as usize }, + maps: read_proc_maps(), + touched: HashMap::new(), + } + } + + fn original_prot(&self, addr: usize) -> Option { + self.maps + .iter() + .find(|m| addr >= m.start && addr < m.end) + .map(|m| m.prot) + } + + /// Make the containing page writable if it isn't already touched, + /// then replace one GOT entry. + unsafe fn override_entry(&mut self, addr: usize, new_value: usize) -> bool { + let aligned = addr & !(self.page_size - 1); + if !self.touched.contains_key(&aligned) { + // If /proc/self/maps isn't available (or the page isn't in + // it, which shouldn't happen for a mapped GOT page) fall + // back to PROT_READ - the RELRO'd default. That's tighter + // than the previous behavior of leaving pages RW. + let orig = self.original_prot(aligned).unwrap_or(PROT_READ); + if mprotect( + aligned as *mut c_void, + self.page_size, + PROT_READ | PROT_WRITE, + ) != 0 + { + return false; + } + self.touched.insert(aligned, orig); + } + std::ptr::write_unaligned(addr as *mut usize, new_value); + true + } + + /// Restore every touched page to its original protection. + fn finish(mut self) { + for (aligned, orig) in self.touched.drain() { + // Best-effort: nothing sensible to do on failure other than + // leave the page RW, which is the pre-fix behavior. + unsafe { mprotect(aligned as *mut c_void, self.page_size, orig) }; + } } - std::ptr::write_unaligned(addr as *mut usize, new_value); - true } /// Read one GOT entry without assuming pointer alignment. @@ -343,9 +447,9 @@ pub fn lookup_symbol(name: &str, not_this_symbol: usize) -> Option /// Per-library revert info: GOT addr -> original value at that addr. #[derive(Default)] struct LibraryRevertInfo { - /// Kept for diagnostics / future logging; we identify libraries by - /// base address elsewhere. - #[allow(dead_code)] + /// 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. library_name: String, old_value_per_address: HashMap, processed: bool, @@ -354,9 +458,11 @@ struct LibraryRevertInfo { /// One registered override entry. struct OverrideInfo { /// Output slot the install path fills with the resolved real symbol - /// address (so hooks can call through it). `AtomicUsize` is what the - /// hook side sees; we just store its address here. - ref_slot: *mut usize, + /// address (so hooks can call through it). Points to a `'static + /// AtomicUsize` supplied by the caller; the install-time write goes + /// through `store(Release)` to pair with the hook-side + /// `load(Acquire)`. + ref_slot: *const AtomicUsize, /// Address of our hook function, written into matching GOT entries. new_symbol: usize, /// If a GOT entry's address equals this, leave it alone. Used to @@ -390,8 +496,9 @@ impl SymbolOverrides { /// Register an override. `ref_slot` is filled in by `apply_overrides` /// with the resolved address of the real symbol so the hook can call - /// through it. - pub fn register(&mut self, name: &str, hook: usize, ref_slot: *mut usize) { + /// through it. Must point to a `'static AtomicUsize`; the install + /// path publishes via `store(Release)`. + pub fn register(&mut self, name: &str, hook: usize, ref_slot: *const AtomicUsize) { self.overrides.insert( name.to_string(), OverrideInfo { @@ -421,7 +528,11 @@ impl SymbolOverrides { .collect(); for (name, addr) in resolved { if let Some(ov) = self.overrides.get_mut(&name) { - unsafe { *ov.ref_slot = addr }; + // SAFETY: `ref_slot` was registered pointing at a + // `'static AtomicUsize`, so the pointer is valid for the + // lifetime of the process. Release pairs with the + // hook-side Acquire load. + unsafe { (*ov.ref_slot).store(addr, Ordering::Release) }; } } self.update_overrides(); @@ -446,10 +557,14 @@ impl SymbolOverrides { v.processed = false; } + let mut guard = PageProtGuard::new(); + // SAFETY: closure runs synchronously inside dl_iterate_phdr. let self_ptr = self as *mut Self as usize; + let guard_ptr = &mut guard as *mut PageProtGuard as usize; iterate_libraries(move |info, _is_exe| unsafe { let this = &mut *(self_ptr as *mut Self); + let g = &mut *(guard_ptr as *mut PageProtGuard); let lib_name = if info.dlpi_name.is_null() { String::new() } else { @@ -461,11 +576,13 @@ impl SymbolOverrides { return false; } if let Some(dyn_info) = DynamicInfo::from_phdr(info) { - this.apply_to_library(&dyn_info, lib_name); + this.apply_to_library(&dyn_info, lib_name, g); } false }); + guard.finish(); + // Drop any tracked libraries that have been unloaded. self.revert_info_per_library.retain(|_, v| v.processed); } @@ -473,29 +590,51 @@ impl SymbolOverrides { /// Restore every GOT entry we touched. pub fn restore_overrides(&mut self) { let info_per_lib = std::mem::take(&mut self.revert_info_per_library); + let mut guard = PageProtGuard::new(); for (_base, revert) in info_per_lib { - unsafe { - for (addr, old) in revert.old_value_per_address { - override_entry(addr, old); - } + for (addr, old) in revert.old_value_per_address { + unsafe { guard.override_entry(addr, old) }; } } + guard.finish(); self.last_seen_nb_libs = -1; } - unsafe fn apply_to_library(&mut self, dyn_info: &DynamicInfo, library_name: String) { - let (entry_is_new, _) = match self.revert_info_per_library.entry(dyn_info.base_address) { + unsafe fn apply_to_library( + &mut self, + dyn_info: &DynamicInfo, + library_name: String, + guard: &mut PageProtGuard, + ) { + // Detect base-address reuse: a previous `dlclose` may have freed + // the load address, and a later `dlopen` can place a different + // library at the same address. If the name differs from what we + // recorded, the stored `old_value_per_address` refers to memory + // that is either unmapped or now belongs to the new library, so + // we must not try to restore it. Drop the stale entry and treat + // this as a fresh library. + let entry_is_new = match self.revert_info_per_library.entry(dyn_info.base_address) { std::collections::hash_map::Entry::Vacant(e) => { e.insert(LibraryRevertInfo { library_name: library_name.clone(), processed: true, ..Default::default() }); - (true, ()) + true } std::collections::hash_map::Entry::Occupied(mut e) => { - e.get_mut().processed = true; - (false, ()) + if e.get().library_name != library_name { + // Base-address reuse: replace the stale entry. + e.insert(LibraryRevertInfo { + library_name: library_name.clone(), + processed: true, + ..Default::default() + }); + true + } else { + e.get_mut().processed = true; + false + } } }; if !entry_is_new { @@ -517,6 +656,7 @@ impl SymbolOverrides { elf64_r_sym(reloc.r_info) as u32, reloc.r_offset as usize, revert, + guard, ); } } @@ -535,6 +675,7 @@ impl SymbolOverrides { elf64_r_sym(reloc.r_info) as u32, reloc.r_offset as usize, revert, + guard, ); } } @@ -546,6 +687,7 @@ impl SymbolOverrides { sym_index: u32, r_offset: usize, revert: &mut LibraryRevertInfo, + guard: &mut PageProtGuard, ) { // st_name -> string in strtab. Walk lazily: we look up the // name in the override map; if it's not there, skip. @@ -562,7 +704,8 @@ impl SymbolOverrides { }; // `ref_slot==0` means we never resolved the real symbol, so the // hook would call a NULL pointer. Skip. - let real = unsafe { *ov.ref_slot }; + // SAFETY: `ref_slot` points at a `'static AtomicUsize`. + let real = unsafe { (*ov.ref_slot).load(Ordering::Acquire) }; if real == 0 { return; } @@ -575,7 +718,7 @@ impl SymbolOverrides { return; } revert.old_value_per_address.insert(addr, read_entry(addr)); - override_entry(addr, ov.new_symbol); + guard.override_entry(addr, ov.new_symbol); } } diff --git a/libdd-heap-gotter/src/hooks.rs b/libdd-heap-gotter/src/hooks.rs index 68b0d267d9..d6fae01345 100644 --- a/libdd-heap-gotter/src/hooks.rs +++ b/libdd-heap-gotter/src/hooks.rs @@ -20,9 +20,11 @@ use std::sync::atomic::{AtomicUsize, Ordering}; use libdd_heap_sampler::{ dd_alloc_req_t, dd_allocation_created, dd_allocation_freed, dd_allocation_requested, - dd_tl_state_get, dd_tl_state_init, + dd_probe_free, dd_sample_flag_peek, dd_tl_state_get, dd_tl_state_init, }; +use crate::realloc_math::sampled_realloc_raw_size; + // Per-thread reentry guard for the gotter shims themselves. Distinct // from the sampler's `dd_tl_state_t::reentry_guard` because that one // lives inside a struct we have to *allocate* (via `calloc`) on first @@ -187,6 +189,18 @@ pub unsafe extern "C" fn gotter_calloc(nmemb: usize, size: usize) -> *mut c_void dd_allocation_created(raw, req) } +/// `realloc` hook. +/// +/// Handled as four disjoint cases: +/// +/// 1. **`ptr == NULL`** — equivalent to `malloc(size)`. Runs the normal sampling path (`request` + +/// `created`). +/// 2. **`size == 0`** — equivalent to `free(ptr)` for the allocators we hook. Consumes the sampler +/// flag via `dd_allocation_freed` (clears the header + fires `ddheap:free`) and forwards. +/// 3. **`ptr` is unsampled** — passthrough to the underlying realloc. We don't newly sample here +/// (see TODO on the branch). +/// 4. **`ptr` is sampled** — MVP: model a successful sampled realloc as `free(old sampled)` +/// followed by a new *unsampled* allocation. See the branch comment for the full rationale. #[no_mangle] pub unsafe extern "C" fn gotter_realloc(ptr: *mut c_void, size: usize) -> *mut c_void { let Some(real): Option = load_fn(&ORIG_REALLOC) else { @@ -197,23 +211,89 @@ pub unsafe extern "C" fn gotter_realloc(ptr: *mut c_void, size: usize) -> *mut c return real(ptr, size); } ensure_tls(); - // realloc decomposes into a free + an alloc from the sampler's - // perspective. If `ptr` was sampled we have to clear its flag so the - // underlying allocator sees the raw pointer; size==0 is a free. - let freed_ptr = if ptr.is_null() { - std::ptr::null_mut() - } else { - dd_allocation_freed(ptr, 0, 0).ptr - }; + // Case 1: realloc(NULL, size) == malloc(size). Normal sampling path. + if ptr.is_null() { + let alignment = core::mem::align_of::() * 2; + let req = dd_allocation_requested(size, alignment); + let raw = real(std::ptr::null_mut(), req.size); + return dd_allocation_created(raw, req); + } + + // Case 2: realloc(p, 0) == free(p). Safe to consume the sampler + // flag before forwarding because the old block will not remain + // live on this path. if size == 0 { - // realloc(p, 0) == free(p). Just forward. - return real(freed_ptr, 0); + let freed = dd_allocation_freed(ptr, 0, 0); + return real(freed.ptr, 0); } - let req = dd_allocation_requested(size, core::mem::align_of::() * 2); - let raw = real(freed_ptr, req.size); - dd_allocation_created(raw, req) + // Peek is non-destructive: if realloc fails below, the old header + // must stay intact so a later gotter_free(ptr) still resolves the + // right raw pointer for libc. + let mut old_raw = std::ptr::null_mut(); + let mut old_offset = 0usize; + if !dd_sample_flag_peek(ptr, &mut old_raw, &mut old_offset) { + // Case 3: unsampled old allocation. Passthrough. We deliberately + // do not opportunistically sample here: we don't know the old + // user-requested size, so we can't safely place a header and + // move user contents without either over-reading the old block + // or shifting user data. + // TODO: revisit once sampled headers carry the original user + // size, or we have another safe bound for copying old contents. + return real(ptr, size); + } + + // Case 4: sampled old allocation. MVP semantics: + // + // successful sampled realloc = ddheap:free(old sampled) + // + new unsampled allocation + // + // Why unsampled on the new side: + // * The new raw pointer picked by the underlying realloc may sit at a different page offset, + // so the sampler's x86 offset picker may want a different `new_offset` than `old_offset`. + // Stamping a header at the new offset before moving user data would overlap the + // still-to-be-copied source region. + // * We also don't know the original user-requested size, so we can't cheaply resize + + // resample + report a coherent alloc event to the profiler. + // + // The overhead of "unsampled from now on" is only paid on + // sampled-then-realloc'd blocks, which are already rare. + // + // TODO: emit `free + new sampled allocation` when sampled headers + // carry the original user size (see project TODO #13 / repo TODO #7). + let Some(raw_size) = sampled_realloc_raw_size(size, old_offset) else { + return std::ptr::null_mut(); + }; + + // Underlying realloc. Contract: on failure, `old_raw` is still + // live and its header is still valid (peek was non-destructive). + let new_raw = real(old_raw, raw_size); + if new_raw.is_null() { + return std::ptr::null_mut(); + } + + // On success, libc copied the old block's bytes into `new_raw` + // starting at offset 0, so the old user data now sits at + // `new_raw + old_offset`. Shift it down to `new_raw` so the caller + // sees the user pointer at offset 0 of an unsampled block. + // memmove (not memcpy) because when the underlying realloc extends + // in place, `new_raw == old_raw` and source/destination overlap. + #[cfg(target_arch = "x86_64")] + { + let copied_user = (new_raw as *mut u8).add(old_offset).cast::(); + if new_raw != copied_user { + libc::memmove(new_raw, copied_user, size); + } + } + + // Fire ddheap:free for the OLD user pointer (the address the + // profiler last saw as live). We call `dd_probe_free` directly, + // not `dd_allocation_freed`, because there's no header left to + // clear: libc has already consumed the old block and we've + // overwritten those bytes above. + dd_probe_free(ptr); + new_raw } #[no_mangle] @@ -232,9 +312,15 @@ pub unsafe extern "C" fn gotter_posix_memalign( ensure_tls(); let req = dd_allocation_requested(size, alignment); let ret = real(memptr, alignment, req.size); + // Always pair with dd_allocation_created, even on failure, so the + // sampler's reentry guard opened by dd_allocation_requested is + // closed. Passing raw == NULL is the documented "allocation + // failed" path: no flag stamped, no USDT fired, guard closed. if ret == 0 && !memptr.is_null() { let raw = *memptr; *memptr = dd_allocation_created(raw, req); + } else { + let _ = dd_allocation_created(std::ptr::null_mut(), req); } ret } @@ -307,5 +393,10 @@ pub unsafe extern "C" fn gotter_pthread_create( // extend hooks later. #[allow(dead_code)] fn _types_anchor() -> dd_alloc_req_t { - dd_alloc_req_t { size: 0, weight: 0 } + dd_alloc_req_t { + size: 0, + user_size: 0, + alignment: 0, + weight: 0, + } } diff --git a/libdd-heap-gotter/src/lib.rs b/libdd-heap-gotter/src/lib.rs index e330107f32..57b8a039f1 100644 --- a/libdd-heap-gotter/src/lib.rs +++ b/libdd-heap-gotter/src/lib.rs @@ -43,6 +43,8 @@ mod elf; #[cfg(target_os = "linux")] mod hooks; +#[cfg(any(target_os = "linux", test))] +mod realloc_math; #[cfg(target_os = "linux")] use std::sync::Mutex; @@ -148,13 +150,11 @@ fn register_all(so: &mut SymbolOverrides) { // Register one entry per supported symbol. The `ref_slot` raw // pointer is to a `'static AtomicUsize`, so it's valid forever. - // AtomicUsize is repr(transparent) over UnsafeCell; we - // intentionally bypass its API for the install-time write because - // we hand the raw `*mut usize` to the ELF GOT scanner. Hooks then - // read it back via `Atomic::load(Acquire)`. + // 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) { - let slot_ptr = slot as *const AtomicUsize as *mut usize; - so.register(name, hook_addr, slot_ptr); + so.register(name, hook_addr, slot as *const AtomicUsize); } reg( diff --git a/libdd-heap-gotter/src/realloc_math.rs b/libdd-heap-gotter/src/realloc_math.rs new file mode 100644 index 0000000000..01f9a092cb --- /dev/null +++ b/libdd-heap-gotter/src/realloc_math.rs @@ -0,0 +1,36 @@ +// Copyright 2025-Present Datadog, Inc. https://www.datadoghq.com/ +// SPDX-License-Identifier: Apache-2.0 + +//! Pure helpers for sampled `realloc` layout arithmetic. +//! +//! Split out from `hooks::gotter_realloc` so the offset/size arithmetic +//! can be unit-tested on any host, without pulling in the Linux-only +//! GOT-hook module. + +/// Size to pass to the underlying `realloc` for a sampled old block. +/// +/// The MVP sampled-realloc path returns an *unsampled* pointer at +/// `new_raw` (offset 0), so the underlying block must be large enough +/// to hold both: +/// +/// 1. the old header + slack in `[0, old_offset)` (untouched by libc's realloc copy), +/// 2. `size` bytes of user data at `[old_offset, old_offset + size)`, which we later `memmove` +/// down to `[0, size)`. +/// +/// Hence `size + old_offset`. Returns `None` on overflow so the caller +/// can fail the realloc rather than truncate. +pub(crate) fn sampled_realloc_raw_size(size: usize, old_offset: usize) -> Option { + size.checked_add(old_offset) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn sampled_realloc_raw_size_reserves_old_offset_for_unsampled_result() { + assert_eq!(sampled_realloc_raw_size(64, 16), Some(80)); + assert_eq!(sampled_realloc_raw_size(64, 64), Some(128)); + assert_eq!(sampled_realloc_raw_size(usize::MAX, 16), None); + } +} diff --git a/libdd-heap-sampler/build.rs b/libdd-heap-sampler/build.rs index 43ad82d406..3131b592be 100644 --- a/libdd-heap-sampler/build.rs +++ b/libdd-heap-sampler/build.rs @@ -109,6 +109,10 @@ mod linux { .flag_if_supported("-mtls-dialect=gnu2"); build.compile("dd_heap_sampler"); + // `allocation_requested.c` calls `log()`; glibc keeps it in libm, + // separate from libc. Harmless on musl (empty stub libm). + println!("cargo:rustc-link-lib=m"); + for f in SOURCES { println!("cargo:rerun-if-changed={f}"); } diff --git a/libdd-heap-sampler/include/datadog/heap/allocation_requested.h b/libdd-heap-sampler/include/datadog/heap/allocation_requested.h index 3de2cfc425..90541e2845 100644 --- a/libdd-heap-sampler/include/datadog/heap/allocation_requested.h +++ b/libdd-heap-sampler/include/datadog/heap/allocation_requested.h @@ -10,12 +10,18 @@ * the cost of drawing a fresh inter-sample interval and setting the reentry * guard. * - * The returned dd_alloc_req_t carries two values the caller must forward: - * - size: the number of bytes to request from the real allocator - * (may be larger than the original on architectures that store - * the sample flag in a header word before the user pointer). - * - weight: 0 if not sampled; otherwise the unbiased size estimator - * (nsamples * interval) to attribute to this allocation. + * The returned dd_alloc_req_t carries the values the caller must forward: + * - size: the number of bytes to request from the real allocator + * (may be larger than the original on architectures that + * store the sample flag in a header word before the user + * pointer). + * - user_size: the original application-requested size, reported to + * the profiler via the ddheap:alloc USDT. + * - alignment: passed through so dd_allocation_created can place the + * user pointer correctly relative to the raw pointer. + * - weight: 0 if not sampled; otherwise the unbiased size + * estimator (nsamples * interval) to attribute to this + * allocation. * * Always pair with dd_allocation_created(), even if the allocator fails. */ @@ -31,18 +37,30 @@ * Return type for dd_allocation_requested, paired with the `req` * argument of dd_allocation_created. * - * size - size the wrapper MUST pass to its underlying allocator. - * Usually the caller's requested size; on sampled - * allocations on architectures that use in-band flagging - * (header magic), this is bumped by the flag overhead. - * weight - 0 if this allocation was not sampled; otherwise the - * unbiased size estimator (nsamples * interval) for - * aggregated reporting. + * size - size the wrapper MUST pass to its underlying allocator. + * Usually the caller's requested size; on sampled + * allocations on architectures that use in-band flagging + * (header magic), this is bumped for the header and any + * alignment slack. + * user_size - the size the application originally asked for. This is + * the value that gets reported to the profiler via the + * ddheap:alloc USDT, so that heap-size distributions are + * not skewed by the sampler's per-allocation overhead. + * alignment - alignment the wrapper MUST pass to its underlying + * allocator. Equals the alignment the caller passed to + * dd_allocation_requested; carried through so + * dd_allocation_created can place the user pointer + * correctly relative to the raw pointer. + * weight - 0 if this allocation was not sampled; otherwise the + * unbiased size estimator (nsamples * interval) for + * aggregated reporting. * - * Kept to 16 bytes so it's returned in registers. + * 32 bytes on 64-bit targets, cache-line-friendly. */ typedef struct { size_t size; + size_t user_size; + size_t alignment; uint64_t weight; } dd_alloc_req_t; @@ -73,7 +91,7 @@ dd_alloc_req_t dd_allocation_requested_slow(dd_tl_state_t *tl, size_t size, */ static inline __attribute__((always_inline)) dd_alloc_req_t dd_allocation_requested(size_t size, size_t alignment) { - dd_alloc_req_t out = { size, 0 }; + dd_alloc_req_t out = { size, size, alignment, 0 }; // If we don't have TLS yet, or the reentry guard is set (meaning a sampled // allocation is already in flight on this thread and something in its slow diff --git a/libdd-heap-sampler/include/datadog/heap/sample_flag.h b/libdd-heap-sampler/include/datadog/heap/sample_flag.h index b7a8db8012..3479f569bc 100644 --- a/libdd-heap-sampler/include/datadog/heap/sample_flag.h +++ b/libdd-heap-sampler/include/datadog/heap/sample_flag.h @@ -23,23 +23,24 @@ * Must be called once per thread before dd_sample_flag_apply/check are used. * On arm64 this issues prctl(PR_SET_TAGGED_ADDR_CTRL) so tagged pointers * survive syscalls; on other architectures this is a no-op. + * + * Returns true when the thread is safe to sample and false when the + * flagging scheme is unavailable (e.g. arm64 kernel/seccomp rejected + * PR_SET_TAGGED_ADDR_CTRL). Callers must treat a false return as + * "disable sampling on this thread" and never call + * dd_sample_flag_apply on it, otherwise tagged pointers will be + * rejected by the kernel with EFAULT on the next syscall. */ -void dd_sample_flag_thread_init(void); - -/* - * Extra bytes the flag reserves on top of the user's requested size. - * dd_allocation_requested adds this to the size it hands back to the - * wrapper so the underlying allocator reserves room for the flag. 0 on - * architectures that flag via pointer bits instead of memory. - */ -size_t dd_sample_flag_overhead(void); +bool dd_sample_flag_thread_init(void); /* * Apply the sampled flag. Takes the raw pointer from the underlying - * allocator (of size `user-requested + dd_sample_flag_overhead()`) and - * returns the user-visible pointer to hand to the application. + * allocator and returns the user-visible pointer to hand to the + * application. On architectures that offset the user pointer inside a + * bumped allocation (x86-64), the offset is picked to satisfy the + * caller-requested alignment. */ -void *dd_sample_flag_apply(void *raw); +void *dd_sample_flag_apply(void *raw, size_t alignment); /* * If `user` was previously returned by dd_sample_flag_apply, write the @@ -48,89 +49,110 @@ void *dd_sample_flag_apply(void *raw); */ bool dd_sample_flag_check(void *user, void **raw_out); +/* + * Non-destructive variant of dd_sample_flag_check. Useful for realloc: + * callers can resolve the raw pointer before calling the underlying + * realloc, while leaving the old allocation's flag intact in case + * realloc fails and the old allocation remains live. + * + * If `user` is sampled, returns true and fills raw_out / offset_out + * without clearing the allocation's flag header. + */ +bool dd_sample_flag_peek(void *user, void **raw_out, size_t *offset_out); + +/* + * Largest alignment the sampler will honor. Above this we pass the + * allocation through unsampled: the header + slack overhead grows with + * alignment and stops being proportionate to any observability gain. + * Sized in bytes; equals one typical 4 KiB page on the supported + * architectures. + */ +#define DD_SAMPLE_ALIGNMENT_CAP 4096 + #if defined(__x86_64__) #define DD_HEADER_BYTES 16 #define DD_PAGE_SIZE 4096 /* - * x86-64 marks sampled allocations with a magic word written 16 bytes - * before the user pointer. The allocator hands us a 16-byte aligned raw - * pointer of size user_size + 32; we pick the user pointer inside that - * surplus so reading 16 bytes behind it always stays within a mapped - * page. + * x86-64 marks sampled allocations with a 16-byte header written + * immediately before the user pointer. The header stores an 8-byte + * magic word plus an 8-byte offset from `raw` to `user`, so recovery + * at free time is direct: `raw = user - offset`. + * + * Layout: + * [raw ... slack ...] [magic(8) | offset(8)] [user_data ...] + * ^ user - 16 ^ user * - * Two layouts are used, distinguished by which magic is present: - * DD_MAGIC_A: user = raw + 16 (common case) - * DD_MAGIC_B: user = raw + 32 (only when raw + 16 would land in - * the first 16 bytes of a page; about - * 1/256 of sampled allocations) + * The user pointer is placed at `raw + N` where * - * With +32 surplus exactly one of {raw+16, raw+32} is always safe: - * both are 16-byte aligned and their page offsets differ by 16, so at - * most one can fall in the unsafe [0, 16) band. The two-magic encoding - * lets us recover raw at free time without any extra metadata. - * This costs us on average 32 bytes on every 512 KiB allocated, or about - * 62 KiB on 1 GiB. + * N = max(alignment, DD_HEADER_BYTES) * - * Alternatives considered: - * - +16 surplus with a free+malloc retry when raw+16 is unsafe. - * Saves 16 bytes per sampled allocation in the common case (~99.6%) - * but plumbs malloc/free callbacks into the sampler and risks an - * OOM on the retry path. Not worth it: sampled allocations are - * rare so the absolute byte cost of always +32 is negligible. - * - Drop samples whose user pointer comes back in the unsafe band. - * Simple, but introduces a bias in the sample set - * (allocations whose raw lands at offset 4080 of a page are never - * reported). + * plus one further `alignment`-sized bump when `raw + N` would land at + * page offset < DD_HEADER_BYTES. That bump preserves the invariant + * that the fast-path filter (`user & (PAGE-1) < 16 => unsampled`) + * relies on to safely read the 16 header bytes without ever + * dereferencing an unmapped previous page. Since `alignment` is a + * power of two, adding another `alignment` bytes preserves the + * requested user-pointer alignment. + * + * Compared with the previous two-magic (A/B) scheme the offset is now + * stored explicitly, which drops one branch on the free path and + * generalises cleanly to arbitrary caller alignments (Rust `Layout`, + * `posix_memalign`, `aligned_alloc`). Overhead in the common + * alignment <= 16 case is 16 bytes per sampled allocation, down from + * the previous 32. */ -#define DD_MAGIC_A 0xfab1eddec0dedca7ULL /* user = raw + 16 */ -#define DD_MAGIC_B 0xfab1eddec0dedca8ULL /* user = raw + 32 */ +#define DD_MAGIC 0xfab1eddec0dedca7ULL /* * Layout helpers. x86_apply and x86_raw_from_user MUST be each other's - * inverse: the magic stamped at apply time is what lets x86_raw_from_user - * recover the same raw at check time. + * inverse: the offset stamped at apply time is what lets + * x86_raw_from_user recover the same raw at check time. */ /* - * Pick the user pointer inside the [raw, raw + 32) surplus and stamp - * the matching magic into the 16-byte header immediately preceding it. - * The chosen layout is encoded in the magic value itself so - * x86_raw_from_user can recover `raw` at free time without any extra - * metadata. + * Pick the user pointer within the bumped allocation backing `raw` + * such that: + * - `user - raw >= DD_HEADER_BYTES` (room for the header), + * - `user` is `alignment`-aligned, + * - `user & (DD_PAGE_SIZE - 1) >= DD_HEADER_BYTES` (fast-path filter + * never treats a sampled allocation as unsampled). * - * raw - 16-byte aligned pointer from the underlying allocator, - * backing a region of size user_size + 32. + * Stamp (magic, offset) at `user - 16` so x86_raw_from_user can + * recover `raw` at free time without any per-allocation metadata + * beyond the header itself. * - * Returns the user-visible pointer (raw+16 or raw+32), chosen so the - * 16-byte header slot never straddles the start of a page. + * The caller is responsible for ensuring `raw` is aligned to + * `alignment` (via aligned_alloc/posix_memalign) or, when + * `alignment <= DD_HEADER_BYTES`, that `raw` is at least + * DD_HEADER_BYTES-aligned (malloc's default on x86-64 glibc/musl). */ static inline __attribute__((always_inline)) -void *x86_apply(void *raw) { +void *x86_apply(void *raw, size_t alignment) { uintptr_t r = (uintptr_t)raw; - uintptr_t u = r + DD_HEADER_BYTES; - uint64_t magic = DD_MAGIC_A; + size_t n = alignment > DD_HEADER_BYTES ? alignment : DD_HEADER_BYTES; + uintptr_t u = r + n; if ((u & (DD_PAGE_SIZE - 1)) < DD_HEADER_BYTES) { - u = r + 2 * DD_HEADER_BYTES; - magic = DD_MAGIC_B; + n += (alignment > DD_HEADER_BYTES ? alignment : DD_HEADER_BYTES); + u = r + n; } + uint64_t magic = DD_MAGIC; + uint64_t offset = (uint64_t)n; memcpy((void *)(u - DD_HEADER_BYTES), &magic, sizeof(magic)); + memcpy((void *)(u - DD_HEADER_BYTES + sizeof(magic)), &offset, + sizeof(offset)); return (void *)u; } /* - * Inverse of x86_apply: given a user pointer previously produced by - * it and the magic read out of the 16-byte header preceding `user` - * (DD_MAGIC_A => user = raw+16, DD_MAGIC_B => raw+32), return the - * original raw pointer to hand back to the underlying free. + * Inverse of x86_apply: given the offset recovered from the header, + * return the original raw pointer to hand back to the underlying free. */ static inline __attribute__((always_inline)) -void *x86_raw_from_user(void *user, uint64_t magic) { - uintptr_t u = (uintptr_t)user; - return (void *)(u - (magic == DD_MAGIC_A ? DD_HEADER_BYTES - : 2 * DD_HEADER_BYTES)); +void *x86_raw_from_user(void *user, uint64_t offset) { + return (void *)((uintptr_t)user - (uintptr_t)offset); } static inline __attribute__((always_inline)) @@ -142,13 +164,21 @@ bool dd_sample_flag_check_fast(void *user, void **raw_out) { void *header = (char *)user - DD_HEADER_BYTES; uint64_t magic; memcpy(&magic, header, sizeof(magic)); - if (magic != DD_MAGIC_A && magic != DD_MAGIC_B) { + if (magic != DD_MAGIC) { return false; } - const uint64_t zero = 0; - memcpy(header, &zero, sizeof(zero)); - *raw_out = x86_raw_from_user(user, magic); + uint64_t offset; + memcpy(&offset, (char *)header + sizeof(magic), sizeof(offset)); + + /* Clear the whole 16-byte header so a re-use of this address + * (e.g. allocator returns the same block to a later, unsampled + * allocation whose user data happens to encode the magic) doesn't + * masquerade as a stale sampled allocation. */ + const uint64_t zeros[2] = { 0, 0 }; + memcpy(header, zeros, sizeof(zeros)); + + *raw_out = x86_raw_from_user(user, offset); return true; } diff --git a/libdd-heap-sampler/src/allocation_created.c b/libdd-heap-sampler/src/allocation_created.c index 0381a973d7..7286dfbed3 100644 --- a/libdd-heap-sampler/src/allocation_created.c +++ b/libdd-heap-sampler/src/allocation_created.c @@ -33,8 +33,11 @@ void *dd_allocation_created_slow(void *raw, dd_alloc_req_t req) { /* Apply the sample flag and fire the USDT. We use the user pointer * (post-flag) as the USDT argument so the profiler sees the same * address the application will. */ - user = dd_sample_flag_apply(raw); - dd_probe_alloc(user, (uint64_t)req.size, req.weight); + user = dd_sample_flag_apply(raw, req.alignment); + /* Report the application-requested size, not the sampler-bumped + * size (`req.size`), so heap-size distributions in the profiler + * aren't skewed by per-sample overhead. */ + dd_probe_alloc(user, (uint64_t)req.user_size, req.weight); } /* Always close the reentry guard, even on allocation failure (raw == NULL), diff --git a/libdd-heap-sampler/src/allocation_freed.c b/libdd-heap-sampler/src/allocation_freed.c index 14b5ca4dee..d9acf34062 100644 --- a/libdd-heap-sampler/src/allocation_freed.c +++ b/libdd-heap-sampler/src/allocation_freed.c @@ -5,6 +5,8 @@ #include #include +#include + /* * Slow path for dd_allocation_freed. We only arrive here when * dd_sample_flag_check_fast confirmed that ptr carries the sample flag, @@ -18,19 +20,37 @@ */ dd_alloc_freed_t dd_allocation_freed_slow(void *ptr, void *raw, size_t size, size_t alignment) { - (void)alignment; - /* Fire with the user-visible pointer, matching what was reported at alloc * time, so the profiler can correlate the two events by address. */ dd_probe_free(ptr); + dd_alloc_freed_t out = { /* Return the raw pointer so the caller passes the real allocation base * to the deallocator, not the user pointer that may be offset or tagged. */ .ptr = raw, - /* Recover the full allocation size including any header reserved at - * alloc time. On arm64 dd_sample_flag_overhead() is 0, so this is - * a no-op there. */ - .size = size + dd_sample_flag_overhead(), + .size = size, }; + +#if defined(__x86_64__) + /* Recover the bumped size the allocator actually holds. On x86-64 + * that's user_size + (user - raw), then rounded up to a multiple + * of alignment to match dd_allocation_requested_slow's bumped-size + * arithmetic. Sized-free callers (sdallocx, operator delete(sz)) + * rely on this being exact. + * + * When the caller doesn't know the alignment (alignment == 0), + * fall back to size + offset. Plain free() ignores out.size so + * this only matters for sized-free variants that must supply an + * alignment. */ + size_t offset = (size_t)((uintptr_t)ptr - (uintptr_t)raw); + size_t bumped = size + offset; + if (alignment > 1) { + size_t mask = alignment - 1; + bumped = (bumped + mask) & ~mask; + } + out.size = bumped; +#else + (void)alignment; +#endif return out; } \ No newline at end of file diff --git a/libdd-heap-sampler/src/allocation_requested.c b/libdd-heap-sampler/src/allocation_requested.c index 66a9f8f5b6..92408feaba 100644 --- a/libdd-heap-sampler/src/allocation_requested.c +++ b/libdd-heap-sampler/src/allocation_requested.c @@ -5,6 +5,8 @@ #include #include +#include +#include /* * Advances the Park-Miller LCG one step and returns the new 31-bit state. @@ -87,10 +89,60 @@ static uint64_t sample(dd_tl_state_t *tl) { * * (ddprof: AllocationTracker::track_allocation / next_sample_interval) */ -dd_alloc_req_t dd_allocation_requested_slow(dd_tl_state_t *tl, size_t size, - size_t alignment) { +/* + * Compute the bumped size to pass to the underlying allocator for a + * sampled allocation. Returns true on success and writes the bumped + * size to *out_size. Returns false when the arithmetic would overflow + * or the alignment exceeds what the sampler is willing to honor, in + * which case the caller must pass this allocation through unsampled. + * + * x86-64 places a 16-byte (magic, offset) header immediately before + * the user pointer, and picks user = raw + max(alignment, 16) (plus + * possibly another `alignment` bump to satisfy the page-boundary + * invariant). The bumped size must reserve room for that offset plus + * the user's requested bytes, and must satisfy aligned_alloc's + * size %% alignment == 0 constraint (a superset of posix_memalign's + * requirements). + * + * arm64 uses TBI tagging with no size bump. + */ +static bool bumped_alloc_size(size_t user_size, size_t alignment, + size_t *out_size) { +#if defined(__x86_64__) + if (alignment > DD_SAMPLE_ALIGNMENT_CAP) return false; + + /* Reserve twice the base offset so x86_apply's page-boundary bump + * (which may push the user pointer another `base` bytes forward) + * always fits inside the allocation. Base is max(alignment, 16): + * the minimum offset needed to seat the 16-byte header before the + * user pointer while staying alignment-aligned. */ + size_t base = alignment > DD_HEADER_BYTES ? alignment : DD_HEADER_BYTES; + if (base > SIZE_MAX / 2) return false; + size_t reserve = base * 2; + + if (reserve > SIZE_MAX - user_size) return false; + size_t bumped = user_size + reserve; + + /* Round up to a multiple of alignment so aligned_alloc callers + * (which require size %% alignment == 0) are satisfied. For + * alignment <= DD_HEADER_BYTES this is already a multiple of + * alignment (reserve = 32, a multiple of 1/2/4/8/16). */ + if (alignment > 1) { + size_t mask = alignment - 1; + if (bumped > SIZE_MAX - mask) return false; + bumped = (bumped + mask) & ~mask; + } + *out_size = bumped; + return true; +#else (void)alignment; + *out_size = user_size; + return true; +#endif +} +dd_alloc_req_t dd_allocation_requested_slow(dd_tl_state_t *tl, size_t size, + size_t alignment) { /* Open the reentry guard before doing anything else. Any allocation that * happens between here and dd_allocation_created_slow (e.g. inside log() * or the USDT machinery) will see the guard set and pass through unsampled. */ @@ -101,15 +153,26 @@ dd_alloc_req_t dd_allocation_requested_slow(dd_tl_state_t *tl, size_t size, /* First-interval miss: no sample this time. Close the guard now since * dd_allocation_created_slow won't be called on the sampled path. */ tl->reentry_guard = false; - dd_alloc_req_t out = { size, 0 }; + dd_alloc_req_t out = { size, size, alignment, 0 }; + return out; + } + + size_t bumped; + if (!bumped_alloc_size(size, alignment, &bumped)) { + /* Alignment too large or arithmetic overflow: pass through as + * an unsampled allocation rather than corrupt the request. The + * guard must be closed here since dd_allocation_created_slow + * won't be reached (weight == 0 fast-path in the header). */ + tl->reentry_guard = false; + dd_alloc_req_t out = { size, size, alignment, 0 }; return out; } dd_alloc_req_t out = { - /* Ask the allocator for extra bytes if the flag scheme needs them - * (16 bytes on x86-64 for the header; 0 on arm64). */ - .size = size + dd_sample_flag_overhead(), - .weight = weight, + .size = bumped, + .user_size = size, + .alignment = alignment, + .weight = weight, }; return out; } \ No newline at end of file diff --git a/libdd-heap-sampler/src/generated/bindings.rs b/libdd-heap-sampler/src/generated/bindings.rs index 5d4078d674..5e972d3397 100644 --- a/libdd-heap-sampler/src/generated/bindings.rs +++ b/libdd-heap-sampler/src/generated/bindings.rs @@ -53,16 +53,22 @@ unsafe extern "C" { #[derive(Debug, Copy, Clone)] pub struct dd_alloc_req_t { pub size: usize, + pub user_size: usize, + pub alignment: usize, pub weight: u64, } #[allow(clippy::unnecessary_operation, clippy::identity_op)] const _: () = { - ["Size of dd_alloc_req_t"][::std::mem::size_of::() - 16usize]; + ["Size of dd_alloc_req_t"][::std::mem::size_of::() - 32usize]; ["Alignment of dd_alloc_req_t"][::std::mem::align_of::() - 8usize]; ["Offset of field: dd_alloc_req_t::size"] [::std::mem::offset_of!(dd_alloc_req_t, size) - 0usize]; + ["Offset of field: dd_alloc_req_t::user_size"] + [::std::mem::offset_of!(dd_alloc_req_t, user_size) - 8usize]; + ["Offset of field: dd_alloc_req_t::alignment"] + [::std::mem::offset_of!(dd_alloc_req_t, alignment) - 16usize]; ["Offset of field: dd_alloc_req_t::weight"] - [::std::mem::offset_of!(dd_alloc_req_t, weight) - 8usize]; + [::std::mem::offset_of!(dd_alloc_req_t, weight) - 24usize]; }; unsafe extern "C" { pub fn dd_allocation_requested_slow( @@ -89,18 +95,25 @@ unsafe extern "C" { ) -> *mut ::std::os::raw::c_void; } unsafe extern "C" { - pub fn dd_sample_flag_thread_init(); + pub fn dd_sample_flag_thread_init() -> bool; } unsafe extern "C" { - pub fn dd_sample_flag_overhead() -> usize; + pub fn dd_sample_flag_apply( + raw: *mut ::std::os::raw::c_void, + alignment: usize, + ) -> *mut ::std::os::raw::c_void; } unsafe extern "C" { - pub fn dd_sample_flag_apply(raw: *mut ::std::os::raw::c_void) -> *mut ::std::os::raw::c_void; + pub fn dd_sample_flag_check( + user: *mut ::std::os::raw::c_void, + raw_out: *mut *mut ::std::os::raw::c_void, + ) -> bool; } unsafe extern "C" { - pub fn dd_sample_flag_check( + pub fn dd_sample_flag_peek( user: *mut ::std::os::raw::c_void, raw_out: *mut *mut ::std::os::raw::c_void, + offset_out: *mut usize, ) -> bool; } unsafe extern "C" { diff --git a/libdd-heap-sampler/src/lib.rs b/libdd-heap-sampler/src/lib.rs index 1d623a7794..07606d7299 100644 --- a/libdd-heap-sampler/src/lib.rs +++ b/libdd-heap-sampler/src/lib.rs @@ -41,7 +41,7 @@ mod tests { #[test] fn alloc_req_layout_matches_c() { - assert_eq!(core::mem::size_of::(), 16); + assert_eq!(core::mem::size_of::(), 32); assert_eq!(core::mem::align_of::(), 8); } @@ -76,6 +76,8 @@ mod tests { let fake = 0xdead_beef_usize as *mut c_void; let req = dd_alloc_req_t { size: 64, + user_size: 64, + alignment: 8, weight: 0, }; assert_eq!(dd_allocation_created(fake, req), fake); diff --git a/libdd-heap-sampler/src/sample_flag.c b/libdd-heap-sampler/src/sample_flag.c index 9ee2004087..58ba9e82f3 100644 --- a/libdd-heap-sampler/src/sample_flag.c +++ b/libdd-heap-sampler/src/sample_flag.c @@ -6,25 +6,45 @@ #if defined(__x86_64__) /* - * x86-64 layout details live in sample_flag.h alongside DD_MAGIC_A/B - * and the x86 helper pair. Summary: allocate user_size + 32, - * let the helper pick the user pointer (raw+16 in the common case, - * raw+32 when raw+16 would land in the unsafe first-16-bytes-of-page - * region), and write the corresponding magic at user-16. On free the - * matched magic tells us which offset to use to recover raw. + * x86-64 layout details live in sample_flag.h alongside DD_MAGIC and + * the x86 helper pair. Summary: bump the size to reserve room for a + * 16-byte (magic, offset) header plus alignment slack, let the helper + * pick the user pointer inside the bumped region such that user is + * alignment-aligned and user & (PAGE-1) >= 16, and stamp (magic, + * offset) at user - 16. On free the offset stored in the header lets + * us recover raw directly. */ -void dd_sample_flag_thread_init(void) { /* nothing needed on x86-64 */ } +bool dd_sample_flag_thread_init(void) { + /* Nothing needed on x86-64; always safe. */ + return true; +} -size_t dd_sample_flag_overhead(void) { - // +32 guarantees both raw+16 and raw+32 are available; one of them - // is always safe to use as the user pointer regardless of where raw - // lands relative to a page boundary. - return 2 * DD_HEADER_BYTES; +void *dd_sample_flag_apply(void *raw, size_t alignment) { + return x86_apply(raw, alignment); } -void *dd_sample_flag_apply(void *raw) { - return x86_apply(raw); +bool dd_sample_flag_peek(void *user, void **raw_out, size_t *offset_out) { + if (((uintptr_t)user & (DD_PAGE_SIZE - 1)) < DD_HEADER_BYTES) { + return false; + } + + void *header = (char *)user - DD_HEADER_BYTES; + uint64_t magic; + memcpy(&magic, header, sizeof(magic)); + if (magic != DD_MAGIC) { + return false; + } + + uint64_t offset; + memcpy(&offset, (char *)header + sizeof(magic), sizeof(offset)); + if (offset < DD_HEADER_BYTES || offset > 2 * DD_SAMPLE_ALIGNMENT_CAP) { + return false; + } + + *raw_out = x86_raw_from_user(user, offset); + *offset_out = (size_t)offset; + return true; } bool dd_sample_flag_check(void *user, void **raw_out) { @@ -52,12 +72,17 @@ bool dd_sample_flag_check(void *user, void **raw_out) { * TODO: audit interaction with HWASan and MTE, which also use the top byte. */ -size_t dd_sample_flag_overhead(void) { - return 0; +void *dd_sample_flag_apply(void *raw, size_t alignment) { + (void)alignment; + return (void *)((uintptr_t)raw | DD_TBI_TAGGED); } -void *dd_sample_flag_apply(void *raw) { - return (void *)((uintptr_t)raw | DD_TBI_TAGGED); +bool dd_sample_flag_peek(void *user, void **raw_out, size_t *offset_out) { + if (!dd_sample_flag_check_fast(user, raw_out)) { + return false; + } + *offset_out = 0; + return true; } bool dd_sample_flag_check(void *user, void **raw_out) { @@ -82,11 +107,19 @@ bool dd_sample_flag_check(void *user, void **raw_out) { # define PR_TAGGED_ADDR_ENABLE (1UL << 0) #endif -void dd_sample_flag_thread_init(void) { - prctl(PR_SET_TAGGED_ADDR_CTRL, PR_TAGGED_ADDR_ENABLE, 0, 0, 0); +bool dd_sample_flag_thread_init(void) { + /* prctl returns 0 on success. On failure (older kernel without + * PR_SET_TAGGED_ADDR_CTRL, seccomp filter blocking it, MTE + * conflict, ...) tagged pointers would be rejected by the kernel + * with EFAULT the next time one crosses a syscall. Report failure + * so the caller disables sampling on this thread. */ + return prctl(PR_SET_TAGGED_ADDR_CTRL, PR_TAGGED_ADDR_ENABLE, 0, 0, 0) == 0; } #else -void dd_sample_flag_thread_init(void) { } +bool dd_sample_flag_thread_init(void) { + /* No tagging on non-Linux arm64 targets; sampling is always safe. */ + return true; +} #endif #else diff --git a/libdd-heap-sampler/src/tl_state.c b/libdd-heap-sampler/src/tl_state.c index 34cd8d39c7..ea2ad3b33e 100644 --- a/libdd-heap-sampler/src/tl_state.c +++ b/libdd-heap-sampler/src/tl_state.c @@ -4,6 +4,7 @@ #include #include +#include #include #include @@ -38,7 +39,7 @@ static void tl_state_populate(dd_tl_state_t *st) { st->initialized = true; st->reentry_guard = true; - dd_sample_flag_thread_init(); + bool flag_ok = dd_sample_flag_thread_init(); struct timespec ts = {0}; (void)clock_gettime(CLOCK_MONOTONIC, &ts); @@ -48,7 +49,15 @@ static void tl_state_populate(dd_tl_state_t *st) { st->rng = seed ? seed : 1u; st->sampling_interval = DD_SAMPLING_INTERVAL_DEFAULT; - st->reentry_guard = false; + /* If the per-thread flagging scheme is unavailable (arm64 prctl + * failure), leave reentry_guard set. The fast path in + * dd_allocation_requested short-circuits on reentry_guard, so this + * thread will pass every allocation through unsampled - no tagged + * pointers get produced, no syscall EFAULTs. Cheaper than adding a + * dedicated "sampling_disabled" field + branch on the hot path. */ + if (flag_ok) { + st->reentry_guard = false; + } } dd_tl_state_t *dd_tl_state_get(void) { From d46e2c48052aee6adb60be083de23c6bbf968087 Mon Sep 17 00:00:00 2001 From: Scott Gerring Date: Wed, 1 Jul 2026 16:55:26 +0200 Subject: [PATCH 03/16] move realloc logic into sampler --- libdd-heap-gotter/src/hooks.rs | 75 +++-------------- libdd-heap-gotter/src/lib.rs | 2 - libdd-heap-gotter/src/realloc_math.rs | 36 --------- libdd-heap-sampler/build.rs | 2 + .../include/datadog/heap/allocation_realloc.h | 80 +++++++++++++++++++ libdd-heap-sampler/src/allocation_realloc.c | 71 ++++++++++++++++ libdd-heap-sampler/src/generated/bindings.rs | 34 ++++++++ .../dd_heap_sampler_static_wrappers.c | 1 + 8 files changed, 199 insertions(+), 102 deletions(-) delete mode 100644 libdd-heap-gotter/src/realloc_math.rs create mode 100644 libdd-heap-sampler/include/datadog/heap/allocation_realloc.h create mode 100644 libdd-heap-sampler/src/allocation_realloc.c diff --git a/libdd-heap-gotter/src/hooks.rs b/libdd-heap-gotter/src/hooks.rs index d6fae01345..fae382602e 100644 --- a/libdd-heap-gotter/src/hooks.rs +++ b/libdd-heap-gotter/src/hooks.rs @@ -19,12 +19,10 @@ use core::ffi::{c_char, c_int, c_void}; use std::sync::atomic::{AtomicUsize, Ordering}; use libdd_heap_sampler::{ - dd_alloc_req_t, dd_allocation_created, dd_allocation_freed, dd_allocation_requested, - dd_probe_free, dd_sample_flag_peek, dd_tl_state_get, dd_tl_state_init, + dd_alloc_req_t, dd_allocation_created, dd_allocation_freed, dd_allocation_realloc_commit, + dd_allocation_realloc_prepare, dd_allocation_requested, dd_tl_state_get, dd_tl_state_init, }; -use crate::realloc_math::sampled_realloc_raw_size; - // Per-thread reentry guard for the gotter shims themselves. Distinct // from the sampler's `dd_tl_state_t::reentry_guard` because that one // lives inside a struct we have to *allocate* (via `calloc`) on first @@ -228,72 +226,21 @@ pub unsafe extern "C" fn gotter_realloc(ptr: *mut c_void, size: usize) -> *mut c return real(freed.ptr, 0); } - // Peek is non-destructive: if realloc fails below, the old header - // must stay intact so a later gotter_free(ptr) still resolves the - // right raw pointer for libc. - let mut old_raw = std::ptr::null_mut(); - let mut old_offset = 0usize; - if !dd_sample_flag_peek(ptr, &mut old_raw, &mut old_offset) { - // Case 3: unsampled old allocation. Passthrough. We deliberately - // do not opportunistically sample here: we don't know the old - // user-requested size, so we can't safely place a header and - // move user contents without either over-reading the old block - // or shifting user data. - // TODO: revisit once sampled headers carry the original user - // size, or we have another safe bound for copying old contents. - return real(ptr, size); - } - - // Case 4: sampled old allocation. MVP semantics: + // Cases 3 & 4: delegate to the sampler. `prepare` is + // non-destructive: if the underlying realloc fails, `ptr` stays + // live with its sampler flag intact for a later free. // + // Sampled path (MVP): commit implements // successful sampled realloc = ddheap:free(old sampled) // + new unsampled allocation - // - // Why unsampled on the new side: - // * The new raw pointer picked by the underlying realloc may sit at a different page offset, - // so the sampler's x86 offset picker may want a different `new_offset` than `old_offset`. - // Stamping a header at the new offset before moving user data would overlap the - // still-to-be-copied source region. - // * We also don't know the original user-requested size, so we can't cheaply resize + - // resample + report a coherent alloc event to the profiler. - // - // The overhead of "unsampled from now on" is only paid on - // sampled-then-realloc'd blocks, which are already rare. + // Unsampled path: prep is a passthrough; commit returns new_raw + // unchanged. // // TODO: emit `free + new sampled allocation` when sampled headers // carry the original user size (see project TODO #13 / repo TODO #7). - let Some(raw_size) = sampled_realloc_raw_size(size, old_offset) else { - return std::ptr::null_mut(); - }; - - // Underlying realloc. Contract: on failure, `old_raw` is still - // live and its header is still valid (peek was non-destructive). - let new_raw = real(old_raw, raw_size); - if new_raw.is_null() { - return std::ptr::null_mut(); - } - - // On success, libc copied the old block's bytes into `new_raw` - // starting at offset 0, so the old user data now sits at - // `new_raw + old_offset`. Shift it down to `new_raw` so the caller - // sees the user pointer at offset 0 of an unsampled block. - // memmove (not memcpy) because when the underlying realloc extends - // in place, `new_raw == old_raw` and source/destination overlap. - #[cfg(target_arch = "x86_64")] - { - let copied_user = (new_raw as *mut u8).add(old_offset).cast::(); - if new_raw != copied_user { - libc::memmove(new_raw, copied_user, size); - } - } - - // Fire ddheap:free for the OLD user pointer (the address the - // profiler last saw as live). We call `dd_probe_free` directly, - // not `dd_allocation_freed`, because there's no header left to - // clear: libc has already consumed the old block and we've - // overwritten those bytes above. - dd_probe_free(ptr); - new_raw + let prep = dd_allocation_realloc_prepare(ptr, size); + let new_raw = real(prep.raw_ptr, prep.raw_size); + dd_allocation_realloc_commit(ptr, new_raw, prep) } #[no_mangle] diff --git a/libdd-heap-gotter/src/lib.rs b/libdd-heap-gotter/src/lib.rs index 57b8a039f1..73bad8e8f0 100644 --- a/libdd-heap-gotter/src/lib.rs +++ b/libdd-heap-gotter/src/lib.rs @@ -43,8 +43,6 @@ mod elf; #[cfg(target_os = "linux")] mod hooks; -#[cfg(any(target_os = "linux", test))] -mod realloc_math; #[cfg(target_os = "linux")] use std::sync::Mutex; diff --git a/libdd-heap-gotter/src/realloc_math.rs b/libdd-heap-gotter/src/realloc_math.rs deleted file mode 100644 index 01f9a092cb..0000000000 --- a/libdd-heap-gotter/src/realloc_math.rs +++ /dev/null @@ -1,36 +0,0 @@ -// Copyright 2025-Present Datadog, Inc. https://www.datadoghq.com/ -// SPDX-License-Identifier: Apache-2.0 - -//! Pure helpers for sampled `realloc` layout arithmetic. -//! -//! Split out from `hooks::gotter_realloc` so the offset/size arithmetic -//! can be unit-tested on any host, without pulling in the Linux-only -//! GOT-hook module. - -/// Size to pass to the underlying `realloc` for a sampled old block. -/// -/// The MVP sampled-realloc path returns an *unsampled* pointer at -/// `new_raw` (offset 0), so the underlying block must be large enough -/// to hold both: -/// -/// 1. the old header + slack in `[0, old_offset)` (untouched by libc's realloc copy), -/// 2. `size` bytes of user data at `[old_offset, old_offset + size)`, which we later `memmove` -/// down to `[0, size)`. -/// -/// Hence `size + old_offset`. Returns `None` on overflow so the caller -/// can fail the realloc rather than truncate. -pub(crate) fn sampled_realloc_raw_size(size: usize, old_offset: usize) -> Option { - size.checked_add(old_offset) -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn sampled_realloc_raw_size_reserves_old_offset_for_unsampled_result() { - assert_eq!(sampled_realloc_raw_size(64, 16), Some(80)); - assert_eq!(sampled_realloc_raw_size(64, 64), Some(128)); - assert_eq!(sampled_realloc_raw_size(usize::MAX, 16), None); - } -} diff --git a/libdd-heap-sampler/build.rs b/libdd-heap-sampler/build.rs index 3131b592be..0f7e1962f8 100644 --- a/libdd-heap-sampler/build.rs +++ b/libdd-heap-sampler/build.rs @@ -27,6 +27,7 @@ mod linux { "src/allocation_requested.c", "src/allocation_created.c", "src/allocation_freed.c", + "src/allocation_realloc.c", "src/probes.c", "src/sample_flag.c", "src/tl_state.c", @@ -126,6 +127,7 @@ mod linux { "include/datadog/heap/allocation_requested.h", "include/datadog/heap/allocation_created.h", "include/datadog/heap/allocation_freed.h", + "include/datadog/heap/allocation_realloc.h", "include/datadog/heap/probes.h", "include/datadog/heap/sample_flag.h", "include/datadog/heap/tl_state.h", diff --git a/libdd-heap-sampler/include/datadog/heap/allocation_realloc.h b/libdd-heap-sampler/include/datadog/heap/allocation_realloc.h new file mode 100644 index 0000000000..c0978fba17 --- /dev/null +++ b/libdd-heap-sampler/include/datadog/heap/allocation_realloc.h @@ -0,0 +1,80 @@ +// Copyright 2025-Present Datadog, Inc. https://www.datadoghq.com/ +// SPDX-License-Identifier: Apache-2.0 + +/** + * @file allocation_realloc.h + * + * Sampler-side helpers to wrap an underlying `realloc` call. + * + * Frontend usage: + * + * dd_realloc_prep_t prep = dd_allocation_realloc_prepare(old_user, new_size); + * void *new_raw = real_realloc(prep.raw_ptr, prep.raw_size); + * return dd_allocation_realloc_commit(old_user, new_raw, prep); + * + * The MVP model for a sampled old allocation is: + * + * successful sampled realloc = ddheap:free(old sampled) + * + new unsampled allocation + * + * See gotter's `gotter_realloc` for the full four-case handling + * (`ptr == NULL`, `size == 0`, unsampled-old, sampled-old); this pair + * covers only the last two cases (regular passthrough vs sampled + * teardown), since malloc/free/realloc(NULL,_) / realloc(_,0) map to + * dedicated sampler primitives. + * + * Not paired with dd_allocation_requested / dd_allocation_created: the + * new block is deliberately unsampled to avoid data-corruption hazards + * around header stamping without knowing the original user-requested + * size. Revisit once sampled headers carry the original user size. + */ +#ifndef DD_SAMPLERS_ALLOCATION_REALLOC_H +#define DD_SAMPLERS_ALLOCATION_REALLOC_H + +#include +#include + +/* + * Snapshot of the sampler state around a sampled realloc. + * + * raw_ptr - pointer the frontend MUST pass to the underlying + * realloc. Equal to old_user on the passthrough path. + * raw_size - size the frontend MUST pass to the underlying + * realloc. Equal to new_size on the passthrough path. + * old_offset - byte offset from raw to user in the OLD sampled + * block. Used by commit() to shift user data down after + * realloc succeeds. 0 on the passthrough path. + * was_sampled - true iff old_user was a sampled allocation. commit() + * only runs the extra teardown work when this is true. + */ +typedef struct { + void *raw_ptr; + size_t raw_size; + size_t old_offset; + bool was_sampled; +} dd_realloc_prep_t; + +/* + * Inspect old_user and compute the request to hand to the underlying + * realloc. Non-destructive: does not clear the sampler flag on + * old_user, so if realloc later returns NULL the old allocation stays + * usable and its flag stays intact for the eventual free. + */ +dd_realloc_prep_t dd_allocation_realloc_prepare(void *old_user, size_t new_size); + +/* + * Finalize the realloc. Given the return value of the underlying + * realloc (`new_raw`, may be NULL on failure) and the prep struct from + * prepare(), returns the user-visible pointer to hand to the + * application. + * + * On the sampled path: shifts old user contents from [old_offset, ...) + * down to [0, ...), fires ddheap:free(old_user), and returns new_raw + * as an unsampled pointer. + * + * On the unsampled/passthrough path: returns new_raw unchanged. + * On realloc failure (new_raw == NULL): returns NULL. + */ +void *dd_allocation_realloc_commit(void *old_user, void *new_raw, dd_realloc_prep_t prep); + +#endif diff --git a/libdd-heap-sampler/src/allocation_realloc.c b/libdd-heap-sampler/src/allocation_realloc.c new file mode 100644 index 0000000000..78002b3390 --- /dev/null +++ b/libdd-heap-sampler/src/allocation_realloc.c @@ -0,0 +1,71 @@ +// Copyright 2025-Present Datadog, Inc. https://www.datadoghq.com/ +// SPDX-License-Identifier: Apache-2.0 + +#include +#include +#include + +#include +#include + +dd_realloc_prep_t dd_allocation_realloc_prepare(void *old_user, size_t new_size) { + dd_realloc_prep_t out = { + .raw_ptr = old_user, + .raw_size = new_size, + .old_offset = 0, + .was_sampled = false, + }; + + void *old_raw = NULL; + size_t old_offset = 0; + if (!dd_sample_flag_peek(old_user, &old_raw, &old_offset)) { + return out; /* passthrough: not sampled */ + } + + /* Reserve room for the old header+slack ([0, old_offset)) plus + * `new_size` bytes of user data at [old_offset, old_offset + new_size). + * commit() shifts the user data down to [0, new_size). Overflow -> + * fall back to passthrough with the caller-supplied size; the + * underlying realloc will likely fail with a huge value, but + * nothing gets silently truncated or misinterpreted. */ + if (new_size > SIZE_MAX - old_offset) { + return out; + } + + out.raw_ptr = old_raw; + out.raw_size = new_size + old_offset; + out.old_offset = old_offset; + out.was_sampled = true; + return out; +} + +void *dd_allocation_realloc_commit(void *old_user, void *new_raw, dd_realloc_prep_t prep) { + /* Underlying realloc failed: C says old_user is still live; its + * sampler flag was left intact by prepare(), so a later free() + * will still resolve the right raw pointer. */ + if (new_raw == NULL) return NULL; + + /* Passthrough: nothing to fix up. */ + if (!prep.was_sampled) return new_raw; + + /* Sampled path. libc realloc copied the old block's bytes into + * new_raw starting at index 0, so the old user data now sits at + * new_raw + old_offset. Shift it down to new_raw = user offset 0 + * so we can hand new_raw back as an unsampled pointer. + * + * memmove (not memcpy) because when realloc extends in place, + * new_raw == old_raw and source/destination overlap. */ + size_t user_size = prep.raw_size - prep.old_offset; + char *src = (char *)new_raw + prep.old_offset; + if ((void *)src != new_raw) { + memmove(new_raw, src, user_size); + } + + /* Report the free of the OLD sampled allocation (the address the + * profiler last saw as live). No matching alloc is fired: the new + * block is unsampled. dd_probe_free just emits the ddheap:free + * USDT; there's no header left to clear because libc consumed the + * old block and its bytes have been overwritten by the memmove. */ + dd_probe_free(old_user); + return new_raw; +} diff --git a/libdd-heap-sampler/src/generated/bindings.rs b/libdd-heap-sampler/src/generated/bindings.rs index 5e972d3397..f6da064057 100644 --- a/libdd-heap-sampler/src/generated/bindings.rs +++ b/libdd-heap-sampler/src/generated/bindings.rs @@ -154,6 +154,40 @@ unsafe extern "C" { alignment: usize, ) -> dd_alloc_freed_t; } +#[repr(C)] +#[derive(Debug, Copy, Clone)] +pub struct dd_realloc_prep_t { + pub raw_ptr: *mut ::std::os::raw::c_void, + pub raw_size: usize, + pub old_offset: usize, + pub was_sampled: bool, +} +#[allow(clippy::unnecessary_operation, clippy::identity_op)] +const _: () = { + ["Size of dd_realloc_prep_t"][::std::mem::size_of::() - 32usize]; + ["Alignment of dd_realloc_prep_t"][::std::mem::align_of::() - 8usize]; + ["Offset of field: dd_realloc_prep_t::raw_ptr"] + [::std::mem::offset_of!(dd_realloc_prep_t, raw_ptr) - 0usize]; + ["Offset of field: dd_realloc_prep_t::raw_size"] + [::std::mem::offset_of!(dd_realloc_prep_t, raw_size) - 8usize]; + ["Offset of field: dd_realloc_prep_t::old_offset"] + [::std::mem::offset_of!(dd_realloc_prep_t, old_offset) - 16usize]; + ["Offset of field: dd_realloc_prep_t::was_sampled"] + [::std::mem::offset_of!(dd_realloc_prep_t, was_sampled) - 24usize]; +}; +unsafe extern "C" { + pub fn dd_allocation_realloc_prepare( + old_user: *mut ::std::os::raw::c_void, + new_size: usize, + ) -> dd_realloc_prep_t; +} +unsafe extern "C" { + pub fn dd_allocation_realloc_commit( + old_user: *mut ::std::os::raw::c_void, + new_raw: *mut ::std::os::raw::c_void, + prep: dd_realloc_prep_t, + ) -> *mut ::std::os::raw::c_void; +} unsafe extern "C" { pub fn dd_probe_alloc(user: *mut ::std::os::raw::c_void, size: u64, weight: u64); } diff --git a/libdd-heap-sampler/src/generated/dd_heap_sampler_static_wrappers.c b/libdd-heap-sampler/src/generated/dd_heap_sampler_static_wrappers.c index 805de31ced..31b2ca8197 100644 --- a/libdd-heap-sampler/src/generated/dd_heap_sampler_static_wrappers.c +++ b/libdd-heap-sampler/src/generated/dd_heap_sampler_static_wrappers.c @@ -6,6 +6,7 @@ #include "include/datadog/heap/allocation_requested.h" #include "include/datadog/heap/allocation_created.h" #include "include/datadog/heap/allocation_freed.h" +#include "include/datadog/heap/allocation_realloc.h" #include "include/datadog/heap/probes.h" #include "include/datadog/heap/sample_flag.h" #include "include/datadog/heap/tl_state.h" From a89cfb7979959ec736bf12672b230271736bf7e0 Mon Sep 17 00:00:00 2001 From: Scott Gerring Date: Wed, 1 Jul 2026 17:19:29 +0200 Subject: [PATCH 04/16] make the gotter demo do wilder things with libc allocation APIs --- .../examples/gotter_usdt_demo.rs | 396 ++++++++++++++++-- 1 file changed, 355 insertions(+), 41 deletions(-) diff --git a/libdd-heap-gotter/examples/gotter_usdt_demo.rs b/libdd-heap-gotter/examples/gotter_usdt_demo.rs index df703e1014..e423aa8ad7 100644 --- a/libdd-heap-gotter/examples/gotter_usdt_demo.rs +++ b/libdd-heap-gotter/examples/gotter_usdt_demo.rs @@ -1,24 +1,32 @@ // Copyright 2025-Present Datadog, Inc. https://www.datadoghq.com/ // SPDX-License-Identifier: Apache-2.0 -//! Sample app exercising `libdd-heap-gotter`: install the GOT overrides -//! at startup, then loop producing allocations through the *unmodified* -//! Rust system allocator. Once `install_heap_overrides()` has run, those -//! `malloc`/`free` calls resolve through the patched GOT entries and -//! flow into `libdd-heap-sampler`, firing `ddheap:alloc` / `ddheap:free` -//! USDTs. +//! Sample app exercising `libdd-heap-gotter`: install the GOT overrides, +//! then drive libc's allocator functions directly (`malloc`, `calloc`, +//! `realloc`, `free`, `posix_memalign`, `aligned_alloc`) from multiple +//! threads with a mix of sizes and alignments — including alignments +//! that exceed the sampler's cap and hit the passthrough path. //! -//! Run (Linux, inside the crate's Lima VM, or natively under a container): +//! Every allocation is filled with a deterministic per-allocation byte +//! pattern; every free/realloc verifies that pattern first, so a bad +//! header stamp, wrong memmove, or misplaced offset shows up as a loud +//! panic rather than a silent corruption. +//! +//! Run (Linux): //! ``` //! cargo run --example gotter_usdt_demo -p libdd-heap-gotter //! ``` -//! Add `-- --stress` to keep a CPU core busy enough for CPU profiles: -//! ``` -//! cargo run --example gotter_usdt_demo -p libdd-heap-gotter -- --stress -//! ``` -//! and in another shell, attach a tracer, e.g. +//! Options: +//! * `--stress` — keep one CPU core hot between iterations (for CPU profiles). +//! * `--secs N` — exit after N seconds instead of looping forever. +//! * `--threads N` — number of worker threads (default: 4). +//! +//! Attach a tracer in another shell, e.g.: //! ``` -//! sudo bpftrace -p -e 'usdt:*:ddheap:alloc { printf("alloc %p %d %d\n", arg0, arg1, arg2); }' +//! sudo bpftrace -p -e ' +//! usdt:*:ddheap:alloc { @allocs = count(); } +//! usdt:*:ddheap:free { @frees = count(); } +//! interval:s:1 { print(@allocs); print(@frees); }' //! ``` //! //! The gotter crate is Linux-only; on other targets the example @@ -36,15 +44,285 @@ fn main() { #[cfg(target_os = "linux")] mod linux { use std::hint::black_box; - use std::thread::sleep; + use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; + use std::sync::Arc; + use std::thread; use std::time::{Duration, Instant}; + /// Tracked live allocation: pointer, its user-visible size, and the + /// seed used to fill it. Content integrity is verified before any + /// realloc/free. + struct LiveAlloc { + ptr: *mut u8, + size: usize, + seed: u64, + } + // Raw pointers aren't Send by default; we're the sole owner in the + // worker thread that produced them, so this is fine. + unsafe impl Send for LiveAlloc {} + unsafe impl Sync for LiveAlloc {} + impl LiveAlloc { + fn as_slice_mut(&self) -> &'static mut [u8] { + // SAFETY: `ptr` is the return of a libc allocator and + // `size` is the user-requested size. Lifetime is scoped by + // the caller of `as_slice_mut` and does not outlive the + // allocation. + unsafe { std::slice::from_raw_parts_mut(self.ptr, self.size) } + } + } + + /// Cheap deterministic PRNG (splitmix64). We use it for both + /// scheduling decisions and content fills so a corruption bug + /// reproduces deterministically for a given (thread, seed) pair. + #[derive(Clone, Copy)] + struct Rng(u64); + impl Rng { + fn new(seed: u64) -> Self { + Rng(seed) + } + fn next(&mut self) -> u64 { + let mut z = self.0.wrapping_add(0x9e37_79b9_7f4a_7c15); + self.0 = z; + z = (z ^ (z >> 30)).wrapping_mul(0xbf58_476d_1ce4_e5b9); + z = (z ^ (z >> 27)).wrapping_mul(0x94d0_49bb_1331_11eb); + z ^ (z >> 31) + } + fn range(&mut self, lo: usize, hi: usize) -> usize { + lo + (self.next() as usize) % (hi - lo).max(1) + } + fn choice<'a, T>(&mut self, xs: &'a [T]) -> &'a T { + &xs[(self.next() as usize) % xs.len()] + } + } + + /// Fill `buf` with a pattern derived from `seed`. Verified later by + /// `verify_content` so we notice content shifts or clobbers. + fn fill_content(buf: &mut [u8], seed: u64) { + let mut r = Rng::new(seed); + for chunk in buf.chunks_mut(8) { + let w = r.next().to_le_bytes(); + let n = chunk.len(); + chunk.copy_from_slice(&w[..n]); + } + } + + fn verify_content(buf: &[u8], seed: u64) { + let mut r = Rng::new(seed); + for (i, chunk) in buf.chunks(8).enumerate() { + let w = r.next().to_le_bytes(); + let n = chunk.len(); + assert_eq!( + chunk, + &w[..n], + "content mismatch at byte offset {} (chunk {})", + i * 8, + i, + ); + } + } + + /// Weighted alignment menu. Small alignments dominate; a rare 4096 + /// exercises the sampler's alignment cap; 8192 exceeds it and + /// forces the passthrough path. + const ALIGNMENTS: &[usize] = &[ + 1, 8, 8, 8, 16, 16, 16, 16, 16, 32, 64, 128, 256, 512, 1024, 4096, 8192, + ]; + + fn pick_size(r: &mut Rng) -> usize { + // Log-uniform bucket, then a uniform offset inside the bucket. + // Skewed toward small allocs (matches typical workloads) but + // occasionally reaches into MB territory. + let bucket = r.range(0, 22); + let hi = 1usize << bucket; + let lo = hi / 2; + r.range(lo.max(1), hi.max(2)) + } + + fn pick_alignment(r: &mut Rng) -> usize { + *r.choice(ALIGNMENTS) + } + + unsafe fn do_malloc(size: usize) -> Option { + let ptr = libc::malloc(size) as *mut u8; + if ptr.is_null() { + return None; + } + Some(LiveAlloc { ptr, size, seed: 0 }) + } + + unsafe fn do_calloc(nmemb: usize, size: usize) -> Option { + let ptr = libc::calloc(nmemb, size) as *mut u8; + if ptr.is_null() { + return None; + } + // calloc zeroes memory; verify that before we overwrite with a + // seed pattern. Catches allocator confusion between raw and + // user pointers. + let total = nmemb.saturating_mul(size); + let slice = std::slice::from_raw_parts(ptr, total); + assert!( + slice.iter().all(|&b| b == 0), + "calloc returned non-zeroed memory" + ); + Some(LiveAlloc { + ptr, + size: total, + seed: 0, + }) + } + + unsafe fn do_aligned_alloc(alignment: usize, size: usize) -> Option { + // aligned_alloc requires size % alignment == 0. Round up. + let rounded = size.div_ceil(alignment) * alignment; + let ptr = libc::aligned_alloc(alignment, rounded) as *mut u8; + if ptr.is_null() { + return None; + } + assert_eq!( + (ptr as usize) % alignment, + 0, + "aligned_alloc returned misaligned pointer" + ); + Some(LiveAlloc { + ptr, + size: rounded, + seed: 0, + }) + } + + unsafe fn do_posix_memalign(alignment: usize, size: usize) -> Option { + // posix_memalign requires alignment to be a power of two and a + // multiple of sizeof(void*). + if alignment < std::mem::size_of::<*mut u8>() || !alignment.is_power_of_two() { + return None; + } + let mut out: *mut libc::c_void = std::ptr::null_mut(); + let rc = libc::posix_memalign(&mut out, alignment, size); + if rc != 0 || out.is_null() { + return None; + } + assert_eq!( + (out as usize) % alignment, + 0, + "posix_memalign returned misaligned pointer" + ); + Some(LiveAlloc { + ptr: out as *mut u8, + size, + seed: 0, + }) + } + + unsafe fn do_realloc(old: LiveAlloc, new_size: usize) -> Option { + // Verify old contents before releasing the block. + verify_content(old.as_slice_mut(), old.seed); + let new_ptr = libc::realloc(old.ptr as *mut libc::c_void, new_size) as *mut u8; + if new_ptr.is_null() { + // Old block is still live on realloc failure. Return it as-is. + return Some(old); + } + // Preserved bytes are `min(old.size, new_size)`; verify them + // against the OLD seed. Any misplaced offset in the sampler's + // realloc path shows up as a mismatch here. + let preserved = old.size.min(new_size); + let preserved_slice = std::slice::from_raw_parts(new_ptr, preserved); + // Re-run the deterministic fill to know what those bytes should be. + let mut expected = vec![0u8; preserved]; + fill_content(&mut expected, old.seed); + assert_eq!( + preserved_slice, + &expected[..], + "realloc did not preserve user contents" + ); + Some(LiveAlloc { + ptr: new_ptr, + size: new_size, + seed: 0, + }) + } + + unsafe fn do_free(a: LiveAlloc) { + verify_content(a.as_slice_mut(), a.seed); + libc::free(a.ptr as *mut libc::c_void); + } + + fn worker( + thread_id: u64, + stop: Arc, + allocs: Arc, + frees: Arc, + reallocs: Arc, + ) { + let mut rng = Rng::new(0xdead_beef_0000_0000 ^ thread_id); + // Cap the working set to keep total RSS bounded; when full, + // subsequent alloc ops replace a random slot (freeing it first). + const MAX_LIVE: usize = 256; + let mut live: Vec = Vec::with_capacity(MAX_LIVE); + + while !stop.load(Ordering::Relaxed) { + let op = rng.range(0, 100); + let want_alloc = live.len() < MAX_LIVE / 2 || op < 40; + + unsafe { + if want_alloc { + let size = pick_size(&mut rng); + let alignment = pick_alignment(&mut rng); + let flavor = rng.range(0, 4); + let a = match flavor { + 0 => do_malloc(size), + 1 => { + let nmemb = rng.range(1, 32); + let each = size.div_ceil(nmemb).max(1); + do_calloc(nmemb, each) + } + 2 => do_aligned_alloc(alignment, size.max(alignment)), + _ => do_posix_memalign(alignment, size), + }; + let Some(mut a) = a else { continue }; + a.seed = rng.next(); + fill_content(a.as_slice_mut(), a.seed); + allocs.fetch_add(1, Ordering::Relaxed); + if live.len() == MAX_LIVE { + let idx = rng.range(0, live.len()); + let victim = live.swap_remove(idx); + do_free(victim); + frees.fetch_add(1, Ordering::Relaxed); + } + live.push(a); + } else if !live.is_empty() && op < 80 { + // realloc a random live block. + let idx = rng.range(0, live.len()); + let old = live.swap_remove(idx); + let new_size = pick_size(&mut rng); + let Some(mut resized) = do_realloc(old, new_size) else { + continue; + }; + // Refill with a fresh seed — old bytes past + // min(old_size,new_size) are undefined per the + // realloc contract. + resized.seed = rng.next(); + fill_content(resized.as_slice_mut(), resized.seed); + reallocs.fetch_add(1, Ordering::Relaxed); + live.push(resized); + } else if !live.is_empty() { + let idx = rng.range(0, live.len()); + let victim = live.swap_remove(idx); + do_free(victim); + frees.fetch_add(1, Ordering::Relaxed); + } + } + } + + // Drain remaining live allocations, verifying contents. + while let Some(a) = live.pop() { + unsafe { do_free(a) }; + frees.fetch_add(1, Ordering::Relaxed); + } + } + fn burn_cpu_for(duration: Duration, mut state: u64) -> u64 { let deadline = Instant::now() + duration; while Instant::now() < deadline { - // A tiny integer workload that is deterministic, dependency-chained, - // and opaque to the optimizer. This keeps one CPU core busy without - // changing the allocation profile this example is meant to exercise. state = state.wrapping_mul(0x9e37_79b9_7f4a_7c15).rotate_left(17) ^ 0xbf58_476d_1ce4_e5b9; black_box(state); @@ -53,47 +331,83 @@ mod linux { } pub fn main() { - let stress = std::env::args().any(|arg| arg == "--stress"); + let args: Vec = std::env::args().collect(); + let stress = args.iter().any(|a| a == "--stress"); + let secs: Option = args + .windows(2) + .find(|w| w[0] == "--secs") + .and_then(|w| w[1].parse().ok()); + let threads: usize = args + .windows(2) + .find(|w| w[0] == "--threads") + .and_then(|w| w[1].parse().ok()) + .unwrap_or(4); println!( - "pid={}; pre-install. Attach a tracer on 'usdt:*:ddheap:*'. stress={stress}", - std::process::id() + "pid={}; pre-install. Attach a tracer on 'usdt:*:ddheap:*'. \ + threads={threads} stress={stress} secs={:?}", + std::process::id(), + secs, ); - // Make a few allocations before install to demonstrate the pre-patch - // baseline — no USDTs should fire for these. + // Baseline noise before install — none of this should emit USDTs. { let warmup: Vec = (0..16).map(|i| format!("warmup-{i}")).collect(); println!("pre-install warmup: {} entries", warmup.len()); } - sleep(Duration::from_secs(2)); + std::thread::sleep(Duration::from_secs(2)); - // Install GOT overrides. After this, malloc/free/calloc/realloc - // calls anywhere in the process (including those issued by the Rust - // System allocator backing `Vec`, `String`, etc.) route through - // `libdd-heap-sampler`. let ok = libdd_heap_gotter::install_heap_overrides(); println!("install_heap_overrides() -> {ok}"); - let mut i: u64 = 0; - let mut cpu_state = 0x1234_5678_9abc_def0; - loop { - // ~1000 small allocations + one larger join: plenty of alloc - // pressure to cross the default 512 KiB sampling interval over - // a handful of iterations. - let parts: Vec = (0..1000) - .map(|j| format!("chunk-{i}-{j}-with-some-padding-to-make-it-meaningful")) - .collect(); - let joined = parts.join(", "); - println!("[{i}] joined {} bytes", joined.len()); - i = i.wrapping_add(1); + let stop = Arc::new(AtomicBool::new(false)); + let allocs = Arc::new(AtomicU64::new(0)); + let frees = Arc::new(AtomicU64::new(0)); + let reallocs = Arc::new(AtomicU64::new(0)); + let handles: Vec<_> = (0..threads) + .map(|t| { + let stop = Arc::clone(&stop); + let allocs = Arc::clone(&allocs); + let frees = Arc::clone(&frees); + let reallocs = Arc::clone(&reallocs); + thread::spawn(move || worker(t as u64, stop, allocs, frees, reallocs)) + }) + .collect(); + + let deadline = secs.map(|s| Instant::now() + Duration::from_secs(s)); + let mut cpu_state = 0x1234_5678_9abc_def0u64; + let mut tick = 0u64; + loop { + if let Some(d) = deadline { + if Instant::now() >= d { + break; + } + } if stress { - cpu_state = burn_cpu_for(Duration::from_secs(1), cpu_state ^ i); + cpu_state = burn_cpu_for(Duration::from_secs(1), cpu_state ^ tick); } else { - sleep(Duration::from_secs(1)); + std::thread::sleep(Duration::from_secs(1)); } + tick += 1; + println!( + "[{tick}s] allocs={} frees={} reallocs={}", + allocs.load(Ordering::Relaxed), + frees.load(Ordering::Relaxed), + reallocs.load(Ordering::Relaxed), + ); } + + stop.store(true, Ordering::Relaxed); + for h in handles { + h.join().unwrap(); + } + println!( + "done. allocs={} frees={} reallocs={}", + allocs.load(Ordering::Relaxed), + frees.load(Ordering::Relaxed), + reallocs.load(Ordering::Relaxed), + ); } } // mod linux From 3db26add78853012ff8270ec96ff461e81bd70ce Mon Sep 17 00:00:00 2001 From: Scott Gerring Date: Wed, 1 Jul 2026 17:28:39 +0200 Subject: [PATCH 05/16] address copilot feedback --- libdd-heap-gotter-ffi/tests/install.rs | 2 +- libdd-heap-gotter/README.md | 4 +- libdd-heap-gotter/src/elf.rs | 44 +++++++++++-------- libdd-heap-gotter/src/hooks.rs | 4 +- libdd-heap-gotter/src/lib.rs | 42 +++++++++--------- libdd-heap-gotter/tests/install.rs | 2 +- libdd-heap-sampler/README.md | 6 +-- libdd-heap-sampler/src/allocation_requested.c | 1 + 8 files changed, 58 insertions(+), 47 deletions(-) diff --git a/libdd-heap-gotter-ffi/tests/install.rs b/libdd-heap-gotter-ffi/tests/install.rs index ac0b6a7a39..795ed8f79f 100644 --- a/libdd-heap-gotter-ffi/tests/install.rs +++ b/libdd-heap-gotter-ffi/tests/install.rs @@ -7,7 +7,7 @@ //! as an integration test in its own binary rather than a unit test. //! Miri can't execute the underlying dl_iterate_phdr/mprotect calls. -#![cfg(all(target_os = "linux", not(miri)))] +#![cfg(all(target_os = "linux", target_pointer_width = "64", not(miri)))] use libdd_common_ffi::VoidResult; use libdd_heap_gotter_ffi::{ diff --git a/libdd-heap-gotter/README.md b/libdd-heap-gotter/README.md index b1a297ee17..1a5865eb01 100644 --- a/libdd-heap-gotter/README.md +++ b/libdd-heap-gotter/README.md @@ -6,7 +6,9 @@ This follows the same approach as `ddprof`, and may prove useful to inject via o It contains: * A set of functions such as `gotter_malloc` that will be used to override the _originals_ of these functions -* Overrides for _other bits_ we need for this to work robustly in a running process - including e.g. `fork` +* Overrides for _other bits_ we need for this to work robustly in a running process (`dlopen` to re-scan on new library load, `pthread_create` to materialise sampler TLS on new threads) * A function to install the overrides in a running process `install_heap_overrides()` +Not yet covered: `operator new` / `operator delete`, `mmap`/`munmap`, jemalloc-specific `*allocx` variants, and `pthread_atfork` child-handler for clean `fork()` state reset. + This will be used in places such as the python profiler to install the heap profiler at runtime to capture native allocations. diff --git a/libdd-heap-gotter/src/elf.rs b/libdd-heap-gotter/src/elf.rs index bf30a6c268..1d727e6a9b 100644 --- a/libdd-heap-gotter/src/elf.rs +++ b/libdd-heap-gotter/src/elf.rs @@ -176,19 +176,19 @@ fn gnu_hash(name: &[u8]) -> u32 { unsafe fn gnu_hash_symbol_count(hashtab: *const u32) -> u32 { let nbuckets = *hashtab; - let symbias = *hashtab.add(1); - let bloom_size = *hashtab.add(2); + let symbias = *hashtab.wrapping_add(1); + let bloom_size = *hashtab.wrapping_add(2); // 4 header words + bloom (one Elf64_Addr per entry == 2 u32s) - let mut p = hashtab.add(4 + 2 * bloom_size as usize); + let mut p = hashtab.wrapping_add(4 + 2 * bloom_size as usize); let buckets = std::slice::from_raw_parts(p, nbuckets as usize); - p = p.add(nbuckets as usize); - let chain_zero = p.offset(-(symbias as isize)); + p = p.wrapping_add(nbuckets as usize); + let chain_zero = p.wrapping_offset(-(symbias as isize)); if nbuckets == 0 { return 0; } let mut idx = *buckets.iter().max().unwrap(); - while *chain_zero.add(idx as usize) & 1 == 0 { + while *chain_zero.wrapping_add(idx as usize) & 1 == 0 { idx += 1; } idx + 1 @@ -197,38 +197,39 @@ unsafe fn gnu_hash_symbol_count(hashtab: *const u32) -> u32 { unsafe fn gnu_hash_lookup(info: &DynamicInfo, name: &[u8]) -> Option { let hashtab = info.gnu_hash; let nbuckets = *hashtab; - let symbias = *hashtab.add(1); - let bloom_size = *hashtab.add(2); - let bloom_shift = *hashtab.add(3); - let bloom = hashtab.add(4) as *const u64; - let mut p = hashtab.add(4 + 2 * bloom_size as usize); + let symbias = *hashtab.wrapping_add(1); + let bloom_size = *hashtab.wrapping_add(2); + let bloom_shift = *hashtab.wrapping_add(3); + let bloom = hashtab.wrapping_add(4) as *const u64; + let mut p = hashtab.wrapping_add(4 + 2 * bloom_size as usize); let buckets = p; - p = p.add(nbuckets as usize); - let chain_zero = p.offset(-(symbias as isize)); + p = p.wrapping_add(nbuckets as usize); + let chain_zero = p.wrapping_offset(-(symbias as isize)); if nbuckets == 0 { return None; } let h = gnu_hash(name); - let word = *bloom.add(((h / 64) & (bloom_size - 1)) as usize); + let word = *bloom.wrapping_add(((h / 64) & (bloom_size - 1)) as usize); let bit1 = h & 63; let bit2 = (h >> bloom_shift) & 63; if ((word >> bit1) & (word >> bit2) & 1) == 0 { return None; } - let mut symidx = *buckets.add((h % nbuckets) as usize); + let mut symidx = *buckets.wrapping_add((h % nbuckets) as usize); if symidx == STN_UNDEF { return None; } loop { - let chain_h = *chain_zero.add(symidx as usize); + let chain_h = *chain_zero.wrapping_add(symidx as usize); if ((chain_h ^ h) >> 1) == 0 { if let Some(sname) = info.sym_name(symidx) { - if sname.to_bytes() == name && check_sym(&*info.symtab.add(symidx as usize)) { - return Some(*info.symtab.add(symidx as usize)); + let sym = info.symtab.wrapping_add(symidx as usize); + if sname.to_bytes() == name && check_sym(&*sym) { + return Some(*sym); } } } @@ -349,8 +350,13 @@ struct PageProtGuard { impl PageProtGuard { fn new() -> Self { + // sysconf can return -1 on error; casting that to usize gives + // a huge value that breaks the alignment/mprotect math. Fall + // back to a conservative 4 KiB default if the query fails. + let raw = unsafe { sysconf(_SC_PAGESIZE) }; + let page_size = if raw > 0 { raw as usize } else { 4096 }; Self { - page_size: unsafe { sysconf(_SC_PAGESIZE) as usize }, + page_size, maps: read_proc_maps(), touched: HashMap::new(), } diff --git a/libdd-heap-gotter/src/hooks.rs b/libdd-heap-gotter/src/hooks.rs index fae382602e..be107025cf 100644 --- a/libdd-heap-gotter/src/hooks.rs +++ b/libdd-heap-gotter/src/hooks.rs @@ -171,7 +171,9 @@ pub unsafe extern "C" fn gotter_calloc(nmemb: usize, size: usize) -> *mut c_void return real(nmemb, size); } ensure_tls(); - let total = nmemb.saturating_mul(size); + let Some(total) = nmemb.checked_mul(size) else { + return real(nmemb, size); + }; let req = dd_allocation_requested(total, core::mem::align_of::() * 2); // calloc takes (nmemb, size); when the sampler bumps `req.size` we // funnel the extra bytes into the size argument (nmemb stays 1's diff --git a/libdd-heap-gotter/src/lib.rs b/libdd-heap-gotter/src/lib.rs index 73bad8e8f0..8f458e6d72 100644 --- a/libdd-heap-gotter/src/lib.rs +++ b/libdd-heap-gotter/src/lib.rs @@ -12,8 +12,8 @@ //! //! The public API is available on every platform so downstream code //! never has to `#[cfg]`-guard its callers. The GOT-patching machinery -//! itself only exists on Linux (where `dl_iterate_phdr` + ELF relocs -//! are well-defined); on every other target the entry points compile +//! itself only exists on 64-bit Linux (where `dl_iterate_phdr` + ELF64 +//! relocs are well-defined); on every other target the entry points compile //! to no-ops and `heap_overrides_are_installed()` always returns //! `false`. //! @@ -39,21 +39,21 @@ //! * jemalloc-specific `mallocx`/`dallocx`/etc. //! * `pthread_atfork` child handler to reset state cleanly across `fork()`. -#[cfg(target_os = "linux")] +#[cfg(all(target_os = "linux", target_pointer_width = "64"))] mod elf; -#[cfg(target_os = "linux")] +#[cfg(all(target_os = "linux", target_pointer_width = "64"))] mod hooks; -#[cfg(target_os = "linux")] +#[cfg(all(target_os = "linux", target_pointer_width = "64"))] use std::sync::Mutex; -#[cfg(target_os = "linux")] +#[cfg(all(target_os = "linux", target_pointer_width = "64"))] use elf::SymbolOverrides; /// Holds the SymbolOverrides registry across calls to `install` / `update` /// / `restore`. ddprof keeps the equivalent state in /// `g_symbol_overrides` guarded by `g_mutex` -#[cfg(target_os = "linux")] +#[cfg(all(target_os = "linux", target_pointer_width = "64"))] static GLOBAL_OVERRIDES: Mutex> = Mutex::new(None); /// Install GOT overrides for the supported allocator and helper symbols. @@ -65,10 +65,10 @@ static GLOBAL_OVERRIDES: Mutex> = Mutex::new(None); /// target process has already been statically linked against a custom /// allocator that doesn't appear in the dynamic symbol table. /// -/// On non-Linux targets this is a no-op that always returns `false` — -/// the GOT-patching path it would otherwise execute has no portable -/// equivalent outside ELF + `dl_iterate_phdr`. -#[cfg(target_os = "linux")] +/// On non-64-bit-Linux targets this is a no-op that always returns +/// `false` — the GOT-patching path it would otherwise execute has no +/// portable equivalent outside ELF64 + `dl_iterate_phdr`. +#[cfg(all(target_os = "linux", target_pointer_width = "64"))] pub fn install_heap_overrides() -> bool { let mut guard = GLOBAL_OVERRIDES.lock().expect("gotter mutex poisoned"); if guard.is_none() { @@ -83,7 +83,7 @@ pub fn install_heap_overrides() -> bool { } /// See the Linux variant above. -#[cfg(not(target_os = "linux"))] +#[cfg(not(all(target_os = "linux", target_pointer_width = "64")))] pub fn install_heap_overrides() -> bool { false } @@ -91,7 +91,7 @@ pub fn install_heap_overrides() -> bool { /// Re-scan loaded libraries and patch any newly-introduced GOT entries. /// Called automatically from the `dlopen` hook; user code typically /// doesn't need to call this directly. No-op on non-Linux targets. -#[cfg(target_os = "linux")] +#[cfg(all(target_os = "linux", target_pointer_width = "64"))] pub fn update_heap_overrides() { // `try_lock` so a dlopen happening on the same thread that owns the // install lock doesn't deadlock - that thread will finish its @@ -104,13 +104,13 @@ pub fn update_heap_overrides() { } /// See the Linux variant above. -#[cfg(not(target_os = "linux"))] +#[cfg(not(all(target_os = "linux", target_pointer_width = "64")))] pub fn update_heap_overrides() {} /// Revert every GOT entry we patched. After this call, the process is /// once again calling the real allocator symbols directly. No-op on /// non-Linux targets. -#[cfg(target_os = "linux")] +#[cfg(all(target_os = "linux", target_pointer_width = "64"))] pub fn restore_heap_overrides() { let mut guard = GLOBAL_OVERRIDES.lock().expect("gotter mutex poisoned"); if let Some(so) = guard.as_mut() { @@ -120,13 +120,13 @@ pub fn restore_heap_overrides() { } /// See the Linux variant above. -#[cfg(not(target_os = "linux"))] +#[cfg(not(all(target_os = "linux", target_pointer_width = "64")))] pub fn restore_heap_overrides() {} /// Return whether heap GOT overrides are currently installed. Always /// returns `false` on non-Linux targets, since `install_heap_overrides` /// is a no-op there. -#[cfg(target_os = "linux")] +#[cfg(all(target_os = "linux", target_pointer_width = "64"))] pub fn heap_overrides_are_installed() -> bool { GLOBAL_OVERRIDES .lock() @@ -135,13 +135,13 @@ pub fn heap_overrides_are_installed() -> bool { } /// See the Linux variant above. -#[cfg(not(target_os = "linux"))] +#[cfg(not(all(target_os = "linux", target_pointer_width = "64")))] pub fn heap_overrides_are_installed() -> bool { false } /// Register GOT overrides for every symbol this crate currently hooks. -#[cfg(target_os = "linux")] +#[cfg(all(target_os = "linux", target_pointer_width = "64"))] fn register_all(so: &mut SymbolOverrides) { use hooks::*; use std::sync::atomic::AtomicUsize; @@ -200,7 +200,7 @@ fn register_all(so: &mut SymbolOverrides) { ); } -#[cfg(target_os = "linux")] +#[cfg(all(target_os = "linux", target_pointer_width = "64"))] fn any_orig_resolved() -> bool { use hooks::*; use std::sync::atomic::Ordering; @@ -221,7 +221,7 @@ fn any_orig_resolved() -> bool { // Tests call into the ELF symbol-lookup path (dl_iterate_phdr + // dynsym parsing of loaded libraries) which miri can't execute, so // skip the whole module under miri. -#[cfg(all(test, target_os = "linux", not(miri)))] +#[cfg(all(test, target_os = "linux", target_pointer_width = "64", not(miri)))] mod tests { use super::*; diff --git a/libdd-heap-gotter/tests/install.rs b/libdd-heap-gotter/tests/install.rs index 4cf496aacb..8a6e813c13 100644 --- a/libdd-heap-gotter/tests/install.rs +++ b/libdd-heap-gotter/tests/install.rs @@ -9,7 +9,7 @@ // Integration test invokes the GOT-patching machinery for real // (dl_iterate_phdr + dlsym + mprotect), which miri can't execute. -#![cfg(all(target_os = "linux", not(miri)))] +#![cfg(all(target_os = "linux", target_pointer_width = "64", not(miri)))] use std::ffi::c_void; diff --git a/libdd-heap-sampler/README.md b/libdd-heap-sampler/README.md index 78c82e99fe..9e277dbc49 100644 --- a/libdd-heap-sampler/README.md +++ b/libdd-heap-sampler/README.md @@ -13,10 +13,10 @@ This profiling infrastructure will initially support these two use cases: **Rust compile-time instrumentation** A crate exposing a `GlobalAlloc` implementation will allow Rust users to, at compile time, opt into heap profiling. We anticipate this will be shipped as a feature of `dd-trace-rs`. This addresses a pain point both internally with the -increase in Rust adoption in Datadog services, and would adress the same pain point within the broader Rust community. +increase in Rust adoption in Datadog services, and would address the same pain point within the broader Rust community. **Python Runtime Instrumentation for _native_ library allocation sampling** -Today our python profiler cannot sample allocations occuring behind the FFI; we will extend `ddtrace-py` to load our dynamic +Today our python profiler cannot sample allocations occurring behind the FFI; we will extend `ddtrace-py` to load our dynamic runtime patching mechanism such that, as native libraries are loaded, we intercept their allocators. As many popular python libraries function largely as API glue around native libraries this will help close the allocation observability gap. @@ -56,7 +56,7 @@ The actual USDTs emitted are: By splitting into `requested` and `created`, these are designed to be generic across different allocation functions (e.g. `malloc`, `operator new`, `aligned_alloc`, etc.). The job of binding these back to concrete callsites in a -process is left to the other components - e.g. `libddd-heap-gotter`, `libdd-heap-allocator`, etc. +process is left to the other components - e.g. `libdd-heap-gotter`, `libdd-heap-allocator`, etc. The allocation-side pair is declared `static inline __attribute__((always_inline))` so the non-sampled fast path inlines into the wrapper with no function-call overhead. diff --git a/libdd-heap-sampler/src/allocation_requested.c b/libdd-heap-sampler/src/allocation_requested.c index 92408feaba..4dc0459340 100644 --- a/libdd-heap-sampler/src/allocation_requested.c +++ b/libdd-heap-sampler/src/allocation_requested.c @@ -48,6 +48,7 @@ static uint64_t next_interval(uint32_t *rng, uint64_t mean) { */ static uint64_t sample(dd_tl_state_t *tl) { uint64_t interval = tl->sampling_interval; + if (interval == 0) return 0; if (!tl->remaining_bytes_initialized) { /* First allocation on this thread: draw the initial interval and From c051bd7f432ccdab1581a97f658670ba9222845c Mon Sep 17 00:00:00 2001 From: Scott Gerring Date: Thu, 2 Jul 2026 14:21:20 +0200 Subject: [PATCH 06/16] address jb robot army review comments --- Cargo.lock | 1 + libdd-heap-gotter/Cargo.toml | 2 + libdd-heap-gotter/src/elf.rs | 325 +++++++++++++++--- libdd-heap-gotter/src/hooks.rs | 95 +---- libdd-heap-gotter/src/lib.rs | 29 +- libdd-heap-gotter/tests/install.rs | 179 ++++++++-- libdd-heap-sampler/README.md | 38 +- libdd-heap-sampler/docs/realloc.md | 66 ++++ libdd-heap-sampler/docs/tagging.md | 102 ++++++ .../include/datadog/heap/sample_flag.h | 3 + .../include/datadog/heap/tl_state.h | 4 +- libdd-heap-sampler/src/allocation_freed.c | 38 +- libdd-heap-sampler/src/allocation_requested.c | 3 +- libdd-heap-sampler/src/lib.rs | 75 ++++ 14 files changed, 755 insertions(+), 205 deletions(-) create mode 100644 libdd-heap-sampler/docs/realloc.md create mode 100644 libdd-heap-sampler/docs/tagging.md diff --git a/Cargo.lock b/Cargo.lock index ccc9d77d4d..fcee1c9201 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3111,6 +3111,7 @@ version = "0.1.0" dependencies = [ "libc", "libdd-heap-sampler", + "serial_test", ] [[package]] diff --git a/libdd-heap-gotter/Cargo.toml b/libdd-heap-gotter/Cargo.toml index c15a334220..9662e61620 100644 --- a/libdd-heap-gotter/Cargo.toml +++ b/libdd-heap-gotter/Cargo.toml @@ -22,3 +22,5 @@ libdd-heap-sampler = { path = "../libdd-heap-sampler" } [target.'cfg(target_os = "linux")'.dev-dependencies] libc = "0.2" +libdd-heap-sampler = { path = "../libdd-heap-sampler" } +serial_test = "3.2" diff --git a/libdd-heap-gotter/src/elf.rs b/libdd-heap-gotter/src/elf.rs index 1d727e6a9b..d0198de7aa 100644 --- a/libdd-heap-gotter/src/elf.rs +++ b/libdd-heap-gotter/src/elf.rs @@ -22,8 +22,8 @@ use std::collections::HashMap; use std::ffi::CStr; use libc::{ - dl_iterate_phdr, dl_phdr_info, mprotect, sysconf, Elf64_Phdr, Elf64_Rel, Elf64_Rela, Elf64_Sym, - PROT_EXEC, PROT_READ, PROT_WRITE, PT_DYNAMIC, PT_LOAD, _SC_PAGESIZE, + dl_iterate_phdr, dl_phdr_info, mprotect, sysconf, Elf64_Rel, Elf64_Rela, Elf64_Sym, PROT_EXEC, + PROT_READ, PROT_WRITE, PT_DYNAMIC, PT_LOAD, _SC_PAGESIZE, }; use std::io::{BufRead, BufReader}; use std::sync::atomic::{AtomicUsize, Ordering}; @@ -64,6 +64,7 @@ struct DynamicInfo { jmprels: *const Elf64_Rela, jmprels_count: usize, gnu_hash: *const u32, + gnu_hash_words: usize, base_address: usize, } @@ -78,6 +79,13 @@ impl DynamicInfo { let dyn_phdr = phdrs.iter().find(|p| p.p_type == PT_DYNAMIC)?; let dyn_begin = (info.dlpi_addr as usize + dyn_phdr.p_vaddr as usize) as *const Elf64_Dyn; let base = info.dlpi_addr as usize; + let containing_load_segment_end = |addr: usize| -> Option { + phdrs.iter().filter(|p| p.p_type == PT_LOAD).find_map(|p| { + let start = base.checked_add(p.p_vaddr as usize)?; + let end = start.checked_add(p.p_memsz as usize)?; + (addr >= start && addr < end).then_some(end) + }) + }; let correct = |a: u64| -> usize { let a = a as usize; if a > base { @@ -133,7 +141,30 @@ impl DynamicInfo { return None; } - let sym_count = gnu_hash_symbol_count(gnu_hash); + let gnu_hash_addr = gnu_hash as usize; + let gnu_hash_words = match containing_load_segment_end(gnu_hash_addr) { + Some(end) => match end.checked_sub(gnu_hash_addr) { + Some(bytes) => bytes / core::mem::size_of::(), + None => return None, + }, + None => return None, + }; + let sym_count = gnu_hash_symbol_count(gnu_hash, gnu_hash_words).unwrap_or_else(|| { + // Fallback for degenerate .gnu.hash (e.g. executables with only + // undefined imports): estimate dynsym entry count from the + // distance between DT_SYMTAB and DT_STRTAB. This works because + // the linker always places .dynsym immediately before .dynstr. + let symtab_addr = symtab as usize; + let strtab_addr = strtab as usize; + if strtab_addr > symtab_addr { + let bytes = strtab_addr - symtab_addr; + (bytes / core::mem::size_of::()) as u32 + } else { + // Can't estimate; allow any index and rely on strtab + // bounds checking in sym_name to catch bad accesses. + u32::MAX + } + }); Some(Self { strtab, @@ -147,6 +178,7 @@ impl DynamicInfo { jmprels, jmprels_count: jmprels_size / core::mem::size_of::(), gnu_hash, + gnu_hash_words, base_address: base, }) } @@ -174,60 +206,108 @@ fn gnu_hash(name: &[u8]) -> u32 { h } -unsafe fn gnu_hash_symbol_count(hashtab: *const u32) -> u32 { +/// Compute the total number of entries in `.dynsym` from the `.gnu.hash` table. +/// +/// Returns `None` only when the table is structurally invalid. When the hash +/// is degenerate (all buckets empty, typical for executables that only import +/// symbols), returns `None` to signal that the caller should use a fallback +/// (e.g. estimate from symtab/strtab distance) since the hash table doesn't +/// tell us how many undefined-import entries precede the hashed region. +unsafe fn gnu_hash_symbol_count(hashtab: *const u32, hashtab_words: usize) -> Option { + if hashtab_words < 4 { + return None; + } + let nbuckets = *hashtab; - let symbias = *hashtab.wrapping_add(1); - let bloom_size = *hashtab.wrapping_add(2); - // 4 header words + bloom (one Elf64_Addr per entry == 2 u32s) - let mut p = hashtab.wrapping_add(4 + 2 * bloom_size as usize); - let buckets = std::slice::from_raw_parts(p, nbuckets as usize); - p = p.wrapping_add(nbuckets as usize); - let chain_zero = p.wrapping_offset(-(symbias as isize)); + let symbias = *hashtab.add(1); + let bloom_size = *hashtab.add(2); + let bloom_size_words = (bloom_size as usize).checked_mul(2)?; + let buckets_start = 4usize.checked_add(bloom_size_words)?; + let chains_start = buckets_start.checked_add(nbuckets as usize)?; + if bloom_size == 0 || buckets_start > hashtab_words || chains_start > hashtab_words { + return None; + } if nbuckets == 0 { - return 0; + // No buckets at all: can't determine symtab size from the hash. + return None; + } + + let buckets = std::slice::from_raw_parts(hashtab.add(buckets_start), nbuckets as usize); + let mut idx = *buckets.iter().max()?; + if idx == STN_UNDEF { + // All buckets empty: hash covers zero defined symbols, but the + // symtab may still have undefined imports. Signal the caller to + // use a fallback. + return None; + } + if idx < symbias { + return None; } - let mut idx = *buckets.iter().max().unwrap(); - while *chain_zero.wrapping_add(idx as usize) & 1 == 0 { - idx += 1; + + let chain_count = hashtab_words - chains_start; + loop { + let chain_idx = (idx - symbias) as usize; + if chain_idx >= chain_count { + return None; + } + if *hashtab.add(chains_start + chain_idx) & 1 != 0 { + return idx.checked_add(1); + } + idx = idx.checked_add(1)?; } - idx + 1 } unsafe fn gnu_hash_lookup(info: &DynamicInfo, name: &[u8]) -> Option { let hashtab = info.gnu_hash; - let nbuckets = *hashtab; - let symbias = *hashtab.wrapping_add(1); - let bloom_size = *hashtab.wrapping_add(2); - let bloom_shift = *hashtab.wrapping_add(3); - let bloom = hashtab.wrapping_add(4) as *const u64; - let mut p = hashtab.wrapping_add(4 + 2 * bloom_size as usize); - let buckets = p; - p = p.wrapping_add(nbuckets as usize); - let chain_zero = p.wrapping_offset(-(symbias as isize)); + if info.gnu_hash_words < 4 { + return None; + } - if nbuckets == 0 { + let nbuckets = *hashtab; + let symbias = *hashtab.add(1); + let bloom_size = *hashtab.add(2); + let bloom_shift = *hashtab.add(3); + let bloom_size_words = (bloom_size as usize).checked_mul(2)?; + let buckets_start = 4usize.checked_add(bloom_size_words)?; + let chains_start = buckets_start.checked_add(nbuckets as usize)?; + + if nbuckets == 0 + || bloom_size == 0 + || buckets_start > info.gnu_hash_words + || chains_start > info.gnu_hash_words + { return None; } let h = gnu_hash(name); - let word = *bloom.wrapping_add(((h / 64) & (bloom_size - 1)) as usize); + let bloom = hashtab.add(4) as *const u64; + let word = *bloom.add(((h / 64) & (bloom_size - 1)) as usize); let bit1 = h & 63; let bit2 = (h >> bloom_shift) & 63; if ((word >> bit1) & (word >> bit2) & 1) == 0 { return None; } - let mut symidx = *buckets.wrapping_add((h % nbuckets) as usize); + let buckets = hashtab.add(buckets_start); + let mut symidx = *buckets.add((h % nbuckets) as usize); if symidx == STN_UNDEF { return None; } + if symidx < symbias { + return None; + } + let chain_count = info.gnu_hash_words - chains_start; loop { - let chain_h = *chain_zero.wrapping_add(symidx as usize); + let chain_idx = (symidx - symbias) as usize; + if chain_idx >= chain_count { + return None; + } + let chain_h = *hashtab.add(chains_start + chain_idx); if ((chain_h ^ h) >> 1) == 0 { if let Some(sname) = info.sym_name(symidx) { - let sym = info.symtab.wrapping_add(symidx as usize); + let sym = info.symtab.add(symidx as usize); if sname.to_bytes() == name && check_sym(&*sym) { return Some(*sym); } @@ -236,7 +316,7 @@ unsafe fn gnu_hash_lookup(info: &DynamicInfo, name: &[u8]) -> Option if chain_h & 1 != 0 { break; } - symidx += 1; + symidx = symidx.checked_add(1)?; } None } @@ -270,13 +350,18 @@ fn iterate_libraries bool>(mut cb: F) { _size: libc::size_t, data: *mut c_void, ) -> c_int { - let ctx = &mut *(data as *mut Ctx); - let is_exe = ctx.is_first; - ctx.is_first = false; - if (ctx.cb)(&*info, is_exe) { - 1 - } else { - 0 + let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { + let ctx = &mut *(data as *mut Ctx); + let is_exe = ctx.is_first; + ctx.is_first = false; + (ctx.cb)(&*info, is_exe) + })); + + match result { + Ok(stop) => i32::from(stop), + // Never unwind a Rust panic through libc's dl_iterate_phdr callback. + // Treat patching as best-effort and stop iteration on panic. + Err(_) => 1, } } @@ -285,6 +370,33 @@ fn iterate_libraries bool>(mut cb: F) { } } +unsafe fn library_name(info: &dl_phdr_info) -> String { + if info.dlpi_name.is_null() { + String::new() + } else { + CStr::from_ptr(info.dlpi_name) + .to_string_lossy() + .into_owned() + } +} + +fn loaded_library_name_at_base(base: usize) -> Option { + let mut found = None; + iterate_libraries(|info, _| unsafe { + if info.dlpi_addr as usize == base { + found = Some(library_name(info)); + true + } else { + false + } + }); + found +} + +fn should_restore_library(expected_name: &str, current_name: Option<&str>) -> bool { + current_name == Some(expected_name) +} + /// A single /proc/self/maps entry: address range + current protection flags. #[derive(Clone, Copy)] struct MapEntry { @@ -563,6 +675,12 @@ impl SymbolOverrides { v.processed = false; } + // TODO: This is intentionally simple but expensive on workloads that + // dlopen many libraries: every change re-walks all loaded objects, + // re-parses their dynamic sections/GNU hash tables, and eagerly reads + // /proc/self/maps via PageProtGuard even if only one new object needs + // patching. Track already-processed libraries and lazily create the + // page-protection guard to avoid repeated heavy work. let mut guard = PageProtGuard::new(); // SAFETY: closure runs synchronously inside dl_iterate_phdr. @@ -595,15 +713,39 @@ impl SymbolOverrides { /// Restore every GOT entry we touched. pub fn restore_overrides(&mut self) { + self.restore_overrides_with_lookup(loaded_library_name_at_base); + } + + fn restore_overrides_with_lookup( + &mut self, + mut loaded_name_at_base: impl FnMut(usize) -> Option, + ) -> usize { let info_per_lib = std::mem::take(&mut self.revert_info_per_library); - let mut guard = PageProtGuard::new(); - for (_base, revert) in info_per_lib { + let mut guard = None; + let mut restored_libraries = 0; + for (base, revert) in info_per_lib { + // A tracked library may have been dlclose'd since the last update, + // and its address range may even have been reused by a different + // mapping. Only write old GOT values back when the same library is + // still loaded at the original base address. + if !should_restore_library( + revert.library_name.as_str(), + loaded_name_at_base(base).as_deref(), + ) { + continue; + } + + restored_libraries += 1; + let guard = guard.get_or_insert_with(PageProtGuard::new); for (addr, old) in revert.old_value_per_address { unsafe { guard.override_entry(addr, old) }; } } - guard.finish(); + if let Some(guard) = guard { + guard.finish(); + } self.last_seen_nb_libs = -1; + restored_libraries } unsafe fn apply_to_library( @@ -696,13 +838,15 @@ impl SymbolOverrides { guard: &mut PageProtGuard, ) { // st_name -> string in strtab. Walk lazily: we look up the - // name in the override map; if it's not there, skip. - let sym = &*dyn_info.symtab.add(sym_index as usize); - let name_off = sym.st_name as usize; - if name_off == 0 || name_off >= dyn_info.strtab_size { + // name in the override map; if it's not there, skip. Relocation + // symbol indices come from the object being inspected, so guard + // them before dereferencing dyn_info.symtab. + let Some(cstr) = dyn_info.sym_name(sym_index) else { + return; + }; + if cstr.to_bytes().is_empty() { return; } - let cstr = CStr::from_ptr(dyn_info.strtab.add(name_off)); let Ok(name) = cstr.to_str() else { return }; let Some(ov) = overrides.get(name) else { @@ -732,12 +876,85 @@ fn elf64_r_sym(info: u64) -> u64 { info >> 32 } -// Phdr/PT_LOAD ranges are unused for now; kept for future "skip self -// library" logic. Silence unused warnings. -#[allow(dead_code)] -const _: () = { - let _ = PT_LOAD; -}; +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn restore_library_requires_same_live_name() { + assert!(should_restore_library( + "/tmp/libfoo.so", + Some("/tmp/libfoo.so") + )); + assert!(!should_restore_library("/tmp/libfoo.so", None)); + assert!(!should_restore_library( + "/tmp/libfoo.so", + Some("/tmp/libbar.so") + )); + } + + #[test] + fn page_prot_guard_finds_original_mapping_protection() { + let guard = PageProtGuard { + page_size: 4096, + maps: vec![ + MapEntry { + start: 0x1000, + end: 0x2000, + prot: PROT_READ, + }, + MapEntry { + start: 0x2000, + end: 0x3000, + prot: PROT_READ | PROT_EXEC, + }, + ], + touched: HashMap::new(), + }; + + assert_eq!(guard.original_prot(0x1000), Some(PROT_READ)); + assert_eq!(guard.original_prot(0x1fff), Some(PROT_READ)); + assert_eq!(guard.original_prot(0x2000), Some(PROT_READ | PROT_EXEC)); + assert_eq!(guard.original_prot(0x3000), None); + } -#[allow(dead_code)] -fn _phdr_marker(_: Elf64_Phdr) {} + #[test] + fn restore_overrides_skips_unloaded_or_reused_libraries() { + let mut overrides = SymbolOverrides::new(); + overrides.revert_info_per_library.insert( + 0x1000, + LibraryRevertInfo { + library_name: "/tmp/libfoo.so".to_string(), + old_value_per_address: HashMap::new(), + processed: true, + }, + ); + overrides.revert_info_per_library.insert( + 0x2000, + LibraryRevertInfo { + library_name: "/tmp/libbar.so".to_string(), + old_value_per_address: HashMap::new(), + processed: true, + }, + ); + overrides.revert_info_per_library.insert( + 0x3000, + LibraryRevertInfo { + library_name: "/tmp/libbaz.so".to_string(), + old_value_per_address: HashMap::new(), + processed: true, + }, + ); + + let restored = overrides.restore_overrides_with_lookup(|base| match base { + 0x1000 => Some("/tmp/libfoo.so".to_string()), + 0x2000 => Some("/tmp/different.so".to_string()), + 0x3000 => None, + _ => unreachable!(), + }); + + assert_eq!(restored, 1); + assert!(overrides.revert_info_per_library.is_empty()); + assert_eq!(overrides.last_seen_nb_libs, -1); + } +} diff --git a/libdd-heap-gotter/src/hooks.rs b/libdd-heap-gotter/src/hooks.rs index be107025cf..5891f7c5b1 100644 --- a/libdd-heap-gotter/src/hooks.rs +++ b/libdd-heap-gotter/src/hooks.rs @@ -7,58 +7,20 @@ //! 1. Calls `dd_allocation_requested` (or skips, for free-side hooks). //! 2. Forwards to the real symbol via its `ORIG_*` slot, which the install path fills in by //! symbol-table lookup. -//! 3. Calls `dd_allocation_created` / `dd_allocation_freed` to fire the USDT and close the reentry -//! guard. +//! 3. Calls `dd_allocation_created` / `dd_allocation_freed` to fire the USDT. //! //! Modeled on ddprof `src/lib/symbol_overrides.cc`, minus the C++ allocator //! family (operator new/delete) and the mmap/munmap pair which aren't //! supported by the sampler yet. -use core::cell::Cell; use core::ffi::{c_char, c_int, c_void}; use std::sync::atomic::{AtomicUsize, Ordering}; use libdd_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_get, dd_tl_state_init, + dd_allocation_realloc_prepare, dd_allocation_requested, dd_tl_state_init, }; -// Per-thread reentry guard for the gotter shims themselves. Distinct -// from the sampler's `dd_tl_state_t::reentry_guard` because that one -// lives inside a struct we have to *allocate* (via `calloc`) on first -// touch - and on a freshly-installed gotter that first `calloc` lands -// right back in `gotter_calloc`. So before we look at the sampler TLS -// at all, we set this flag; any reentry while it's set forwards -// straight through to the real allocator with no sampling. -// -// `const { Cell::new(false) }` keeps the TLS slot lazy-init-free - -// macOS and glibc both initialise it without an allocation. -std::thread_local! { - static IN_HOOK: Cell = const { Cell::new(false) }; -} - -struct GotterReentry(bool); - -impl GotterReentry { - fn enter() -> Self { - let was = IN_HOOK.with(|c| { - let prev = c.get(); - c.set(true); - prev - }); - GotterReentry(was) - } - fn reentered(&self) -> bool { - self.0 - } -} - -impl Drop for GotterReentry { - fn drop(&mut self) { - IN_HOOK.with(|c| c.set(self.0)); - } -} - /// Resolved address of the real `malloc`; filled by `install_heap_overrides`. /// The bare `usize` payload is a function pointer. pub(crate) static ORIG_MALLOC: AtomicUsize = AtomicUsize::new(0); @@ -77,19 +39,6 @@ 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); -/// Ensure the sampler's per-thread state exists before recording an allocation. -#[inline] -fn ensure_tls() { - unsafe { - if dd_tl_state_get().is_null() { - // dd_tl_state_init calls calloc internally; the gotter - // reentry guard around the caller stops us from re-entering - // this path through our own gotter_calloc hook. - dd_tl_state_init(); - } - } -} - /// Load a resolved function pointer from one of the `ORIG_*` slots. #[inline] unsafe fn load_fn(slot: &AtomicUsize) -> Option { @@ -133,11 +82,6 @@ pub unsafe extern "C" fn gotter_malloc(size: usize) -> *mut c_void { let Some(real): Option = load_fn(&ORIG_MALLOC) else { return std::ptr::null_mut(); }; - let guard = GotterReentry::enter(); - if guard.reentered() { - return real(size); - } - ensure_tls(); // Default alignment for malloc on glibc is 2*sizeof(void*) == 16. let req = dd_allocation_requested(size, core::mem::align_of::() * 2); let raw = real(req.size); @@ -152,11 +96,6 @@ pub unsafe extern "C" fn gotter_free(ptr: *mut c_void) { if ptr.is_null() { return; } - let guard = GotterReentry::enter(); - if guard.reentered() { - real(ptr); - return; - } let freed = dd_allocation_freed(ptr, 0, 0); real(freed.ptr); } @@ -166,11 +105,6 @@ pub unsafe extern "C" fn gotter_calloc(nmemb: usize, size: usize) -> *mut c_void let Some(real): Option = load_fn(&ORIG_CALLOC) else { return std::ptr::null_mut(); }; - let guard = GotterReentry::enter(); - if guard.reentered() { - return real(nmemb, size); - } - ensure_tls(); let Some(total) = nmemb.checked_mul(size) else { return real(nmemb, size); }; @@ -201,17 +135,16 @@ pub unsafe extern "C" fn gotter_calloc(nmemb: usize, size: usize) -> *mut c_void /// (see TODO on the branch). /// 4. **`ptr` is sampled** — MVP: model a successful sampled realloc as `free(old sampled)` /// followed by a new *unsampled* allocation. See the branch comment for the full rationale. +/// +/// This is a lot of logic that might end up being duplicated across other interception mechanisms +/// on top of libdd-heap-sampler. If that _is_ the case, we should consider adding a special case +/// `realloc` to the underlying sampler library and using that here; we've not done it for now to +/// preserve the general pre/post pattern used by the other functions. #[no_mangle] pub unsafe extern "C" fn gotter_realloc(ptr: *mut c_void, size: usize) -> *mut c_void { let Some(real): Option = load_fn(&ORIG_REALLOC) else { return std::ptr::null_mut(); }; - let guard = GotterReentry::enter(); - if guard.reentered() { - return real(ptr, size); - } - ensure_tls(); - // Case 1: realloc(NULL, size) == malloc(size). Normal sampling path. if ptr.is_null() { let alignment = core::mem::align_of::() * 2; @@ -254,11 +187,6 @@ pub unsafe extern "C" fn gotter_posix_memalign( let Some(real): Option = load_fn(&ORIG_POSIX_MEMALIGN) else { return libc::ENOMEM; }; - let guard = GotterReentry::enter(); - if guard.reentered() { - return real(memptr, alignment, size); - } - ensure_tls(); let req = dd_allocation_requested(size, alignment); let ret = real(memptr, alignment, req.size); // Always pair with dd_allocation_created, even on failure, so the @@ -279,11 +207,6 @@ pub unsafe extern "C" fn gotter_aligned_alloc(alignment: usize, size: usize) -> let Some(real): Option = load_fn(&ORIG_ALIGNED_ALLOC) else { return std::ptr::null_mut(); }; - let guard = GotterReentry::enter(); - if guard.reentered() { - return real(alignment, size); - } - ensure_tls(); let req = dd_allocation_requested(size, alignment); let raw = real(alignment, req.size); dd_allocation_created(raw, req) @@ -298,7 +221,9 @@ pub unsafe extern "C" fn gotter_dlopen(filename: *const c_char, flags: c_int) -> }; let handle = real(filename, flags); // New library may have introduced new GOT entries that need patching. - crate::update_heap_overrides(); + // This hook is an extern "C" boundary, so never let a Rust panic from + // best-effort ELF parsing/GOT patching unwind into the caller. + let _ = std::panic::catch_unwind(crate::update_heap_overrides); handle } diff --git a/libdd-heap-gotter/src/lib.rs b/libdd-heap-gotter/src/lib.rs index 8f458e6d72..49a05b41d7 100644 --- a/libdd-heap-gotter/src/lib.rs +++ b/libdd-heap-gotter/src/lib.rs @@ -45,7 +45,7 @@ mod elf; mod hooks; #[cfg(all(target_os = "linux", target_pointer_width = "64"))] -use std::sync::Mutex; +use std::sync::{Mutex, MutexGuard, TryLockError}; #[cfg(all(target_os = "linux", target_pointer_width = "64"))] use elf::SymbolOverrides; @@ -56,6 +56,13 @@ use elf::SymbolOverrides; #[cfg(all(target_os = "linux", target_pointer_width = "64"))] static GLOBAL_OVERRIDES: Mutex> = Mutex::new(None); +#[cfg(all(target_os = "linux", target_pointer_width = "64"))] +fn lock_global_overrides() -> MutexGuard<'static, Option> { + GLOBAL_OVERRIDES + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) +} + /// Install GOT overrides for the supported allocator and helper symbols. /// Safe to call more than once: the registry is rebuilt and re-applied, /// which also picks up any libraries loaded since the last call. @@ -70,7 +77,7 @@ static GLOBAL_OVERRIDES: Mutex> = Mutex::new(None); /// portable equivalent outside ELF64 + `dl_iterate_phdr`. #[cfg(all(target_os = "linux", target_pointer_width = "64"))] pub fn install_heap_overrides() -> bool { - let mut guard = GLOBAL_OVERRIDES.lock().expect("gotter mutex poisoned"); + let mut guard = lock_global_overrides(); if guard.is_none() { let mut so = SymbolOverrides::new(); register_all(&mut so); @@ -96,10 +103,13 @@ pub fn update_heap_overrides() { // `try_lock` so a dlopen happening on the same thread that owns the // install lock doesn't deadlock - that thread will finish its // outer apply_overrides, which already walks every library. - if let Ok(mut guard) = GLOBAL_OVERRIDES.try_lock() { - if let Some(so) = guard.as_mut() { - so.update_overrides(); - } + let mut guard = match GLOBAL_OVERRIDES.try_lock() { + Ok(guard) => guard, + Err(TryLockError::Poisoned(poisoned)) => poisoned.into_inner(), + Err(TryLockError::WouldBlock) => return, + }; + if let Some(so) = guard.as_mut() { + so.update_overrides(); } } @@ -112,7 +122,7 @@ pub fn update_heap_overrides() {} /// non-Linux targets. #[cfg(all(target_os = "linux", target_pointer_width = "64"))] pub fn restore_heap_overrides() { - let mut guard = GLOBAL_OVERRIDES.lock().expect("gotter mutex poisoned"); + let mut guard = lock_global_overrides(); if let Some(so) = guard.as_mut() { so.restore_overrides(); } @@ -128,10 +138,7 @@ pub fn restore_heap_overrides() {} /// is a no-op there. #[cfg(all(target_os = "linux", target_pointer_width = "64"))] pub fn heap_overrides_are_installed() -> bool { - GLOBAL_OVERRIDES - .lock() - .map(|guard| guard.is_some()) - .unwrap_or(false) + lock_global_overrides().is_some() } /// See the Linux variant above. diff --git a/libdd-heap-gotter/tests/install.rs b/libdd-heap-gotter/tests/install.rs index 8a6e813c13..fc418eaa4b 100644 --- a/libdd-heap-gotter/tests/install.rs +++ b/libdd-heap-gotter/tests/install.rs @@ -1,25 +1,28 @@ // Copyright 2025-Present Datadog, Inc. https://www.datadoghq.com/ // SPDX-License-Identifier: Apache-2.0 -//! End-to-end smoke test: install the GOT overrides into the live test -//! process, do a libc `malloc`/`free`, and check that the hook ran. +//! End-to-end smoke tests: install the GOT overrides into the live test +//! process, exercise libc allocator functions, and confirm the sampler +//! is actually intercepting them. //! -//! Single test only — installing the overrides mutates global process -//! state, so we keep this isolated from the unit-test binary. +//! These tests mutate global process state (GOT entries), so they must +//! not run in parallel with each other. The `#[serial]` attribute +//! enforces that. -// Integration test invokes the GOT-patching machinery for real -// (dl_iterate_phdr + dlsym + mprotect), which miri can't execute. +// Integration tests invoke GOT-patching machinery for real +// (dl_iterate_phdr + mprotect), which miri can't execute. #![cfg(all(target_os = "linux", target_pointer_width = "64", not(miri)))] use std::ffi::c_void; -// We don't have a clean way to instrument the gotter's own hooks from -// outside the crate, so this test goes one level lower: it confirms -// that after install, the heap is still functional and that no -// recursive crash has occurred when malloc/free go through the patched -// GOT. +use libdd_heap_sampler::{dd_sample_flag_check_fast, dd_tl_state_get, dd_tl_state_get_or_init}; +use serial_test::serial; + +/// After install the heap should still be functional and no recursive +/// crash should occur when malloc/free go through the patched GOT. #[test] -fn install_changes_malloc_dispatch() { +#[serial] +fn install_and_restore_keeps_heap_functional() { extern "C" { fn malloc(size: usize) -> *mut c_void; } @@ -30,13 +33,6 @@ fn install_changes_malloc_dispatch() { "expected install_heap_overrides to find at least one symbol" ); - // Touch the heap to exercise the patched GOT. We can't assert that - // the address differs from the original libc malloc address from - // Rust directly — the `malloc` extern fn item resolves via the same - // GOT we just patched, so reading its address here either gives us - // the PLT entry (still the same) or the post-resolution function - // pointer (now our hook). What we *can* assert is that the heap is - // still functional and that no recursive crash has occurred. unsafe { let p = malloc(64); assert!(!p.is_null(), "malloc returned NULL post-install"); @@ -45,10 +41,153 @@ fn install_changes_malloc_dispatch() { libdd_heap_gotter::restore_heap_overrides(); - // A second alloc post-restore should also be fine. unsafe { let p = malloc(64); assert!(!p.is_null(), "malloc returned NULL post-restore"); libc::free(p); } } + +/// Confirm that after install, allocations actually flow through the +/// sampler and produce tagged pointers. We force sampling_interval=1 +/// so every allocation is sampled, then check that libc::malloc returns +/// a pointer carrying the sample flag. +#[test] +#[serial] +fn install_produces_sampled_allocations() { + let installed = libdd_heap_gotter::install_heap_overrides(); + assert!(installed); + + unsafe { + // Ensure sampler TLS is materialised on this thread. + let tl = dd_tl_state_get_or_init(); + assert!(!tl.is_null()); + + // Force every allocation to be sampled. + (*tl).sampling_interval = 1; + (*tl).remaining_bytes = 0; + (*tl).remaining_bytes_initialized = true; + + // Allocate. The gotter hook should intercept this, the sampler + // should decide to sample (interval=1), and the returned + // pointer should carry the sample flag. + let p = libc::malloc(128); + assert!(!p.is_null()); + + let mut raw: *mut c_void = std::ptr::null_mut(); + let sampled = dd_sample_flag_check_fast(p, &mut raw); + assert!( + sampled, + "expected malloc to return a sampled pointer with interval=1" + ); + assert!(!raw.is_null()); + + // Free via libc::free which goes through gotter_free; it + // should handle the tagged pointer correctly. + libc::free(p); + + // Restore the default interval so we don't mess with anything + // after the test. + (*tl).sampling_interval = 512 * 1024; + } + + libdd_heap_gotter::restore_heap_overrides(); +} + +/// Same as above but for realloc: confirm a sampled allocation that +/// gets reallocated comes back as a valid (unsampled) pointer with the +/// user data intact. +#[test] +#[serial] +fn realloc_of_sampled_allocation_preserves_data() { + let installed = libdd_heap_gotter::install_heap_overrides(); + assert!(installed); + + unsafe { + let tl = dd_tl_state_get_or_init(); + assert!(!tl.is_null()); + (*tl).sampling_interval = 1; + (*tl).remaining_bytes = 0; + (*tl).remaining_bytes_initialized = true; + + // Allocate and write a pattern. + let p = libc::malloc(64) as *mut u8; + assert!(!p.is_null()); + for i in 0..64u8 { + *p.add(i as usize) = i; + } + + // check_fast destructively clears the flag, so we can't peek + // and then realloc the same pointer. Just free this one and + // allocate a fresh one for the realloc test. + libc::free(p as *mut c_void); + + // Fresh sampled allocation for the realloc test. + let p = libc::malloc(64) as *mut u8; + assert!(!p.is_null()); + for i in 0..64u8 { + *p.add(i as usize) = 0xAB ^ i; + } + + // Realloc to a larger size. After realloc the pointer should be + // valid (possibly unsampled per the MVP model) and the original + // data should be preserved. + let p2 = libc::realloc(p as *mut c_void, 256) as *mut u8; + assert!(!p2.is_null(), "realloc returned NULL"); + + // Verify data integrity. + for i in 0..64u8 { + let got = *p2.add(i as usize); + assert_eq!( + got, + 0xAB ^ i, + "data corruption at byte {i}: expected 0x{:02x}, got 0x{got:02x}", + 0xAB ^ i + ); + } + + libc::free(p2 as *mut c_void); + (*tl).sampling_interval = 512 * 1024; + } + + libdd_heap_gotter::restore_heap_overrides(); +} + +/// Confirm that after restore, allocations are no longer sampled. +#[test] +#[serial] +fn restore_stops_sampling() { + let installed = libdd_heap_gotter::install_heap_overrides(); + assert!(installed); + + unsafe { + let tl = dd_tl_state_get_or_init(); + assert!(!tl.is_null()); + (*tl).sampling_interval = 1; + (*tl).remaining_bytes = 0; + (*tl).remaining_bytes_initialized = true; + } + + libdd_heap_gotter::restore_heap_overrides(); + + // After restore, malloc should return a plain pointer with no + // sample flag, even though the sampler TLS still has interval=1. + unsafe { + let p = libc::malloc(128); + assert!(!p.is_null()); + + let mut raw: *mut c_void = std::ptr::null_mut(); + let sampled = dd_sample_flag_check_fast(p, &mut raw); + assert!( + !sampled, + "expected malloc to return an unsampled pointer after restore" + ); + + libc::free(p); + + let tl = dd_tl_state_get(); + if !tl.is_null() { + (*tl).sampling_interval = 512 * 1024; + } + } +} diff --git a/libdd-heap-sampler/README.md b/libdd-heap-sampler/README.md index 9e277dbc49..f65313d51f 100644 --- a/libdd-heap-sampler/README.md +++ b/libdd-heap-sampler/README.md @@ -6,6 +6,13 @@ For allocations that are sampled as well as the corresponding frees of these all emitted such that an external process such as the [eBPF full host profiler](https://github.com/open-telemetry/opentelemetry-ebpf-profiler) can collect the samples as well as the stack trace at the time they are emitted to ultimately emit as a heap profiling event stream. +## Docs + +Two parts of this crate are fiddly enough that they get their own writeup: + +* [docs/tagging.md](docs/tagging.md) - how we mark an allocation as sampled and recognise it again at free time +* [docs/realloc.md](docs/realloc.md) - how we handle `realloc` + ## Use Cases This profiling infrastructure will initially support these two use cases: @@ -101,19 +108,10 @@ the common public surface, so those internal arch-specific helpers and constants are not emitted into `bindings.rs` and a single checked-in binding set works for both supported Linux architectures. -### Requirements - -- **`libclang`** (`libclang-dev` on Debian/Ubuntu, `libclang-devel` on - RPM distros, `brew install llvm` on macOS). This is the only hard - requirement. - -You do **not** need `bindgen-cli`, a cross libc, a cross Rust toolchain, -or a cross linker — `build.rs` invokes the `bindgen` crate directly and -only emits the arch-independent Rust-facing ABI. - -### Regenerating - ```bash +# We use an env var and not a feature, as several parts of the libdatadog +# build turn on all features, and we don't want everything to need the extra +# build tooling. LIBDD_HEAP_SAMPLER_REGEN=1 cargo build -p libdd-heap-sampler ``` @@ -122,10 +120,12 @@ This refreshes `src/generated/bindings.rs` and `verify-heap-sampler-bindings` workflow runs the same command on every PR and fails if the checked-in files are stale. -> **Why an env var, not a `cargo` feature?** Several of libdatadog's -> CI jobs (clippy, coverage, workspace tests) run with -> `--all-features`. Cargo has no way to hide a feature from that flag, -> so a `regen-bindings` feature would be silently enabled everywhere -> and drag `libclang` into every build. An env var is -> invisible to `--all-features`, so `build.rs` only invokes bindgen -> when a human (or the verify-bindings workflow) explicitly asks. +### Requirements + +- **`libclang`** (`libclang-dev` on Debian/Ubuntu, `libclang-devel` on + RPM distros, `brew install llvm` on macOS). This is the only hard + requirement. + +You do **not** need `bindgen-cli`, a cross libc, a cross Rust toolchain, +or a cross linker — `build.rs` invokes the `bindgen` crate directly and +only emits the arch-independent Rust-facing ABI. diff --git a/libdd-heap-sampler/docs/realloc.md b/libdd-heap-sampler/docs/realloc.md new file mode 100644 index 0000000000..d68109b4d7 --- /dev/null +++ b/libdd-heap-sampler/docs/realloc.md @@ -0,0 +1,66 @@ +# Realloc handling + +When we get a realloc, we treat it as a separate free and a new allocation +to keep things simple. This means a resized allocation that was sampled +before the resize is not sampled after it: we report the free and just let +the new block go untracked, rather than trying to carry the sample +forward. Doing that properly would mean stamping a fresh header with the +original user-requested size, which we don't currently store, so it's a +fair chunk of extra work for a case that's rare enough not to skew results +much. Worth revisiting later. + +## What we need to preserve + +Whatever we do here has to keep a couple of properties intact: + +* If the real `realloc` fails, the old pointer must still be exactly as + usable as it was before we touched anything, including still being + recognised as sampled if it was sampled going in. +* We must never hand back a pointer with a stale or wrong header +* The common, unsampled path should stay as close as possible to a plain + passthrough to the real `realloc` + +## The algorithm + +`gotter_realloc` (`libdd-heap-gotter/src/hooks.rs`) - and any other users of the samplers - +sorts every call into one of four cases before deciding what to do: + +* `ptr == NULL` is just `malloc(size)`, so it runs through the normal + allocation sampling path. +* `size == 0` is just `free(ptr)`, so it consumes the sampler flag (if + there is one) and forwards to the real allocator. +* An unsampled `ptr` is a plain passthrough to the real `realloc`. We + don't start sampling the resulting block here either, for the same + reason as above: we didn't want realloc behaviour to depend on which + side of the sampling coin flip the original allocation happened to land + on. +* A sampled `ptr` is where the interesting work happens, split into a + `prepare` / `commit` pair in `allocation_realloc.c` so we can respect + the "old pointer stays valid on failure" rule. + +`dd_allocation_realloc_prepare(old_user, new_size)` looks at the old +pointer using the non-destructive `peek` (see [tagging.md](tagging.md)), +which tells us where the real allocation starts (`old_raw`) and how far +into it the user pointer was (`old_offset`). +It computes the size to actually ask the real `realloc` for: +`new_size + old_offset`. That extra `old_offset` bytes exist +because libc's realloc will copy the old data forward from the old raw +pointer, so the old header ends up occupying the front of the new block +too unless we make room for it. Importantly, `prepare` leaves the old +allocation's flag alone, so if the real realloc fails, the pointer is +untouched and a later `free` on it still behaves correctly. + +The caller then calls the real `realloc` with the pointer and size +`prepare` computed, and passes whatever comes back into +`dd_allocation_realloc_commit(old_user, new_raw, prep)`: + +* If `new_raw` is `NULL`, realloc failed. We return `NULL` and leave + everything else alone. +* Otherwise, the old user bytes copied by libc are sitting at + `new_raw + old_offset` instead of at the start of the block, so we + `memmove` them down to the front. It has to be `memmove`, not `memcpy`, + because when realloc grows a block in place, `new_raw` and `old_raw` are + the same address and the ranges overlap. +* We fire the `ddheap:free` USDT for the old address, since as far as the + profiler is concerned that allocation no longer exists. +* We return `new_raw` as a plain, unsampled pointer. diff --git a/libdd-heap-sampler/docs/tagging.md b/libdd-heap-sampler/docs/tagging.md new file mode 100644 index 0000000000..f7e496117c --- /dev/null +++ b/libdd-heap-sampler/docs/tagging.md @@ -0,0 +1,102 @@ +# Address tagging + +When we sample an allocation we need a way to recognise, at free time, that +the pointer we're being handed back is one we sampled. That's what +`sample_flag.h` does. There are two completely different implementations, +one per architecture, because the best mechanism available differs a lot +between x86-64 and arm64. + +## arm64: pointer tagging + +arm64 supports Top-Byte Ignore (TBI). The CPU lets you set the top 8 bits +of a 64-bit pointer to whatever you like and still dereference it normally, +because those bits are ignored by the hardware. So instead of a header, +we just set the top byte to a fixed marker (`0xDD`) when we sample an allocation. +This gives us the lowest possible cost on the `free` side, as we simply +have to mask the pointer against our chosen flag bit. + +The catch is that the kernel needs to be told, per thread, that it's OK for +tagged pointers to come back through syscalls without being rejected. That +is done through `prctl(PR_SET_TAGGED_ADDR_CTRL)` in +`dd_sample_flag_thread_init`, which must run once per thread before any +tagging happens on that thread. If the kernel refuses (older kernels, +seccomp policies, etc.) we disable sampling entirely on that thread. We might +want to revisit this in the future if we see that we consistently encounter +this case. + +## x86-64: a header hidden before the user pointer + +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 +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. + +The header holds two 8-byte values: + +* a magic constant (`DD_MAGIC`) +* the offset from the user pointer back to the raw pointer the allocator + actually returned + +At free time we look at the 16 bytes before the pointer we were given. If +the magic matches, we treat this as a sampled allocation, read the offset, +and use it to recover the raw pointer to hand to the real `free`. If the +magic doesn't match, it's an ordinary allocation and we pass it straight +through. + +**Page Alignment** +On the `free` side, if our pointer is within 16 bytes of the previous page +boundary we cannot safely read beneath it, as we may read into unmapped memory. + +This means that when we sample an allocation, we must ensure that the resulting +pointer is _not_ within these 16 bytes. Preserving this property _and_ satisfying +the user's requested alignment is where the complexity in this mechanism comes from. + +We handle this by asking for enough extra space that we can always place the +user pointer at a safe offset. The base offset is `max(alignment, 16)` (room +for the header while staying aligned), and we reserve *twice* that in the +bumped allocation. The second copy is there so that if the first landing spot +would put us within 16 bytes of a page boundary, we can bump forward by +another `alignment` bytes and still fit inside the allocation. + +**Zeroing the header on free** + +Once we've confirmed a match and recovered the raw pointer, we zero out the +header. The underlying allocator can hand the same memory back out later for +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** + +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. + +**Alignment cap** + +We don't sample allocations with alignment above 4096 bytes +(`DD_SAMPLE_ALIGNMENT_CAP`). Because the bumped size is +`user_size + 2 * max(alignment, 16)`, high alignments mean a lot of wasted +space just to hold a 16-byte header. + +The worry is that some workloads/allocators may use big alignments to encourage the +system to hand back huge pages - see +[this blog post](https://mazzo.li/posts/check-huge-page.html) for context on +how common large-alignment allocations can be in practice. We need to monitor +how much we see this happening in real workloads, because if it's a lot we're +going to need to work out what to do with it. + +## Why two check functions + +You'll notice `dd_sample_flag_check` and `dd_sample_flag_peek` both exist. +`check` destroys the flag on the way out (clears the x86-64 header), which +is correct for a normal free: once we've recovered the raw pointer we never +need the flag on that address again. `peek` leaves the flag alone, which we +need for realloc; see `realloc.md` for why that distinction matters there. diff --git a/libdd-heap-sampler/include/datadog/heap/sample_flag.h b/libdd-heap-sampler/include/datadog/heap/sample_flag.h index 3479f569bc..10e60165cd 100644 --- a/libdd-heap-sampler/include/datadog/heap/sample_flag.h +++ b/libdd-heap-sampler/include/datadog/heap/sample_flag.h @@ -170,6 +170,9 @@ bool dd_sample_flag_check_fast(void *user, void **raw_out) { uint64_t offset; memcpy(&offset, (char *)header + sizeof(magic), sizeof(offset)); + if (offset < DD_HEADER_BYTES || offset > 2 * DD_SAMPLE_ALIGNMENT_CAP) { + return false; + } /* Clear the whole 16-byte header so a re-use of this address * (e.g. allocator returns the same block to a later, unsampled diff --git a/libdd-heap-sampler/include/datadog/heap/tl_state.h b/libdd-heap-sampler/include/datadog/heap/tl_state.h index a99ce5819e..3faee19534 100644 --- a/libdd-heap-sampler/include/datadog/heap/tl_state.h +++ b/libdd-heap-sampler/include/datadog/heap/tl_state.h @@ -86,8 +86,8 @@ typedef struct { This _will probably_ be constant, but if we drop it in the TLS we afford the eBPF profiler the opportunity to tune it to adjust overhead - dynamically. Whether or not this turns out to be - a clever idea remains to be seen. */ + dynamically. A value of 0 explicitly disables + sampling for this thread. */ int64_t remaining_bytes; /* signed counter; sample when >= 0 */ bool remaining_bytes_initialized; /* false until first interval drawn */ bool initialized; /* false until dd_tl_state_init() has run; diff --git a/libdd-heap-sampler/src/allocation_freed.c b/libdd-heap-sampler/src/allocation_freed.c index d9acf34062..609fc06f6c 100644 --- a/libdd-heap-sampler/src/allocation_freed.c +++ b/libdd-heap-sampler/src/allocation_freed.c @@ -32,23 +32,35 @@ dd_alloc_freed_t dd_allocation_freed_slow(void *ptr, void *raw, size_t size, }; #if defined(__x86_64__) - /* Recover the bumped size the allocator actually holds. On x86-64 - * that's user_size + (user - raw), then rounded up to a multiple - * of alignment to match dd_allocation_requested_slow's bumped-size - * arithmetic. Sized-free callers (sdallocx, operator delete(sz)) - * rely on this being exact. + /* Recover the bumped size the allocator actually holds. This must + * exactly mirror allocation_requested.c's bumped_alloc_size(): + * + * base = max(alignment, DD_HEADER_BYTES) + * reserve = 2 * base + * bumped = round_up(size + reserve, alignment) + * + * Do not use (user - raw) as the reserve. That offset is usually + * only one `base`, while allocation reserved two so x86_apply() has + * room for its optional page-boundary bump. Sized-free callers + * (Rust GlobalAlloc::dealloc, sdallocx, operator delete(sz)) rely + * on this being exact. * * When the caller doesn't know the alignment (alignment == 0), - * fall back to size + offset. Plain free() ignores out.size so - * this only matters for sized-free variants that must supply an - * alignment. */ + * fall back to size + offset. Plain free() ignores out.size so this + * only matters for sized-free variants that must supply an alignment. */ size_t offset = (size_t)((uintptr_t)ptr - (uintptr_t)raw); - size_t bumped = size + offset; - if (alignment > 1) { - size_t mask = alignment - 1; - bumped = (bumped + mask) & ~mask; + if (alignment == 0) { + out.size = size + offset; + } else { + size_t base = alignment > DD_HEADER_BYTES ? alignment : DD_HEADER_BYTES; + size_t reserve = base * 2; + size_t bumped = size + reserve; + if (alignment > 1) { + size_t mask = alignment - 1; + bumped = (bumped + mask) & ~mask; + } + out.size = bumped; } - out.size = bumped; #else (void)alignment; #endif diff --git a/libdd-heap-sampler/src/allocation_requested.c b/libdd-heap-sampler/src/allocation_requested.c index 4dc0459340..b9bcca8636 100644 --- a/libdd-heap-sampler/src/allocation_requested.c +++ b/libdd-heap-sampler/src/allocation_requested.c @@ -36,7 +36,8 @@ static uint64_t next_interval(uint32_t *rng, uint64_t mean) { * Called when remaining_bytes has crossed zero, meaning at least one sample * is owed. Draws fresh intervals until the counter is negative again, counting * how many samples fired. Returns nsamples * interval as the unbiased weight - * estimator to attribute to this allocation. + * estimator to attribute to this allocation. A sampling_interval of 0 is the + * documented "do not sample this thread" value and returns 0 immediately. * * On the very first call for a thread, remaining_bytes_initialized is false * and we draw the initial interval from scratch. If that interval exceeds the diff --git a/libdd-heap-sampler/src/lib.rs b/libdd-heap-sampler/src/lib.rs index 07606d7299..00b482bda9 100644 --- a/libdd-heap-sampler/src/lib.rs +++ b/libdd-heap-sampler/src/lib.rs @@ -84,6 +84,26 @@ mod tests { } } + #[test] + fn zero_sampling_interval_disables_sampling() { + let mut tl = dd_tl_state_t { + sampling_interval: 0, + remaining_bytes: 0, + remaining_bytes_initialized: true, + initialized: true, + reentry_guard: false, + rng: 1, + }; + + let req = unsafe { dd_allocation_requested_slow(&mut tl, 64, 8) }; + + assert_eq!(req.size, 64); + assert_eq!(req.user_size, 64); + assert_eq!(req.alignment, 8); + assert_eq!(req.weight, 0); + assert!(!tl.reentry_guard); + } + #[test] fn freed_unsampled_returns_inputs_unchanged() { // The C side now reads 8 bytes at ptr-16 looking for DD_MAGIC, @@ -98,6 +118,61 @@ mod tests { } } + #[cfg(target_arch = "x86_64")] + #[test] + fn sample_flag_check_fast_rejects_invalid_offset() { + const MAGIC: u64 = 0xfab1eddec0dedca7; + const HEADER_BYTES: usize = 16; + + let mut buf = vec![0u8; 8192]; + let base = buf.as_mut_ptr() as usize; + let user_addr = (base + 128 + 4095) & !4095; + let user_addr = user_addr + 64; + assert!(user_addr + 64 <= base + buf.len()); + + let header = user_addr - HEADER_BYTES; + let header_idx = header - base; + buf[header_idx..header_idx + 8].copy_from_slice(&MAGIC.to_ne_bytes()); + // Too small to have been produced by x86_apply(). + buf[header_idx + 8..header_idx + 16].copy_from_slice(&8u64.to_ne_bytes()); + + let mut raw = core::ptr::null_mut(); + let sampled = unsafe { dd_sample_flag_check_fast(user_addr as *mut c_void, &mut raw) }; + + assert!(!sampled); + assert!(raw.is_null()); + assert_eq!(&buf[header_idx..header_idx + 8], &MAGIC.to_ne_bytes()); + } + + #[cfg(target_arch = "x86_64")] + #[test] + fn freed_slow_size_matches_requested_bump_formula() { + const HEADER_BYTES: usize = 16; + + fn requested_bumped_size(user_size: usize, alignment: usize) -> usize { + let base = alignment.max(HEADER_BYTES); + let reserve = 2 * base; + let bumped = user_size + reserve; + if alignment > 1 { + let mask = alignment - 1; + (bumped + mask) & !mask + } else { + bumped + } + } + + let user_size = 100; + let alignment = 16; + let base = alignment.max(HEADER_BYTES); + let raw = 0x100000usize as *mut c_void; + let ptr = (0x100000usize + base) as *mut c_void; + + let freed = unsafe { dd_allocation_freed_slow(ptr, raw, user_size, alignment) }; + + assert_eq!(freed.ptr, raw); + assert_eq!(freed.size, requested_bumped_size(user_size, alignment)); + } + #[test] fn tl_state_init_then_get_returns_same_non_null() { std::thread::spawn(|| unsafe { From f0a58dff53c29b68d07f0a556171595af782563a Mon Sep 17 00:00:00 2001 From: Scott Gerring Date: Fri, 3 Jul 2026 10:55:29 +0200 Subject: [PATCH 07/16] address a correctness issue --- libdd-heap-gotter/tests/install.rs | 35 +++++++++++++++++++++--------- 1 file changed, 25 insertions(+), 10 deletions(-) diff --git a/libdd-heap-gotter/tests/install.rs b/libdd-heap-gotter/tests/install.rs index fc418eaa4b..7900d74916 100644 --- a/libdd-heap-gotter/tests/install.rs +++ b/libdd-heap-gotter/tests/install.rs @@ -15,7 +15,7 @@ use std::ffi::c_void; -use libdd_heap_sampler::{dd_sample_flag_check_fast, dd_tl_state_get, dd_tl_state_get_or_init}; +use libdd_heap_sampler::{dd_sample_flag_peek, dd_tl_state_get_or_init}; use serial_test::serial; /// After install the heap should still be functional and no recursive @@ -74,8 +74,12 @@ fn install_produces_sampled_allocations() { let p = libc::malloc(128); assert!(!p.is_null()); + // Use peek (non-destructive) to verify the flag is set without + // clearing it. gotter_free needs the flag intact to recover the + // raw pointer. let mut raw: *mut c_void = std::ptr::null_mut(); - let sampled = dd_sample_flag_check_fast(p, &mut raw); + let mut offset: usize = 0; + let sampled = dd_sample_flag_peek(p, &mut raw, &mut offset); assert!( sampled, "expected malloc to return a sampled pointer with interval=1" @@ -83,7 +87,7 @@ fn install_produces_sampled_allocations() { assert!(!raw.is_null()); // Free via libc::free which goes through gotter_free; it - // should handle the tagged pointer correctly. + // handles the tagged pointer correctly (check + free raw). libc::free(p); // Restore the default interval so we don't mess with anything @@ -163,31 +167,42 @@ fn restore_stops_sampling() { unsafe { let tl = dd_tl_state_get_or_init(); assert!(!tl.is_null()); + + // Confirm sampling works while installed. (*tl).sampling_interval = 1; (*tl).remaining_bytes = 0; (*tl).remaining_bytes_initialized = true; + + let p = libc::malloc(64); + assert!(!p.is_null()); + let mut raw: *mut c_void = std::ptr::null_mut(); + let mut offset: usize = 0; + let sampled = dd_sample_flag_peek(p, &mut raw, &mut offset); + assert!(sampled, "expected sampling while installed"); + libc::free(p); + + // Disable sampling before restore so internal allocations + // during restore don't get tagged (they'd be freed through + // the unpatched GOT afterwards, causing SIGABRT on x86_64). + (*tl).sampling_interval = 512 * 1024; } libdd_heap_gotter::restore_heap_overrides(); // After restore, malloc should return a plain pointer with no - // sample flag, even though the sampler TLS still has interval=1. + // sample flag. unsafe { let p = libc::malloc(128); assert!(!p.is_null()); let mut raw: *mut c_void = std::ptr::null_mut(); - let sampled = dd_sample_flag_check_fast(p, &mut raw); + let mut offset: usize = 0; + let sampled = dd_sample_flag_peek(p, &mut raw, &mut offset); assert!( !sampled, "expected malloc to return an unsampled pointer after restore" ); libc::free(p); - - let tl = dd_tl_state_get(); - if !tl.is_null() { - (*tl).sampling_interval = 512 * 1024; - } } } From 2533727e4d409f14240760a555c8f567f8cbe658 Mon Sep 17 00:00:00 2001 From: Scott Gerring Date: Fri, 3 Jul 2026 11:39:44 +0200 Subject: [PATCH 08/16] x86-64 can't sample 4 KiB alignment; add a test to confirm this --- libdd-heap-gotter/tests/install.rs | 37 +++++++++++++++++++ libdd-heap-sampler/docs/tagging.md | 14 +++++-- .../include/datadog/heap/sample_flag.h | 7 ++-- 3 files changed, 51 insertions(+), 7 deletions(-) diff --git a/libdd-heap-gotter/tests/install.rs b/libdd-heap-gotter/tests/install.rs index 7900d74916..fe4467c6ef 100644 --- a/libdd-heap-gotter/tests/install.rs +++ b/libdd-heap-gotter/tests/install.rs @@ -98,6 +98,43 @@ fn install_produces_sampled_allocations() { libdd_heap_gotter::restore_heap_overrides(); } +/// On x86-64, page-aligned allocations must pass through unsampled. +/// The header checker refuses pointers in the first 16 bytes of a page, +/// so a sampled 4096-aligned pointer could not be recognised later by +/// free or realloc. +#[cfg(target_arch = "x86_64")] +#[test] +#[serial] +fn page_aligned_allocations_are_unsampled() { + let installed = libdd_heap_gotter::install_heap_overrides(); + assert!(installed); + + unsafe { + let tl = dd_tl_state_get_or_init(); + assert!(!tl.is_null()); + (*tl).sampling_interval = 1; + (*tl).remaining_bytes = 0; + (*tl).remaining_bytes_initialized = true; + + let p = libc::aligned_alloc(4096, 4096); + assert!(!p.is_null()); + assert_eq!((p as usize) % 4096, 0); + + let mut raw: *mut c_void = std::ptr::null_mut(); + let mut offset: usize = 0; + let sampled = dd_sample_flag_peek(p, &mut raw, &mut offset); + assert!( + !sampled, + "page-aligned allocation must pass through unsampled" + ); + + libc::free(p); + (*tl).sampling_interval = 512 * 1024; + } + + libdd_heap_gotter::restore_heap_overrides(); +} + /// Same as above but for realloc: confirm a sampled allocation that /// gets reallocated comes back as a valid (unsampled) pointer with the /// user data intact. diff --git a/libdd-heap-sampler/docs/tagging.md b/libdd-heap-sampler/docs/tagging.md index f7e496117c..653c2bd3f4 100644 --- a/libdd-heap-sampler/docs/tagging.md +++ b/libdd-heap-sampler/docs/tagging.md @@ -81,13 +81,19 @@ one, check the other two. **Alignment cap** -We don't sample allocations with alignment above 4096 bytes +We don't sample allocations with alignment above 1024 bytes (`DD_SAMPLE_ALIGNMENT_CAP`). Because the bumped size is `user_size + 2 * max(alignment, 16)`, high alignments mean a lot of wasted -space just to hold a 16-byte header. +space just to hold a 16-byte header. -The worry is that some workloads/allocators may use big alignments to encourage the -system to hand back huge pages - see +There is also a correctness reason for keeping the cap below a page. The x86-64 +fast check refuses to read `user - 16` when `user` is in the first 16 bytes of a +page. That avoids touching an unmapped previous page for ordinary unsampled +pointers. A 4096-byte aligned pointer always has page offset 0, so a sampled +page-aligned allocation would not be recognised by free or realloc. + +Some workloads/allocators may use big alignments to encourage the system to +hand back huge pages - see [this blog post](https://mazzo.li/posts/check-huge-page.html) for context on how common large-alignment allocations can be in practice. We need to monitor how much we see this happening in real workloads, because if it's a lot we're diff --git a/libdd-heap-sampler/include/datadog/heap/sample_flag.h b/libdd-heap-sampler/include/datadog/heap/sample_flag.h index 10e60165cd..ae852923ff 100644 --- a/libdd-heap-sampler/include/datadog/heap/sample_flag.h +++ b/libdd-heap-sampler/include/datadog/heap/sample_flag.h @@ -64,10 +64,11 @@ bool dd_sample_flag_peek(void *user, void **raw_out, size_t *offset_out); * Largest alignment the sampler will honor. Above this we pass the * allocation through unsampled: the header + slack overhead grows with * alignment and stops being proportionate to any observability gain. - * Sized in bytes; equals one typical 4 KiB page on the supported - * architectures. + * Sized in bytes; kept below one typical 4 KiB page so x86-64 never + * returns a sampled pointer at page offset 0, which its fast-path + * checker deliberately rejects. */ -#define DD_SAMPLE_ALIGNMENT_CAP 4096 +#define DD_SAMPLE_ALIGNMENT_CAP 1024 #if defined(__x86_64__) From 99cfed4527407c1c0c3b9a9dcc59bfbe31f55608 Mon Sep 17 00:00:00 2001 From: Scott Gerring Date: Fri, 3 Jul 2026 13:00:04 +0200 Subject: [PATCH 09/16] skip deepbind --- libdd-heap-gotter/src/hooks.rs | 52 ++--------- libdd-heap-gotter/tests/install.rs | 36 +++++++- libdd-heap-sampler/docs/realloc.md | 46 +++++----- .../include/datadog/heap/allocation_realloc.h | 89 +++++++++++-------- libdd-heap-sampler/src/allocation_realloc.c | 48 ++++++++-- libdd-heap-sampler/src/generated/bindings.rs | 16 +++- 6 files changed, 173 insertions(+), 114 deletions(-) diff --git a/libdd-heap-gotter/src/hooks.rs b/libdd-heap-gotter/src/hooks.rs index 5891f7c5b1..d44bf7a765 100644 --- a/libdd-heap-gotter/src/hooks.rs +++ b/libdd-heap-gotter/src/hooks.rs @@ -125,54 +125,15 @@ pub unsafe extern "C" fn gotter_calloc(nmemb: usize, size: usize) -> *mut c_void /// `realloc` hook. /// -/// Handled as four disjoint cases: -/// -/// 1. **`ptr == NULL`** — equivalent to `malloc(size)`. Runs the normal sampling path (`request` + -/// `created`). -/// 2. **`size == 0`** — equivalent to `free(ptr)` for the allocators we hook. Consumes the sampler -/// flag via `dd_allocation_freed` (clears the header + fires `ddheap:free`) and forwards. -/// 3. **`ptr` is unsampled** — passthrough to the underlying realloc. We don't newly sample here -/// (see TODO on the branch). -/// 4. **`ptr` is sampled** — MVP: model a successful sampled realloc as `free(old sampled)` -/// followed by a new *unsampled* allocation. See the branch comment for the full rationale. -/// -/// This is a lot of logic that might end up being duplicated across other interception mechanisms -/// on top of libdd-heap-sampler. If that _is_ the case, we should consider adding a special case -/// `realloc` to the underlying sampler library and using that here; we've not done it for now to -/// preserve the general pre/post pattern used by the other functions. +/// The sampler owns the realloc cases and pointer math. Gotter only does +/// the pre/post split: ask the sampler what raw call to make, call the +/// real realloc symbol, then ask the sampler to turn the result back into +/// the user-visible pointer. #[no_mangle] pub unsafe extern "C" fn gotter_realloc(ptr: *mut c_void, size: usize) -> *mut c_void { let Some(real): Option = load_fn(&ORIG_REALLOC) else { return std::ptr::null_mut(); }; - // Case 1: realloc(NULL, size) == malloc(size). Normal sampling path. - if ptr.is_null() { - let alignment = core::mem::align_of::() * 2; - let req = dd_allocation_requested(size, alignment); - let raw = real(std::ptr::null_mut(), req.size); - return dd_allocation_created(raw, req); - } - - // Case 2: realloc(p, 0) == free(p). Safe to consume the sampler - // flag before forwarding because the old block will not remain - // live on this path. - if size == 0 { - let freed = dd_allocation_freed(ptr, 0, 0); - return real(freed.ptr, 0); - } - - // Cases 3 & 4: delegate to the sampler. `prepare` is - // non-destructive: if the underlying realloc fails, `ptr` stays - // live with its sampler flag intact for a later free. - // - // Sampled path (MVP): commit implements - // successful sampled realloc = ddheap:free(old sampled) - // + new unsampled allocation - // Unsampled path: prep is a passthrough; commit returns new_raw - // unchanged. - // - // TODO: emit `free + new sampled allocation` when sampled headers - // carry the original user size (see project TODO #13 / repo TODO #7). let prep = dd_allocation_realloc_prepare(ptr, size); let new_raw = real(prep.raw_ptr, prep.raw_size); dd_allocation_realloc_commit(ptr, new_raw, prep) @@ -220,6 +181,11 @@ pub unsafe extern "C" fn gotter_dlopen(filename: *const c_char, flags: c_int) -> return libc::dlopen(filename, flags); }; let handle = real(filename, flags); + if flags & libc::RTLD_DEEPBIND != 0 { + // DEEPBIND changes symbol resolution order and causes issues with + // GOT patching, so skip newly-loaded deep-bound libraries for now. + return handle; + } // New library may have introduced new GOT entries that need patching. // This hook is an extern "C" boundary, so never let a Rust panic from // best-effort ELF parsing/GOT patching unwind into the caller. diff --git a/libdd-heap-gotter/tests/install.rs b/libdd-heap-gotter/tests/install.rs index fe4467c6ef..be3c668988 100644 --- a/libdd-heap-gotter/tests/install.rs +++ b/libdd-heap-gotter/tests/install.rs @@ -98,6 +98,39 @@ fn install_produces_sampled_allocations() { libdd_heap_gotter::restore_heap_overrides(); } +/// Confirm realloc(NULL, size) goes through the sampler-side allocation +/// case, not a gotter-specific special case. +#[test] +#[serial] +fn realloc_null_produces_sampled_allocation() { + let installed = libdd_heap_gotter::install_heap_overrides(); + assert!(installed); + + unsafe { + let tl = dd_tl_state_get_or_init(); + assert!(!tl.is_null()); + (*tl).sampling_interval = 1; + (*tl).remaining_bytes = 0; + (*tl).remaining_bytes_initialized = true; + + let p = libc::realloc(std::ptr::null_mut(), 128); + assert!(!p.is_null()); + + let mut raw: *mut c_void = std::ptr::null_mut(); + let mut offset: usize = 0; + let sampled = dd_sample_flag_peek(p, &mut raw, &mut offset); + assert!( + sampled, + "realloc(NULL, size) should use allocation sampling" + ); + + libc::free(p); + (*tl).sampling_interval = 512 * 1024; + } + + libdd_heap_gotter::restore_heap_overrides(); +} + /// On x86-64, page-aligned allocations must pass through unsampled. /// The header checker refuses pointers in the first 16 bytes of a page, /// so a sampled 4096-aligned pointer could not be recognised later by @@ -171,8 +204,7 @@ fn realloc_of_sampled_allocation_preserves_data() { } // Realloc to a larger size. After realloc the pointer should be - // valid (possibly unsampled per the MVP model) and the original - // data should be preserved. + // valid and the original data should be preserved. let p2 = libc::realloc(p as *mut c_void, 256) as *mut u8; assert!(!p2.is_null(), "realloc returned NULL"); diff --git a/libdd-heap-sampler/docs/realloc.md b/libdd-heap-sampler/docs/realloc.md index d68109b4d7..bbf18be54c 100644 --- a/libdd-heap-sampler/docs/realloc.md +++ b/libdd-heap-sampler/docs/realloc.md @@ -22,37 +22,39 @@ Whatever we do here has to keep a couple of properties intact: ## The algorithm -`gotter_realloc` (`libdd-heap-gotter/src/hooks.rs`) - and any other users of the samplers - -sorts every call into one of four cases before deciding what to do: +The sampler sorts every call into one of four cases in +`dd_allocation_realloc_prepare(old_user, new_size)`: * `ptr == NULL` is just `malloc(size)`, so it runs through the normal - allocation sampling path. -* `size == 0` is just `free(ptr)`, so it consumes the sampler flag (if - there is one) and forwards to the real allocator. + allocation sampling path. `prepare` returns the raw size to pass to the + real allocator, and `commit` pairs the result with + `dd_allocation_created`. +* `size == 0` is just `free(ptr)` for the allocators we hook, so + `prepare` consumes the sampler flag if there is one and returns the raw + pointer to forward to the real allocator. * An unsampled `ptr` is a plain passthrough to the real `realloc`. We don't start sampling the resulting block here either, for the same reason as above: we didn't want realloc behaviour to depend on which side of the sampling coin flip the original allocation happened to land on. -* A sampled `ptr` is where the interesting work happens, split into a - `prepare` / `commit` pair in `allocation_realloc.c` so we can respect - the "old pointer stays valid on failure" rule. +* A sampled `ptr` is where the interesting work happens. `prepare` uses + the non-destructive `peek` (see [tagging.md](tagging.md)) to find the + real allocation start (`old_raw`) and the offset of the user pointer + (`old_offset`). It asks the frontend to call the real allocator with + `old_raw` and `new_size + old_offset`. -`dd_allocation_realloc_prepare(old_user, new_size)` looks at the old -pointer using the non-destructive `peek` (see [tagging.md](tagging.md)), -which tells us where the real allocation starts (`old_raw`) and how far -into it the user pointer was (`old_offset`). -It computes the size to actually ask the real `realloc` for: -`new_size + old_offset`. That extra `old_offset` bytes exist -because libc's realloc will copy the old data forward from the old raw -pointer, so the old header ends up occupying the front of the new block -too unless we make room for it. Importantly, `prepare` leaves the old -allocation's flag alone, so if the real realloc fails, the pointer is -untouched and a later `free` on it still behaves correctly. +That extra `old_offset` bytes exist because libc's realloc will copy the +old data forward from the old raw pointer, so the old header ends up +occupying the front of the new block too unless we make room for it. +Importantly, the sampled-old prepare path leaves the old allocation's flag +alone, so if the real realloc fails, the pointer is untouched and a later +`free` on it still behaves correctly. -The caller then calls the real `realloc` with the pointer and size -`prepare` computed, and passes whatever comes back into -`dd_allocation_realloc_commit(old_user, new_raw, prep)`: +The frontend calls the real `realloc` with the pointer and size `prepare` +computed, and passes whatever comes back into +`dd_allocation_realloc_commit(old_user, new_raw, prep)`. + +For the sampled-old case: * If `new_raw` is `NULL`, realloc failed. We return `NULL` and leave everything else alone. diff --git a/libdd-heap-sampler/include/datadog/heap/allocation_realloc.h b/libdd-heap-sampler/include/datadog/heap/allocation_realloc.h index c0978fba17..cc7dd2ca8e 100644 --- a/libdd-heap-sampler/include/datadog/heap/allocation_realloc.h +++ b/libdd-heap-sampler/include/datadog/heap/allocation_realloc.h @@ -12,53 +12,67 @@ * void *new_raw = real_realloc(prep.raw_ptr, prep.raw_size); * return dd_allocation_realloc_commit(old_user, new_raw, prep); * - * The MVP model for a sampled old allocation is: + * The prep/commit split mirrors dd_allocation_requested/created and + * dd_allocation_freed: the sampler owns the pointer tagging policy and + * the frontend owns the call to the real allocator. * - * successful sampled realloc = ddheap:free(old sampled) - * + new unsampled allocation + * For a sampled old allocation, a successful realloc is reported as: * - * See gotter's `gotter_realloc` for the full four-case handling - * (`ptr == NULL`, `size == 0`, unsampled-old, sampled-old); this pair - * covers only the last two cases (regular passthrough vs sampled - * teardown), since malloc/free/realloc(NULL,_) / realloc(_,0) map to - * dedicated sampler primitives. + * ddheap:free(old sampled) + new unsampled allocation * - * Not paired with dd_allocation_requested / dd_allocation_created: the - * new block is deliberately unsampled to avoid data-corruption hazards - * around header stamping without knowing the original user-requested - * size. Revisit once sampled headers carry the original user size. + * New blocks from realloc(NULL, size) use the normal allocation sampling + * path. Existing unsampled blocks pass through unchanged. Existing + * sampled blocks are torn down as above. */ #ifndef DD_SAMPLERS_ALLOCATION_REALLOC_H #define DD_SAMPLERS_ALLOCATION_REALLOC_H -#include +#include + #include +#include + +/* Which realloc case prepare() classified. */ +typedef enum { + DD_REALLOC_KIND_PASSTHROUGH = 0, + DD_REALLOC_KIND_ALLOC = 1, + DD_REALLOC_KIND_FREE = 2, + DD_REALLOC_KIND_SAMPLED = 3, +} dd_realloc_kind_t; /* - * Snapshot of the sampler state around a sampled realloc. + * Snapshot of the sampler state around realloc. * - * raw_ptr - pointer the frontend MUST pass to the underlying - * realloc. Equal to old_user on the passthrough path. - * raw_size - size the frontend MUST pass to the underlying - * realloc. Equal to new_size on the passthrough path. - * old_offset - byte offset from raw to user in the OLD sampled - * block. Used by commit() to shift user data down after - * realloc succeeds. 0 on the passthrough path. - * was_sampled - true iff old_user was a sampled allocation. commit() - * only runs the extra teardown work when this is true. + * raw_ptr - pointer the frontend MUST pass to the underlying + * realloc. NULL for realloc(NULL, size). Equal to + * old_user on the passthrough path. + * raw_size - size the frontend MUST pass to the underlying realloc. + * old_offset - byte offset from raw to user in the OLD sampled block. + * Used by commit() to shift user data down after realloc + * succeeds. 0 except on the sampled-old path. + * alloc_req - allocation request state for realloc(NULL, size), so + * commit() can pair the real realloc result with + * dd_allocation_created and close the sampler guard. + * kind - which realloc case prepare() selected. */ typedef struct { - void *raw_ptr; - size_t raw_size; - size_t old_offset; - bool was_sampled; + void *raw_ptr; + size_t raw_size; + size_t old_offset; + dd_alloc_req_t alloc_req; + dd_realloc_kind_t kind; } dd_realloc_prep_t; /* * Inspect old_user and compute the request to hand to the underlying - * realloc. Non-destructive: does not clear the sampler flag on - * old_user, so if realloc later returns NULL the old allocation stays - * usable and its flag stays intact for the eventual free. + * realloc. For sampled old allocations this is non-destructive: it does + * not clear the sampler flag on old_user, so if realloc later returns + * NULL the old allocation stays usable and its flag stays intact for the + * eventual free. + * + * realloc(old_user, 0) is destructive by definition for the allocators + * we hook. prepare() consumes the sampler flag in that case before the + * frontend forwards the raw pointer to realloc(raw, 0). */ dd_realloc_prep_t dd_allocation_realloc_prepare(void *old_user, size_t new_size); @@ -68,12 +82,17 @@ dd_realloc_prep_t dd_allocation_realloc_prepare(void *old_user, size_t new_size) * prepare(), returns the user-visible pointer to hand to the * application. * - * On the sampled path: shifts old user contents from [old_offset, ...) - * down to [0, ...), fires ddheap:free(old_user), and returns new_raw - * as an unsampled pointer. + * On realloc(NULL, size): pairs new_raw with dd_allocation_created and + * returns the possibly tagged user pointer. + * + * On the sampled-old path: shifts old user contents from [old_offset, ...) + * down to [0, ...), fires ddheap:free(old_user), and returns new_raw as + * an unsampled pointer. * - * On the unsampled/passthrough path: returns new_raw unchanged. - * On realloc failure (new_raw == NULL): returns NULL. + * On the unsampled/passthrough and realloc(ptr, 0) paths: returns new_raw + * unchanged. + * On sampled realloc failure (new_raw == NULL): returns NULL and leaves + * old_user live with its sampler flag intact. */ void *dd_allocation_realloc_commit(void *old_user, void *new_raw, dd_realloc_prep_t prep); diff --git a/libdd-heap-sampler/src/allocation_realloc.c b/libdd-heap-sampler/src/allocation_realloc.c index 78002b3390..a1c919265c 100644 --- a/libdd-heap-sampler/src/allocation_realloc.c +++ b/libdd-heap-sampler/src/allocation_realloc.c @@ -1,6 +1,8 @@ // Copyright 2025-Present Datadog, Inc. https://www.datadoghq.com/ // SPDX-License-Identifier: Apache-2.0 +#include +#include #include #include #include @@ -8,14 +10,35 @@ #include #include +#define DD_REALLOC_DEFAULT_ALIGNMENT (sizeof(void *) * 2) + dd_realloc_prep_t dd_allocation_realloc_prepare(void *old_user, size_t new_size) { dd_realloc_prep_t out = { .raw_ptr = old_user, .raw_size = new_size, .old_offset = 0, - .was_sampled = false, + .alloc_req = { 0, 0, 0, 0 }, + .kind = DD_REALLOC_KIND_PASSTHROUGH, }; + if (old_user == NULL) { + dd_alloc_req_t req = dd_allocation_requested( + new_size, DD_REALLOC_DEFAULT_ALIGNMENT); + out.raw_ptr = NULL; + out.raw_size = req.size; + out.alloc_req = req; + out.kind = DD_REALLOC_KIND_ALLOC; + return out; + } + + if (new_size == 0) { + dd_alloc_freed_t freed = dd_allocation_freed(old_user, 0, 0); + out.raw_ptr = freed.ptr; + out.raw_size = 0; + out.kind = DD_REALLOC_KIND_FREE; + return out; + } + void *old_raw = NULL; size_t old_offset = 0; if (!dd_sample_flag_peek(old_user, &old_raw, &old_offset)) { @@ -32,22 +55,31 @@ dd_realloc_prep_t dd_allocation_realloc_prepare(void *old_user, size_t new_size) return out; } - out.raw_ptr = old_raw; - out.raw_size = new_size + old_offset; - out.old_offset = old_offset; - out.was_sampled = true; + out.raw_ptr = old_raw; + out.raw_size = new_size + old_offset; + out.old_offset = old_offset; + out.kind = DD_REALLOC_KIND_SAMPLED; return out; } void *dd_allocation_realloc_commit(void *old_user, void *new_raw, dd_realloc_prep_t prep) { + if (prep.kind == DD_REALLOC_KIND_ALLOC) { + return dd_allocation_created(new_raw, prep.alloc_req); + } + + if (prep.kind == DD_REALLOC_KIND_FREE) { + return new_raw; + } + + if (prep.kind == DD_REALLOC_KIND_PASSTHROUGH) { + return new_raw; + } + /* Underlying realloc failed: C says old_user is still live; its * sampler flag was left intact by prepare(), so a later free() * will still resolve the right raw pointer. */ if (new_raw == NULL) return NULL; - /* Passthrough: nothing to fix up. */ - if (!prep.was_sampled) return new_raw; - /* Sampled path. libc realloc copied the old block's bytes into * new_raw starting at index 0, so the old user data now sits at * new_raw + old_offset. Shift it down to new_raw = user offset 0 diff --git a/libdd-heap-sampler/src/generated/bindings.rs b/libdd-heap-sampler/src/generated/bindings.rs index f6da064057..d8cab0e81b 100644 --- a/libdd-heap-sampler/src/generated/bindings.rs +++ b/libdd-heap-sampler/src/generated/bindings.rs @@ -154,17 +154,23 @@ unsafe extern "C" { alignment: usize, ) -> dd_alloc_freed_t; } +pub const dd_realloc_kind_t_DD_REALLOC_KIND_PASSTHROUGH: dd_realloc_kind_t = 0; +pub const dd_realloc_kind_t_DD_REALLOC_KIND_ALLOC: dd_realloc_kind_t = 1; +pub const dd_realloc_kind_t_DD_REALLOC_KIND_FREE: dd_realloc_kind_t = 2; +pub const dd_realloc_kind_t_DD_REALLOC_KIND_SAMPLED: dd_realloc_kind_t = 3; +pub type dd_realloc_kind_t = ::std::os::raw::c_uint; #[repr(C)] #[derive(Debug, Copy, Clone)] pub struct dd_realloc_prep_t { pub raw_ptr: *mut ::std::os::raw::c_void, pub raw_size: usize, pub old_offset: usize, - pub was_sampled: bool, + pub alloc_req: dd_alloc_req_t, + pub kind: dd_realloc_kind_t, } #[allow(clippy::unnecessary_operation, clippy::identity_op)] const _: () = { - ["Size of dd_realloc_prep_t"][::std::mem::size_of::() - 32usize]; + ["Size of dd_realloc_prep_t"][::std::mem::size_of::() - 64usize]; ["Alignment of dd_realloc_prep_t"][::std::mem::align_of::() - 8usize]; ["Offset of field: dd_realloc_prep_t::raw_ptr"] [::std::mem::offset_of!(dd_realloc_prep_t, raw_ptr) - 0usize]; @@ -172,8 +178,10 @@ const _: () = { [::std::mem::offset_of!(dd_realloc_prep_t, raw_size) - 8usize]; ["Offset of field: dd_realloc_prep_t::old_offset"] [::std::mem::offset_of!(dd_realloc_prep_t, old_offset) - 16usize]; - ["Offset of field: dd_realloc_prep_t::was_sampled"] - [::std::mem::offset_of!(dd_realloc_prep_t, was_sampled) - 24usize]; + ["Offset of field: dd_realloc_prep_t::alloc_req"] + [::std::mem::offset_of!(dd_realloc_prep_t, alloc_req) - 24usize]; + ["Offset of field: dd_realloc_prep_t::kind"] + [::std::mem::offset_of!(dd_realloc_prep_t, kind) - 56usize]; }; unsafe extern "C" { pub fn dd_allocation_realloc_prepare( From b7e1573cca64b3f16eecc727981f9a64ad47a66f Mon Sep 17 00:00:00 2001 From: Scott Gerring Date: Tue, 30 Jun 2026 17:10:55 +0200 Subject: [PATCH 10/16] feat(libdd-heap-allocator): add sampler_overhead criterion benchmark Adds a Criterion benchmark comparing direct System allocation, the SampledAllocator hot path, and direct sampler calls, plus a small bench.sh helper for libdd-heap-sampler. Kept on a separate change because it pulls criterion (and thus bindgen via libdd-heap-sampler's build) into the bench compile set, which the libdatadog benchmarking-platform image cannot build without libclang-dev. --- Cargo.lock | 1 + libdd-heap-allocator/Cargo.toml | 7 + libdd-heap-allocator/README.md | 22 +- .../benches/sampler_overhead.rs | 234 ++++++++++++++++++ libdd-heap-sampler/bench.sh | 77 ++++++ 5 files changed, 340 insertions(+), 1 deletion(-) create mode 100644 libdd-heap-allocator/benches/sampler_overhead.rs create mode 100755 libdd-heap-sampler/bench.sh diff --git a/Cargo.lock b/Cargo.lock index fcee1c9201..fe91fd8153 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3102,6 +3102,7 @@ dependencies = [ name = "libdd-heap-allocator" version = "0.1.0" dependencies = [ + "criterion", "libdd-heap-sampler", ] diff --git a/libdd-heap-allocator/Cargo.toml b/libdd-heap-allocator/Cargo.toml index 56f9d5ff19..b7931753e2 100644 --- a/libdd-heap-allocator/Cargo.toml +++ b/libdd-heap-allocator/Cargo.toml @@ -17,3 +17,10 @@ bench = false [dependencies] libdd-heap-sampler = { path = "../libdd-heap-sampler" } + +[dev-dependencies] +criterion = "0.5.1" + +[[bench]] +name = "sampler_overhead" +harness = false diff --git a/libdd-heap-allocator/README.md b/libdd-heap-allocator/README.md index 8e07d5838d..affe32b12d 100644 --- a/libdd-heap-allocator/README.md +++ b/libdd-heap-allocator/README.md @@ -22,4 +22,24 @@ To wrap a custom allocator instead; note that this is kind of ill-advised; we wa static ALLOC: SampledAllocator = SampledAllocator::new(MyAllocator::new()); ``` -See [`examples/usdt_demo.rs`](examples/usdt_demo.rs) for a runnable demo that fires USDT probes in a loop for `bpftrace` to observe. \ No newline at end of file +See [`examples/usdt_demo.rs`](examples/usdt_demo.rs) for a runnable demo that fires USDT probes in a loop for `bpftrace` to observe. + +## Benchmarking sampler overhead + +The `sampler_overhead` Criterion benchmark measures the allocator/sampler hot path without installing `SampledAllocator` as the process global allocator. It compares direct `System` allocation, `SampledAllocator`, a no-op allocator, `SampledAllocator`, and direct sampler calls. + +```sh +cargo bench -p libdd-heap-allocator --bench sampler_overhead +``` + +One quick validation run produced these fast-path results: + +| Size | Base: `System` alloc/free | Sampled fast path | Overhead | Overhead % | +|---:|---:|---:|---:|---:| +| 16 B | 5.9719 ns | 10.916 ns | +4.9441 ns | +82.8% | +| 64 B | 5.9309 ns | 12.405 ns | +6.4741 ns | +109.2% | +| 256 B | 5.9639 ns | 10.827 ns | +4.8631 ns | +81.5% | +| 4096 B | 23.237 ns | 29.402 ns | +6.1650 ns | +26.5% | +| 65536 B | 23.880 ns | 28.496 ns | +4.6160 ns | +19.3% | + +The no-op allocator comparison isolates the wrapper/sampler fast path at roughly **+5–6 ns per alloc/free pair**. diff --git a/libdd-heap-allocator/benches/sampler_overhead.rs b/libdd-heap-allocator/benches/sampler_overhead.rs new file mode 100644 index 0000000000..bc2cc5e97a --- /dev/null +++ b/libdd-heap-allocator/benches/sampler_overhead.rs @@ -0,0 +1,234 @@ +// Copyright 2025-Present Datadog, Inc. https://www.datadoghq.com/ +// SPDX-License-Identifier: Apache-2.0 + +// The sampler/allocator items exercised here are Linux-only; on other +// targets the bench compiles to a no-op `main` so workspace-wide +// `cargo check --all-targets` doesn't fail. + +#[cfg(not(target_os = "linux"))] +fn main() {} + +#[cfg(target_os = "linux")] +criterion::criterion_main!(linux_bench::benches); + +#[cfg(target_os = "linux")] +mod linux_bench { + use criterion::{criterion_group, BenchmarkId, Criterion, Throughput}; + use libdd_heap_allocator::SampledAllocator; + use libdd_heap_sampler::{ + dd_allocation_created, dd_allocation_freed, dd_allocation_requested, dd_tl_state_get, + dd_tl_state_init, + }; + use std::alloc::{GlobalAlloc, Layout, System}; + use std::hint::black_box; + use std::ptr; + + const SIZES: &[usize] = &[16, 64, 256, 4096, 65_536]; + const ALIGN: usize = 8; + + #[repr(align(4096))] + struct AlignedBuffer([u8; 128 * 1024]); + + static mut NOOP_BUFFER: AlignedBuffer = AlignedBuffer([0; 128 * 1024]); + + struct NoopAllocator; + + unsafe impl GlobalAlloc for NoopAllocator { + unsafe fn alloc(&self, _layout: Layout) -> *mut u8 { + // 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. + unsafe { ptr::addr_of_mut!(NOOP_BUFFER.0).cast::().add(4096) } + } + + unsafe fn dealloc(&self, _ptr: *mut u8, _layout: Layout) {} + } + + unsafe fn noop_user_ptr() -> *mut u8 { + unsafe { ptr::addr_of_mut!(NOOP_BUFFER.0).cast::().add(4096) } + } + + unsafe fn sampler_tl_state() -> *mut libdd_heap_sampler::dd_tl_state_t { + let mut tl = unsafe { dd_tl_state_get() }; + if tl.is_null() { + let _ = unsafe { dd_tl_state_init() }; + tl = unsafe { dd_tl_state_get() }; + } + tl + } + + unsafe fn pin_sampler_to_fast_path() { + let tl = unsafe { sampler_tl_state() }; + if !tl.is_null() { + unsafe { + (*tl).sampling_interval = u64::MAX / 4; + (*tl).remaining_bytes = i64::MIN / 4; + (*tl).remaining_bytes_initialized = true; + (*tl).reentry_guard = false; + } + } + } + + unsafe fn force_next_allocation_to_sample() { + let tl = unsafe { sampler_tl_state() }; + if !tl.is_null() { + unsafe { + (*tl).sampling_interval = 512 * 1024; + (*tl).remaining_bytes = 0; + (*tl).remaining_bytes_initialized = true; + (*tl).reentry_guard = false; + } + } + } + + fn bench_system_alloc_free(c: &mut Criterion) { + let mut group = c.benchmark_group("alloc_free/system"); + for &size in SIZES { + let layout = Layout::from_size_align(size, ALIGN).unwrap(); + group.throughput(Throughput::Bytes(size as u64)); + group.bench_with_input(BenchmarkId::from_parameter(size), &layout, |b, &layout| { + b.iter(|| unsafe { + let ptr = System.alloc(layout); + black_box(ptr); + System.dealloc(ptr, layout); + }); + }); + } + group.finish(); + } + + fn bench_sampled_system_alloc_free(c: &mut Criterion) { + let alloc = SampledAllocator::new(System); + let mut group = c.benchmark_group("alloc_free/sampled_system_fast_path"); + for &size in SIZES { + let layout = Layout::from_size_align(size, ALIGN).unwrap(); + group.throughput(Throughput::Bytes(size as u64)); + group.bench_with_input(BenchmarkId::from_parameter(size), &layout, |b, &layout| { + unsafe { pin_sampler_to_fast_path() }; + b.iter(|| unsafe { + let ptr = alloc.alloc(layout); + black_box(ptr); + alloc.dealloc(ptr, layout); + }); + }); + } + group.finish(); + } + + fn bench_noop_alloc_free(c: &mut Criterion) { + let alloc = NoopAllocator; + let mut group = c.benchmark_group("alloc_free/noop"); + for &size in SIZES { + let layout = Layout::from_size_align(size, ALIGN).unwrap(); + group.throughput(Throughput::Bytes(size as u64)); + group.bench_with_input(BenchmarkId::from_parameter(size), &layout, |b, &layout| { + b.iter(|| unsafe { + let ptr = alloc.alloc(layout); + black_box(ptr); + alloc.dealloc(ptr, layout); + }); + }); + } + group.finish(); + } + + fn bench_sampled_noop_alloc_free(c: &mut Criterion) { + let alloc = SampledAllocator::new(NoopAllocator); + let mut group = c.benchmark_group("alloc_free/sampled_noop_fast_path"); + for &size in SIZES { + let layout = Layout::from_size_align(size, ALIGN).unwrap(); + group.throughput(Throughput::Bytes(size as u64)); + group.bench_with_input(BenchmarkId::from_parameter(size), &layout, |b, &layout| { + unsafe { pin_sampler_to_fast_path() }; + b.iter(|| unsafe { + let ptr = alloc.alloc(layout); + black_box(ptr); + alloc.dealloc(ptr, layout); + }); + }); + } + group.finish(); + } + + fn bench_sampler_only(c: &mut Criterion) { + let mut group = c.benchmark_group("sampler_only/fast_path"); + for &size in SIZES { + group.throughput(Throughput::Bytes(size as u64)); + group.bench_with_input(BenchmarkId::from_parameter(size), &size, |b, &size| { + unsafe { pin_sampler_to_fast_path() }; + b.iter(|| unsafe { + let req = dd_allocation_requested(black_box(size), black_box(ALIGN)); + let user = dd_allocation_created(black_box(noop_user_ptr()).cast(), req); + let freed = dd_allocation_freed(user, black_box(size), black_box(ALIGN)); + black_box(freed); + }); + }); + } + group.finish(); + } + + fn bench_sampled_system_slow_path(c: &mut Criterion) { + let alloc = SampledAllocator::new(System); + let mut group = c.benchmark_group("alloc_free/sampled_system_slow_path"); + for &size in SIZES { + let layout = Layout::from_size_align(size, ALIGN).unwrap(); + group.throughput(Throughput::Bytes(size as u64)); + group.bench_with_input(BenchmarkId::from_parameter(size), &layout, |b, &layout| { + b.iter(|| unsafe { + force_next_allocation_to_sample(); + let ptr = alloc.alloc(layout); + black_box(ptr); + alloc.dealloc(ptr, layout); + }); + }); + } + group.finish(); + } + + fn bench_sampled_noop_slow_path(c: &mut Criterion) { + let alloc = SampledAllocator::new(NoopAllocator); + let mut group = c.benchmark_group("alloc_free/sampled_noop_slow_path"); + for &size in SIZES { + let layout = Layout::from_size_align(size, ALIGN).unwrap(); + group.throughput(Throughput::Bytes(size as u64)); + group.bench_with_input(BenchmarkId::from_parameter(size), &layout, |b, &layout| { + b.iter(|| unsafe { + force_next_allocation_to_sample(); + let ptr = alloc.alloc(layout); + black_box(ptr); + alloc.dealloc(ptr, layout); + }); + }); + } + group.finish(); + } + + fn bench_sampler_only_slow_path(c: &mut Criterion) { + let mut group = c.benchmark_group("sampler_only/slow_path"); + for &size in SIZES { + group.throughput(Throughput::Bytes(size as u64)); + group.bench_with_input(BenchmarkId::from_parameter(size), &size, |b, &size| { + b.iter(|| unsafe { + force_next_allocation_to_sample(); + let req = dd_allocation_requested(black_box(size), black_box(ALIGN)); + let user = dd_allocation_created(black_box(noop_user_ptr()).cast(), req); + let freed = dd_allocation_freed(user, black_box(size), black_box(ALIGN)); + black_box(freed); + }); + }); + } + group.finish(); + } + + criterion_group!( + benches, + bench_system_alloc_free, + bench_sampled_system_alloc_free, + bench_noop_alloc_free, + bench_sampled_noop_alloc_free, + bench_sampler_only, + bench_sampled_system_slow_path, + bench_sampled_noop_slow_path, + bench_sampler_only_slow_path, + ); +} // mod linux_bench diff --git a/libdd-heap-sampler/bench.sh b/libdd-heap-sampler/bench.sh new file mode 100755 index 0000000000..94910a46ab --- /dev/null +++ b/libdd-heap-sampler/bench.sh @@ -0,0 +1,77 @@ +#!/usr/bin/env bash +# Copyright 2025-Present Datadog, Inc. https://www.datadoghq.com/ +# SPDX-License-Identifier: Apache-2.0 + +set -euo pipefail +set -x + +# Run from inside heap-profiling-amd64. This script assumes the repo is mounted at +# /workspaces/libdatadog, but also works if launched from the repo root. +if [ -d /workspaces/libdatadog ]; then + cd /workspaces/libdatadog +fi + +OUT_DIR="libdd-heap-sampler" +mkdir -p "$OUT_DIR" + +# Rebuild the benchmark so build.rs changes such as -fcf-protection=none are used. +cargo bench -p libdd-heap-allocator --bench sampler_overhead --no-run + +# Pick the x86-64 bench binary; the shared target dir may contain stale arm64 bins. +BENCH_BIN=$( + for f in target/release/deps/sampler_overhead-*; do + [ -x "$f" ] || continue + file "$f" | grep -q 'x86-64' && { echo "$f"; break; } + done +) + +echo "BENCH_BIN=$BENCH_BIN" +file "$BENCH_BIN" + +# Check whether endbr64 is gone from the sampler entrypoints. +objdump -d --demangle "$BENCH_BIN" \ + | grep -A12 ':' \ + > "$OUT_DIR/objdump-cet-none-snippet.txt" || true +objdump -d --demangle "$BENCH_BIN" \ + | grep -A12 ':' \ + >> "$OUT_DIR/objdump-cet-none-snippet.txt" || true +objdump -d --demangle "$BENCH_BIN" \ + | grep -A16 ':' \ + >> "$OUT_DIR/objdump-cet-none-snippet.txt" || true +cat "$OUT_DIR/objdump-cet-none-snippet.txt" + +# Focused timing comparison. +cargo bench -p libdd-heap-allocator --bench sampler_overhead -- \ + --warm-up-time 0.3 \ + --measurement-time 0.3 \ + --sample-size 10 \ + 'alloc_free/noop/64|alloc_free/sampled_noop_fast_path/64|sampler_only/fast_path/64|alloc_free/system/64|alloc_free/sampled_system_fast_path/64' \ + 2>&1 | tee "$OUT_DIR/bench-amd64-cet-none" + +# Optional perf profile of the fast path, if perf is available/allowed. +if command -v perf >/dev/null 2>&1; then + sudo perf record -o "$OUT_DIR/perf-sampled-noop-fast-path-cet-none.data" \ + -F 997 \ + -e cpu-clock \ + -g --call-graph dwarf \ + "$BENCH_BIN" \ + --bench 'alloc_free/sampled_noop_fast_path/64' \ + --profile-time 10 + + sudo perf report -f \ + -i "$OUT_DIR/perf-sampled-noop-fast-path-cet-none.data" \ + --stdio \ + --no-children \ + --sort symbol,dso \ + > "$OUT_DIR/perf-sampled-noop-fast-path-cet-none.report.txt" + + for sym in dd_allocation_requested__extern dd_allocation_created__extern dd_allocation_freed__extern; do + sudo perf annotate -f \ + -i "$OUT_DIR/perf-sampled-noop-fast-path-cet-none.data" \ + --stdio \ + --symbol "$sym" \ + > "$OUT_DIR/perf-annotate-$sym-cet-none.txt" || true + done +fi + +ls -lh "$OUT_DIR"/*cet-none* "$OUT_DIR"/objdump-cet-none-snippet.txt 2>/dev/null || true From 00e15fd98114182395297d3b010c82d4bf12791c Mon Sep 17 00:00:00 2001 From: Scott Gerring Date: Fri, 3 Jul 2026 13:39:51 +0200 Subject: [PATCH 11/16] fix issue in test --- libdd-heap-gotter/tests/install.rs | 25 +++++++++++++++++-------- 1 file changed, 17 insertions(+), 8 deletions(-) diff --git a/libdd-heap-gotter/tests/install.rs b/libdd-heap-gotter/tests/install.rs index be3c668988..aa8ecaf86b 100644 --- a/libdd-heap-gotter/tests/install.rs +++ b/libdd-heap-gotter/tests/install.rs @@ -90,9 +90,10 @@ fn install_produces_sampled_allocations() { // handles the tagged pointer correctly (check + free raw). libc::free(p); - // Restore the default interval so we don't mess with anything - // after the test. - (*tl).sampling_interval = 512 * 1024; + // Disable sampling before restore. The forced interval=1 leaves + // remaining_bytes close to zero, so simply restoring the default + // interval can still let restore-time internal allocations sample. + (*tl).sampling_interval = 0; } libdd_heap_gotter::restore_heap_overrides(); @@ -125,7 +126,7 @@ fn realloc_null_produces_sampled_allocation() { ); libc::free(p); - (*tl).sampling_interval = 512 * 1024; + (*tl).sampling_interval = 0; } libdd_heap_gotter::restore_heap_overrides(); @@ -162,7 +163,7 @@ fn page_aligned_allocations_are_unsampled() { ); libc::free(p); - (*tl).sampling_interval = 512 * 1024; + (*tl).sampling_interval = 0; } libdd_heap_gotter::restore_heap_overrides(); @@ -220,7 +221,7 @@ fn realloc_of_sampled_allocation_preserves_data() { } libc::free(p2 as *mut c_void); - (*tl).sampling_interval = 512 * 1024; + (*tl).sampling_interval = 0; } libdd_heap_gotter::restore_heap_overrides(); @@ -253,14 +254,22 @@ fn restore_stops_sampling() { // Disable sampling before restore so internal allocations // during restore don't get tagged (they'd be freed through // the unpatched GOT afterwards, causing SIGABRT on x86_64). - (*tl).sampling_interval = 512 * 1024; + (*tl).sampling_interval = 0; } libdd_heap_gotter::restore_heap_overrides(); // After restore, malloc should return a plain pointer with no - // sample flag. + // sample flag, even if the sampler TLS is configured to sample + // every allocation. The GOT should no longer route malloc through + // the sampler. unsafe { + let tl = dd_tl_state_get_or_init(); + assert!(!tl.is_null()); + (*tl).sampling_interval = 1; + (*tl).remaining_bytes = 0; + (*tl).remaining_bytes_initialized = true; + let p = libc::malloc(128); assert!(!p.is_null()); From adc2f14380346f00208edd5302da8fcb3949a22b Mon Sep 17 00:00:00 2001 From: Scott Gerring Date: Fri, 3 Jul 2026 14:18:44 +0200 Subject: [PATCH 12/16] cleanup APIs --- libdd-heap-gotter/src/hooks.rs | 6 +++++- .../include/datadog/heap/allocation_created.h | 5 +++-- .../include/datadog/heap/allocation_freed.h | 2 +- .../include/datadog/heap/allocation_realloc.h | 6 ++++-- .../datadog/heap/allocation_requested.h | 9 ++++---- .../include/datadog/heap/sample_flag.h | 21 +++++++++++-------- .../include/datadog/heap/tl_state.h | 8 ++----- libdd-heap-sampler/src/allocation_freed.c | 2 +- libdd-heap-sampler/src/generated/bindings.rs | 20 ++++-------------- .../dd_heap_sampler_static_wrappers.c | 6 +++--- libdd-heap-sampler/src/lib.rs | 2 +- libdd-heap-sampler/src/sample_flag.c | 10 +-------- libdd-heap-sampler/src/tl_state.c | 8 ------- 13 files changed, 42 insertions(+), 63 deletions(-) diff --git a/libdd-heap-gotter/src/hooks.rs b/libdd-heap-gotter/src/hooks.rs index d44bf7a765..976f8a535c 100644 --- a/libdd-heap-gotter/src/hooks.rs +++ b/libdd-heap-gotter/src/hooks.rs @@ -67,6 +67,10 @@ type PosixMemalignFn = unsafe extern "C" fn(*mut *mut c_void, usize, usize) -> c type AlignedAllocFn = unsafe extern "C" fn(usize, usize) -> *mut c_void; /// Signature of the real `dlopen`. type DlopenFn = unsafe extern "C" fn(*const c_char, c_int) -> *mut c_void; +/// Linux RTLD_DEEPBIND. Some libcs we build against (notably musl on Alpine) +/// don't expose this constant through the Rust libc crate, but the flag value +/// is stable Linux ABI from . +const RTLD_DEEPBIND: c_int = 0x00008; /// Signature of a `pthread_create` start routine. type StartRoutine = unsafe extern "C" fn(*mut c_void) -> *mut c_void; /// Signature of the real `pthread_create`. @@ -181,7 +185,7 @@ pub unsafe extern "C" fn gotter_dlopen(filename: *const c_char, flags: c_int) -> return libc::dlopen(filename, flags); }; let handle = real(filename, flags); - if flags & libc::RTLD_DEEPBIND != 0 { + if flags & RTLD_DEEPBIND != 0 { // DEEPBIND changes symbol resolution order and causes issues with // GOT patching, so skip newly-loaded deep-bound libraries for now. return handle; diff --git a/libdd-heap-sampler/include/datadog/heap/allocation_created.h b/libdd-heap-sampler/include/datadog/heap/allocation_created.h index 1d7d86135f..715e8ce13c 100644 --- a/libdd-heap-sampler/include/datadog/heap/allocation_created.h +++ b/libdd-heap-sampler/include/datadog/heap/allocation_created.h @@ -24,7 +24,8 @@ * from dd_allocation_created so that we don't bloat the instruction * cache for the fast path. */ -void *dd_allocation_created_slow(void *raw, dd_alloc_req_t req); +void *dd_allocation_created_slow(void *raw, dd_alloc_req_t req) + __attribute__((warn_unused_result)); /* * Post-allocation hook. Pair with dd_allocation_requested. @@ -41,7 +42,7 @@ void *dd_allocation_created_slow(void *raw, dd_alloc_req_t req); * * Safe when raw == NULL (allocator failed): no USDT, guard still closed. */ -static inline __attribute__((always_inline)) +static inline __attribute__((always_inline, warn_unused_result)) void *dd_allocation_created(void *raw, dd_alloc_req_t req) { if (__builtin_expect(req.weight == 0, 1)) return raw; return dd_allocation_created_slow(raw, req); diff --git a/libdd-heap-sampler/include/datadog/heap/allocation_freed.h b/libdd-heap-sampler/include/datadog/heap/allocation_freed.h index 0e63bda996..d4daf18d9a 100644 --- a/libdd-heap-sampler/include/datadog/heap/allocation_freed.h +++ b/libdd-heap-sampler/include/datadog/heap/allocation_freed.h @@ -61,7 +61,7 @@ dd_alloc_freed_t dd_allocation_freed_slow(void *ptr, void *raw, size_t size, static inline __attribute__((always_inline)) dd_alloc_freed_t dd_allocation_freed(void *ptr, size_t size, size_t alignment) { void *raw; - if (__builtin_expect(dd_sample_flag_check_fast(ptr, &raw), 0)) { + if (__builtin_expect(dd_sample_flag_check(ptr, &raw), 0)) { return dd_allocation_freed_slow(ptr, raw, size, alignment); } diff --git a/libdd-heap-sampler/include/datadog/heap/allocation_realloc.h b/libdd-heap-sampler/include/datadog/heap/allocation_realloc.h index cc7dd2ca8e..6a9a5a2b91 100644 --- a/libdd-heap-sampler/include/datadog/heap/allocation_realloc.h +++ b/libdd-heap-sampler/include/datadog/heap/allocation_realloc.h @@ -74,7 +74,8 @@ typedef struct { * we hook. prepare() consumes the sampler flag in that case before the * frontend forwards the raw pointer to realloc(raw, 0). */ -dd_realloc_prep_t dd_allocation_realloc_prepare(void *old_user, size_t new_size); +dd_realloc_prep_t dd_allocation_realloc_prepare(void *old_user, size_t new_size) + __attribute__((warn_unused_result)); /* * Finalize the realloc. Given the return value of the underlying @@ -94,6 +95,7 @@ dd_realloc_prep_t dd_allocation_realloc_prepare(void *old_user, size_t new_size) * On sampled realloc failure (new_raw == NULL): returns NULL and leaves * old_user live with its sampler flag intact. */ -void *dd_allocation_realloc_commit(void *old_user, void *new_raw, dd_realloc_prep_t prep); +void *dd_allocation_realloc_commit(void *old_user, void *new_raw, dd_realloc_prep_t prep) + __attribute__((warn_unused_result)); #endif diff --git a/libdd-heap-sampler/include/datadog/heap/allocation_requested.h b/libdd-heap-sampler/include/datadog/heap/allocation_requested.h index 90541e2845..35421c5042 100644 --- a/libdd-heap-sampler/include/datadog/heap/allocation_requested.h +++ b/libdd-heap-sampler/include/datadog/heap/allocation_requested.h @@ -66,10 +66,11 @@ typedef struct { /* Slow path for an allocation request. This is only taken when we think we * need to sample, and is declared as a separate function to avoid bloating - the instruction cache of the fast path + * the instruction cache of the fast path */ dd_alloc_req_t dd_allocation_requested_slow(dd_tl_state_t *tl, size_t size, - size_t alignment); + size_t alignment) + __attribute__((warn_unused_result)); /* * Pre-allocation hook. Call BEFORE invoking the underlying allocator. @@ -89,7 +90,7 @@ dd_alloc_req_t dd_allocation_requested_slow(dd_tl_state_t *tl, size_t size, * by the paired dd_allocation_created call. ALWAYS pair them, even on * allocator failure (pass raw=NULL). */ -static inline __attribute__((always_inline)) +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 }; @@ -97,7 +98,7 @@ dd_alloc_req_t dd_allocation_requested(size_t size, size_t alignment) { // allocation is already in flight on this thread and something in its slow // path triggered another allocation), pass through without sampling. // Either condition is rare on a hot path, so mark the branch unlikely. - dd_tl_state_t *tl = dd_tl_state_get_or_init_fast(); + dd_tl_state_t *tl = dd_tl_state_get_or_init(); if (__builtin_expect(!tl || tl->reentry_guard, 0)) return out; // If we haven't crossed the sampling boundary, do nothing and bail. diff --git a/libdd-heap-sampler/include/datadog/heap/sample_flag.h b/libdd-heap-sampler/include/datadog/heap/sample_flag.h index ae852923ff..2d5a4dcb23 100644 --- a/libdd-heap-sampler/include/datadog/heap/sample_flag.h +++ b/libdd-heap-sampler/include/datadog/heap/sample_flag.h @@ -42,13 +42,6 @@ bool dd_sample_flag_thread_init(void); */ void *dd_sample_flag_apply(void *raw, size_t alignment); -/* - * If `user` was previously returned by dd_sample_flag_apply, write the - * raw pointer (to pass to the underlying free) into *raw_out and return - * true. Otherwise leave *raw_out untouched and return false. - */ -bool dd_sample_flag_check(void *user, void **raw_out); - /* * Non-destructive variant of dd_sample_flag_check. Useful for realloc: * callers can resolve the raw pointer before calling the underlying @@ -156,8 +149,13 @@ void *x86_raw_from_user(void *user, uint64_t offset) { return (void *)((uintptr_t)user - (uintptr_t)offset); } +/* + * If `user` was previously returned by dd_sample_flag_apply, write the + * raw pointer (to pass to the underlying free) into *raw_out and return + * true. Otherwise leave *raw_out untouched and return false. + */ static inline __attribute__((always_inline)) -bool dd_sample_flag_check_fast(void *user, void **raw_out) { +bool dd_sample_flag_check(void *user, void **raw_out) { if (((uintptr_t)user & (DD_PAGE_SIZE - 1)) < DD_HEADER_BYTES) { return false; } @@ -192,8 +190,13 @@ bool dd_sample_flag_check_fast(void *user, void **raw_out) { #define DD_TBI_TAG_MASK ((uintptr_t)0xFFu << 56) #define DD_TBI_TAGGED ((uintptr_t)DD_TBI_TAG << 56) +/* + * If `user` was previously returned by dd_sample_flag_apply, write the + * raw pointer (to pass to the underlying free) into *raw_out and return + * true. Otherwise leave *raw_out untouched and return false. + */ static inline __attribute__((always_inline)) -bool dd_sample_flag_check_fast(void *user, void **raw_out) { +bool dd_sample_flag_check(void *user, void **raw_out) { uintptr_t addr = (uintptr_t)user; if ((addr & DD_TBI_TAG_MASK) != DD_TBI_TAGGED) { return false; diff --git a/libdd-heap-sampler/include/datadog/heap/tl_state.h b/libdd-heap-sampler/include/datadog/heap/tl_state.h index 3faee19534..17c81f795f 100644 --- a/libdd-heap-sampler/include/datadog/heap/tl_state.h +++ b/libdd-heap-sampler/include/datadog/heap/tl_state.h @@ -121,10 +121,8 @@ extern _Thread_local dd_tl_state_t dd_tl_state_storage; * Returns the current thread's state, or NULL if not yet initialised. * Never allocates. Callers must treat NULL as "don't sample". */ -dd_tl_state_t *dd_tl_state_get(void); - static inline __attribute__((always_inline)) -dd_tl_state_t *dd_tl_state_get_fast(void) { +dd_tl_state_t *dd_tl_state_get(void) { if (__builtin_expect(!dd_tl_state_storage.initialized, 0)) return NULL; return &dd_tl_state_storage; } @@ -138,10 +136,8 @@ dd_tl_state_t *dd_tl_state_init(void); /* * Returns the current thread's tracking state, initializing it on first use. */ -dd_tl_state_t *dd_tl_state_get_or_init(void); - static inline __attribute__((always_inline)) -dd_tl_state_t *dd_tl_state_get_or_init_fast(void) { +dd_tl_state_t *dd_tl_state_get_or_init(void) { if (__builtin_expect(dd_tl_state_storage.initialized, 1)) return &dd_tl_state_storage; return dd_tl_state_init(); } diff --git a/libdd-heap-sampler/src/allocation_freed.c b/libdd-heap-sampler/src/allocation_freed.c index 609fc06f6c..bf50e83bd3 100644 --- a/libdd-heap-sampler/src/allocation_freed.c +++ b/libdd-heap-sampler/src/allocation_freed.c @@ -9,7 +9,7 @@ /* * Slow path for dd_allocation_freed. We only arrive here when - * dd_sample_flag_check_fast confirmed that ptr carries the sample flag, + * dd_sample_flag_check confirmed that ptr carries the sample flag, * meaning this allocation was previously sampled. * * Fires the ddheap:free USDT with the user-visible pointer, then returns diff --git a/libdd-heap-sampler/src/generated/bindings.rs b/libdd-heap-sampler/src/generated/bindings.rs index d8cab0e81b..8bb5650161 100644 --- a/libdd-heap-sampler/src/generated/bindings.rs +++ b/libdd-heap-sampler/src/generated/bindings.rs @@ -33,22 +33,16 @@ const _: () = { ["Offset of field: dd_tl_state_t::rng"][::std::mem::offset_of!(dd_tl_state_t, rng) - 20usize]; }; unsafe extern "C" { + #[link_name = "dd_tl_state_get__extern"] pub fn dd_tl_state_get() -> *mut dd_tl_state_t; } -unsafe extern "C" { - #[link_name = "dd_tl_state_get_fast__extern"] - pub fn dd_tl_state_get_fast() -> *mut dd_tl_state_t; -} unsafe extern "C" { pub fn dd_tl_state_init() -> *mut dd_tl_state_t; } unsafe extern "C" { + #[link_name = "dd_tl_state_get_or_init__extern"] pub fn dd_tl_state_get_or_init() -> *mut dd_tl_state_t; } -unsafe extern "C" { - #[link_name = "dd_tl_state_get_or_init_fast__extern"] - pub fn dd_tl_state_get_or_init_fast() -> *mut dd_tl_state_t; -} #[repr(C)] #[derive(Debug, Copy, Clone)] pub struct dd_alloc_req_t { @@ -103,12 +97,6 @@ unsafe extern "C" { alignment: usize, ) -> *mut ::std::os::raw::c_void; } -unsafe extern "C" { - pub fn dd_sample_flag_check( - user: *mut ::std::os::raw::c_void, - raw_out: *mut *mut ::std::os::raw::c_void, - ) -> bool; -} unsafe extern "C" { pub fn dd_sample_flag_peek( user: *mut ::std::os::raw::c_void, @@ -117,8 +105,8 @@ unsafe extern "C" { ) -> bool; } unsafe extern "C" { - #[link_name = "dd_sample_flag_check_fast__extern"] - pub fn dd_sample_flag_check_fast( + #[link_name = "dd_sample_flag_check__extern"] + pub fn dd_sample_flag_check( user: *mut ::std::os::raw::c_void, raw_out: *mut *mut ::std::os::raw::c_void, ) -> bool; diff --git a/libdd-heap-sampler/src/generated/dd_heap_sampler_static_wrappers.c b/libdd-heap-sampler/src/generated/dd_heap_sampler_static_wrappers.c index 31b2ca8197..321d0d06ae 100644 --- a/libdd-heap-sampler/src/generated/dd_heap_sampler_static_wrappers.c +++ b/libdd-heap-sampler/src/generated/dd_heap_sampler_static_wrappers.c @@ -13,9 +13,9 @@ // Static wrappers -dd_tl_state_t * dd_tl_state_get_fast__extern(void) { return dd_tl_state_get_fast(); } -dd_tl_state_t * dd_tl_state_get_or_init_fast__extern(void) { return dd_tl_state_get_or_init_fast(); } +dd_tl_state_t * dd_tl_state_get__extern(void) { return dd_tl_state_get(); } +dd_tl_state_t * dd_tl_state_get_or_init__extern(void) { return dd_tl_state_get_or_init(); } dd_alloc_req_t dd_allocation_requested__extern(size_t size, size_t alignment) { return dd_allocation_requested(size, alignment); } void * dd_allocation_created__extern(void *raw, dd_alloc_req_t req) { return dd_allocation_created(raw, req); } -bool dd_sample_flag_check_fast__extern(void *user, void **raw_out) { return dd_sample_flag_check_fast(user, raw_out); } +bool dd_sample_flag_check__extern(void *user, void **raw_out) { return dd_sample_flag_check(user, raw_out); } dd_alloc_freed_t dd_allocation_freed__extern(void *ptr, size_t size, size_t alignment) { return dd_allocation_freed(ptr, size, alignment); } diff --git a/libdd-heap-sampler/src/lib.rs b/libdd-heap-sampler/src/lib.rs index 00b482bda9..e44447528a 100644 --- a/libdd-heap-sampler/src/lib.rs +++ b/libdd-heap-sampler/src/lib.rs @@ -137,7 +137,7 @@ mod tests { buf[header_idx + 8..header_idx + 16].copy_from_slice(&8u64.to_ne_bytes()); let mut raw = core::ptr::null_mut(); - let sampled = unsafe { dd_sample_flag_check_fast(user_addr as *mut c_void, &mut raw) }; + let sampled = unsafe { dd_sample_flag_check(user_addr as *mut c_void, &mut raw) }; assert!(!sampled); assert!(raw.is_null()); diff --git a/libdd-heap-sampler/src/sample_flag.c b/libdd-heap-sampler/src/sample_flag.c index 58ba9e82f3..3de1f80d11 100644 --- a/libdd-heap-sampler/src/sample_flag.c +++ b/libdd-heap-sampler/src/sample_flag.c @@ -47,10 +47,6 @@ bool dd_sample_flag_peek(void *user, void **raw_out, size_t *offset_out) { return true; } -bool dd_sample_flag_check(void *user, void **raw_out) { - return dd_sample_flag_check_fast(user, raw_out); -} - #elif defined(__aarch64__) /* @@ -78,17 +74,13 @@ void *dd_sample_flag_apply(void *raw, size_t alignment) { } bool dd_sample_flag_peek(void *user, void **raw_out, size_t *offset_out) { - if (!dd_sample_flag_check_fast(user, raw_out)) { + if (!dd_sample_flag_check(user, raw_out)) { return false; } *offset_out = 0; return true; } -bool dd_sample_flag_check(void *user, void **raw_out) { - return dd_sample_flag_check_fast(user, raw_out); -} - #if defined(__linux__) #include diff --git a/libdd-heap-sampler/src/tl_state.c b/libdd-heap-sampler/src/tl_state.c index ea2ad3b33e..49f4a06309 100644 --- a/libdd-heap-sampler/src/tl_state.c +++ b/libdd-heap-sampler/src/tl_state.c @@ -60,10 +60,6 @@ static void tl_state_populate(dd_tl_state_t *st) { } } -dd_tl_state_t *dd_tl_state_get(void) { - return dd_tl_state_get_fast(); -} - /* * Initialises TLS for this thread on first call; returns NULL on subsequent * calls (state already exists). Callers that need the pointer regardless @@ -75,7 +71,3 @@ dd_tl_state_t *dd_tl_state_init(void) { tl_state_populate(&dd_tl_state_storage); return &dd_tl_state_storage; } - -dd_tl_state_t *dd_tl_state_get_or_init(void) { - return dd_tl_state_get_or_init_fast(); -} From 9fc56030320f048d74b21296998edceb79196566 Mon Sep 17 00:00:00 2001 From: Scott Gerring Date: Fri, 3 Jul 2026 15:55:17 +0200 Subject: [PATCH 13/16] Apply suggestions from code review Co-authored-by: Yann Hamdaoui --- libdd-heap-allocator/examples/usdt_demo.rs | 2 +- libdd-heap-allocator/src/allocator.rs | 6 +++++- libdd-heap-gotter-ffi/README.md | 1 - libdd-heap-gotter-ffi/src/lib.rs | 1 + libdd-heap-gotter-ffi/tests/install.rs | 1 + 5 files changed, 8 insertions(+), 3 deletions(-) diff --git a/libdd-heap-allocator/examples/usdt_demo.rs b/libdd-heap-allocator/examples/usdt_demo.rs index ebd4829dcd..e30dbada29 100644 --- a/libdd-heap-allocator/examples/usdt_demo.rs +++ b/libdd-heap-allocator/examples/usdt_demo.rs @@ -58,4 +58,4 @@ mod linux { sleep(Duration::from_secs(1)); } } -} // mod linux +} diff --git a/libdd-heap-allocator/src/allocator.rs b/libdd-heap-allocator/src/allocator.rs index ac996e0f5d..9ed8b45598 100644 --- a/libdd-heap-allocator/src/allocator.rs +++ b/libdd-heap-allocator/src/allocator.rs @@ -36,7 +36,8 @@ impl SampledAllocator { } unsafe impl GlobalAlloc for SampledAllocator { - #[cfg(target_os = "linux")] + #[cfg(target_os = "linux")] + #[inline] unsafe fn alloc(&self, layout: Layout) -> *mut u8 { let req = dd_allocation_requested(layout.size(), layout.align()); // Sampled paths may bump the size for inline flag storage; @@ -47,11 +48,13 @@ unsafe impl GlobalAlloc for SampledAllocator { } #[cfg(not(target_os = "linux"))] + #[inline(always)] unsafe fn alloc(&self, layout: Layout) -> *mut u8 { self.inner.alloc(layout) } #[cfg(target_os = "linux")] + #[inline] unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) { let freed = dd_allocation_freed(ptr.cast(), layout.size(), layout.align()); let inner_layout = Layout::from_size_align_unchecked(freed.size, layout.align()); @@ -59,6 +62,7 @@ unsafe impl GlobalAlloc for SampledAllocator { } #[cfg(not(target_os = "linux"))] + #[inline(always)] unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) { self.inner.dealloc(ptr, layout); } diff --git a/libdd-heap-gotter-ffi/README.md b/libdd-heap-gotter-ffi/README.md index 8b1e5e90d5..547b4d9f42 100644 --- a/libdd-heap-gotter-ffi/README.md +++ b/libdd-heap-gotter-ffi/README.md @@ -32,7 +32,6 @@ cargo check -p libdd-heap-gotter-ffi The `cdylib_demo` example loads the generated shared library with `dlopen`, resolves the C ABI symbols with `dlsym`, installs the GOT hooks, and produces allocation pressure. ```bash -cargo build -p libdd-heap-gotter-ffi cargo run -p libdd-heap-gotter-ffi --example cdylib_demo ``` diff --git a/libdd-heap-gotter-ffi/src/lib.rs b/libdd-heap-gotter-ffi/src/lib.rs index 0f02e5489c..e69f00c061 100644 --- a/libdd-heap-gotter-ffi/src/lib.rs +++ b/libdd-heap-gotter-ffi/src/lib.rs @@ -11,6 +11,7 @@ #![cfg_attr(not(test), deny(clippy::expect_used))] #![cfg_attr(not(test), deny(clippy::todo))] #![cfg_attr(not(test), deny(clippy::unimplemented))] +#![cfg_attr(not(test), deny(clippy::unreachable))] use function_name::named; use libdd_common_ffi::{wrap_with_void_ffi_result, VoidResult}; diff --git a/libdd-heap-gotter-ffi/tests/install.rs b/libdd-heap-gotter-ffi/tests/install.rs index 795ed8f79f..2361d9ba25 100644 --- a/libdd-heap-gotter-ffi/tests/install.rs +++ b/libdd-heap-gotter-ffi/tests/install.rs @@ -14,6 +14,7 @@ use libdd_heap_gotter_ffi::{ ddog_heap_gotter_install, ddog_heap_gotter_is_installed, ddog_heap_gotter_restore, }; +#[track_caller] fn assert_ok(result: VoidResult, what: &str) { match result { VoidResult::Ok => {} From ff846ec85da8dcb4094a012f818a23079213e6cf Mon Sep 17 00:00:00 2001 From: Scott Gerring Date: Fri, 3 Jul 2026 16:00:09 +0200 Subject: [PATCH 14/16] more code review work --- libdd-heap-allocator/README.md | 6 +++++- libdd-heap-allocator/benches/sampler_overhead.rs | 11 +++-------- libdd-heap-gotter-ffi/README.md | 8 ++++++-- libdd-heap-gotter-ffi/src/lib.rs | 1 + libdd-heap-sampler/include/datadog/heap/tl_state.h | 13 ++++++++----- libdd-heap-sampler/src/generated/bindings.rs | 2 +- libdd-heap-sampler/src/lib.rs | 11 +++++++---- libdd-heap-sampler/src/tl_state.c | 11 +++++------ 8 files changed, 36 insertions(+), 27 deletions(-) diff --git a/libdd-heap-allocator/README.md b/libdd-heap-allocator/README.md index affe32b12d..ab9af1f708 100644 --- a/libdd-heap-allocator/README.md +++ b/libdd-heap-allocator/README.md @@ -15,13 +15,17 @@ use std::alloc::System; static ALLOC: SampledAllocator = SampledAllocator::::DEFAULT; ``` -To wrap a custom allocator instead; note that this is kind of ill-advised; we want to see _all_ allocations for the process: +To wrap a custom allocator instead: ```rust #[global_allocator] static ALLOC: SampledAllocator = SampledAllocator::new(MyAllocator::new()); ``` +For profiling, prefer wrapping the allocator that is actually installed as the +process global allocator. Heap profiling is most useful when all allocations in +the process pass through the sampled wrapper. + See [`examples/usdt_demo.rs`](examples/usdt_demo.rs) for a runnable demo that fires USDT probes in a loop for `bpftrace` to observe. ## Benchmarking sampler overhead diff --git a/libdd-heap-allocator/benches/sampler_overhead.rs b/libdd-heap-allocator/benches/sampler_overhead.rs index bc2cc5e97a..f4a022b6ea 100644 --- a/libdd-heap-allocator/benches/sampler_overhead.rs +++ b/libdd-heap-allocator/benches/sampler_overhead.rs @@ -16,8 +16,8 @@ mod linux_bench { use criterion::{criterion_group, BenchmarkId, Criterion, Throughput}; use libdd_heap_allocator::SampledAllocator; use libdd_heap_sampler::{ - dd_allocation_created, dd_allocation_freed, dd_allocation_requested, dd_tl_state_get, - dd_tl_state_init, + dd_allocation_created, dd_allocation_freed, dd_allocation_requested, + dd_tl_state_get_or_init, }; use std::alloc::{GlobalAlloc, Layout, System}; use std::hint::black_box; @@ -49,12 +49,7 @@ mod linux_bench { } unsafe fn sampler_tl_state() -> *mut libdd_heap_sampler::dd_tl_state_t { - let mut tl = unsafe { dd_tl_state_get() }; - if tl.is_null() { - let _ = unsafe { dd_tl_state_init() }; - tl = unsafe { dd_tl_state_get() }; - } - tl + unsafe { dd_tl_state_get_or_init() } } unsafe fn pin_sampler_to_fast_path() { diff --git a/libdd-heap-gotter-ffi/README.md b/libdd-heap-gotter-ffi/README.md index 547b4d9f42..ce790cd565 100644 --- a/libdd-heap-gotter-ffi/README.md +++ b/libdd-heap-gotter-ffi/README.md @@ -19,14 +19,18 @@ The API installs hooks for supported allocator symbols, updates hooks after new The shared library containing these hooks must remain loaded while overrides are installed. Patched GOT entries point at functions in this library, so unloading it before calling `ddog_heap_gotter_restore()` can leave dangling function pointers and crash the process. +TODO: consider hooking `dlclose` too, so we can at least partially protect against unloading this library while its hooks are still installed. + ## Building -This crate follows the standard libdatadog FFI layout: it builds `staticlib` and `cdylib` artifacts and generates a C header with cbindgen. +This crate follows the standard libdatadog FFI layout and can produce `staticlib` and `cdylib` artifacts. ```bash -cargo check -p libdd-heap-gotter-ffi +cargo build -p libdd-heap-gotter-ffi ``` +The C header is generated with cbindgen by the libdatadog release tooling, not by ordinary `cargo build`. + ## Dynamic-loading demo The `cdylib_demo` example loads the generated shared library with `dlopen`, resolves the C ABI symbols with `dlsym`, installs the GOT hooks, and produces allocation pressure. diff --git a/libdd-heap-gotter-ffi/src/lib.rs b/libdd-heap-gotter-ffi/src/lib.rs index e69f00c061..f4e171fe7b 100644 --- a/libdd-heap-gotter-ffi/src/lib.rs +++ b/libdd-heap-gotter-ffi/src/lib.rs @@ -13,6 +13,7 @@ #![cfg_attr(not(test), deny(clippy::unimplemented))] #![cfg_attr(not(test), deny(clippy::unreachable))] +// `wrap_with_void_ffi_result!` uses `function_name!()` below. use function_name::named; use libdd_common_ffi::{wrap_with_void_ffi_result, VoidResult}; diff --git a/libdd-heap-sampler/include/datadog/heap/tl_state.h b/libdd-heap-sampler/include/datadog/heap/tl_state.h index 17c81f795f..ec9f0a1d42 100644 --- a/libdd-heap-sampler/include/datadog/heap/tl_state.h +++ b/libdd-heap-sampler/include/datadog/heap/tl_state.h @@ -128,18 +128,21 @@ dd_tl_state_t *dd_tl_state_get(void) { } /* - * Initializes the current thread's tracking state if not already present. - * Returns the new state, or NULL if state already exists on this thread. + * Ensures the current thread's tracking state exists, initializing it on the + * first call and doing nothing on subsequent calls. This is a fire-and-forget + * command: call it to eagerly warm TLS (e.g. from a thread-start hook) so the + * first tracked allocation on the thread doesn't pay the init cost. Callers + * that need the pointer should use dd_tl_state_get_or_init() instead. */ -dd_tl_state_t *dd_tl_state_init(void); +void dd_tl_state_init(void); /* * Returns the current thread's tracking state, initializing it on first use. */ static inline __attribute__((always_inline)) dd_tl_state_t *dd_tl_state_get_or_init(void) { - if (__builtin_expect(dd_tl_state_storage.initialized, 1)) return &dd_tl_state_storage; - return dd_tl_state_init(); + if (__builtin_expect(!dd_tl_state_storage.initialized, 0)) dd_tl_state_init(); + return &dd_tl_state_storage; } #endif diff --git a/libdd-heap-sampler/src/generated/bindings.rs b/libdd-heap-sampler/src/generated/bindings.rs index 8bb5650161..3859f3614d 100644 --- a/libdd-heap-sampler/src/generated/bindings.rs +++ b/libdd-heap-sampler/src/generated/bindings.rs @@ -37,7 +37,7 @@ unsafe extern "C" { pub fn dd_tl_state_get() -> *mut dd_tl_state_t; } unsafe extern "C" { - pub fn dd_tl_state_init() -> *mut dd_tl_state_t; + pub fn dd_tl_state_init(); } unsafe extern "C" { #[link_name = "dd_tl_state_get_or_init__extern"] diff --git a/libdd-heap-sampler/src/lib.rs b/libdd-heap-sampler/src/lib.rs index e44447528a..be5d9ad41d 100644 --- a/libdd-heap-sampler/src/lib.rs +++ b/libdd-heap-sampler/src/lib.rs @@ -174,12 +174,15 @@ mod tests { } #[test] - fn tl_state_init_then_get_returns_same_non_null() { + fn tl_state_init_populates_state() { std::thread::spawn(|| unsafe { - let initialized = dd_tl_state_init(); - assert!(!initialized.is_null(), "init returned NULL on fresh thread"); + assert!( + dd_tl_state_get().is_null(), + "fresh thread should start without sampler TLS" + ); + dd_tl_state_init(); let got = dd_tl_state_get(); - assert_eq!(initialized, got); + assert!(!got.is_null(), "init should populate sampler TLS"); assert_eq!((*got).sampling_interval, 512 * 1024); }) .join() diff --git a/libdd-heap-sampler/src/tl_state.c b/libdd-heap-sampler/src/tl_state.c index 49f4a06309..3bdc306c84 100644 --- a/libdd-heap-sampler/src/tl_state.c +++ b/libdd-heap-sampler/src/tl_state.c @@ -61,13 +61,12 @@ static void tl_state_populate(dd_tl_state_t *st) { } /* - * Initialises TLS for this thread on first call; returns NULL on subsequent - * calls (state already exists). Callers that need the pointer regardless - * should use dd_tl_state_get_or_init(). + * Initialises TLS for this thread on the first call and is a no-op on + * subsequent calls. Fire-and-forget: callers that need the pointer should + * use dd_tl_state_get_or_init(). */ -dd_tl_state_t *dd_tl_state_init(void) { - if (dd_tl_state_storage.initialized) return NULL; +void dd_tl_state_init(void) { + if (dd_tl_state_storage.initialized) return; tl_state_populate(&dd_tl_state_storage); - return &dd_tl_state_storage; } From df27ec1a11cdbd89c909fc9558da1f2cb78f0937 Mon Sep 17 00:00:00 2001 From: Scott Gerring Date: Fri, 3 Jul 2026 16:41:35 +0200 Subject: [PATCH 15/16] better tests --- libdd-heap-gotter/tests/install.rs | 127 +++++++++++++++++++++++++++++ 1 file changed, 127 insertions(+) diff --git a/libdd-heap-gotter/tests/install.rs b/libdd-heap-gotter/tests/install.rs index aa8ecaf86b..68ae16b8a8 100644 --- a/libdd-heap-gotter/tests/install.rs +++ b/libdd-heap-gotter/tests/install.rs @@ -227,6 +227,133 @@ fn realloc_of_sampled_allocation_preserves_data() { libdd_heap_gotter::restore_heap_overrides(); } +/// Allocate `size` bytes at `align` through the (hooked) libc. Uses +/// posix_memalign for real alignments and malloc for align <= 1; +/// posix_memalign requires alignment to be a power-of-two multiple of +/// sizeof(void*), which every alignment in the menu except 1 satisfies, +/// and (unlike aligned_alloc) it places no multiple-of-alignment +/// constraint on the size. +unsafe fn alloc_aligned(align: usize, size: usize) -> *mut c_void { + if align <= 1 { + libc::malloc(size) + } else { + let mut out: *mut c_void = std::ptr::null_mut(); + if libc::posix_memalign(&mut out, align, size) != 0 { + std::ptr::null_mut() + } else { + out + } + } +} + +/// Stress the realloc + free paths across a matrix of alignments and +/// sizes with sampling forced on. Mirrors examples/gotter_usdt_demo.rs's +/// allocation menu, which straddles DD_SAMPLE_ALIGNMENT_CAP (1024 == cap, +/// 2048/4096/8192 above it), and is a regression guard for the class of +/// crash where a sampled, page-aligned pointer is mis-recovered on +/// free/realloc and an invalid pointer is handed to the real allocator +/// (munmap_chunk(): invalid pointer / SIGABRT). +/// +/// The headline assertion is that the whole matrix completes without +/// aborting; per-iteration checks confirm the surviving prefix is +/// preserved across each realloc, and `saw_sampled` confirms the tagged +/// pointer / raw-recovery path was actually exercised (so a regression +/// that silently stopped sampling can't make this pass trivially). +#[test] +#[serial] +fn realloc_stress_across_alignments_preserves_data() { + // 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]; + // A realloc walk per allocation: grow, grow, shrink, grow. + const SIZE_WALK: &[usize] = &[48, 200, 1000, 64, 4096]; + + // Deterministic, per-alignment-distinct pattern so we can verify the + // surviving prefix after each realloc. Pure arithmetic - no allocation. + fn pattern_byte(seed: usize, i: usize) -> u8 { + (i as u8) + .wrapping_mul(31) + .wrapping_add((seed as u8).wrapping_mul(101)) + ^ 0x5A + } + + let installed = libdd_heap_gotter::install_heap_overrides(); + assert!(installed); + + let mut saw_sampled = false; + + unsafe { + let tl = dd_tl_state_get_or_init(); + assert!(!tl.is_null()); + // Force every allocation to sample so the tagged-pointer / + // raw-recovery paths are exercised on every eligible iteration. + // Everything below this point avoids incidental Rust allocation + // (const inputs, raw-pointer writes only) so the sampled window + // isn't polluted by the test's own bookkeeping. + (*tl).sampling_interval = 1; + (*tl).remaining_bytes = 0; + (*tl).remaining_bytes_initialized = true; + + for (seed, &align) in ALIGNMENTS.iter().enumerate() { + let mut size = SIZE_WALK[0]; + + let mut p = alloc_aligned(align, size) as *mut u8; + assert!(!p.is_null(), "alloc failed (align={align}, size={size})"); + + // Record whether this block was actually sampled (peek is + // non-destructive, so it's safe before the realloc walk). + let mut raw: *mut c_void = std::ptr::null_mut(); + let mut offset: usize = 0; + if dd_sample_flag_peek(p as *mut c_void, &mut raw, &mut offset) { + saw_sampled = true; + } + + for i in 0..size { + *p.add(i) = pattern_byte(seed, i); + } + + for &new_size in &SIZE_WALK[1..] { + let p2 = libc::realloc(p as *mut c_void, new_size) as *mut u8; + assert!( + !p2.is_null(), + "realloc failed (align={align}, {size}->{new_size})" + ); + + // The overlapping prefix must survive the realloc. + let preserved = size.min(new_size); + for i in 0..preserved { + assert_eq!( + *p2.add(i), + pattern_byte(seed, i), + "corruption at byte {i} (align={align}, {size}->{new_size})" + ); + } + // Repaint the whole new region for the next realloc to keep. + for i in 0..new_size { + *p2.add(i) = pattern_byte(seed, i); + } + + p = p2; + size = new_size; + } + + libc::free(p as *mut c_void); + } + + // Disable sampling before restore so restore-time internal + // allocations aren't tagged (freed through the unpatched GOT + // afterwards would SIGABRT on x86_64). + (*tl).sampling_interval = 0; + } + + libdd_heap_gotter::restore_heap_overrides(); + + assert!( + saw_sampled, + "expected at least one allocation to be sampled with interval=1" + ); +} + /// Confirm that after restore, allocations are no longer sampled. #[test] #[serial] From bebaa383fbef69ac25880c7b9954dabc44259b67 Mon Sep 17 00:00:00 2001 From: Scott Gerring Date: Fri, 3 Jul 2026 17:09:45 +0200 Subject: [PATCH 16/16] more review comments & a correctness fix --- .../examples/gotter_usdt_demo.rs | 26 +++-- libdd-heap-gotter/src/elf.rs | 98 ++++++++----------- libdd-heap-gotter/src/hooks.rs | 2 +- libdd-heap-gotter/src/lib.rs | 16 ++- libdd-heap-gotter/tests/install.rs | 85 +++++++++------- libdd-heap-sampler/src/allocation_realloc.c | 55 +++++++++-- 6 files changed, 163 insertions(+), 119 deletions(-) diff --git a/libdd-heap-gotter/examples/gotter_usdt_demo.rs b/libdd-heap-gotter/examples/gotter_usdt_demo.rs index e423aa8ad7..30c7aa2553 100644 --- a/libdd-heap-gotter/examples/gotter_usdt_demo.rs +++ b/libdd-heap-gotter/examples/gotter_usdt_demo.rs @@ -62,11 +62,18 @@ mod linux { unsafe impl Send for LiveAlloc {} unsafe impl Sync for LiveAlloc {} impl LiveAlloc { - fn as_slice_mut(&self) -> &'static mut [u8] { + fn as_slice(&self) -> &[u8] { // SAFETY: `ptr` is the return of a libc allocator and // `size` is the user-requested size. Lifetime is scoped by - // the caller of `as_slice_mut` and does not outlive the - // allocation. + // the caller and does not outlive the allocation. + unsafe { std::slice::from_raw_parts(self.ptr, self.size) } + } + + fn as_slice_mut(&mut self) -> &mut [u8] { + // SAFETY: `ptr` is the return of a libc allocator and + // `size` is the user-requested size. `&mut self` prevents + // callers from creating two mutable slices to the same live + // allocation through this helper. unsafe { std::slice::from_raw_parts_mut(self.ptr, self.size) } } } @@ -80,6 +87,7 @@ mod linux { fn new(seed: u64) -> Self { Rng(seed) } + fn next(&mut self) -> u64 { let mut z = self.0.wrapping_add(0x9e37_79b9_7f4a_7c15); self.0 = z; @@ -87,9 +95,11 @@ mod linux { z = (z ^ (z >> 27)).wrapping_mul(0x94d0_49bb_1331_11eb); z ^ (z >> 31) } + fn range(&mut self, lo: usize, hi: usize) -> usize { lo + (self.next() as usize) % (hi - lo).max(1) } + fn choice<'a, T>(&mut self, xs: &'a [T]) -> &'a T { &xs[(self.next() as usize) % xs.len()] } @@ -215,7 +225,7 @@ mod linux { unsafe fn do_realloc(old: LiveAlloc, new_size: usize) -> Option { // Verify old contents before releasing the block. - verify_content(old.as_slice_mut(), old.seed); + verify_content(old.as_slice(), old.seed); let new_ptr = libc::realloc(old.ptr as *mut libc::c_void, new_size) as *mut u8; if new_ptr.is_null() { // Old block is still live on realloc failure. Return it as-is. @@ -242,7 +252,7 @@ mod linux { } unsafe fn do_free(a: LiveAlloc) { - verify_content(a.as_slice_mut(), a.seed); + verify_content(a.as_slice(), a.seed); libc::free(a.ptr as *mut libc::c_void); } @@ -280,7 +290,8 @@ mod linux { }; let Some(mut a) = a else { continue }; a.seed = rng.next(); - fill_content(a.as_slice_mut(), a.seed); + let seed = a.seed; + fill_content(a.as_slice_mut(), seed); allocs.fetch_add(1, Ordering::Relaxed); if live.len() == MAX_LIVE { let idx = rng.range(0, live.len()); @@ -301,7 +312,8 @@ mod linux { // min(old_size,new_size) are undefined per the // realloc contract. resized.seed = rng.next(); - fill_content(resized.as_slice_mut(), resized.seed); + let seed = resized.seed; + fill_content(resized.as_slice_mut(), seed); reallocs.fetch_add(1, Ordering::Relaxed); live.push(resized); } else if !live.is_empty() { diff --git a/libdd-heap-gotter/src/elf.rs b/libdd-heap-gotter/src/elf.rs index d0198de7aa..6b255bb92a 100644 --- a/libdd-heap-gotter/src/elf.rs +++ b/libdd-heap-gotter/src/elf.rs @@ -13,8 +13,8 @@ //! Scope: //! * 64-bit ELF only (`Elf64_*`). Other targets are gated out at compile time via `#[cfg]` on the //! parent module. -//! * GNU hash tables only - `DT_HASH` is skipped, matching ddprof (sometimes points into kernel -//! space on old glibcs). +//! * GNU hash tables only - `DT_HASH` is skipped because it has caused problems on older glibc +//! systems. //! * REL / RELA / JMPREL relocation arrays. use core::ffi::{c_char, c_int, c_void}; @@ -69,11 +69,10 @@ struct DynamicInfo { } impl DynamicInfo { - /// Read DT_* entries out of a PT_DYNAMIC array. ddprof's - /// `retrieve_dynamic_info` handles the glibc-vs-musl quirk where - /// glibc stores absolute addresses in DT entries while musl stores - /// load-relative offsets; we use the same `addr > base ? addr : base + addr` - /// heuristic. + /// Read DT_* entries out of a PT_DYNAMIC array. Handles the + /// glibc-vs-musl quirk where glibc stores absolute addresses in DT + /// entries while musl stores load-relative offsets; we use the + /// `addr > base ? addr : base + addr` heuristic. unsafe fn from_phdr(info: &dl_phdr_info) -> Option { let phdrs = std::slice::from_raw_parts(info.dlpi_phdr, info.dlpi_phnum as usize); let dyn_phdr = phdrs.iter().find(|p| p.p_type == PT_DYNAMIC)?; @@ -151,9 +150,11 @@ impl DynamicInfo { }; let sym_count = gnu_hash_symbol_count(gnu_hash, gnu_hash_words).unwrap_or_else(|| { // Fallback for degenerate .gnu.hash (e.g. executables with only - // undefined imports): estimate dynsym entry count from the - // distance between DT_SYMTAB and DT_STRTAB. This works because - // the linker always places .dynsym immediately before .dynstr. + // undefined imports): estimate dynsym entry count from the common + // .dynsym-before-.dynstr layout. This is a heuristic, not an ELF + // guarantee. If it underestimates we may skip patching some + // relocations; valid relocation indexes should still keep an + // overestimate from faulting on normal loaded objects. let symtab_addr = symtab as usize; let strtab_addr = strtab as usize; if strtab_addr > symtab_addr { @@ -321,27 +322,25 @@ unsafe fn gnu_hash_lookup(info: &DynamicInfo, name: &[u8]) -> Option None } -/// Mirror of ddprof's `check`: defining symbol, function/object/notype. +/// Return whether this is a defining function/object/notype symbol. fn check_sym(sym: &Elf64_Sym) -> bool { const SHN_ABS: u16 = 0xfff1; let stt = sym.st_info & 0xf; - if sym.st_value == 0 && sym.st_shndx != SHN_ABS { - return false; - } - // STT_NOTYPE(0), STT_OBJECT(1), STT_FUNC(2), STT_GNU_IFUNC(10) - matches!(stt, 0 | 1 | 2 | 10) + (sym.st_value != 0 || sym.st_shndx == SHN_ABS) && + // STT_NOTYPE(0), STT_OBJECT(1), STT_FUNC(2), STT_GNU_IFUNC(10) + matches!(stt, 0 | 1 | 2 | 10) } /// Visit each loaded ELF object once. `is_exe` is true only on the /// first callback (the main executable). The callback returns `true` to /// stop iteration. -fn iterate_libraries bool>(mut cb: F) { +fn iterate_libraries(mut callback: impl FnMut(&dl_phdr_info, bool) -> bool) { struct Ctx<'a> { - cb: &'a mut dyn FnMut(&dl_phdr_info, bool) -> bool, + callback: &'a mut dyn FnMut(&dl_phdr_info, bool) -> bool, is_first: bool, } let mut ctx = Ctx { - cb: &mut cb, + callback: &mut callback, is_first: true, }; @@ -354,15 +353,12 @@ fn iterate_libraries bool>(mut cb: F) { let ctx = &mut *(data as *mut Ctx); let is_exe = ctx.is_first; ctx.is_first = false; - (ctx.cb)(&*info, is_exe) + (ctx.callback)(&*info, is_exe) })); - match result { - Ok(stop) => i32::from(stop), - // Never unwind a Rust panic through libc's dl_iterate_phdr callback. - // Treat patching as best-effort and stop iteration on panic. - Err(_) => 1, - } + // Never unwind a Rust panic through libc's dl_iterate_phdr callback. + // Treat patching as best-effort and stop iteration on panic. + result.map(i32::from).unwrap_or(1) } unsafe { @@ -452,7 +448,7 @@ fn read_proc_maps() -> Vec { /// restored them, weakening RELRO for the process lifetime. This guard /// mprotects each unique page once (RW), lets the caller write as /// many entries as it needs, then mprotects each page back to what -/// /proc/self/maps reported at guard-construction time. +/// `/proc/self/maps` reported at guard-construction time. struct PageProtGuard { page_size: usize, maps: Vec, @@ -462,11 +458,10 @@ struct PageProtGuard { impl PageProtGuard { fn new() -> Self { - // sysconf can return -1 on error; casting that to usize gives - // a huge value that breaks the alignment/mprotect math. Fall - // back to a conservative 4 KiB default if the query fails. + // sysconf can return -1 on error; fall back to a conservative + // 4 KiB default if the query fails. let raw = unsafe { sysconf(_SC_PAGESIZE) }; - let page_size = if raw > 0 { raw as usize } else { 4096 }; + let page_size = usize::try_from(raw).unwrap_or(4096); Self { page_size, maps: read_proc_maps(), @@ -523,13 +518,11 @@ unsafe fn read_entry(addr: usize) -> usize { #[derive(Clone, Copy)] pub struct LookupResult { pub address: usize, - #[allow(dead_code)] // exposed for diagnostics / future filtering - pub size: u64, } /// Look up a symbol across loaded objects, returning the first /// non-zero-sized definition whose address is not `not_this_symbol`. -/// Mirrors ddprof's `lookup_symbol(..., accept_null_sized_symbol=false)`. +/// Null-sized symbols are ignored so hooks resolve to callable definitions. pub fn lookup_symbol(name: &str, not_this_symbol: usize) -> Option { let needle = name.as_bytes(); let mut found: Option = None; @@ -549,10 +542,7 @@ pub fn lookup_symbol(name: &str, not_this_symbol: usize) -> Option if sym.st_size > 0 { let addr = sym.st_value as usize + dyn_info.base_address; if addr != not_this_symbol { - found = Some(LookupResult { - address: addr, - size: sym.st_size, - }); + found = Some(LookupResult { address: addr }); return true; // stop } } @@ -576,11 +566,10 @@ struct LibraryRevertInfo { /// One registered override entry. struct OverrideInfo { /// Output slot the install path fills with the resolved real symbol - /// address (so hooks can call through it). Points to a `'static - /// AtomicUsize` supplied by the caller; the install-time write goes - /// through `store(Release)` to pair with the hook-side - /// `load(Acquire)`. - ref_slot: *const AtomicUsize, + /// address (so hooks can call through it). This is a shared static + /// atomic supplied by the caller; the install-time write goes through + /// `store(Release)` to pair with the hook-side `load(Acquire)`. + ref_slot: &'static AtomicUsize, /// Address of our hook function, written into matching GOT entries. new_symbol: usize, /// If a GOT entry's address equals this, leave it alone. Used to @@ -591,11 +580,7 @@ struct OverrideInfo { do_not_override_this_symbol: usize, } -unsafe impl Send for OverrideInfo {} -unsafe impl Sync for OverrideInfo {} - -/// Mirror of ddprof's `SymbolOverrides`. Holds the override table and -/// the per-library revert info needed to undo writes. +/// Holds the override table and the per-library revert info needed to undo writes. #[derive(Default)] pub struct SymbolOverrides { overrides: HashMap, @@ -614,9 +599,8 @@ impl SymbolOverrides { /// Register an override. `ref_slot` is filled in by `apply_overrides` /// with the resolved address of the real symbol so the hook can call - /// through it. Must point to a `'static AtomicUsize`; the install - /// path publishes via `store(Release)`. - pub fn register(&mut self, name: &str, hook: usize, ref_slot: *const AtomicUsize) { + /// through it. The install path publishes via `store(Release)`. + pub fn register(&mut self, name: &str, hook: usize, ref_slot: &'static AtomicUsize) { self.overrides.insert( name.to_string(), OverrideInfo { @@ -646,11 +630,8 @@ impl SymbolOverrides { .collect(); for (name, addr) in resolved { if let Some(ov) = self.overrides.get_mut(&name) { - // SAFETY: `ref_slot` was registered pointing at a - // `'static AtomicUsize`, so the pointer is valid for the - // lifetime of the process. Release pairs with the - // hook-side Acquire load. - unsafe { (*ov.ref_slot).store(addr, Ordering::Release) }; + // Release pairs with the hook-side Acquire load. + ov.ref_slot.store(addr, Ordering::Release); } } self.update_overrides(); @@ -660,7 +641,7 @@ impl SymbolOverrides { /// No-op if the loaded-library count hasn't changed. pub fn update_overrides(&mut self) { // `dl_phdr_info::dlpi_adds` is incremented on every dlopen. - // ddprof uses it as a cheap "did anything change?" probe. + // Use it as a cheap "did anything change?" probe. let mut nb_loaded: i32 = -1; iterate_libraries(|info, _| { nb_loaded = info.dlpi_adds as i32; @@ -854,8 +835,7 @@ impl SymbolOverrides { }; // `ref_slot==0` means we never resolved the real symbol, so the // hook would call a NULL pointer. Skip. - // SAFETY: `ref_slot` points at a `'static AtomicUsize`. - let real = unsafe { (*ov.ref_slot).load(Ordering::Acquire) }; + let real = ov.ref_slot.load(Ordering::Acquire); if real == 0 { return; } diff --git a/libdd-heap-gotter/src/hooks.rs b/libdd-heap-gotter/src/hooks.rs index 976f8a535c..5d114500c1 100644 --- a/libdd-heap-gotter/src/hooks.rs +++ b/libdd-heap-gotter/src/hooks.rs @@ -198,7 +198,7 @@ pub unsafe extern "C" fn gotter_dlopen(filename: *const c_char, flags: c_int) -> } /// Args we package up so the wrapped start routine sees its original -/// arg through our trampoline. Matches ddprof's `Args = tuple<...>`. +/// arg through our trampoline. struct PthreadCreateArgs { start: StartRoutine, arg: *mut c_void, diff --git a/libdd-heap-gotter/src/lib.rs b/libdd-heap-gotter/src/lib.rs index 49a05b41d7..233b05d86c 100644 --- a/libdd-heap-gotter/src/lib.rs +++ b/libdd-heap-gotter/src/lib.rs @@ -51,8 +51,8 @@ use std::sync::{Mutex, MutexGuard, TryLockError}; use elf::SymbolOverrides; /// Holds the SymbolOverrides registry across calls to `install` / `update` -/// / `restore`. ddprof keeps the equivalent state in -/// `g_symbol_overrides` guarded by `g_mutex` +/// / `restore`. Guarded globally because GOT patching mutates process-wide +/// state. #[cfg(all(target_os = "linux", target_pointer_width = "64"))] static GLOBAL_OVERRIDES: Mutex> = Mutex::new(None); @@ -153,13 +153,12 @@ fn register_all(so: &mut SymbolOverrides) { use hooks::*; use std::sync::atomic::AtomicUsize; - // Register one entry per supported symbol. The `ref_slot` raw - // pointer is to a `'static AtomicUsize`, so it's valid forever. - // 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. + // 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 as *const AtomicUsize); + so.register(name, hook_addr, slot); } reg( @@ -241,7 +240,6 @@ mod tests { assert!(r.is_some(), "expected to find malloc in loaded libraries"); let r = r.unwrap(); assert!(r.address != 0); - assert!(r.size > 0); } #[test] diff --git a/libdd-heap-gotter/tests/install.rs b/libdd-heap-gotter/tests/install.rs index 68ae16b8a8..9f7f49bd68 100644 --- a/libdd-heap-gotter/tests/install.rs +++ b/libdd-heap-gotter/tests/install.rs @@ -277,6 +277,17 @@ fn realloc_stress_across_alignments_preserves_data() { ^ 0x5A } + // Number of times to repeat the whole alignment matrix. Default 1 keeps + // the committed test fast; a soak run can crank it (the crash this guards + // against is heap-reuse-dependent and only shows up after tens of + // thousands of realloc/free operations) via + // REALLOC_STRESS_REPEAT=20000. Read before the sampled window so the + // env lookup's own allocation isn't sampled. + let repeat: usize = std::env::var("REALLOC_STRESS_REPEAT") + .ok() + .and_then(|v| v.parse().ok()) + .unwrap_or(1); + let installed = libdd_heap_gotter::install_heap_overrides(); assert!(installed); @@ -294,50 +305,52 @@ fn realloc_stress_across_alignments_preserves_data() { (*tl).remaining_bytes = 0; (*tl).remaining_bytes_initialized = true; - for (seed, &align) in ALIGNMENTS.iter().enumerate() { - let mut size = SIZE_WALK[0]; + for _rep in 0..repeat { + for (seed, &align) in ALIGNMENTS.iter().enumerate() { + let mut size = SIZE_WALK[0]; - let mut p = alloc_aligned(align, size) as *mut u8; - assert!(!p.is_null(), "alloc failed (align={align}, size={size})"); + let mut p = alloc_aligned(align, size) as *mut u8; + assert!(!p.is_null(), "alloc failed (align={align}, size={size})"); - // Record whether this block was actually sampled (peek is - // non-destructive, so it's safe before the realloc walk). - let mut raw: *mut c_void = std::ptr::null_mut(); - let mut offset: usize = 0; - if dd_sample_flag_peek(p as *mut c_void, &mut raw, &mut offset) { - saw_sampled = true; - } + // Record whether this block was actually sampled (peek is + // non-destructive, so it's safe before the realloc walk). + let mut raw: *mut c_void = std::ptr::null_mut(); + let mut offset: usize = 0; + if dd_sample_flag_peek(p as *mut c_void, &mut raw, &mut offset) { + saw_sampled = true; + } - for i in 0..size { - *p.add(i) = pattern_byte(seed, i); - } + for i in 0..size { + *p.add(i) = pattern_byte(seed, i); + } - for &new_size in &SIZE_WALK[1..] { - let p2 = libc::realloc(p as *mut c_void, new_size) as *mut u8; - assert!( - !p2.is_null(), - "realloc failed (align={align}, {size}->{new_size})" - ); - - // The overlapping prefix must survive the realloc. - let preserved = size.min(new_size); - for i in 0..preserved { - assert_eq!( - *p2.add(i), - pattern_byte(seed, i), - "corruption at byte {i} (align={align}, {size}->{new_size})" + for &new_size in &SIZE_WALK[1..] { + let p2 = libc::realloc(p as *mut c_void, new_size) as *mut u8; + assert!( + !p2.is_null(), + "realloc failed (align={align}, {size}->{new_size})" ); - } - // Repaint the whole new region for the next realloc to keep. - for i in 0..new_size { - *p2.add(i) = pattern_byte(seed, i); + + // The overlapping prefix must survive the realloc. + let preserved = size.min(new_size); + for i in 0..preserved { + assert_eq!( + *p2.add(i), + pattern_byte(seed, i), + "corruption at byte {i} (align={align}, {size}->{new_size})" + ); + } + // Repaint the whole new region for the next realloc to keep. + for i in 0..new_size { + *p2.add(i) = pattern_byte(seed, i); + } + + p = p2; + size = new_size; } - p = p2; - size = new_size; + libc::free(p as *mut c_void); } - - libc::free(p as *mut c_void); } // Disable sampling before restore so restore-time internal diff --git a/libdd-heap-sampler/src/allocation_realloc.c b/libdd-heap-sampler/src/allocation_realloc.c index a1c919265c..de11128634 100644 --- a/libdd-heap-sampler/src/allocation_realloc.c +++ b/libdd-heap-sampler/src/allocation_realloc.c @@ -45,13 +45,45 @@ dd_realloc_prep_t dd_allocation_realloc_prepare(void *old_user, size_t new_size) return out; /* passthrough: not sampled */ } + /* Clear the magic header NOW, while we still own the block. Once the + * real realloc runs it may free the old block internally; if we left + * the magic intact, a future allocation that reuses that memory could + * be falsely detected as sampled by dd_sample_flag_check, leading to + * a bogus raw pointer and a heap corruption on free. + * + * This is safe even if the real realloc subsequently fails (returns + * NULL): in that case old_user is still live but now unsampled. + * commit() returns NULL to the caller, so the application retains + * old_user; a later free(old_user) will take the unsampled fast path + * and pass old_user directly to the underlying free — which is + * incorrect (it should pass old_raw). To handle this, commit() + * re-stamps the header when realloc fails. */ +#if defined(__x86_64__) + { + void *old_header = (char *)old_user - DD_HEADER_BYTES; + const uint64_t zero = 0; + memcpy(old_header, &zero, sizeof(zero)); + } +#endif + /* Reserve room for the old header+slack ([0, old_offset)) plus * `new_size` bytes of user data at [old_offset, old_offset + new_size). * commit() shifts the user data down to [0, new_size). Overflow -> * fall back to passthrough with the caller-supplied size; the * underlying realloc will likely fail with a huge value, but - * nothing gets silently truncated or misinterpreted. */ + * nothing gets silently truncated or misinterpreted. + * + * We already cleared the header above, so we must re-stamp it + * before falling through to passthrough, otherwise a later free + * would not recover old_raw. */ if (new_size > SIZE_MAX - old_offset) { +#if defined(__x86_64__) + void *hdr = (char *)old_user - DD_HEADER_BYTES; + uint64_t m = DD_MAGIC; + uint64_t o = (uint64_t)old_offset; + memcpy(hdr, &m, sizeof(m)); + memcpy((char *)hdr + sizeof(m), &o, sizeof(o)); +#endif return out; } @@ -75,10 +107,20 @@ void *dd_allocation_realloc_commit(void *old_user, void *new_raw, dd_realloc_pre return new_raw; } - /* Underlying realloc failed: C says old_user is still live; its - * sampler flag was left intact by prepare(), so a later free() - * will still resolve the right raw pointer. */ - if (new_raw == NULL) return NULL; + /* Underlying realloc failed: C says old_user is still live. + * prepare() cleared the header optimistically, so re-stamp it + * now so that a later free(old_user) correctly resolves the raw + * pointer via dd_sample_flag_check. */ + if (new_raw == NULL) { +#if defined(__x86_64__) + void *old_header = (char *)old_user - DD_HEADER_BYTES; + uint64_t magic = DD_MAGIC; + uint64_t offset = (uint64_t)prep.old_offset; + memcpy(old_header, &magic, sizeof(magic)); + memcpy((char *)old_header + sizeof(magic), &offset, sizeof(offset)); +#endif + return NULL; + } /* Sampled path. libc realloc copied the old block's bytes into * new_raw starting at index 0, so the old user data now sits at @@ -96,8 +138,7 @@ void *dd_allocation_realloc_commit(void *old_user, void *new_raw, dd_realloc_pre /* Report the free of the OLD sampled allocation (the address the * profiler last saw as live). No matching alloc is fired: the new * block is unsampled. dd_probe_free just emits the ddheap:free - * USDT; there's no header left to clear because libc consumed the - * old block and its bytes have been overwritten by the memmove. */ + * USDT so the profiler can close the live-heap entry. */ dd_probe_free(old_user); return new_raw; }