misc. code cleanup for Python native-heap profiling

Author: vlad-scherbich

.gitlab/scripts/build-wheel-helpers.sh

@@ -151,22 +151,49 @@ PY
   # Repair wheel (ONLY PLATFORM-SPECIFIC CODE)
   section_start "repair_wheel" "Repairing wheel"
   if [[ "$(uname -s)" == "Linux" ]]; then
-    # The opt-in heap-gotter cdylib (DD_PROFILING_NATIVE_HEAP_BUILD=1) has
-    # non-standard ELF versioning sections that trip auditwheel's iter_versions
-    # parser. --exclude does not help: it only drops a SONAME from dependency
-    # grafting, while repair still parses every ELF listed in the wheel's RECORD.
-    # So the cdylib has to leave the wheel entirely and be reinserted after.
+    # Heap-gotter's ELF versioning trips auditwheel iter_versions; --exclude
+    # only skips SONAME grafting, so stash .so (+ .so.debug) out of the wheel,
+    # repair, then reinsert the runtime .so. Python zipfile: Info-ZIP globs are
+    # unreliable on archive paths.
     GOTTER_STASH_DIR="${WORK_DIR}/heap_gotter_stash"
-    GOTTER_PATTERN='*libdd_heap_gotter*.so'
-    if unzip -l "${BUILT_WHEEL_FILE}" | grep -q 'libdd_heap_gotter.*\.so$'; then
-      mkdir -p "${GOTTER_STASH_DIR}"
-      unzip -q "${BUILT_WHEEL_FILE}" "${GOTTER_PATTERN}" -d "${GOTTER_STASH_DIR}"
-      uv run --no-project scripts/zip_filter.py "${BUILT_WHEEL_FILE}" "${GOTTER_PATTERN}"
-    fi
+    mkdir -p "${GOTTER_STASH_DIR}"
+    BUILT_WHEEL_FILE="${BUILT_WHEEL_FILE}" GOTTER_STASH_DIR="${GOTTER_STASH_DIR}" \
+      uv run --no-project python - <<'PY'
+import os
+import zipfile
+from pathlib import Path
+
+import subprocess
+import sys
+
+wheel = Path(os.environ["BUILT_WHEEL_FILE"])
+stash = Path(os.environ["GOTTER_STASH_DIR"])
+marker = "libdd_heap_gotter"
+with zipfile.ZipFile(wheel, "r") as zf:
+    gotter = [
+        n
+        for n in zf.namelist()
+        if marker in Path(n).name and (n.endswith(".so") or n.endswith(".so.debug"))
+    ]
+    for name in gotter:
+        dest = stash / name
+        dest.parent.mkdir(parents=True, exist_ok=True)
+        dest.write_bytes(zf.read(name))
+        print(f"Stashed heap-gotter artifact: {name}")
+if gotter:
+    # zip_filter keeps RECORD consistent.
+    patterns = [f"*{marker}*.so", f"*{marker}*.so.debug", f"*/{marker}*"]
+    subprocess.check_call([sys.executable, "scripts/zip_filter.py", str(wheel), *patterns])
+with zipfile.ZipFile(wheel, "r") as zf:
+    leftover = [n for n in zf.namelist() if marker in Path(n).name]
+if leftover:
+    raise SystemExit(f"heap-gotter still in wheel before auditwheel: {leftover}")
+print(f"heap-gotter stash count before auditwheel: {len(gotter)}")
+PY
 
     auditwheel repair -w "${TMP_WHEEL_DIR}" "${BUILT_WHEEL_FILE}"
 
-    if [[ -d "${GOTTER_STASH_DIR}" ]]; then
+    if find "${GOTTER_STASH_DIR}" \( -name 'libdd_heap_gotter*.so' -o -name 'libdd_heap_gotter*.so.debug' \) 2>/dev/null | grep -q .; then
       REPAIRED_WHEEL_FILE=$(ls "${TMP_WHEEL_DIR}"/*.whl | head -n 1)
       GOTTER_STASH_DIR="${GOTTER_STASH_DIR}" REPAIRED_WHEEL_FILE="${REPAIRED_WHEEL_FILE}" \
         uv run --no-project python - <<'PY'
@@ -181,7 +208,12 @@ from pathlib import Path
 wheel = Path(os.environ["REPAIRED_WHEEL_FILE"])
 stash = Path(os.environ["GOTTER_STASH_DIR"])
 
-additions = {str(p.relative_to(stash)): p for p in sorted(stash.rglob("*")) if p.is_file()}
+# Runtime .so only; .so.debug stays in debugwheelhouse.
+additions = {
+    str(p.relative_to(stash)): p
+    for p in sorted(stash.rglob("*"))
+    if p.is_file() and p.name.endswith(".so")
+}
 if not additions:
     print("No stashed heap-gotter cdylib to reinsert")
     raise SystemExit(0)

ddtrace/internal/datadog/profiling/heap_gotter/__init__.py

@@ -1,31 +1,28 @@
 """Activator for native (C/C++) heap allocation profiling via GOT rewriting.
 
-This module dlopen's the `libdd_heap_gotter` cdylib (built out-of-band from
-libdatadog's `libdd-profiling-heap-gotter-ffi`; see `src/native_heap_gotter`)
-and drives it through a tiny, stable C ABI:
-
-    bool ddtrace_heap_gotter_install(void);      # install + report success
-    bool ddtrace_heap_gotter_is_installed(void); # current install state
-
-Calling `install()` patches the process's GOT entries for heap allocation
-symbols so that Datadog's `ddheap:alloc` USDT probe sites fire on sampled allocations.
-The Full Host eBPF profiler then attaches uprobes to those sites to collect native allocation stacks.
-
-If the cdylib is missing or anything goes wrong loading it, `is_available` is `False`
-and `install()` is a no-op. Loading this module must never raise.
-
-Installation cannot be undone (the patched GOT entries point at functions inside the cdylib),
-so the library must stay mapped for the life of the process. We keep the `ctypes.CDLL` handle
-at module scope and never unload it.
-
-After a successful `install()`, a child of `fork()` inherits the mapping and the patched GOT.
-Re-entering `install()` in that child is therefore unnecessary; we skip the native call when
-this module already recorded a successful arm (Python module state is also inherited). That
-avoids re-locking upstream's process-global registry mutex. Upstream does not yet implement a
-`pthread_atfork` child reset, so forking *during* an in-flight `install()`/`update()` can still
-leave that mutex locked in the child — prefer arming on the main thread, or after fork in the
-worker (gunicorn/uWSGI-style), and treat mid-install fork as unsafe until libdatadog lands
-atfork handling.
+Dlopens ``libdd_heap_gotter`` (see ``src/native_heap_gotter``) and drives:
+
+    bool ddtrace_heap_gotter_install(void);
+    bool ddtrace_heap_gotter_is_installed(void);
+    bool ddtrace_heap_gotter_live_heap_enabled(void);  # True if built with ddheap:free
+
+``install()`` patches GOT entries so ``ddheap:alloc`` (and, on live-heap builds,
+``ddheap:free``) USDT sites fire; the Full Host eBPF profiler attaches uprobes.
+Nothing is collected or uploaded from Python.
+
+``live_heap_enabled()`` is a compile-time property of the loaded artifact
+(default-on ``live-heap`` feature): True when the cdylib stamps retain flags and
+emits ``ddheap:free``. False if the cdylib is missing, alloc-only
+(``--no-default-features``), or predates this symbol (bound defensively).
+
+Missing/broken load → ``is_available`` False, ``install()`` no-op; import never
+raises. Install is permanent (GOT points into the cdylib), so the CDLL handle
+stays mapped for the process lifetime.
+
+After a successful ``install()``, fork children inherit the mapping and patched
+GOT; ``_armed`` skips re-entering the native installer. Upstream has no
+``pthread_atfork`` reset for its registry mutex — prefer arming on the main
+thread or in the worker after fork; mid-install fork is unsafe.
 """
 
 from __future__ import annotations
@@ -35,27 +32,26 @@
 import sysconfig
 
 
-# Mirror the ddup/stack modules: importers (notably settings/profiling.py) read
-# these two attributes to decide whether the feature can run.
+# Mirror ddup/stack: settings/profiling.py reads these to gate the feature.
 is_available: bool = False
 failure_msg: str = ""
 
-_lib: ctypes.CDLL | None = None  # kept alive for process lifetime; never dlclose'd
-# Set when install() has succeeded in this process (inherited across fork).
-_armed: bool = False
+# Compile-time live-heap capability of the loaded artifact (see module docstring).
+_live_heap_available: bool = False
+
+_lib: ctypes.CDLL | None = None  # process lifetime; never dlclose'd
+_armed: bool = False  # successful install(); inherited across fork
 
 
 def _library_path() -> str:
-    # The cdylib is staged next to libdd_wrapper in the profiling package and
-    # carries the interpreter EXT_SUFFIX, matching setup.py's naming.
+    # Staged next to libdd_wrapper with the interpreter EXT_SUFFIX (setup.py).
     suffix: str = sysconfig.get_config_var("EXT_SUFFIX") or ".so"
     profiling_dir: str = os.path.dirname(os.path.dirname(__file__))
     return os.path.join(profiling_dir, "libdd_heap_gotter" + suffix)
 
 
 try:
-    # Native heap profiling via the gotter is Linux-only; on every other
-    # platform the underlying library is a no-op, so don't even try to load.
+    # Linux-only; elsewhere the gotter is a no-op.
     sysname = os.uname().sysname if os.name == "posix" else os.name
     if sysname != "Linux":
         raise OSError(f"Native heap gotter is only supported on Linux. Running on {sysname}")
@@ -64,8 +60,7 @@ def _library_path() -> str:
     if not os.path.exists(_path):
         raise FileNotFoundError(_path)
 
-    # RTLD_GLOBAL so the loaded code is unambiguously resolvable; RTLD_NOW so any
-    # unresolved symbol fails here (fail-closed) rather than at first call.
+    # RTLD_GLOBAL for resolvability; RTLD_NOW fail-closed on unresolved symbols.
     _lib = ctypes.CDLL(_path, mode=ctypes.RTLD_GLOBAL | getattr(os, "RTLD_NOW", 0))
 
     _lib.ddtrace_heap_gotter_install.argtypes = []
@@ -75,18 +70,25 @@ def _library_path() -> str:
 
     is_available = True
 
+    # Optional symbol (pre-Phase-2 cdylibs); failure must not disable install().
+    try:
+        _lib.ddtrace_heap_gotter_live_heap_enabled.argtypes = []
+        _lib.ddtrace_heap_gotter_live_heap_enabled.restype = ctypes.c_bool
+        _live_heap_available = bool(_lib.ddtrace_heap_gotter_live_heap_enabled())
+    except Exception:
+        _live_heap_available = False
+
 except Exception as e:
     failure_msg = str(e)
     _lib = None
 
 
 def install() -> bool:
-    """Install the native heap GOT overrides. Returns True if now installed; False otherwise.
+    """Install native heap GOT overrides. True if installed; False otherwise.
 
-    Idempotent at the Python layer: once a call has succeeded, further calls
-    (including in a forked child that inherited ``_armed``) return True without
-    re-entering the native installer. No-op that returns False when the cdylib
-    is unavailable. See the module docstring for fork-safety limits.
+    Idempotent at the Python layer: after success (including in a fork child that
+    inherited ``_armed``), further calls return True without re-entering the native
+    installer. See module docstring for fork-safety limits.
     """
     global _armed
     if not is_available or _lib is None:
@@ -112,3 +114,12 @@ def is_installed() -> bool:
         return bool(_lib.ddtrace_heap_gotter_is_installed())
     except Exception:
         return False
+
+
+def live_heap_enabled() -> bool:
+    """Return whether the loaded cdylib was built with live-heap tracking.
+
+    Compile-time property of the artifact (default-on feature). False when the
+    cdylib is missing, alloc-only, or predates this symbol.
+    """
+    return _live_heap_available

ddtrace/profiling/profiler.py

@@ -386,7 +386,8 @@ def _start_service(self) -> None:
 
             try:
                 if heap_gotter.install():
-                    LOG.info("Native heap profiling armed (GOT overrides installed)")
+                    mode: str = "live-heap" if heap_gotter.live_heap_enabled() else "allocation-only"
+                    LOG.info("Native heap profiling armed (GOT overrides installed, %s)", mode)
                 else:
                     LOG.warning("Native heap profiling requested but GOT overrides were not installed")
             except Exception:

releasenotes/notes/native-heap-profiling-live-heap-336bc489c93989da.yaml

@@ -0,0 +1,6 @@
+---
+features:
+  - |
+    profiling: Enable collection of samples for native (C/C++) allocations 
+    in Python processes. Enable with ``DD_PROFILING_NATIVE_HEAP_ENABLED=true``.
+    This setting works for Linux only.

setup.py

@@ -135,7 +135,14 @@
 # Opt-in build of the native heap-gotter cdylib.
 # Off by default so normal builds don't pay the extra cargo fetch/compile and
 # mainline wheels don't ship the artifact until it GA's.
-BUILD_NATIVE_HEAP_GOTTER: bool = os.getenv("DD_PROFILING_NATIVE_HEAP_BUILD", "0").lower() in ("1", "yes", "on", "true")
+# Same env var as runtime arming (ProfilingConfigNativeHeap.enabled); setup.py
+# reads it via os.getenv during the package build, independent of DDConfig.
+BUILD_NATIVE_HEAP_GOTTER: bool = os.getenv("DD_PROFILING_NATIVE_HEAP_ENABLED", "0").lower() in (
+    "1",
+    "yes",
+    "on",
+    "true",
+)
 # Keep the staged cdylib unstripped when building with the upstream test-support
 # feature (hook-hit counter for e2e / integration tests).
 BUILD_NATIVE_HEAP_GOTTER_TEST_SUPPORT: bool = os.getenv("DD_PROFILING_NATIVE_HEAP_TEST_SUPPORT", "0").lower() in (

src/native_heap_gotter/Cargo.toml

@@ -31,12 +31,17 @@ debug = "line-tables-only"
 codegen-units = 1
 
 [features]
+default = ["live-heap"]
 # Forwards to the upstream crate's test-support surface so integration tests can
 # assert the patched GOT hooks actually ran (hit counter) without needing a
 # live eBPF uprobe attach. Never enabled for shipped wheels.
 test-support = ["libdd-profiling-heap-gotter/test-support"]
 
+# Default-on live-heap (ddheap:free + retain flags). Sole enablement path for
+# upstream live-heap — keeps cfg! / live_heap_enabled() in lockstep with the
+# artifact. Omit via --no-default-features for alloc-only builds.
+live-heap = ["libdd-profiling-heap-gotter/live-heap"]
+
 [dependencies]
-# Pinned to the published crates.io release. `live-heap` enables balanced
-# alloc/free sampling so the profiler can report retained (live) heap.
-libdd-profiling-heap-gotter = { version = "1.0.0", features = ["live-heap"] }
+# Live-heap only via the feature above, not unconditionally.
+libdd-profiling-heap-gotter = { version = "1.0.0" }

src/native_heap_gotter/lib.rs

@@ -62,6 +62,18 @@ pub extern "C" fn ddtrace_heap_gotter_update() {
     libdd_profiling_heap_gotter::update_heap_overrides();
 }
 
+/// Compile-time live-heap capability: `true` when built with the `live-heap`
+/// feature (`ddheap:free` + retain flagging). Not a runtime toggle; keep in
+/// lockstep with Cargo.toml's forward to `libdd-profiling-heap-gotter/live-heap`.
+///
+/// # Safety
+///
+/// C ABI entry point with no arguments and no pointers; always safe to call.
+#[no_mangle]
+pub extern "C" fn ddtrace_heap_gotter_live_heap_enabled() -> bool {
+    cfg!(feature = "live-heap")
+}
+
 /// Test-only: number of times a patched hook has run in this process. Lets
 /// integration tests prove the patched GOT was actually exercised without a
 /// live eBPF attach. Only present when built with the `test-support` feature;

tests/profiling/test_native_heap_gotter.py

@@ -2,7 +2,7 @@
 
 The activator (``ddtrace.internal.datadog.profiling.heap_gotter``) is fail-closed
 and must behave correctly whether or not the opt-in gotter cdylib was built into
-the wheel (``DD_PROFILING_NATIVE_HEAP_BUILD=1``):
+the wheel (``DD_PROFILING_NATIVE_HEAP_ENABLED=1`` at build time):
 
 * If the library is absent (the default), ``install()``/``is_installed()`` are
   no-ops returning ``False``.
@@ -27,17 +27,17 @@ def test_native_heap_gotter_smoke() -> None:
     from ddtrace.internal.datadog.profiling import heap_gotter
 
     if not heap_gotter.is_available:
-        # Wheel built without the gotter cdylib: strictly a no-op.
         assert heap_gotter.install() is False
         assert heap_gotter.is_installed() is False
+        assert heap_gotter.live_heap_enabled() is False
     else:
-        # Native-heap build: arming must take effect and be idempotent.
         assert heap_gotter.is_installed() is False
         assert heap_gotter.install() is True
         assert heap_gotter.is_installed() is True
-        assert heap_gotter.install() is True
+        assert heap_gotter.install() is True  # idempotent
+        # Default gotter builds enable the live-heap Cargo feature (ddheap:free).
+        assert heap_gotter.live_heap_enabled() is True
 
-        # Generate allocation pressure; this must not crash with the patched GOT.
         blobs: list[tuple[str, int]] = []
         for i in range(200):
             blobs.append(("x" * 4096, i))