feat(profiling): add heap-live allocation tracking to Profile

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Author: realFlowControl

datadog-profiling-replayer/src/main.rs

@@ -209,7 +209,7 @@ fn main() -> anyhow::Result<()> {
 
     let before = Instant::now();
     for (timestamp, sample) in samples {
-        outprof.try_add_sample(sample, timestamp)?;
+        outprof.try_add_sample(sample, timestamp, None)?;
     }
     let duration = before.elapsed();
 

examples/ffi/exporter.cpp

@@ -66,7 +66,7 @@ int main(int argc, char *argv[]) {
       .values = {&value, 1},
       .labels = {&label, 1},
   };
-  auto add_result = ddog_prof_Profile_add(profile.get(), sample, 0);
+  auto add_result = ddog_prof_Profile_add(profile.get(), sample, 0, 0);
   if (add_result.tag != DDOG_PROF_PROFILE_RESULT_OK) {
     print_error("Failed to add sample to profile: ", add_result.err);
     ddog_Error_drop(&add_result.err);

examples/ffi/exporter_manager.c

@@ -60,7 +60,7 @@ ddog_prof_Profile *create_profile_with_sample(void) {
     };
 
     // Pass 0 as the timestamp parameter
-    ddog_prof_Profile_Result add_result = ddog_prof_Profile_add(profile, sample, 0);
+    ddog_prof_Profile_Result add_result = ddog_prof_Profile_add(profile, sample, 0, 0);
     if (add_result.tag != DDOG_PROF_PROFILE_RESULT_OK) {
         print_error("Failed to add sample to profile", &add_result.err);
         ddog_Error_drop(&add_result.err);

examples/ffi/profiles.c

@@ -61,7 +61,7 @@ int main(void) {
   for (int i = 0; i < NUM_SAMPLES; i++) {
     label.num = i;
 
-    ddog_prof_Profile_Result add_result = ddog_prof_Profile_add(&profile, sample, 0);
+    ddog_prof_Profile_Result add_result = ddog_prof_Profile_add(&profile, sample, 0, 0);
     if (add_result.tag != DDOG_PROF_PROFILE_RESULT_OK) {
       ddog_CharSlice message = ddog_Error_message(&add_result.err);
       fprintf(stderr, "%.*s", (int)message.len, message.ptr);

libdd-profiling-ffi/src/profiles/datatypes.rs

@@ -583,15 +583,26 @@ impl From<ProfileNewResult> for Result<Profile, Error> {
 /// The `profile` ptr must point to a valid Profile object created by this
 /// module.
 /// This call is _NOT_ thread-safe.
+///
+/// # Arguments
+/// * `track_ptr` - If non-zero, also track this allocation for heap-live
+///   profiling. The sample data is copied into owned storage and will be
+///   automatically injected during profile reset. Use 0 to skip tracking.
 #[must_use]
 #[no_mangle]
 pub unsafe extern "C" fn ddog_prof_Profile_add(
     profile: *mut Profile,
     sample: Sample,
     timestamp: Option<NonZeroI64>,
+    track_ptr: usize,
 ) -> ProfileResult {
     (|| {
         let profile = profile_ptr_to_inner(profile)?;
+        let track = if track_ptr != 0 {
+            Some(track_ptr as u64)
+        } else {
+            None
+        };
         let uses_string_ids = sample
             .labels
             .first()
@@ -600,7 +611,7 @@ pub unsafe extern "C" fn ddog_prof_Profile_add(
         if uses_string_ids {
             profile.add_string_id_sample(sample.into(), timestamp)
         } else {
-            profile.try_add_sample(sample.try_into()?, timestamp)
+            profile.try_add_sample(sample.try_into()?, timestamp, track)
         }
     })()
     .context("ddog_prof_Profile_add failed")
@@ -930,6 +941,76 @@ pub unsafe extern "C" fn ddog_prof_Profile_reset(profile: *mut Profile) -> Profi
     .into()
 }
 
+/// Enable heap-live allocation tracking on this profile. When enabled,
+/// calls to `ddog_prof_Profile_add` with a non-zero `track_ptr` will copy
+/// the sample data into owned storage. Tracked allocations are automatically
+/// injected during profile reset and survive across resets.
+///
+/// # Arguments
+/// * `profile` - A mutable reference to the profile.
+/// * `max_tracked` - Maximum number of allocations to track simultaneously.
+/// * `excluded_labels` - Label keys to strip from tracked allocations
+///   (e.g., high-cardinality labels like "span id").
+/// * `excluded_labels_len` - Number of elements in `excluded_labels`.
+/// * `alloc_size_idx` - Index of alloc-size in the sample values array.
+/// * `heap_live_samples_idx` - Index of heap-live-samples in the sample values array.
+/// * `heap_live_size_idx` - Index of heap-live-size in the sample values array.
+///
+/// # Safety
+/// The `profile` ptr must point to a valid Profile object created by this
+/// module. `excluded_labels` must be valid for `excluded_labels_len` elements.
+/// This call is _NOT_ thread-safe.
+#[no_mangle]
+pub unsafe extern "C" fn ddog_prof_Profile_enable_heap_live_tracking(
+    profile: *mut Profile,
+    max_tracked: usize,
+    excluded_labels: *const CharSlice<'_>,
+    excluded_labels_len: usize,
+    alloc_size_idx: usize,
+    heap_live_samples_idx: usize,
+    heap_live_size_idx: usize,
+) -> ProfileResult {
+    (|| {
+        let profile = profile_ptr_to_inner(profile)?;
+        let labels: Vec<&str> = if excluded_labels.is_null() || excluded_labels_len == 0 {
+            Vec::new()
+        } else {
+            let slice = std::slice::from_raw_parts(excluded_labels, excluded_labels_len);
+            slice
+                .iter()
+                .map(|cs| cs.try_to_utf8())
+                .collect::<Result<Vec<_>, _>>()?
+        };
+        profile.enable_heap_live_tracking(
+            max_tracked,
+            &labels,
+            alloc_size_idx,
+            heap_live_samples_idx,
+            heap_live_size_idx,
+        );
+        anyhow::Ok(())
+    })()
+    .context("ddog_prof_Profile_enable_heap_live_tracking failed")
+    .into()
+}
+
+/// Remove a tracked heap-live allocation by pointer. No-op if heap-live
+/// tracking is disabled or the pointer is not tracked.
+///
+/// # Arguments
+/// * `profile` - A mutable reference to the profile.
+/// * `ptr` - The pointer value of the allocation to untrack.
+///
+/// # Safety
+/// The `profile` ptr must point to a valid Profile object created by this
+/// module. This call is _NOT_ thread-safe.
+#[no_mangle]
+pub unsafe extern "C" fn ddog_prof_Profile_untrack_allocation(profile: *mut Profile, ptr: usize) {
+    if let Ok(profile) = profile_ptr_to_inner(profile) {
+        profile.untrack_allocation(ptr as u64);
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -965,7 +1046,7 @@ mod tests {
                 labels: Slice::empty(),
             };
 
-            let result = Result::from(ddog_prof_Profile_add(&mut profile, sample, None));
+            let result = Result::from(ddog_prof_Profile_add(&mut profile, sample, None, 0));
             result.unwrap_err();
             ddog_prof_Profile_drop(&mut profile);
             Ok(())
@@ -1013,7 +1094,7 @@ mod tests {
                 labels: Slice::from(&labels),
             };
 
-            Result::from(ddog_prof_Profile_add(&mut profile, sample, None))?;
+            Result::from(ddog_prof_Profile_add(&mut profile, sample, None, 0))?;
             assert_eq!(
                 profile
                     .inner
@@ -1023,7 +1104,7 @@ mod tests {
                 1
             );
 
-            Result::from(ddog_prof_Profile_add(&mut profile, sample, None))?;
+            Result::from(ddog_prof_Profile_add(&mut profile, sample, None, 0))?;
             assert_eq!(
                 profile
                     .inner
@@ -1099,7 +1180,7 @@ mod tests {
             labels: Slice::from(labels.as_slice()),
         };
 
-        Result::from(ddog_prof_Profile_add(&mut profile, main_sample, None)).unwrap();
+        Result::from(ddog_prof_Profile_add(&mut profile, main_sample, None, 0)).unwrap();
         assert_eq!(
             profile
                 .inner
@@ -1109,7 +1190,7 @@ mod tests {
             1
         );
 
-        Result::from(ddog_prof_Profile_add(&mut profile, test_sample, None)).unwrap();
+        Result::from(ddog_prof_Profile_add(&mut profile, test_sample, None, 0)).unwrap();
         assert_eq!(
             profile
                 .inner

libdd-profiling/benches/add_samples.rs

@@ -129,7 +129,7 @@ pub fn bench_add_sample_vs_add2(c: &mut Criterion) {
                     values: &values,
                     labels: vec![],
                 };
-                black_box(profile.try_add_sample(sample, None)).unwrap();
+                black_box(profile.try_add_sample(sample, None, None)).unwrap();
             }
             black_box(profile.only_for_testing_num_aggregated_samples())
         })

libdd-profiling/examples/profiles.rs

@@ -49,7 +49,7 @@ fn main() {
     // Intentionally use the current time.
     let mut profile = Profile::try_new(&sample_types, Some(period)).unwrap();
 
-    match profile.try_add_sample(sample, None) {
+    match profile.try_add_sample(sample, None, None) {
         Ok(_) => {}
         Err(_) => exit(1),
     }

libdd-profiling/src/cxx.rs

@@ -410,7 +410,7 @@ impl Profile {
         };
 
         // Profile interns the strings
-        self.inner.try_add_sample(api_sample, None)?;
+        self.inner.try_add_sample(api_sample, None, None)?;
         Ok(())
     }
 

libdd-profiling/src/internal/heap_live.rs

@@ -0,0 +1,128 @@
+// Copyright 2021-Present Datadog, Inc. https://www.datadoghq.com/
+// SPDX-License-Identifier: Apache-2.0
+
+use crate::api;
+use crate::internal::owned_types::{OwnedFrame, OwnedLabel};
+use std::collections::HashMap;
+use std::hash::BuildHasherDefault;
+
+/// Tracks live heap allocations inside a Profile.
+/// Stores owned copies of sample data (frames, labels, values) so tracked
+/// allocations survive across profile resets. Injected into the Profile's
+/// observations automatically before serialization via
+/// `reset_and_return_previous()`.
+type FxBuildHasher = BuildHasherDefault<rustc_hash::FxHasher>;
+
+pub(crate) struct HeapLiveState {
+    pub tracked: HashMap<u64, TrackedAlloc, FxBuildHasher>,
+    pub max_tracked: usize,
+    pub excluded_labels: Vec<Box<str>>,
+    /// Index of alloc-size in the sample values array (to read allocation size).
+    alloc_size_idx: usize,
+    /// Index of heap-live-samples in the sample values array (to write 1).
+    heap_live_samples_idx: usize,
+    /// Index of heap-live-size in the sample values array (to write the size).
+    heap_live_size_idx: usize,
+    /// Total number of values per sample.
+    num_values: usize,
+}
+
+/// A single tracked live allocation with owned frame/label/values data.
+pub(crate) struct TrackedAlloc {
+    pub frames: Vec<OwnedFrame>,
+    pub labels: Vec<OwnedLabel>,
+    pub values: Vec<i64>,
+}
+
+impl HeapLiveState {
+    pub fn new(
+        max_tracked: usize,
+        excluded_labels: &[&str],
+        alloc_size_idx: usize,
+        heap_live_samples_idx: usize,
+        heap_live_size_idx: usize,
+        num_values: usize,
+    ) -> Self {
+        Self {
+            tracked: HashMap::with_capacity_and_hasher(max_tracked, FxBuildHasher::default()),
+            max_tracked,
+            excluded_labels: excluded_labels.iter().map(|s| Box::from(*s)).collect(),
+            alloc_size_idx,
+            heap_live_samples_idx,
+            heap_live_size_idx,
+            num_values,
+        }
+    }
+
+    /// Track a new allocation. Copies borrowed strings from the sample into
+    /// owned storage. Constructs heap-live-only values: all zeros except
+    /// heap-live-samples=1 and heap-live-size=alloc_size.
+    /// Returns false if tracker is full.
+    pub fn track(&mut self, ptr: u64, sample: &api::Sample) -> bool {
+        if self.tracked.len() >= self.max_tracked {
+            return false;
+        }
+        let alloc = TrackedAlloc::from_api_sample(
+            sample,
+            &self.excluded_labels,
+            self.alloc_size_idx,
+            self.heap_live_samples_idx,
+            self.heap_live_size_idx,
+            self.num_values,
+        );
+        self.tracked.insert(ptr, alloc);
+        true
+    }
+
+    /// Remove a tracked allocation. No-op if ptr is not tracked.
+    pub fn untrack(&mut self, ptr: u64) {
+        self.tracked.remove(&ptr);
+    }
+}
+
+impl TrackedAlloc {
+    fn from_api_sample(
+        sample: &api::Sample,
+        excluded_labels: &[Box<str>],
+        alloc_size_idx: usize,
+        heap_live_samples_idx: usize,
+        heap_live_size_idx: usize,
+        num_values: usize,
+    ) -> Self {
+        let frames = sample
+            .locations
+            .iter()
+            .map(|loc| OwnedFrame {
+                function_name: loc.function.name.into(),
+                filename: loc.function.filename.into(),
+                line: loc.line,
+            })
+            .collect();
+
+        let labels = sample
+            .labels
+            .iter()
+            .filter(|l| !excluded_labels.iter().any(|ex| ex.as_ref() == l.key))
+            .map(|l| OwnedLabel {
+                key: l.key.into(),
+                str_value: l.str.into(),
+                num: l.num,
+                num_unit: l.num_unit.into(),
+            })
+            .collect();
+
+        // Construct heap-live-only values: all zeros except the heap-live
+        // fields. The allocation size is read from the original sample's
+        // alloc-size slot.
+        let alloc_size = sample.values.get(alloc_size_idx).copied().unwrap_or(0);
+        let mut values = vec![0i64; num_values];
+        values[heap_live_samples_idx] = 1;
+        values[heap_live_size_idx] = alloc_size;
+
+        TrackedAlloc {
+            frames,
+            labels,
+            values,
+        }
+    }
+}

libdd-profiling/src/internal/mod.rs

@@ -4,6 +4,7 @@
 mod endpoint_stats;
 mod endpoints;
 mod function;
+pub(crate) mod heap_live;
 mod label;
 mod location;
 mod mapping;