From cbf7fd0c874be25bdf88c9b295a3ede04504d841 Mon Sep 17 00:00:00 2001 From: Florian Engelhardt Date: Mon, 9 Feb 2026 13:36:19 +0100 Subject: [PATCH 1/2] feat(profiling): add heap-live allocation tracking to Profile Co-Authored-By: Claude Opus 4.6 --- datadog-profiling-replayer/src/main.rs | 2 +- examples/ffi/exporter.cpp | 2 +- examples/ffi/exporter_manager.c | 2 +- examples/ffi/profiles.c | 2 +- libdd-profiling-ffi/src/profiles/datatypes.rs | 84 ++- libdd-profiling/benches/add_samples.rs | 2 +- libdd-profiling/examples/profiles.rs | 2 +- libdd-profiling/src/cxx.rs | 2 +- libdd-profiling/src/internal/heap_live.rs | 271 ++++++++ libdd-profiling/src/internal/mod.rs | 1 + .../src/internal/profile/fuzz_tests.rs | 2 +- libdd-profiling/src/internal/profile/mod.rs | 628 +++++++++++++++++- 12 files changed, 949 insertions(+), 51 deletions(-) create mode 100644 libdd-profiling/src/internal/heap_live.rs diff --git a/datadog-profiling-replayer/src/main.rs b/datadog-profiling-replayer/src/main.rs index 6a8ae2b32f..0ee889fec1 100644 --- a/datadog-profiling-replayer/src/main.rs +++ b/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(); diff --git a/examples/ffi/exporter.cpp b/examples/ffi/exporter.cpp index 5b02debf7d..65e8b4c750 100644 --- a/examples/ffi/exporter.cpp +++ b/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); diff --git a/examples/ffi/exporter_manager.c b/examples/ffi/exporter_manager.c index 765c11d2de..f75ad3a3da 100644 --- a/examples/ffi/exporter_manager.c +++ b/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); diff --git a/examples/ffi/profiles.c b/examples/ffi/profiles.c index a4b0eb2800..054d984d27 100644 --- a/examples/ffi/profiles.c +++ b/examples/ffi/profiles.c @@ -62,7 +62,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); diff --git a/libdd-profiling-ffi/src/profiles/datatypes.rs b/libdd-profiling-ffi/src/profiles/datatypes.rs index c7ec589d5c..c2a93f3120 100644 --- a/libdd-profiling-ffi/src/profiles/datatypes.rs +++ b/libdd-profiling-ffi/src/profiles/datatypes.rs @@ -549,24 +549,31 @@ impl From for Result { /// 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, + track_ptr: usize, ) -> ProfileResult { (|| { let profile = profile_ptr_to_inner(profile)?; + let track = (track_ptr != 0).then_some(track_ptr as u64); 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") @@ -896,6 +903,69 @@ 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"). +/// * `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 a valid `Slice`. +/// 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: Slice, + 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> = excluded_labels + .iter() + .map(|cs| cs.try_to_utf8()) + .collect::, _>>()?; + 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::*; @@ -931,7 +1001,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(()) @@ -979,7 +1049,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 @@ -989,7 +1059,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 @@ -1065,7 +1135,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 @@ -1075,7 +1145,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 diff --git a/libdd-profiling/benches/add_samples.rs b/libdd-profiling/benches/add_samples.rs index fc27d886b3..5ca0b07a57 100644 --- a/libdd-profiling/benches/add_samples.rs +++ b/libdd-profiling/benches/add_samples.rs @@ -138,7 +138,7 @@ pub fn bench_add_sample_vs_add2(c: &mut Criterion) { values: &values, labels: labels_api.clone(), }; - 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()) }) diff --git a/libdd-profiling/examples/profiles.rs b/libdd-profiling/examples/profiles.rs index 2d4fed25d4..e72f2642c7 100644 --- a/libdd-profiling/examples/profiles.rs +++ b/libdd-profiling/examples/profiles.rs @@ -48,7 +48,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), } diff --git a/libdd-profiling/src/cxx.rs b/libdd-profiling/src/cxx.rs index 312b1c636d..d281d5d256 100644 --- a/libdd-profiling/src/cxx.rs +++ b/libdd-profiling/src/cxx.rs @@ -549,7 +549,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(()) } diff --git a/libdd-profiling/src/internal/heap_live.rs b/libdd-profiling/src/internal/heap_live.rs new file mode 100644 index 0000000000..884d1a85a4 --- /dev/null +++ b/libdd-profiling/src/internal/heap_live.rs @@ -0,0 +1,271 @@ +// Copyright 2021-Present Datadog, Inc. https://www.datadoghq.com/ +// SPDX-License-Identifier: Apache-2.0 + +use crate::api; +use crate::api::ManagedStringId; +use crate::collections::string_storage::ManagedStringStorage; +use std::collections::HashMap; +use std::hash::BuildHasherDefault; + +/// Owned frame data for heap-live tracking. +/// Stores copies of borrowed strings so tracked allocations survive across +/// profile resets. +pub(crate) struct OwnedFrame { + pub function_name: Box, + pub filename: Box, + pub line: i64, +} + +/// Owned label for heap-live tracking. +pub(crate) struct OwnedLabel { + pub key: Box, + pub str_value: Box, + pub num: i64, + pub num_unit: Box, +} + +impl OwnedFrame { + pub fn as_api_location(&self) -> api::Location<'_> { + api::Location { + function: api::Function { + name: &self.function_name, + system_name: "", + filename: &self.filename, + }, + line: self.line, + ..api::Location::default() + } + } +} + +impl OwnedLabel { + pub fn as_api_label(&self) -> api::Label<'_> { + api::Label { + key: &self.key, + str: &self.str_value, + num: self.num, + num_unit: &self.num_unit, + } + } +} + +type FxBuildHasher = BuildHasherDefault; + +/// Indices into the sample values array for reading/writing heap-live fields. +struct ValueIndices { + /// Index of alloc-size in the sample values array (to read allocation size). + alloc_size: usize, + /// Index of heap-live-samples in the sample values array (to write 1). + heap_live_samples: usize, + /// Index of heap-live-size in the sample values array (to write the size). + heap_live_size: usize, + /// Total number of values per sample. + num_values: usize, +} + +impl ValueIndices { + /// Build a heap-live values vector: all zeros except heap-live-samples=1 + /// and heap-live-size=alloc_size (read from the original sample). + fn build_values(&self, sample_values: &[i64]) -> Vec { + let alloc_size = sample_values.get(self.alloc_size).copied().unwrap_or(0); + let mut values = vec![0i64; self.num_values]; + values[self.heap_live_samples] = 1; + values[self.heap_live_size] = alloc_size; + values + } +} + +/// 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()`. +pub(crate) struct HeapLiveState { + pub tracked: HashMap, + pub max_tracked: usize, + pub excluded_labels: Vec>, + indices: ValueIndices, +} + +/// A single tracked live allocation with owned frame/label/values data. +pub(crate) struct TrackedAlloc { + pub frames: Vec, + pub labels: Vec, + pub values: Vec, +} + +impl HeapLiveState { + /// Create a new heap-live tracker. + /// + /// # Panics + /// Panics if any of the value indices (`alloc_size_idx`, + /// `heap_live_samples_idx`, `heap_live_size_idx`) are >= `num_values`. + 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 { + assert!( + alloc_size_idx < num_values, + "alloc_size_idx ({alloc_size_idx}) must be < num_values ({num_values})" + ); + assert!( + heap_live_samples_idx < num_values, + "heap_live_samples_idx ({heap_live_samples_idx}) must be < num_values ({num_values})" + ); + assert!( + heap_live_size_idx < num_values, + "heap_live_size_idx ({heap_live_size_idx}) must be < num_values ({num_values})" + ); + Self { + tracked: HashMap::with_capacity_and_hasher(max_tracked, FxBuildHasher::default()), + max_tracked, + excluded_labels: excluded_labels.iter().map(|s| Box::from(*s)).collect(), + indices: ValueIndices { + alloc_size: alloc_size_idx, + heap_live_samples: heap_live_samples_idx, + heap_live_size: 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 the tracker is at capacity **and** `ptr` is not + /// already tracked. When `ptr` is already present the entry is replaced + /// unconditionally. + pub fn track(&mut self, ptr: u64, sample: &api::Sample) -> bool { + if self.tracked.len() >= self.max_tracked && !self.tracked.contains_key(&ptr) { + return false; + } + let alloc = TrackedAlloc::from_api_sample(sample, &self.excluded_labels, &self.indices); + self.tracked.insert(ptr, alloc); + true + } + + /// Track a new allocation from a StringIdSample. Resolves ManagedStringIds + /// via the provided string storage into owned strings. + /// + /// Returns false if the tracker is at capacity **and** `ptr` is not + /// already tracked. When `ptr` is already present the entry is replaced + /// unconditionally. + pub fn track_string_id( + &mut self, + ptr: u64, + sample: &api::StringIdSample, + storage: &ManagedStringStorage, + ) -> anyhow::Result { + if self.tracked.len() >= self.max_tracked && !self.tracked.contains_key(&ptr) { + return Ok(false); + } + let alloc = TrackedAlloc::from_string_id_sample( + sample, + storage, + &self.excluded_labels, + &self.indices, + )?; + self.tracked.insert(ptr, alloc); + Ok(true) + } + + /// Remove a tracked allocation. No-op if ptr is not tracked. + pub fn untrack(&mut self, ptr: u64) { + self.tracked.remove(&ptr); + } +} + +fn is_excluded(excluded_labels: &[Box], key: &str) -> bool { + excluded_labels.iter().any(|ex| ex.as_ref() == key) +} + +impl TrackedAlloc { + fn from_api_sample( + sample: &api::Sample, + excluded_labels: &[Box], + indices: &ValueIndices, + ) -> 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| !is_excluded(excluded_labels, l.key)) + .map(|l| OwnedLabel { + key: l.key.into(), + str_value: l.str.into(), + num: l.num, + num_unit: l.num_unit.into(), + }) + .collect(); + + TrackedAlloc { + frames, + labels, + values: indices.build_values(sample.values), + } + } + + fn from_string_id_sample( + sample: &api::StringIdSample, + storage: &ManagedStringStorage, + excluded_labels: &[Box], + indices: &ValueIndices, + ) -> anyhow::Result { + let frames = sample + .locations + .iter() + .map(|loc| -> anyhow::Result { + Ok(OwnedFrame { + function_name: resolve_managed_string(storage, loc.function.name)?, + filename: resolve_managed_string(storage, loc.function.filename)?, + line: loc.line, + }) + }) + .collect::, _>>()?; + + let mut labels = Vec::with_capacity(sample.labels.len()); + for l in &sample.labels { + let key = resolve_managed_string(storage, l.key)?; + if is_excluded(excluded_labels, &key) { + continue; + } + labels.push(OwnedLabel { + key, + str_value: resolve_managed_string(storage, l.str)?, + num: l.num, + num_unit: resolve_managed_string(storage, l.num_unit)?, + }); + } + + Ok(TrackedAlloc { + frames, + labels, + values: indices.build_values(sample.values), + }) + } +} + +fn resolve_managed_string( + storage: &ManagedStringStorage, + id: ManagedStringId, +) -> anyhow::Result> { + if id.value == 0 { + return Ok(Box::from("")); + } + let rc_str = storage.get_string(id.value)?; + Ok(Box::from(rc_str.as_ref())) +} diff --git a/libdd-profiling/src/internal/mod.rs b/libdd-profiling/src/internal/mod.rs index 56e4fe65c6..2b245eec28 100644 --- a/libdd-profiling/src/internal/mod.rs +++ b/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; diff --git a/libdd-profiling/src/internal/profile/fuzz_tests.rs b/libdd-profiling/src/internal/profile/fuzz_tests.rs index a134462254..43f7341b65 100644 --- a/libdd-profiling/src/internal/profile/fuzz_tests.rs +++ b/libdd-profiling/src/internal/profile/fuzz_tests.rs @@ -430,7 +430,7 @@ fn fuzz_add_sample<'a>( samples_with_timestamps: &mut Vec<&'a Sample>, samples_without_timestamps: &mut HashMap<(&'a [Location], &'a [Label]), Vec>, ) { - let r = profile.try_add_sample(sample.into(), *timestamp); + let r = profile.try_add_sample(sample.into(), *timestamp, None); if expected_sample_types.len() == sample.values.len() { assert!(r.is_ok()); if timestamp.is_some() { diff --git a/libdd-profiling/src/internal/profile/mod.rs b/libdd-profiling/src/internal/profile/mod.rs index 280cebcd7a..ffd54115f6 100644 --- a/libdd-profiling/src/internal/profile/mod.rs +++ b/libdd-profiling/src/internal/profile/mod.rs @@ -15,6 +15,7 @@ use crate::api::ManagedStringId; use crate::collections::identifiable::*; use crate::collections::string_storage::{CachedProfileId, ManagedStringStorage}; use crate::collections::string_table::{self, StringTable}; +use crate::internal::heap_live::{HeapLiveState, OwnedFrame, OwnedLabel}; use crate::iter::{IntoLendingIterator, LendingIterator}; use crate::profiles::collections::ArcOverflow; use crate::profiles::datatypes::ProfilesDictionary; @@ -53,6 +54,10 @@ pub struct Profile { string_storage_cached_profile_id: Option, timestamp_key: StringId, upscaling_rules: UpscalingRules, + /// Optional heap-live tracking state. Survives profile resets. + /// When enabled, tracks live allocations and auto-injects them + /// into observations during `reset_and_return_previous()`. + heap_live: Option, } pub struct EncodedProfile { @@ -122,10 +127,21 @@ impl Profile { Ok(()) } + /// Add a sample to the profile. If `track_ptr` is `Some(ptr)`, the sample + /// data is also copied into owned storage for heap-live tracking. The + /// tracked allocation will be automatically injected into the profile's + /// observations during `reset_and_return_previous()` and can be removed + /// earlier via `untrack_allocation()`. + /// + /// If heap-live tracking is at capacity, the tracking request is silently + /// discarded, the sample is still added to the profile normally. + /// If `ptr` matches an already-tracked allocation, the old entry is + /// replaced with the new sample data. pub fn try_add_sample( &mut self, sample: api::Sample, timestamp: Option, + track_ptr: Option, ) -> anyhow::Result<()> { #[cfg(debug_assertions)] { @@ -162,7 +178,18 @@ impl Profile { locations.push(self.try_add_location(location)?); } - self.try_add_sample_internal(sample.values, labels, locations, timestamp) + self.try_add_sample_internal(sample.values, labels, locations, timestamp)?; + + // Track for heap-live only after the sample was successfully added to + // the profile. This avoids phantom tracked allocations when adding + // fails (e.g. due to interning or dedup errors). + if let Some(ptr) = track_ptr { + if let Some(hl) = self.heap_live.as_mut() { + hl.track(ptr, &sample); + } + } + + Ok(()) } /// Tries to add a sample using `api2` structures. @@ -261,6 +288,7 @@ impl Profile { &mut self, sample: api::StringIdSample, timestamp: Option, + track_ptr: Option, ) -> anyhow::Result<()> { anyhow::ensure!( self.string_storage.is_some(), @@ -293,7 +321,22 @@ impl Profile { locations.push(self.add_string_id_location(location)?); } - self.try_add_sample_internal(sample.values, labels, locations, timestamp) + self.try_add_sample_internal(sample.values, labels, locations, timestamp)?; + + // Track for heap-live only after the sample was successfully added. + if let Some(ptr) = track_ptr { + if let Some(hl) = self.heap_live.as_mut() { + // SAFETY: the ensure! at the top of this function guarantees + // string_storage is Some, so unwrap cannot panic here. + let storage = unsafe { self.string_storage.as_ref().unwrap_unchecked() }; + let locked = storage + .lock() + .map_err(|_| anyhow::anyhow!("string storage lock was poisoned"))?; + hl.track_string_id(ptr, &sample, &locked)?; + } + } + + Ok(()) } fn try_add_sample_internal( @@ -338,6 +381,43 @@ impl Profile { Ok(()) } + /// Enable heap-live allocation tracking on this profile. When enabled, + /// calls to `try_add_sample()` with `track_ptr: Some(ptr)` will copy + /// the sample's data into owned storage. Tracked allocations are + /// automatically injected into the profile during + /// `reset_and_return_previous()`, and the tracker moves to the new + /// active profile. + /// + /// # Arguments + /// * `max_tracked` - Maximum number of allocations to track simultaneously + /// * `excluded_labels` - Label keys to strip when copying into the tracker (e.g., + /// high-cardinality labels like "span id", "trace id") + pub fn enable_heap_live_tracking( + &mut self, + max_tracked: usize, + excluded_labels: &[&str], + alloc_size_idx: usize, + heap_live_samples_idx: usize, + heap_live_size_idx: usize, + ) { + self.heap_live = Some(HeapLiveState::new( + max_tracked, + excluded_labels, + alloc_size_idx, + heap_live_samples_idx, + heap_live_size_idx, + self.sample_types.len(), + )); + } + + /// Remove a tracked heap-live allocation by pointer. No-op if heap-live + /// tracking is disabled or the pointer is not tracked. + pub fn untrack_allocation(&mut self, ptr: u64) { + if let Some(hl) = self.heap_live.as_mut() { + hl.untrack(ptr); + } + } + pub fn get_generation(&self) -> anyhow::Result { Ok(self.generation) } @@ -428,6 +508,10 @@ impl Profile { /// Returns the previous Profile on success. #[inline] pub fn reset_and_return_previous(&mut self) -> anyhow::Result { + // Inject heap-live samples into current profile before swapping. + // This makes tracked allocations appear in the serialized output. + self.inject_heap_live_samples()?; + let current_active_samples = self.sample_block()?; anyhow::ensure!( current_active_samples == 0, @@ -451,6 +535,11 @@ impl Profile { ) .context("failed to initialize new profile")?; + // Move heap-live tracker to the new active profile so it survives + // the reset. The old profile (returned) has heap-live data in its + // observations but no tracker. + profile.heap_live = self.heap_live.take(); + std::mem::swap(&mut *self, &mut profile); Ok(profile) } @@ -822,6 +911,36 @@ impl Profile { } } + /// Inject all tracked heap-live allocations into the profile's + /// observations. Called automatically by `reset_and_return_previous()`. + /// Takes the heap_live state out temporarily to avoid borrow conflicts. + fn inject_heap_live_samples(&mut self) -> anyhow::Result<()> { + let Some(hl) = self.heap_live.take() else { + return Ok(()); + }; + let result = (|| { + for alloc in hl.tracked.values() { + let locations: Vec> = alloc + .frames + .iter() + .map(OwnedFrame::as_api_location) + .collect(); + let labels: Vec> = + alloc.labels.iter().map(OwnedLabel::as_api_label).collect(); + let sample = api::Sample { + locations, + values: &alloc.values, + labels, + }; + // Pass None for track_ptr to avoid re-tracking during injection + self.try_add_sample(sample, None, None)?; + } + anyhow::Ok(()) + })(); + self.heap_live = Some(hl); + result + } + /// Fetches the endpoint information for the label. There may be errors, /// but there may also be no endpoint information for a given endpoint. /// Hence, the return type of Result, _>. @@ -939,6 +1058,7 @@ impl Profile { * CachedProfileId for why. */ timestamp_key: Default::default(), upscaling_rules: Default::default(), + heap_live: None, }; let _id = profile.intern(""); @@ -1055,6 +1175,433 @@ impl Profile { } } +#[cfg(test)] +mod heap_live_tests { + use super::*; + use crate::pprof::test_utils::{roundtrip_to_pprof, sorted_samples, string_table_fetch}; + + // Sample types matching the dd-trace-php layout when allocation + heap-live + // profiling are enabled (without cpu-time). + // Indices: 0=alloc-samples, 1=alloc-size, 2=heap-live-samples, 3=heap-live-size + const ALLOC_SIZE_IDX: usize = 1; + const HEAP_LIVE_SAMPLES_IDX: usize = 2; + const HEAP_LIVE_SIZE_IDX: usize = 3; + + fn heap_live_sample_types() -> [api::SampleType; 4] { + [ + api::SampleType::AllocSamples, + api::SampleType::AllocSize, + api::SampleType::HeapLiveSamples, + api::SampleType::HeapLiveSize, + ] + } + + fn make_mapping() -> api::Mapping<'static> { + api::Mapping { + filename: "php", + ..Default::default() + } + } + + fn make_locations(mapping: api::Mapping<'static>) -> Vec> { + vec![api::Location { + mapping, + function: api::Function { + name: "malloc", + system_name: "malloc", + filename: "test.php", + }, + line: 10, + ..Default::default() + }] + } + + #[test] + fn heap_live_tracked_sample_appears_in_serialized_pprof() { + let sample_types = heap_live_sample_types(); + let mut profile = Profile::new(&sample_types, None); + profile.enable_heap_live_tracking( + 4096, + &[], + ALLOC_SIZE_IDX, + HEAP_LIVE_SAMPLES_IDX, + HEAP_LIVE_SIZE_IDX, + ); + + let mapping = make_mapping(); + let locations = make_locations(mapping); + let labels = vec![api::Label { + key: "thread id", + num: 42, + ..Default::default() + }]; + + // values: [alloc-samples=1, alloc-size=1024, heap-live-samples=0, heap-live-size=0] + let sample = api::Sample { + locations, + values: &[1, 1024, 0, 0], + labels, + }; + + profile + .try_add_sample(sample, None, Some(0x1000)) + .expect("add to succeed"); + + let old = profile + .reset_and_return_previous() + .expect("reset to succeed"); + + let pprof = roundtrip_to_pprof(old).unwrap(); + // Original sample + 1 injected heap-live copy + assert_eq!(pprof.samples.len(), 2); + + let samples = sorted_samples(&pprof); + let mut found_values: Vec> = samples.iter().map(|s| s.values.clone()).collect(); + found_values.sort(); + // Original: [1, 1024, 0, 0], Injected: [0, 0, 1, 1024] + assert_eq!(found_values, vec![vec![0, 0, 1, 1024], vec![1, 1024, 0, 0]]); + + // Verify function name appears + assert!(pprof.string_table.iter().any(|s| s == "malloc")); + assert!(pprof.string_table.iter().any(|s| s == "test.php")); + + // Verify label on all samples + for s in &samples { + assert!(!s.labels.is_empty()); + let label = &s.labels[0]; + assert_eq!(string_table_fetch(&pprof, label.key), "thread id"); + assert_eq!(label.num, 42); + } + } + + #[test] + fn heap_live_untracked_sample_does_not_appear() { + let sample_types = heap_live_sample_types(); + let mut profile = Profile::new(&sample_types, None); + profile.enable_heap_live_tracking( + 4096, + &[], + ALLOC_SIZE_IDX, + HEAP_LIVE_SAMPLES_IDX, + HEAP_LIVE_SIZE_IDX, + ); + + let mapping = make_mapping(); + let locations = make_locations(mapping); + + let sample = api::Sample { + locations, + values: &[1, 512, 0, 0], + labels: vec![], + }; + + profile + .try_add_sample(sample, None, Some(0x2000)) + .expect("add to succeed"); + + // Untrack before reset — heap-live injection should skip this ptr + profile.untrack_allocation(0x2000); + + let old = profile + .reset_and_return_previous() + .expect("reset to succeed"); + + let pprof = roundtrip_to_pprof(old).unwrap(); + // Only the original sample, no heap-live injection + assert_eq!(pprof.samples.len(), 1); + } + + #[test] + fn heap_live_survives_reset() { + let sample_types = heap_live_sample_types(); + let mut profile = Profile::new(&sample_types, None); + profile.enable_heap_live_tracking( + 4096, + &[], + ALLOC_SIZE_IDX, + HEAP_LIVE_SAMPLES_IDX, + HEAP_LIVE_SIZE_IDX, + ); + + let mapping = make_mapping(); + let locations = make_locations(mapping); + + // values: [alloc-samples=1, alloc-size=2048, heap-live-samples=0, heap-live-size=0] + let sample = api::Sample { + locations: locations.clone(), + values: &[1, 2048, 0, 0], + labels: vec![], + }; + + profile + .try_add_sample(sample, None, Some(0x3000)) + .expect("add to succeed"); + + // First reset: original + heap-live copy + let old1 = profile + .reset_and_return_previous() + .expect("first reset to succeed"); + + let pprof1 = roundtrip_to_pprof(old1).unwrap(); + assert_eq!(pprof1.samples.len(), 2); + + // Add a new regular sample to the reset profile + let regular_sample = api::Sample { + locations, + values: &[5, 9999, 0, 0], + labels: vec![], + }; + + profile + .try_add_sample(regular_sample, None, None) + .expect("add regular to succeed"); + + // Second reset: the still-tracked heap-live + the new regular + let old2 = profile + .reset_and_return_previous() + .expect("second reset to succeed"); + + let pprof2 = roundtrip_to_pprof(old2).unwrap(); + assert_eq!(pprof2.samples.len(), 2); + + let samples = sorted_samples(&pprof2); + let mut found_values: Vec> = samples.iter().map(|s| s.values.clone()).collect(); + found_values.sort(); + // Injected heap-live: [0, 0, 1, 2048], Regular: [5, 9999, 0, 0] + assert_eq!(found_values, vec![vec![0, 0, 1, 2048], vec![5, 9999, 0, 0]]); + } + + #[test] + fn heap_live_excluded_labels_are_filtered() { + let sample_types = heap_live_sample_types(); + let mut profile = Profile::new(&sample_types, None); + profile.enable_heap_live_tracking( + 4096, + &["span id", "trace id"], + ALLOC_SIZE_IDX, + HEAP_LIVE_SAMPLES_IDX, + HEAP_LIVE_SIZE_IDX, + ); + + let mapping = make_mapping(); + let locations = make_locations(mapping); + + let labels = vec![ + api::Label { + key: "thread id", + num: 1, + ..Default::default() + }, + api::Label { + key: "span id", + num: 999, + ..Default::default() + }, + api::Label { + key: "trace id", + num: 888, + ..Default::default() + }, + ]; + + let sample = api::Sample { + locations, + values: &[1, 256, 0, 0], + labels, + }; + + profile + .try_add_sample(sample, None, Some(0x4000)) + .expect("add to succeed"); + + let old = profile + .reset_and_return_previous() + .expect("reset to succeed"); + + let pprof = roundtrip_to_pprof(old).unwrap(); + // Original (3 labels) + heap-live copy (filtered labels) + assert_eq!(pprof.samples.len(), 2); + + let samples = sorted_samples(&pprof); + + // Find the sample with fewer labels — that's the injected heap-live copy + let injected = samples + .iter() + .find(|s| s.labels.len() == 1) + .expect("should find heap-live sample with 1 label"); + + let label = &injected.labels[0]; + assert_eq!(string_table_fetch(&pprof, label.key), "thread id"); + assert_eq!(label.num, 1); + + // The original sample should still have all 3 labels + let original = samples + .iter() + .find(|s| s.labels.len() == 3) + .expect("should find original sample with 3 labels"); + assert_eq!(original.labels.len(), 3); + } + + #[test] + fn heap_live_capacity_limit() { + let sample_types = heap_live_sample_types(); + let mut profile = Profile::new(&sample_types, None); + // Only allow 2 tracked allocations + profile.enable_heap_live_tracking( + 2, + &[], + ALLOC_SIZE_IDX, + HEAP_LIVE_SAMPLES_IDX, + HEAP_LIVE_SIZE_IDX, + ); + + let mapping = make_mapping(); + + // Add 3 samples with different track_ptrs + for (i, ptr) in [1u64, 2, 3].iter().enumerate() { + let locations = vec![api::Location { + mapping, + function: api::Function { + name: "alloc_fn", + system_name: "alloc_fn", + filename: "test.php", + }, + line: (i + 1) as i64, + ..Default::default() + }]; + + let values: Vec = vec![1, (i as i64 + 1) * 100, 0, 0]; + let sample = api::Sample { + locations, + values: &values, + labels: vec![], + }; + + profile + .try_add_sample(sample, None, Some(*ptr)) + .expect("add to succeed"); + } + + let old = profile + .reset_and_return_previous() + .expect("reset to succeed"); + + let pprof = roundtrip_to_pprof(old).unwrap(); + // 3 original samples + 2 heap-live injected (third rejected at capacity) + assert_eq!(pprof.samples.len(), 5); + } + + #[test] + fn heap_live_disabled_by_default() { + let sample_types = heap_live_sample_types(); + let mut profile = Profile::new(&sample_types, None); + // Do NOT call enable_heap_live_tracking + + let mapping = make_mapping(); + let locations = make_locations(mapping); + + let sample = api::Sample { + locations, + values: &[1, 4096, 0, 0], + labels: vec![], + }; + + profile + .try_add_sample(sample, None, Some(0x5000)) + .expect("add to succeed"); + + let old = profile + .reset_and_return_previous() + .expect("reset to succeed"); + + let pprof = roundtrip_to_pprof(old).unwrap(); + // Only the original sample — track_ptr silently ignored + assert_eq!(pprof.samples.len(), 1); + } + + #[test] + fn heap_live_mixed_tracked_and_regular_samples() { + let sample_types = heap_live_sample_types(); + let mut profile = Profile::new(&sample_types, None); + profile.enable_heap_live_tracking( + 4096, + &[], + ALLOC_SIZE_IDX, + HEAP_LIVE_SAMPLES_IDX, + HEAP_LIVE_SIZE_IDX, + ); + + let mapping = make_mapping(); + + // Tracked sample + let tracked_locations = vec![api::Location { + mapping, + function: api::Function { + name: "tracked_fn", + system_name: "tracked_fn", + filename: "alloc.php", + }, + line: 1, + ..Default::default() + }]; + + let tracked_sample = api::Sample { + locations: tracked_locations, + values: &[1, 512, 0, 0], + labels: vec![], + }; + + profile + .try_add_sample(tracked_sample, None, Some(0x6000)) + .expect("add tracked to succeed"); + + // Regular sample (no tracking) + let regular_locations = vec![api::Location { + mapping, + function: api::Function { + name: "regular_fn", + system_name: "regular_fn", + filename: "main.php", + }, + line: 5, + ..Default::default() + }]; + + let regular_sample = api::Sample { + locations: regular_locations, + values: &[3, 7777, 0, 0], + labels: vec![], + }; + + profile + .try_add_sample(regular_sample, None, None) + .expect("add regular to succeed"); + + let old = profile + .reset_and_return_previous() + .expect("reset to succeed"); + + let pprof = roundtrip_to_pprof(old).unwrap(); + // 2 original + 1 injected heap-live copy of the tracked sample + assert_eq!(pprof.samples.len(), 3); + + // Verify we have both function names + assert!(pprof.string_table.iter().any(|s| s == "tracked_fn")); + assert!(pprof.string_table.iter().any(|s| s == "regular_fn")); + + let samples = sorted_samples(&pprof); + let mut found_values: Vec> = samples.iter().map(|s| s.values.clone()).collect(); + found_values.sort(); + // Original tracked: [1, 512, 0, 0] + // Injected heap-live: [0, 0, 1, 512] + // Regular: [3, 7777, 0, 0] + assert_eq!( + found_values, + vec![vec![0, 0, 1, 512], vec![1, 512, 0, 0], vec![3, 7777, 0, 0]] + ); + } +} + #[cfg(test)] mod api_tests { use super::*; @@ -1120,6 +1667,7 @@ mod api_tests { labels: vec![], }, None, + None, ) .expect("add to succeed"); @@ -1191,18 +1739,18 @@ mod api_tests { assert_eq!(profile.only_for_testing_num_aggregated_samples(), 0); profile - .try_add_sample(main_sample, None) + .try_add_sample(main_sample, None, None) .expect("profile to not be full"); assert_eq!(profile.only_for_testing_num_aggregated_samples(), 1); profile - .try_add_sample(test_sample, None) + .try_add_sample(test_sample, None, None) .expect("profile to not be full"); assert_eq!(profile.only_for_testing_num_aggregated_samples(), 2); assert_eq!(profile.only_for_testing_num_timestamped_samples(), 0); profile - .try_add_sample(timestamp_sample, Timestamp::new(42)) + .try_add_sample(timestamp_sample, Timestamp::new(42), None) .expect("profile to not be full"); assert_eq!(profile.only_for_testing_num_timestamped_samples(), 1); profile @@ -1372,7 +1920,7 @@ mod api_tests { labels: vec![id_label], }; - assert!(profile.try_add_sample(sample, None).is_err()); + assert!(profile.try_add_sample(sample, None, None).is_err()); } #[test] @@ -1415,11 +1963,11 @@ mod api_tests { }; profile - .try_add_sample(sample1, None) + .try_add_sample(sample1, None, None) .expect("add to success"); profile - .try_add_sample(sample2, None) + .try_add_sample(sample2, None, None) .expect("add to success"); profile.add_endpoint(10, Cow::from("my endpoint"))?; @@ -1535,7 +2083,7 @@ mod api_tests { labels, }; - profile.try_add_sample(sample, None).unwrap_err(); + profile.try_add_sample(sample, None, None).unwrap_err(); } #[test] @@ -1558,7 +2106,7 @@ mod api_tests { }; profile - .try_add_sample(sample1, None) + .try_add_sample(sample1, None, None) .expect("add to success"); let serialized_profile = roundtrip_to_pprof(profile).unwrap(); @@ -1600,7 +2148,7 @@ mod api_tests { }; profile - .try_add_sample(sample1, None) + .try_add_sample(sample1, None, None) .expect("add to success"); let upscaling_info = UpscalingInfo::Proportional { scale: 2.0 }; @@ -1630,7 +2178,7 @@ mod api_tests { }; profile - .try_add_sample(sample1, None) + .try_add_sample(sample1, None, None) .expect("add to success"); let upscaling_info = UpscalingInfo::Proportional { scale: 2.7 }; @@ -1660,7 +2208,7 @@ mod api_tests { }; profile - .try_add_sample(sample1, None) + .try_add_sample(sample1, None, None) .expect("add to success"); let upscaling_info = UpscalingInfo::Poisson { @@ -1694,7 +2242,7 @@ mod api_tests { }; profile - .try_add_sample(sample1, None) + .try_add_sample(sample1, None, None) .expect("add to success"); let upscaling_info = UpscalingInfo::PoissonNonSampleTypeCount { @@ -1728,7 +2276,7 @@ mod api_tests { }; profile - .try_add_sample(sample1, None) + .try_add_sample(sample1, None, None) .expect("add to success"); let upscaling_info = UpscalingInfo::Poisson { @@ -1762,7 +2310,7 @@ mod api_tests { }; profile - .try_add_sample(sample1, None) + .try_add_sample(sample1, None, None) .expect("add to success"); // invalid sampling_distance value @@ -1833,10 +2381,10 @@ mod api_tests { }; profile - .try_add_sample(sample1, None) + .try_add_sample(sample1, None, None) .expect("add to success"); profile - .try_add_sample(sample2, None) + .try_add_sample(sample2, None, None) .expect("add to success"); // upscale the first value and the last one @@ -1892,10 +2440,10 @@ mod api_tests { }; profile - .try_add_sample(sample1, None) + .try_add_sample(sample1, None, None) .expect("add to success"); profile - .try_add_sample(sample2, None) + .try_add_sample(sample2, None, None) .expect("add to success"); let mut values_offset: Vec = vec![0]; @@ -1941,7 +2489,7 @@ mod api_tests { }; profile - .try_add_sample(sample1, None) + .try_add_sample(sample1, None, None) .expect("add to success"); let values_offset: Vec = vec![0]; @@ -1999,7 +2547,7 @@ mod api_tests { }; profile - .try_add_sample(sample1, None) + .try_add_sample(sample1, None, None) .expect("add to success"); let upscaling_info = UpscalingInfo::Proportional { scale: 2.0 }; @@ -2057,10 +2605,10 @@ mod api_tests { }; profile - .try_add_sample(sample1, None) + .try_add_sample(sample1, None, None) .expect("add to success"); profile - .try_add_sample(sample2, None) + .try_add_sample(sample2, None, None) .expect("add to success"); let upscaling_info = UpscalingInfo::Proportional { scale: 2.0 }; @@ -2131,10 +2679,10 @@ mod api_tests { }; profile - .try_add_sample(sample1, None) + .try_add_sample(sample1, None, None) .expect("add to success"); profile - .try_add_sample(sample2, None) + .try_add_sample(sample2, None, None) .expect("add to success"); // add rule for the first sample on the 1st value @@ -2188,7 +2736,7 @@ mod api_tests { }; profile - .try_add_sample(sample1, None) + .try_add_sample(sample1, None, None) .expect("add to success"); // upscale samples and wall-time values @@ -2226,7 +2774,7 @@ mod api_tests { }; profile - .try_add_sample(sample1, None) + .try_add_sample(sample1, None, None) .expect("add to success"); let upscaling_info = UpscalingInfo::Proportional { scale: 2.0 }; @@ -2500,7 +3048,7 @@ mod api_tests { }; profile - .try_add_sample(sample1, None) + .try_add_sample(sample1, None, None) .expect("add to success"); let mut value_offsets: Vec = vec![0, 1]; @@ -2566,10 +3114,10 @@ mod api_tests { }; profile - .try_add_sample(sample1, None) + .try_add_sample(sample1, None, None) .expect("add to success"); profile - .try_add_sample(sample2, None) + .try_add_sample(sample2, None, None) .expect("add to success"); profile.add_endpoint(10, Cow::from("endpoint 10"))?; @@ -2868,8 +3416,12 @@ mod api_tests { labels: vec![], }; - profile.add_string_id_sample(sample.clone(), None).unwrap(); - profile.add_string_id_sample(sample.clone(), None).unwrap(); + profile + .add_string_id_sample(sample.clone(), None, None) + .unwrap(); + profile + .add_string_id_sample(sample.clone(), None, None) + .unwrap(); let pprof_first_profile = roundtrip_to_pprof(profile.reset_and_return_previous().unwrap()).unwrap(); @@ -2886,8 +3438,12 @@ mod api_tests { // If the cache invalidation on the managed string table is working correctly, these strings // get correctly re-added to the profile's string table - profile.add_string_id_sample(sample.clone(), None).unwrap(); - profile.add_string_id_sample(sample.clone(), None).unwrap(); + profile + .add_string_id_sample(sample.clone(), None, None) + .unwrap(); + profile + .add_string_id_sample(sample.clone(), None, None) + .unwrap(); let pprof_second_profile = roundtrip_to_pprof(profile).unwrap(); assert!(pprof_second_profile @@ -2932,7 +3488,7 @@ mod api_tests { }; profile - .try_add_sample(sample, None) // No timestamp = aggregated + .try_add_sample(sample, None, None) // No timestamp = aggregated .expect("profile to not be full"); // We want to trigger an error inside the loop over observations. From f51c09961d3620fba36ff6574077f45be9694586 Mon Sep 17 00:00:00 2001 From: Florian Engelhardt Date: Wed, 11 Mar 2026 13:58:46 +0100 Subject: [PATCH 2/2] update --- datadog-profiling-replayer/src/main.rs | 2 +- examples/ffi/exporter.cpp | 2 +- examples/ffi/exporter_manager.c | 2 +- examples/ffi/profiles.c | 2 +- libdd-profiling-ffi/src/profiles/datatypes.rs | 202 +++++++++--- libdd-profiling/benches/add_samples.rs | 2 +- libdd-profiling/examples/profiles.rs | 2 +- libdd-profiling/src/cxx.rs | 2 +- libdd-profiling/src/internal/heap_live.rs | 151 ++++++--- .../src/internal/profile/fuzz_tests.rs | 2 +- libdd-profiling/src/internal/profile/mod.rs | 310 +++++++++--------- 11 files changed, 424 insertions(+), 255 deletions(-) diff --git a/datadog-profiling-replayer/src/main.rs b/datadog-profiling-replayer/src/main.rs index 0ee889fec1..6a8ae2b32f 100644 --- a/datadog-profiling-replayer/src/main.rs +++ b/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, None)?; + outprof.try_add_sample(sample, timestamp)?; } let duration = before.elapsed(); diff --git a/examples/ffi/exporter.cpp b/examples/ffi/exporter.cpp index 65e8b4c750..5b02debf7d 100644 --- a/examples/ffi/exporter.cpp +++ b/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, 0); + auto add_result = ddog_prof_Profile_add(profile.get(), sample, 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); diff --git a/examples/ffi/exporter_manager.c b/examples/ffi/exporter_manager.c index f75ad3a3da..765c11d2de 100644 --- a/examples/ffi/exporter_manager.c +++ b/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, 0); + ddog_prof_Profile_Result add_result = ddog_prof_Profile_add(profile, sample, 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); diff --git a/examples/ffi/profiles.c b/examples/ffi/profiles.c index 054d984d27..a4b0eb2800 100644 --- a/examples/ffi/profiles.c +++ b/examples/ffi/profiles.c @@ -62,7 +62,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, 0); + ddog_prof_Profile_Result add_result = ddog_prof_Profile_add(&profile, sample, 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); diff --git a/libdd-profiling-ffi/src/profiles/datatypes.rs b/libdd-profiling-ffi/src/profiles/datatypes.rs index c2a93f3120..3dc86bf132 100644 --- a/libdd-profiling-ffi/src/profiles/datatypes.rs +++ b/libdd-profiling-ffi/src/profiles/datatypes.rs @@ -242,6 +242,13 @@ pub struct Sample2<'a> { pub labels: Slice<'a, Label2<'a>>, } +#[repr(C)] +#[derive(Copy, Clone)] +pub struct HeapLiveConfig<'a> { + pub max_tracked_allocations: usize, + pub excluded_label_keys: Slice<'a, CharSlice<'a>>, +} + impl<'a> TryFrom<&'a Mapping<'a>> for api::Mapping<'a> { type Error = Utf8Error; @@ -549,36 +556,64 @@ impl From for Result { /// 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, - track_ptr: usize, ) -> ProfileResult { (|| { let profile = profile_ptr_to_inner(profile)?; - let track = (track_ptr != 0).then_some(track_ptr as u64); 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, track) + profile.add_string_id_sample(sample.into(), timestamp) } else { - profile.try_add_sample(sample.try_into()?, timestamp, track) + profile.try_add_sample(sample.try_into()?, timestamp) } })() .context("ddog_prof_Profile_add failed") .into() } + +/// # Safety +/// The `profile` ptr must point to a valid Profile object created by this +/// module. All pointers inside the `sample` need to be valid for the duration +/// of this call. +/// +/// If successful, it returns the Ok variant. +/// On error, it holds an error message in the error variant. +/// +/// # Arguments +/// * `ptr` - The allocation pointer to retain for heap-live profiling. +#[must_use] +#[no_mangle] +pub unsafe extern "C" fn ddog_prof_Profile_add_tracked_allocation( + profile: *mut Profile, + sample: Sample, + timestamp: Option, + ptr: usize, +) -> ProfileResult { + (|| { + let profile = profile_ptr_to_inner(profile)?; + 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_tracked_string_id_allocation(sample.into(), timestamp, ptr as u64) + } else { + profile.add_tracked_allocation(sample.try_into()?, timestamp, ptr as u64) + } + })() + .context("ddog_prof_Profile_add_tracked_allocation failed") + .into() +} /// # Safety /// The `profile` ptr must point to a valid Profile object created by this /// module. All pointers inside the `sample` need to be valid for the duration @@ -903,49 +938,33 @@ 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. +/// Configure heap-live allocation tracking on this profile. 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"). -/// * `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. +/// * `config` - Heap-live tracking configuration. Required sample types are derived from the +/// profile schema and validated here. /// /// # Safety /// The `profile` ptr must point to a valid Profile object created by this -/// module. `excluded_labels` must be a valid `Slice`. +/// module. `config.excluded_label_keys` must be a valid `Slice`. /// This call is _NOT_ thread-safe. #[no_mangle] -pub unsafe extern "C" fn ddog_prof_Profile_enable_heap_live_tracking( +pub unsafe extern "C" fn ddog_prof_Profile_configure_heap_live( profile: *mut Profile, - max_tracked: usize, - excluded_labels: Slice, - alloc_size_idx: usize, - heap_live_samples_idx: usize, - heap_live_size_idx: usize, + config: HeapLiveConfig, ) -> ProfileResult { (|| { let profile = profile_ptr_to_inner(profile)?; - let labels: Vec<&str> = excluded_labels + let labels: Vec<&str> = config + .excluded_label_keys .iter() .map(|cs| cs.try_to_utf8()) .collect::, _>>()?; - profile.enable_heap_live_tracking( - max_tracked, - &labels, - alloc_size_idx, - heap_live_samples_idx, - heap_live_size_idx, - ); - anyhow::Ok(()) + profile.configure_heap_live(config.max_tracked_allocations, &labels) })() - .context("ddog_prof_Profile_enable_heap_live_tracking failed") + .context("ddog_prof_Profile_configure_heap_live failed") .into() } @@ -1001,7 +1020,7 @@ mod tests { labels: Slice::empty(), }; - let result = Result::from(ddog_prof_Profile_add(&mut profile, sample, None, 0)); + let result = Result::from(ddog_prof_Profile_add(&mut profile, sample, None)); result.unwrap_err(); ddog_prof_Profile_drop(&mut profile); Ok(()) @@ -1049,7 +1068,7 @@ mod tests { labels: Slice::from(&labels), }; - Result::from(ddog_prof_Profile_add(&mut profile, sample, None, 0))?; + Result::from(ddog_prof_Profile_add(&mut profile, sample, None))?; assert_eq!( profile .inner @@ -1059,7 +1078,7 @@ mod tests { 1 ); - Result::from(ddog_prof_Profile_add(&mut profile, sample, None, 0))?; + Result::from(ddog_prof_Profile_add(&mut profile, sample, None))?; assert_eq!( profile .inner @@ -1135,7 +1154,7 @@ mod tests { labels: Slice::from(labels.as_slice()), }; - Result::from(ddog_prof_Profile_add(&mut profile, main_sample, None, 0)).unwrap(); + Result::from(ddog_prof_Profile_add(&mut profile, main_sample, None)).unwrap(); assert_eq!( profile .inner @@ -1145,7 +1164,7 @@ mod tests { 1 ); - Result::from(ddog_prof_Profile_add(&mut profile, test_sample, None, 0)).unwrap(); + Result::from(ddog_prof_Profile_add(&mut profile, test_sample, None)).unwrap(); assert_eq!( profile .inner @@ -1164,4 +1183,105 @@ mod tests { ddog_prof_Profile_drop(&mut provide_distinct_locations_ffi()); } } + + #[test] + fn configure_heap_live_ffi_requires_expected_sample_types() -> Result<(), Error> { + unsafe { + let sample_types = [SampleType::AllocSamples, SampleType::AllocSize]; + let mut profile = Result::from(ddog_prof_Profile_new( + Slice::from_raw_parts(sample_types.as_ptr(), sample_types.len()), + None, + ))?; + + let result = Result::from(ddog_prof_Profile_configure_heap_live( + &mut profile, + HeapLiveConfig { + max_tracked_allocations: 16, + excluded_label_keys: Slice::empty(), + }, + )); + assert!(result.unwrap_err().to_string().contains("HeapLiveSamples")); + + ddog_prof_Profile_drop(&mut profile); + Ok(()) + } + } + + #[test] + fn add_tracked_allocation_ffi_requires_configuration() -> Result<(), Error> { + unsafe { + let sample_types = [ + SampleType::AllocSamples, + SampleType::AllocSize, + SampleType::HeapLiveSamples, + SampleType::HeapLiveSize, + ]; + let mut profile = Result::from(ddog_prof_Profile_new( + Slice::from_raw_parts(sample_types.as_ptr(), sample_types.len()), + None, + ))?; + + let values = [1, 128, 0, 0]; + let sample = Sample { + locations: Slice::empty(), + values: Slice::from(values.as_slice()), + labels: Slice::empty(), + }; + + let result = Result::from(ddog_prof_Profile_add_tracked_allocation( + &mut profile, + sample, + None, + 0x1234, + )); + assert!(result + .unwrap_err() + .to_string() + .contains("heap-live tracking is not configured for this profile")); + + ddog_prof_Profile_drop(&mut profile); + Ok(()) + } + } + + #[test] + fn configure_heap_live_and_add_tracked_allocation_ffi() -> Result<(), Error> { + unsafe { + let sample_types = [ + SampleType::AllocSamples, + SampleType::AllocSize, + SampleType::HeapLiveSamples, + SampleType::HeapLiveSize, + ]; + let mut profile = Result::from(ddog_prof_Profile_new( + Slice::from_raw_parts(sample_types.as_ptr(), sample_types.len()), + None, + ))?; + + Result::from(ddog_prof_Profile_configure_heap_live( + &mut profile, + HeapLiveConfig { + max_tracked_allocations: 16, + excluded_label_keys: Slice::empty(), + }, + ))?; + + let values = [1, 128, 0, 0]; + let sample = Sample { + locations: Slice::empty(), + values: Slice::from(values.as_slice()), + labels: Slice::empty(), + }; + + Result::from(ddog_prof_Profile_add_tracked_allocation( + &mut profile, + sample, + None, + 0x1234, + ))?; + + ddog_prof_Profile_drop(&mut profile); + Ok(()) + } + } } diff --git a/libdd-profiling/benches/add_samples.rs b/libdd-profiling/benches/add_samples.rs index 5ca0b07a57..fc27d886b3 100644 --- a/libdd-profiling/benches/add_samples.rs +++ b/libdd-profiling/benches/add_samples.rs @@ -138,7 +138,7 @@ pub fn bench_add_sample_vs_add2(c: &mut Criterion) { values: &values, labels: labels_api.clone(), }; - black_box(profile.try_add_sample(sample, None, None)).unwrap(); + black_box(profile.try_add_sample(sample, None)).unwrap(); } black_box(profile.only_for_testing_num_aggregated_samples()) }) diff --git a/libdd-profiling/examples/profiles.rs b/libdd-profiling/examples/profiles.rs index e72f2642c7..2d4fed25d4 100644 --- a/libdd-profiling/examples/profiles.rs +++ b/libdd-profiling/examples/profiles.rs @@ -48,7 +48,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, None) { + match profile.try_add_sample(sample, None) { Ok(_) => {} Err(_) => exit(1), } diff --git a/libdd-profiling/src/cxx.rs b/libdd-profiling/src/cxx.rs index d281d5d256..312b1c636d 100644 --- a/libdd-profiling/src/cxx.rs +++ b/libdd-profiling/src/cxx.rs @@ -549,7 +549,7 @@ impl Profile { }; // Profile interns the strings - self.inner.try_add_sample(api_sample, None, None)?; + self.inner.try_add_sample(api_sample, None)?; Ok(()) } diff --git a/libdd-profiling/src/internal/heap_live.rs b/libdd-profiling/src/internal/heap_live.rs index 884d1a85a4..31c1afe1db 100644 --- a/libdd-profiling/src/internal/heap_live.rs +++ b/libdd-profiling/src/internal/heap_live.rs @@ -3,16 +3,32 @@ use crate::api; use crate::api::ManagedStringId; +use crate::api::SampleType; use crate::collections::string_storage::ManagedStringStorage; use std::collections::HashMap; use std::hash::BuildHasherDefault; -/// Owned frame data for heap-live tracking. +/// Owned location data for heap-live tracking. /// Stores copies of borrowed strings so tracked allocations survive across /// profile resets. -pub(crate) struct OwnedFrame { +pub(crate) struct OwnedMapping { + pub memory_start: u64, + pub memory_limit: u64, + pub file_offset: u64, + pub filename: Box, + pub build_id: Box, +} + +pub(crate) struct OwnedFunction { pub function_name: Box, + pub system_name: Box, pub filename: Box, +} + +pub(crate) struct OwnedLocation { + pub mapping: OwnedMapping, + pub function: OwnedFunction, + pub address: u64, pub line: i64, } @@ -24,16 +40,35 @@ pub(crate) struct OwnedLabel { pub num_unit: Box, } -impl OwnedFrame { +impl OwnedMapping { + pub fn as_api_mapping(&self) -> api::Mapping<'_> { + api::Mapping { + memory_start: self.memory_start, + memory_limit: self.memory_limit, + file_offset: self.file_offset, + filename: &self.filename, + build_id: &self.build_id, + } + } +} + +impl OwnedFunction { + pub fn as_api_function(&self) -> api::Function<'_> { + api::Function { + name: &self.function_name, + system_name: &self.system_name, + filename: &self.filename, + } + } +} + +impl OwnedLocation { pub fn as_api_location(&self) -> api::Location<'_> { api::Location { - function: api::Function { - name: &self.function_name, - system_name: "", - filename: &self.filename, - }, + mapping: self.mapping.as_api_mapping(), + function: self.function.as_api_function(), + address: self.address, line: self.line, - ..api::Location::default() } } } @@ -64,6 +99,24 @@ struct ValueIndices { } impl ValueIndices { + fn for_sample_types(sample_types: &[SampleType]) -> anyhow::Result { + fn index_of(sample_types: &[SampleType], target: SampleType) -> anyhow::Result { + sample_types + .iter() + .position(|sample_type| *sample_type == target) + .ok_or_else(|| { + anyhow::anyhow!("heap-live tracking requires sample type {target:?}") + }) + } + + Ok(Self { + alloc_size: index_of(sample_types, SampleType::AllocSize)?, + heap_live_samples: index_of(sample_types, SampleType::HeapLiveSamples)?, + heap_live_size: index_of(sample_types, SampleType::HeapLiveSize)?, + num_values: sample_types.len(), + }) + } + /// Build a heap-live values vector: all zeros except heap-live-samples=1 /// and heap-live-size=alloc_size (read from the original sample). fn build_values(&self, sample_values: &[i64]) -> Vec { @@ -89,48 +142,24 @@ pub(crate) struct HeapLiveState { /// A single tracked live allocation with owned frame/label/values data. pub(crate) struct TrackedAlloc { - pub frames: Vec, + pub locations: Vec, pub labels: Vec, pub values: Vec, } impl HeapLiveState { /// Create a new heap-live tracker. - /// - /// # Panics - /// Panics if any of the value indices (`alloc_size_idx`, - /// `heap_live_samples_idx`, `heap_live_size_idx`) are >= `num_values`. 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 { - assert!( - alloc_size_idx < num_values, - "alloc_size_idx ({alloc_size_idx}) must be < num_values ({num_values})" - ); - assert!( - heap_live_samples_idx < num_values, - "heap_live_samples_idx ({heap_live_samples_idx}) must be < num_values ({num_values})" - ); - assert!( - heap_live_size_idx < num_values, - "heap_live_size_idx ({heap_live_size_idx}) must be < num_values ({num_values})" - ); - Self { + sample_types: &[SampleType], + ) -> anyhow::Result { + Ok(Self { tracked: HashMap::with_capacity_and_hasher(max_tracked, FxBuildHasher::default()), max_tracked, excluded_labels: excluded_labels.iter().map(|s| Box::from(*s)).collect(), - indices: ValueIndices { - alloc_size: alloc_size_idx, - heap_live_samples: heap_live_samples_idx, - heap_live_size: heap_live_size_idx, - num_values, - }, - } + indices: ValueIndices::for_sample_types(sample_types)?, + }) } /// Track a new allocation. Copies borrowed strings from the sample into @@ -190,12 +219,23 @@ impl TrackedAlloc { excluded_labels: &[Box], indices: &ValueIndices, ) -> Self { - let frames = sample + let locations = sample .locations .iter() - .map(|loc| OwnedFrame { - function_name: loc.function.name.into(), - filename: loc.function.filename.into(), + .map(|loc| OwnedLocation { + mapping: OwnedMapping { + memory_start: loc.mapping.memory_start, + memory_limit: loc.mapping.memory_limit, + file_offset: loc.mapping.file_offset, + filename: loc.mapping.filename.into(), + build_id: loc.mapping.build_id.into(), + }, + function: OwnedFunction { + function_name: loc.function.name.into(), + system_name: loc.function.system_name.into(), + filename: loc.function.filename.into(), + }, + address: loc.address, line: loc.line, }) .collect(); @@ -213,7 +253,7 @@ impl TrackedAlloc { .collect(); TrackedAlloc { - frames, + locations, labels, values: indices.build_values(sample.values), } @@ -225,13 +265,24 @@ impl TrackedAlloc { excluded_labels: &[Box], indices: &ValueIndices, ) -> anyhow::Result { - let frames = sample + let locations = sample .locations .iter() - .map(|loc| -> anyhow::Result { - Ok(OwnedFrame { - function_name: resolve_managed_string(storage, loc.function.name)?, - filename: resolve_managed_string(storage, loc.function.filename)?, + .map(|loc| -> anyhow::Result { + Ok(OwnedLocation { + mapping: OwnedMapping { + memory_start: loc.mapping.memory_start, + memory_limit: loc.mapping.memory_limit, + file_offset: loc.mapping.file_offset, + filename: resolve_managed_string(storage, loc.mapping.filename)?, + build_id: resolve_managed_string(storage, loc.mapping.build_id)?, + }, + function: OwnedFunction { + function_name: resolve_managed_string(storage, loc.function.name)?, + system_name: resolve_managed_string(storage, loc.function.system_name)?, + filename: resolve_managed_string(storage, loc.function.filename)?, + }, + address: loc.address, line: loc.line, }) }) @@ -252,7 +303,7 @@ impl TrackedAlloc { } Ok(TrackedAlloc { - frames, + locations, labels, values: indices.build_values(sample.values), }) diff --git a/libdd-profiling/src/internal/profile/fuzz_tests.rs b/libdd-profiling/src/internal/profile/fuzz_tests.rs index 43f7341b65..a134462254 100644 --- a/libdd-profiling/src/internal/profile/fuzz_tests.rs +++ b/libdd-profiling/src/internal/profile/fuzz_tests.rs @@ -430,7 +430,7 @@ fn fuzz_add_sample<'a>( samples_with_timestamps: &mut Vec<&'a Sample>, samples_without_timestamps: &mut HashMap<(&'a [Location], &'a [Label]), Vec>, ) { - let r = profile.try_add_sample(sample.into(), *timestamp, None); + let r = profile.try_add_sample(sample.into(), *timestamp); if expected_sample_types.len() == sample.values.len() { assert!(r.is_ok()); if timestamp.is_some() { diff --git a/libdd-profiling/src/internal/profile/mod.rs b/libdd-profiling/src/internal/profile/mod.rs index ffd54115f6..36b7789b52 100644 --- a/libdd-profiling/src/internal/profile/mod.rs +++ b/libdd-profiling/src/internal/profile/mod.rs @@ -15,7 +15,7 @@ use crate::api::ManagedStringId; use crate::collections::identifiable::*; use crate::collections::string_storage::{CachedProfileId, ManagedStringStorage}; use crate::collections::string_table::{self, StringTable}; -use crate::internal::heap_live::{HeapLiveState, OwnedFrame, OwnedLabel}; +use crate::internal::heap_live::{HeapLiveState, OwnedLabel, OwnedLocation}; use crate::iter::{IntoLendingIterator, LendingIterator}; use crate::profiles::collections::ArcOverflow; use crate::profiles::datatypes::ProfilesDictionary; @@ -127,20 +127,34 @@ impl Profile { Ok(()) } - /// Add a sample to the profile. If `track_ptr` is `Some(ptr)`, the sample - /// data is also copied into owned storage for heap-live tracking. The - /// tracked allocation will be automatically injected into the profile's - /// observations during `reset_and_return_previous()` and can be removed - /// earlier via `untrack_allocation()`. - /// - /// If heap-live tracking is at capacity, the tracking request is silently - /// discarded, the sample is still added to the profile normally. - /// If `ptr` matches an already-tracked allocation, the old entry is - /// replaced with the new sample data. + /// Add a sample to the profile. pub fn try_add_sample( &mut self, sample: api::Sample, timestamp: Option, + ) -> anyhow::Result<()> { + self.try_add_sample_impl(sample, timestamp, None) + } + + /// Add a sample to the profile and retain it as a tracked allocation for + /// heap-live profiling. + pub fn add_tracked_allocation( + &mut self, + sample: api::Sample, + timestamp: Option, + ptr: u64, + ) -> anyhow::Result<()> { + anyhow::ensure!( + self.heap_live.is_some(), + "heap-live tracking is not configured for this profile" + ); + self.try_add_sample_impl(sample, timestamp, Some(ptr)) + } + + fn try_add_sample_impl( + &mut self, + sample: api::Sample, + timestamp: Option, track_ptr: Option, ) -> anyhow::Result<()> { #[cfg(debug_assertions)] @@ -288,6 +302,27 @@ impl Profile { &mut self, sample: api::StringIdSample, timestamp: Option, + ) -> anyhow::Result<()> { + self.add_string_id_sample_impl(sample, timestamp, None) + } + + pub fn add_tracked_string_id_allocation( + &mut self, + sample: api::StringIdSample, + timestamp: Option, + ptr: u64, + ) -> anyhow::Result<()> { + anyhow::ensure!( + self.heap_live.is_some(), + "heap-live tracking is not configured for this profile" + ); + self.add_string_id_sample_impl(sample, timestamp, Some(ptr)) + } + + fn add_string_id_sample_impl( + &mut self, + sample: api::StringIdSample, + timestamp: Option, track_ptr: Option, ) -> anyhow::Result<()> { anyhow::ensure!( @@ -381,33 +416,26 @@ impl Profile { Ok(()) } - /// Enable heap-live allocation tracking on this profile. When enabled, - /// calls to `try_add_sample()` with `track_ptr: Some(ptr)` will copy - /// the sample's data into owned storage. Tracked allocations are - /// automatically injected into the profile during - /// `reset_and_return_previous()`, and the tracker moves to the new - /// active profile. + /// Configure heap-live allocation tracking on this profile. Tracked + /// allocations are automatically injected into the profile during + /// `reset_and_return_previous()`, and the tracker moves to the new active + /// profile. /// /// # Arguments /// * `max_tracked` - Maximum number of allocations to track simultaneously /// * `excluded_labels` - Label keys to strip when copying into the tracker (e.g., /// high-cardinality labels like "span id", "trace id") - pub fn enable_heap_live_tracking( + pub fn configure_heap_live( &mut self, max_tracked: usize, excluded_labels: &[&str], - alloc_size_idx: usize, - heap_live_samples_idx: usize, - heap_live_size_idx: usize, - ) { + ) -> anyhow::Result<()> { self.heap_live = Some(HeapLiveState::new( max_tracked, excluded_labels, - alloc_size_idx, - heap_live_samples_idx, - heap_live_size_idx, - self.sample_types.len(), - )); + &self.sample_types, + )?); + Ok(()) } /// Remove a tracked heap-live allocation by pointer. No-op if heap-live @@ -921,9 +949,9 @@ impl Profile { let result = (|| { for alloc in hl.tracked.values() { let locations: Vec> = alloc - .frames + .locations .iter() - .map(OwnedFrame::as_api_location) + .map(OwnedLocation::as_api_location) .collect(); let labels: Vec> = alloc.labels.iter().map(OwnedLabel::as_api_label).collect(); @@ -932,8 +960,7 @@ impl Profile { values: &alloc.values, labels, }; - // Pass None for track_ptr to avoid re-tracking during injection - self.try_add_sample(sample, None, None)?; + self.try_add_sample(sample, None)?; } anyhow::Ok(()) })(); @@ -1180,13 +1207,6 @@ mod heap_live_tests { use super::*; use crate::pprof::test_utils::{roundtrip_to_pprof, sorted_samples, string_table_fetch}; - // Sample types matching the dd-trace-php layout when allocation + heap-live - // profiling are enabled (without cpu-time). - // Indices: 0=alloc-samples, 1=alloc-size, 2=heap-live-samples, 3=heap-live-size - const ALLOC_SIZE_IDX: usize = 1; - const HEAP_LIVE_SAMPLES_IDX: usize = 2; - const HEAP_LIVE_SIZE_IDX: usize = 3; - fn heap_live_sample_types() -> [api::SampleType; 4] { [ api::SampleType::AllocSamples, @@ -1198,8 +1218,11 @@ mod heap_live_tests { fn make_mapping() -> api::Mapping<'static> { api::Mapping { + memory_start: 0x1000, + memory_limit: 0x2000, + file_offset: 0x80, filename: "php", - ..Default::default() + build_id: "build-id", } } @@ -1220,13 +1243,7 @@ mod heap_live_tests { fn heap_live_tracked_sample_appears_in_serialized_pprof() { let sample_types = heap_live_sample_types(); let mut profile = Profile::new(&sample_types, None); - profile.enable_heap_live_tracking( - 4096, - &[], - ALLOC_SIZE_IDX, - HEAP_LIVE_SAMPLES_IDX, - HEAP_LIVE_SIZE_IDX, - ); + profile.configure_heap_live(4096, &[]).unwrap(); let mapping = make_mapping(); let locations = make_locations(mapping); @@ -1244,7 +1261,7 @@ mod heap_live_tests { }; profile - .try_add_sample(sample, None, Some(0x1000)) + .add_tracked_allocation(sample, None, 0x1000) .expect("add to succeed"); let old = profile @@ -1252,18 +1269,29 @@ mod heap_live_tests { .expect("reset to succeed"); let pprof = roundtrip_to_pprof(old).unwrap(); - // Original sample + 1 injected heap-live copy - assert_eq!(pprof.samples.len(), 2); + // The injected heap-live values aggregate into the original sample + // because labels and stacktrace match. + assert_eq!(pprof.samples.len(), 1); let samples = sorted_samples(&pprof); let mut found_values: Vec> = samples.iter().map(|s| s.values.clone()).collect(); found_values.sort(); - // Original: [1, 1024, 0, 0], Injected: [0, 0, 1, 1024] - assert_eq!(found_values, vec![vec![0, 0, 1, 1024], vec![1, 1024, 0, 0]]); + assert_eq!(found_values, vec![vec![1, 1024, 1, 1024]]); // Verify function name appears assert!(pprof.string_table.iter().any(|s| s == "malloc")); assert!(pprof.string_table.iter().any(|s| s == "test.php")); + assert!(pprof.string_table.iter().any(|s| s == "php")); + assert!(pprof.string_table.iter().any(|s| s == "build-id")); + + let sample = &samples[0]; + let location = &pprof.locations[(sample.location_ids[0] - 1) as usize]; + let function = &pprof.functions[(location.lines[0].function_id - 1) as usize]; + let mapping = &pprof.mappings[(location.mapping_id - 1) as usize]; + assert_eq!(location.address, 0); + assert_eq!(string_table_fetch(&pprof, function.system_name), "malloc"); + assert_eq!(string_table_fetch(&pprof, mapping.filename), "php"); + assert_eq!(string_table_fetch(&pprof, mapping.build_id), "build-id"); // Verify label on all samples for s in &samples { @@ -1278,13 +1306,7 @@ mod heap_live_tests { fn heap_live_untracked_sample_does_not_appear() { let sample_types = heap_live_sample_types(); let mut profile = Profile::new(&sample_types, None); - profile.enable_heap_live_tracking( - 4096, - &[], - ALLOC_SIZE_IDX, - HEAP_LIVE_SAMPLES_IDX, - HEAP_LIVE_SIZE_IDX, - ); + profile.configure_heap_live(4096, &[]).unwrap(); let mapping = make_mapping(); let locations = make_locations(mapping); @@ -1296,7 +1318,7 @@ mod heap_live_tests { }; profile - .try_add_sample(sample, None, Some(0x2000)) + .add_tracked_allocation(sample, None, 0x2000) .expect("add to succeed"); // Untrack before reset — heap-live injection should skip this ptr @@ -1315,13 +1337,7 @@ mod heap_live_tests { fn heap_live_survives_reset() { let sample_types = heap_live_sample_types(); let mut profile = Profile::new(&sample_types, None); - profile.enable_heap_live_tracking( - 4096, - &[], - ALLOC_SIZE_IDX, - HEAP_LIVE_SAMPLES_IDX, - HEAP_LIVE_SIZE_IDX, - ); + profile.configure_heap_live(4096, &[]).unwrap(); let mapping = make_mapping(); let locations = make_locations(mapping); @@ -1334,7 +1350,7 @@ mod heap_live_tests { }; profile - .try_add_sample(sample, None, Some(0x3000)) + .add_tracked_allocation(sample, None, 0x3000) .expect("add to succeed"); // First reset: original + heap-live copy @@ -1343,7 +1359,7 @@ mod heap_live_tests { .expect("first reset to succeed"); let pprof1 = roundtrip_to_pprof(old1).unwrap(); - assert_eq!(pprof1.samples.len(), 2); + assert_eq!(pprof1.samples.len(), 1); // Add a new regular sample to the reset profile let regular_sample = api::Sample { @@ -1353,7 +1369,7 @@ mod heap_live_tests { }; profile - .try_add_sample(regular_sample, None, None) + .try_add_sample(regular_sample, None) .expect("add regular to succeed"); // Second reset: the still-tracked heap-live + the new regular @@ -1362,26 +1378,21 @@ mod heap_live_tests { .expect("second reset to succeed"); let pprof2 = roundtrip_to_pprof(old2).unwrap(); - assert_eq!(pprof2.samples.len(), 2); + assert_eq!(pprof2.samples.len(), 1); let samples = sorted_samples(&pprof2); let mut found_values: Vec> = samples.iter().map(|s| s.values.clone()).collect(); found_values.sort(); - // Injected heap-live: [0, 0, 1, 2048], Regular: [5, 9999, 0, 0] - assert_eq!(found_values, vec![vec![0, 0, 1, 2048], vec![5, 9999, 0, 0]]); + assert_eq!(found_values, vec![vec![5, 9999, 1, 2048]]); } #[test] fn heap_live_excluded_labels_are_filtered() { let sample_types = heap_live_sample_types(); let mut profile = Profile::new(&sample_types, None); - profile.enable_heap_live_tracking( - 4096, - &["span id", "trace id"], - ALLOC_SIZE_IDX, - HEAP_LIVE_SAMPLES_IDX, - HEAP_LIVE_SIZE_IDX, - ); + profile + .configure_heap_live(4096, &["span id", "trace id"]) + .unwrap(); let mapping = make_mapping(); let locations = make_locations(mapping); @@ -1411,7 +1422,7 @@ mod heap_live_tests { }; profile - .try_add_sample(sample, None, Some(0x4000)) + .add_tracked_allocation(sample, None, 0x4000) .expect("add to succeed"); let old = profile @@ -1447,13 +1458,7 @@ mod heap_live_tests { let sample_types = heap_live_sample_types(); let mut profile = Profile::new(&sample_types, None); // Only allow 2 tracked allocations - profile.enable_heap_live_tracking( - 2, - &[], - ALLOC_SIZE_IDX, - HEAP_LIVE_SAMPLES_IDX, - HEAP_LIVE_SIZE_IDX, - ); + profile.configure_heap_live(2, &[]).unwrap(); let mapping = make_mapping(); @@ -1478,7 +1483,7 @@ mod heap_live_tests { }; profile - .try_add_sample(sample, None, Some(*ptr)) + .add_tracked_allocation(sample, None, *ptr) .expect("add to succeed"); } @@ -1487,15 +1492,15 @@ mod heap_live_tests { .expect("reset to succeed"); let pprof = roundtrip_to_pprof(old).unwrap(); - // 3 original samples + 2 heap-live injected (third rejected at capacity) - assert_eq!(pprof.samples.len(), 5); + // The first two tracked samples aggregate with their injected + // heap-live values. The third sample remains untracked. + assert_eq!(pprof.samples.len(), 3); } #[test] - fn heap_live_disabled_by_default() { + fn tracked_allocation_requires_configuration() { let sample_types = heap_live_sample_types(); let mut profile = Profile::new(&sample_types, None); - // Do NOT call enable_heap_live_tracking let mapping = make_mapping(); let locations = make_locations(mapping); @@ -1506,30 +1511,32 @@ mod heap_live_tests { labels: vec![], }; - profile - .try_add_sample(sample, None, Some(0x5000)) - .expect("add to succeed"); + let err = profile + .add_tracked_allocation(sample, None, 0x5000) + .expect_err("tracked allocations should require explicit configuration"); + assert!(err + .to_string() + .contains("heap-live tracking is not configured for this profile")); + } - let old = profile - .reset_and_return_previous() - .expect("reset to succeed"); + #[test] + fn heap_live_requires_expected_sample_types() { + let sample_types = [api::SampleType::AllocSamples, api::SampleType::AllocSize]; + let mut profile = Profile::new(&sample_types, None); - let pprof = roundtrip_to_pprof(old).unwrap(); - // Only the original sample — track_ptr silently ignored - assert_eq!(pprof.samples.len(), 1); + let err = profile + .configure_heap_live(4096, &[]) + .expect_err("heap-live configuration should validate sample types"); + assert!(err + .to_string() + .contains("heap-live tracking requires sample type HeapLiveSamples")); } #[test] fn heap_live_mixed_tracked_and_regular_samples() { let sample_types = heap_live_sample_types(); let mut profile = Profile::new(&sample_types, None); - profile.enable_heap_live_tracking( - 4096, - &[], - ALLOC_SIZE_IDX, - HEAP_LIVE_SAMPLES_IDX, - HEAP_LIVE_SIZE_IDX, - ); + profile.configure_heap_live(4096, &[]).unwrap(); let mapping = make_mapping(); @@ -1552,7 +1559,7 @@ mod heap_live_tests { }; profile - .try_add_sample(tracked_sample, None, Some(0x6000)) + .add_tracked_allocation(tracked_sample, None, 0x6000) .expect("add tracked to succeed"); // Regular sample (no tracking) @@ -1574,7 +1581,7 @@ mod heap_live_tests { }; profile - .try_add_sample(regular_sample, None, None) + .try_add_sample(regular_sample, None) .expect("add regular to succeed"); let old = profile @@ -1582,8 +1589,8 @@ mod heap_live_tests { .expect("reset to succeed"); let pprof = roundtrip_to_pprof(old).unwrap(); - // 2 original + 1 injected heap-live copy of the tracked sample - assert_eq!(pprof.samples.len(), 3); + // The tracked sample aggregates with its injected heap-live copy. + assert_eq!(pprof.samples.len(), 2); // Verify we have both function names assert!(pprof.string_table.iter().any(|s| s == "tracked_fn")); @@ -1597,7 +1604,7 @@ mod heap_live_tests { // Regular: [3, 7777, 0, 0] assert_eq!( found_values, - vec![vec![0, 0, 1, 512], vec![1, 512, 0, 0], vec![3, 7777, 0, 0]] + vec![vec![1, 512, 1, 512], vec![3, 7777, 0, 0]] ); } } @@ -1667,7 +1674,6 @@ mod api_tests { labels: vec![], }, None, - None, ) .expect("add to succeed"); @@ -1739,18 +1745,18 @@ mod api_tests { assert_eq!(profile.only_for_testing_num_aggregated_samples(), 0); profile - .try_add_sample(main_sample, None, None) + .try_add_sample(main_sample, None) .expect("profile to not be full"); assert_eq!(profile.only_for_testing_num_aggregated_samples(), 1); profile - .try_add_sample(test_sample, None, None) + .try_add_sample(test_sample, None) .expect("profile to not be full"); assert_eq!(profile.only_for_testing_num_aggregated_samples(), 2); assert_eq!(profile.only_for_testing_num_timestamped_samples(), 0); profile - .try_add_sample(timestamp_sample, Timestamp::new(42), None) + .try_add_sample(timestamp_sample, Timestamp::new(42)) .expect("profile to not be full"); assert_eq!(profile.only_for_testing_num_timestamped_samples(), 1); profile @@ -1920,7 +1926,7 @@ mod api_tests { labels: vec![id_label], }; - assert!(profile.try_add_sample(sample, None, None).is_err()); + assert!(profile.try_add_sample(sample, None).is_err()); } #[test] @@ -1963,11 +1969,11 @@ mod api_tests { }; profile - .try_add_sample(sample1, None, None) + .try_add_sample(sample1, None) .expect("add to success"); profile - .try_add_sample(sample2, None, None) + .try_add_sample(sample2, None) .expect("add to success"); profile.add_endpoint(10, Cow::from("my endpoint"))?; @@ -2083,7 +2089,7 @@ mod api_tests { labels, }; - profile.try_add_sample(sample, None, None).unwrap_err(); + profile.try_add_sample(sample, None).unwrap_err(); } #[test] @@ -2106,7 +2112,7 @@ mod api_tests { }; profile - .try_add_sample(sample1, None, None) + .try_add_sample(sample1, None) .expect("add to success"); let serialized_profile = roundtrip_to_pprof(profile).unwrap(); @@ -2148,7 +2154,7 @@ mod api_tests { }; profile - .try_add_sample(sample1, None, None) + .try_add_sample(sample1, None) .expect("add to success"); let upscaling_info = UpscalingInfo::Proportional { scale: 2.0 }; @@ -2178,7 +2184,7 @@ mod api_tests { }; profile - .try_add_sample(sample1, None, None) + .try_add_sample(sample1, None) .expect("add to success"); let upscaling_info = UpscalingInfo::Proportional { scale: 2.7 }; @@ -2208,7 +2214,7 @@ mod api_tests { }; profile - .try_add_sample(sample1, None, None) + .try_add_sample(sample1, None) .expect("add to success"); let upscaling_info = UpscalingInfo::Poisson { @@ -2242,7 +2248,7 @@ mod api_tests { }; profile - .try_add_sample(sample1, None, None) + .try_add_sample(sample1, None) .expect("add to success"); let upscaling_info = UpscalingInfo::PoissonNonSampleTypeCount { @@ -2276,7 +2282,7 @@ mod api_tests { }; profile - .try_add_sample(sample1, None, None) + .try_add_sample(sample1, None) .expect("add to success"); let upscaling_info = UpscalingInfo::Poisson { @@ -2310,7 +2316,7 @@ mod api_tests { }; profile - .try_add_sample(sample1, None, None) + .try_add_sample(sample1, None) .expect("add to success"); // invalid sampling_distance value @@ -2381,10 +2387,10 @@ mod api_tests { }; profile - .try_add_sample(sample1, None, None) + .try_add_sample(sample1, None) .expect("add to success"); profile - .try_add_sample(sample2, None, None) + .try_add_sample(sample2, None) .expect("add to success"); // upscale the first value and the last one @@ -2440,10 +2446,10 @@ mod api_tests { }; profile - .try_add_sample(sample1, None, None) + .try_add_sample(sample1, None) .expect("add to success"); profile - .try_add_sample(sample2, None, None) + .try_add_sample(sample2, None) .expect("add to success"); let mut values_offset: Vec = vec![0]; @@ -2489,7 +2495,7 @@ mod api_tests { }; profile - .try_add_sample(sample1, None, None) + .try_add_sample(sample1, None) .expect("add to success"); let values_offset: Vec = vec![0]; @@ -2547,7 +2553,7 @@ mod api_tests { }; profile - .try_add_sample(sample1, None, None) + .try_add_sample(sample1, None) .expect("add to success"); let upscaling_info = UpscalingInfo::Proportional { scale: 2.0 }; @@ -2605,10 +2611,10 @@ mod api_tests { }; profile - .try_add_sample(sample1, None, None) + .try_add_sample(sample1, None) .expect("add to success"); profile - .try_add_sample(sample2, None, None) + .try_add_sample(sample2, None) .expect("add to success"); let upscaling_info = UpscalingInfo::Proportional { scale: 2.0 }; @@ -2679,10 +2685,10 @@ mod api_tests { }; profile - .try_add_sample(sample1, None, None) + .try_add_sample(sample1, None) .expect("add to success"); profile - .try_add_sample(sample2, None, None) + .try_add_sample(sample2, None) .expect("add to success"); // add rule for the first sample on the 1st value @@ -2736,7 +2742,7 @@ mod api_tests { }; profile - .try_add_sample(sample1, None, None) + .try_add_sample(sample1, None) .expect("add to success"); // upscale samples and wall-time values @@ -2774,7 +2780,7 @@ mod api_tests { }; profile - .try_add_sample(sample1, None, None) + .try_add_sample(sample1, None) .expect("add to success"); let upscaling_info = UpscalingInfo::Proportional { scale: 2.0 }; @@ -3048,7 +3054,7 @@ mod api_tests { }; profile - .try_add_sample(sample1, None, None) + .try_add_sample(sample1, None) .expect("add to success"); let mut value_offsets: Vec = vec![0, 1]; @@ -3114,10 +3120,10 @@ mod api_tests { }; profile - .try_add_sample(sample1, None, None) + .try_add_sample(sample1, None) .expect("add to success"); profile - .try_add_sample(sample2, None, None) + .try_add_sample(sample2, None) .expect("add to success"); profile.add_endpoint(10, Cow::from("endpoint 10"))?; @@ -3416,12 +3422,8 @@ mod api_tests { labels: vec![], }; - profile - .add_string_id_sample(sample.clone(), None, None) - .unwrap(); - profile - .add_string_id_sample(sample.clone(), None, None) - .unwrap(); + profile.add_string_id_sample(sample.clone(), None).unwrap(); + profile.add_string_id_sample(sample.clone(), None).unwrap(); let pprof_first_profile = roundtrip_to_pprof(profile.reset_and_return_previous().unwrap()).unwrap(); @@ -3438,12 +3440,8 @@ mod api_tests { // If the cache invalidation on the managed string table is working correctly, these strings // get correctly re-added to the profile's string table - profile - .add_string_id_sample(sample.clone(), None, None) - .unwrap(); - profile - .add_string_id_sample(sample.clone(), None, None) - .unwrap(); + profile.add_string_id_sample(sample.clone(), None).unwrap(); + profile.add_string_id_sample(sample.clone(), None).unwrap(); let pprof_second_profile = roundtrip_to_pprof(profile).unwrap(); assert!(pprof_second_profile @@ -3488,7 +3486,7 @@ mod api_tests { }; profile - .try_add_sample(sample, None, None) // No timestamp = aggregated + .try_add_sample(sample, None) // No timestamp = aggregated .expect("profile to not be full"); // We want to trigger an error inside the loop over observations.