fix up local ffi; type-annotate

Author: vlad-scherbich

setup.py

@@ -1107,10 +1107,10 @@ def build_heap_gotter(self) -> None:
             subprocess.run(["install_name_tool", "-id", gotter_name, gotter_library], check=True)
 
         if self._should_strip_heap_gotter():
-            debug_sidecar = self._extract_and_strip_staged_debug_symbols(gotter_library)
+            debug_sidecar: t.Optional[Path] = self._extract_and_strip_staged_debug_symbols(gotter_library)
             if debug_sidecar:
-                unstripped_size = built.stat().st_size
-                stripped_size = gotter_library.stat().st_size
+                unstripped_size: int = built.stat().st_size
+                stripped_size: int = gotter_library.stat().st_size
                 print(
                     f"Stripped heap-gotter cdylib for wheel packaging: "
                     f"{unstripped_size} -> {stripped_size} bytes "
@@ -1225,15 +1225,15 @@ def _extract_and_strip_staged_debug_symbols(so_file: Path) -> t.Optional[Path]:
         are stripped post-build; heap-gotter is stripped at staging time because it
         is built out-of-band from setuptools extensions).
         """
-        objcopy = shutil.which("objcopy")
-        strip_bin = shutil.which("strip")
+        objcopy: t.Optional[str] = shutil.which("objcopy")
+        strip_bin: t.Optional[str] = shutil.which("strip")
         if not objcopy or not strip_bin:
             print("WARNING: objcopy/strip not found, skipping heap-gotter symbol stripping", flush=True)
             return None
 
-        so_path = str(so_file)
+        so_path: str = str(so_file)
         subprocess.run([objcopy, "--remove-section", ".llvmbc", so_path], check=False)
-        debug_out = f"{so_path}.debug"
+        debug_out: str = f"{so_path}.debug"
         try:
             subprocess.run([objcopy, "--only-keep-debug", so_path, debug_out], check=True)
             if not Path(debug_out).is_file() or Path(debug_out).stat().st_size == 0:

src/native_heap_gotter/Cargo.toml

@@ -37,7 +37,6 @@ codegen-units = 1
 test-support = ["libdd-profiling-heap-gotter/test-support"]
 
 [dependencies]
-# Pinned to the published crates.io release (Phase 1.5 migration off the
-# libdatadog `main` git pin). `default-features = false` keeps the
-# allocation-only surface: `live-heap` stays opt-in and off for Phase 1.
-libdd-profiling-heap-gotter = { version = "1.0.0", default-features = false }
+# 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"] }

src/native_heap_gotter/lib.rs

@@ -1,24 +1,8 @@
 // Copyright 2025-Present Datadog, Inc. https://www.datadoghq.com/
 // SPDX-License-Identifier: Apache-2.0
 
-//! Thin cdylib that wraps libdatadog's published `libdd-profiling-heap-gotter`
-//! crate under stable, ddtrace-owned C symbols. The Python ctypes activator
-//! (`ddtrace/internal/datadog/profiling/heap_gotter`) dlopen's the resulting
-//! `libdd_heap_gotter.<ext-suffix>.so` and drives GOT-based native heap
-//! profiling from it.
-//!
-//! The upstream crate exposes a pure-Rust API (`install_heap_overrides`,
-//! `heap_overrides_are_installed`); it is not a C-ABI surface. Here we re-export
-//! those calls as fixed, unmangled `extern "C"` entry points returning a plain
-//! `bool`, so the Python ctypes side links against stable symbol names and gets
-//! a trivial success signal.
-//!
-//! Installation is permanent and process-global: the GOT entries patched by
-//! `install` point at functions inside the linked-in gotter code, so this
-//! library must stay loaded for the life of the process. The Python activator
-//! loads it into the global namespace and never unloads it. After `fork()` the
-//! child inherits both the loaded library and the patched GOT, so a re-install
-//! in the child is a harmless no-op.
+//! Thin cdylib wrapping `libdd-profiling-heap-gotter` (crates.io) under stable
+//! `extern "C"` symbols for the Python ctypes activator to dlopen.
 
 /// Install GOT overrides for supported heap-allocation symbols and report
 /// whether the install actually took effect.
@@ -50,6 +34,29 @@ pub extern "C" fn ddtrace_heap_gotter_is_installed() -> bool {
     libdd_profiling_heap_gotter::heap_overrides_are_installed()
 }
 
+/// Set the mean sample distance (bytes between samples) for the heap sampler.
+/// Must be called before `ddtrace_heap_gotter_install` to take effect.
+///
+/// # Safety
+///
+/// C ABI entry point; always safe to call.
+#[no_mangle]
+pub extern "C" fn ddtrace_heap_gotter_set_sampling_distance(distance: u64) {
+    libdd_profiling_heap_gotter::set_default_sampling_distance(distance);
+}
+
+/// Re-scan loaded libraries and patch any newly-introduced GOT entries.
+/// Normally called automatically from the internal `dlopen` hook; exposed here
+/// for cases where the Python side loads a `.so` and wants immediate coverage.
+///
+/// # Safety
+///
+/// C ABI entry point with no arguments and no pointers; always safe to call.
+#[no_mangle]
+pub extern "C" fn ddtrace_heap_gotter_update() {
+    libdd_profiling_heap_gotter::update_heap_overrides();
+}
+
 /// 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;