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,24 +583,35 @@ 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()
             .is_some_and(|label| label.key.is_empty() && label.key_id.value > 0);
 
         if uses_string_ids {
-            profile.add_string_id_sample(sample.into(), timestamp)
+            profile.add_string_id_sample(sample.into(), timestamp, track)
         } 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(())
     }