diff --git a/libdd-profiling-ffi/src/profiles/datatypes.rs b/libdd-profiling-ffi/src/profiles/datatypes.rs index c7ec589d5c..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; @@ -572,6 +579,41 @@ pub unsafe extern "C" fn ddog_prof_Profile_add( .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 @@ -896,6 +938,53 @@ pub unsafe extern "C" fn ddog_prof_Profile_reset(profile: *mut Profile) -> Profi .into() } +/// 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. +/// * `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. `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_configure_heap_live( + profile: *mut Profile, + config: HeapLiveConfig, +) -> ProfileResult { + (|| { + let profile = profile_ptr_to_inner(profile)?; + let labels: Vec<&str> = config + .excluded_label_keys + .iter() + .map(|cs| cs.try_to_utf8()) + .collect::, _>>()?; + profile.configure_heap_live(config.max_tracked_allocations, &labels) + })() + .context("ddog_prof_Profile_configure_heap_live 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::*; @@ -1094,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/src/internal/heap_live.rs b/libdd-profiling/src/internal/heap_live.rs new file mode 100644 index 0000000000..31c1afe1db --- /dev/null +++ b/libdd-profiling/src/internal/heap_live.rs @@ -0,0 +1,322 @@ +// Copyright 2021-Present Datadog, Inc. https://www.datadoghq.com/ +// SPDX-License-Identifier: Apache-2.0 + +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 location data for heap-live tracking. +/// Stores copies of borrowed strings so tracked allocations survive across +/// profile resets. +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, +} + +/// 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 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 { + mapping: self.mapping.as_api_mapping(), + function: self.function.as_api_function(), + address: self.address, + line: self.line, + } + } +} + +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 { + 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 { + 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 locations: Vec, + pub labels: Vec, + pub values: Vec, +} + +impl HeapLiveState { + /// Create a new heap-live tracker. + pub fn new( + max_tracked: usize, + excluded_labels: &[&str], + 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::for_sample_types(sample_types)?, + }) + } + + /// 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 locations = sample + .locations + .iter() + .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(); + + 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 { + locations, + 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 locations = sample + .locations + .iter() + .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, + }) + }) + .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 { + locations, + 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/mod.rs b/libdd-profiling/src/internal/profile/mod.rs index 280cebcd7a..36b7789b52 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, OwnedLabel, OwnedLocation}; 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,35 @@ impl Profile { Ok(()) } + /// 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)] { @@ -162,7 +192,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 +302,28 @@ 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!( self.string_storage.is_some(), @@ -293,7 +356,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 +416,36 @@ impl Profile { Ok(()) } + /// 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 configure_heap_live( + &mut self, + max_tracked: usize, + excluded_labels: &[&str], + ) -> anyhow::Result<()> { + self.heap_live = Some(HeapLiveState::new( + max_tracked, + excluded_labels, + &self.sample_types, + )?); + Ok(()) + } + + /// 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 +536,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 +563,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 +939,35 @@ 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 + .locations + .iter() + .map(OwnedLocation::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, + }; + self.try_add_sample(sample, 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 +1085,7 @@ impl Profile { * CachedProfileId for why. */ timestamp_key: Default::default(), upscaling_rules: Default::default(), + heap_live: None, }; let _id = profile.intern(""); @@ -1055,6 +1202,413 @@ impl Profile { } } +#[cfg(test)] +mod heap_live_tests { + use super::*; + use crate::pprof::test_utils::{roundtrip_to_pprof, sorted_samples, string_table_fetch}; + + 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 { + memory_start: 0x1000, + memory_limit: 0x2000, + file_offset: 0x80, + filename: "php", + build_id: "build-id", + } + } + + 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.configure_heap_live(4096, &[]).unwrap(); + + 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 + .add_tracked_allocation(sample, None, 0x1000) + .expect("add to succeed"); + + let old = profile + .reset_and_return_previous() + .expect("reset to succeed"); + + let pprof = roundtrip_to_pprof(old).unwrap(); + // 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(); + 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 { + 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.configure_heap_live(4096, &[]).unwrap(); + + let mapping = make_mapping(); + let locations = make_locations(mapping); + + let sample = api::Sample { + locations, + values: &[1, 512, 0, 0], + labels: vec![], + }; + + profile + .add_tracked_allocation(sample, None, 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.configure_heap_live(4096, &[]).unwrap(); + + 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 + .add_tracked_allocation(sample, None, 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(), 1); + + // 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) + .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(), 1); + + let samples = sorted_samples(&pprof2); + let mut found_values: Vec> = samples.iter().map(|s| s.values.clone()).collect(); + found_values.sort(); + 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 + .configure_heap_live(4096, &["span id", "trace id"]) + .unwrap(); + + 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 + .add_tracked_allocation(sample, None, 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.configure_heap_live(2, &[]).unwrap(); + + 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 + .add_tracked_allocation(sample, None, *ptr) + .expect("add to succeed"); + } + + let old = profile + .reset_and_return_previous() + .expect("reset to succeed"); + + let pprof = roundtrip_to_pprof(old).unwrap(); + // 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 tracked_allocation_requires_configuration() { + let sample_types = heap_live_sample_types(); + let mut profile = Profile::new(&sample_types, None); + + let mapping = make_mapping(); + let locations = make_locations(mapping); + + let sample = api::Sample { + locations, + values: &[1, 4096, 0, 0], + labels: vec![], + }; + + 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")); + } + + #[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 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.configure_heap_live(4096, &[]).unwrap(); + + 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 + .add_tracked_allocation(tracked_sample, None, 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) + .expect("add regular to succeed"); + + let old = profile + .reset_and_return_previous() + .expect("reset to succeed"); + + let pprof = roundtrip_to_pprof(old).unwrap(); + // 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")); + 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![1, 512, 1, 512], vec![3, 7777, 0, 0]] + ); + } +} + #[cfg(test)] mod api_tests { use super::*;