From 725b26fc465d5fc9144d8d19bb885397cda0d05b Mon Sep 17 00:00:00 2001 From: Daniel Schwartz-Narbonne Date: Thu, 29 May 2025 16:40:41 -0400 Subject: [PATCH 01/52] [profiling] Add a channel based profiling manager --- Cargo.lock | 1 + datadog-profiling-ffi/Cargo.toml | 1 + datadog-profiling-ffi/src/lib.rs | 3 + datadog-profiling-ffi/src/manager/mod.rs | 110 ++++++++++++++++++ datadog-profiling-ffi/src/manager/uploader.rs | 2 + .../src/profiles/datatypes.rs | 15 +++ datadog-profiling-ffi/src/profiles/mod.rs | 2 +- 7 files changed, 133 insertions(+), 1 deletion(-) create mode 100644 datadog-profiling-ffi/src/manager/mod.rs create mode 100644 datadog-profiling-ffi/src/manager/uploader.rs diff --git a/Cargo.lock b/Cargo.lock index dc8eb745f1..e5aa2214fb 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1600,6 +1600,7 @@ version = "18.1.0" dependencies = [ "anyhow", "build_common", + "crossbeam-channel", "data-pipeline-ffi", "datadog-crashtracker-ffi", "datadog-library-config-ffi", diff --git a/datadog-profiling-ffi/Cargo.toml b/datadog-profiling-ffi/Cargo.toml index dc665df1ae..a6a1ee6225 100644 --- a/datadog-profiling-ffi/Cargo.toml +++ b/datadog-profiling-ffi/Cargo.toml @@ -34,6 +34,7 @@ build_common = { path = "../build-common" } [dependencies] anyhow = "1.0" +crossbeam-channel = "0.5.15" data-pipeline-ffi = { path = "../data-pipeline-ffi", default-features = false, optional = true } datadog-crashtracker-ffi = { path = "../datadog-crashtracker-ffi", default-features = false, optional = true} datadog-library-config-ffi = { path = "../datadog-library-config-ffi", default-features = false, optional = true } diff --git a/datadog-profiling-ffi/src/lib.rs b/datadog-profiling-ffi/src/lib.rs index 17b9a70c50..fa45424843 100644 --- a/datadog-profiling-ffi/src/lib.rs +++ b/datadog-profiling-ffi/src/lib.rs @@ -11,6 +11,7 @@ pub use symbolizer_ffi::*; mod exporter; +mod manager; mod profiles; mod string_storage; @@ -33,3 +34,5 @@ pub use datadog_library_config_ffi::*; // re-export tracer metadata functions #[cfg(feature = "ddcommon-ffi")] pub use ddcommon_ffi::*; + +pub use manager::*; diff --git a/datadog-profiling-ffi/src/manager/mod.rs b/datadog-profiling-ffi/src/manager/mod.rs new file mode 100644 index 0000000000..8803fd2780 --- /dev/null +++ b/datadog-profiling-ffi/src/manager/mod.rs @@ -0,0 +1,110 @@ +// Copyright 2025-Present Datadog, Inc. https://www.datadoghq.com/ +// SPDX-License-Identifier: Apache-2.0 +#![allow(unused_variables)] +#![allow(dead_code)] +#![allow(clippy::todo)] + +pub mod uploader; + +use std::{ + ffi::c_void, + time::{Duration, Instant}, +}; + +use crate::profiles::datatypes::Profile; +use crossbeam_channel::{select, tick, Receiver, Sender}; +use datadog_profiling::{api, internal}; + +#[repr(C)] +pub struct ManagedSample { + data: *mut c_void, + callback: extern "C" fn(*mut c_void, *mut Profile), +} + +impl ManagedSample { + /// Creates a new ManagedSample from a raw pointer and callback. + /// + /// # Safety + /// + /// The caller must ensure that: + /// - The pointer is valid and points to memory that will remain valid for the lifetime of the + /// ManagedSample + /// - The pointer is properly aligned + /// - The memory is properly initialized + #[no_mangle] + pub unsafe extern "C" fn new( + data: *mut c_void, + callback: extern "C" fn(*mut c_void, *mut Profile), + ) -> Self { + Self { data, callback } + } + + /// Returns the raw pointer to the underlying data. + #[no_mangle] + pub extern "C" fn as_ptr(&self) -> *mut c_void { + self.data + } + + #[allow(clippy::todo)] + pub fn add(self, profile: &mut Profile) { + (self.callback)(self.data, profile); + } +} + +pub struct ProfilerManager { + samples_receiver: Receiver, + samples_sender: Sender, + cpu_ticker: Receiver, + upload_ticker: Receiver, + active: internal::Profile, + standby: internal::Profile, +} + +impl ProfilerManager { + pub fn new(sample_types: &[api::ValueType], period: Option) -> Self { + let (samples_sender, samples_receiver) = crossbeam_channel::bounded(10); + let active = internal::Profile::new(sample_types, period); + let standby = internal::Profile::new(sample_types, period); + let cpu_ticker = tick(Duration::from_millis(100)); + let upload_ticker = tick(Duration::from_secs(1)); + Self { + cpu_ticker, + upload_ticker, + samples_receiver, + samples_sender, + active, + standby, + } + } + + #[allow(clippy::todo)] + pub fn main(&mut self) -> anyhow::Result<()> { + // This is just here to allow us to easily bail out. + let done = tick(Duration::from_secs(5)); + loop { + select! { + recv(self.samples_receiver) -> sample => { + let mut ffi_profile = unsafe { Profile::from_pointer(&mut self.active as *mut _) }; + sample?.add(&mut ffi_profile); + }, + recv(self.cpu_ticker) -> msg => println!("{msg:?} call echion"), + recv(self.upload_ticker) -> msg => println!("{msg:?} swap and upload the old one"), + recv(done) -> msg => return Ok(()), + } + } + } +} + +#[cfg(test)] +mod tests { + use crate::manager::ProfilerManager; + + #[test] + fn test_the_thing() { + let sample_types = []; + let period = None; + let mut profile_manager = ProfilerManager::new(&sample_types, period); + println!("start"); + profile_manager.main().unwrap(); + } +} diff --git a/datadog-profiling-ffi/src/manager/uploader.rs b/datadog-profiling-ffi/src/manager/uploader.rs new file mode 100644 index 0000000000..5a97e72498 --- /dev/null +++ b/datadog-profiling-ffi/src/manager/uploader.rs @@ -0,0 +1,2 @@ +// Copyright 2025-Present Datadog, Inc. https://www.datadoghq.com/ +// SPDX-License-Identifier: Apache-2.0 diff --git a/datadog-profiling-ffi/src/profiles/datatypes.rs b/datadog-profiling-ffi/src/profiles/datatypes.rs index 4956f1a34c..689aaf6a9a 100644 --- a/datadog-profiling-ffi/src/profiles/datatypes.rs +++ b/datadog-profiling-ffi/src/profiles/datatypes.rs @@ -28,6 +28,21 @@ impl Profile { } } + /// Creates a Profile from a raw pointer to an internal::Profile. + /// THIS IS UNSAFE, DO NOT USE IT. + /// IT ONLY EXISTS TO ALLOW US TO CREATE A PROFILE FROM A RAW POINTER FOR RAPID PROTOTYPING FOR + /// THE MANAGER. + /// + /// # Safety + /// + /// The caller must ensure that: + /// - The pointer is valid and points to a valid internal::Profile + /// - The pointer is properly aligned + /// - The memory is properly initialized + pub unsafe fn from_pointer(ptr: *mut internal::Profile) -> Self { + Profile { inner: ptr } + } + fn take(&mut self) -> Option> { // Leaving a null will help with double-free issues that can // arise in C. Of course, it's best to never get there in the diff --git a/datadog-profiling-ffi/src/profiles/mod.rs b/datadog-profiling-ffi/src/profiles/mod.rs index 86136a6fbf..8df9a23029 100644 --- a/datadog-profiling-ffi/src/profiles/mod.rs +++ b/datadog-profiling-ffi/src/profiles/mod.rs @@ -1,5 +1,5 @@ // Copyright 2021-Present Datadog, Inc. https://www.datadoghq.com/ // SPDX-License-Identifier: Apache-2.0 -mod datatypes; +pub(crate) mod datatypes; mod interning_api; From aba6e6a1c104a307288205ab0d4c6409ff8f5183 Mon Sep 17 00:00:00 2001 From: Daniel Schwartz-Narbonne Date: Fri, 30 May 2025 16:17:27 -0400 Subject: [PATCH 02/52] more callbacks --- datadog-profiling-ffi/src/manager/mod.rs | 107 +++++++++++++++--- .../src/profiles/datatypes.rs | 15 --- 2 files changed, 90 insertions(+), 32 deletions(-) diff --git a/datadog-profiling-ffi/src/manager/mod.rs b/datadog-profiling-ffi/src/manager/mod.rs index 8803fd2780..225c33ff59 100644 --- a/datadog-profiling-ffi/src/manager/mod.rs +++ b/datadog-profiling-ffi/src/manager/mod.rs @@ -8,19 +8,27 @@ pub mod uploader; use std::{ ffi::c_void, + num::NonZeroI64, time::{Duration, Instant}, }; -use crate::profiles::datatypes::Profile; +use crate::profiles::datatypes::{ProfileResult, Sample}; +use anyhow::Context; use crossbeam_channel::{select, tick, Receiver, Sender}; use datadog_profiling::{api, internal}; #[repr(C)] pub struct ManagedSample { data: *mut c_void, - callback: extern "C" fn(*mut c_void, *mut Profile), + callback: extern "C" fn(*mut c_void, *mut internal::Profile), } +// void add_sample(void *data, Profile *profile) { +// Sample *sample = (Sample *)data; +// Profile *profile = (Profile *)data; +// profile->add_sample(sample); +// } + impl ManagedSample { /// Creates a new ManagedSample from a raw pointer and callback. /// @@ -34,7 +42,7 @@ impl ManagedSample { #[no_mangle] pub unsafe extern "C" fn new( data: *mut c_void, - callback: extern "C" fn(*mut c_void, *mut Profile), + callback: extern "C" fn(*mut c_void, *mut internal::Profile), ) -> Self { Self { data, callback } } @@ -46,7 +54,7 @@ impl ManagedSample { } #[allow(clippy::todo)] - pub fn add(self, profile: &mut Profile) { + pub fn add(self, profile: &mut internal::Profile) { (self.callback)(self.data, profile); } } @@ -56,15 +64,22 @@ pub struct ProfilerManager { samples_sender: Sender, cpu_ticker: Receiver, upload_ticker: Receiver, - active: internal::Profile, - standby: internal::Profile, + profile: internal::Profile, + cpu_sampler_callback: extern "C" fn(*mut internal::Profile), + sample_converter: extern "C" fn(ManagedSample) -> Sample<'static>, + reset_callback: extern "C" fn(ManagedSample) -> bool, } impl ProfilerManager { - pub fn new(sample_types: &[api::ValueType], period: Option) -> Self { + pub fn new( + sample_types: &[api::ValueType], + period: Option, + cpu_sampler_callback: extern "C" fn(*mut internal::Profile), + sample_converter: extern "C" fn(ManagedSample) -> Sample<'static>, + reset_callback: extern "C" fn(ManagedSample) -> bool, + ) -> Self { let (samples_sender, samples_receiver) = crossbeam_channel::bounded(10); - let active = internal::Profile::new(sample_types, period); - let standby = internal::Profile::new(sample_types, period); + let profile = internal::Profile::new(sample_types, period); let cpu_ticker = tick(Duration::from_millis(100)); let upload_ticker = tick(Duration::from_secs(1)); Self { @@ -72,8 +87,10 @@ impl ProfilerManager { upload_ticker, samples_receiver, samples_sender, - active, - standby, + profile, + cpu_sampler_callback, + sample_converter, + reset_callback, } } @@ -84,26 +101,82 @@ impl ProfilerManager { loop { select! { recv(self.samples_receiver) -> sample => { - let mut ffi_profile = unsafe { Profile::from_pointer(&mut self.active as *mut _) }; - sample?.add(&mut ffi_profile); + sample?.add(&mut self.profile); + }, + recv(self.cpu_ticker) -> msg => { + (self.cpu_sampler_callback)(&mut self.profile); + }, + recv(self.upload_ticker) -> msg => { + let old_profile = self.profile.reset_and_return_previous()?; + std::thread::spawn(move || { + if let Ok(encoded) = old_profile.serialize_into_compressed_pprof(None, None) { + println!("Successfully serialized profile"); + } + }); }, - recv(self.cpu_ticker) -> msg => println!("{msg:?} call echion"), - recv(self.upload_ticker) -> msg => println!("{msg:?} swap and upload the old one"), recv(done) -> msg => return Ok(()), } } } } +/// # Safety +/// The `profile` ptr must point to a valid internal::Profile object. +/// All pointers inside the `sample` need to be valid for the duration of this call. +/// This call is _NOT_ thread-safe. +#[must_use] +#[no_mangle] +pub unsafe extern "C" fn ddog_prof_Profile_add_internal( + profile: *mut internal::Profile, + sample: Sample, + timestamp: Option, +) -> ProfileResult { + (|| { + let profile = profile + .as_mut() + .ok_or_else(|| anyhow::anyhow!("profile pointer was null"))?; + 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) + } else { + profile.add_sample(sample.try_into()?, timestamp) + } + })() + .context("ddog_prof_Profile_add_internal failed") + .into() +} + #[cfg(test)] mod tests { - use crate::manager::ProfilerManager; + use super::*; + + extern "C" fn test_cpu_sampler_callback(_: *mut datadog_profiling::internal::Profile) {} + extern "C" fn test_sample_converter(_: ManagedSample) -> Sample<'static> { + Sample { + locations: ddcommon_ffi::Slice::empty(), + values: ddcommon_ffi::Slice::empty(), + labels: ddcommon_ffi::Slice::empty(), + } + } + extern "C" fn test_reset_callback(_: ManagedSample) -> bool { + false + } #[test] fn test_the_thing() { let sample_types = []; let period = None; - let mut profile_manager = ProfilerManager::new(&sample_types, period); + let mut profile_manager = ProfilerManager::new( + &sample_types, + period, + test_cpu_sampler_callback, + test_sample_converter, + test_reset_callback, + ); println!("start"); profile_manager.main().unwrap(); } diff --git a/datadog-profiling-ffi/src/profiles/datatypes.rs b/datadog-profiling-ffi/src/profiles/datatypes.rs index 689aaf6a9a..4956f1a34c 100644 --- a/datadog-profiling-ffi/src/profiles/datatypes.rs +++ b/datadog-profiling-ffi/src/profiles/datatypes.rs @@ -28,21 +28,6 @@ impl Profile { } } - /// Creates a Profile from a raw pointer to an internal::Profile. - /// THIS IS UNSAFE, DO NOT USE IT. - /// IT ONLY EXISTS TO ALLOW US TO CREATE A PROFILE FROM A RAW POINTER FOR RAPID PROTOTYPING FOR - /// THE MANAGER. - /// - /// # Safety - /// - /// The caller must ensure that: - /// - The pointer is valid and points to a valid internal::Profile - /// - The pointer is properly aligned - /// - The memory is properly initialized - pub unsafe fn from_pointer(ptr: *mut internal::Profile) -> Self { - Profile { inner: ptr } - } - fn take(&mut self) -> Option> { // Leaving a null will help with double-free issues that can // arise in C. Of course, it's best to never get there in the From 2c21e5ee16404fa2cfcbc6f0be59b7205eb8d3ae Mon Sep 17 00:00:00 2001 From: Daniel Schwartz-Narbonne Date: Fri, 30 May 2025 16:33:53 -0400 Subject: [PATCH 03/52] recycled samples --- datadog-profiling-ffi/src/manager/mod.rs | 75 +++++++----------------- 1 file changed, 22 insertions(+), 53 deletions(-) diff --git a/datadog-profiling-ffi/src/manager/mod.rs b/datadog-profiling-ffi/src/manager/mod.rs index 225c33ff59..8ee1153546 100644 --- a/datadog-profiling-ffi/src/manager/mod.rs +++ b/datadog-profiling-ffi/src/manager/mod.rs @@ -17,57 +17,19 @@ use anyhow::Context; use crossbeam_channel::{select, tick, Receiver, Sender}; use datadog_profiling::{api, internal}; -#[repr(C)] -pub struct ManagedSample { - data: *mut c_void, - callback: extern "C" fn(*mut c_void, *mut internal::Profile), -} - -// void add_sample(void *data, Profile *profile) { -// Sample *sample = (Sample *)data; -// Profile *profile = (Profile *)data; -// profile->add_sample(sample); -// } - -impl ManagedSample { - /// Creates a new ManagedSample from a raw pointer and callback. - /// - /// # Safety - /// - /// The caller must ensure that: - /// - The pointer is valid and points to memory that will remain valid for the lifetime of the - /// ManagedSample - /// - The pointer is properly aligned - /// - The memory is properly initialized - #[no_mangle] - pub unsafe extern "C" fn new( - data: *mut c_void, - callback: extern "C" fn(*mut c_void, *mut internal::Profile), - ) -> Self { - Self { data, callback } - } - - /// Returns the raw pointer to the underlying data. - #[no_mangle] - pub extern "C" fn as_ptr(&self) -> *mut c_void { - self.data - } - - #[allow(clippy::todo)] - pub fn add(self, profile: &mut internal::Profile) { - (self.callback)(self.data, profile); - } -} - pub struct ProfilerManager { - samples_receiver: Receiver, - samples_sender: Sender, + samples_receiver: Receiver<*mut c_void>, + samples_sender: Sender<*mut c_void>, + recycled_samples_receiver: Receiver<*mut c_void>, + recycled_samples_sender: Sender<*mut c_void>, cpu_ticker: Receiver, upload_ticker: Receiver, profile: internal::Profile, cpu_sampler_callback: extern "C" fn(*mut internal::Profile), - sample_converter: extern "C" fn(ManagedSample) -> Sample<'static>, - reset_callback: extern "C" fn(ManagedSample) -> bool, + // Static is probably the wrong type here, but worry about that later. + sample_converter: extern "C" fn(*mut c_void) -> Sample<'static>, + // True if the sample should be reusused after reset, false otherwise. + reset_callback: extern "C" fn(*mut c_void) -> bool, } impl ProfilerManager { @@ -75,10 +37,11 @@ impl ProfilerManager { sample_types: &[api::ValueType], period: Option, cpu_sampler_callback: extern "C" fn(*mut internal::Profile), - sample_converter: extern "C" fn(ManagedSample) -> Sample<'static>, - reset_callback: extern "C" fn(ManagedSample) -> bool, + sample_converter: extern "C" fn(*mut c_void) -> Sample<'static>, + reset_callback: extern "C" fn(*mut c_void) -> bool, ) -> Self { let (samples_sender, samples_receiver) = crossbeam_channel::bounded(10); + let (recycled_samples_sender, recycled_samples_receiver) = crossbeam_channel::bounded(10); let profile = internal::Profile::new(sample_types, period); let cpu_ticker = tick(Duration::from_millis(100)); let upload_ticker = tick(Duration::from_secs(1)); @@ -87,6 +50,8 @@ impl ProfilerManager { upload_ticker, samples_receiver, samples_sender, + recycled_samples_receiver, + recycled_samples_sender, profile, cpu_sampler_callback, sample_converter, @@ -94,14 +59,18 @@ impl ProfilerManager { } } - #[allow(clippy::todo)] pub fn main(&mut self) -> anyhow::Result<()> { // This is just here to allow us to easily bail out. let done = tick(Duration::from_secs(5)); loop { select! { - recv(self.samples_receiver) -> sample => { - sample?.add(&mut self.profile); + recv(self.samples_receiver) -> raw_sample => { + let data = raw_sample?; + let sample = (self.sample_converter)(data); + self.profile.add_sample(sample.try_into()?, None)?; + if (self.reset_callback)(data) { + let _ = self.recycled_samples_sender.send(data); + } }, recv(self.cpu_ticker) -> msg => { (self.cpu_sampler_callback)(&mut self.profile); @@ -155,14 +124,14 @@ mod tests { use super::*; extern "C" fn test_cpu_sampler_callback(_: *mut datadog_profiling::internal::Profile) {} - extern "C" fn test_sample_converter(_: ManagedSample) -> Sample<'static> { + extern "C" fn test_sample_converter(_: *mut c_void) -> Sample<'static> { Sample { locations: ddcommon_ffi::Slice::empty(), values: ddcommon_ffi::Slice::empty(), labels: ddcommon_ffi::Slice::empty(), } } - extern "C" fn test_reset_callback(_: ManagedSample) -> bool { + extern "C" fn test_reset_callback(_: *mut c_void) -> bool { false } From e2d7c7cdfa9db281fed67212201437bae29eddcf Mon Sep 17 00:00:00 2001 From: Daniel Schwartz-Narbonne Date: Fri, 30 May 2025 17:04:42 -0400 Subject: [PATCH 04/52] more callbacks --- datadog-profiling-ffi/src/manager/mod.rs | 70 ++++++++++++++++++------ 1 file changed, 53 insertions(+), 17 deletions(-) diff --git a/datadog-profiling-ffi/src/manager/mod.rs b/datadog-profiling-ffi/src/manager/mod.rs index 8ee1153546..7f488ac485 100644 --- a/datadog-profiling-ffi/src/manager/mod.rs +++ b/datadog-profiling-ffi/src/manager/mod.rs @@ -17,6 +17,30 @@ use anyhow::Context; use crossbeam_channel::{select, tick, Receiver, Sender}; use datadog_profiling::{api, internal}; +#[repr(C)] +pub struct ManagedSampleCallbacks { + // Static is probably the wrong type here, but worry about that later. + converter: extern "C" fn(*mut c_void) -> Sample<'static>, + // Resets the sample for reuse. + reset: extern "C" fn(*mut c_void), + // Called when a sample is dropped (not recycled) + drop: extern "C" fn(*mut c_void), +} + +impl ManagedSampleCallbacks { + pub fn new( + converter: extern "C" fn(*mut c_void) -> Sample<'static>, + reset: extern "C" fn(*mut c_void), + drop: extern "C" fn(*mut c_void), + ) -> Self { + Self { + converter, + reset, + drop, + } + } +} + pub struct ProfilerManager { samples_receiver: Receiver<*mut c_void>, samples_sender: Sender<*mut c_void>, @@ -26,10 +50,8 @@ pub struct ProfilerManager { upload_ticker: Receiver, profile: internal::Profile, cpu_sampler_callback: extern "C" fn(*mut internal::Profile), - // Static is probably the wrong type here, but worry about that later. - sample_converter: extern "C" fn(*mut c_void) -> Sample<'static>, - // True if the sample should be reusused after reset, false otherwise. - reset_callback: extern "C" fn(*mut c_void) -> bool, + upload_callback: extern "C" fn(*mut internal::Profile), + sample_callbacks: ManagedSampleCallbacks, } impl ProfilerManager { @@ -37,8 +59,10 @@ impl ProfilerManager { sample_types: &[api::ValueType], period: Option, cpu_sampler_callback: extern "C" fn(*mut internal::Profile), + upload_callback: extern "C" fn(*mut internal::Profile), sample_converter: extern "C" fn(*mut c_void) -> Sample<'static>, - reset_callback: extern "C" fn(*mut c_void) -> bool, + reset_callback: extern "C" fn(*mut c_void), + drop_callback: extern "C" fn(*mut c_void), ) -> Self { let (samples_sender, samples_receiver) = crossbeam_channel::bounded(10); let (recycled_samples_sender, recycled_samples_receiver) = crossbeam_channel::bounded(10); @@ -54,8 +78,8 @@ impl ProfilerManager { recycled_samples_sender, profile, cpu_sampler_callback, - sample_converter, - reset_callback, + upload_callback, + sample_callbacks: ManagedSampleCallbacks::new(sample_converter, reset_callback, drop_callback), } } @@ -66,21 +90,22 @@ impl ProfilerManager { select! { recv(self.samples_receiver) -> raw_sample => { let data = raw_sample?; - let sample = (self.sample_converter)(data); + let sample = (self.sample_callbacks.converter)(data); self.profile.add_sample(sample.try_into()?, None)?; - if (self.reset_callback)(data) { - let _ = self.recycled_samples_sender.send(data); + (self.sample_callbacks.reset)(data); + if self.recycled_samples_sender.send(data).is_err() { + (self.sample_callbacks.drop)(data); } }, recv(self.cpu_ticker) -> msg => { (self.cpu_sampler_callback)(&mut self.profile); }, recv(self.upload_ticker) -> msg => { - let old_profile = self.profile.reset_and_return_previous()?; + let mut old_profile = self.profile.reset_and_return_previous()?; + let upload_callback = self.upload_callback; std::thread::spawn(move || { - if let Ok(encoded) = old_profile.serialize_into_compressed_pprof(None, None) { - println!("Successfully serialized profile"); - } + (upload_callback)(&mut old_profile); + // TODO: make sure we cleanup the profile. }); }, recv(done) -> msg => return Ok(()), @@ -123,16 +148,25 @@ pub unsafe extern "C" fn ddog_prof_Profile_add_internal( mod tests { use super::*; - extern "C" fn test_cpu_sampler_callback(_: *mut datadog_profiling::internal::Profile) {} + extern "C" fn test_cpu_sampler_callback(_: *mut datadog_profiling::internal::Profile) { + println!("cpu sampler callback"); + } + extern "C" fn test_upload_callback(_: *mut datadog_profiling::internal::Profile) { + println!("upload callback"); + } extern "C" fn test_sample_converter(_: *mut c_void) -> Sample<'static> { + println!("sample converter"); Sample { locations: ddcommon_ffi::Slice::empty(), values: ddcommon_ffi::Slice::empty(), labels: ddcommon_ffi::Slice::empty(), } } - extern "C" fn test_reset_callback(_: *mut c_void) -> bool { - false + extern "C" fn test_reset_callback(_: *mut c_void) { + println!("reset callback"); + } + extern "C" fn test_drop_callback(_: *mut c_void) { + println!("drop callback"); } #[test] @@ -143,8 +177,10 @@ mod tests { &sample_types, period, test_cpu_sampler_callback, + test_upload_callback, test_sample_converter, test_reset_callback, + test_drop_callback, ); println!("start"); profile_manager.main().unwrap(); From 3d4bb81b88d35f1f8f72a42ca9abc7605d2e9942 Mon Sep 17 00:00:00 2001 From: Daniel Schwartz-Narbonne Date: Mon, 2 Jun 2025 13:44:59 -0400 Subject: [PATCH 05/52] remove file I wasn't using --- datadog-profiling-ffi/src/manager/mod.rs | 8 +++++--- datadog-profiling-ffi/src/manager/uploader.rs | 2 -- 2 files changed, 5 insertions(+), 5 deletions(-) delete mode 100644 datadog-profiling-ffi/src/manager/uploader.rs diff --git a/datadog-profiling-ffi/src/manager/mod.rs b/datadog-profiling-ffi/src/manager/mod.rs index 7f488ac485..0826de4b93 100644 --- a/datadog-profiling-ffi/src/manager/mod.rs +++ b/datadog-profiling-ffi/src/manager/mod.rs @@ -4,8 +4,6 @@ #![allow(dead_code)] #![allow(clippy::todo)] -pub mod uploader; - use std::{ ffi::c_void, num::NonZeroI64, @@ -79,7 +77,11 @@ impl ProfilerManager { profile, cpu_sampler_callback, upload_callback, - sample_callbacks: ManagedSampleCallbacks::new(sample_converter, reset_callback, drop_callback), + sample_callbacks: ManagedSampleCallbacks::new( + sample_converter, + reset_callback, + drop_callback, + ), } } diff --git a/datadog-profiling-ffi/src/manager/uploader.rs b/datadog-profiling-ffi/src/manager/uploader.rs deleted file mode 100644 index 5a97e72498..0000000000 --- a/datadog-profiling-ffi/src/manager/uploader.rs +++ /dev/null @@ -1,2 +0,0 @@ -// Copyright 2025-Present Datadog, Inc. https://www.datadoghq.com/ -// SPDX-License-Identifier: Apache-2.0 From a0e8f25a59ee9758499d35ffca018e6f1dd67d4c Mon Sep 17 00:00:00 2001 From: Daniel Schwartz-Narbonne Date: Mon, 2 Jun 2025 15:56:31 -0400 Subject: [PATCH 06/52] send_sample api --- datadog-profiling-ffi/src/manager/mod.rs | 114 +++++++++++++++++------ 1 file changed, 86 insertions(+), 28 deletions(-) diff --git a/datadog-profiling-ffi/src/manager/mod.rs b/datadog-profiling-ffi/src/manager/mod.rs index 0826de4b93..c1030b4f2c 100644 --- a/datadog-profiling-ffi/src/manager/mod.rs +++ b/datadog-profiling-ffi/src/manager/mod.rs @@ -39,11 +39,67 @@ impl ManagedSampleCallbacks { } } +#[repr(transparent)] +pub struct SendSample(*mut c_void); + +// SAFETY: This type is used to transfer ownership of a sample between threads via channels. +// The sample is only accessed by one thread at a time, and ownership is transferred along +// with the SendSample wrapper. The sample is either processed by the manager thread or +// recycled back to the original thread. +unsafe impl Send for SendSample {} + +impl SendSample { + pub fn new(ptr: *mut c_void) -> Self { + Self(ptr) + } + + pub fn as_ptr(&self) -> *mut c_void { + self.0 + } +} + +pub struct SampleChannels { + samples_sender: Sender, + recycled_samples_receiver: Receiver, +} + +impl SampleChannels { + pub fn new() -> (Self, Receiver, Sender) { + let (samples_sender, samples_receiver) = crossbeam_channel::bounded(10); + let (recycled_samples_sender, recycled_samples_receiver) = crossbeam_channel::bounded(10); + ( + Self { + samples_sender, + recycled_samples_receiver, + }, + samples_receiver, + recycled_samples_sender, + ) + } + + /// # Safety + /// The caller must ensure that: + /// 1. The sample pointer is valid and points to a properly initialized sample + /// 2. The caller transfers ownership of the sample to this function + /// - The sample is not being used by any other thread + /// - The sample must not be accessed by the caller after this call + /// - The manager will either free the sample or recycle it back + /// 3. The sample will be properly cleaned up if it cannot be sent + pub unsafe fn send_sample( + &self, + sample: *mut c_void, + ) -> Result<(), crossbeam_channel::SendError> { + self.samples_sender.send(SendSample::new(sample)) + } + + pub fn try_recv_recycled(&self) -> Result<*mut c_void, crossbeam_channel::TryRecvError> { + self.recycled_samples_receiver.try_recv().map(|s| s.as_ptr()) + } +} + pub struct ProfilerManager { - samples_receiver: Receiver<*mut c_void>, - samples_sender: Sender<*mut c_void>, - recycled_samples_receiver: Receiver<*mut c_void>, - recycled_samples_sender: Sender<*mut c_void>, + samples_receiver: Receiver, + recycled_samples_sender: Sender, cpu_ticker: Receiver, upload_ticker: Receiver, profile: internal::Profile, @@ -53,49 +109,48 @@ pub struct ProfilerManager { } impl ProfilerManager { - pub fn new( + pub fn start( sample_types: &[api::ValueType], period: Option, cpu_sampler_callback: extern "C" fn(*mut internal::Profile), upload_callback: extern "C" fn(*mut internal::Profile), - sample_converter: extern "C" fn(*mut c_void) -> Sample<'static>, - reset_callback: extern "C" fn(*mut c_void), - drop_callback: extern "C" fn(*mut c_void), - ) -> Self { - let (samples_sender, samples_receiver) = crossbeam_channel::bounded(10); - let (recycled_samples_sender, recycled_samples_receiver) = crossbeam_channel::bounded(10); + sample_callbacks: ManagedSampleCallbacks, + ) -> SampleChannels { + let (channels, samples_receiver, recycled_samples_sender) = SampleChannels::new(); let profile = internal::Profile::new(sample_types, period); let cpu_ticker = tick(Duration::from_millis(100)); let upload_ticker = tick(Duration::from_secs(1)); - Self { + let mut manager = Self { cpu_ticker, upload_ticker, samples_receiver, - samples_sender, - recycled_samples_receiver, recycled_samples_sender, profile, cpu_sampler_callback, upload_callback, - sample_callbacks: ManagedSampleCallbacks::new( - sample_converter, - reset_callback, - drop_callback, - ), - } + sample_callbacks, + }; + + std::thread::spawn(move || { + if let Err(e) = manager.main() { + eprintln!("ProfilerManager error: {}", e); + } + }); + + channels } - pub fn main(&mut self) -> anyhow::Result<()> { + fn main(&mut self) -> anyhow::Result<()> { // This is just here to allow us to easily bail out. let done = tick(Duration::from_secs(5)); loop { select! { recv(self.samples_receiver) -> raw_sample => { - let data = raw_sample?; + let data = raw_sample?.as_ptr(); let sample = (self.sample_callbacks.converter)(data); self.profile.add_sample(sample.try_into()?, None)?; (self.sample_callbacks.reset)(data); - if self.recycled_samples_sender.send(data).is_err() { + if self.recycled_samples_sender.send(SendSample::new(data)).is_err() { (self.sample_callbacks.drop)(data); } }, @@ -175,16 +230,19 @@ mod tests { fn test_the_thing() { let sample_types = []; let period = None; - let mut profile_manager = ProfilerManager::new( + let sample_callbacks = ManagedSampleCallbacks::new( + test_sample_converter, + test_reset_callback, + test_drop_callback, + ); + let _channels = ProfilerManager::start( &sample_types, period, test_cpu_sampler_callback, test_upload_callback, - test_sample_converter, - test_reset_callback, - test_drop_callback, + sample_callbacks, ); println!("start"); - profile_manager.main().unwrap(); + std::thread::sleep(Duration::from_secs(5)); } } From 9f2545bfc8ef7b6b73aedf94701cedc2cc4d34b3 Mon Sep 17 00:00:00 2001 From: Daniel Schwartz-Narbonne Date: Mon, 2 Jun 2025 16:26:51 -0400 Subject: [PATCH 07/52] handle to allow us to join on the profiler thread --- datadog-profiling-ffi/src/manager/mod.rs | 44 ++++++++++++++++++++---- 1 file changed, 38 insertions(+), 6 deletions(-) diff --git a/datadog-profiling-ffi/src/manager/mod.rs b/datadog-profiling-ffi/src/manager/mod.rs index c1030b4f2c..118e97f77b 100644 --- a/datadog-profiling-ffi/src/manager/mod.rs +++ b/datadog-profiling-ffi/src/manager/mod.rs @@ -93,7 +93,9 @@ impl SampleChannels { } pub fn try_recv_recycled(&self) -> Result<*mut c_void, crossbeam_channel::TryRecvError> { - self.recycled_samples_receiver.try_recv().map(|s| s.as_ptr()) + self.recycled_samples_receiver + .try_recv() + .map(|s| s.as_ptr()) } } @@ -115,7 +117,7 @@ impl ProfilerManager { cpu_sampler_callback: extern "C" fn(*mut internal::Profile), upload_callback: extern "C" fn(*mut internal::Profile), sample_callbacks: ManagedSampleCallbacks, - ) -> SampleChannels { + ) -> ProfilerHandle { let (channels, samples_receiver, recycled_samples_sender) = SampleChannels::new(); let profile = internal::Profile::new(sample_types, period); let cpu_ticker = tick(Duration::from_millis(100)); @@ -131,13 +133,13 @@ impl ProfilerManager { sample_callbacks, }; - std::thread::spawn(move || { + let handle = std::thread::spawn(move || { if let Err(e) = manager.main() { eprintln!("ProfilerManager error: {}", e); } }); - channels + ProfilerHandle { channels, handle } } fn main(&mut self) -> anyhow::Result<()> { @@ -171,6 +173,36 @@ impl ProfilerManager { } } +pub struct ProfilerHandle { + channels: SampleChannels, + handle: std::thread::JoinHandle<()>, +} + +impl ProfilerHandle { + /// # Safety + /// The caller must ensure that: + /// 1. The sample pointer is valid and points to a properly initialized sample + /// 2. The caller transfers ownership of the sample to this function + /// - The sample is not being used by any other thread + /// - The sample must not be accessed by the caller after this call + /// - The manager will either free the sample or recycle it back + /// 3. The sample will be properly cleaned up if it cannot be sent + pub unsafe fn send_sample( + &self, + sample: *mut c_void, + ) -> Result<(), crossbeam_channel::SendError> { + self.channels.send_sample(sample) + } + + pub fn try_recv_recycled(&self) -> Result<*mut c_void, crossbeam_channel::TryRecvError> { + self.channels.try_recv_recycled() + } + + pub fn join(self) -> std::thread::Result<()> { + self.handle.join() + } +} + /// # Safety /// The `profile` ptr must point to a valid internal::Profile object. /// All pointers inside the `sample` need to be valid for the duration of this call. @@ -235,7 +267,7 @@ mod tests { test_reset_callback, test_drop_callback, ); - let _channels = ProfilerManager::start( + let handle = ProfilerManager::start( &sample_types, period, test_cpu_sampler_callback, @@ -243,6 +275,6 @@ mod tests { sample_callbacks, ); println!("start"); - std::thread::sleep(Duration::from_secs(5)); + handle.join().unwrap(); } } From d01c28d220a6a6d8a8dcc2c8d3e866f4939c1a51 Mon Sep 17 00:00:00 2001 From: Daniel Schwartz-Narbonne Date: Mon, 2 Jun 2025 16:34:30 -0400 Subject: [PATCH 08/52] shutdown --- datadog-profiling-ffi/src/manager/mod.rs | 20 ++++++++++++++------ 1 file changed, 14 insertions(+), 6 deletions(-) diff --git a/datadog-profiling-ffi/src/manager/mod.rs b/datadog-profiling-ffi/src/manager/mod.rs index 118e97f77b..40c54b4304 100644 --- a/datadog-profiling-ffi/src/manager/mod.rs +++ b/datadog-profiling-ffi/src/manager/mod.rs @@ -104,6 +104,7 @@ pub struct ProfilerManager { recycled_samples_sender: Sender, cpu_ticker: Receiver, upload_ticker: Receiver, + shutdown_receiver: Receiver<()>, profile: internal::Profile, cpu_sampler_callback: extern "C" fn(*mut internal::Profile), upload_callback: extern "C" fn(*mut internal::Profile), @@ -119,6 +120,7 @@ impl ProfilerManager { sample_callbacks: ManagedSampleCallbacks, ) -> ProfilerHandle { let (channels, samples_receiver, recycled_samples_sender) = SampleChannels::new(); + let (shutdown_sender, shutdown_receiver) = crossbeam_channel::bounded(1); let profile = internal::Profile::new(sample_types, period); let cpu_ticker = tick(Duration::from_millis(100)); let upload_ticker = tick(Duration::from_secs(1)); @@ -127,6 +129,7 @@ impl ProfilerManager { upload_ticker, samples_receiver, recycled_samples_sender, + shutdown_receiver, profile, cpu_sampler_callback, upload_callback, @@ -139,12 +142,14 @@ impl ProfilerManager { } }); - ProfilerHandle { channels, handle } + ProfilerHandle { + channels, + handle, + shutdown_sender, + } } fn main(&mut self) -> anyhow::Result<()> { - // This is just here to allow us to easily bail out. - let done = tick(Duration::from_secs(5)); loop { select! { recv(self.samples_receiver) -> raw_sample => { @@ -167,7 +172,7 @@ impl ProfilerManager { // TODO: make sure we cleanup the profile. }); }, - recv(done) -> msg => return Ok(()), + recv(self.shutdown_receiver) -> _ => return Ok(()), } } } @@ -176,6 +181,7 @@ impl ProfilerManager { pub struct ProfilerHandle { channels: SampleChannels, handle: std::thread::JoinHandle<()>, + shutdown_sender: Sender<()>, } impl ProfilerHandle { @@ -198,7 +204,8 @@ impl ProfilerHandle { self.channels.try_recv_recycled() } - pub fn join(self) -> std::thread::Result<()> { + pub fn shutdown(self) -> std::thread::Result<()> { + let _ = self.shutdown_sender.send(()); self.handle.join() } } @@ -275,6 +282,7 @@ mod tests { sample_callbacks, ); println!("start"); - handle.join().unwrap(); + std::thread::sleep(Duration::from_secs(5)); + handle.shutdown().unwrap(); } } From 979e243e2353c729fac0df6640021e04527892c2 Mon Sep 17 00:00:00 2001 From: Daniel Schwartz-Narbonne Date: Mon, 2 Jun 2025 16:39:37 -0400 Subject: [PATCH 09/52] unsafe --- datadog-profiling-ffi/src/manager/mod.rs | 20 +++++++++++++++++--- 1 file changed, 17 insertions(+), 3 deletions(-) diff --git a/datadog-profiling-ffi/src/manager/mod.rs b/datadog-profiling-ffi/src/manager/mod.rs index 40c54b4304..ee95734a53 100644 --- a/datadog-profiling-ffi/src/manager/mod.rs +++ b/datadog-profiling-ffi/src/manager/mod.rs @@ -49,7 +49,12 @@ pub struct SendSample(*mut c_void); unsafe impl Send for SendSample {} impl SendSample { - pub fn new(ptr: *mut c_void) -> Self { + /// # Safety + /// The caller must ensure that: + /// 1. The sample pointer is valid and points to a properly initialized sample + /// 2. The sample is not being used by any other thread + /// 3. The caller transfers ownership of the sample to this function + pub unsafe fn new(ptr: *mut c_void) -> Self { Self(ptr) } @@ -157,7 +162,10 @@ impl ProfilerManager { let sample = (self.sample_callbacks.converter)(data); self.profile.add_sample(sample.try_into()?, None)?; (self.sample_callbacks.reset)(data); - if self.recycled_samples_sender.send(SendSample::new(data)).is_err() { + // SAFETY: The sample pointer is valid because it came from the samples channel + // and was just processed by the converter and reset callbacks. We have exclusive + // access to it since we're the only thread that can receive from the samples channel. + if self.recycled_samples_sender.send(unsafe { SendSample::new(data) }).is_err() { (self.sample_callbacks.drop)(data); } }, @@ -172,7 +180,13 @@ impl ProfilerManager { // TODO: make sure we cleanup the profile. }); }, - recv(self.shutdown_receiver) -> _ => return Ok(()), + recv(self.shutdown_receiver) -> _ => { + // Drain any remaining samples and drop them + while let Ok(sample) = self.samples_receiver.try_recv() { + (self.sample_callbacks.drop)(sample.as_ptr()); + } + return Ok(()); + }, } } } From 29d65b644bbbd8b84a97ace35cfc60e8a71004bd Mon Sep 17 00:00:00 2001 From: Daniel Schwartz-Narbonne Date: Mon, 2 Jun 2025 17:15:12 -0400 Subject: [PATCH 10/52] cancellation token --- Cargo.lock | 1 + datadog-profiling-ffi/Cargo.toml | 1 + datadog-profiling-ffi/src/manager/mod.rs | 23 +++++++++++++++++++---- 3 files changed, 21 insertions(+), 4 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index e5aa2214fb..04f6ad9aed 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1615,6 +1615,7 @@ dependencies = [ "libc", "serde_json", "symbolizer-ffi", + "tokio", "tokio-util", ] diff --git a/datadog-profiling-ffi/Cargo.toml b/datadog-profiling-ffi/Cargo.toml index a6a1ee6225..48c8ed799f 100644 --- a/datadog-profiling-ffi/Cargo.toml +++ b/datadog-profiling-ffi/Cargo.toml @@ -49,4 +49,5 @@ hyper = { version = "1.6", features = ["http1", "client"] } libc = "0.2" serde_json = { version = "1.0" } symbolizer-ffi = { path = "../symbolizer-ffi", optional = true, default-features = false } +tokio = { version = "1.36", features = ["sync", "rt"] } tokio-util = "0.7.1" diff --git a/datadog-profiling-ffi/src/manager/mod.rs b/datadog-profiling-ffi/src/manager/mod.rs index ee95734a53..6cae4ec7ea 100644 --- a/datadog-profiling-ffi/src/manager/mod.rs +++ b/datadog-profiling-ffi/src/manager/mod.rs @@ -14,6 +14,7 @@ use crate::profiles::datatypes::{ProfileResult, Sample}; use anyhow::Context; use crossbeam_channel::{select, tick, Receiver, Sender}; use datadog_profiling::{api, internal}; +use tokio_util::sync::CancellationToken; #[repr(C)] pub struct ManagedSampleCallbacks { @@ -112,8 +113,9 @@ pub struct ProfilerManager { shutdown_receiver: Receiver<()>, profile: internal::Profile, cpu_sampler_callback: extern "C" fn(*mut internal::Profile), - upload_callback: extern "C" fn(*mut internal::Profile), + upload_callback: extern "C" fn(*mut internal::Profile, &mut Option), sample_callbacks: ManagedSampleCallbacks, + cancellation_token: Option, } impl ProfilerManager { @@ -121,7 +123,7 @@ impl ProfilerManager { sample_types: &[api::ValueType], period: Option, cpu_sampler_callback: extern "C" fn(*mut internal::Profile), - upload_callback: extern "C" fn(*mut internal::Profile), + upload_callback: extern "C" fn(*mut internal::Profile, &mut Option), sample_callbacks: ManagedSampleCallbacks, ) -> ProfilerHandle { let (channels, samples_receiver, recycled_samples_sender) = SampleChannels::new(); @@ -139,6 +141,7 @@ impl ProfilerManager { cpu_sampler_callback, upload_callback, sample_callbacks, + cancellation_token: None, }; let handle = std::thread::spawn(move || { @@ -175,8 +178,13 @@ impl ProfilerManager { recv(self.upload_ticker) -> msg => { let mut old_profile = self.profile.reset_and_return_previous()?; let upload_callback = self.upload_callback; + // Create a new cancellation token for this upload + let token = CancellationToken::new(); + // Store a clone of the token in the manager + self.cancellation_token = Some(token.clone()); + let mut cancellation_token = Some(token); std::thread::spawn(move || { - (upload_callback)(&mut old_profile); + (upload_callback)(&mut old_profile, &mut cancellation_token); // TODO: make sure we cleanup the profile. }); }, @@ -185,6 +193,10 @@ impl ProfilerManager { while let Ok(sample) = self.samples_receiver.try_recv() { (self.sample_callbacks.drop)(sample.as_ptr()); } + // Cancel any ongoing upload and drop the token + if let Some(token) = self.cancellation_token.take() { + token.cancel(); + } return Ok(()); }, } @@ -261,7 +273,10 @@ mod tests { extern "C" fn test_cpu_sampler_callback(_: *mut datadog_profiling::internal::Profile) { println!("cpu sampler callback"); } - extern "C" fn test_upload_callback(_: *mut datadog_profiling::internal::Profile) { + extern "C" fn test_upload_callback( + _: *mut datadog_profiling::internal::Profile, + _: &mut Option, + ) { println!("upload callback"); } extern "C" fn test_sample_converter(_: *mut c_void) -> Sample<'static> { From 1b9c6ab181747a477b4c55c7ac0a33d5f6de2de5 Mon Sep 17 00:00:00 2001 From: Daniel Schwartz-Narbonne Date: Mon, 2 Jun 2025 17:33:29 -0400 Subject: [PATCH 11/52] cleaner main loop --- datadog-profiling-ffi/src/manager/mod.rs | 82 +++++++++++++++--------- 1 file changed, 51 insertions(+), 31 deletions(-) diff --git a/datadog-profiling-ffi/src/manager/mod.rs b/datadog-profiling-ffi/src/manager/mod.rs index 6cae4ec7ea..afcc34562d 100644 --- a/datadog-profiling-ffi/src/manager/mod.rs +++ b/datadog-profiling-ffi/src/manager/mod.rs @@ -157,47 +157,67 @@ impl ProfilerManager { } } + fn handle_sample(&mut self, raw_sample: Result) -> anyhow::Result<()> { + let data = raw_sample?.as_ptr(); + let sample = (self.sample_callbacks.converter)(data); + self.profile.add_sample(sample.try_into()?, None)?; + (self.sample_callbacks.reset)(data); + // SAFETY: The sample pointer is valid because it came from the samples channel + // and was just processed by the converter and reset callbacks. We have exclusive + // access to it since we're the only thread that can receive from the samples channel. + if self.recycled_samples_sender.send(unsafe { SendSample::new(data) }).is_err() { + (self.sample_callbacks.drop)(data); + } + Ok(()) + } + + fn handle_cpu_tick(&mut self) { + (self.cpu_sampler_callback)(&mut self.profile); + } + + fn handle_upload_tick(&mut self) -> anyhow::Result<()> { + let mut old_profile = self.profile.reset_and_return_previous()?; + let upload_callback = self.upload_callback; + // Create a new cancellation token for this upload + let token = CancellationToken::new(); + // Store a clone of the token in the manager + self.cancellation_token = Some(token.clone()); + let mut cancellation_token = Some(token); + std::thread::spawn(move || { + (upload_callback)(&mut old_profile, &mut cancellation_token); + // TODO: make sure we cleanup the profile. + }); + Ok(()) + } + + fn handle_shutdown(&mut self) -> anyhow::Result<()> { + // Drain any remaining samples and drop them + while let Ok(sample) = self.samples_receiver.try_recv() { + (self.sample_callbacks.drop)(sample.as_ptr()); + } + // Cancel any ongoing upload and drop the token + if let Some(token) = self.cancellation_token.take() { + token.cancel(); + } + Ok(()) + } + fn main(&mut self) -> anyhow::Result<()> { loop { select! { recv(self.samples_receiver) -> raw_sample => { - let data = raw_sample?.as_ptr(); - let sample = (self.sample_callbacks.converter)(data); - self.profile.add_sample(sample.try_into()?, None)?; - (self.sample_callbacks.reset)(data); - // SAFETY: The sample pointer is valid because it came from the samples channel - // and was just processed by the converter and reset callbacks. We have exclusive - // access to it since we're the only thread that can receive from the samples channel. - if self.recycled_samples_sender.send(unsafe { SendSample::new(data) }).is_err() { - (self.sample_callbacks.drop)(data); - } + let _ = self.handle_sample(raw_sample) + .map_err(|e| eprintln!("Failed to process sample: {}", e)); }, recv(self.cpu_ticker) -> msg => { - (self.cpu_sampler_callback)(&mut self.profile); + self.handle_cpu_tick(); }, recv(self.upload_ticker) -> msg => { - let mut old_profile = self.profile.reset_and_return_previous()?; - let upload_callback = self.upload_callback; - // Create a new cancellation token for this upload - let token = CancellationToken::new(); - // Store a clone of the token in the manager - self.cancellation_token = Some(token.clone()); - let mut cancellation_token = Some(token); - std::thread::spawn(move || { - (upload_callback)(&mut old_profile, &mut cancellation_token); - // TODO: make sure we cleanup the profile. - }); + let _ = self.handle_upload_tick() + .map_err(|e| eprintln!("Failed to handle upload: {}", e)); }, recv(self.shutdown_receiver) -> _ => { - // Drain any remaining samples and drop them - while let Ok(sample) = self.samples_receiver.try_recv() { - (self.sample_callbacks.drop)(sample.as_ptr()); - } - // Cancel any ongoing upload and drop the token - if let Some(token) = self.cancellation_token.take() { - token.cancel(); - } - return Ok(()); + return self.handle_shutdown(); }, } } From be2d58322a6ca0928fd10a306a0f137cf16e4128 Mon Sep 17 00:00:00 2001 From: Daniel Schwartz-Narbonne Date: Mon, 2 Jun 2025 17:40:17 -0400 Subject: [PATCH 12/52] todo --- datadog-profiling-ffi/src/manager/mod.rs | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) diff --git a/datadog-profiling-ffi/src/manager/mod.rs b/datadog-profiling-ffi/src/manager/mod.rs index afcc34562d..a15638a55d 100644 --- a/datadog-profiling-ffi/src/manager/mod.rs +++ b/datadog-profiling-ffi/src/manager/mod.rs @@ -157,7 +157,10 @@ impl ProfilerManager { } } - fn handle_sample(&mut self, raw_sample: Result) -> anyhow::Result<()> { + fn handle_sample( + &mut self, + raw_sample: Result, + ) -> anyhow::Result<()> { let data = raw_sample?.as_ptr(); let sample = (self.sample_callbacks.converter)(data); self.profile.add_sample(sample.try_into()?, None)?; @@ -165,7 +168,11 @@ impl ProfilerManager { // SAFETY: The sample pointer is valid because it came from the samples channel // and was just processed by the converter and reset callbacks. We have exclusive // access to it since we're the only thread that can receive from the samples channel. - if self.recycled_samples_sender.send(unsafe { SendSample::new(data) }).is_err() { + if self + .recycled_samples_sender + .send(unsafe { SendSample::new(data) }) + .is_err() + { (self.sample_callbacks.drop)(data); } Ok(()) @@ -199,6 +206,7 @@ impl ProfilerManager { if let Some(token) = self.cancellation_token.take() { token.cancel(); } + // TODO: cleanup the recycled samples. Ok(()) } From baeecb9879640be894846372a7234a758d68a6bd Mon Sep 17 00:00:00 2001 From: Daniel Schwartz-Narbonne Date: Tue, 3 Jun 2025 14:44:59 -0400 Subject: [PATCH 13/52] client gets its own file --- datadog-profiling-ffi/src/manager/client.rs | 53 +++++++++++++++++++++ datadog-profiling-ffi/src/manager/mod.rs | 45 +++-------------- 2 files changed, 60 insertions(+), 38 deletions(-) create mode 100644 datadog-profiling-ffi/src/manager/client.rs diff --git a/datadog-profiling-ffi/src/manager/client.rs b/datadog-profiling-ffi/src/manager/client.rs new file mode 100644 index 0000000000..54d093d04c --- /dev/null +++ b/datadog-profiling-ffi/src/manager/client.rs @@ -0,0 +1,53 @@ +use std::{ + ffi::c_void, + thread::JoinHandle, +}; + +use crossbeam_channel::{Sender, SendError, TryRecvError}; + +use super::SampleChannels; +use super::SendSample; + +pub struct ManagedProfilerClient { + channels: SampleChannels, + handle: JoinHandle<()>, + shutdown_sender: Sender<()>, +} + +impl ManagedProfilerClient { + pub(crate) fn new( + channels: SampleChannels, + handle: JoinHandle<()>, + shutdown_sender: Sender<()>, + ) -> Self { + Self { + channels, + handle, + shutdown_sender, + } + } + + /// # Safety + /// The caller must ensure that: + /// 1. The sample pointer is valid and points to a properly initialized sample + /// 2. The caller transfers ownership of the sample to this function + /// - The sample is not being used by any other thread + /// - The sample must not be accessed by the caller after this call + /// - The manager will either free the sample or recycle it back + /// 3. The sample will be properly cleaned up if it cannot be sent + pub unsafe fn send_sample( + &self, + sample: *mut c_void, + ) -> Result<(), SendError> { + self.channels.send_sample(sample) + } + + pub fn try_recv_recycled(&self) -> Result<*mut c_void, TryRecvError> { + self.channels.try_recv_recycled() + } + + pub fn shutdown(self) -> std::thread::Result<()> { + let _ = self.shutdown_sender.send(()); + self.handle.join() + } +} \ No newline at end of file diff --git a/datadog-profiling-ffi/src/manager/mod.rs b/datadog-profiling-ffi/src/manager/mod.rs index a15638a55d..870dcfea8a 100644 --- a/datadog-profiling-ffi/src/manager/mod.rs +++ b/datadog-profiling-ffi/src/manager/mod.rs @@ -16,6 +16,9 @@ use crossbeam_channel::{select, tick, Receiver, Sender}; use datadog_profiling::{api, internal}; use tokio_util::sync::CancellationToken; +mod client; +pub use client::ManagedProfilerClient; + #[repr(C)] pub struct ManagedSampleCallbacks { // Static is probably the wrong type here, but worry about that later. @@ -40,6 +43,7 @@ impl ManagedSampleCallbacks { } } +// TODO: this owns the memory. It should probably be a full wrapper, with a destructor. #[repr(transparent)] pub struct SendSample(*mut c_void); @@ -125,7 +129,7 @@ impl ProfilerManager { cpu_sampler_callback: extern "C" fn(*mut internal::Profile), upload_callback: extern "C" fn(*mut internal::Profile, &mut Option), sample_callbacks: ManagedSampleCallbacks, - ) -> ProfilerHandle { + ) -> ManagedProfilerClient { let (channels, samples_receiver, recycled_samples_sender) = SampleChannels::new(); let (shutdown_sender, shutdown_receiver) = crossbeam_channel::bounded(1); let profile = internal::Profile::new(sample_types, period); @@ -150,11 +154,7 @@ impl ProfilerManager { } }); - ProfilerHandle { - channels, - handle, - shutdown_sender, - } + ManagedProfilerClient::new(channels, handle, shutdown_sender) } fn handle_sample( @@ -198,6 +198,7 @@ impl ProfilerManager { } fn handle_shutdown(&mut self) -> anyhow::Result<()> { + // TODO: a mechanism to force threads to wait to write to the channel. // Drain any remaining samples and drop them while let Ok(sample) = self.samples_receiver.try_recv() { (self.sample_callbacks.drop)(sample.as_ptr()); @@ -232,38 +233,6 @@ impl ProfilerManager { } } -pub struct ProfilerHandle { - channels: SampleChannels, - handle: std::thread::JoinHandle<()>, - shutdown_sender: Sender<()>, -} - -impl ProfilerHandle { - /// # Safety - /// The caller must ensure that: - /// 1. The sample pointer is valid and points to a properly initialized sample - /// 2. The caller transfers ownership of the sample to this function - /// - The sample is not being used by any other thread - /// - The sample must not be accessed by the caller after this call - /// - The manager will either free the sample or recycle it back - /// 3. The sample will be properly cleaned up if it cannot be sent - pub unsafe fn send_sample( - &self, - sample: *mut c_void, - ) -> Result<(), crossbeam_channel::SendError> { - self.channels.send_sample(sample) - } - - pub fn try_recv_recycled(&self) -> Result<*mut c_void, crossbeam_channel::TryRecvError> { - self.channels.try_recv_recycled() - } - - pub fn shutdown(self) -> std::thread::Result<()> { - let _ = self.shutdown_sender.send(()); - self.handle.join() - } -} - /// # Safety /// The `profile` ptr must point to a valid internal::Profile object. /// All pointers inside the `sample` need to be valid for the duration of this call. From 9397edbc1035e3056d68601b5fc262263cb2fff5 Mon Sep 17 00:00:00 2001 From: Daniel Schwartz-Narbonne Date: Tue, 3 Jun 2025 14:51:02 -0400 Subject: [PATCH 14/52] refactored files --- datadog-profiling-ffi/src/manager/client.rs | 14 +- datadog-profiling-ffi/src/manager/mod.rs | 232 +----------------- .../src/manager/profiler_manager.rs | 160 ++++++++++++ datadog-profiling-ffi/src/manager/samples.rs | 66 +++++ 4 files changed, 240 insertions(+), 232 deletions(-) create mode 100644 datadog-profiling-ffi/src/manager/profiler_manager.rs create mode 100644 datadog-profiling-ffi/src/manager/samples.rs diff --git a/datadog-profiling-ffi/src/manager/client.rs b/datadog-profiling-ffi/src/manager/client.rs index 54d093d04c..ce526c19fa 100644 --- a/datadog-profiling-ffi/src/manager/client.rs +++ b/datadog-profiling-ffi/src/manager/client.rs @@ -1,9 +1,6 @@ -use std::{ - ffi::c_void, - thread::JoinHandle, -}; +use std::{ffi::c_void, thread::JoinHandle}; -use crossbeam_channel::{Sender, SendError, TryRecvError}; +use crossbeam_channel::{SendError, Sender, TryRecvError}; use super::SampleChannels; use super::SendSample; @@ -35,10 +32,7 @@ impl ManagedProfilerClient { /// - The sample must not be accessed by the caller after this call /// - The manager will either free the sample or recycle it back /// 3. The sample will be properly cleaned up if it cannot be sent - pub unsafe fn send_sample( - &self, - sample: *mut c_void, - ) -> Result<(), SendError> { + pub unsafe fn send_sample(&self, sample: *mut c_void) -> Result<(), SendError> { self.channels.send_sample(sample) } @@ -50,4 +44,4 @@ impl ManagedProfilerClient { let _ = self.shutdown_sender.send(()); self.handle.join() } -} \ No newline at end of file +} diff --git a/datadog-profiling-ffi/src/manager/mod.rs b/datadog-profiling-ffi/src/manager/mod.rs index 870dcfea8a..2174e30dce 100644 --- a/datadog-profiling-ffi/src/manager/mod.rs +++ b/datadog-profiling-ffi/src/manager/mod.rs @@ -4,234 +4,19 @@ #![allow(dead_code)] #![allow(clippy::todo)] -use std::{ - ffi::c_void, - num::NonZeroI64, - time::{Duration, Instant}, -}; +use std::num::NonZeroI64; use crate::profiles::datatypes::{ProfileResult, Sample}; use anyhow::Context; -use crossbeam_channel::{select, tick, Receiver, Sender}; -use datadog_profiling::{api, internal}; -use tokio_util::sync::CancellationToken; +use datadog_profiling::internal; mod client; -pub use client::ManagedProfilerClient; - -#[repr(C)] -pub struct ManagedSampleCallbacks { - // Static is probably the wrong type here, but worry about that later. - converter: extern "C" fn(*mut c_void) -> Sample<'static>, - // Resets the sample for reuse. - reset: extern "C" fn(*mut c_void), - // Called when a sample is dropped (not recycled) - drop: extern "C" fn(*mut c_void), -} - -impl ManagedSampleCallbacks { - pub fn new( - converter: extern "C" fn(*mut c_void) -> Sample<'static>, - reset: extern "C" fn(*mut c_void), - drop: extern "C" fn(*mut c_void), - ) -> Self { - Self { - converter, - reset, - drop, - } - } -} - -// TODO: this owns the memory. It should probably be a full wrapper, with a destructor. -#[repr(transparent)] -pub struct SendSample(*mut c_void); - -// SAFETY: This type is used to transfer ownership of a sample between threads via channels. -// The sample is only accessed by one thread at a time, and ownership is transferred along -// with the SendSample wrapper. The sample is either processed by the manager thread or -// recycled back to the original thread. -unsafe impl Send for SendSample {} - -impl SendSample { - /// # Safety - /// The caller must ensure that: - /// 1. The sample pointer is valid and points to a properly initialized sample - /// 2. The sample is not being used by any other thread - /// 3. The caller transfers ownership of the sample to this function - pub unsafe fn new(ptr: *mut c_void) -> Self { - Self(ptr) - } - - pub fn as_ptr(&self) -> *mut c_void { - self.0 - } -} - -pub struct SampleChannels { - samples_sender: Sender, - recycled_samples_receiver: Receiver, -} - -impl SampleChannels { - pub fn new() -> (Self, Receiver, Sender) { - let (samples_sender, samples_receiver) = crossbeam_channel::bounded(10); - let (recycled_samples_sender, recycled_samples_receiver) = crossbeam_channel::bounded(10); - ( - Self { - samples_sender, - recycled_samples_receiver, - }, - samples_receiver, - recycled_samples_sender, - ) - } - - /// # Safety - /// The caller must ensure that: - /// 1. The sample pointer is valid and points to a properly initialized sample - /// 2. The caller transfers ownership of the sample to this function - /// - The sample is not being used by any other thread - /// - The sample must not be accessed by the caller after this call - /// - The manager will either free the sample or recycle it back - /// 3. The sample will be properly cleaned up if it cannot be sent - pub unsafe fn send_sample( - &self, - sample: *mut c_void, - ) -> Result<(), crossbeam_channel::SendError> { - self.samples_sender.send(SendSample::new(sample)) - } - - pub fn try_recv_recycled(&self) -> Result<*mut c_void, crossbeam_channel::TryRecvError> { - self.recycled_samples_receiver - .try_recv() - .map(|s| s.as_ptr()) - } -} +mod profiler_manager; +mod samples; -pub struct ProfilerManager { - samples_receiver: Receiver, - recycled_samples_sender: Sender, - cpu_ticker: Receiver, - upload_ticker: Receiver, - shutdown_receiver: Receiver<()>, - profile: internal::Profile, - cpu_sampler_callback: extern "C" fn(*mut internal::Profile), - upload_callback: extern "C" fn(*mut internal::Profile, &mut Option), - sample_callbacks: ManagedSampleCallbacks, - cancellation_token: Option, -} - -impl ProfilerManager { - pub fn start( - sample_types: &[api::ValueType], - period: Option, - cpu_sampler_callback: extern "C" fn(*mut internal::Profile), - upload_callback: extern "C" fn(*mut internal::Profile, &mut Option), - sample_callbacks: ManagedSampleCallbacks, - ) -> ManagedProfilerClient { - let (channels, samples_receiver, recycled_samples_sender) = SampleChannels::new(); - let (shutdown_sender, shutdown_receiver) = crossbeam_channel::bounded(1); - let profile = internal::Profile::new(sample_types, period); - let cpu_ticker = tick(Duration::from_millis(100)); - let upload_ticker = tick(Duration::from_secs(1)); - let mut manager = Self { - cpu_ticker, - upload_ticker, - samples_receiver, - recycled_samples_sender, - shutdown_receiver, - profile, - cpu_sampler_callback, - upload_callback, - sample_callbacks, - cancellation_token: None, - }; - - let handle = std::thread::spawn(move || { - if let Err(e) = manager.main() { - eprintln!("ProfilerManager error: {}", e); - } - }); - - ManagedProfilerClient::new(channels, handle, shutdown_sender) - } - - fn handle_sample( - &mut self, - raw_sample: Result, - ) -> anyhow::Result<()> { - let data = raw_sample?.as_ptr(); - let sample = (self.sample_callbacks.converter)(data); - self.profile.add_sample(sample.try_into()?, None)?; - (self.sample_callbacks.reset)(data); - // SAFETY: The sample pointer is valid because it came from the samples channel - // and was just processed by the converter and reset callbacks. We have exclusive - // access to it since we're the only thread that can receive from the samples channel. - if self - .recycled_samples_sender - .send(unsafe { SendSample::new(data) }) - .is_err() - { - (self.sample_callbacks.drop)(data); - } - Ok(()) - } - - fn handle_cpu_tick(&mut self) { - (self.cpu_sampler_callback)(&mut self.profile); - } - - fn handle_upload_tick(&mut self) -> anyhow::Result<()> { - let mut old_profile = self.profile.reset_and_return_previous()?; - let upload_callback = self.upload_callback; - // Create a new cancellation token for this upload - let token = CancellationToken::new(); - // Store a clone of the token in the manager - self.cancellation_token = Some(token.clone()); - let mut cancellation_token = Some(token); - std::thread::spawn(move || { - (upload_callback)(&mut old_profile, &mut cancellation_token); - // TODO: make sure we cleanup the profile. - }); - Ok(()) - } - - fn handle_shutdown(&mut self) -> anyhow::Result<()> { - // TODO: a mechanism to force threads to wait to write to the channel. - // Drain any remaining samples and drop them - while let Ok(sample) = self.samples_receiver.try_recv() { - (self.sample_callbacks.drop)(sample.as_ptr()); - } - // Cancel any ongoing upload and drop the token - if let Some(token) = self.cancellation_token.take() { - token.cancel(); - } - // TODO: cleanup the recycled samples. - Ok(()) - } - - fn main(&mut self) -> anyhow::Result<()> { - loop { - select! { - recv(self.samples_receiver) -> raw_sample => { - let _ = self.handle_sample(raw_sample) - .map_err(|e| eprintln!("Failed to process sample: {}", e)); - }, - recv(self.cpu_ticker) -> msg => { - self.handle_cpu_tick(); - }, - recv(self.upload_ticker) -> msg => { - let _ = self.handle_upload_tick() - .map_err(|e| eprintln!("Failed to handle upload: {}", e)); - }, - recv(self.shutdown_receiver) -> _ => { - return self.handle_shutdown(); - }, - } - } - } -} +pub use client::ManagedProfilerClient; +pub use profiler_manager::{ManagedSampleCallbacks, ProfilerManager}; +pub use samples::{SampleChannels, SendSample}; /// # Safety /// The `profile` ptr must point to a valid internal::Profile object. @@ -265,7 +50,10 @@ pub unsafe extern "C" fn ddog_prof_Profile_add_internal( #[cfg(test)] mod tests { + use std::{ffi::c_void, time::Duration}; + use super::*; + use tokio_util::sync::CancellationToken; extern "C" fn test_cpu_sampler_callback(_: *mut datadog_profiling::internal::Profile) { println!("cpu sampler callback"); diff --git a/datadog-profiling-ffi/src/manager/profiler_manager.rs b/datadog-profiling-ffi/src/manager/profiler_manager.rs new file mode 100644 index 0000000000..d849f12ab4 --- /dev/null +++ b/datadog-profiling-ffi/src/manager/profiler_manager.rs @@ -0,0 +1,160 @@ +use std::{ + ffi::c_void, + time::{Duration, Instant}, +}; + +use anyhow::Result; +use crossbeam_channel::{select, tick, Receiver, Sender}; +use datadog_profiling::{api, internal}; +use tokio_util::sync::CancellationToken; + +use super::samples::{SampleChannels, SendSample}; +use crate::profiles::datatypes::Sample; + +#[repr(C)] +pub struct ManagedSampleCallbacks { + // Static is probably the wrong type here, but worry about that later. + converter: extern "C" fn(*mut c_void) -> Sample<'static>, + // Resets the sample for reuse. + reset: extern "C" fn(*mut c_void), + // Called when a sample is dropped (not recycled) + drop: extern "C" fn(*mut c_void), +} + +impl ManagedSampleCallbacks { + pub fn new( + converter: extern "C" fn(*mut c_void) -> Sample<'static>, + reset: extern "C" fn(*mut c_void), + drop: extern "C" fn(*mut c_void), + ) -> Self { + Self { + converter, + reset, + drop, + } + } +} + +pub struct ProfilerManager { + samples_receiver: Receiver, + recycled_samples_sender: Sender, + cpu_ticker: Receiver, + upload_ticker: Receiver, + shutdown_receiver: Receiver<()>, + profile: internal::Profile, + cpu_sampler_callback: extern "C" fn(*mut internal::Profile), + upload_callback: extern "C" fn(*mut internal::Profile, &mut Option), + sample_callbacks: ManagedSampleCallbacks, + cancellation_token: Option, +} + +impl ProfilerManager { + pub fn start( + sample_types: &[api::ValueType], + period: Option, + cpu_sampler_callback: extern "C" fn(*mut internal::Profile), + upload_callback: extern "C" fn(*mut internal::Profile, &mut Option), + sample_callbacks: ManagedSampleCallbacks, + ) -> super::client::ManagedProfilerClient { + let (channels, samples_receiver, recycled_samples_sender) = SampleChannels::new(); + let (shutdown_sender, shutdown_receiver) = crossbeam_channel::bounded(1); + let profile = internal::Profile::new(sample_types, period); + let cpu_ticker = tick(Duration::from_millis(100)); + let upload_ticker = tick(Duration::from_secs(1)); + let mut manager = Self { + cpu_ticker, + upload_ticker, + samples_receiver, + recycled_samples_sender, + shutdown_receiver, + profile, + cpu_sampler_callback, + upload_callback, + sample_callbacks, + cancellation_token: None, + }; + + let handle = std::thread::spawn(move || { + if let Err(e) = manager.main() { + eprintln!("ProfilerManager error: {}", e); + } + }); + + super::client::ManagedProfilerClient::new(channels, handle, shutdown_sender) + } + + fn handle_sample( + &mut self, + raw_sample: Result, + ) -> Result<()> { + let data = raw_sample?.as_ptr(); + let sample = (self.sample_callbacks.converter)(data); + self.profile.add_sample(sample.try_into()?, None)?; + (self.sample_callbacks.reset)(data); + // SAFETY: The sample pointer is valid because it came from the samples channel + // and was just processed by the converter and reset callbacks. We have exclusive + // access to it since we're the only thread that can receive from the samples channel. + if self + .recycled_samples_sender + .send(unsafe { SendSample::new(data) }) + .is_err() + { + (self.sample_callbacks.drop)(data); + } + Ok(()) + } + + fn handle_cpu_tick(&mut self) { + (self.cpu_sampler_callback)(&mut self.profile); + } + + fn handle_upload_tick(&mut self) -> Result<()> { + let mut old_profile = self.profile.reset_and_return_previous()?; + let upload_callback = self.upload_callback; + // Create a new cancellation token for this upload + let token = CancellationToken::new(); + // Store a clone of the token in the manager + self.cancellation_token = Some(token.clone()); + let mut cancellation_token = Some(token); + std::thread::spawn(move || { + (upload_callback)(&mut old_profile, &mut cancellation_token); + // TODO: make sure we cleanup the profile. + }); + Ok(()) + } + + fn handle_shutdown(&mut self) -> Result<()> { + // TODO: a mechanism to force threads to wait to write to the channel. + // Drain any remaining samples and drop them + while let Ok(sample) = self.samples_receiver.try_recv() { + (self.sample_callbacks.drop)(sample.as_ptr()); + } + // Cancel any ongoing upload and drop the token + if let Some(token) = self.cancellation_token.take() { + token.cancel(); + } + // TODO: cleanup the recycled samples. + Ok(()) + } + + fn main(&mut self) -> Result<()> { + loop { + select! { + recv(self.samples_receiver) -> raw_sample => { + let _ = self.handle_sample(raw_sample) + .map_err(|e| eprintln!("Failed to process sample: {}", e)); + }, + recv(self.cpu_ticker) -> msg => { + self.handle_cpu_tick(); + }, + recv(self.upload_ticker) -> msg => { + let _ = self.handle_upload_tick() + .map_err(|e| eprintln!("Failed to handle upload: {}", e)); + }, + recv(self.shutdown_receiver) -> _ => { + return self.handle_shutdown(); + }, + } + } + } +} diff --git a/datadog-profiling-ffi/src/manager/samples.rs b/datadog-profiling-ffi/src/manager/samples.rs new file mode 100644 index 0000000000..13766ddcf7 --- /dev/null +++ b/datadog-profiling-ffi/src/manager/samples.rs @@ -0,0 +1,66 @@ +use std::ffi::c_void; + +use crossbeam_channel::{Receiver, SendError, Sender, TryRecvError}; + +// TODO: this owns the memory. It should probably be a full wrapper, with a destructor. +#[repr(transparent)] +pub struct SendSample(*mut c_void); + +// SAFETY: This type is used to transfer ownership of a sample between threads via channels. +// The sample is only accessed by one thread at a time, and ownership is transferred along +// with the SendSample wrapper. The sample is either processed by the manager thread or +// recycled back to the original thread. +unsafe impl Send for SendSample {} + +impl SendSample { + /// # Safety + /// The caller must ensure that: + /// 1. The sample pointer is valid and points to a properly initialized sample + /// 2. The sample is not being used by any other thread + /// 3. The caller transfers ownership of the sample to this function + pub unsafe fn new(ptr: *mut c_void) -> Self { + Self(ptr) + } + + pub fn as_ptr(&self) -> *mut c_void { + self.0 + } +} + +pub struct SampleChannels { + samples_sender: Sender, + recycled_samples_receiver: Receiver, +} + +impl SampleChannels { + pub fn new() -> (Self, Receiver, Sender) { + let (samples_sender, samples_receiver) = crossbeam_channel::bounded(10); + let (recycled_samples_sender, recycled_samples_receiver) = crossbeam_channel::bounded(10); + ( + Self { + samples_sender, + recycled_samples_receiver, + }, + samples_receiver, + recycled_samples_sender, + ) + } + + /// # Safety + /// The caller must ensure that: + /// 1. The sample pointer is valid and points to a properly initialized sample + /// 2. The caller transfers ownership of the sample to this function + /// - The sample is not being used by any other thread + /// - The sample must not be accessed by the caller after this call + /// - The manager will either free the sample or recycle it back + /// 3. The sample will be properly cleaned up if it cannot be sent + pub unsafe fn send_sample(&self, sample: *mut c_void) -> Result<(), SendError> { + self.samples_sender.send(SendSample::new(sample)) + } + + pub fn try_recv_recycled(&self) -> Result<*mut c_void, TryRecvError> { + self.recycled_samples_receiver + .try_recv() + .map(|sample| sample.as_ptr()) + } +} From cdcce6d32412c39da9aeb401b909eab67f698d1c Mon Sep 17 00:00:00 2001 From: Daniel Schwartz-Narbonne Date: Tue, 3 Jun 2025 15:01:35 -0400 Subject: [PATCH 15/52] continue moving files --- .../src/manager/ffi_utils.rs | 35 +++++++ datadog-profiling-ffi/src/manager/mod.rs | 92 +------------------ datadog-profiling-ffi/src/manager/tests.rs | 51 ++++++++++ 3 files changed, 89 insertions(+), 89 deletions(-) create mode 100644 datadog-profiling-ffi/src/manager/ffi_utils.rs create mode 100644 datadog-profiling-ffi/src/manager/tests.rs diff --git a/datadog-profiling-ffi/src/manager/ffi_utils.rs b/datadog-profiling-ffi/src/manager/ffi_utils.rs new file mode 100644 index 0000000000..297a3c1819 --- /dev/null +++ b/datadog-profiling-ffi/src/manager/ffi_utils.rs @@ -0,0 +1,35 @@ +use std::num::NonZeroI64; + +use crate::profiles::datatypes::{ProfileResult, Sample}; +use anyhow::Context; +use datadog_profiling::internal; + +/// # Safety +/// The `profile` ptr must point to a valid internal::Profile object. +/// All pointers inside the `sample` need to be valid for the duration of this call. +/// This call is _NOT_ thread-safe. +#[must_use] +#[no_mangle] +pub unsafe extern "C" fn ddog_prof_Profile_add_internal( + profile: *mut internal::Profile, + sample: Sample, + timestamp: Option, +) -> ProfileResult { + (|| { + let profile = profile + .as_mut() + .ok_or_else(|| anyhow::anyhow!("profile pointer was null"))?; + 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) + } else { + profile.add_sample(sample.try_into()?, timestamp) + } + })() + .context("ddog_prof_Profile_add_internal failed") + .into() +} diff --git a/datadog-profiling-ffi/src/manager/mod.rs b/datadog-profiling-ffi/src/manager/mod.rs index 2174e30dce..6c26e6f812 100644 --- a/datadog-profiling-ffi/src/manager/mod.rs +++ b/datadog-profiling-ffi/src/manager/mod.rs @@ -4,99 +4,13 @@ #![allow(dead_code)] #![allow(clippy::todo)] -use std::num::NonZeroI64; - -use crate::profiles::datatypes::{ProfileResult, Sample}; -use anyhow::Context; -use datadog_profiling::internal; - mod client; +mod ffi_utils; mod profiler_manager; mod samples; +#[cfg(test)] +mod tests; pub use client::ManagedProfilerClient; pub use profiler_manager::{ManagedSampleCallbacks, ProfilerManager}; pub use samples::{SampleChannels, SendSample}; - -/// # Safety -/// The `profile` ptr must point to a valid internal::Profile object. -/// All pointers inside the `sample` need to be valid for the duration of this call. -/// This call is _NOT_ thread-safe. -#[must_use] -#[no_mangle] -pub unsafe extern "C" fn ddog_prof_Profile_add_internal( - profile: *mut internal::Profile, - sample: Sample, - timestamp: Option, -) -> ProfileResult { - (|| { - let profile = profile - .as_mut() - .ok_or_else(|| anyhow::anyhow!("profile pointer was null"))?; - 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) - } else { - profile.add_sample(sample.try_into()?, timestamp) - } - })() - .context("ddog_prof_Profile_add_internal failed") - .into() -} - -#[cfg(test)] -mod tests { - use std::{ffi::c_void, time::Duration}; - - use super::*; - use tokio_util::sync::CancellationToken; - - extern "C" fn test_cpu_sampler_callback(_: *mut datadog_profiling::internal::Profile) { - println!("cpu sampler callback"); - } - extern "C" fn test_upload_callback( - _: *mut datadog_profiling::internal::Profile, - _: &mut Option, - ) { - println!("upload callback"); - } - extern "C" fn test_sample_converter(_: *mut c_void) -> Sample<'static> { - println!("sample converter"); - Sample { - locations: ddcommon_ffi::Slice::empty(), - values: ddcommon_ffi::Slice::empty(), - labels: ddcommon_ffi::Slice::empty(), - } - } - extern "C" fn test_reset_callback(_: *mut c_void) { - println!("reset callback"); - } - extern "C" fn test_drop_callback(_: *mut c_void) { - println!("drop callback"); - } - - #[test] - fn test_the_thing() { - let sample_types = []; - let period = None; - let sample_callbacks = ManagedSampleCallbacks::new( - test_sample_converter, - test_reset_callback, - test_drop_callback, - ); - let handle = ProfilerManager::start( - &sample_types, - period, - test_cpu_sampler_callback, - test_upload_callback, - sample_callbacks, - ); - println!("start"); - std::thread::sleep(Duration::from_secs(5)); - handle.shutdown().unwrap(); - } -} diff --git a/datadog-profiling-ffi/src/manager/tests.rs b/datadog-profiling-ffi/src/manager/tests.rs new file mode 100644 index 0000000000..17ee2749a1 --- /dev/null +++ b/datadog-profiling-ffi/src/manager/tests.rs @@ -0,0 +1,51 @@ +use std::{ffi::c_void, time::Duration}; + +use crate::profiles::datatypes::Sample; +use tokio_util::sync::CancellationToken; + +use super::{ManagedSampleCallbacks, ProfilerManager}; + +extern "C" fn test_cpu_sampler_callback(_: *mut datadog_profiling::internal::Profile) { + println!("cpu sampler callback"); +} +extern "C" fn test_upload_callback( + _: *mut datadog_profiling::internal::Profile, + _: &mut Option, +) { + println!("upload callback"); +} +extern "C" fn test_sample_converter(_: *mut c_void) -> Sample<'static> { + println!("sample converter"); + Sample { + locations: ddcommon_ffi::Slice::empty(), + values: ddcommon_ffi::Slice::empty(), + labels: ddcommon_ffi::Slice::empty(), + } +} +extern "C" fn test_reset_callback(_: *mut c_void) { + println!("reset callback"); +} +extern "C" fn test_drop_callback(_: *mut c_void) { + println!("drop callback"); +} + +#[test] +fn test_the_thing() { + let sample_types = []; + let period = None; + let sample_callbacks = ManagedSampleCallbacks::new( + test_sample_converter, + test_reset_callback, + test_drop_callback, + ); + let handle = ProfilerManager::start( + &sample_types, + period, + test_cpu_sampler_callback, + test_upload_callback, + sample_callbacks, + ); + println!("start"); + std::thread::sleep(Duration::from_secs(5)); + handle.shutdown().unwrap(); +} From b3baacb8bbde74f58cf4d3756f7787513a569a87 Mon Sep 17 00:00:00 2001 From: Daniel Schwartz-Narbonne Date: Tue, 3 Jun 2025 15:05:31 -0400 Subject: [PATCH 16/52] ffi_api file --- datadog-profiling-ffi/src/manager/ffi_api.rs | 2 ++ datadog-profiling-ffi/src/manager/mod.rs | 1 + 2 files changed, 3 insertions(+) create mode 100644 datadog-profiling-ffi/src/manager/ffi_api.rs diff --git a/datadog-profiling-ffi/src/manager/ffi_api.rs b/datadog-profiling-ffi/src/manager/ffi_api.rs new file mode 100644 index 0000000000..5a97e72498 --- /dev/null +++ b/datadog-profiling-ffi/src/manager/ffi_api.rs @@ -0,0 +1,2 @@ +// Copyright 2025-Present Datadog, Inc. https://www.datadoghq.com/ +// SPDX-License-Identifier: Apache-2.0 diff --git a/datadog-profiling-ffi/src/manager/mod.rs b/datadog-profiling-ffi/src/manager/mod.rs index 6c26e6f812..6bcbf76cf1 100644 --- a/datadog-profiling-ffi/src/manager/mod.rs +++ b/datadog-profiling-ffi/src/manager/mod.rs @@ -5,6 +5,7 @@ #![allow(clippy::todo)] mod client; +mod ffi_api; mod ffi_utils; mod profiler_manager; mod samples; From 9e5b53372452807d166df97e9a16c007febd5f10 Mon Sep 17 00:00:00 2001 From: Daniel Schwartz-Narbonne Date: Tue, 3 Jun 2025 15:48:33 -0400 Subject: [PATCH 17/52] starting on an FFI api --- datadog-profiling-ffi/src/manager/client.rs | 8 ++- datadog-profiling-ffi/src/manager/ffi_api.rs | 71 +++++++++++++++++++ .../src/manager/profiler_manager.rs | 10 ++- datadog-profiling-ffi/src/manager/tests.rs | 3 +- 4 files changed, 87 insertions(+), 5 deletions(-) diff --git a/datadog-profiling-ffi/src/manager/client.rs b/datadog-profiling-ffi/src/manager/client.rs index ce526c19fa..520913752c 100644 --- a/datadog-profiling-ffi/src/manager/client.rs +++ b/datadog-profiling-ffi/src/manager/client.rs @@ -40,8 +40,12 @@ impl ManagedProfilerClient { self.channels.try_recv_recycled() } - pub fn shutdown(self) -> std::thread::Result<()> { + pub fn shutdown(self) -> anyhow::Result<()> { + // Todo: Should we report if there was an error sending the shutdown signal? let _ = self.shutdown_sender.send(()); - self.handle.join() + self.handle + .join() + .map_err(|e| anyhow::anyhow!("Failed to join handle: {:?}", e))?; + Ok(()) } } diff --git a/datadog-profiling-ffi/src/manager/ffi_api.rs b/datadog-profiling-ffi/src/manager/ffi_api.rs index 5a97e72498..69cde51eaf 100644 --- a/datadog-profiling-ffi/src/manager/ffi_api.rs +++ b/datadog-profiling-ffi/src/manager/ffi_api.rs @@ -1,2 +1,73 @@ // Copyright 2025-Present Datadog, Inc. https://www.datadoghq.com/ // SPDX-License-Identifier: Apache-2.0 + +use crate::manager::{ + profiler_manager::{ManagedSampleCallbacks, ProfilerManager}, + ManagedProfilerClient, +}; +use crate::profiles::datatypes::{Period, ValueType}; +use anyhow::Ok; +use datadog_profiling::internal; +use ddcommon_ffi::{ + wrap_with_ffi_result, wrap_with_void_ffi_result, Handle, Result as FFIResult, Slice, ToInner, + VoidResult, +}; +use function_name::named; +use tokio_util::sync::CancellationToken; + +/// # Safety +/// - The caller is responsible for eventually calling the appropriate shutdown and cleanup +/// functions. +/// - The sample_callbacks must remain valid for the lifetime of the profiler. +/// - This function is not thread-safe. +#[no_mangle] +#[named] +pub unsafe extern "C" fn ddog_prof_ProfilerManager_start( + sample_types: Slice, + period: Option<&Period>, + cpu_sampler_callback: extern "C" fn(*mut internal::Profile), + upload_callback: extern "C" fn(*mut internal::Profile, &mut Option), + sample_callbacks: ManagedSampleCallbacks, +) -> FFIResult> { + wrap_with_ffi_result!({ + let sample_types_vec: Vec<_> = sample_types.into_slice().iter().map(Into::into).collect(); + let period_opt = period.map(Into::into); + let client: ManagedProfilerClient = ProfilerManager::start( + &sample_types_vec, + period_opt, + cpu_sampler_callback, + upload_callback, + sample_callbacks, + )?; + Ok(Handle::from(client)) + }) +} + +/// # Safety +/// - The handle must have been returned by ddog_prof_ProfilerManager_start and not yet dropped. +#[no_mangle] +#[named] +pub unsafe extern "C" fn ddog_prof_ProfilerManager_shutdown( + handle: *mut Handle, +) -> VoidResult { + wrap_with_void_ffi_result!({ + let handle = handle.as_mut().context("Invalid handle")?; + let client = handle + .take() + .context("Failed to take ownership of client")?; + client + .shutdown() + .map_err(|e| anyhow::anyhow!("Failed to shutdown client: {:?}", e))?; + }) +} + +/// # Safety +/// - The handle must have been returned by ddog_prof_ProfilerManager_start and not yet dropped. +#[no_mangle] +pub unsafe extern "C" fn ddog_prof_ProfilerManager_drop( + handle: *mut Handle, +) { + if let Some(handle) = handle.as_mut() { + let _ = handle.take(); + } +} diff --git a/datadog-profiling-ffi/src/manager/profiler_manager.rs b/datadog-profiling-ffi/src/manager/profiler_manager.rs index d849f12ab4..077101651d 100644 --- a/datadog-profiling-ffi/src/manager/profiler_manager.rs +++ b/datadog-profiling-ffi/src/manager/profiler_manager.rs @@ -8,10 +8,12 @@ use crossbeam_channel::{select, tick, Receiver, Sender}; use datadog_profiling::{api, internal}; use tokio_util::sync::CancellationToken; +use super::client::ManagedProfilerClient; use super::samples::{SampleChannels, SendSample}; use crate::profiles::datatypes::Sample; #[repr(C)] +#[derive(Clone)] pub struct ManagedSampleCallbacks { // Static is probably the wrong type here, but worry about that later. converter: extern "C" fn(*mut c_void) -> Sample<'static>, @@ -55,7 +57,7 @@ impl ProfilerManager { cpu_sampler_callback: extern "C" fn(*mut internal::Profile), upload_callback: extern "C" fn(*mut internal::Profile, &mut Option), sample_callbacks: ManagedSampleCallbacks, - ) -> super::client::ManagedProfilerClient { + ) -> Result { let (channels, samples_receiver, recycled_samples_sender) = SampleChannels::new(); let (shutdown_sender, shutdown_receiver) = crossbeam_channel::bounded(1); let profile = internal::Profile::new(sample_types, period); @@ -80,7 +82,11 @@ impl ProfilerManager { } }); - super::client::ManagedProfilerClient::new(channels, handle, shutdown_sender) + Ok(ManagedProfilerClient::new( + channels, + handle, + shutdown_sender, + )) } fn handle_sample( diff --git a/datadog-profiling-ffi/src/manager/tests.rs b/datadog-profiling-ffi/src/manager/tests.rs index 17ee2749a1..05deb57d6b 100644 --- a/datadog-profiling-ffi/src/manager/tests.rs +++ b/datadog-profiling-ffi/src/manager/tests.rs @@ -44,7 +44,8 @@ fn test_the_thing() { test_cpu_sampler_callback, test_upload_callback, sample_callbacks, - ); + ) + .expect("Failed to start profiler"); println!("start"); std::thread::sleep(Duration::from_secs(5)); handle.shutdown().unwrap(); From d303f9f892f0ee6056d7849811ef0262fec973f8 Mon Sep 17 00:00:00 2001 From: Daniel Schwartz-Narbonne Date: Tue, 3 Jun 2025 16:33:05 -0400 Subject: [PATCH 18/52] ffi api functions --- datadog-profiling-ffi/src/manager/ffi_api.rs | 57 +++++++++++++++++++- 1 file changed, 55 insertions(+), 2 deletions(-) diff --git a/datadog-profiling-ffi/src/manager/ffi_api.rs b/datadog-profiling-ffi/src/manager/ffi_api.rs index 69cde51eaf..8c96f0d03b 100644 --- a/datadog-profiling-ffi/src/manager/ffi_api.rs +++ b/datadog-profiling-ffi/src/manager/ffi_api.rs @@ -6,13 +6,14 @@ use crate::manager::{ ManagedProfilerClient, }; use crate::profiles::datatypes::{Period, ValueType}; -use anyhow::Ok; +use crossbeam_channel::TryRecvError; use datadog_profiling::internal; use ddcommon_ffi::{ wrap_with_ffi_result, wrap_with_void_ffi_result, Handle, Result as FFIResult, Slice, ToInner, VoidResult, }; use function_name::named; +use std::ffi::c_void; use tokio_util::sync::CancellationToken; /// # Safety @@ -39,7 +40,58 @@ pub unsafe extern "C" fn ddog_prof_ProfilerManager_start( upload_callback, sample_callbacks, )?; - Ok(Handle::from(client)) + anyhow::Ok(Handle::from(client)) + }) +} + +/// # Safety +/// - The handle must have been returned by ddog_prof_ProfilerManager_start and not yet dropped. +/// - The caller must ensure that: +/// 1. The sample pointer is valid and points to a properly initialized sample +/// 2. The caller transfers ownership of the sample to this function +/// - The sample is not being used by any other thread +/// - The sample must not be accessed by the caller after this call +/// - The manager will either free the sample or recycle it back +#[no_mangle] +#[named] +pub unsafe extern "C" fn ddog_prof_ProfilerManager_enqueue_sample( + handle: *mut Handle, + sample_ptr: *mut c_void, +) -> VoidResult { + wrap_with_void_ffi_result!({ + let handle = handle.as_mut().context("Invalid handle")?; + let client = handle.to_inner_mut()?; + client + .send_sample(sample_ptr) + .map_err(|e| anyhow::anyhow!("Failed to send sample: {:?}", e))?; + }) +} + +/// Attempts to receive a recycled sample from the profiler manager. +/// +/// This function will: +/// - Return a valid sample pointer if a recycled sample is available +/// - Return a null pointer if the queue is empty (this is a valid success case) +/// - Return an error if the channel is disconnected +/// +/// The caller should check if the returned pointer is null to determine if there were no samples +/// available. +/// +/// # Safety +/// - The handle must have been returned by ddog_prof_ProfilerManager_start and not yet dropped. +#[no_mangle] +#[named] +pub unsafe extern "C" fn ddog_prof_ProfilerManager_try_recv_recycled( + handle: *mut Handle, +) -> FFIResult<*mut c_void> { + wrap_with_ffi_result!({ + let handle = handle.as_mut().context("Invalid handle")?; + let client = handle.to_inner_mut()?; + match client.try_recv_recycled() { + Ok(sample_ptr) => anyhow::Ok(sample_ptr), + Err(TryRecvError::Empty) => anyhow::Ok(std::ptr::null_mut()), + Err(TryRecvError::Disconnected) => Err(anyhow::anyhow!("Channel disconnected")), + } }) } @@ -64,6 +116,7 @@ pub unsafe extern "C" fn ddog_prof_ProfilerManager_shutdown( /// # Safety /// - The handle must have been returned by ddog_prof_ProfilerManager_start and not yet dropped. #[no_mangle] +// TODO: Do we want drop and shutdown to be separate functions? Or should it always be shutdown? pub unsafe extern "C" fn ddog_prof_ProfilerManager_drop( handle: *mut Handle, ) { From e76015b87280db839d8525eb7f26d01bcd44e079 Mon Sep 17 00:00:00 2001 From: Daniel Schwartz-Narbonne Date: Wed, 4 Jun 2025 14:06:30 -0400 Subject: [PATCH 19/52] todos --- datadog-profiling-ffi/src/manager/profiler_manager.rs | 3 +++ 1 file changed, 3 insertions(+) diff --git a/datadog-profiling-ffi/src/manager/profiler_manager.rs b/datadog-profiling-ffi/src/manager/profiler_manager.rs index 077101651d..6473a86f9d 100644 --- a/datadog-profiling-ffi/src/manager/profiler_manager.rs +++ b/datadog-profiling-ffi/src/manager/profiler_manager.rs @@ -61,7 +61,10 @@ impl ProfilerManager { let (channels, samples_receiver, recycled_samples_sender) = SampleChannels::new(); let (shutdown_sender, shutdown_receiver) = crossbeam_channel::bounded(1); let profile = internal::Profile::new(sample_types, period); + // For adaptive sampling, we need to be able to adjust this duration. Look into how to do + // this. let cpu_ticker = tick(Duration::from_millis(100)); + // one second for testing, make this 1 minute in production let upload_ticker = tick(Duration::from_secs(1)); let mut manager = Self { cpu_ticker, From 4c54b020a6027ed5a46d33f55b055e88f891b0d5 Mon Sep 17 00:00:00 2001 From: Daniel Schwartz-Narbonne Date: Wed, 4 Jun 2025 16:08:53 -0400 Subject: [PATCH 20/52] directly take a profile when starting the profilemanager --- datadog-profiling-ffi/src/manager/ffi_api.rs | 14 +++--- .../src/manager/profiler_manager.rs | 6 +-- datadog-profiling-ffi/src/manager/tests.rs | 5 +- .../src/profiles/datatypes.rs | 49 ++++++++++++++----- datadog-trace-protobuf/src/remoteconfig.rs | 30 ++++++++---- 5 files changed, 67 insertions(+), 37 deletions(-) diff --git a/datadog-profiling-ffi/src/manager/ffi_api.rs b/datadog-profiling-ffi/src/manager/ffi_api.rs index 8c96f0d03b..01ca9b274d 100644 --- a/datadog-profiling-ffi/src/manager/ffi_api.rs +++ b/datadog-profiling-ffi/src/manager/ffi_api.rs @@ -5,11 +5,11 @@ use crate::manager::{ profiler_manager::{ManagedSampleCallbacks, ProfilerManager}, ManagedProfilerClient, }; -use crate::profiles::datatypes::{Period, ValueType}; +use crate::profiles::datatypes::{Profile, ProfilePtrExt}; use crossbeam_channel::TryRecvError; use datadog_profiling::internal; use ddcommon_ffi::{ - wrap_with_ffi_result, wrap_with_void_ffi_result, Handle, Result as FFIResult, Slice, ToInner, + wrap_with_ffi_result, wrap_with_void_ffi_result, Handle, Result as FFIResult, ToInner, VoidResult, }; use function_name::named; @@ -21,21 +21,19 @@ use tokio_util::sync::CancellationToken; /// functions. /// - The sample_callbacks must remain valid for the lifetime of the profiler. /// - This function is not thread-safe. +/// - This function takes ownership of the profile. The profile must not be used after this call. #[no_mangle] #[named] pub unsafe extern "C" fn ddog_prof_ProfilerManager_start( - sample_types: Slice, - period: Option<&Period>, + profile: *mut Profile, cpu_sampler_callback: extern "C" fn(*mut internal::Profile), upload_callback: extern "C" fn(*mut internal::Profile, &mut Option), sample_callbacks: ManagedSampleCallbacks, ) -> FFIResult> { wrap_with_ffi_result!({ - let sample_types_vec: Vec<_> = sample_types.into_slice().iter().map(Into::into).collect(); - let period_opt = period.map(Into::into); + let internal_profile = *profile.take()?; let client: ManagedProfilerClient = ProfilerManager::start( - &sample_types_vec, - period_opt, + internal_profile, cpu_sampler_callback, upload_callback, sample_callbacks, diff --git a/datadog-profiling-ffi/src/manager/profiler_manager.rs b/datadog-profiling-ffi/src/manager/profiler_manager.rs index 6473a86f9d..6e6c4bab34 100644 --- a/datadog-profiling-ffi/src/manager/profiler_manager.rs +++ b/datadog-profiling-ffi/src/manager/profiler_manager.rs @@ -5,7 +5,7 @@ use std::{ use anyhow::Result; use crossbeam_channel::{select, tick, Receiver, Sender}; -use datadog_profiling::{api, internal}; +use datadog_profiling::internal; use tokio_util::sync::CancellationToken; use super::client::ManagedProfilerClient; @@ -52,15 +52,13 @@ pub struct ProfilerManager { impl ProfilerManager { pub fn start( - sample_types: &[api::ValueType], - period: Option, + profile: internal::Profile, cpu_sampler_callback: extern "C" fn(*mut internal::Profile), upload_callback: extern "C" fn(*mut internal::Profile, &mut Option), sample_callbacks: ManagedSampleCallbacks, ) -> Result { let (channels, samples_receiver, recycled_samples_sender) = SampleChannels::new(); let (shutdown_sender, shutdown_receiver) = crossbeam_channel::bounded(1); - let profile = internal::Profile::new(sample_types, period); // For adaptive sampling, we need to be able to adjust this duration. Look into how to do // this. let cpu_ticker = tick(Duration::from_millis(100)); diff --git a/datadog-profiling-ffi/src/manager/tests.rs b/datadog-profiling-ffi/src/manager/tests.rs index 05deb57d6b..dcbf5310db 100644 --- a/datadog-profiling-ffi/src/manager/tests.rs +++ b/datadog-profiling-ffi/src/manager/tests.rs @@ -1,6 +1,7 @@ use std::{ffi::c_void, time::Duration}; use crate::profiles::datatypes::Sample; +use datadog_profiling::internal; use tokio_util::sync::CancellationToken; use super::{ManagedSampleCallbacks, ProfilerManager}; @@ -33,14 +34,14 @@ extern "C" fn test_drop_callback(_: *mut c_void) { fn test_the_thing() { let sample_types = []; let period = None; + let profile = internal::Profile::new(&sample_types, period); let sample_callbacks = ManagedSampleCallbacks::new( test_sample_converter, test_reset_callback, test_drop_callback, ); let handle = ProfilerManager::start( - &sample_types, - period, + profile, test_cpu_sampler_callback, test_upload_callback, sample_callbacks, diff --git a/datadog-profiling-ffi/src/profiles/datatypes.rs b/datadog-profiling-ffi/src/profiles/datatypes.rs index 4956f1a34c..5a097ed68c 100644 --- a/datadog-profiling-ffi/src/profiles/datatypes.rs +++ b/datadog-profiling-ffi/src/profiles/datatypes.rs @@ -27,24 +27,47 @@ impl Profile { inner: Box::into_raw(Box::new(profile)), } } +} - fn take(&mut self) -> Option> { - // Leaving a null will help with double-free issues that can - // arise in C. Of course, it's best to never get there in the - // first place! - let raw = std::mem::replace(&mut self.inner, std::ptr::null_mut()); +impl Drop for Profile { + fn drop(&mut self) { + // SAFETY: Profile's inner pointer is only set in new() and take(), and take() ensures + // the pointer is null after taking ownership. Since this is Drop, we know the Profile + // is being destroyed and won't be used again. + unsafe { drop(self.take()) } + } +} - if raw.is_null() { - None - } else { - Some(unsafe { Box::from_raw(raw) }) - } +impl ToInner for Profile { + unsafe fn to_inner_mut(&mut self) -> anyhow::Result<&mut internal::Profile> { + self.inner + .as_mut() + .context("inner pointer was null, indicates use after free") + } + + unsafe fn take(&mut self) -> anyhow::Result> { + let raw = std::mem::replace(&mut self.inner, std::ptr::null_mut()); + anyhow::ensure!( + !raw.is_null(), + "inner pointer was null, indicates use after free" + ); + Ok(Box::from_raw(raw)) } } -impl Drop for Profile { - fn drop(&mut self) { - drop(self.take()) +/// Extension trait for raw Profile pointers. +/// We need this trait because Rust's orphan rules prevent us from implementing methods directly +/// on raw pointers (*mut Profile). This trait provides a safe way to take ownership of a Profile +/// from a raw pointer while maintaining proper error handling. +pub trait ProfilePtrExt { + /// # Safety + /// The pointer must be non-null and point to a valid Profile that hasn't been dropped. + unsafe fn take(self) -> anyhow::Result>; +} + +impl ProfilePtrExt for *mut Profile { + unsafe fn take(self) -> anyhow::Result> { + self.as_mut().context("Null pointer")?.take() } } diff --git a/datadog-trace-protobuf/src/remoteconfig.rs b/datadog-trace-protobuf/src/remoteconfig.rs index a8a4b66524..fdc89ee871 100644 --- a/datadog-trace-protobuf/src/remoteconfig.rs +++ b/datadog-trace-protobuf/src/remoteconfig.rs @@ -3,7 +3,8 @@ use serde::{Deserialize, Serialize}; // This file is @generated by prost-build. -#[derive(Deserialize, Serialize, Clone, PartialEq, ::prost::Message)] +#[derive(Deserialize, Serialize)] +#[derive(Clone, PartialEq, ::prost::Message)] pub struct File { #[prost(string, tag = "1")] pub path: ::prost::alloc::string::String, @@ -11,7 +12,8 @@ pub struct File { #[serde(with = "serde_bytes")] pub raw: ::prost::alloc::vec::Vec, } -#[derive(Deserialize, Serialize, Clone, PartialEq, ::prost::Message)] +#[derive(Deserialize, Serialize)] +#[derive(Clone, PartialEq, ::prost::Message)] pub struct Client { #[prost(message, optional, tag = "1")] pub state: ::core::option::Option, @@ -33,7 +35,8 @@ pub struct Client { #[prost(bytes = "vec", tag = "11")] pub capabilities: ::prost::alloc::vec::Vec, } -#[derive(Deserialize, Serialize, Clone, PartialEq, ::prost::Message)] +#[derive(Deserialize, Serialize)] +#[derive(Clone, PartialEq, ::prost::Message)] pub struct ClientTracer { #[prost(string, tag = "1")] pub runtime_id: ::prost::alloc::string::String, @@ -52,7 +55,8 @@ pub struct ClientTracer { #[prost(string, repeated, tag = "7")] pub tags: ::prost::alloc::vec::Vec<::prost::alloc::string::String>, } -#[derive(Deserialize, Serialize, Clone, PartialEq, ::prost::Message)] +#[derive(Deserialize, Serialize)] +#[derive(Clone, PartialEq, ::prost::Message)] pub struct ClientAgent { #[prost(string, tag = "1")] pub name: ::prost::alloc::string::String, @@ -65,7 +69,8 @@ pub struct ClientAgent { #[prost(string, repeated, tag = "5")] pub cws_workloads: ::prost::alloc::vec::Vec<::prost::alloc::string::String>, } -#[derive(Deserialize, Serialize, Clone, PartialEq, ::prost::Message)] +#[derive(Deserialize, Serialize)] +#[derive(Clone, PartialEq, ::prost::Message)] pub struct ConfigState { #[prost(string, tag = "1")] pub id: ::prost::alloc::string::String, @@ -78,7 +83,8 @@ pub struct ConfigState { #[prost(string, tag = "5")] pub apply_error: ::prost::alloc::string::String, } -#[derive(Deserialize, Serialize, Clone, PartialEq, ::prost::Message)] +#[derive(Deserialize, Serialize)] +#[derive(Clone, PartialEq, ::prost::Message)] pub struct ClientState { #[prost(uint64, tag = "1")] pub root_version: u64, @@ -93,14 +99,16 @@ pub struct ClientState { #[prost(bytes = "vec", tag = "6")] pub backend_client_state: ::prost::alloc::vec::Vec, } -#[derive(Deserialize, Serialize, Clone, PartialEq, ::prost::Message)] +#[derive(Deserialize, Serialize)] +#[derive(Clone, PartialEq, ::prost::Message)] pub struct TargetFileHash { #[prost(string, tag = "1")] pub algorithm: ::prost::alloc::string::String, #[prost(string, tag = "3")] pub hash: ::prost::alloc::string::String, } -#[derive(Deserialize, Serialize, Clone, PartialEq, ::prost::Message)] +#[derive(Deserialize, Serialize)] +#[derive(Clone, PartialEq, ::prost::Message)] pub struct TargetFileMeta { #[prost(string, tag = "1")] pub path: ::prost::alloc::string::String, @@ -109,14 +117,16 @@ pub struct TargetFileMeta { #[prost(message, repeated, tag = "3")] pub hashes: ::prost::alloc::vec::Vec, } -#[derive(Deserialize, Serialize, Clone, PartialEq, ::prost::Message)] +#[derive(Deserialize, Serialize)] +#[derive(Clone, PartialEq, ::prost::Message)] pub struct ClientGetConfigsRequest { #[prost(message, optional, tag = "1")] pub client: ::core::option::Option, #[prost(message, repeated, tag = "2")] pub cached_target_files: ::prost::alloc::vec::Vec, } -#[derive(Deserialize, Serialize, Clone, PartialEq, ::prost::Message)] +#[derive(Deserialize, Serialize)] +#[derive(Clone, PartialEq, ::prost::Message)] pub struct ClientGetConfigsResponse { #[prost(bytes = "vec", repeated, tag = "1")] #[serde(with = "crate::serde")] From a2ea1b8b4dfe9ba4df1547c402955769e203e739 Mon Sep 17 00:00:00 2001 From: Daniel Schwartz-Narbonne Date: Wed, 4 Jun 2025 16:27:11 -0400 Subject: [PATCH 21/52] shutting down return the profile --- datadog-profiling-ffi/src/manager/client.rs | 10 +++++----- datadog-profiling-ffi/src/manager/ffi_api.rs | 7 ++++--- .../src/manager/profiler_manager.rs | 17 +++++++++-------- datadog-profiling-ffi/src/profiles/datatypes.rs | 2 +- 4 files changed, 19 insertions(+), 17 deletions(-) diff --git a/datadog-profiling-ffi/src/manager/client.rs b/datadog-profiling-ffi/src/manager/client.rs index 520913752c..494f315d25 100644 --- a/datadog-profiling-ffi/src/manager/client.rs +++ b/datadog-profiling-ffi/src/manager/client.rs @@ -1,20 +1,21 @@ use std::{ffi::c_void, thread::JoinHandle}; use crossbeam_channel::{SendError, Sender, TryRecvError}; +use datadog_profiling::internal; use super::SampleChannels; use super::SendSample; pub struct ManagedProfilerClient { channels: SampleChannels, - handle: JoinHandle<()>, + handle: JoinHandle>, shutdown_sender: Sender<()>, } impl ManagedProfilerClient { pub(crate) fn new( channels: SampleChannels, - handle: JoinHandle<()>, + handle: JoinHandle>, shutdown_sender: Sender<()>, ) -> Self { Self { @@ -40,12 +41,11 @@ impl ManagedProfilerClient { self.channels.try_recv_recycled() } - pub fn shutdown(self) -> anyhow::Result<()> { + pub fn shutdown(self) -> anyhow::Result { // Todo: Should we report if there was an error sending the shutdown signal? let _ = self.shutdown_sender.send(()); self.handle .join() - .map_err(|e| anyhow::anyhow!("Failed to join handle: {:?}", e))?; - Ok(()) + .map_err(|e| anyhow::anyhow!("Failed to join handle: {:?}", e))? } } diff --git a/datadog-profiling-ffi/src/manager/ffi_api.rs b/datadog-profiling-ffi/src/manager/ffi_api.rs index 01ca9b274d..2bfaea4181 100644 --- a/datadog-profiling-ffi/src/manager/ffi_api.rs +++ b/datadog-profiling-ffi/src/manager/ffi_api.rs @@ -99,15 +99,16 @@ pub unsafe extern "C" fn ddog_prof_ProfilerManager_try_recv_recycled( #[named] pub unsafe extern "C" fn ddog_prof_ProfilerManager_shutdown( handle: *mut Handle, -) -> VoidResult { - wrap_with_void_ffi_result!({ +) -> FFIResult { + wrap_with_ffi_result!({ let handle = handle.as_mut().context("Invalid handle")?; let client = handle .take() .context("Failed to take ownership of client")?; - client + let profile = client .shutdown() .map_err(|e| anyhow::anyhow!("Failed to shutdown client: {:?}", e))?; + anyhow::Ok(Profile::new(profile)) }) } diff --git a/datadog-profiling-ffi/src/manager/profiler_manager.rs b/datadog-profiling-ffi/src/manager/profiler_manager.rs index 6e6c4bab34..6fe1062a50 100644 --- a/datadog-profiling-ffi/src/manager/profiler_manager.rs +++ b/datadog-profiling-ffi/src/manager/profiler_manager.rs @@ -64,7 +64,7 @@ impl ProfilerManager { let cpu_ticker = tick(Duration::from_millis(100)); // one second for testing, make this 1 minute in production let upload_ticker = tick(Duration::from_secs(1)); - let mut manager = Self { + let manager = Self { cpu_ticker, upload_ticker, samples_receiver, @@ -77,11 +77,7 @@ impl ProfilerManager { cancellation_token: None, }; - let handle = std::thread::spawn(move || { - if let Err(e) = manager.main() { - eprintln!("ProfilerManager error: {}", e); - } - }); + let handle = std::thread::spawn(move || manager.main()); Ok(ManagedProfilerClient::new( channels, @@ -144,7 +140,10 @@ impl ProfilerManager { Ok(()) } - fn main(&mut self) -> Result<()> { + /// # Safety + /// - The caller must ensure that the callbacks remain valid for the lifetime of the profiler. + /// - The callbacks must be thread-safe. + pub fn main(mut self) -> Result { loop { select! { recv(self.samples_receiver) -> raw_sample => { @@ -159,9 +158,11 @@ impl ProfilerManager { .map_err(|e| eprintln!("Failed to handle upload: {}", e)); }, recv(self.shutdown_receiver) -> _ => { - return self.handle_shutdown(); + self.handle_shutdown()?; + break; }, } } + Ok(self.profile) } } diff --git a/datadog-profiling-ffi/src/profiles/datatypes.rs b/datadog-profiling-ffi/src/profiles/datatypes.rs index 5a097ed68c..acba94cd1d 100644 --- a/datadog-profiling-ffi/src/profiles/datatypes.rs +++ b/datadog-profiling-ffi/src/profiles/datatypes.rs @@ -22,7 +22,7 @@ pub struct Profile { } impl Profile { - fn new(profile: internal::Profile) -> Self { + pub(crate) fn new(profile: internal::Profile) -> Self { Profile { inner: Box::into_raw(Box::new(profile)), } From 24f38870586f7b09cd744da6638d79092644a41f Mon Sep 17 00:00:00 2001 From: Daniel Schwartz-Narbonne Date: Wed, 4 Jun 2025 16:39:15 -0400 Subject: [PATCH 22/52] block sending samples when shutdown has started --- datadog-profiling-ffi/src/manager/client.rs | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/datadog-profiling-ffi/src/manager/client.rs b/datadog-profiling-ffi/src/manager/client.rs index 494f315d25..cd06451c84 100644 --- a/datadog-profiling-ffi/src/manager/client.rs +++ b/datadog-profiling-ffi/src/manager/client.rs @@ -1,4 +1,4 @@ -use std::{ffi::c_void, thread::JoinHandle}; +use std::{ffi::c_void, sync::atomic::AtomicBool, thread::JoinHandle}; use crossbeam_channel::{SendError, Sender, TryRecvError}; use datadog_profiling::internal; @@ -10,6 +10,7 @@ pub struct ManagedProfilerClient { channels: SampleChannels, handle: JoinHandle>, shutdown_sender: Sender<()>, + is_shutdown: AtomicBool, } impl ManagedProfilerClient { @@ -22,6 +23,7 @@ impl ManagedProfilerClient { channels, handle, shutdown_sender, + is_shutdown: AtomicBool::new(false), } } @@ -34,6 +36,9 @@ impl ManagedProfilerClient { /// - The manager will either free the sample or recycle it back /// 3. The sample will be properly cleaned up if it cannot be sent pub unsafe fn send_sample(&self, sample: *mut c_void) -> Result<(), SendError> { + if self.is_shutdown.load(std::sync::atomic::Ordering::SeqCst) { + return Err(SendError(unsafe { SendSample::new(sample) })); + } self.channels.send_sample(sample) } @@ -42,6 +47,8 @@ impl ManagedProfilerClient { } pub fn shutdown(self) -> anyhow::Result { + self.is_shutdown + .store(true, std::sync::atomic::Ordering::SeqCst); // Todo: Should we report if there was an error sending the shutdown signal? let _ = self.shutdown_sender.send(()); self.handle From a5fcfab6644f46e07d9b40a324e418209fbb2373 Mon Sep 17 00:00:00 2001 From: Daniel Schwartz-Narbonne Date: Wed, 4 Jun 2025 16:44:06 -0400 Subject: [PATCH 23/52] handle any last samples before cancelling --- .../src/manager/profiler_manager.rs | 14 +++++++++++--- 1 file changed, 11 insertions(+), 3 deletions(-) diff --git a/datadog-profiling-ffi/src/manager/profiler_manager.rs b/datadog-profiling-ffi/src/manager/profiler_manager.rs index 6fe1062a50..5cf6152fe8 100644 --- a/datadog-profiling-ffi/src/manager/profiler_manager.rs +++ b/datadog-profiling-ffi/src/manager/profiler_manager.rs @@ -127,10 +127,18 @@ impl ProfilerManager { } fn handle_shutdown(&mut self) -> Result<()> { - // TODO: a mechanism to force threads to wait to write to the channel. - // Drain any remaining samples and drop them + // Try to process any remaining samples before dropping them while let Ok(sample) = self.samples_receiver.try_recv() { - (self.sample_callbacks.drop)(sample.as_ptr()); + let data = sample.as_ptr(); + let sample = (self.sample_callbacks.converter)(data); + if let Ok(converted_sample) = sample.try_into() { + if let Err(e) = self.profile.add_sample(converted_sample, None) { + eprintln!("Failed to add sample during shutdown: {}", e); + } + } else { + eprintln!("Failed to convert sample during shutdown"); + } + (self.sample_callbacks.drop)(data); } // Cancel any ongoing upload and drop the token if let Some(token) = self.cancellation_token.take() { From 319ee147b9ff71d468333f539714762cc13cb86f Mon Sep 17 00:00:00 2001 From: Daniel Schwartz-Narbonne Date: Wed, 4 Jun 2025 17:12:32 -0400 Subject: [PATCH 24/52] clear the recycled channels --- datadog-profiling-ffi/src/manager/client.rs | 7 +++-- datadog-profiling-ffi/src/manager/mod.rs | 2 +- .../src/manager/profiler_manager.rs | 28 +++++++++++-------- datadog-profiling-ffi/src/manager/samples.rs | 19 +++++++++---- 4 files changed, 36 insertions(+), 20 deletions(-) diff --git a/datadog-profiling-ffi/src/manager/client.rs b/datadog-profiling-ffi/src/manager/client.rs index cd06451c84..93f7a1eb3b 100644 --- a/datadog-profiling-ffi/src/manager/client.rs +++ b/datadog-profiling-ffi/src/manager/client.rs @@ -1,13 +1,14 @@ use std::{ffi::c_void, sync::atomic::AtomicBool, thread::JoinHandle}; +use anyhow::Result; use crossbeam_channel::{SendError, Sender, TryRecvError}; use datadog_profiling::internal; -use super::SampleChannels; +use super::ClientSampleChannels; use super::SendSample; pub struct ManagedProfilerClient { - channels: SampleChannels, + channels: ClientSampleChannels, handle: JoinHandle>, shutdown_sender: Sender<()>, is_shutdown: AtomicBool, @@ -15,7 +16,7 @@ pub struct ManagedProfilerClient { impl ManagedProfilerClient { pub(crate) fn new( - channels: SampleChannels, + channels: ClientSampleChannels, handle: JoinHandle>, shutdown_sender: Sender<()>, ) -> Self { diff --git a/datadog-profiling-ffi/src/manager/mod.rs b/datadog-profiling-ffi/src/manager/mod.rs index 6bcbf76cf1..a9c81ac3ba 100644 --- a/datadog-profiling-ffi/src/manager/mod.rs +++ b/datadog-profiling-ffi/src/manager/mod.rs @@ -14,4 +14,4 @@ mod tests; pub use client::ManagedProfilerClient; pub use profiler_manager::{ManagedSampleCallbacks, ProfilerManager}; -pub use samples::{SampleChannels, SendSample}; +pub use samples::{ClientSampleChannels, SendSample}; diff --git a/datadog-profiling-ffi/src/manager/profiler_manager.rs b/datadog-profiling-ffi/src/manager/profiler_manager.rs index 5cf6152fe8..c9d9e57f43 100644 --- a/datadog-profiling-ffi/src/manager/profiler_manager.rs +++ b/datadog-profiling-ffi/src/manager/profiler_manager.rs @@ -4,12 +4,12 @@ use std::{ }; use anyhow::Result; -use crossbeam_channel::{select, tick, Receiver, Sender}; +use crossbeam_channel::{select, tick, Receiver}; use datadog_profiling::internal; use tokio_util::sync::CancellationToken; use super::client::ManagedProfilerClient; -use super::samples::{SampleChannels, SendSample}; +use super::samples::{ClientSampleChannels, ManagerSampleChannels, SendSample}; use crate::profiles::datatypes::Sample; #[repr(C)] @@ -38,8 +38,7 @@ impl ManagedSampleCallbacks { } pub struct ProfilerManager { - samples_receiver: Receiver, - recycled_samples_sender: Sender, + channels: ManagerSampleChannels, cpu_ticker: Receiver, upload_ticker: Receiver, shutdown_receiver: Receiver<()>, @@ -57,18 +56,19 @@ impl ProfilerManager { upload_callback: extern "C" fn(*mut internal::Profile, &mut Option), sample_callbacks: ManagedSampleCallbacks, ) -> Result { - let (channels, samples_receiver, recycled_samples_sender) = SampleChannels::new(); + let (client_channels, manager_channels) = ClientSampleChannels::new(); let (shutdown_sender, shutdown_receiver) = crossbeam_channel::bounded(1); + // For adaptive sampling, we need to be able to adjust this duration. Look into how to do // this. let cpu_ticker = tick(Duration::from_millis(100)); // one second for testing, make this 1 minute in production let upload_ticker = tick(Duration::from_secs(1)); + let manager = Self { + channels: manager_channels, cpu_ticker, upload_ticker, - samples_receiver, - recycled_samples_sender, shutdown_receiver, profile, cpu_sampler_callback, @@ -80,7 +80,7 @@ impl ProfilerManager { let handle = std::thread::spawn(move || manager.main()); Ok(ManagedProfilerClient::new( - channels, + client_channels, handle, shutdown_sender, )) @@ -98,6 +98,7 @@ impl ProfilerManager { // and was just processed by the converter and reset callbacks. We have exclusive // access to it since we're the only thread that can receive from the samples channel. if self + .channels .recycled_samples_sender .send(unsafe { SendSample::new(data) }) .is_err() @@ -128,7 +129,7 @@ impl ProfilerManager { fn handle_shutdown(&mut self) -> Result<()> { // Try to process any remaining samples before dropping them - while let Ok(sample) = self.samples_receiver.try_recv() { + while let Ok(sample) = self.channels.samples_receiver.try_recv() { let data = sample.as_ptr(); let sample = (self.sample_callbacks.converter)(data); if let Ok(converted_sample) = sample.try_into() { @@ -140,11 +141,16 @@ impl ProfilerManager { } (self.sample_callbacks.drop)(data); } + + // Drain any recycled samples + while let Ok(sample) = self.channels.recycled_samples_receiver.try_recv() { + (self.sample_callbacks.drop)(sample.as_ptr()); + } + // Cancel any ongoing upload and drop the token if let Some(token) = self.cancellation_token.take() { token.cancel(); } - // TODO: cleanup the recycled samples. Ok(()) } @@ -154,7 +160,7 @@ impl ProfilerManager { pub fn main(mut self) -> Result { loop { select! { - recv(self.samples_receiver) -> raw_sample => { + recv(self.channels.samples_receiver) -> raw_sample => { let _ = self.handle_sample(raw_sample) .map_err(|e| eprintln!("Failed to process sample: {}", e)); }, diff --git a/datadog-profiling-ffi/src/manager/samples.rs b/datadog-profiling-ffi/src/manager/samples.rs index 13766ddcf7..7c34bfa9db 100644 --- a/datadog-profiling-ffi/src/manager/samples.rs +++ b/datadog-profiling-ffi/src/manager/samples.rs @@ -27,22 +27,31 @@ impl SendSample { } } -pub struct SampleChannels { +pub struct ClientSampleChannels { samples_sender: Sender, recycled_samples_receiver: Receiver, } -impl SampleChannels { - pub fn new() -> (Self, Receiver, Sender) { +pub struct ManagerSampleChannels { + pub samples_receiver: Receiver, + pub recycled_samples_sender: Sender, + pub recycled_samples_receiver: Receiver, +} + +impl ClientSampleChannels { + pub fn new() -> (Self, ManagerSampleChannels) { let (samples_sender, samples_receiver) = crossbeam_channel::bounded(10); let (recycled_samples_sender, recycled_samples_receiver) = crossbeam_channel::bounded(10); ( Self { samples_sender, + recycled_samples_receiver: recycled_samples_receiver.clone(), + }, + ManagerSampleChannels { + samples_receiver, + recycled_samples_sender, recycled_samples_receiver, }, - samples_receiver, - recycled_samples_sender, ) } From 516385b65596eaed77d0e0401d28d23a772ea511 Mon Sep 17 00:00:00 2001 From: Daniel Schwartz-Narbonne Date: Wed, 4 Jun 2025 17:17:58 -0400 Subject: [PATCH 25/52] use handle.take directly --- datadog-profiling-ffi/src/manager/ffi_api.rs | 28 +++++++------------- 1 file changed, 10 insertions(+), 18 deletions(-) diff --git a/datadog-profiling-ffi/src/manager/ffi_api.rs b/datadog-profiling-ffi/src/manager/ffi_api.rs index 2bfaea4181..e24d054b7a 100644 --- a/datadog-profiling-ffi/src/manager/ffi_api.rs +++ b/datadog-profiling-ffi/src/manager/ffi_api.rs @@ -53,13 +53,12 @@ pub unsafe extern "C" fn ddog_prof_ProfilerManager_start( #[no_mangle] #[named] pub unsafe extern "C" fn ddog_prof_ProfilerManager_enqueue_sample( - handle: *mut Handle, + mut handle: *mut Handle, sample_ptr: *mut c_void, ) -> VoidResult { wrap_with_void_ffi_result!({ - let handle = handle.as_mut().context("Invalid handle")?; - let client = handle.to_inner_mut()?; - client + handle + .to_inner_mut()? .send_sample(sample_ptr) .map_err(|e| anyhow::anyhow!("Failed to send sample: {:?}", e))?; }) @@ -80,12 +79,10 @@ pub unsafe extern "C" fn ddog_prof_ProfilerManager_enqueue_sample( #[no_mangle] #[named] pub unsafe extern "C" fn ddog_prof_ProfilerManager_try_recv_recycled( - handle: *mut Handle, + mut handle: *mut Handle, ) -> FFIResult<*mut c_void> { wrap_with_ffi_result!({ - let handle = handle.as_mut().context("Invalid handle")?; - let client = handle.to_inner_mut()?; - match client.try_recv_recycled() { + match handle.to_inner_mut()?.try_recv_recycled() { Ok(sample_ptr) => anyhow::Ok(sample_ptr), Err(TryRecvError::Empty) => anyhow::Ok(std::ptr::null_mut()), Err(TryRecvError::Disconnected) => Err(anyhow::anyhow!("Channel disconnected")), @@ -98,14 +95,11 @@ pub unsafe extern "C" fn ddog_prof_ProfilerManager_try_recv_recycled( #[no_mangle] #[named] pub unsafe extern "C" fn ddog_prof_ProfilerManager_shutdown( - handle: *mut Handle, + mut handle: *mut Handle, ) -> FFIResult { wrap_with_ffi_result!({ - let handle = handle.as_mut().context("Invalid handle")?; - let client = handle - .take() - .context("Failed to take ownership of client")?; - let profile = client + let profile = handle + .take()? .shutdown() .map_err(|e| anyhow::anyhow!("Failed to shutdown client: {:?}", e))?; anyhow::Ok(Profile::new(profile)) @@ -117,9 +111,7 @@ pub unsafe extern "C" fn ddog_prof_ProfilerManager_shutdown( #[no_mangle] // TODO: Do we want drop and shutdown to be separate functions? Or should it always be shutdown? pub unsafe extern "C" fn ddog_prof_ProfilerManager_drop( - handle: *mut Handle, + mut handle: *mut Handle, ) { - if let Some(handle) = handle.as_mut() { - let _ = handle.take(); - } + let _ = handle.take(); } From 9cf7c3c9bc993936203418d69c6e427e84e985f4 Mon Sep 17 00:00:00 2001 From: Daniel Schwartz-Narbonne Date: Thu, 5 Jun 2025 11:25:20 -0400 Subject: [PATCH 26/52] config --- datadog-profiling-ffi/src/manager/ffi_api.rs | 4 +++- .../src/manager/profiler_manager.rs | 14 +++++++++++++- datadog-profiling-ffi/src/manager/samples.rs | 7 ++++--- datadog-profiling-ffi/src/manager/tests.rs | 3 +++ 4 files changed, 23 insertions(+), 5 deletions(-) diff --git a/datadog-profiling-ffi/src/manager/ffi_api.rs b/datadog-profiling-ffi/src/manager/ffi_api.rs index e24d054b7a..53c6e5cc79 100644 --- a/datadog-profiling-ffi/src/manager/ffi_api.rs +++ b/datadog-profiling-ffi/src/manager/ffi_api.rs @@ -2,7 +2,7 @@ // SPDX-License-Identifier: Apache-2.0 use crate::manager::{ - profiler_manager::{ManagedSampleCallbacks, ProfilerManager}, + profiler_manager::{ManagedSampleCallbacks, ProfilerManager, ProfilerManagerConfig}, ManagedProfilerClient, }; use crate::profiles::datatypes::{Profile, ProfilePtrExt}; @@ -29,6 +29,7 @@ pub unsafe extern "C" fn ddog_prof_ProfilerManager_start( cpu_sampler_callback: extern "C" fn(*mut internal::Profile), upload_callback: extern "C" fn(*mut internal::Profile, &mut Option), sample_callbacks: ManagedSampleCallbacks, + config: ProfilerManagerConfig, ) -> FFIResult> { wrap_with_ffi_result!({ let internal_profile = *profile.take()?; @@ -37,6 +38,7 @@ pub unsafe extern "C" fn ddog_prof_ProfilerManager_start( cpu_sampler_callback, upload_callback, sample_callbacks, + config, )?; anyhow::Ok(Handle::from(client)) }) diff --git a/datadog-profiling-ffi/src/manager/profiler_manager.rs b/datadog-profiling-ffi/src/manager/profiler_manager.rs index c9d9e57f43..c8aa602c67 100644 --- a/datadog-profiling-ffi/src/manager/profiler_manager.rs +++ b/datadog-profiling-ffi/src/manager/profiler_manager.rs @@ -12,6 +12,17 @@ use super::client::ManagedProfilerClient; use super::samples::{ClientSampleChannels, ManagerSampleChannels, SendSample}; use crate::profiles::datatypes::Sample; +#[repr(C)] +pub struct ProfilerManagerConfig { + pub channel_depth: usize, +} + +impl Default for ProfilerManagerConfig { + fn default() -> Self { + Self { channel_depth: 10 } + } +} + #[repr(C)] #[derive(Clone)] pub struct ManagedSampleCallbacks { @@ -55,8 +66,9 @@ impl ProfilerManager { cpu_sampler_callback: extern "C" fn(*mut internal::Profile), upload_callback: extern "C" fn(*mut internal::Profile, &mut Option), sample_callbacks: ManagedSampleCallbacks, + config: ProfilerManagerConfig, ) -> Result { - let (client_channels, manager_channels) = ClientSampleChannels::new(); + let (client_channels, manager_channels) = ClientSampleChannels::new(config.channel_depth); let (shutdown_sender, shutdown_receiver) = crossbeam_channel::bounded(1); // For adaptive sampling, we need to be able to adjust this duration. Look into how to do diff --git a/datadog-profiling-ffi/src/manager/samples.rs b/datadog-profiling-ffi/src/manager/samples.rs index 7c34bfa9db..f3c3b40ff6 100644 --- a/datadog-profiling-ffi/src/manager/samples.rs +++ b/datadog-profiling-ffi/src/manager/samples.rs @@ -39,9 +39,10 @@ pub struct ManagerSampleChannels { } impl ClientSampleChannels { - pub fn new() -> (Self, ManagerSampleChannels) { - let (samples_sender, samples_receiver) = crossbeam_channel::bounded(10); - let (recycled_samples_sender, recycled_samples_receiver) = crossbeam_channel::bounded(10); + pub fn new(channel_depth: usize) -> (Self, ManagerSampleChannels) { + let (samples_sender, samples_receiver) = crossbeam_channel::bounded(channel_depth); + let (recycled_samples_sender, recycled_samples_receiver) = + crossbeam_channel::bounded(channel_depth); ( Self { samples_sender, diff --git a/datadog-profiling-ffi/src/manager/tests.rs b/datadog-profiling-ffi/src/manager/tests.rs index dcbf5310db..ee131f23e1 100644 --- a/datadog-profiling-ffi/src/manager/tests.rs +++ b/datadog-profiling-ffi/src/manager/tests.rs @@ -5,6 +5,7 @@ use datadog_profiling::internal; use tokio_util::sync::CancellationToken; use super::{ManagedSampleCallbacks, ProfilerManager}; +use crate::manager::profiler_manager::ProfilerManagerConfig; extern "C" fn test_cpu_sampler_callback(_: *mut datadog_profiling::internal::Profile) { println!("cpu sampler callback"); @@ -40,11 +41,13 @@ fn test_the_thing() { test_reset_callback, test_drop_callback, ); + let config = ProfilerManagerConfig::default(); let handle = ProfilerManager::start( profile, test_cpu_sampler_callback, test_upload_callback, sample_callbacks, + config, ) .expect("Failed to start profiler"); println!("start"); From 268f7cd128377f2e9546fd4c68f0693b5258152b Mon Sep 17 00:00:00 2001 From: Daniel Schwartz-Narbonne Date: Thu, 5 Jun 2025 11:28:59 -0400 Subject: [PATCH 27/52] more config --- .../src/manager/profiler_manager.rs | 15 +++++++++------ datadog-profiling-ffi/src/manager/tests.rs | 6 +++++- 2 files changed, 14 insertions(+), 7 deletions(-) diff --git a/datadog-profiling-ffi/src/manager/profiler_manager.rs b/datadog-profiling-ffi/src/manager/profiler_manager.rs index c8aa602c67..d3f18bbae9 100644 --- a/datadog-profiling-ffi/src/manager/profiler_manager.rs +++ b/datadog-profiling-ffi/src/manager/profiler_manager.rs @@ -15,11 +15,17 @@ use crate::profiles::datatypes::Sample; #[repr(C)] pub struct ProfilerManagerConfig { pub channel_depth: usize, + pub cpu_sampling_interval_ms: u64, + pub upload_interval_ms: u64, } impl Default for ProfilerManagerConfig { fn default() -> Self { - Self { channel_depth: 10 } + Self { + channel_depth: 10, + cpu_sampling_interval_ms: 1000, // 1 second + upload_interval_ms: 10000, // 10 seconds + } } } @@ -71,11 +77,8 @@ impl ProfilerManager { let (client_channels, manager_channels) = ClientSampleChannels::new(config.channel_depth); let (shutdown_sender, shutdown_receiver) = crossbeam_channel::bounded(1); - // For adaptive sampling, we need to be able to adjust this duration. Look into how to do - // this. - let cpu_ticker = tick(Duration::from_millis(100)); - // one second for testing, make this 1 minute in production - let upload_ticker = tick(Duration::from_secs(1)); + let cpu_ticker = tick(Duration::from_millis(config.cpu_sampling_interval_ms)); + let upload_ticker = tick(Duration::from_millis(config.upload_interval_ms)); let manager = Self { channels: manager_channels, diff --git a/datadog-profiling-ffi/src/manager/tests.rs b/datadog-profiling-ffi/src/manager/tests.rs index ee131f23e1..ee9dea34bf 100644 --- a/datadog-profiling-ffi/src/manager/tests.rs +++ b/datadog-profiling-ffi/src/manager/tests.rs @@ -41,7 +41,11 @@ fn test_the_thing() { test_reset_callback, test_drop_callback, ); - let config = ProfilerManagerConfig::default(); + let config = ProfilerManagerConfig { + channel_depth: 10, + cpu_sampling_interval_ms: 100, // 100ms for testing + upload_interval_ms: 1000, // 1 second for testing + }; let handle = ProfilerManager::start( profile, test_cpu_sampler_callback, From d5fc23825f037fc79558ed3d73ed0ec24053a2ac Mon Sep 17 00:00:00 2001 From: Daniel Schwartz-Narbonne Date: Thu, 5 Jun 2025 17:10:06 -0400 Subject: [PATCH 28/52] better types for the callbacks etc --- .../src/manager/profiler_manager.rs | 53 +++++-------- datadog-profiling-ffi/src/manager/tests.rs | 78 ++++++++++--------- ddcommon/src/unix_utils.rs | 5 +- 3 files changed, 65 insertions(+), 71 deletions(-) diff --git a/datadog-profiling-ffi/src/manager/profiler_manager.rs b/datadog-profiling-ffi/src/manager/profiler_manager.rs index d3f18bbae9..cc2dce9ce7 100644 --- a/datadog-profiling-ffi/src/manager/profiler_manager.rs +++ b/datadog-profiling-ffi/src/manager/profiler_manager.rs @@ -1,7 +1,4 @@ -use std::{ - ffi::c_void, - time::{Duration, Instant}, -}; +use std::time::{Duration, Instant}; use anyhow::Result; use crossbeam_channel::{select, tick, Receiver}; @@ -23,8 +20,8 @@ impl Default for ProfilerManagerConfig { fn default() -> Self { Self { channel_depth: 10, - cpu_sampling_interval_ms: 1000, // 1 second - upload_interval_ms: 10000, // 10 seconds + cpu_sampling_interval_ms: 100, // 100ms + upload_interval_ms: 60000, // 1 minute } } } @@ -33,18 +30,18 @@ impl Default for ProfilerManagerConfig { #[derive(Clone)] pub struct ManagedSampleCallbacks { // Static is probably the wrong type here, but worry about that later. - converter: extern "C" fn(*mut c_void) -> Sample<'static>, + converter: extern "C" fn(&SendSample) -> Sample, // Resets the sample for reuse. - reset: extern "C" fn(*mut c_void), + reset: extern "C" fn(&mut SendSample), // Called when a sample is dropped (not recycled) - drop: extern "C" fn(*mut c_void), + drop: extern "C" fn(SendSample), } impl ManagedSampleCallbacks { pub fn new( - converter: extern "C" fn(*mut c_void) -> Sample<'static>, - reset: extern "C" fn(*mut c_void), - drop: extern "C" fn(*mut c_void), + converter: extern "C" fn(&SendSample) -> Sample, + reset: extern "C" fn(&mut SendSample), + drop: extern "C" fn(SendSample), ) -> Self { Self { converter, @@ -105,22 +102,15 @@ impl ProfilerManager { &mut self, raw_sample: Result, ) -> Result<()> { - let data = raw_sample?.as_ptr(); - let sample = (self.sample_callbacks.converter)(data); - self.profile.add_sample(sample.try_into()?, None)?; - (self.sample_callbacks.reset)(data); - // SAFETY: The sample pointer is valid because it came from the samples channel - // and was just processed by the converter and reset callbacks. We have exclusive - // access to it since we're the only thread that can receive from the samples channel. - if self - .channels + let mut sample = raw_sample?; + let converted_sample = (self.sample_callbacks.converter)(&sample); + let add_result = self.profile.add_sample(converted_sample.try_into()?, None); + (self.sample_callbacks.reset)(&mut sample); + self.channels .recycled_samples_sender - .send(unsafe { SendSample::new(data) }) - .is_err() - { - (self.sample_callbacks.drop)(data); - } - Ok(()) + .send(sample) + .map_or_else(|e| (self.sample_callbacks.drop)(e.0), |_| ()); + add_result } fn handle_cpu_tick(&mut self) { @@ -145,21 +135,20 @@ impl ProfilerManager { fn handle_shutdown(&mut self) -> Result<()> { // Try to process any remaining samples before dropping them while let Ok(sample) = self.channels.samples_receiver.try_recv() { - let data = sample.as_ptr(); - let sample = (self.sample_callbacks.converter)(data); - if let Ok(converted_sample) = sample.try_into() { + let converted_sample = (self.sample_callbacks.converter)(&sample); + if let Ok(converted_sample) = converted_sample.try_into() { if let Err(e) = self.profile.add_sample(converted_sample, None) { eprintln!("Failed to add sample during shutdown: {}", e); } } else { eprintln!("Failed to convert sample during shutdown"); } - (self.sample_callbacks.drop)(data); + (self.sample_callbacks.drop)(sample); } // Drain any recycled samples while let Ok(sample) = self.channels.recycled_samples_receiver.try_recv() { - (self.sample_callbacks.drop)(sample.as_ptr()); + (self.sample_callbacks.drop)(sample); } // Cancel any ongoing upload and drop the token diff --git a/datadog-profiling-ffi/src/manager/tests.rs b/datadog-profiling-ffi/src/manager/tests.rs index ee9dea34bf..d3194ac795 100644 --- a/datadog-profiling-ffi/src/manager/tests.rs +++ b/datadog-profiling-ffi/src/manager/tests.rs @@ -1,60 +1,62 @@ -use std::{ffi::c_void, time::Duration}; +use std::ffi::c_void; use crate::profiles::datatypes::Sample; -use datadog_profiling::internal; +use datadog_profiling::internal::Profile; use tokio_util::sync::CancellationToken; use super::{ManagedSampleCallbacks, ProfilerManager}; use crate::manager::profiler_manager::ProfilerManagerConfig; +use crate::manager::samples::SendSample; +use ddcommon_ffi::Slice; -extern "C" fn test_cpu_sampler_callback(_: *mut datadog_profiling::internal::Profile) { - println!("cpu sampler callback"); -} -extern "C" fn test_upload_callback( - _: *mut datadog_profiling::internal::Profile, - _: &mut Option, -) { - println!("upload callback"); +extern "C" fn test_cpu_sampler_callback(_profile: *mut Profile) {} + +extern "C" fn test_upload_callback(_profile: *mut Profile, _token: &mut Option) { } -extern "C" fn test_sample_converter(_: *mut c_void) -> Sample<'static> { - println!("sample converter"); + +extern "C" fn test_converter(sample: &SendSample) -> Sample<'static> { + static VALUES: [i64; 1] = [42]; Sample { - locations: ddcommon_ffi::Slice::empty(), - values: ddcommon_ffi::Slice::empty(), - labels: ddcommon_ffi::Slice::empty(), + locations: Slice::empty(), + values: Slice::from(&VALUES[..]), + labels: Slice::empty(), } } -extern "C" fn test_reset_callback(_: *mut c_void) { - println!("reset callback"); -} -extern "C" fn test_drop_callback(_: *mut c_void) { - println!("drop callback"); -} + +extern "C" fn test_reset(_sample: &mut SendSample) {} + +extern "C" fn test_drop(_sample: SendSample) {} #[test] -fn test_the_thing() { - let sample_types = []; - let period = None; - let profile = internal::Profile::new(&sample_types, period); - let sample_callbacks = ManagedSampleCallbacks::new( - test_sample_converter, - test_reset_callback, - test_drop_callback, - ); +fn test_profiler_manager() { let config = ProfilerManagerConfig { - channel_depth: 10, - cpu_sampling_interval_ms: 100, // 100ms for testing - upload_interval_ms: 1000, // 1 second for testing + channel_depth: 1, + cpu_sampling_interval_ms: 100, // 100ms for faster testing + upload_interval_ms: 500, // 500ms for faster testing }; - let handle = ProfilerManager::start( + + let sample_callbacks = ManagedSampleCallbacks::new(test_converter, test_reset, test_drop); + + let profile = Profile::new(&[], None); + let client = ProfilerManager::start( profile, test_cpu_sampler_callback, test_upload_callback, sample_callbacks, config, ) - .expect("Failed to start profiler"); - println!("start"); - std::thread::sleep(Duration::from_secs(5)); - handle.shutdown().unwrap(); + .unwrap(); + + // Send a sample + let sample_ptr = Box::into_raw(Box::new(42)) as *mut c_void; + unsafe { + client.send_sample(sample_ptr).unwrap(); + } + + // Receive a recycled sample + let recycled = client.try_recv_recycled().unwrap(); + assert_eq!(unsafe { *(recycled as *const i32) }, 42); + + // Shutdown + let _profile = client.shutdown().unwrap(); } diff --git a/ddcommon/src/unix_utils.rs b/ddcommon/src/unix_utils.rs index 2fc4ed6617..c280756747 100644 --- a/ddcommon/src/unix_utils.rs +++ b/ddcommon/src/unix_utils.rs @@ -4,10 +4,13 @@ #![cfg(unix)] use anyhow::Context; -use libc::{_exit, execve, nfds_t, pid_t, poll, pollfd, EXIT_FAILURE, POLLHUP}; +#[cfg(target_os = "linux")] +use libc::pid_t; +use libc::{_exit, execve, nfds_t, poll, pollfd, EXIT_FAILURE, POLLHUP}; use nix::errno::Errno; use nix::sys::wait::{waitpid, WaitPidFlag, WaitStatus}; use nix::unistd::Pid; +#[cfg(target_os = "linux")] use std::io::{self, BufRead, BufReader}; use std::os::fd::IntoRawFd; use std::time::{Duration, Instant}; From eb872c2e34a39b5186f68bbbbff697f6595c7fc4 Mon Sep 17 00:00:00 2001 From: Daniel Schwartz-Narbonne Date: Thu, 5 Jun 2025 17:12:17 -0400 Subject: [PATCH 29/52] fix test --- datadog-profiling-ffi/src/manager/tests.rs | 3 +++ 1 file changed, 3 insertions(+) diff --git a/datadog-profiling-ffi/src/manager/tests.rs b/datadog-profiling-ffi/src/manager/tests.rs index d3194ac795..7536d982e8 100644 --- a/datadog-profiling-ffi/src/manager/tests.rs +++ b/datadog-profiling-ffi/src/manager/tests.rs @@ -53,6 +53,9 @@ fn test_profiler_manager() { client.send_sample(sample_ptr).unwrap(); } + // Give the manager thread time to process and recycle the sample + std::thread::sleep(std::time::Duration::from_millis(10)); + // Receive a recycled sample let recycled = client.try_recv_recycled().unwrap(); assert_eq!(unsafe { *(recycled as *const i32) }, 42); From 5748b32e84fd8c39f341189950042197c3f9cc1f Mon Sep 17 00:00:00 2001 From: Daniel Schwartz-Narbonne Date: Thu, 5 Jun 2025 17:30:39 -0400 Subject: [PATCH 30/52] better test --- datadog-profiling-ffi/src/manager/tests.rs | 88 ++++++++++++++++++---- 1 file changed, 72 insertions(+), 16 deletions(-) diff --git a/datadog-profiling-ffi/src/manager/tests.rs b/datadog-profiling-ffi/src/manager/tests.rs index 7536d982e8..818ddc9b66 100644 --- a/datadog-profiling-ffi/src/manager/tests.rs +++ b/datadog-profiling-ffi/src/manager/tests.rs @@ -1,6 +1,10 @@ use std::ffi::c_void; +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::sync::Arc; +use std::sync::LazyLock; use crate::profiles::datatypes::Sample; +use datadog_profiling::api::ValueType; use datadog_profiling::internal::Profile; use tokio_util::sync::CancellationToken; @@ -11,14 +15,62 @@ use ddcommon_ffi::Slice; extern "C" fn test_cpu_sampler_callback(_profile: *mut Profile) {} -extern "C" fn test_upload_callback(_profile: *mut Profile, _token: &mut Option) { +static UPLOAD_COUNT: AtomicUsize = AtomicUsize::new(0); +static SAMPLE_COUNT: AtomicUsize = AtomicUsize::new(0); + +extern "C" fn test_upload_callback( + profile: *mut Profile, + _token: &mut Option, +) { + let profile = unsafe { &*profile }; + let count = profile.only_for_testing_num_aggregated_samples(); + SAMPLE_COUNT.store(count, Ordering::SeqCst); + UPLOAD_COUNT.fetch_add(1, Ordering::SeqCst); +} + +#[repr(C)] +struct TestSample<'a> { + values: [i64; 1], + locations: [crate::profiles::datatypes::Location<'a>; 1], +} + +fn create_test_sample(value: i64) -> TestSample<'static> { + let function = crate::profiles::datatypes::Function { + name: match value { + 42 => "function_1", + 43 => "function_2", + 44 => "function_3", + 45 => "function_4", + 46 => "function_5", + _ => "unknown_function", + }.into(), + system_name: match value { + 42 => "function_1", + 43 => "function_2", + 44 => "function_3", + 45 => "function_4", + 46 => "function_5", + _ => "unknown_function", + }.into(), + filename: "test.rs".into(), + ..Default::default() + }; + + TestSample { + values: [value], + locations: [crate::profiles::datatypes::Location { + function, + ..Default::default() + }], + } } -extern "C" fn test_converter(sample: &SendSample) -> Sample<'static> { - static VALUES: [i64; 1] = [42]; +extern "C" fn test_converter(sample: &SendSample) -> Sample { + let test_sample = unsafe { &*(sample.as_ptr() as *const TestSample) }; + Sample { - locations: Slice::empty(), - values: Slice::from(&VALUES[..]), + locations: Slice::from(&test_sample.locations[..]), + values: Slice::from(&test_sample.values[..]), labels: Slice::empty(), } } @@ -30,14 +82,15 @@ extern "C" fn test_drop(_sample: SendSample) {} #[test] fn test_profiler_manager() { let config = ProfilerManagerConfig { - channel_depth: 1, + channel_depth: 10, cpu_sampling_interval_ms: 100, // 100ms for faster testing upload_interval_ms: 500, // 500ms for faster testing }; let sample_callbacks = ManagedSampleCallbacks::new(test_converter, test_reset, test_drop); - let profile = Profile::new(&[], None); + let sample_types = [ValueType::new("samples", "count")]; + let profile = Profile::new(&sample_types, None); let client = ProfilerManager::start( profile, test_cpu_sampler_callback, @@ -47,18 +100,21 @@ fn test_profiler_manager() { ) .unwrap(); - // Send a sample - let sample_ptr = Box::into_raw(Box::new(42)) as *mut c_void; - unsafe { - client.send_sample(sample_ptr).unwrap(); + // Send multiple samples + for i in 0..5 { + let test_sample = create_test_sample(42 + i as i64); + let sample_ptr = Box::into_raw(Box::new(test_sample)) as *mut c_void; + unsafe { + client.send_sample(sample_ptr).unwrap(); + } } - // Give the manager thread time to process and recycle the sample - std::thread::sleep(std::time::Duration::from_millis(10)); + // Give the manager thread time to process samples and trigger an upload + std::thread::sleep(std::time::Duration::from_millis(600)); - // Receive a recycled sample - let recycled = client.try_recv_recycled().unwrap(); - assert_eq!(unsafe { *(recycled as *const i32) }, 42); + // Verify samples were uploaded + assert_eq!(UPLOAD_COUNT.load(Ordering::SeqCst), 1); + assert_eq!(SAMPLE_COUNT.load(Ordering::SeqCst), 5); // Shutdown let _profile = client.shutdown().unwrap(); From 084377d2a2b8b005476ecac712d101ecddc6f3de Mon Sep 17 00:00:00 2001 From: Daniel Schwartz-Narbonne Date: Fri, 6 Jun 2025 13:42:45 -0400 Subject: [PATCH 31/52] fix tests, found lifetime issue --- Cargo.lock | 14 +- datadog-profiling-ffi/Cargo.toml | 5 + .../src/manager/profiler_manager.rs | 9 +- datadog-profiling-ffi/src/manager/tests.rs | 126 +++++++++++++++--- 4 files changed, 131 insertions(+), 23 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index a87a426819..63f7064d3c 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1597,7 +1597,7 @@ dependencies = [ "hyper 1.6.0", "hyper-multipart-rfc7578", "indexmap 2.6.0", - "lz4_flex", + "lz4_flex 0.9.5", "mime", "prost", "rustc-hash 1.1.0", @@ -1620,6 +1620,7 @@ dependencies = [ "datadog-library-config-ffi", "datadog-log-ffi", "datadog-profiling", + "datadog-profiling-protobuf", "ddcommon", "ddcommon-ffi", "ddtelemetry-ffi", @@ -1628,6 +1629,8 @@ dependencies = [ "http-body-util", "hyper 1.6.0", "libc", + "lz4_flex 0.11.3", + "prost", "serde_json", "symbolizer-ffi", "tokio", @@ -3362,6 +3365,15 @@ dependencies = [ "twox-hash", ] +[[package]] +name = "lz4_flex" +version = "0.11.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "75761162ae2b0e580d7e7c390558127e5f01b4194debd6221fd8c207fc80e3f5" +dependencies = [ + "twox-hash", +] + [[package]] name = "manual_future" version = "0.1.1" diff --git a/datadog-profiling-ffi/Cargo.toml b/datadog-profiling-ffi/Cargo.toml index b29160d560..b57eb26b51 100644 --- a/datadog-profiling-ffi/Cargo.toml +++ b/datadog-profiling-ffi/Cargo.toml @@ -53,3 +53,8 @@ serde_json = { version = "1.0" } symbolizer-ffi = { path = "../symbolizer-ffi", optional = true, default-features = false } tokio = { version = "1.36", features = ["sync", "rt"] } tokio-util = "0.7.1" + +[dev-dependencies] +lz4_flex = "0.11.3" +prost = "0.13.5" +datadog-profiling-protobuf = { path = "../datadog-profiling-protobuf" } diff --git a/datadog-profiling-ffi/src/manager/profiler_manager.rs b/datadog-profiling-ffi/src/manager/profiler_manager.rs index cc2dce9ce7..8df3d399c4 100644 --- a/datadog-profiling-ffi/src/manager/profiler_manager.rs +++ b/datadog-profiling-ffi/src/manager/profiler_manager.rs @@ -118,7 +118,7 @@ impl ProfilerManager { } fn handle_upload_tick(&mut self) -> Result<()> { - let mut old_profile = self.profile.reset_and_return_previous()?; + let old_profile = self.profile.reset_and_return_previous()?; let upload_callback = self.upload_callback; // Create a new cancellation token for this upload let token = CancellationToken::new(); @@ -126,8 +126,10 @@ impl ProfilerManager { self.cancellation_token = Some(token.clone()); let mut cancellation_token = Some(token); std::thread::spawn(move || { - (upload_callback)(&mut old_profile, &mut cancellation_token); - // TODO: make sure we cleanup the profile. + let mut profile = old_profile; + (upload_callback)(&mut profile, &mut cancellation_token); + // The profile is consumed by the callback, so we don't drop it here + std::mem::forget(profile); }); Ok(()) } @@ -155,6 +157,7 @@ impl ProfilerManager { if let Some(token) = self.cancellation_token.take() { token.cancel(); } + // Ok(()) } diff --git a/datadog-profiling-ffi/src/manager/tests.rs b/datadog-profiling-ffi/src/manager/tests.rs index 818ddc9b66..e4fbc0f809 100644 --- a/datadog-profiling-ffi/src/manager/tests.rs +++ b/datadog-profiling-ffi/src/manager/tests.rs @@ -1,7 +1,5 @@ use std::ffi::c_void; use std::sync::atomic::{AtomicUsize, Ordering}; -use std::sync::Arc; -use std::sync::LazyLock; use crate::profiles::datatypes::Sample; use datadog_profiling::api::ValueType; @@ -11,21 +9,23 @@ use tokio_util::sync::CancellationToken; use super::{ManagedSampleCallbacks, ProfilerManager}; use crate::manager::profiler_manager::ProfilerManagerConfig; use crate::manager::samples::SendSample; +use datadog_profiling_protobuf::prost_impls::Profile as ProstProfile; use ddcommon_ffi::Slice; +use prost::Message; extern "C" fn test_cpu_sampler_callback(_profile: *mut Profile) {} static UPLOAD_COUNT: AtomicUsize = AtomicUsize::new(0); -static SAMPLE_COUNT: AtomicUsize = AtomicUsize::new(0); -extern "C" fn test_upload_callback( - profile: *mut Profile, - _token: &mut Option, -) { +extern "C" fn test_upload_callback(profile: *mut Profile, _token: &mut Option) { let profile = unsafe { &*profile }; - let count = profile.only_for_testing_num_aggregated_samples(); - SAMPLE_COUNT.store(count, Ordering::SeqCst); - UPLOAD_COUNT.fetch_add(1, Ordering::SeqCst); + let upload_count = UPLOAD_COUNT.fetch_add(1, Ordering::SeqCst); + + // On the first upload (when upload_count is 0), verify the samples + if upload_count == 0 { + let profile = unsafe { std::ptr::read(profile) }; + verify_samples(profile); + } } #[repr(C)] @@ -43,7 +43,8 @@ fn create_test_sample(value: i64) -> TestSample<'static> { 45 => "function_4", 46 => "function_5", _ => "unknown_function", - }.into(), + } + .into(), system_name: match value { 42 => "function_1", 43 => "function_2", @@ -51,11 +52,12 @@ fn create_test_sample(value: i64) -> TestSample<'static> { 45 => "function_4", 46 => "function_5", _ => "unknown_function", - }.into(), + } + .into(), filename: "test.rs".into(), ..Default::default() }; - + TestSample { values: [value], locations: [crate::profiles::datatypes::Location { @@ -67,7 +69,7 @@ fn create_test_sample(value: i64) -> TestSample<'static> { extern "C" fn test_converter(sample: &SendSample) -> Sample { let test_sample = unsafe { &*(sample.as_ptr() as *const TestSample) }; - + Sample { locations: Slice::from(&test_sample.locations[..]), values: Slice::from(&test_sample.values[..]), @@ -75,9 +77,94 @@ extern "C" fn test_converter(sample: &SendSample) -> Sample { } } -extern "C" fn test_reset(_sample: &mut SendSample) {} +extern "C" fn test_reset(sample: &mut SendSample) { + let test_sample = unsafe { &mut *(sample.as_ptr() as *mut TestSample) }; + test_sample.values[0] = 0; + test_sample.locations[0] = crate::profiles::datatypes::Location { + function: crate::profiles::datatypes::Function { + name: "".into(), + system_name: "".into(), + filename: "".into(), + ..Default::default() + }, + ..Default::default() + }; +} + +extern "C" fn test_drop(sample: SendSample) { + let test_sample = unsafe { Box::from_raw(sample.as_ptr() as *mut TestSample) }; + // Box will be dropped here, freeing the memory +} + +fn decode_pprof(encoded: &[u8]) -> ProstProfile { + let mut decoder = lz4_flex::frame::FrameDecoder::new(encoded); + let mut buf = Vec::new(); + use std::io::Read; + decoder.read_to_end(&mut buf).unwrap(); + ProstProfile::decode(buf.as_slice()).unwrap() +} + +fn roundtrip_to_pprof(profile: datadog_profiling::internal::Profile) -> ProstProfile { + let encoded = profile.serialize_into_compressed_pprof(None, None).unwrap(); + decode_pprof(&encoded.buffer) +} -extern "C" fn test_drop(_sample: SendSample) {} +fn string_table_fetch(profile: &ProstProfile, id: i64) -> &str { + profile + .string_table + .get(id as usize) + .map(|s| s.as_str()) + .unwrap_or("") +} + +fn verify_samples(profile: datadog_profiling::internal::Profile) { + let pprof = roundtrip_to_pprof(profile); + println!("Number of samples in profile: {}", pprof.samples.len()); + println!( + "Sample values: {:?}", + pprof + .samples + .iter() + .map(|s| s.values[0]) + .collect::>() + ); + assert_eq!(pprof.samples.len(), 5); + + // Sort samples by their first value + let mut samples = pprof.samples.clone(); + samples.sort_by_key(|s| s.values[0]); + + // Check each sample's value and function name + for (i, sample) in samples.iter().enumerate() { + let value = 42 + i as i64; + assert_eq!(sample.values[0], value); + + // Get the function name from the location + let location_id = sample.location_ids[0]; + let location = pprof + .locations + .iter() + .find(|l| l.id == location_id) + .unwrap(); + let function_id = location.lines[0].function_id; + let function = pprof + .functions + .iter() + .find(|f| f.id == function_id) + .unwrap(); + let function_name = string_table_fetch(&pprof, function.name); + + let expected_function = match value { + 42 => "function_1", + 43 => "function_2", + 44 => "function_3", + 45 => "function_4", + 46 => "function_5", + _ => "unknown_function", + }; + assert_eq!(function_name, expected_function); + } +} #[test] fn test_profiler_manager() { @@ -114,8 +201,9 @@ fn test_profiler_manager() { // Verify samples were uploaded assert_eq!(UPLOAD_COUNT.load(Ordering::SeqCst), 1); - assert_eq!(SAMPLE_COUNT.load(Ordering::SeqCst), 5); - // Shutdown - let _profile = client.shutdown().unwrap(); + // Get the profile and verify it has no samples (they were consumed by the upload) + let profile = client.shutdown().unwrap(); + let pprof = roundtrip_to_pprof(profile); + assert_eq!(pprof.samples.len(), 0); } From 9cd6bd709956749b7e7d3ab351a285ec1ae5d813 Mon Sep 17 00:00:00 2001 From: Daniel Schwartz-Narbonne Date: Fri, 6 Jun 2025 15:40:21 -0400 Subject: [PATCH 32/52] use Handle for the profile callback --- datadog-profiling-ffi/src/manager/ffi_api.rs | 2 +- .../src/manager/profiler_manager.rs | 14 ++++++++------ datadog-profiling-ffi/src/manager/tests.rs | 9 ++++++--- 3 files changed, 15 insertions(+), 10 deletions(-) diff --git a/datadog-profiling-ffi/src/manager/ffi_api.rs b/datadog-profiling-ffi/src/manager/ffi_api.rs index 53c6e5cc79..47620e5454 100644 --- a/datadog-profiling-ffi/src/manager/ffi_api.rs +++ b/datadog-profiling-ffi/src/manager/ffi_api.rs @@ -27,7 +27,7 @@ use tokio_util::sync::CancellationToken; pub unsafe extern "C" fn ddog_prof_ProfilerManager_start( profile: *mut Profile, cpu_sampler_callback: extern "C" fn(*mut internal::Profile), - upload_callback: extern "C" fn(*mut internal::Profile, &mut Option), + upload_callback: extern "C" fn(*mut Handle, &mut Option), sample_callbacks: ManagedSampleCallbacks, config: ProfilerManagerConfig, ) -> FFIResult> { diff --git a/datadog-profiling-ffi/src/manager/profiler_manager.rs b/datadog-profiling-ffi/src/manager/profiler_manager.rs index 8df3d399c4..871ebe3e95 100644 --- a/datadog-profiling-ffi/src/manager/profiler_manager.rs +++ b/datadog-profiling-ffi/src/manager/profiler_manager.rs @@ -3,6 +3,7 @@ use std::time::{Duration, Instant}; use anyhow::Result; use crossbeam_channel::{select, tick, Receiver}; use datadog_profiling::internal; +use ddcommon_ffi::Handle; use tokio_util::sync::CancellationToken; use super::client::ManagedProfilerClient; @@ -58,7 +59,7 @@ pub struct ProfilerManager { shutdown_receiver: Receiver<()>, profile: internal::Profile, cpu_sampler_callback: extern "C" fn(*mut internal::Profile), - upload_callback: extern "C" fn(*mut internal::Profile, &mut Option), + upload_callback: extern "C" fn(*mut Handle, &mut Option), sample_callbacks: ManagedSampleCallbacks, cancellation_token: Option, } @@ -67,7 +68,10 @@ impl ProfilerManager { pub fn start( profile: internal::Profile, cpu_sampler_callback: extern "C" fn(*mut internal::Profile), - upload_callback: extern "C" fn(*mut internal::Profile, &mut Option), + upload_callback: extern "C" fn( + *mut Handle, + &mut Option, + ), sample_callbacks: ManagedSampleCallbacks, config: ProfilerManagerConfig, ) -> Result { @@ -126,10 +130,8 @@ impl ProfilerManager { self.cancellation_token = Some(token.clone()); let mut cancellation_token = Some(token); std::thread::spawn(move || { - let mut profile = old_profile; - (upload_callback)(&mut profile, &mut cancellation_token); - // The profile is consumed by the callback, so we don't drop it here - std::mem::forget(profile); + let mut handle = Handle::from(old_profile); + (upload_callback)(&mut handle, &mut cancellation_token); }); Ok(()) } diff --git a/datadog-profiling-ffi/src/manager/tests.rs b/datadog-profiling-ffi/src/manager/tests.rs index e4fbc0f809..6851e5cdb9 100644 --- a/datadog-profiling-ffi/src/manager/tests.rs +++ b/datadog-profiling-ffi/src/manager/tests.rs @@ -4,6 +4,7 @@ use std::sync::atomic::{AtomicUsize, Ordering}; use crate::profiles::datatypes::Sample; use datadog_profiling::api::ValueType; use datadog_profiling::internal::Profile; +use ddcommon_ffi::{Handle, ToInner}; use tokio_util::sync::CancellationToken; use super::{ManagedSampleCallbacks, ProfilerManager}; @@ -17,13 +18,15 @@ extern "C" fn test_cpu_sampler_callback(_profile: *mut Profile) {} static UPLOAD_COUNT: AtomicUsize = AtomicUsize::new(0); -extern "C" fn test_upload_callback(profile: *mut Profile, _token: &mut Option) { - let profile = unsafe { &*profile }; +extern "C" fn test_upload_callback( + profile: *mut Handle, + _token: &mut Option, +) { let upload_count = UPLOAD_COUNT.fetch_add(1, Ordering::SeqCst); // On the first upload (when upload_count is 0), verify the samples if upload_count == 0 { - let profile = unsafe { std::ptr::read(profile) }; + let profile = unsafe { *(*profile).take().unwrap() }; verify_samples(profile); } } From 4a37931de75f90bd662ee201a7fbed66766c3a31 Mon Sep 17 00:00:00 2001 From: Daniel Schwartz-Narbonne Date: Fri, 6 Jun 2025 16:20:12 -0400 Subject: [PATCH 33/52] better upload handling --- .../src/manager/profiler_manager.rs | 52 +++++++++++++------ 1 file changed, 35 insertions(+), 17 deletions(-) diff --git a/datadog-profiling-ffi/src/manager/profiler_manager.rs b/datadog-profiling-ffi/src/manager/profiler_manager.rs index 871ebe3e95..5e453e87db 100644 --- a/datadog-profiling-ffi/src/manager/profiler_manager.rs +++ b/datadog-profiling-ffi/src/manager/profiler_manager.rs @@ -1,7 +1,8 @@ +use std::thread::JoinHandle; use std::time::{Duration, Instant}; use anyhow::Result; -use crossbeam_channel::{select, tick, Receiver}; +use crossbeam_channel::{select, tick, Receiver, Sender}; use datadog_profiling::internal; use ddcommon_ffi::Handle; use tokio_util::sync::CancellationToken; @@ -61,7 +62,9 @@ pub struct ProfilerManager { cpu_sampler_callback: extern "C" fn(*mut internal::Profile), upload_callback: extern "C" fn(*mut Handle, &mut Option), sample_callbacks: ManagedSampleCallbacks, - cancellation_token: Option, + cancellation_token: CancellationToken, + upload_sender: Sender, + upload_thread: JoinHandle<()>, } impl ProfilerManager { @@ -77,10 +80,23 @@ impl ProfilerManager { ) -> Result { let (client_channels, manager_channels) = ClientSampleChannels::new(config.channel_depth); let (shutdown_sender, shutdown_receiver) = crossbeam_channel::bounded(1); + let (upload_sender, upload_receiver) = crossbeam_channel::bounded(2); let cpu_ticker = tick(Duration::from_millis(config.cpu_sampling_interval_ms)); let upload_ticker = tick(Duration::from_millis(config.upload_interval_ms)); + // Create a single cancellation token for all uploads + let cancellation_token = CancellationToken::new(); + + // Spawn the upload thread + let mut token = Some(cancellation_token.clone()); + let upload_thread = std::thread::spawn(move || { + while let Ok(profile) = upload_receiver.recv() { + let mut handle = Handle::from(profile); + (upload_callback)(&mut handle, &mut token); + } + }); + let manager = Self { channels: manager_channels, cpu_ticker, @@ -90,7 +106,9 @@ impl ProfilerManager { cpu_sampler_callback, upload_callback, sample_callbacks, - cancellation_token: None, + cancellation_token, + upload_sender, + upload_thread, }; let handle = std::thread::spawn(move || manager.main()); @@ -123,16 +141,9 @@ impl ProfilerManager { fn handle_upload_tick(&mut self) -> Result<()> { let old_profile = self.profile.reset_and_return_previous()?; - let upload_callback = self.upload_callback; - // Create a new cancellation token for this upload - let token = CancellationToken::new(); - // Store a clone of the token in the manager - self.cancellation_token = Some(token.clone()); - let mut cancellation_token = Some(token); - std::thread::spawn(move || { - let mut handle = Handle::from(old_profile); - (upload_callback)(&mut handle, &mut cancellation_token); - }); + self.upload_sender + .send(old_profile) + .map_err(|e| anyhow::anyhow!("Failed to send profile for upload: {}", e))?; Ok(()) } @@ -155,11 +166,18 @@ impl ProfilerManager { (self.sample_callbacks.drop)(sample); } - // Cancel any ongoing upload and drop the token - if let Some(token) = self.cancellation_token.take() { - token.cancel(); + // Cancel any ongoing upload + self.cancellation_token.cancel(); + + // Take ownership of the upload thread and sender + let upload_thread = std::mem::replace(&mut self.upload_thread, std::thread::spawn(|| {})); + let _ = std::mem::replace(&mut self.upload_sender, crossbeam_channel::bounded(1).0); + + // Wait for the upload thread to finish + if let Err(e) = upload_thread.join() { + eprintln!("Error joining upload thread: {:?}", e); } - // + Ok(()) } From 6cc6abc5fb04d1b0c9f8f3303f0d616232455f53 Mon Sep 17 00:00:00 2001 From: Daniel Schwartz-Narbonne Date: Tue, 10 Jun 2025 15:24:18 -0400 Subject: [PATCH 34/52] handle shutdown should just consume the object --- .../src/manager/profiler_manager.rs | 12 +++++------- 1 file changed, 5 insertions(+), 7 deletions(-) diff --git a/datadog-profiling-ffi/src/manager/profiler_manager.rs b/datadog-profiling-ffi/src/manager/profiler_manager.rs index 5e453e87db..e857b0a05f 100644 --- a/datadog-profiling-ffi/src/manager/profiler_manager.rs +++ b/datadog-profiling-ffi/src/manager/profiler_manager.rs @@ -147,7 +147,7 @@ impl ProfilerManager { Ok(()) } - fn handle_shutdown(&mut self) -> Result<()> { + fn handle_shutdown(mut self) -> Result { // Try to process any remaining samples before dropping them while let Ok(sample) = self.channels.samples_receiver.try_recv() { let converted_sample = (self.sample_callbacks.converter)(&sample); @@ -170,15 +170,15 @@ impl ProfilerManager { self.cancellation_token.cancel(); // Take ownership of the upload thread and sender - let upload_thread = std::mem::replace(&mut self.upload_thread, std::thread::spawn(|| {})); - let _ = std::mem::replace(&mut self.upload_sender, crossbeam_channel::bounded(1).0); + let upload_thread = self.upload_thread; + drop(self.upload_sender); // Wait for the upload thread to finish if let Err(e) = upload_thread.join() { eprintln!("Error joining upload thread: {:?}", e); } - Ok(()) + Ok(self.profile) } /// # Safety @@ -199,11 +199,9 @@ impl ProfilerManager { .map_err(|e| eprintln!("Failed to handle upload: {}", e)); }, recv(self.shutdown_receiver) -> _ => { - self.handle_shutdown()?; - break; + return self.handle_shutdown(); }, } } - Ok(self.profile) } } From 4cb37b5893c87b4d8682177c9562e6bb28657278 Mon Sep 17 00:00:00 2001 From: Daniel Schwartz-Narbonne Date: Wed, 11 Jun 2025 15:58:23 -0400 Subject: [PATCH 35/52] code cleanup --- .../src/manager/profiler_manager.rs | 18 +++++++----------- 1 file changed, 7 insertions(+), 11 deletions(-) diff --git a/datadog-profiling-ffi/src/manager/profiler_manager.rs b/datadog-profiling-ffi/src/manager/profiler_manager.rs index e857b0a05f..cf1b721e1f 100644 --- a/datadog-profiling-ffi/src/manager/profiler_manager.rs +++ b/datadog-profiling-ffi/src/manager/profiler_manager.rs @@ -148,20 +148,16 @@ impl ProfilerManager { } fn handle_shutdown(mut self) -> Result { - // Try to process any remaining samples before dropping them + // Process any remaining samples while let Ok(sample) = self.channels.samples_receiver.try_recv() { let converted_sample = (self.sample_callbacks.converter)(&sample); - if let Ok(converted_sample) = converted_sample.try_into() { - if let Err(e) = self.profile.add_sample(converted_sample, None) { - eprintln!("Failed to add sample during shutdown: {}", e); - } - } else { - eprintln!("Failed to convert sample during shutdown"); + if let Ok(s) = converted_sample.try_into() { + let _ = self.profile.add_sample(s, None); } (self.sample_callbacks.drop)(sample); } - // Drain any recycled samples + // Drain recycled samples while let Ok(sample) = self.channels.recycled_samples_receiver.try_recv() { (self.sample_callbacks.drop)(sample); } @@ -169,12 +165,12 @@ impl ProfilerManager { // Cancel any ongoing upload self.cancellation_token.cancel(); - // Take ownership of the upload thread and sender - let upload_thread = self.upload_thread; + // Drop the sender to signal the upload thread that no more messages will be sent + // This is necessary to allow the upload thread to exit its message processing loop drop(self.upload_sender); // Wait for the upload thread to finish - if let Err(e) = upload_thread.join() { + if let Err(e) = self.upload_thread.join() { eprintln!("Error joining upload thread: {:?}", e); } From a5ac0e51b3f0aa05c8386afacf23c1ec65f60133 Mon Sep 17 00:00:00 2001 From: Daniel Schwartz-Narbonne Date: Fri, 13 Jun 2025 15:36:28 -0400 Subject: [PATCH 36/52] don't reformat vs main --- datadog-trace-protobuf/src/remoteconfig.rs | 30 ++++++++-------------- 1 file changed, 10 insertions(+), 20 deletions(-) diff --git a/datadog-trace-protobuf/src/remoteconfig.rs b/datadog-trace-protobuf/src/remoteconfig.rs index fdc89ee871..a8a4b66524 100644 --- a/datadog-trace-protobuf/src/remoteconfig.rs +++ b/datadog-trace-protobuf/src/remoteconfig.rs @@ -3,8 +3,7 @@ use serde::{Deserialize, Serialize}; // This file is @generated by prost-build. -#[derive(Deserialize, Serialize)] -#[derive(Clone, PartialEq, ::prost::Message)] +#[derive(Deserialize, Serialize, Clone, PartialEq, ::prost::Message)] pub struct File { #[prost(string, tag = "1")] pub path: ::prost::alloc::string::String, @@ -12,8 +11,7 @@ pub struct File { #[serde(with = "serde_bytes")] pub raw: ::prost::alloc::vec::Vec, } -#[derive(Deserialize, Serialize)] -#[derive(Clone, PartialEq, ::prost::Message)] +#[derive(Deserialize, Serialize, Clone, PartialEq, ::prost::Message)] pub struct Client { #[prost(message, optional, tag = "1")] pub state: ::core::option::Option, @@ -35,8 +33,7 @@ pub struct Client { #[prost(bytes = "vec", tag = "11")] pub capabilities: ::prost::alloc::vec::Vec, } -#[derive(Deserialize, Serialize)] -#[derive(Clone, PartialEq, ::prost::Message)] +#[derive(Deserialize, Serialize, Clone, PartialEq, ::prost::Message)] pub struct ClientTracer { #[prost(string, tag = "1")] pub runtime_id: ::prost::alloc::string::String, @@ -55,8 +52,7 @@ pub struct ClientTracer { #[prost(string, repeated, tag = "7")] pub tags: ::prost::alloc::vec::Vec<::prost::alloc::string::String>, } -#[derive(Deserialize, Serialize)] -#[derive(Clone, PartialEq, ::prost::Message)] +#[derive(Deserialize, Serialize, Clone, PartialEq, ::prost::Message)] pub struct ClientAgent { #[prost(string, tag = "1")] pub name: ::prost::alloc::string::String, @@ -69,8 +65,7 @@ pub struct ClientAgent { #[prost(string, repeated, tag = "5")] pub cws_workloads: ::prost::alloc::vec::Vec<::prost::alloc::string::String>, } -#[derive(Deserialize, Serialize)] -#[derive(Clone, PartialEq, ::prost::Message)] +#[derive(Deserialize, Serialize, Clone, PartialEq, ::prost::Message)] pub struct ConfigState { #[prost(string, tag = "1")] pub id: ::prost::alloc::string::String, @@ -83,8 +78,7 @@ pub struct ConfigState { #[prost(string, tag = "5")] pub apply_error: ::prost::alloc::string::String, } -#[derive(Deserialize, Serialize)] -#[derive(Clone, PartialEq, ::prost::Message)] +#[derive(Deserialize, Serialize, Clone, PartialEq, ::prost::Message)] pub struct ClientState { #[prost(uint64, tag = "1")] pub root_version: u64, @@ -99,16 +93,14 @@ pub struct ClientState { #[prost(bytes = "vec", tag = "6")] pub backend_client_state: ::prost::alloc::vec::Vec, } -#[derive(Deserialize, Serialize)] -#[derive(Clone, PartialEq, ::prost::Message)] +#[derive(Deserialize, Serialize, Clone, PartialEq, ::prost::Message)] pub struct TargetFileHash { #[prost(string, tag = "1")] pub algorithm: ::prost::alloc::string::String, #[prost(string, tag = "3")] pub hash: ::prost::alloc::string::String, } -#[derive(Deserialize, Serialize)] -#[derive(Clone, PartialEq, ::prost::Message)] +#[derive(Deserialize, Serialize, Clone, PartialEq, ::prost::Message)] pub struct TargetFileMeta { #[prost(string, tag = "1")] pub path: ::prost::alloc::string::String, @@ -117,16 +109,14 @@ pub struct TargetFileMeta { #[prost(message, repeated, tag = "3")] pub hashes: ::prost::alloc::vec::Vec, } -#[derive(Deserialize, Serialize)] -#[derive(Clone, PartialEq, ::prost::Message)] +#[derive(Deserialize, Serialize, Clone, PartialEq, ::prost::Message)] pub struct ClientGetConfigsRequest { #[prost(message, optional, tag = "1")] pub client: ::core::option::Option, #[prost(message, repeated, tag = "2")] pub cached_target_files: ::prost::alloc::vec::Vec, } -#[derive(Deserialize, Serialize)] -#[derive(Clone, PartialEq, ::prost::Message)] +#[derive(Deserialize, Serialize, Clone, PartialEq, ::prost::Message)] pub struct ClientGetConfigsResponse { #[prost(bytes = "vec", repeated, tag = "1")] #[serde(with = "crate::serde")] From 955049177770b87e36ed83573dba6950316e4d2a Mon Sep 17 00:00:00 2001 From: Daniel Schwartz-Narbonne Date: Fri, 13 Jun 2025 16:15:20 -0400 Subject: [PATCH 37/52] start at handling forks --- datadog-profiling-ffi/src/manager/ffi_api.rs | 12 ++- .../src/manager/fork_handler.rs | 101 ++++++++++++++++++ datadog-profiling-ffi/src/manager/mod.rs | 1 + .../src/manager/profiler_manager.rs | 44 ++++---- datadog-profiling-ffi/src/manager/tests.rs | 13 ++- datadog-trace-protobuf/src/remoteconfig.rs | 30 ++++-- 6 files changed, 161 insertions(+), 40 deletions(-) create mode 100644 datadog-profiling-ffi/src/manager/fork_handler.rs diff --git a/datadog-profiling-ffi/src/manager/ffi_api.rs b/datadog-profiling-ffi/src/manager/ffi_api.rs index 47620e5454..335b3502cf 100644 --- a/datadog-profiling-ffi/src/manager/ffi_api.rs +++ b/datadog-profiling-ffi/src/manager/ffi_api.rs @@ -2,7 +2,9 @@ // SPDX-License-Identifier: Apache-2.0 use crate::manager::{ - profiler_manager::{ManagedSampleCallbacks, ProfilerManager, ProfilerManagerConfig}, + profiler_manager::{ + ManagedSampleCallbacks, ManagerCallbacks, ProfilerManager, ProfilerManagerConfig, + }, ManagedProfilerClient, }; use crate::profiles::datatypes::{Profile, ProfilePtrExt}; @@ -35,9 +37,11 @@ pub unsafe extern "C" fn ddog_prof_ProfilerManager_start( let internal_profile = *profile.take()?; let client: ManagedProfilerClient = ProfilerManager::start( internal_profile, - cpu_sampler_callback, - upload_callback, - sample_callbacks, + ManagerCallbacks { + cpu_sampler_callback, + upload_callback, + sample_callbacks, + }, config, )?; anyhow::Ok(Handle::from(client)) diff --git a/datadog-profiling-ffi/src/manager/fork_handler.rs b/datadog-profiling-ffi/src/manager/fork_handler.rs new file mode 100644 index 0000000000..6abc7c95b1 --- /dev/null +++ b/datadog-profiling-ffi/src/manager/fork_handler.rs @@ -0,0 +1,101 @@ +#![allow(clippy::unwrap_used)] + +use super::client::ManagedProfilerClient; +use super::profiler_manager::{ + ManagedSampleCallbacks, ManagerCallbacks, ProfilerManager, ProfilerManagerConfig, +}; +use anyhow::{Context, Result}; +use datadog_profiling::internal::Profile; +use std::sync::Mutex; + +// Global state for the profile manager +static MANAGER: Mutex> = Mutex::new(None); +static PROFILE: Mutex> = Mutex::new(None); +static CONFIG: Mutex> = Mutex::new(None); +static CALLBACKS: Mutex> = Mutex::new(None); + +/// Stores the current profile manager state globally. +/// This should be called before forking to ensure clean state. +pub fn store_manager_state( + manager: ProfilerManager, + config: ProfilerManagerConfig, + cpu_sampler_callback: extern "C" fn(*mut Profile), + upload_callback: extern "C" fn( + *mut ddcommon_ffi::Handle, + &mut Option, + ), + sample_callbacks: ManagedSampleCallbacks, +) -> Result<()> { + // Store the manager + *MANAGER.lock().unwrap() = Some(manager); + + // Store the config and callbacks separately + *CONFIG.lock().unwrap() = Some(config); + *CALLBACKS.lock().unwrap() = Some(ManagerCallbacks { + cpu_sampler_callback, + upload_callback, + sample_callbacks, + }); + + Ok(()) +} + +/// Shuts down the stored profile manager and stores its profile. +/// This should be called before forking to ensure clean state. +pub fn shutdown_stored_manager() -> Result<()> { + let manager = MANAGER + .lock() + .unwrap() + .take() + .context("No profile manager stored")?; + + let profile = manager.handle_shutdown()?; + + // Store the profile + *PROFILE.lock().unwrap() = Some(profile); + + Ok(()) +} + +/// Restarts the profile manager in the parent process with the stored profile. +/// This should be called after fork in the parent process. +pub fn restart_manager_in_parent() -> Result { + let profile = PROFILE + .lock() + .unwrap() + .take() + .context("No profile stored")?; + + let config = CONFIG.lock().unwrap().take().context("No config stored")?; + + let callbacks = CALLBACKS + .lock() + .unwrap() + .take() + .context("No callbacks stored")?; + + ProfilerManager::start(profile, callbacks, config) +} + +/// Restarts the profile manager in the child process with a fresh profile. +/// This should be called after fork in the child process. +pub fn restart_manager_in_child() -> Result { + let config = CONFIG.lock().unwrap().take().context("No config stored")?; + + let callbacks = CALLBACKS + .lock() + .unwrap() + .take() + .context("No callbacks stored")?; + + let mut profile = PROFILE + .lock() + .unwrap() + .take() + .context("No profile stored")?; + + // Reset the profile, discarding the previous one + let _ = profile.reset_and_return_previous()?; + + ProfilerManager::start(profile, callbacks, config) +} diff --git a/datadog-profiling-ffi/src/manager/mod.rs b/datadog-profiling-ffi/src/manager/mod.rs index a9c81ac3ba..c4529f8c18 100644 --- a/datadog-profiling-ffi/src/manager/mod.rs +++ b/datadog-profiling-ffi/src/manager/mod.rs @@ -7,6 +7,7 @@ mod client; mod ffi_api; mod ffi_utils; +mod fork_handler; mod profiler_manager; mod samples; #[cfg(test)] diff --git a/datadog-profiling-ffi/src/manager/profiler_manager.rs b/datadog-profiling-ffi/src/manager/profiler_manager.rs index cf1b721e1f..d0fd78f491 100644 --- a/datadog-profiling-ffi/src/manager/profiler_manager.rs +++ b/datadog-profiling-ffi/src/manager/profiler_manager.rs @@ -11,6 +11,14 @@ use super::client::ManagedProfilerClient; use super::samples::{ClientSampleChannels, ManagerSampleChannels, SendSample}; use crate::profiles::datatypes::Sample; +/// Holds the callbacks needed to restart the profile manager +pub struct ManagerCallbacks { + pub cpu_sampler_callback: extern "C" fn(*mut internal::Profile), + pub upload_callback: + extern "C" fn(*mut Handle, &mut Option), + pub sample_callbacks: ManagedSampleCallbacks, +} + #[repr(C)] pub struct ProfilerManagerConfig { pub channel_depth: usize, @@ -59,9 +67,7 @@ pub struct ProfilerManager { upload_ticker: Receiver, shutdown_receiver: Receiver<()>, profile: internal::Profile, - cpu_sampler_callback: extern "C" fn(*mut internal::Profile), - upload_callback: extern "C" fn(*mut Handle, &mut Option), - sample_callbacks: ManagedSampleCallbacks, + callbacks: ManagerCallbacks, cancellation_token: CancellationToken, upload_sender: Sender, upload_thread: JoinHandle<()>, @@ -70,12 +76,7 @@ pub struct ProfilerManager { impl ProfilerManager { pub fn start( profile: internal::Profile, - cpu_sampler_callback: extern "C" fn(*mut internal::Profile), - upload_callback: extern "C" fn( - *mut Handle, - &mut Option, - ), - sample_callbacks: ManagedSampleCallbacks, + callbacks: ManagerCallbacks, config: ProfilerManagerConfig, ) -> Result { let (client_channels, manager_channels) = ClientSampleChannels::new(config.channel_depth); @@ -93,7 +94,7 @@ impl ProfilerManager { let upload_thread = std::thread::spawn(move || { while let Ok(profile) = upload_receiver.recv() { let mut handle = Handle::from(profile); - (upload_callback)(&mut handle, &mut token); + (callbacks.upload_callback)(&mut handle, &mut token); } }); @@ -103,9 +104,7 @@ impl ProfilerManager { upload_ticker, shutdown_receiver, profile, - cpu_sampler_callback, - upload_callback, - sample_callbacks, + callbacks, cancellation_token, upload_sender, upload_thread, @@ -125,18 +124,18 @@ impl ProfilerManager { raw_sample: Result, ) -> Result<()> { let mut sample = raw_sample?; - let converted_sample = (self.sample_callbacks.converter)(&sample); + let converted_sample = (self.callbacks.sample_callbacks.converter)(&sample); let add_result = self.profile.add_sample(converted_sample.try_into()?, None); - (self.sample_callbacks.reset)(&mut sample); + (self.callbacks.sample_callbacks.reset)(&mut sample); self.channels .recycled_samples_sender .send(sample) - .map_or_else(|e| (self.sample_callbacks.drop)(e.0), |_| ()); + .map_or_else(|e| (self.callbacks.sample_callbacks.drop)(e.0), |_| ()); add_result } fn handle_cpu_tick(&mut self) { - (self.cpu_sampler_callback)(&mut self.profile); + (self.callbacks.cpu_sampler_callback)(&mut self.profile); } fn handle_upload_tick(&mut self) -> Result<()> { @@ -147,19 +146,22 @@ impl ProfilerManager { Ok(()) } - fn handle_shutdown(mut self) -> Result { + /// # Safety + /// - The caller must ensure that the callbacks remain valid for the lifetime of the profiler. + /// - The callbacks must be thread-safe. + pub fn handle_shutdown(mut self) -> Result { // Process any remaining samples while let Ok(sample) = self.channels.samples_receiver.try_recv() { - let converted_sample = (self.sample_callbacks.converter)(&sample); + let converted_sample = (self.callbacks.sample_callbacks.converter)(&sample); if let Ok(s) = converted_sample.try_into() { let _ = self.profile.add_sample(s, None); } - (self.sample_callbacks.drop)(sample); + (self.callbacks.sample_callbacks.drop)(sample); } // Drain recycled samples while let Ok(sample) = self.channels.recycled_samples_receiver.try_recv() { - (self.sample_callbacks.drop)(sample); + (self.callbacks.sample_callbacks.drop)(sample); } // Cancel any ongoing upload diff --git a/datadog-profiling-ffi/src/manager/tests.rs b/datadog-profiling-ffi/src/manager/tests.rs index 6851e5cdb9..b1e4220a65 100644 --- a/datadog-profiling-ffi/src/manager/tests.rs +++ b/datadog-profiling-ffi/src/manager/tests.rs @@ -7,8 +7,9 @@ use datadog_profiling::internal::Profile; use ddcommon_ffi::{Handle, ToInner}; use tokio_util::sync::CancellationToken; -use super::{ManagedSampleCallbacks, ProfilerManager}; -use crate::manager::profiler_manager::ProfilerManagerConfig; +use super::profiler_manager::{ + ManagedSampleCallbacks, ManagerCallbacks, ProfilerManager, ProfilerManagerConfig, +}; use crate::manager::samples::SendSample; use datadog_profiling_protobuf::prost_impls::Profile as ProstProfile; use ddcommon_ffi::Slice; @@ -183,9 +184,11 @@ fn test_profiler_manager() { let profile = Profile::new(&sample_types, None); let client = ProfilerManager::start( profile, - test_cpu_sampler_callback, - test_upload_callback, - sample_callbacks, + ManagerCallbacks { + cpu_sampler_callback: test_cpu_sampler_callback, + upload_callback: test_upload_callback, + sample_callbacks, + }, config, ) .unwrap(); diff --git a/datadog-trace-protobuf/src/remoteconfig.rs b/datadog-trace-protobuf/src/remoteconfig.rs index a8a4b66524..fdc89ee871 100644 --- a/datadog-trace-protobuf/src/remoteconfig.rs +++ b/datadog-trace-protobuf/src/remoteconfig.rs @@ -3,7 +3,8 @@ use serde::{Deserialize, Serialize}; // This file is @generated by prost-build. -#[derive(Deserialize, Serialize, Clone, PartialEq, ::prost::Message)] +#[derive(Deserialize, Serialize)] +#[derive(Clone, PartialEq, ::prost::Message)] pub struct File { #[prost(string, tag = "1")] pub path: ::prost::alloc::string::String, @@ -11,7 +12,8 @@ pub struct File { #[serde(with = "serde_bytes")] pub raw: ::prost::alloc::vec::Vec, } -#[derive(Deserialize, Serialize, Clone, PartialEq, ::prost::Message)] +#[derive(Deserialize, Serialize)] +#[derive(Clone, PartialEq, ::prost::Message)] pub struct Client { #[prost(message, optional, tag = "1")] pub state: ::core::option::Option, @@ -33,7 +35,8 @@ pub struct Client { #[prost(bytes = "vec", tag = "11")] pub capabilities: ::prost::alloc::vec::Vec, } -#[derive(Deserialize, Serialize, Clone, PartialEq, ::prost::Message)] +#[derive(Deserialize, Serialize)] +#[derive(Clone, PartialEq, ::prost::Message)] pub struct ClientTracer { #[prost(string, tag = "1")] pub runtime_id: ::prost::alloc::string::String, @@ -52,7 +55,8 @@ pub struct ClientTracer { #[prost(string, repeated, tag = "7")] pub tags: ::prost::alloc::vec::Vec<::prost::alloc::string::String>, } -#[derive(Deserialize, Serialize, Clone, PartialEq, ::prost::Message)] +#[derive(Deserialize, Serialize)] +#[derive(Clone, PartialEq, ::prost::Message)] pub struct ClientAgent { #[prost(string, tag = "1")] pub name: ::prost::alloc::string::String, @@ -65,7 +69,8 @@ pub struct ClientAgent { #[prost(string, repeated, tag = "5")] pub cws_workloads: ::prost::alloc::vec::Vec<::prost::alloc::string::String>, } -#[derive(Deserialize, Serialize, Clone, PartialEq, ::prost::Message)] +#[derive(Deserialize, Serialize)] +#[derive(Clone, PartialEq, ::prost::Message)] pub struct ConfigState { #[prost(string, tag = "1")] pub id: ::prost::alloc::string::String, @@ -78,7 +83,8 @@ pub struct ConfigState { #[prost(string, tag = "5")] pub apply_error: ::prost::alloc::string::String, } -#[derive(Deserialize, Serialize, Clone, PartialEq, ::prost::Message)] +#[derive(Deserialize, Serialize)] +#[derive(Clone, PartialEq, ::prost::Message)] pub struct ClientState { #[prost(uint64, tag = "1")] pub root_version: u64, @@ -93,14 +99,16 @@ pub struct ClientState { #[prost(bytes = "vec", tag = "6")] pub backend_client_state: ::prost::alloc::vec::Vec, } -#[derive(Deserialize, Serialize, Clone, PartialEq, ::prost::Message)] +#[derive(Deserialize, Serialize)] +#[derive(Clone, PartialEq, ::prost::Message)] pub struct TargetFileHash { #[prost(string, tag = "1")] pub algorithm: ::prost::alloc::string::String, #[prost(string, tag = "3")] pub hash: ::prost::alloc::string::String, } -#[derive(Deserialize, Serialize, Clone, PartialEq, ::prost::Message)] +#[derive(Deserialize, Serialize)] +#[derive(Clone, PartialEq, ::prost::Message)] pub struct TargetFileMeta { #[prost(string, tag = "1")] pub path: ::prost::alloc::string::String, @@ -109,14 +117,16 @@ pub struct TargetFileMeta { #[prost(message, repeated, tag = "3")] pub hashes: ::prost::alloc::vec::Vec, } -#[derive(Deserialize, Serialize, Clone, PartialEq, ::prost::Message)] +#[derive(Deserialize, Serialize)] +#[derive(Clone, PartialEq, ::prost::Message)] pub struct ClientGetConfigsRequest { #[prost(message, optional, tag = "1")] pub client: ::core::option::Option, #[prost(message, repeated, tag = "2")] pub cached_target_files: ::prost::alloc::vec::Vec, } -#[derive(Deserialize, Serialize, Clone, PartialEq, ::prost::Message)] +#[derive(Deserialize, Serialize)] +#[derive(Clone, PartialEq, ::prost::Message)] pub struct ClientGetConfigsResponse { #[prost(bytes = "vec", repeated, tag = "1")] #[serde(with = "crate::serde")] From 1cec23f06ad92ee2643684305ff457119a64aa5f Mon Sep 17 00:00:00 2001 From: Daniel Schwartz-Narbonne Date: Fri, 13 Jun 2025 16:23:07 -0400 Subject: [PATCH 38/52] smarter way of handling shutdown --- datadog-profiling-ffi/src/manager/client.rs | 36 +++++++++++-------- .../src/manager/profiler_manager.rs | 12 +++++++ 2 files changed, 34 insertions(+), 14 deletions(-) diff --git a/datadog-profiling-ffi/src/manager/client.rs b/datadog-profiling-ffi/src/manager/client.rs index 93f7a1eb3b..ac55cb3738 100644 --- a/datadog-profiling-ffi/src/manager/client.rs +++ b/datadog-profiling-ffi/src/manager/client.rs @@ -1,7 +1,8 @@ +use std::sync::Arc; use std::{ffi::c_void, sync::atomic::AtomicBool, thread::JoinHandle}; -use anyhow::Result; -use crossbeam_channel::{SendError, Sender, TryRecvError}; +use anyhow::{ensure, Result}; +use crossbeam_channel::{SendError, Sender}; use datadog_profiling::internal; use super::ClientSampleChannels; @@ -9,22 +10,23 @@ use super::SendSample; pub struct ManagedProfilerClient { channels: ClientSampleChannels, - handle: JoinHandle>, + handle: JoinHandle>, shutdown_sender: Sender<()>, - is_shutdown: AtomicBool, + is_shutdown: Arc, } impl ManagedProfilerClient { - pub(crate) fn new( + pub fn new( channels: ClientSampleChannels, - handle: JoinHandle>, + handle: JoinHandle>, shutdown_sender: Sender<()>, + is_shutdown: Arc, ) -> Self { Self { channels, handle, shutdown_sender, - is_shutdown: AtomicBool::new(false), + is_shutdown, } } @@ -43,17 +45,23 @@ impl ManagedProfilerClient { self.channels.send_sample(sample) } - pub fn try_recv_recycled(&self) -> Result<*mut c_void, TryRecvError> { + pub fn try_recv_recycled( + &self, + ) -> Result<*mut std::ffi::c_void, crossbeam_channel::TryRecvError> { + if self.is_shutdown.load(std::sync::atomic::Ordering::SeqCst) { + return Err(crossbeam_channel::TryRecvError::Disconnected); + } self.channels.try_recv_recycled() } - pub fn shutdown(self) -> anyhow::Result { - self.is_shutdown - .store(true, std::sync::atomic::Ordering::SeqCst); - // Todo: Should we report if there was an error sending the shutdown signal? - let _ = self.shutdown_sender.send(()); + pub fn shutdown(self) -> Result { + ensure!( + !self.is_shutdown.load(std::sync::atomic::Ordering::SeqCst), + "Profiler manager is already shutdown" + ); + self.shutdown_sender.send(())?; self.handle .join() - .map_err(|e| anyhow::anyhow!("Failed to join handle: {:?}", e))? + .map_err(|e| anyhow::anyhow!("Failed to join manager thread: {:?}", e))? } } diff --git a/datadog-profiling-ffi/src/manager/profiler_manager.rs b/datadog-profiling-ffi/src/manager/profiler_manager.rs index d0fd78f491..06260e5ba6 100644 --- a/datadog-profiling-ffi/src/manager/profiler_manager.rs +++ b/datadog-profiling-ffi/src/manager/profiler_manager.rs @@ -1,3 +1,5 @@ +use std::sync::atomic::AtomicBool; +use std::sync::Arc; use std::thread::JoinHandle; use std::time::{Duration, Instant}; @@ -71,6 +73,7 @@ pub struct ProfilerManager { cancellation_token: CancellationToken, upload_sender: Sender, upload_thread: JoinHandle<()>, + is_shutdown: Arc, } impl ProfilerManager { @@ -89,6 +92,9 @@ impl ProfilerManager { // Create a single cancellation token for all uploads let cancellation_token = CancellationToken::new(); + // Create shared shutdown state + let is_shutdown = Arc::new(AtomicBool::new(false)); + // Spawn the upload thread let mut token = Some(cancellation_token.clone()); let upload_thread = std::thread::spawn(move || { @@ -108,6 +114,7 @@ impl ProfilerManager { cancellation_token, upload_sender, upload_thread, + is_shutdown: is_shutdown.clone(), }; let handle = std::thread::spawn(move || manager.main()); @@ -116,6 +123,7 @@ impl ProfilerManager { client_channels, handle, shutdown_sender, + is_shutdown, )) } @@ -150,6 +158,10 @@ impl ProfilerManager { /// - The caller must ensure that the callbacks remain valid for the lifetime of the profiler. /// - The callbacks must be thread-safe. pub fn handle_shutdown(mut self) -> Result { + // Mark as shutdown + self.is_shutdown + .store(true, std::sync::atomic::Ordering::SeqCst); + // Process any remaining samples while let Ok(sample) = self.channels.samples_receiver.try_recv() { let converted_sample = (self.callbacks.sample_callbacks.converter)(&sample); From 64ec9c0f1b2e6b03f2c7023fba8929ddd9d431a2 Mon Sep 17 00:00:00 2001 From: Daniel Schwartz-Narbonne Date: Tue, 1 Jul 2025 13:35:43 -0400 Subject: [PATCH 39/52] fork handler is cleaner --- datadog-profiling-ffi/src/manager/client.rs | 28 +++- datadog-profiling-ffi/src/manager/ffi_api.rs | 28 ++-- .../src/manager/fork_handler.rs | 157 +++++++++++------- datadog-profiling-ffi/src/manager/mod.rs | 2 +- .../src/manager/profiler_manager.rs | 38 ++--- datadog-profiling-ffi/src/manager/samples.rs | 18 ++ datadog-profiling-ffi/src/manager/tests.rs | 4 +- datadog-trace-protobuf/src/remoteconfig.rs | 30 ++-- 8 files changed, 173 insertions(+), 132 deletions(-) diff --git a/datadog-profiling-ffi/src/manager/client.rs b/datadog-profiling-ffi/src/manager/client.rs index ac55cb3738..2e6b6bb67a 100644 --- a/datadog-profiling-ffi/src/manager/client.rs +++ b/datadog-profiling-ffi/src/manager/client.rs @@ -8,24 +8,22 @@ use datadog_profiling::internal; use super::ClientSampleChannels; use super::SendSample; +#[derive(Debug, Clone)] pub struct ManagedProfilerClient { channels: ClientSampleChannels, + is_shutdown: Arc, +} + +pub struct ManagedProfilerController { handle: JoinHandle>, shutdown_sender: Sender<()>, is_shutdown: Arc, } impl ManagedProfilerClient { - pub fn new( - channels: ClientSampleChannels, - handle: JoinHandle>, - shutdown_sender: Sender<()>, - is_shutdown: Arc, - ) -> Self { + pub fn new(channels: ClientSampleChannels, is_shutdown: Arc) -> Self { Self { channels, - handle, - shutdown_sender, is_shutdown, } } @@ -53,6 +51,20 @@ impl ManagedProfilerClient { } self.channels.try_recv_recycled() } +} + +impl ManagedProfilerController { + pub fn new( + handle: JoinHandle>, + shutdown_sender: Sender<()>, + is_shutdown: Arc, + ) -> Self { + Self { + handle, + shutdown_sender, + is_shutdown, + } + } pub fn shutdown(self) -> Result { ensure!( diff --git a/datadog-profiling-ffi/src/manager/ffi_api.rs b/datadog-profiling-ffi/src/manager/ffi_api.rs index 335b3502cf..1b1259ff1c 100644 --- a/datadog-profiling-ffi/src/manager/ffi_api.rs +++ b/datadog-profiling-ffi/src/manager/ffi_api.rs @@ -2,9 +2,8 @@ // SPDX-License-Identifier: Apache-2.0 use crate::manager::{ - profiler_manager::{ - ManagedSampleCallbacks, ManagerCallbacks, ProfilerManager, ProfilerManagerConfig, - }, + fork_handler, + profiler_manager::{ManagedSampleCallbacks, ProfilerManagerConfig}, ManagedProfilerClient, }; use crate::profiles::datatypes::{Profile, ProfilePtrExt}; @@ -35,14 +34,12 @@ pub unsafe extern "C" fn ddog_prof_ProfilerManager_start( ) -> FFIResult> { wrap_with_ffi_result!({ let internal_profile = *profile.take()?; - let client: ManagedProfilerClient = ProfilerManager::start( + let client = fork_handler::start( internal_profile, - ManagerCallbacks { - cpu_sampler_callback, - upload_callback, - sample_callbacks, - }, config, + cpu_sampler_callback, + upload_callback, + sample_callbacks, )?; anyhow::Ok(Handle::from(client)) }) @@ -100,15 +97,10 @@ pub unsafe extern "C" fn ddog_prof_ProfilerManager_try_recv_recycled( /// - The handle must have been returned by ddog_prof_ProfilerManager_start and not yet dropped. #[no_mangle] #[named] -pub unsafe extern "C" fn ddog_prof_ProfilerManager_shutdown( - mut handle: *mut Handle, -) -> FFIResult { - wrap_with_ffi_result!({ - let profile = handle - .take()? - .shutdown() - .map_err(|e| anyhow::anyhow!("Failed to shutdown client: {:?}", e))?; - anyhow::Ok(Profile::new(profile)) +pub unsafe extern "C" fn ddog_prof_ProfilerManager_shutdown() -> VoidResult { + wrap_with_void_ffi_result!({ + fork_handler::shutdown_global_manager() + .map_err(|e| anyhow::anyhow!("Failed to shutdown global manager: {:?}", e))?; }) } diff --git a/datadog-profiling-ffi/src/manager/fork_handler.rs b/datadog-profiling-ffi/src/manager/fork_handler.rs index 6abc7c95b1..72fa8a4354 100644 --- a/datadog-profiling-ffi/src/manager/fork_handler.rs +++ b/datadog-profiling-ffi/src/manager/fork_handler.rs @@ -1,23 +1,39 @@ #![allow(clippy::unwrap_used)] -use super::client::ManagedProfilerClient; +use super::client::{ManagedProfilerClient, ManagedProfilerController}; use super::profiler_manager::{ ManagedSampleCallbacks, ManagerCallbacks, ProfilerManager, ProfilerManagerConfig, }; -use anyhow::{Context, Result}; +use anyhow::Result; use datadog_profiling::internal::Profile; use std::sync::Mutex; +/// Global state for the profile manager during fork operations +enum ManagerState { + /// Manager has not been initialized yet + Uninitialized, + /// Manager is running with active profiler client and controller + Running { + client: ManagedProfilerClient, + controller: ManagedProfilerController, + config: ProfilerManagerConfig, + callbacks: ManagerCallbacks, + }, + /// Manager is paused (shutdown) with stored profile + Paused { + profile: Box, + config: ProfilerManagerConfig, + callbacks: ManagerCallbacks, + }, +} + // Global state for the profile manager -static MANAGER: Mutex> = Mutex::new(None); -static PROFILE: Mutex> = Mutex::new(None); -static CONFIG: Mutex> = Mutex::new(None); -static CALLBACKS: Mutex> = Mutex::new(None); +static MANAGER_STATE: Mutex = Mutex::new(ManagerState::Uninitialized); -/// Stores the current profile manager state globally. -/// This should be called before forking to ensure clean state. -pub fn store_manager_state( - manager: ProfilerManager, +/// Starts a new profile manager and stores the global state. +/// Returns the client for external use. +pub fn start( + profile: datadog_profiling::internal::Profile, config: ProfilerManagerConfig, cpu_sampler_callback: extern "C" fn(*mut Profile), upload_callback: extern "C" fn( @@ -25,74 +41,87 @@ pub fn store_manager_state( &mut Option, ), sample_callbacks: ManagedSampleCallbacks, -) -> Result<()> { - // Store the manager - *MANAGER.lock().unwrap() = Some(manager); - - // Store the config and callbacks separately - *CONFIG.lock().unwrap() = Some(config); - *CALLBACKS.lock().unwrap() = Some(ManagerCallbacks { - cpu_sampler_callback, - upload_callback, - sample_callbacks, - }); - - Ok(()) +) -> Result { + let (client, controller) = ProfilerManager::start( + profile, + ManagerCallbacks { + cpu_sampler_callback, + upload_callback, + sample_callbacks: sample_callbacks.clone(), + }, + config, + )?; + + let mut state = MANAGER_STATE.lock().unwrap(); + *state = ManagerState::Running { + client: client.clone(), + controller, + config, + callbacks: ManagerCallbacks { + cpu_sampler_callback, + upload_callback, + sample_callbacks, + }, + }; + + Ok(client) } /// Shuts down the stored profile manager and stores its profile. /// This should be called before forking to ensure clean state. -pub fn shutdown_stored_manager() -> Result<()> { - let manager = MANAGER - .lock() - .unwrap() - .take() - .context("No profile manager stored")?; - - let profile = manager.handle_shutdown()?; - - // Store the profile - *PROFILE.lock().unwrap() = Some(profile); - - Ok(()) +pub fn shutdown_global_manager() -> Result<()> { + let mut state = MANAGER_STATE.lock().unwrap(); + match std::mem::replace(&mut *state, ManagerState::Uninitialized) { + ManagerState::Running { + client: _, + controller, + config, + callbacks, + } => { + let profile = controller.shutdown()?; + *state = ManagerState::Paused { + profile: Box::new(profile), + config, + callbacks, + }; + Ok(()) + } + _ => Err(anyhow::anyhow!("Manager is not in running state")), + } } /// Restarts the profile manager in the parent process with the stored profile. /// This should be called after fork in the parent process. -pub fn restart_manager_in_parent() -> Result { - let profile = PROFILE - .lock() - .unwrap() - .take() - .context("No profile stored")?; - - let config = CONFIG.lock().unwrap().take().context("No config stored")?; - - let callbacks = CALLBACKS - .lock() - .unwrap() - .take() - .context("No callbacks stored")?; +pub fn restart_manager_in_parent() -> Result<(ManagedProfilerClient, ManagedProfilerController)> { + let mut state = MANAGER_STATE.lock().unwrap(); + + let (profile, config, callbacks) = + match std::mem::replace(&mut *state, ManagerState::Uninitialized) { + ManagerState::Paused { + profile, + config, + callbacks, + } => (*profile, config, callbacks), + _ => return Err(anyhow::anyhow!("Manager is not in paused state")), + }; ProfilerManager::start(profile, callbacks, config) } /// Restarts the profile manager in the child process with a fresh profile. /// This should be called after fork in the child process. -pub fn restart_manager_in_child() -> Result { - let config = CONFIG.lock().unwrap().take().context("No config stored")?; - - let callbacks = CALLBACKS - .lock() - .unwrap() - .take() - .context("No callbacks stored")?; - - let mut profile = PROFILE - .lock() - .unwrap() - .take() - .context("No profile stored")?; +pub fn restart_manager_in_child() -> Result<(ManagedProfilerClient, ManagedProfilerController)> { + let mut state = MANAGER_STATE.lock().unwrap(); + + let (mut profile, config, callbacks) = + match std::mem::replace(&mut *state, ManagerState::Uninitialized) { + ManagerState::Paused { + profile, + config, + callbacks, + } => (*profile, config, callbacks), + _ => return Err(anyhow::anyhow!("Manager is not in paused state")), + }; // Reset the profile, discarding the previous one let _ = profile.reset_and_return_previous()?; diff --git a/datadog-profiling-ffi/src/manager/mod.rs b/datadog-profiling-ffi/src/manager/mod.rs index c4529f8c18..5736992145 100644 --- a/datadog-profiling-ffi/src/manager/mod.rs +++ b/datadog-profiling-ffi/src/manager/mod.rs @@ -13,6 +13,6 @@ mod samples; #[cfg(test)] mod tests; -pub use client::ManagedProfilerClient; +pub use client::{ManagedProfilerClient, ManagedProfilerController}; pub use profiler_manager::{ManagedSampleCallbacks, ProfilerManager}; pub use samples::{ClientSampleChannels, SendSample}; diff --git a/datadog-profiling-ffi/src/manager/profiler_manager.rs b/datadog-profiling-ffi/src/manager/profiler_manager.rs index 06260e5ba6..0605bc0d34 100644 --- a/datadog-profiling-ffi/src/manager/profiler_manager.rs +++ b/datadog-profiling-ffi/src/manager/profiler_manager.rs @@ -9,7 +9,7 @@ use datadog_profiling::internal; use ddcommon_ffi::Handle; use tokio_util::sync::CancellationToken; -use super::client::ManagedProfilerClient; +use super::client::{ManagedProfilerClient, ManagedProfilerController}; use super::samples::{ClientSampleChannels, ManagerSampleChannels, SendSample}; use crate::profiles::datatypes::Sample; @@ -22,6 +22,7 @@ pub struct ManagerCallbacks { } #[repr(C)] +#[derive(Copy, Clone)] pub struct ProfilerManagerConfig { pub channel_depth: usize, pub cpu_sampling_interval_ms: u64, @@ -81,7 +82,7 @@ impl ProfilerManager { profile: internal::Profile, callbacks: ManagerCallbacks, config: ProfilerManagerConfig, - ) -> Result { + ) -> Result<(ManagedProfilerClient, ManagedProfilerController)> { let (client_channels, manager_channels) = ClientSampleChannels::new(config.channel_depth); let (shutdown_sender, shutdown_receiver) = crossbeam_channel::bounded(1); let (upload_sender, upload_receiver) = crossbeam_channel::bounded(2); @@ -119,12 +120,11 @@ impl ProfilerManager { let handle = std::thread::spawn(move || manager.main()); - Ok(ManagedProfilerClient::new( - client_channels, - handle, - shutdown_sender, - is_shutdown, - )) + let client = ManagedProfilerClient::new(client_channels, is_shutdown.clone()); + + let controller = ManagedProfilerController::new(handle, shutdown_sender, is_shutdown); + + Ok((client, controller)) } fn handle_sample( @@ -150,7 +150,7 @@ impl ProfilerManager { let old_profile = self.profile.reset_and_return_previous()?; self.upload_sender .send(old_profile) - .map_err(|e| anyhow::anyhow!("Failed to send profile for upload: {}", e))?; + .map_err(|e| anyhow::anyhow!("Failed to send profile for upload: {e}"))?; Ok(()) } @@ -162,6 +162,13 @@ impl ProfilerManager { self.is_shutdown .store(true, std::sync::atomic::Ordering::SeqCst); + // Cancel any ongoing upload + self.cancellation_token.cancel(); + + // Drop the sender to signal the upload thread that no more messages will be sent + // This is necessary to allow the upload thread to exit its message processing loop + drop(self.upload_sender); + // Process any remaining samples while let Ok(sample) = self.channels.samples_receiver.try_recv() { let converted_sample = (self.callbacks.sample_callbacks.converter)(&sample); @@ -176,16 +183,9 @@ impl ProfilerManager { (self.callbacks.sample_callbacks.drop)(sample); } - // Cancel any ongoing upload - self.cancellation_token.cancel(); - - // Drop the sender to signal the upload thread that no more messages will be sent - // This is necessary to allow the upload thread to exit its message processing loop - drop(self.upload_sender); - // Wait for the upload thread to finish if let Err(e) = self.upload_thread.join() { - eprintln!("Error joining upload thread: {:?}", e); + eprintln!("Error joining upload thread: {e:?}"); } Ok(self.profile) @@ -199,14 +199,14 @@ impl ProfilerManager { select! { recv(self.channels.samples_receiver) -> raw_sample => { let _ = self.handle_sample(raw_sample) - .map_err(|e| eprintln!("Failed to process sample: {}", e)); + .map_err(|e| eprintln!("Failed to process sample: {e}")); }, recv(self.cpu_ticker) -> msg => { self.handle_cpu_tick(); }, recv(self.upload_ticker) -> msg => { let _ = self.handle_upload_tick() - .map_err(|e| eprintln!("Failed to handle upload: {}", e)); + .map_err(|e| eprintln!("Failed to handle upload: {e}")); }, recv(self.shutdown_receiver) -> _ => { return self.handle_shutdown(); diff --git a/datadog-profiling-ffi/src/manager/samples.rs b/datadog-profiling-ffi/src/manager/samples.rs index f3c3b40ff6..dd38d484af 100644 --- a/datadog-profiling-ffi/src/manager/samples.rs +++ b/datadog-profiling-ffi/src/manager/samples.rs @@ -32,6 +32,24 @@ pub struct ClientSampleChannels { recycled_samples_receiver: Receiver, } +impl std::fmt::Debug for ClientSampleChannels { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("ClientSampleChannels") + .field("samples_sender", &"Sender") + .field("recycled_samples_receiver", &"Receiver") + .finish() + } +} + +impl Clone for ClientSampleChannels { + fn clone(&self) -> Self { + Self { + samples_sender: self.samples_sender.clone(), + recycled_samples_receiver: self.recycled_samples_receiver.clone(), + } + } +} + pub struct ManagerSampleChannels { pub samples_receiver: Receiver, pub recycled_samples_sender: Sender, diff --git a/datadog-profiling-ffi/src/manager/tests.rs b/datadog-profiling-ffi/src/manager/tests.rs index b1e4220a65..ea5a4d3d49 100644 --- a/datadog-profiling-ffi/src/manager/tests.rs +++ b/datadog-profiling-ffi/src/manager/tests.rs @@ -182,7 +182,7 @@ fn test_profiler_manager() { let sample_types = [ValueType::new("samples", "count")]; let profile = Profile::new(&sample_types, None); - let client = ProfilerManager::start( + let (client, controller) = ProfilerManager::start( profile, ManagerCallbacks { cpu_sampler_callback: test_cpu_sampler_callback, @@ -209,7 +209,7 @@ fn test_profiler_manager() { assert_eq!(UPLOAD_COUNT.load(Ordering::SeqCst), 1); // Get the profile and verify it has no samples (they were consumed by the upload) - let profile = client.shutdown().unwrap(); + let profile = controller.shutdown().unwrap(); let pprof = roundtrip_to_pprof(profile); assert_eq!(pprof.samples.len(), 0); } diff --git a/datadog-trace-protobuf/src/remoteconfig.rs b/datadog-trace-protobuf/src/remoteconfig.rs index fdc89ee871..a8a4b66524 100644 --- a/datadog-trace-protobuf/src/remoteconfig.rs +++ b/datadog-trace-protobuf/src/remoteconfig.rs @@ -3,8 +3,7 @@ use serde::{Deserialize, Serialize}; // This file is @generated by prost-build. -#[derive(Deserialize, Serialize)] -#[derive(Clone, PartialEq, ::prost::Message)] +#[derive(Deserialize, Serialize, Clone, PartialEq, ::prost::Message)] pub struct File { #[prost(string, tag = "1")] pub path: ::prost::alloc::string::String, @@ -12,8 +11,7 @@ pub struct File { #[serde(with = "serde_bytes")] pub raw: ::prost::alloc::vec::Vec, } -#[derive(Deserialize, Serialize)] -#[derive(Clone, PartialEq, ::prost::Message)] +#[derive(Deserialize, Serialize, Clone, PartialEq, ::prost::Message)] pub struct Client { #[prost(message, optional, tag = "1")] pub state: ::core::option::Option, @@ -35,8 +33,7 @@ pub struct Client { #[prost(bytes = "vec", tag = "11")] pub capabilities: ::prost::alloc::vec::Vec, } -#[derive(Deserialize, Serialize)] -#[derive(Clone, PartialEq, ::prost::Message)] +#[derive(Deserialize, Serialize, Clone, PartialEq, ::prost::Message)] pub struct ClientTracer { #[prost(string, tag = "1")] pub runtime_id: ::prost::alloc::string::String, @@ -55,8 +52,7 @@ pub struct ClientTracer { #[prost(string, repeated, tag = "7")] pub tags: ::prost::alloc::vec::Vec<::prost::alloc::string::String>, } -#[derive(Deserialize, Serialize)] -#[derive(Clone, PartialEq, ::prost::Message)] +#[derive(Deserialize, Serialize, Clone, PartialEq, ::prost::Message)] pub struct ClientAgent { #[prost(string, tag = "1")] pub name: ::prost::alloc::string::String, @@ -69,8 +65,7 @@ pub struct ClientAgent { #[prost(string, repeated, tag = "5")] pub cws_workloads: ::prost::alloc::vec::Vec<::prost::alloc::string::String>, } -#[derive(Deserialize, Serialize)] -#[derive(Clone, PartialEq, ::prost::Message)] +#[derive(Deserialize, Serialize, Clone, PartialEq, ::prost::Message)] pub struct ConfigState { #[prost(string, tag = "1")] pub id: ::prost::alloc::string::String, @@ -83,8 +78,7 @@ pub struct ConfigState { #[prost(string, tag = "5")] pub apply_error: ::prost::alloc::string::String, } -#[derive(Deserialize, Serialize)] -#[derive(Clone, PartialEq, ::prost::Message)] +#[derive(Deserialize, Serialize, Clone, PartialEq, ::prost::Message)] pub struct ClientState { #[prost(uint64, tag = "1")] pub root_version: u64, @@ -99,16 +93,14 @@ pub struct ClientState { #[prost(bytes = "vec", tag = "6")] pub backend_client_state: ::prost::alloc::vec::Vec, } -#[derive(Deserialize, Serialize)] -#[derive(Clone, PartialEq, ::prost::Message)] +#[derive(Deserialize, Serialize, Clone, PartialEq, ::prost::Message)] pub struct TargetFileHash { #[prost(string, tag = "1")] pub algorithm: ::prost::alloc::string::String, #[prost(string, tag = "3")] pub hash: ::prost::alloc::string::String, } -#[derive(Deserialize, Serialize)] -#[derive(Clone, PartialEq, ::prost::Message)] +#[derive(Deserialize, Serialize, Clone, PartialEq, ::prost::Message)] pub struct TargetFileMeta { #[prost(string, tag = "1")] pub path: ::prost::alloc::string::String, @@ -117,16 +109,14 @@ pub struct TargetFileMeta { #[prost(message, repeated, tag = "3")] pub hashes: ::prost::alloc::vec::Vec, } -#[derive(Deserialize, Serialize)] -#[derive(Clone, PartialEq, ::prost::Message)] +#[derive(Deserialize, Serialize, Clone, PartialEq, ::prost::Message)] pub struct ClientGetConfigsRequest { #[prost(message, optional, tag = "1")] pub client: ::core::option::Option, #[prost(message, repeated, tag = "2")] pub cached_target_files: ::prost::alloc::vec::Vec, } -#[derive(Deserialize, Serialize)] -#[derive(Clone, PartialEq, ::prost::Message)] +#[derive(Deserialize, Serialize, Clone, PartialEq, ::prost::Message)] pub struct ClientGetConfigsResponse { #[prost(bytes = "vec", repeated, tag = "1")] #[serde(with = "crate::serde")] From 905f2ad2a5560e18ea004601a87ebc437c121b27 Mon Sep 17 00:00:00 2001 From: Daniel Schwartz-Narbonne Date: Tue, 1 Jul 2025 14:59:55 -0400 Subject: [PATCH 40/52] feels like a decent api --- datadog-profiling-ffi/src/manager/client.rs | 37 +-- datadog-profiling-ffi/src/manager/ffi_api.rs | 16 +- .../src/manager/fork_handler.rs | 130 -------- datadog-profiling-ffi/src/manager/mod.rs | 5 +- .../src/manager/profiler_manager.rs | 308 ++++++++++++++---- datadog-profiling-ffi/src/manager/tests.rs | 4 +- 6 files changed, 266 insertions(+), 234 deletions(-) delete mode 100644 datadog-profiling-ffi/src/manager/fork_handler.rs diff --git a/datadog-profiling-ffi/src/manager/client.rs b/datadog-profiling-ffi/src/manager/client.rs index 2e6b6bb67a..d0db6229eb 100644 --- a/datadog-profiling-ffi/src/manager/client.rs +++ b/datadog-profiling-ffi/src/manager/client.rs @@ -1,9 +1,7 @@ use std::sync::Arc; -use std::{ffi::c_void, sync::atomic::AtomicBool, thread::JoinHandle}; +use std::{ffi::c_void, sync::atomic::AtomicBool}; -use anyhow::{ensure, Result}; -use crossbeam_channel::{SendError, Sender}; -use datadog_profiling::internal; +use crossbeam_channel::SendError; use super::ClientSampleChannels; use super::SendSample; @@ -14,12 +12,6 @@ pub struct ManagedProfilerClient { is_shutdown: Arc, } -pub struct ManagedProfilerController { - handle: JoinHandle>, - shutdown_sender: Sender<()>, - is_shutdown: Arc, -} - impl ManagedProfilerClient { pub fn new(channels: ClientSampleChannels, is_shutdown: Arc) -> Self { Self { @@ -52,28 +44,3 @@ impl ManagedProfilerClient { self.channels.try_recv_recycled() } } - -impl ManagedProfilerController { - pub fn new( - handle: JoinHandle>, - shutdown_sender: Sender<()>, - is_shutdown: Arc, - ) -> Self { - Self { - handle, - shutdown_sender, - is_shutdown, - } - } - - pub fn shutdown(self) -> Result { - ensure!( - !self.is_shutdown.load(std::sync::atomic::Ordering::SeqCst), - "Profiler manager is already shutdown" - ); - self.shutdown_sender.send(())?; - self.handle - .join() - .map_err(|e| anyhow::anyhow!("Failed to join manager thread: {:?}", e))? - } -} diff --git a/datadog-profiling-ffi/src/manager/ffi_api.rs b/datadog-profiling-ffi/src/manager/ffi_api.rs index 1b1259ff1c..5a05e51e84 100644 --- a/datadog-profiling-ffi/src/manager/ffi_api.rs +++ b/datadog-profiling-ffi/src/manager/ffi_api.rs @@ -2,8 +2,9 @@ // SPDX-License-Identifier: Apache-2.0 use crate::manager::{ - fork_handler, - profiler_manager::{ManagedSampleCallbacks, ProfilerManagerConfig}, + profiler_manager::{ + ManagedSampleCallbacks, ManagerCallbacks, ProfilerManager, ProfilerManagerConfig, + }, ManagedProfilerClient, }; use crate::profiles::datatypes::{Profile, ProfilePtrExt}; @@ -34,13 +35,12 @@ pub unsafe extern "C" fn ddog_prof_ProfilerManager_start( ) -> FFIResult> { wrap_with_ffi_result!({ let internal_profile = *profile.take()?; - let client = fork_handler::start( - internal_profile, - config, + let callbacks = ManagerCallbacks { cpu_sampler_callback, upload_callback, sample_callbacks, - )?; + }; + let client = ProfilerManager::start(internal_profile, callbacks, config)?; anyhow::Ok(Handle::from(client)) }) } @@ -99,8 +99,8 @@ pub unsafe extern "C" fn ddog_prof_ProfilerManager_try_recv_recycled( #[named] pub unsafe extern "C" fn ddog_prof_ProfilerManager_shutdown() -> VoidResult { wrap_with_void_ffi_result!({ - fork_handler::shutdown_global_manager() - .map_err(|e| anyhow::anyhow!("Failed to shutdown global manager: {:?}", e))?; + ProfilerManager::pause() + .map_err(|e| anyhow::anyhow!("Failed to pause global manager: {:?}", e))?; }) } diff --git a/datadog-profiling-ffi/src/manager/fork_handler.rs b/datadog-profiling-ffi/src/manager/fork_handler.rs deleted file mode 100644 index 72fa8a4354..0000000000 --- a/datadog-profiling-ffi/src/manager/fork_handler.rs +++ /dev/null @@ -1,130 +0,0 @@ -#![allow(clippy::unwrap_used)] - -use super::client::{ManagedProfilerClient, ManagedProfilerController}; -use super::profiler_manager::{ - ManagedSampleCallbacks, ManagerCallbacks, ProfilerManager, ProfilerManagerConfig, -}; -use anyhow::Result; -use datadog_profiling::internal::Profile; -use std::sync::Mutex; - -/// Global state for the profile manager during fork operations -enum ManagerState { - /// Manager has not been initialized yet - Uninitialized, - /// Manager is running with active profiler client and controller - Running { - client: ManagedProfilerClient, - controller: ManagedProfilerController, - config: ProfilerManagerConfig, - callbacks: ManagerCallbacks, - }, - /// Manager is paused (shutdown) with stored profile - Paused { - profile: Box, - config: ProfilerManagerConfig, - callbacks: ManagerCallbacks, - }, -} - -// Global state for the profile manager -static MANAGER_STATE: Mutex = Mutex::new(ManagerState::Uninitialized); - -/// Starts a new profile manager and stores the global state. -/// Returns the client for external use. -pub fn start( - profile: datadog_profiling::internal::Profile, - config: ProfilerManagerConfig, - cpu_sampler_callback: extern "C" fn(*mut Profile), - upload_callback: extern "C" fn( - *mut ddcommon_ffi::Handle, - &mut Option, - ), - sample_callbacks: ManagedSampleCallbacks, -) -> Result { - let (client, controller) = ProfilerManager::start( - profile, - ManagerCallbacks { - cpu_sampler_callback, - upload_callback, - sample_callbacks: sample_callbacks.clone(), - }, - config, - )?; - - let mut state = MANAGER_STATE.lock().unwrap(); - *state = ManagerState::Running { - client: client.clone(), - controller, - config, - callbacks: ManagerCallbacks { - cpu_sampler_callback, - upload_callback, - sample_callbacks, - }, - }; - - Ok(client) -} - -/// Shuts down the stored profile manager and stores its profile. -/// This should be called before forking to ensure clean state. -pub fn shutdown_global_manager() -> Result<()> { - let mut state = MANAGER_STATE.lock().unwrap(); - match std::mem::replace(&mut *state, ManagerState::Uninitialized) { - ManagerState::Running { - client: _, - controller, - config, - callbacks, - } => { - let profile = controller.shutdown()?; - *state = ManagerState::Paused { - profile: Box::new(profile), - config, - callbacks, - }; - Ok(()) - } - _ => Err(anyhow::anyhow!("Manager is not in running state")), - } -} - -/// Restarts the profile manager in the parent process with the stored profile. -/// This should be called after fork in the parent process. -pub fn restart_manager_in_parent() -> Result<(ManagedProfilerClient, ManagedProfilerController)> { - let mut state = MANAGER_STATE.lock().unwrap(); - - let (profile, config, callbacks) = - match std::mem::replace(&mut *state, ManagerState::Uninitialized) { - ManagerState::Paused { - profile, - config, - callbacks, - } => (*profile, config, callbacks), - _ => return Err(anyhow::anyhow!("Manager is not in paused state")), - }; - - ProfilerManager::start(profile, callbacks, config) -} - -/// Restarts the profile manager in the child process with a fresh profile. -/// This should be called after fork in the child process. -pub fn restart_manager_in_child() -> Result<(ManagedProfilerClient, ManagedProfilerController)> { - let mut state = MANAGER_STATE.lock().unwrap(); - - let (mut profile, config, callbacks) = - match std::mem::replace(&mut *state, ManagerState::Uninitialized) { - ManagerState::Paused { - profile, - config, - callbacks, - } => (*profile, config, callbacks), - _ => return Err(anyhow::anyhow!("Manager is not in paused state")), - }; - - // Reset the profile, discarding the previous one - let _ = profile.reset_and_return_previous()?; - - ProfilerManager::start(profile, callbacks, config) -} diff --git a/datadog-profiling-ffi/src/manager/mod.rs b/datadog-profiling-ffi/src/manager/mod.rs index 5736992145..34d9382ba6 100644 --- a/datadog-profiling-ffi/src/manager/mod.rs +++ b/datadog-profiling-ffi/src/manager/mod.rs @@ -7,12 +7,11 @@ mod client; mod ffi_api; mod ffi_utils; -mod fork_handler; mod profiler_manager; mod samples; #[cfg(test)] mod tests; -pub use client::{ManagedProfilerClient, ManagedProfilerController}; -pub use profiler_manager::{ManagedSampleCallbacks, ProfilerManager}; +pub use client::ManagedProfilerClient; +pub use profiler_manager::{ManagedProfilerController, ManagedSampleCallbacks, ProfilerManager}; pub use samples::{ClientSampleChannels, SendSample}; diff --git a/datadog-profiling-ffi/src/manager/profiler_manager.rs b/datadog-profiling-ffi/src/manager/profiler_manager.rs index 0605bc0d34..9edc4c9abe 100644 --- a/datadog-profiling-ffi/src/manager/profiler_manager.rs +++ b/datadog-profiling-ffi/src/manager/profiler_manager.rs @@ -1,19 +1,22 @@ +#![allow(clippy::unwrap_used)] + use std::sync::atomic::AtomicBool; -use std::sync::Arc; +use std::sync::{Arc, Mutex}; use std::thread::JoinHandle; use std::time::{Duration, Instant}; -use anyhow::Result; +use anyhow::{Context, Result}; use crossbeam_channel::{select, tick, Receiver, Sender}; use datadog_profiling::internal; use ddcommon_ffi::Handle; use tokio_util::sync::CancellationToken; -use super::client::{ManagedProfilerClient, ManagedProfilerController}; +use super::client::ManagedProfilerClient; use super::samples::{ClientSampleChannels, ManagerSampleChannels, SendSample}; use crate::profiles::datatypes::Sample; /// Holds the callbacks needed to restart the profile manager +#[derive(Copy, Clone)] pub struct ManagerCallbacks { pub cpu_sampler_callback: extern "C" fn(*mut internal::Profile), pub upload_callback: @@ -40,13 +43,10 @@ impl Default for ProfilerManagerConfig { } #[repr(C)] -#[derive(Clone)] +#[derive(Copy, Clone)] pub struct ManagedSampleCallbacks { - // Static is probably the wrong type here, but worry about that later. converter: extern "C" fn(&SendSample) -> Sample, - // Resets the sample for reuse. reset: extern "C" fn(&mut SendSample), - // Called when a sample is dropped (not recycled) drop: extern "C" fn(SendSample), } @@ -64,6 +64,60 @@ impl ManagedSampleCallbacks { } } +/// Controller for managing the profiler manager lifecycle +pub struct ManagedProfilerController { + handle: JoinHandle>, + shutdown_sender: Sender<()>, + is_shutdown: Arc, +} + +impl ManagedProfilerController { + pub fn new( + handle: JoinHandle>, + shutdown_sender: Sender<()>, + is_shutdown: Arc, + ) -> Self { + Self { + handle, + shutdown_sender, + is_shutdown, + } + } + + pub fn shutdown(self) -> Result { + anyhow::ensure!( + !self.is_shutdown.load(std::sync::atomic::Ordering::SeqCst), + "Profiler manager is already shutdown" + ); + self.shutdown_sender.send(())?; + self.handle + .join() + .map_err(|e| anyhow::anyhow!("Failed to join manager thread: {:?}", e))? + } +} + +/// Global state for the profile manager during fork operations +enum ManagerState { + /// Manager has not been initialized yet + Uninitialized, + /// Manager is running with active profiler client and controller + Running { + client: ManagedProfilerClient, + controller: ManagedProfilerController, + config: ProfilerManagerConfig, + callbacks: ManagerCallbacks, + }, + /// Manager is paused (shutdown) with stored profile + Paused { + profile: Box, + config: ProfilerManagerConfig, + callbacks: ManagerCallbacks, + }, +} + +// Global state for the profile manager +static MANAGER_STATE: Mutex = Mutex::new(ManagerState::Uninitialized); + pub struct ProfilerManager { channels: ManagerSampleChannels, cpu_ticker: Receiver, @@ -78,55 +132,7 @@ pub struct ProfilerManager { } impl ProfilerManager { - pub fn start( - profile: internal::Profile, - callbacks: ManagerCallbacks, - config: ProfilerManagerConfig, - ) -> Result<(ManagedProfilerClient, ManagedProfilerController)> { - let (client_channels, manager_channels) = ClientSampleChannels::new(config.channel_depth); - let (shutdown_sender, shutdown_receiver) = crossbeam_channel::bounded(1); - let (upload_sender, upload_receiver) = crossbeam_channel::bounded(2); - - let cpu_ticker = tick(Duration::from_millis(config.cpu_sampling_interval_ms)); - let upload_ticker = tick(Duration::from_millis(config.upload_interval_ms)); - - // Create a single cancellation token for all uploads - let cancellation_token = CancellationToken::new(); - - // Create shared shutdown state - let is_shutdown = Arc::new(AtomicBool::new(false)); - - // Spawn the upload thread - let mut token = Some(cancellation_token.clone()); - let upload_thread = std::thread::spawn(move || { - while let Ok(profile) = upload_receiver.recv() { - let mut handle = Handle::from(profile); - (callbacks.upload_callback)(&mut handle, &mut token); - } - }); - - let manager = Self { - channels: manager_channels, - cpu_ticker, - upload_ticker, - shutdown_receiver, - profile, - callbacks, - cancellation_token, - upload_sender, - upload_thread, - is_shutdown: is_shutdown.clone(), - }; - - let handle = std::thread::spawn(move || manager.main()); - - let client = ManagedProfilerClient::new(client_channels, is_shutdown.clone()); - - let controller = ManagedProfilerController::new(handle, shutdown_sender, is_shutdown); - - Ok((client, controller)) - } - + // --- Member functions (instance methods) --- fn handle_sample( &mut self, raw_sample: Result, @@ -215,3 +221,193 @@ impl ProfilerManager { } } } + +impl ProfilerManager { + // --- Global functions (static methods) --- + /// Starts a new profile manager and stores the global state. + /// Returns the client for external use. + pub fn start( + profile: internal::Profile, + callbacks: ManagerCallbacks, + config: ProfilerManagerConfig, + ) -> Result { + let mut state = MANAGER_STATE + .lock() + .map_err(|e| anyhow::anyhow!("{}", e)) + .context("Failed to lock MANAGER_STATE")?; + + // Check if manager is already initialized + anyhow::ensure!( + matches!(&*state, ManagerState::Uninitialized), + "Manager is already initialized" + ); + + let (client_channels, manager_channels) = ClientSampleChannels::new(config.channel_depth); + let (shutdown_sender, shutdown_receiver) = crossbeam_channel::bounded(1); + let (upload_sender, upload_receiver) = crossbeam_channel::bounded(2); + + let cpu_ticker = tick(Duration::from_millis(config.cpu_sampling_interval_ms)); + let upload_ticker = tick(Duration::from_millis(config.upload_interval_ms)); + + // Create a single cancellation token for all uploads + let cancellation_token = CancellationToken::new(); + + // Create shared shutdown state + let is_shutdown = Arc::new(AtomicBool::new(false)); + + // Spawn the upload thread + let mut token = Some(cancellation_token.clone()); + let upload_thread = std::thread::spawn(move || { + while let Ok(profile) = upload_receiver.recv() { + let mut handle = Handle::from(profile); + (callbacks.upload_callback)(&mut handle, &mut token); + } + }); + + let manager = Self { + channels: manager_channels, + cpu_ticker, + upload_ticker, + shutdown_receiver, + profile, + callbacks, + cancellation_token, + upload_sender, + upload_thread, + is_shutdown: is_shutdown.clone(), + }; + + let handle = std::thread::spawn(move || manager.main()); + + let client = ManagedProfilerClient::new(client_channels, is_shutdown.clone()); + + let controller = ManagedProfilerController::new(handle, shutdown_sender, is_shutdown); + + *state = ManagerState::Running { + client: client.clone(), + controller, + config, + callbacks, + }; + + Ok(client) + } + + pub fn pause() -> Result<()> { + let mut state = MANAGER_STATE + .lock() + .map_err(|e| anyhow::anyhow!("{}", e)) + .context("Failed to lock MANAGER_STATE")?; + + // Check if manager is in running state + anyhow::ensure!( + matches!(&*state, ManagerState::Running { .. }), + "Manager is not in running state" + ); + + // Extract the running state and replace with uninitialized + let running_state = match std::mem::replace(&mut *state, ManagerState::Uninitialized) { + ManagerState::Running { + client: _, + controller, + config, + callbacks, + } => (controller, config, callbacks), + _ => unreachable!(), // We already checked above + }; + + let (controller, config, callbacks) = running_state; + let profile = controller.shutdown()?; + + *state = ManagerState::Paused { + profile: Box::new(profile), + config, + callbacks, + }; + + Ok(()) + } + + pub fn restart_in_parent() -> Result { + let mut state = MANAGER_STATE + .lock() + .map_err(|e| anyhow::anyhow!("{}", e)) + .context("Failed to lock MANAGER_STATE")?; + + let (profile, config, callbacks) = + match std::mem::replace(&mut *state, ManagerState::Uninitialized) { + ManagerState::Paused { + profile, + config, + callbacks, + } => (*profile, config, callbacks), + _ => return Err(anyhow::anyhow!("Manager is not in paused state")), + }; + + Self::start(profile, callbacks, config) + } + + pub fn restart_in_child() -> Result { + let mut state = MANAGER_STATE + .lock() + .map_err(|e| anyhow::anyhow!("{}", e)) + .context("Failed to lock MANAGER_STATE")?; + + let (mut profile, config, callbacks) = + match std::mem::replace(&mut *state, ManagerState::Uninitialized) { + ManagerState::Paused { + profile, + config, + callbacks, + } => (*profile, config, callbacks), + _ => return Err(anyhow::anyhow!("Manager is not in paused state")), + }; + + // Reset the profile, discarding the previous one + let _ = profile.reset_and_return_previous()?; + + Self::start(profile, callbacks, config) + } + + /// Terminates the global profile manager and returns the final profile. + /// This should be called when the profiler is no longer needed. + pub fn terminate() -> Result { + let mut state = MANAGER_STATE + .lock() + .map_err(|e| anyhow::anyhow!("{}", e)) + .context("Failed to lock MANAGER_STATE")?; + + // Check if manager is in running or paused state + anyhow::ensure!( + matches!( + &*state, + ManagerState::Running { .. } | ManagerState::Paused { .. } + ), + "Manager is not in running or paused state" + ); + + // Extract the profile and replace with uninitialized + let profile = match std::mem::replace(&mut *state, ManagerState::Uninitialized) { + ManagerState::Running { + client: _, + controller, + config: _, + callbacks: _, + } => { + // Shutdown the controller and get the profile + controller.shutdown()? + } + ManagerState::Paused { + profile, + config: _, + callbacks: _, + } => { + // Return the stored profile + *profile + } + _ => unreachable!(), // We already checked above + }; + + Ok(profile) + } +} diff --git a/datadog-profiling-ffi/src/manager/tests.rs b/datadog-profiling-ffi/src/manager/tests.rs index ea5a4d3d49..d924117163 100644 --- a/datadog-profiling-ffi/src/manager/tests.rs +++ b/datadog-profiling-ffi/src/manager/tests.rs @@ -182,7 +182,7 @@ fn test_profiler_manager() { let sample_types = [ValueType::new("samples", "count")]; let profile = Profile::new(&sample_types, None); - let (client, controller) = ProfilerManager::start( + let client = ProfilerManager::start( profile, ManagerCallbacks { cpu_sampler_callback: test_cpu_sampler_callback, @@ -209,7 +209,7 @@ fn test_profiler_manager() { assert_eq!(UPLOAD_COUNT.load(Ordering::SeqCst), 1); // Get the profile and verify it has no samples (they were consumed by the upload) - let profile = controller.shutdown().unwrap(); + let profile = ProfilerManager::terminate().unwrap(); let pprof = roundtrip_to_pprof(profile); assert_eq!(pprof.samples.len(), 0); } From d84d2bbc39b2a04d492d427483b6ccc7335d6cb0 Mon Sep 17 00:00:00 2001 From: Daniel Schwartz-Narbonne Date: Tue, 1 Jul 2025 17:28:59 -0400 Subject: [PATCH 41/52] cleanup api --- datadog-profiling-ffi/src/manager/ffi_api.rs | 79 +++++++++++++++++-- .../src/manager/profiler_manager.rs | 31 ++------ 2 files changed, 79 insertions(+), 31 deletions(-) diff --git a/datadog-profiling-ffi/src/manager/ffi_api.rs b/datadog-profiling-ffi/src/manager/ffi_api.rs index 5a05e51e84..92376a538a 100644 --- a/datadog-profiling-ffi/src/manager/ffi_api.rs +++ b/datadog-profiling-ffi/src/manager/ffi_api.rs @@ -93,23 +93,86 @@ pub unsafe extern "C" fn ddog_prof_ProfilerManager_try_recv_recycled( }) } +/// Pauses the global profiler manager, shutting down the current instance and storing the profile. +/// The manager can be restarted later using restart functions. +/// /// # Safety -/// - The handle must have been returned by ddog_prof_ProfilerManager_start and not yet dropped. +/// - This function is thread-safe and can be called from any thread. +/// - The manager must be in running state. #[no_mangle] #[named] -pub unsafe extern "C" fn ddog_prof_ProfilerManager_shutdown() -> VoidResult { +pub unsafe extern "C" fn ddog_prof_ProfilerManager_pause() -> VoidResult { wrap_with_void_ffi_result!({ - ProfilerManager::pause() - .map_err(|e| anyhow::anyhow!("Failed to pause global manager: {:?}", e))?; + ProfilerManager::pause().context("Failed to pause global manager")?; + }) +} + +/// Restarts the profiler manager in the parent process after a fork. +/// This preserves the profile data from before the pause. +/// +/// # Safety +/// - This function is thread-safe and can be called from any thread. +/// - The manager must be in paused state. +/// - This should be called in the parent process after a fork. +#[no_mangle] +#[named] +pub unsafe extern "C" fn ddog_prof_ProfilerManager_restart_in_parent( +) -> FFIResult> { + wrap_with_ffi_result!({ + let client = + ProfilerManager::restart_in_parent().context("Failed to restart manager in parent")?; + anyhow::Ok(Handle::from(client)) + }) +} + +/// Restarts the profiler manager in the child process after a fork. +/// This discards the profile data from before the pause and starts fresh. +/// +/// # Safety +/// - This function is thread-safe and can be called from any thread. +/// - The manager must be in paused state. +/// - This should be called in the child process after a fork. +#[no_mangle] +#[named] +pub unsafe extern "C" fn ddog_prof_ProfilerManager_restart_in_child( +) -> FFIResult> { + wrap_with_ffi_result!({ + let client = + ProfilerManager::restart_in_child().context("Failed to restart manager in child")?; + anyhow::Ok(Handle::from(client)) + }) +} + +/// Terminates the global profiler manager and returns the final profile. +/// This should be called when the profiler is no longer needed. +/// +/// # Safety +/// - This function is thread-safe and can be called from any thread. +/// - The manager must be in running or paused state. +/// - The returned profile handle must be properly managed by the caller. +#[no_mangle] +#[named] +pub unsafe extern "C" fn ddog_prof_ProfilerManager_terminate( +) -> FFIResult> { + wrap_with_ffi_result!({ + let profile = ProfilerManager::terminate().context("Failed to terminate global manager")?; + anyhow::Ok(Handle::from(profile)) }) } +/// Drops a profiler client handle. +/// This only drops the client handle and does not affect the global manager state. +/// /// # Safety /// - The handle must have been returned by ddog_prof_ProfilerManager_start and not yet dropped. #[no_mangle] -// TODO: Do we want drop and shutdown to be separate functions? Or should it always be shutdown? -pub unsafe extern "C" fn ddog_prof_ProfilerManager_drop( +#[named] +pub unsafe extern "C" fn ddog_prof_ProfilerClient_drop( mut handle: *mut Handle, -) { - let _ = handle.take(); +) -> VoidResult { + wrap_with_void_ffi_result!({ + handle + .take() + .context("Failed to drop profiler client handle")?; + }) } diff --git a/datadog-profiling-ffi/src/manager/profiler_manager.rs b/datadog-profiling-ffi/src/manager/profiler_manager.rs index 9edc4c9abe..58e59dda60 100644 --- a/datadog-profiling-ffi/src/manager/profiler_manager.rs +++ b/datadog-profiling-ffi/src/manager/profiler_manager.rs @@ -5,7 +5,7 @@ use std::sync::{Arc, Mutex}; use std::thread::JoinHandle; use std::time::{Duration, Instant}; -use anyhow::{Context, Result}; +use anyhow::Result; use crossbeam_channel::{select, tick, Receiver, Sender}; use datadog_profiling::internal; use ddcommon_ffi::Handle; @@ -231,10 +231,7 @@ impl ProfilerManager { callbacks: ManagerCallbacks, config: ProfilerManagerConfig, ) -> Result { - let mut state = MANAGER_STATE - .lock() - .map_err(|e| anyhow::anyhow!("{}", e)) - .context("Failed to lock MANAGER_STATE")?; + let mut state = MANAGER_STATE.lock().map_err(|e| anyhow::anyhow!("{}", e))?; // Check if manager is already initialized anyhow::ensure!( @@ -294,10 +291,7 @@ impl ProfilerManager { } pub fn pause() -> Result<()> { - let mut state = MANAGER_STATE - .lock() - .map_err(|e| anyhow::anyhow!("{}", e)) - .context("Failed to lock MANAGER_STATE")?; + let mut state = MANAGER_STATE.lock().map_err(|e| anyhow::anyhow!("{}", e))?; // Check if manager is in running state anyhow::ensure!( @@ -329,10 +323,7 @@ impl ProfilerManager { } pub fn restart_in_parent() -> Result { - let mut state = MANAGER_STATE - .lock() - .map_err(|e| anyhow::anyhow!("{}", e)) - .context("Failed to lock MANAGER_STATE")?; + let mut state = MANAGER_STATE.lock().map_err(|e| anyhow::anyhow!("{}", e))?; let (profile, config, callbacks) = match std::mem::replace(&mut *state, ManagerState::Uninitialized) { @@ -341,17 +332,14 @@ impl ProfilerManager { config, callbacks, } => (*profile, config, callbacks), - _ => return Err(anyhow::anyhow!("Manager is not in paused state")), + _ => anyhow::bail!("Manager is not in paused state"), }; Self::start(profile, callbacks, config) } pub fn restart_in_child() -> Result { - let mut state = MANAGER_STATE - .lock() - .map_err(|e| anyhow::anyhow!("{}", e)) - .context("Failed to lock MANAGER_STATE")?; + let mut state = MANAGER_STATE.lock().map_err(|e| anyhow::anyhow!("{}", e))?; let (mut profile, config, callbacks) = match std::mem::replace(&mut *state, ManagerState::Uninitialized) { @@ -360,7 +348,7 @@ impl ProfilerManager { config, callbacks, } => (*profile, config, callbacks), - _ => return Err(anyhow::anyhow!("Manager is not in paused state")), + _ => anyhow::bail!("Manager is not in paused state"), }; // Reset the profile, discarding the previous one @@ -372,10 +360,7 @@ impl ProfilerManager { /// Terminates the global profile manager and returns the final profile. /// This should be called when the profiler is no longer needed. pub fn terminate() -> Result { - let mut state = MANAGER_STATE - .lock() - .map_err(|e| anyhow::anyhow!("{}", e)) - .context("Failed to lock MANAGER_STATE")?; + let mut state = MANAGER_STATE.lock().map_err(|e| anyhow::anyhow!("{}", e))?; // Check if manager is in running or paused state anyhow::ensure!( From 51ed0f695c4501821f915fbae02173dc0a3b17df Mon Sep 17 00:00:00 2001 From: Daniel Schwartz-Narbonne Date: Wed, 2 Jul 2025 15:28:08 -0400 Subject: [PATCH 42/52] new tests, fix the issue with restart hanging --- datadog-profiling-ffi/src/lib.rs | 14 +- datadog-profiling-ffi/src/manager/ffi_api.rs | 15 ++ datadog-profiling-ffi/src/manager/mod.rs | 9 +- .../src/manager/profiler_manager.rs | 95 +++++++++---- .../src/profiles/datatypes.rs | 16 ++- datadog-profiling-ffi/src/profiles/mod.rs | 2 +- .../tests/test_ffi_lifecycle_basic.rs | 111 +++++++++++++++ .../tests/test_ffi_lifecycle_with_data.rs | 132 ++++++++++++++++++ .../tests/test_ffi_start_pause_terminate.rs | 112 +++++++++++++++ .../tests/test_ffi_start_terminate.rs | 103 ++++++++++++++ datadog-profiling-ffi/tests/test_utils.rs | 87 ++++++++++++ 11 files changed, 665 insertions(+), 31 deletions(-) create mode 100644 datadog-profiling-ffi/tests/test_ffi_lifecycle_basic.rs create mode 100644 datadog-profiling-ffi/tests/test_ffi_lifecycle_with_data.rs create mode 100644 datadog-profiling-ffi/tests/test_ffi_start_pause_terminate.rs create mode 100644 datadog-profiling-ffi/tests/test_ffi_start_terminate.rs create mode 100644 datadog-profiling-ffi/tests/test_utils.rs diff --git a/datadog-profiling-ffi/src/lib.rs b/datadog-profiling-ffi/src/lib.rs index 38fffda91c..0391c97d43 100644 --- a/datadog-profiling-ffi/src/lib.rs +++ b/datadog-profiling-ffi/src/lib.rs @@ -12,7 +12,7 @@ pub use symbolizer_ffi::*; mod exporter; mod manager; -mod profiles; +pub mod profiles; mod string_storage; // re-export crashtracker ffi @@ -40,3 +40,15 @@ pub use datadog_log_ffi::*; pub use ddcommon_ffi::*; pub use manager::*; + +// Re-export for integration tests +pub use crate::manager::ffi_api::{ + ddog_prof_ProfilerClient_drop, ddog_prof_ProfilerManager_enqueue_sample, + ddog_prof_ProfilerManager_pause, ddog_prof_ProfilerManager_reset_for_testing, + ddog_prof_ProfilerManager_restart_in_parent, ddog_prof_ProfilerManager_start, + ddog_prof_ProfilerManager_terminate, +}; +pub use crate::manager::{ManagedSampleCallbacks, ProfilerManagerConfig, SendSample}; +pub use crate::profiles::datatypes::{ + ddog_prof_Profile_new, Function, Location, ProfileNewResult, Sample, ValueType, +}; diff --git a/datadog-profiling-ffi/src/manager/ffi_api.rs b/datadog-profiling-ffi/src/manager/ffi_api.rs index 92376a538a..2ad7e274b1 100644 --- a/datadog-profiling-ffi/src/manager/ffi_api.rs +++ b/datadog-profiling-ffi/src/manager/ffi_api.rs @@ -176,3 +176,18 @@ pub unsafe extern "C" fn ddog_prof_ProfilerClient_drop( .context("Failed to drop profiler client handle")?; }) } + +/// Resets the global profiler manager state to uninitialized. +/// This is intended for testing purposes only and should not be used in production. +/// +/// # Safety +/// - This function is thread-safe and can be called from any thread. +/// - This function will forcefully reset the state without proper cleanup. +/// - This should only be used in test environments. +#[no_mangle] +#[named] +pub unsafe extern "C" fn ddog_prof_ProfilerManager_reset_for_testing() -> VoidResult { + wrap_with_void_ffi_result!({ + ProfilerManager::reset_for_testing().map_err(|msg| anyhow::anyhow!(msg))?; + }) +} diff --git a/datadog-profiling-ffi/src/manager/mod.rs b/datadog-profiling-ffi/src/manager/mod.rs index 34d9382ba6..c51d454436 100644 --- a/datadog-profiling-ffi/src/manager/mod.rs +++ b/datadog-profiling-ffi/src/manager/mod.rs @@ -5,7 +5,7 @@ #![allow(clippy::todo)] mod client; -mod ffi_api; +pub mod ffi_api; mod ffi_utils; mod profiler_manager; mod samples; @@ -13,5 +13,10 @@ mod samples; mod tests; pub use client::ManagedProfilerClient; -pub use profiler_manager::{ManagedProfilerController, ManagedSampleCallbacks, ProfilerManager}; +pub use profiler_manager::{ + ManagedProfilerController, ManagedSampleCallbacks, ProfilerManager, ProfilerManagerConfig, +}; pub use samples::{ClientSampleChannels, SendSample}; + +// Re-export FFI functions for integration tests +pub use ffi_api::*; diff --git a/datadog-profiling-ffi/src/manager/profiler_manager.rs b/datadog-profiling-ffi/src/manager/profiler_manager.rs index 58e59dda60..64625a7346 100644 --- a/datadog-profiling-ffi/src/manager/profiler_manager.rs +++ b/datadog-profiling-ffi/src/manager/profiler_manager.rs @@ -6,7 +6,7 @@ use std::thread::JoinHandle; use std::time::{Duration, Instant}; use anyhow::Result; -use crossbeam_channel::{select, tick, Receiver, Sender}; +use crossbeam_channel::{select_biased, tick, Receiver, Sender}; use datadog_profiling::internal; use ddcommon_ffi::Handle; use tokio_util::sync::CancellationToken; @@ -66,20 +66,23 @@ impl ManagedSampleCallbacks { /// Controller for managing the profiler manager lifecycle pub struct ManagedProfilerController { - handle: JoinHandle>, + handle: JoinHandle<()>, shutdown_sender: Sender<()>, + profile_result_receiver: Receiver, is_shutdown: Arc, } impl ManagedProfilerController { pub fn new( - handle: JoinHandle>, + handle: JoinHandle<()>, shutdown_sender: Sender<()>, + profile_result_receiver: Receiver, is_shutdown: Arc, ) -> Self { Self { handle, shutdown_sender, + profile_result_receiver, is_shutdown, } } @@ -90,9 +93,19 @@ impl ManagedProfilerController { "Profiler manager is already shutdown" ); self.shutdown_sender.send(())?; + + // Wait for the profile result from the manager thread + let profile = self + .profile_result_receiver + .recv() + .map_err(|e| anyhow::anyhow!("Failed to receive profile result: {e}"))?; + + // Wait for the manager thread to finish self.handle .join() - .map_err(|e| anyhow::anyhow!("Failed to join manager thread: {:?}", e))? + .map_err(|e| anyhow::anyhow!("Failed to join manager thread: {:?}", e))?; + + Ok(profile) } } @@ -129,6 +142,7 @@ pub struct ProfilerManager { upload_sender: Sender, upload_thread: JoinHandle<()>, is_shutdown: Arc, + profile_result_sender: Sender, } impl ProfilerManager { @@ -163,7 +177,7 @@ impl ProfilerManager { /// # Safety /// - The caller must ensure that the callbacks remain valid for the lifetime of the profiler. /// - The callbacks must be thread-safe. - pub fn handle_shutdown(mut self) -> Result { + pub fn handle_shutdown(mut self) { // Mark as shutdown self.is_shutdown .store(true, std::sync::atomic::Ordering::SeqCst); @@ -194,29 +208,32 @@ impl ProfilerManager { eprintln!("Error joining upload thread: {e:?}"); } - Ok(self.profile) + // Send the profile through the channel + let _ = self.profile_result_sender.send(self.profile); } /// # Safety /// - The caller must ensure that the callbacks remain valid for the lifetime of the profiler. /// - The callbacks must be thread-safe. - pub fn main(mut self) -> Result { + pub fn main(mut self) { loop { - select! { - recv(self.channels.samples_receiver) -> raw_sample => { - let _ = self.handle_sample(raw_sample) - .map_err(|e| eprintln!("Failed to process sample: {e}")); + select_biased! { + // Prioritize shutdown signal to ensure quick response to shutdown requests + recv(self.shutdown_receiver) -> _ => { + self.handle_shutdown(); + return; }, recv(self.cpu_ticker) -> msg => { self.handle_cpu_tick(); }, + recv(self.channels.samples_receiver) -> raw_sample => { + let _ = self.handle_sample(raw_sample) + .map_err(|e| eprintln!("Failed to process sample: {e}")); + }, recv(self.upload_ticker) -> msg => { let _ = self.handle_upload_tick() .map_err(|e| eprintln!("Failed to handle upload: {e}")); }, - recv(self.shutdown_receiver) -> _ => { - return self.handle_shutdown(); - }, } } } @@ -242,6 +259,7 @@ impl ProfilerManager { let (client_channels, manager_channels) = ClientSampleChannels::new(config.channel_depth); let (shutdown_sender, shutdown_receiver) = crossbeam_channel::bounded(1); let (upload_sender, upload_receiver) = crossbeam_channel::bounded(2); + let (profile_result_sender, profile_result_receiver) = crossbeam_channel::bounded(1); let cpu_ticker = tick(Duration::from_millis(config.cpu_sampling_interval_ms)); let upload_ticker = tick(Duration::from_millis(config.upload_interval_ms)); @@ -272,13 +290,19 @@ impl ProfilerManager { upload_sender, upload_thread, is_shutdown: is_shutdown.clone(), + profile_result_sender, }; let handle = std::thread::spawn(move || manager.main()); let client = ManagedProfilerClient::new(client_channels, is_shutdown.clone()); - let controller = ManagedProfilerController::new(handle, shutdown_sender, is_shutdown); + let controller = ManagedProfilerController::new( + handle, + shutdown_sender, + profile_result_receiver, + is_shutdown, + ); *state = ManagerState::Running { client: client.clone(), @@ -323,17 +347,20 @@ impl ProfilerManager { } pub fn restart_in_parent() -> Result { - let mut state = MANAGER_STATE.lock().map_err(|e| anyhow::anyhow!("{}", e))?; - - let (profile, config, callbacks) = - match std::mem::replace(&mut *state, ManagerState::Uninitialized) { - ManagerState::Paused { - profile, - config, - callbacks, - } => (*profile, config, callbacks), - _ => anyhow::bail!("Manager is not in paused state"), - }; + let (profile, config, callbacks) = { + let mut state = MANAGER_STATE.lock().map_err(|e| anyhow::anyhow!("{}", e))?; + + let (profile, config, callbacks) = + match std::mem::replace(&mut *state, ManagerState::Uninitialized) { + ManagerState::Paused { + profile, + config, + callbacks, + } => (*profile, config, callbacks), + _ => anyhow::bail!("Manager is not in paused state"), + }; + (profile, config, callbacks) + }; Self::start(profile, callbacks, config) } @@ -395,4 +422,20 @@ impl ProfilerManager { Ok(profile) } + + /// Resets the global state to uninitialized. + /// This is useful for testing to ensure clean state between tests. + /// # Safety + /// This function should only be used in tests. It will forcefully reset the state + /// without proper cleanup, which could lead to resource leaks in production code. + /// If the lock cannot be acquired, this function will return an error (test-only behavior). + pub fn reset_for_testing() -> Result<(), String> { + match MANAGER_STATE.lock() { + Ok(mut state) => { + *state = ManagerState::Uninitialized; + Ok(()) + } + Err(e) => Err(format!("Failed to acquire state lock: {e}")), + } + } } diff --git a/datadog-profiling-ffi/src/profiles/datatypes.rs b/datadog-profiling-ffi/src/profiles/datatypes.rs index acba94cd1d..90e561a7e7 100644 --- a/datadog-profiling-ffi/src/profiles/datatypes.rs +++ b/datadog-profiling-ffi/src/profiles/datatypes.rs @@ -425,9 +425,17 @@ pub unsafe extern "C" fn ddog_prof_Profile_new( } /// Same as `ddog_profile_new` but also configures a `string_storage` for the profile. +/// +/// # Safety +/// - `sample_types` must be a valid slice of ValueType. +/// - `period`, if provided, must be a valid reference. +/// - `string_storage` must be a valid ManagedStringStorage. +/// - The caller is responsible for ensuring that all pointers remain valid for the duration of the +/// call. +/// +/// TODO: @ivoanjo Should this take a `*mut ManagedStringStorage` like Profile APIs do? #[no_mangle] #[must_use] -/// TODO: @ivoanjo Should this take a `*mut ManagedStringStorage` like Profile APIs do? pub unsafe extern "C" fn ddog_prof_Profile_with_string_storage( sample_types: Slice, period: Option<&Period>, @@ -782,6 +790,12 @@ pub unsafe extern "C" fn ddog_prof_Profile_serialize( .into() } +/// Returns a slice view of the given Vec. +/// +/// # Safety +/// - `vec` must be a valid reference to a ddcommon_ffi::Vec. +/// - The returned slice is only valid as long as the original Vec is valid and not mutated or +/// dropped. #[must_use] #[no_mangle] pub unsafe extern "C" fn ddog_Vec_U8_as_slice(vec: &ddcommon_ffi::Vec) -> Slice { diff --git a/datadog-profiling-ffi/src/profiles/mod.rs b/datadog-profiling-ffi/src/profiles/mod.rs index 8df9a23029..b8fceee8d0 100644 --- a/datadog-profiling-ffi/src/profiles/mod.rs +++ b/datadog-profiling-ffi/src/profiles/mod.rs @@ -1,5 +1,5 @@ // Copyright 2021-Present Datadog, Inc. https://www.datadoghq.com/ // SPDX-License-Identifier: Apache-2.0 -pub(crate) mod datatypes; +pub mod datatypes; mod interning_api; diff --git a/datadog-profiling-ffi/tests/test_ffi_lifecycle_basic.rs b/datadog-profiling-ffi/tests/test_ffi_lifecycle_basic.rs new file mode 100644 index 0000000000..202eadd256 --- /dev/null +++ b/datadog-profiling-ffi/tests/test_ffi_lifecycle_basic.rs @@ -0,0 +1,111 @@ +mod test_utils; +use datadog_profiling_ffi::{ + ddog_prof_Profile_new, ddog_prof_ProfilerClient_drop, ddog_prof_ProfilerManager_pause, + ddog_prof_ProfilerManager_reset_for_testing, ddog_prof_ProfilerManager_restart_in_parent, + ddog_prof_ProfilerManager_start, ddog_prof_ProfilerManager_terminate, ManagedSampleCallbacks, + ProfileNewResult, ProfilerManagerConfig, Slice, ValueType, VoidResult, +}; +use test_utils::*; + +#[test] +fn test_ffi_lifecycle_basic() { + println!("[test] Starting basic lifecycle test"); + // Reset global state for this test + unsafe { ddog_prof_ProfilerManager_reset_for_testing() }.unwrap(); + // Reset upload count for this test + UPLOAD_COUNT.store(0, std::sync::atomic::Ordering::SeqCst); + + // Create a profile + println!("[test] Creating profile"); + let sample_types = [ValueType::new("samples", "count")]; + let profile_result = unsafe { ddog_prof_Profile_new(Slice::from(&sample_types[..]), None) }; + println!("[test] Profile created"); + let mut profile = match profile_result { + ProfileNewResult::Ok(p) => p, + ProfileNewResult::Err(e) => { + panic!("Failed to create profile: {e}") + } + }; + + // Create sample callbacks + println!("[test] Creating sample callbacks"); + let sample_callbacks = ManagedSampleCallbacks::new(test_converter, test_reset, test_drop); + + // Create config with very long intervals to avoid timer issues + println!("[test] Creating config"); + let config = ProfilerManagerConfig { + channel_depth: 10, + cpu_sampling_interval_ms: 10000, // 10 seconds - very long + upload_interval_ms: 10000, // 10 seconds - very long + }; + + // Start the profiler manager using FFI + println!("[test] Calling ddog_prof_ProfilerManager_start"); + let client_result = unsafe { + ddog_prof_ProfilerManager_start( + &mut profile, + test_cpu_sampler_callback, + test_upload_callback, + sample_callbacks, + config, + ) + }; + println!("[test] ddog_prof_ProfilerManager_start returned"); + + let mut client_handle = match client_result { + ddcommon_ffi::Result::Ok(handle) => handle, + ddcommon_ffi::Result::Err(e) => panic!("Failed to start profiler manager: {e}"), + }; + + println!("[test] Profiler manager started successfully"); + + // Pause immediately without sending any samples + println!("[test] Calling ddog_prof_ProfilerManager_pause"); + let pause_result = unsafe { ddog_prof_ProfilerManager_pause() }; + println!("[test] ddog_prof_ProfilerManager_pause returned"); + match pause_result { + VoidResult::Ok => println!("[test] Profiler manager paused successfully"), + VoidResult::Err(e) => panic!("Failed to pause profiler manager: {e}"), + } + + // Restart the profiler manager in parent (preserves profile data) + println!("[test] Calling ddog_prof_ProfilerManager_restart_in_parent"); + let restart_result = unsafe { ddog_prof_ProfilerManager_restart_in_parent() }; + println!("[test] ddog_prof_ProfilerManager_restart_in_parent returned"); + let mut new_client_handle = match restart_result { + ddcommon_ffi::Result::Ok(handle) => { + println!("[test] Profiler manager restarted successfully"); + handle + } + ddcommon_ffi::Result::Err(e) => panic!("Failed to restart profiler manager: {e}"), + }; + + // Terminate the profiler manager immediately + println!("[test] Calling ddog_prof_ProfilerManager_terminate"); + let terminate_result = unsafe { ddog_prof_ProfilerManager_terminate() }; + println!("[test] ddog_prof_ProfilerManager_terminate returned"); + let _final_profile_handle = match terminate_result { + ddcommon_ffi::Result::Ok(handle) => { + println!("[test] Profiler manager terminated successfully"); + handle + } + ddcommon_ffi::Result::Err(e) => panic!("Failed to terminate profiler manager: {e}"), + }; + + // Drop the client handles + println!("[test] Dropping first client handle"); + let drop_result = unsafe { ddog_prof_ProfilerClient_drop(&mut client_handle) }; + match drop_result { + VoidResult::Ok => println!("[test] First client handle dropped successfully"), + VoidResult::Err(e) => println!("Warning: failed to drop first client handle: {e}"), + } + + println!("[test] Dropping second client handle"); + let drop_result2 = unsafe { ddog_prof_ProfilerClient_drop(&mut new_client_handle) }; + match drop_result2 { + VoidResult::Ok => println!("[test] Second client handle dropped successfully"), + VoidResult::Err(e) => println!("Warning: failed to drop second client handle: {e}"), + } + + println!("[test] Basic lifecycle test completed successfully"); +} diff --git a/datadog-profiling-ffi/tests/test_ffi_lifecycle_with_data.rs b/datadog-profiling-ffi/tests/test_ffi_lifecycle_with_data.rs new file mode 100644 index 0000000000..c82551aa0e --- /dev/null +++ b/datadog-profiling-ffi/tests/test_ffi_lifecycle_with_data.rs @@ -0,0 +1,132 @@ +mod test_utils; +use datadog_profiling_ffi::{ + ddog_prof_Profile_new, ddog_prof_ProfilerClient_drop, ddog_prof_ProfilerManager_enqueue_sample, + ddog_prof_ProfilerManager_pause, ddog_prof_ProfilerManager_reset_for_testing, + ddog_prof_ProfilerManager_restart_in_parent, ddog_prof_ProfilerManager_start, + ddog_prof_ProfilerManager_terminate, ManagedSampleCallbacks, ProfileNewResult, + ProfilerManagerConfig, Slice, ValueType, VoidResult, +}; +use std::ffi::c_void; +use test_utils::*; + +#[test] +fn test_ffi_lifecycle_with_data() { + println!("[test] Starting lifecycle test with data processing"); + // Reset global state for this test + unsafe { ddog_prof_ProfilerManager_reset_for_testing() }.unwrap(); + // Reset upload count for this test + UPLOAD_COUNT.store(0, std::sync::atomic::Ordering::SeqCst); + + // Create a profile + println!("[test] Creating profile"); + let sample_types = [ValueType::new("samples", "count")]; + let profile_result = unsafe { ddog_prof_Profile_new(Slice::from(&sample_types[..]), None) }; + println!("[test] Profile created"); + let mut profile = match profile_result { + ProfileNewResult::Ok(p) => p, + ProfileNewResult::Err(e) => { + panic!("Failed to create profile: {e}") + } + }; + + // Create sample callbacks + println!("[test] Creating sample callbacks"); + let sample_callbacks = ManagedSampleCallbacks::new(test_converter, test_reset, test_drop); + + // Create config with very short intervals for testing + println!("[test] Creating config"); + let config = ProfilerManagerConfig { + channel_depth: 10, + cpu_sampling_interval_ms: 50, // 50ms for faster testing + upload_interval_ms: 100, // 100ms for faster testing + }; + + // Start the profiler manager using FFI + println!("[test] Calling ddog_prof_ProfilerManager_start"); + let client_result = unsafe { + ddog_prof_ProfilerManager_start( + &mut profile, + test_cpu_sampler_callback, + test_upload_callback, + sample_callbacks, + config, + ) + }; + println!("[test] ddog_prof_ProfilerManager_start returned"); + + let mut client_handle = match client_result { + ddcommon_ffi::Result::Ok(handle) => handle, + ddcommon_ffi::Result::Err(e) => panic!("Failed to start profiler manager: {e}"), + }; + + println!("[test] Profiler manager started successfully"); + + // Send a sample using FFI + println!("[test] Sending sample"); + let test_sample = create_test_sample(42); + let sample_ptr = Box::into_raw(Box::new(test_sample)) as *mut c_void; + + let enqueue_result = + unsafe { ddog_prof_ProfilerManager_enqueue_sample(&mut client_handle, sample_ptr) }; + println!("[test] ddog_prof_ProfilerManager_enqueue_sample returned"); + + match enqueue_result { + VoidResult::Ok => println!("[test] Sample enqueued successfully"), + VoidResult::Err(e) => panic!("Failed to enqueue sample: {e}"), + } + + // Give the manager a very short time to process + println!("[test] Sleeping briefly"); + std::thread::sleep(std::time::Duration::from_millis(50)); + println!("[test] Woke up"); + + // Pause the profiler manager using FFI + println!("[test] Calling ddog_prof_ProfilerManager_pause"); + let pause_result = unsafe { ddog_prof_ProfilerManager_pause() }; + println!("[test] ddog_prof_ProfilerManager_pause returned"); + match pause_result { + VoidResult::Ok => println!("[test] Profiler manager paused successfully"), + VoidResult::Err(e) => panic!("Failed to pause profiler manager: {e}"), + } + + // Restart the profiler manager in parent (preserves profile data) + println!("[test] Calling ddog_prof_ProfilerManager_restart_in_parent"); + let restart_result = unsafe { ddog_prof_ProfilerManager_restart_in_parent() }; + println!("[test] ddog_prof_ProfilerManager_restart_in_parent returned"); + let mut new_client_handle = match restart_result { + ddcommon_ffi::Result::Ok(handle) => { + println!("[test] Profiler manager restarted successfully"); + handle + } + ddcommon_ffi::Result::Err(e) => panic!("Failed to restart profiler manager: {e}"), + }; + + // Terminate the profiler manager immediately + println!("[test] Calling ddog_prof_ProfilerManager_terminate"); + let terminate_result = unsafe { ddog_prof_ProfilerManager_terminate() }; + println!("[test] ddog_prof_ProfilerManager_terminate returned"); + let _final_profile_handle = match terminate_result { + ddcommon_ffi::Result::Ok(handle) => { + println!("[test] Profiler manager terminated successfully"); + handle + } + ddcommon_ffi::Result::Err(e) => panic!("Failed to terminate profiler manager: {e}"), + }; + + // Drop the client handles + println!("[test] Dropping first client handle"); + let drop_result = unsafe { ddog_prof_ProfilerClient_drop(&mut client_handle) }; + match drop_result { + VoidResult::Ok => println!("[test] First client handle dropped successfully"), + VoidResult::Err(e) => println!("Warning: failed to drop first client handle: {e}"), + } + + println!("[test] Dropping second client handle"); + let drop_result2 = unsafe { ddog_prof_ProfilerClient_drop(&mut new_client_handle) }; + match drop_result2 { + VoidResult::Ok => println!("[test] Second client handle dropped successfully"), + VoidResult::Err(e) => println!("Warning: failed to drop second client handle: {e}"), + } + + println!("[test] Lifecycle test with data processing completed successfully"); +} diff --git a/datadog-profiling-ffi/tests/test_ffi_start_pause_terminate.rs b/datadog-profiling-ffi/tests/test_ffi_start_pause_terminate.rs new file mode 100644 index 0000000000..71cd212cbd --- /dev/null +++ b/datadog-profiling-ffi/tests/test_ffi_start_pause_terminate.rs @@ -0,0 +1,112 @@ +mod test_utils; +use datadog_profiling_ffi::{ + ddog_prof_Profile_new, ddog_prof_ProfilerClient_drop, ddog_prof_ProfilerManager_enqueue_sample, + ddog_prof_ProfilerManager_pause, ddog_prof_ProfilerManager_reset_for_testing, + ddog_prof_ProfilerManager_start, ddog_prof_ProfilerManager_terminate, ManagedSampleCallbacks, + ProfileNewResult, ProfilerManagerConfig, Slice, ValueType, VoidResult, +}; +use std::ffi::c_void; +use test_utils::*; + +#[test] +fn test_ffi_start_pause_terminate() { + println!("[test] Starting start/pause/terminate test"); + // Reset global state for this test + unsafe { ddog_prof_ProfilerManager_reset_for_testing() }.unwrap(); + // Reset upload count for this test + UPLOAD_COUNT.store(0, std::sync::atomic::Ordering::SeqCst); + + // Create a profile + println!("[test] Creating profile"); + let sample_types = [ValueType::new("samples", "count")]; + let profile_result = unsafe { ddog_prof_Profile_new(Slice::from(&sample_types[..]), None) }; + println!("[test] Profile created"); + let mut profile = match profile_result { + ProfileNewResult::Ok(p) => p, + ProfileNewResult::Err(e) => { + panic!("Failed to create profile: {e}") + } + }; + + // Create sample callbacks + println!("[test] Creating sample callbacks"); + let sample_callbacks = ManagedSampleCallbacks::new(test_converter, test_reset, test_drop); + + // Create config with very short intervals for testing + println!("[test] Creating config"); + let config = ProfilerManagerConfig { + channel_depth: 10, + cpu_sampling_interval_ms: 50, // 50ms for faster testing + upload_interval_ms: 100, // 100ms for faster testing + }; + + // Start the profiler manager using FFI + println!("[test] Calling ddog_prof_ProfilerManager_start"); + let client_result = unsafe { + ddog_prof_ProfilerManager_start( + &mut profile, + test_cpu_sampler_callback, + test_upload_callback, + sample_callbacks, + config, + ) + }; + println!("[test] ddog_prof_ProfilerManager_start returned"); + + let mut client_handle = match client_result { + ddcommon_ffi::Result::Ok(handle) => handle, + ddcommon_ffi::Result::Err(e) => panic!("Failed to start profiler manager: {e}"), + }; + + println!("[test] Profiler manager started successfully"); + + // Send a sample using FFI + println!("[test] Sending sample"); + let test_sample = create_test_sample(42); + let sample_ptr = Box::into_raw(Box::new(test_sample)) as *mut c_void; + + let enqueue_result = + unsafe { ddog_prof_ProfilerManager_enqueue_sample(&mut client_handle, sample_ptr) }; + println!("[test] ddog_prof_ProfilerManager_enqueue_sample returned"); + + match enqueue_result { + VoidResult::Ok => println!("[test] Sample enqueued successfully"), + VoidResult::Err(e) => panic!("Failed to enqueue sample: {e}"), + } + + // Give the manager a very short time to process + println!("[test] Sleeping briefly"); + std::thread::sleep(std::time::Duration::from_millis(50)); + println!("[test] Woke up"); + + // Pause the profiler manager using FFI + println!("[test] Calling ddog_prof_ProfilerManager_pause"); + let pause_result = unsafe { ddog_prof_ProfilerManager_pause() }; + println!("[test] ddog_prof_ProfilerManager_pause returned"); + match pause_result { + VoidResult::Ok => println!("[test] Profiler manager paused successfully"), + VoidResult::Err(e) => panic!("Failed to pause profiler manager: {e}"), + } + + // Terminate the profiler manager immediately + println!("[test] Calling ddog_prof_ProfilerManager_terminate"); + let terminate_result = unsafe { ddog_prof_ProfilerManager_terminate() }; + println!("[test] ddog_prof_ProfilerManager_terminate returned"); + let _final_profile_handle = match terminate_result { + ddcommon_ffi::Result::Ok(handle) => { + println!("[test] Profiler manager terminated successfully"); + handle + } + ddcommon_ffi::Result::Err(e) => panic!("Failed to terminate profiler manager: {e}"), + }; + + // Drop the client handle + println!("[test] Dropping client handle"); + let drop_result = unsafe { ddog_prof_ProfilerClient_drop(&mut client_handle) }; + match drop_result { + VoidResult::Ok => println!("[test] Client handle dropped successfully"), + VoidResult::Err(e) => println!("Warning: failed to drop client handle: {e}"), + } + + println!("[test] Start/pause/terminate test completed successfully"); +} diff --git a/datadog-profiling-ffi/tests/test_ffi_start_terminate.rs b/datadog-profiling-ffi/tests/test_ffi_start_terminate.rs new file mode 100644 index 0000000000..4201fa03ac --- /dev/null +++ b/datadog-profiling-ffi/tests/test_ffi_start_terminate.rs @@ -0,0 +1,103 @@ +mod test_utils; +use datadog_profiling_ffi::{ + ddog_prof_Profile_new, ddog_prof_ProfilerClient_drop, ddog_prof_ProfilerManager_enqueue_sample, + ddog_prof_ProfilerManager_reset_for_testing, ddog_prof_ProfilerManager_start, + ddog_prof_ProfilerManager_terminate, ManagedSampleCallbacks, ProfileNewResult, + ProfilerManagerConfig, Slice, ValueType, VoidResult, +}; +use std::ffi::c_void; +use test_utils::*; + +#[test] +fn test_ffi_start_terminate() { + println!("[test] Starting simple start/terminate test"); + // Reset global state for this test + unsafe { ddog_prof_ProfilerManager_reset_for_testing() }.unwrap(); + // Reset upload count for this test + UPLOAD_COUNT.store(0, std::sync::atomic::Ordering::SeqCst); + + // Create a profile + println!("[test] Creating profile"); + let sample_types = [ValueType::new("samples", "count")]; + let profile_result = unsafe { ddog_prof_Profile_new(Slice::from(&sample_types[..]), None) }; + println!("[test] Profile created"); + let mut profile = match profile_result { + ProfileNewResult::Ok(p) => p, + ProfileNewResult::Err(e) => { + panic!("Failed to create profile: {e}") + } + }; + + // Create sample callbacks + println!("[test] Creating sample callbacks"); + let sample_callbacks = ManagedSampleCallbacks::new(test_converter, test_reset, test_drop); + + // Create config with very short intervals for testing + println!("[test] Creating config"); + let config = ProfilerManagerConfig { + channel_depth: 10, + cpu_sampling_interval_ms: 50, // 50ms for faster testing + upload_interval_ms: 100, // 100ms for faster testing + }; + + // Start the profiler manager using FFI + println!("[test] Calling ddog_prof_ProfilerManager_start"); + let client_result = unsafe { + ddog_prof_ProfilerManager_start( + &mut profile, + test_cpu_sampler_callback, + test_upload_callback, + sample_callbacks, + config, + ) + }; + println!("[test] ddog_prof_ProfilerManager_start returned"); + + let mut client_handle = match client_result { + ddcommon_ffi::Result::Ok(handle) => handle, + ddcommon_ffi::Result::Err(e) => panic!("Failed to start profiler manager: {e}"), + }; + + println!("[test] Profiler manager started successfully"); + + // Send a sample using FFI + println!("[test] Sending sample"); + let test_sample = create_test_sample(42); + let sample_ptr = Box::into_raw(Box::new(test_sample)) as *mut c_void; + + let enqueue_result = + unsafe { ddog_prof_ProfilerManager_enqueue_sample(&mut client_handle, sample_ptr) }; + println!("[test] ddog_prof_ProfilerManager_enqueue_sample returned"); + + match enqueue_result { + VoidResult::Ok => println!("[test] Sample enqueued successfully"), + VoidResult::Err(e) => panic!("Failed to enqueue sample: {e}"), + } + + // Give the manager a very short time to process + println!("[test] Sleeping briefly"); + std::thread::sleep(std::time::Duration::from_millis(50)); + println!("[test] Woke up"); + + // Terminate the profiler manager immediately + println!("[test] Calling ddog_prof_ProfilerManager_terminate"); + let terminate_result = unsafe { ddog_prof_ProfilerManager_terminate() }; + println!("[test] ddog_prof_ProfilerManager_terminate returned"); + let _final_profile_handle = match terminate_result { + ddcommon_ffi::Result::Ok(handle) => { + println!("[test] Profiler manager terminated successfully"); + handle + } + ddcommon_ffi::Result::Err(e) => panic!("Failed to terminate profiler manager: {e}"), + }; + + // Drop the client handle + println!("[test] Dropping client handle"); + let drop_result = unsafe { ddog_prof_ProfilerClient_drop(&mut client_handle) }; + match drop_result { + VoidResult::Ok => println!("[test] Client handle dropped successfully"), + VoidResult::Err(e) => println!("Warning: failed to drop client handle: {e}"), + } + + println!("[test] Simple start/terminate test completed successfully"); +} diff --git a/datadog-profiling-ffi/tests/test_utils.rs b/datadog-profiling-ffi/tests/test_utils.rs new file mode 100644 index 0000000000..9196d4a967 --- /dev/null +++ b/datadog-profiling-ffi/tests/test_utils.rs @@ -0,0 +1,87 @@ +use std::sync::atomic::{AtomicUsize, Ordering}; + +use datadog_profiling::internal::Profile; +use datadog_profiling_ffi::*; +use ddcommon_ffi::{Handle, Slice}; +use tokio_util::sync::CancellationToken; + +pub extern "C" fn test_cpu_sampler_callback(_profile: *mut Profile) {} + +pub static UPLOAD_COUNT: AtomicUsize = AtomicUsize::new(0); + +pub extern "C" fn test_upload_callback( + _profile: *mut Handle, + _token: &mut std::option::Option, +) { + let upload_count = UPLOAD_COUNT.fetch_add(1, Ordering::SeqCst); + println!("[upload_callback] called, count: {upload_count}"); +} + +#[repr(C)] +pub struct TestSample<'a> { + pub values: [i64; 1], + pub locations: [profiles::datatypes::Location<'a>; 1], +} + +#[allow(dead_code)] +pub fn create_test_sample(value: i64) -> TestSample<'static> { + let function = profiles::datatypes::Function { + name: match value { + 42 => "function_1", + 43 => "function_2", + 44 => "function_3", + 45 => "function_4", + 46 => "function_5", + _ => "unknown_function", + } + .into(), + system_name: match value { + 42 => "function_1", + 43 => "function_2", + 44 => "function_3", + 45 => "function_4", + 46 => "function_5", + _ => "unknown_function", + } + .into(), + filename: "test.rs".into(), + ..Default::default() + }; + + TestSample { + values: [value], + locations: [profiles::datatypes::Location { + function, + ..Default::default() + }], + } +} + +pub extern "C" fn test_converter(sample: &SendSample) -> profiles::datatypes::Sample { + let test_sample = unsafe { &*(sample.as_ptr() as *const TestSample) }; + + profiles::datatypes::Sample { + locations: Slice::from(&test_sample.locations[..]), + values: Slice::from(&test_sample.values[..]), + labels: Slice::empty(), + } +} + +pub extern "C" fn test_reset(sample: &mut SendSample) { + let test_sample = unsafe { &mut *(sample.as_ptr() as *mut TestSample) }; + test_sample.values[0] = 0; + test_sample.locations[0] = profiles::datatypes::Location { + function: profiles::datatypes::Function { + name: "".into(), + system_name: "".into(), + filename: "".into(), + ..Default::default() + }, + ..Default::default() + }; +} + +pub extern "C" fn test_drop(sample: SendSample) { + let _test_sample = unsafe { Box::from_raw(sample.as_ptr() as *mut TestSample) }; + // Box will be dropped here, freeing the memory +} From d42cbc8075ea59cd9c37d8d35307fee9c18a98f3 Mon Sep 17 00:00:00 2001 From: Daniel Schwartz-Narbonne Date: Wed, 2 Jul 2025 22:47:27 -0400 Subject: [PATCH 43/52] moving along, cleaning up the manager --- .../src/manager/profiler_manager.rs | 78 +++++++++++++------ 1 file changed, 54 insertions(+), 24 deletions(-) diff --git a/datadog-profiling-ffi/src/manager/profiler_manager.rs b/datadog-profiling-ffi/src/manager/profiler_manager.rs index 64625a7346..b034d7d3ff 100644 --- a/datadog-profiling-ffi/src/manager/profiler_manager.rs +++ b/datadog-profiling-ffi/src/manager/profiler_manager.rs @@ -126,6 +126,8 @@ enum ManagerState { config: ProfilerManagerConfig, callbacks: ManagerCallbacks, }, + /// Manager is in a temporarily invalid state during transitions + Invalid, } // Global state for the profile manager @@ -253,9 +255,22 @@ impl ProfilerManager { // Check if manager is already initialized anyhow::ensure!( matches!(&*state, ManagerState::Uninitialized), - "Manager is already initialized" + "Manager is already initialized or in invalid state" ); + let (client, running_state) = Self::start_internal(profile, callbacks, config)?; + *state = running_state; + + Ok(client) + } + + /// Internal function that handles the actual startup logic. + /// This function does not acquire the lock and returns the state to be stored. + fn start_internal( + profile: internal::Profile, + callbacks: ManagerCallbacks, + config: ProfilerManagerConfig, + ) -> Result<(ManagedProfilerClient, ManagerState)> { let (client_channels, manager_channels) = ClientSampleChannels::new(config.channel_depth); let (shutdown_sender, shutdown_receiver) = crossbeam_channel::bounded(1); let (upload_sender, upload_receiver) = crossbeam_channel::bounded(2); @@ -304,14 +319,14 @@ impl ProfilerManager { is_shutdown, ); - *state = ManagerState::Running { + let running_state = ManagerState::Running { client: client.clone(), controller, config, callbacks, }; - Ok(client) + Ok((client, running_state)) } pub fn pause() -> Result<()> { @@ -323,8 +338,8 @@ impl ProfilerManager { "Manager is not in running state" ); - // Extract the running state and replace with uninitialized - let running_state = match std::mem::replace(&mut *state, ManagerState::Uninitialized) { + // Extract the running state and replace with invalid during transition + let running_state = match std::mem::replace(&mut *state, ManagerState::Invalid) { ManagerState::Running { client: _, controller, @@ -347,29 +362,29 @@ impl ProfilerManager { } pub fn restart_in_parent() -> Result { - let (profile, config, callbacks) = { - let mut state = MANAGER_STATE.lock().map_err(|e| anyhow::anyhow!("{}", e))?; - - let (profile, config, callbacks) = - match std::mem::replace(&mut *state, ManagerState::Uninitialized) { - ManagerState::Paused { - profile, - config, - callbacks, - } => (*profile, config, callbacks), - _ => anyhow::bail!("Manager is not in paused state"), - }; - (profile, config, callbacks) - }; + let mut state = MANAGER_STATE.lock().map_err(|e| anyhow::anyhow!("{}", e))?; + + let (profile, config, callbacks) = + match std::mem::replace(&mut *state, ManagerState::Invalid) { + ManagerState::Paused { + profile, + config, + callbacks, + } => (*profile, config, callbacks), + _ => anyhow::bail!("Manager is not in paused state"), + }; - Self::start(profile, callbacks, config) + let (client, running_state) = Self::start_internal(profile, callbacks, config)?; + *state = running_state; + + Ok(client) } pub fn restart_in_child() -> Result { let mut state = MANAGER_STATE.lock().map_err(|e| anyhow::anyhow!("{}", e))?; let (mut profile, config, callbacks) = - match std::mem::replace(&mut *state, ManagerState::Uninitialized) { + match std::mem::replace(&mut *state, ManagerState::Invalid) { ManagerState::Paused { profile, config, @@ -381,7 +396,10 @@ impl ProfilerManager { // Reset the profile, discarding the previous one let _ = profile.reset_and_return_previous()?; - Self::start(profile, callbacks, config) + let (client, running_state) = Self::start_internal(profile, callbacks, config)?; + *state = running_state; + + Ok(client) } /// Terminates the global profile manager and returns the final profile. @@ -398,8 +416,8 @@ impl ProfilerManager { "Manager is not in running or paused state" ); - // Extract the profile and replace with uninitialized - let profile = match std::mem::replace(&mut *state, ManagerState::Uninitialized) { + // Extract the profile and replace with invalid during transition + let profile = match std::mem::replace(&mut *state, ManagerState::Invalid) { ManagerState::Running { client: _, controller, @@ -420,6 +438,9 @@ impl ProfilerManager { _ => unreachable!(), // We already checked above }; + // Set the final state to uninitialized + *state = ManagerState::Uninitialized; + Ok(profile) } @@ -438,4 +459,13 @@ impl ProfilerManager { Err(e) => Err(format!("Failed to acquire state lock: {e}")), } } + + /// Checks if the manager is in an invalid state. + /// This can be used to detect when the manager is in a transitional state. + pub fn is_invalid() -> Result { + match MANAGER_STATE.lock() { + Ok(state) => Ok(matches!(&*state, ManagerState::Invalid)), + Err(e) => Err(format!("Failed to acquire state lock: {e}")), + } + } } From d286fff6f7c97210e498986c014199cb725d8496 Mon Sep 17 00:00:00 2001 From: Daniel Schwartz-Narbonne Date: Wed, 2 Jul 2025 22:53:27 -0400 Subject: [PATCH 44/52] nicer state extraction code --- .../src/manager/profiler_manager.rs | 85 +++++++------------ 1 file changed, 29 insertions(+), 56 deletions(-) diff --git a/datadog-profiling-ffi/src/manager/profiler_manager.rs b/datadog-profiling-ffi/src/manager/profiler_manager.rs index b034d7d3ff..e787ad52f4 100644 --- a/datadog-profiling-ffi/src/manager/profiler_manager.rs +++ b/datadog-profiling-ffi/src/manager/profiler_manager.rs @@ -332,24 +332,17 @@ impl ProfilerManager { pub fn pause() -> Result<()> { let mut state = MANAGER_STATE.lock().map_err(|e| anyhow::anyhow!("{}", e))?; - // Check if manager is in running state - anyhow::ensure!( - matches!(&*state, ManagerState::Running { .. }), - "Manager is not in running state" - ); - // Extract the running state and replace with invalid during transition - let running_state = match std::mem::replace(&mut *state, ManagerState::Invalid) { - ManagerState::Running { - client: _, - controller, - config, - callbacks, - } => (controller, config, callbacks), - _ => unreachable!(), // We already checked above + let ManagerState::Running { + controller, + config, + callbacks, + .. + } = std::mem::replace(&mut *state, ManagerState::Invalid) + else { + anyhow::bail!("Manager is not in running state"); }; - let (controller, config, callbacks) = running_state; let profile = controller.shutdown()?; *state = ManagerState::Paused { @@ -364,17 +357,16 @@ impl ProfilerManager { pub fn restart_in_parent() -> Result { let mut state = MANAGER_STATE.lock().map_err(|e| anyhow::anyhow!("{}", e))?; - let (profile, config, callbacks) = - match std::mem::replace(&mut *state, ManagerState::Invalid) { - ManagerState::Paused { - profile, - config, - callbacks, - } => (*profile, config, callbacks), - _ => anyhow::bail!("Manager is not in paused state"), - }; + let ManagerState::Paused { + profile, + config, + callbacks, + } = std::mem::replace(&mut *state, ManagerState::Invalid) + else { + anyhow::bail!("Manager is not in paused state"); + }; - let (client, running_state) = Self::start_internal(profile, callbacks, config)?; + let (client, running_state) = Self::start_internal(*profile, callbacks, config)?; *state = running_state; Ok(client) @@ -383,20 +375,19 @@ impl ProfilerManager { pub fn restart_in_child() -> Result { let mut state = MANAGER_STATE.lock().map_err(|e| anyhow::anyhow!("{}", e))?; - let (mut profile, config, callbacks) = - match std::mem::replace(&mut *state, ManagerState::Invalid) { - ManagerState::Paused { - profile, - config, - callbacks, - } => (*profile, config, callbacks), - _ => anyhow::bail!("Manager is not in paused state"), - }; + let ManagerState::Paused { + mut profile, + config, + callbacks, + } = std::mem::replace(&mut *state, ManagerState::Invalid) + else { + anyhow::bail!("Manager is not in paused state"); + }; // Reset the profile, discarding the previous one let _ = profile.reset_and_return_previous()?; - let (client, running_state) = Self::start_internal(profile, callbacks, config)?; + let (client, running_state) = Self::start_internal(*profile, callbacks, config)?; *state = running_state; Ok(client) @@ -407,35 +398,17 @@ impl ProfilerManager { pub fn terminate() -> Result { let mut state = MANAGER_STATE.lock().map_err(|e| anyhow::anyhow!("{}", e))?; - // Check if manager is in running or paused state - anyhow::ensure!( - matches!( - &*state, - ManagerState::Running { .. } | ManagerState::Paused { .. } - ), - "Manager is not in running or paused state" - ); - // Extract the profile and replace with invalid during transition let profile = match std::mem::replace(&mut *state, ManagerState::Invalid) { - ManagerState::Running { - client: _, - controller, - config: _, - callbacks: _, - } => { + ManagerState::Running { controller, .. } => { // Shutdown the controller and get the profile controller.shutdown()? } - ManagerState::Paused { - profile, - config: _, - callbacks: _, - } => { + ManagerState::Paused { profile, .. } => { // Return the stored profile *profile } - _ => unreachable!(), // We already checked above + _ => anyhow::bail!("Manager is not in running or paused state"), }; // Set the final state to uninitialized From b0804df33f0097c29f6c170e1ce12e96d9e59587 Mon Sep 17 00:00:00 2001 From: Daniel Schwartz-Narbonne Date: Thu, 3 Jul 2025 13:36:33 -0400 Subject: [PATCH 45/52] fork tests, almost work --- .../src/manager/profiler_manager.rs | 12 +- .../tests/test_ffi_fork_child_only.rs | 184 +++++++++++ .../tests/test_ffi_fork_data_preservation.rs | 294 ++++++++++++++++++ .../tests/test_ffi_fork_parent_child.rs | 242 ++++++++++++++ .../tests/test_ffi_fork_parent_only.rs | 184 +++++++++++ .../tests/test_ffi_lifecycle_basic.rs | 9 +- .../tests/test_ffi_lifecycle_with_data.rs | 8 +- datadog-profiling-ffi/tests/test_utils.rs | 43 +++ 8 files changed, 971 insertions(+), 5 deletions(-) create mode 100644 datadog-profiling-ffi/tests/test_ffi_fork_child_only.rs create mode 100644 datadog-profiling-ffi/tests/test_ffi_fork_data_preservation.rs create mode 100644 datadog-profiling-ffi/tests/test_ffi_fork_parent_child.rs create mode 100644 datadog-profiling-ffi/tests/test_ffi_fork_parent_only.rs diff --git a/datadog-profiling-ffi/src/manager/profiler_manager.rs b/datadog-profiling-ffi/src/manager/profiler_manager.rs index e787ad52f4..0dfa4ae321 100644 --- a/datadog-profiling-ffi/src/manager/profiler_manager.rs +++ b/datadog-profiling-ffi/src/manager/profiler_manager.rs @@ -340,6 +340,8 @@ impl ProfilerManager { .. } = std::mem::replace(&mut *state, ManagerState::Invalid) else { + // TODO: Consider cleanup when global state is unexpected (e.g., if state is Invalid or + // contains stale resources) anyhow::bail!("Manager is not in running state"); }; @@ -363,6 +365,8 @@ impl ProfilerManager { callbacks, } = std::mem::replace(&mut *state, ManagerState::Invalid) else { + // TODO: Consider cleanup when global state is unexpected (e.g., if state is Invalid or + // contains stale resources) anyhow::bail!("Manager is not in paused state"); }; @@ -381,6 +385,8 @@ impl ProfilerManager { callbacks, } = std::mem::replace(&mut *state, ManagerState::Invalid) else { + // TODO: Consider cleanup when global state is unexpected (e.g., if state is Invalid or + // contains stale resources) anyhow::bail!("Manager is not in paused state"); }; @@ -408,7 +414,11 @@ impl ProfilerManager { // Return the stored profile *profile } - _ => anyhow::bail!("Manager is not in running or paused state"), + _ => { + // TODO: Consider cleanup when global state is unexpected (e.g., if state is Invalid + // or contains stale resources) + anyhow::bail!("Manager is not in running or paused state") + } }; // Set the final state to uninitialized diff --git a/datadog-profiling-ffi/tests/test_ffi_fork_child_only.rs b/datadog-profiling-ffi/tests/test_ffi_fork_child_only.rs new file mode 100644 index 0000000000..f9ffdb96dc --- /dev/null +++ b/datadog-profiling-ffi/tests/test_ffi_fork_child_only.rs @@ -0,0 +1,184 @@ +mod test_utils; +use datadog_profiling_ffi::{ + ddog_prof_Profile_new, ddog_prof_ProfilerClient_drop, ddog_prof_ProfilerManager_enqueue_sample, + ddog_prof_ProfilerManager_pause, ddog_prof_ProfilerManager_reset_for_testing, + ddog_prof_ProfilerManager_restart_in_child, ddog_prof_ProfilerManager_start, + ddog_prof_ProfilerManager_terminate, ManagedSampleCallbacks, ProfileNewResult, + ProfilerManagerConfig, Slice, ValueType, VoidResult, +}; +use ddcommon_ffi::ToInner; +use std::ffi::c_void; +use test_utils::*; + +#[test] +fn test_ffi_fork_child_only() { + println!("[test] Starting fork child-only test"); + + // Reset global state for this test + unsafe { ddog_prof_ProfilerManager_reset_for_testing() }.unwrap(); + // Reset upload count for this test + UPLOAD_COUNT.store(0, std::sync::atomic::Ordering::SeqCst); + + // Create a profile + println!("[test] Creating profile"); + let sample_types = [ValueType::new("samples", "count")]; + let profile_result = unsafe { ddog_prof_Profile_new(Slice::from(&sample_types[..]), None) }; + let mut profile = match profile_result { + ProfileNewResult::Ok(p) => p, + ProfileNewResult::Err(e) => { + panic!("Failed to create profile: {e}") + } + }; + + // Create sample callbacks + let sample_callbacks = ManagedSampleCallbacks::new(test_converter, test_reset, test_drop); + + // Create config with short intervals for testing + let config = ProfilerManagerConfig { + channel_depth: 10, + cpu_sampling_interval_ms: 50, // 50ms for faster testing + upload_interval_ms: 100, // 100ms for faster testing + }; + + // Start the profiler manager + println!("[test] Starting profiler manager"); + let client_result = unsafe { + ddog_prof_ProfilerManager_start( + &mut profile, + test_cpu_sampler_callback, + test_upload_callback, + sample_callbacks, + config, + ) + }; + + let mut client_handle = match client_result { + ddcommon_ffi::Result::Ok(handle) => handle, + ddcommon_ffi::Result::Err(e) => panic!("Failed to start profiler manager: {e}"), + }; + + // Send a sample before forking + println!("[test] Sending sample before fork"); + let test_sample = create_test_sample(42); + let sample_ptr = Box::into_raw(Box::new(test_sample)) as *mut c_void; + + let enqueue_result = + unsafe { ddog_prof_ProfilerManager_enqueue_sample(&mut client_handle, sample_ptr) }; + match enqueue_result { + VoidResult::Ok => println!("[test] Sample enqueued successfully before fork"), + VoidResult::Err(e) => panic!("Failed to enqueue sample before fork: {e}"), + } + + // Give the manager time to process + std::thread::sleep(std::time::Duration::from_millis(50)); + + // Pause the profiler manager before forking + println!("[test] Pausing profiler manager before fork"); + let pause_result = unsafe { ddog_prof_ProfilerManager_pause() }; + match pause_result { + VoidResult::Ok => println!("[test] Profiler manager paused successfully"), + VoidResult::Err(e) => panic!("Failed to pause profiler manager: {e}"), + } + + // Fork the process + println!("[test] Forking process"); + match unsafe { libc::fork() } { + -1 => panic!("Failed to fork"), + 0 => { + // Child process - test restart_in_child + println!("[child] Child process started"); + + // Child should restart with fresh profile (discards previous data) + println!("[child] Restarting profiler manager in child"); + let restart_result = unsafe { ddog_prof_ProfilerManager_restart_in_child() }; + let mut child_client_handle = match restart_result { + ddcommon_ffi::Result::Ok(handle) => { + println!("[child] Profiler manager restarted successfully in child"); + handle + } + ddcommon_ffi::Result::Err(e) => { + panic!("[child] Failed to restart profiler manager in child: {e}") + } + }; + + // Send a sample in child process + println!("[child] Sending sample in child process"); + let child_sample = create_test_sample(100); + let child_sample_ptr = Box::into_raw(Box::new(child_sample)) as *mut c_void; + + let child_enqueue_result = unsafe { + ddog_prof_ProfilerManager_enqueue_sample(&mut child_client_handle, child_sample_ptr) + }; + match child_enqueue_result { + VoidResult::Ok => println!("[child] Sample enqueued successfully in child"), + VoidResult::Err(e) => panic!("[child] Failed to enqueue sample in child: {e}"), + } + + // Give the manager time to process + std::thread::sleep(std::time::Duration::from_millis(50)); + + // Terminate the profiler manager in child + println!("[child] Terminating profiler manager in child"); + let terminate_result = unsafe { ddog_prof_ProfilerManager_terminate() }; + let mut final_profile_handle = match terminate_result { + ddcommon_ffi::Result::Ok(handle) => { + println!("[child] Profiler manager terminated successfully in child"); + handle + } + ddcommon_ffi::Result::Err(e) => { + panic!("[child] Failed to terminate profiler manager in child: {e}") + } + }; + + // Check that the expected sample is present in the final profile + let profile_result = unsafe { final_profile_handle.take() }; + assert_profile_has_sample_values(profile_result, &[100]); + + // Drop the child client handle + let drop_result = unsafe { ddog_prof_ProfilerClient_drop(&mut child_client_handle) }; + match drop_result { + VoidResult::Ok => println!("[child] Child client handle dropped successfully"), + VoidResult::Err(e) => { + println!("[child] Warning: failed to drop child client handle: {e}") + } + } + + println!("[child] Child process completed successfully"); + std::process::exit(0); + } + child_pid => { + // Parent process - just wait for child and clean up + println!( + "[parent] Parent process continuing, child PID: {}", + child_pid + ); + + // Wait for child to complete + println!("[parent] Waiting for child process to complete"); + let mut status = 0; + let wait_result = unsafe { libc::waitpid(child_pid, &mut status, 0) }; + if wait_result == -1 { + panic!("[parent] Failed to wait for child process"); + } + + if libc::WIFEXITED(status) { + let exit_code = libc::WEXITSTATUS(status); + println!("[parent] Child process exited with code: {}", exit_code); + assert_eq!(exit_code, 0, "Child process should exit successfully"); + } else { + println!("[parent] Child process terminated by signal"); + } + + // Drop the original client handle + let drop_result = unsafe { ddog_prof_ProfilerClient_drop(&mut client_handle) }; + match drop_result { + VoidResult::Ok => println!("[parent] Original client handle dropped successfully"), + VoidResult::Err(e) => { + println!("[parent] Warning: failed to drop original client handle: {e}") + } + } + + println!("[parent] Parent process completed successfully"); + } + } +} diff --git a/datadog-profiling-ffi/tests/test_ffi_fork_data_preservation.rs b/datadog-profiling-ffi/tests/test_ffi_fork_data_preservation.rs new file mode 100644 index 0000000000..ba43c30b8f --- /dev/null +++ b/datadog-profiling-ffi/tests/test_ffi_fork_data_preservation.rs @@ -0,0 +1,294 @@ +mod test_utils; +use datadog_profiling_ffi::{ + ddog_prof_Profile_new, ddog_prof_ProfilerClient_drop, ddog_prof_ProfilerManager_enqueue_sample, + ddog_prof_ProfilerManager_pause, ddog_prof_ProfilerManager_reset_for_testing, + ddog_prof_ProfilerManager_restart_in_child, ddog_prof_ProfilerManager_restart_in_parent, + ddog_prof_ProfilerManager_start, ddog_prof_ProfilerManager_terminate, ManagedSampleCallbacks, + ProfileNewResult, ProfilerManagerConfig, Slice, ValueType, VoidResult, +}; +use ddcommon_ffi::ToInner; +use std::ffi::c_void; +use std::sync::atomic::{AtomicUsize, Ordering}; +use test_utils::*; + +// Global counter to track uploads in different processes +static PARENT_UPLOAD_COUNT: AtomicUsize = AtomicUsize::new(0); +static CHILD_UPLOAD_COUNT: AtomicUsize = AtomicUsize::new(0); + +pub extern "C" fn parent_upload_callback( + _profile: *mut ddcommon_ffi::Handle, + _token: &mut std::option::Option, +) { + let upload_count = PARENT_UPLOAD_COUNT.fetch_add(1, Ordering::SeqCst); + println!("[parent_upload_callback] called, count: {upload_count}"); +} + +pub extern "C" fn child_upload_callback( + _profile: *mut ddcommon_ffi::Handle, + _token: &mut std::option::Option, +) { + let upload_count = CHILD_UPLOAD_COUNT.fetch_add(1, Ordering::SeqCst); + println!("[child_upload_callback] called, count: {upload_count}"); +} + +#[test] +fn test_ffi_fork_data_preservation() { + println!("[test] Starting fork data preservation test"); + + // Reset global state for this test + unsafe { ddog_prof_ProfilerManager_reset_for_testing() }.unwrap(); + // Reset upload counts for this test + UPLOAD_COUNT.store(0, std::sync::atomic::Ordering::SeqCst); + PARENT_UPLOAD_COUNT.store(0, std::sync::atomic::Ordering::SeqCst); + CHILD_UPLOAD_COUNT.store(0, std::sync::atomic::Ordering::SeqCst); + + // Create a profile + println!("[test] Creating profile"); + let sample_types = [ValueType::new("samples", "count")]; + let profile_result = unsafe { ddog_prof_Profile_new(Slice::from(&sample_types[..]), None) }; + let mut profile = match profile_result { + ProfileNewResult::Ok(p) => p, + ProfileNewResult::Err(e) => { + panic!("Failed to create profile: {e}") + } + }; + + // Create sample callbacks + let sample_callbacks = ManagedSampleCallbacks::new(test_converter, test_reset, test_drop); + + // Create config with very short intervals to trigger uploads quickly + let config = ProfilerManagerConfig { + channel_depth: 10, + cpu_sampling_interval_ms: 10, // 10ms for very fast testing + upload_interval_ms: 20, // 20ms for very fast testing + }; + + // Start the profiler manager + println!("[test] Starting profiler manager"); + let client_result = unsafe { + ddog_prof_ProfilerManager_start( + &mut profile, + test_cpu_sampler_callback, + test_upload_callback, + sample_callbacks, + config, + ) + }; + + let mut client_handle = match client_result { + ddcommon_ffi::Result::Ok(handle) => handle, + ddcommon_ffi::Result::Err(e) => panic!("Failed to start profiler manager: {e}"), + }; + + // Send multiple samples before forking to accumulate data + println!("[test] Sending samples before fork"); + for i in 0..5 { + let test_sample = create_test_sample(42 + i); + let sample_ptr = Box::into_raw(Box::new(test_sample)) as *mut c_void; + + let enqueue_result = + unsafe { ddog_prof_ProfilerManager_enqueue_sample(&mut client_handle, sample_ptr) }; + match enqueue_result { + VoidResult::Ok => println!("[test] Sample {} enqueued successfully before fork", i), + VoidResult::Err(e) => panic!("Failed to enqueue sample {} before fork: {e}", i), + } + } + + // Give the manager time to process and potentially upload + println!("[test] Waiting for processing before fork"); + std::thread::sleep(std::time::Duration::from_millis(100)); + + // Pause the profiler manager before forking + println!("[test] Pausing profiler manager before fork"); + let pause_result = unsafe { ddog_prof_ProfilerManager_pause() }; + match pause_result { + VoidResult::Ok => println!("[test] Profiler manager paused successfully"), + VoidResult::Err(e) => panic!("Failed to pause profiler manager: {e}"), + } + + // Fork the process + println!("[test] Forking process"); + match unsafe { libc::fork() } { + -1 => panic!("Failed to fork"), + 0 => { + // Child process - should restart with fresh profile (discards previous data) + println!("[child] Child process started"); + + // Child should restart with fresh profile (discards previous data) + println!("[child] Restarting profiler manager in child"); + let restart_result = unsafe { ddog_prof_ProfilerManager_restart_in_child() }; + let mut child_client_handle = match restart_result { + ddcommon_ffi::Result::Ok(handle) => { + println!("[child] Profiler manager restarted successfully in child"); + handle + } + ddcommon_ffi::Result::Err(e) => { + panic!("[child] Failed to restart profiler manager in child: {e}") + } + }; + + // Send a few samples in child process + println!("[child] Sending samples in child process"); + for i in 0..3 { + let child_sample = create_test_sample(100 + i); + let child_sample_ptr = Box::into_raw(Box::new(child_sample)) as *mut c_void; + + let child_enqueue_result = unsafe { + ddog_prof_ProfilerManager_enqueue_sample( + &mut child_client_handle, + child_sample_ptr, + ) + }; + match child_enqueue_result { + VoidResult::Ok => { + println!("[child] Sample {} enqueued successfully in child", i) + } + VoidResult::Err(e) => { + panic!("[child] Failed to enqueue sample {} in child: {e}", i) + } + } + } + + // Give the manager time to process and potentially upload + println!("[child] Waiting for processing in child"); + std::thread::sleep(std::time::Duration::from_millis(100)); + + // Terminate the profiler manager in child + println!("[child] Terminating profiler manager in child"); + let terminate_result = unsafe { ddog_prof_ProfilerManager_terminate() }; + let mut final_profile_handle = match terminate_result { + ddcommon_ffi::Result::Ok(handle) => { + println!("[child] Profiler manager terminated successfully in child"); + handle + } + ddcommon_ffi::Result::Err(e) => { + panic!("[child] Failed to terminate profiler manager in child: {e}") + } + }; + + // Check that the expected sample is present in the final profile + let profile_result = unsafe { final_profile_handle.take() }; + assert_profile_has_sample_values(profile_result, &[100, 101, 102]); + + // Drop the child client handle + let drop_result = unsafe { ddog_prof_ProfilerClient_drop(&mut child_client_handle) }; + match drop_result { + VoidResult::Ok => println!("[child] Child client handle dropped successfully"), + VoidResult::Err(e) => { + println!("[child] Warning: failed to drop child client handle: {e}") + } + } + + println!("[child] Child process completed successfully"); + std::process::exit(0); + } + child_pid => { + // Parent process - should restart preserving profile data + println!( + "[parent] Parent process continuing, child PID: {}", + child_pid + ); + + // Parent should restart preserving profile data + println!("[parent] Restarting profiler manager in parent"); + let restart_result = unsafe { ddog_prof_ProfilerManager_restart_in_parent() }; + let mut parent_client_handle = match restart_result { + ddcommon_ffi::Result::Ok(handle) => { + println!("[parent] Profiler manager restarted successfully in parent"); + handle + } + ddcommon_ffi::Result::Err(e) => { + panic!("[parent] Failed to restart profiler manager in parent: {e}") + } + }; + + // Send a few more samples in parent process + println!("[parent] Sending samples in parent process"); + for i in 0..3 { + let parent_sample = create_test_sample(200 + i); + let parent_sample_ptr = Box::into_raw(Box::new(parent_sample)) as *mut c_void; + + let parent_enqueue_result = unsafe { + ddog_prof_ProfilerManager_enqueue_sample( + &mut parent_client_handle, + parent_sample_ptr, + ) + }; + match parent_enqueue_result { + VoidResult::Ok => { + println!("[parent] Sample {} enqueued successfully in parent", i) + } + VoidResult::Err(e) => { + panic!("[parent] Failed to enqueue sample {} in parent: {e}", i) + } + } + } + + // Give the manager time to process and potentially upload + println!("[parent] Waiting for processing in parent"); + std::thread::sleep(std::time::Duration::from_millis(100)); + + // Wait for child to complete + println!("[parent] Waiting for child process to complete"); + let mut status = 0; + let wait_result = unsafe { libc::waitpid(child_pid, &mut status, 0) }; + if wait_result == -1 { + panic!("[parent] Failed to wait for child process"); + } + + if libc::WIFEXITED(status) { + let exit_code = libc::WEXITSTATUS(status); + println!("[parent] Child process exited with code: {}", exit_code); + assert_eq!(exit_code, 0, "Child process should exit successfully"); + } else { + println!("[parent] Child process terminated by signal"); + } + + // Terminate the profiler manager in parent + println!("[parent] Terminating profiler manager in parent"); + let terminate_result = unsafe { ddog_prof_ProfilerManager_terminate() }; + let mut final_profile_handle = match terminate_result { + ddcommon_ffi::Result::Ok(handle) => { + println!("[parent] Profiler manager terminated successfully in parent"); + handle + } + ddcommon_ffi::Result::Err(e) => { + panic!("[parent] Failed to terminate profiler manager in parent: {e}") + } + }; + + // Check that the expected sample is present in the final profile + let profile_result = unsafe { final_profile_handle.take() }; + assert_profile_has_sample_values(profile_result, &[42, 43, 44, 45, 46, 200, 201, 202]); + + // Drop the client handles + let drop_result = unsafe { ddog_prof_ProfilerClient_drop(&mut client_handle) }; + match drop_result { + VoidResult::Ok => println!("[parent] Original client handle dropped successfully"), + VoidResult::Err(e) => { + println!("[parent] Warning: failed to drop original client handle: {e}") + } + } + + let drop_result2 = unsafe { ddog_prof_ProfilerClient_drop(&mut parent_client_handle) }; + match drop_result2 { + VoidResult::Ok => println!("[parent] Parent client handle dropped successfully"), + VoidResult::Err(e) => { + println!("[parent] Warning: failed to drop parent client handle: {e}") + } + } + + // Verify that we had uploads (indicating data was processed) + let total_uploads = UPLOAD_COUNT.load(Ordering::SeqCst); + println!( + "[parent] Total uploads across all processes: {}", + total_uploads + ); + + // The exact number of uploads may vary due to timing, but we should have some + assert!(total_uploads > 0, "Should have had at least some uploads"); + + println!("[parent] Parent process completed successfully"); + } + } +} diff --git a/datadog-profiling-ffi/tests/test_ffi_fork_parent_child.rs b/datadog-profiling-ffi/tests/test_ffi_fork_parent_child.rs new file mode 100644 index 0000000000..9c1ec7a4f6 --- /dev/null +++ b/datadog-profiling-ffi/tests/test_ffi_fork_parent_child.rs @@ -0,0 +1,242 @@ +mod test_utils; +use datadog_profiling_ffi::{ + ddog_prof_Profile_new, ddog_prof_ProfilerClient_drop, ddog_prof_ProfilerManager_enqueue_sample, + ddog_prof_ProfilerManager_pause, ddog_prof_ProfilerManager_reset_for_testing, + ddog_prof_ProfilerManager_restart_in_child, ddog_prof_ProfilerManager_restart_in_parent, + ddog_prof_ProfilerManager_start, ddog_prof_ProfilerManager_terminate, ManagedSampleCallbacks, + ProfileNewResult, ProfilerManagerConfig, Slice, ValueType, VoidResult, +}; +use ddcommon_ffi::ToInner; +use std::ffi::c_void; + +use test_utils::*; + +#[test] +fn test_ffi_fork_parent_child() { + println!("[test] Starting fork parent/child test"); + + // Reset global state for this test + unsafe { ddog_prof_ProfilerManager_reset_for_testing() }.unwrap(); + // Reset upload count for this test + UPLOAD_COUNT.store(0, std::sync::atomic::Ordering::SeqCst); + + // Create a profile + println!("[test] Creating profile"); + let sample_types = [ValueType::new("samples", "count")]; + let profile_result = unsafe { ddog_prof_Profile_new(Slice::from(&sample_types[..]), None) }; + let mut profile = match profile_result { + ProfileNewResult::Ok(p) => p, + ProfileNewResult::Err(e) => { + panic!("Failed to create profile: {e}") + } + }; + + // Create sample callbacks + let sample_callbacks = ManagedSampleCallbacks::new(test_converter, test_reset, test_drop); + + // Create config with short intervals for testing + let config = ProfilerManagerConfig { + channel_depth: 10, + cpu_sampling_interval_ms: 50, // 50ms for faster testing + upload_interval_ms: 100, // 100ms for faster testing + }; + + // Start the profiler manager + println!("[test] Starting profiler manager"); + let client_result = unsafe { + ddog_prof_ProfilerManager_start( + &mut profile, + test_cpu_sampler_callback, + test_upload_callback, + sample_callbacks, + config, + ) + }; + + let mut client_handle = match client_result { + ddcommon_ffi::Result::Ok(handle) => handle, + ddcommon_ffi::Result::Err(e) => panic!("Failed to start profiler manager: {e}"), + }; + + // Send a sample before forking + println!("[test] Sending sample before fork"); + let test_sample = create_test_sample(42); + let sample_ptr = Box::into_raw(Box::new(test_sample)) as *mut c_void; + + let enqueue_result = + unsafe { ddog_prof_ProfilerManager_enqueue_sample(&mut client_handle, sample_ptr) }; + match enqueue_result { + VoidResult::Ok => println!("[test] Sample enqueued successfully before fork"), + VoidResult::Err(e) => panic!("Failed to enqueue sample before fork: {e}"), + } + + // Give the manager time to process + std::thread::sleep(std::time::Duration::from_millis(50)); + + // Pause the profiler manager before forking + println!("[test] Pausing profiler manager before fork"); + let pause_result = unsafe { ddog_prof_ProfilerManager_pause() }; + match pause_result { + VoidResult::Ok => println!("[test] Profiler manager paused successfully"), + VoidResult::Err(e) => panic!("Failed to pause profiler manager: {e}"), + } + + // Fork the process + println!("[test] Forking process"); + match unsafe { libc::fork() } { + -1 => panic!("Failed to fork"), + 0 => { + // Child process + println!("[child] Child process started"); + + // Child should restart with fresh profile (discards previous data) + println!("[child] Restarting profiler manager in child"); + let restart_result = unsafe { ddog_prof_ProfilerManager_restart_in_child() }; + let mut child_client_handle = match restart_result { + ddcommon_ffi::Result::Ok(handle) => { + println!("[child] Profiler manager restarted successfully in child"); + handle + } + ddcommon_ffi::Result::Err(e) => { + panic!("[child] Failed to restart profiler manager in child: {e}") + } + }; + + // Send a sample in child process + println!("[child] Sending sample in child process"); + let child_sample = create_test_sample(100); + let child_sample_ptr = Box::into_raw(Box::new(child_sample)) as *mut c_void; + + let child_enqueue_result = unsafe { + ddog_prof_ProfilerManager_enqueue_sample(&mut child_client_handle, child_sample_ptr) + }; + match child_enqueue_result { + VoidResult::Ok => println!("[child] Sample enqueued successfully in child"), + VoidResult::Err(e) => panic!("[child] Failed to enqueue sample in child: {e}"), + } + + // Give the manager time to process + std::thread::sleep(std::time::Duration::from_millis(50)); + + // Terminate the profiler manager in child + println!("[child] Terminating profiler manager in child"); + let terminate_result = unsafe { ddog_prof_ProfilerManager_terminate() }; + let mut final_profile_handle = match terminate_result { + ddcommon_ffi::Result::Ok(handle) => { + println!("[child] Profiler manager terminated successfully in child"); + handle + } + ddcommon_ffi::Result::Err(e) => { + panic!("[child] Failed to terminate profiler manager in child: {e}") + } + }; + + // Check that the expected sample is present in the final profile + let profile_result = unsafe { final_profile_handle.take() }; + assert_profile_has_sample_values(profile_result, &[100]); + + // Drop the child client handle + let drop_result = unsafe { ddog_prof_ProfilerClient_drop(&mut child_client_handle) }; + match drop_result { + VoidResult::Ok => println!("[child] Child client handle dropped successfully"), + VoidResult::Err(e) => { + println!("[child] Warning: failed to drop child client handle: {e}") + } + } + + println!("[child] Child process completed successfully"); + std::process::exit(0); + } + child_pid => { + // Parent process + println!( + "[parent] Parent process continuing, child PID: {}", + child_pid + ); + + // Parent should restart preserving profile data + println!("[parent] Restarting profiler manager in parent"); + let restart_result = unsafe { ddog_prof_ProfilerManager_restart_in_parent() }; + let mut parent_client_handle = match restart_result { + ddcommon_ffi::Result::Ok(handle) => { + println!("[parent] Profiler manager restarted successfully in parent"); + handle + } + ddcommon_ffi::Result::Err(e) => { + panic!("[parent] Failed to restart profiler manager in parent: {e}") + } + }; + + // Send another sample in parent process + println!("[parent] Sending sample in parent process"); + let parent_sample = create_test_sample(200); + let parent_sample_ptr = Box::into_raw(Box::new(parent_sample)) as *mut c_void; + + let parent_enqueue_result = unsafe { + ddog_prof_ProfilerManager_enqueue_sample( + &mut parent_client_handle, + parent_sample_ptr, + ) + }; + match parent_enqueue_result { + VoidResult::Ok => println!("[parent] Sample enqueued successfully in parent"), + VoidResult::Err(e) => panic!("[parent] Failed to enqueue sample in parent: {e}"), + } + + // Give the manager time to process + std::thread::sleep(std::time::Duration::from_millis(50)); + + // Wait for child to complete + println!("[parent] Waiting for child process to complete"); + let mut status = 0; + let wait_result = unsafe { libc::waitpid(child_pid, &mut status, 0) }; + if wait_result == -1 { + panic!("[parent] Failed to wait for child process"); + } + + if libc::WIFEXITED(status) { + let exit_code = libc::WEXITSTATUS(status); + println!("[parent] Child process exited with code: {}", exit_code); + assert_eq!(exit_code, 0, "Child process should exit successfully"); + } else { + println!("[parent] Child process terminated by signal"); + } + + // Terminate the profiler manager in parent + println!("[parent] Terminating profiler manager in parent"); + let terminate_result = unsafe { ddog_prof_ProfilerManager_terminate() }; + let mut final_profile_handle = match terminate_result { + ddcommon_ffi::Result::Ok(handle) => { + println!("[parent] Profiler manager terminated successfully in parent"); + handle + } + ddcommon_ffi::Result::Err(e) => { + panic!("[parent] Failed to terminate profiler manager in parent: {e}") + } + }; + + // Check that the expected sample is present in the final profile + let profile_result = unsafe { final_profile_handle.take() }; + assert_profile_has_sample_values(profile_result, &[42, 200]); + + // Drop the client handles + let drop_result = unsafe { ddog_prof_ProfilerClient_drop(&mut client_handle) }; + match drop_result { + VoidResult::Ok => println!("[parent] Original client handle dropped successfully"), + VoidResult::Err(e) => { + println!("[parent] Warning: failed to drop original client handle: {e}") + } + } + + let drop_result2 = unsafe { ddog_prof_ProfilerClient_drop(&mut parent_client_handle) }; + match drop_result2 { + VoidResult::Ok => println!("[parent] Parent client handle dropped successfully"), + VoidResult::Err(e) => { + println!("[parent] Warning: failed to drop parent client handle: {e}") + } + } + + println!("[parent] Parent process completed successfully"); + } + } +} diff --git a/datadog-profiling-ffi/tests/test_ffi_fork_parent_only.rs b/datadog-profiling-ffi/tests/test_ffi_fork_parent_only.rs new file mode 100644 index 0000000000..7258388d5e --- /dev/null +++ b/datadog-profiling-ffi/tests/test_ffi_fork_parent_only.rs @@ -0,0 +1,184 @@ +mod test_utils; +use datadog_profiling_ffi::{ + ddog_prof_Profile_new, ddog_prof_ProfilerClient_drop, ddog_prof_ProfilerManager_enqueue_sample, + ddog_prof_ProfilerManager_pause, ddog_prof_ProfilerManager_reset_for_testing, + ddog_prof_ProfilerManager_restart_in_parent, ddog_prof_ProfilerManager_start, + ddog_prof_ProfilerManager_terminate, ManagedSampleCallbacks, ProfileNewResult, + ProfilerManagerConfig, Slice, ValueType, VoidResult, +}; +use ddcommon_ffi::ToInner; +use std::ffi::c_void; +use test_utils::*; + +#[test] +fn test_ffi_fork_parent_only() { + println!("[test] Starting fork parent-only test"); + + // Reset global state for this test + unsafe { ddog_prof_ProfilerManager_reset_for_testing() }.unwrap(); + // Reset upload count for this test + UPLOAD_COUNT.store(0, std::sync::atomic::Ordering::SeqCst); + + // Create a profile + println!("[test] Creating profile"); + let sample_types = [ValueType::new("samples", "count")]; + let profile_result = unsafe { ddog_prof_Profile_new(Slice::from(&sample_types[..]), None) }; + let mut profile = match profile_result { + ProfileNewResult::Ok(p) => p, + ProfileNewResult::Err(e) => { + panic!("Failed to create profile: {e}") + } + }; + + // Create sample callbacks + let sample_callbacks = ManagedSampleCallbacks::new(test_converter, test_reset, test_drop); + + // Create config with short intervals for testing + let config = ProfilerManagerConfig { + channel_depth: 10, + cpu_sampling_interval_ms: 50, // 50ms for faster testing + upload_interval_ms: 100, // 100ms for faster testing + }; + + // Start the profiler manager + println!("[test] Starting profiler manager"); + let client_result = unsafe { + ddog_prof_ProfilerManager_start( + &mut profile, + test_cpu_sampler_callback, + test_upload_callback, + sample_callbacks, + config, + ) + }; + + let mut client_handle = match client_result { + ddcommon_ffi::Result::Ok(handle) => handle, + ddcommon_ffi::Result::Err(e) => panic!("Failed to start profiler manager: {e}"), + }; + + // Send a sample before forking + println!("[test] Sending sample before fork"); + let test_sample = create_test_sample(42); + let sample_ptr = Box::into_raw(Box::new(test_sample)) as *mut c_void; + + let enqueue_result = + unsafe { ddog_prof_ProfilerManager_enqueue_sample(&mut client_handle, sample_ptr) }; + match enqueue_result { + VoidResult::Ok => println!("[test] Sample enqueued successfully before fork"), + VoidResult::Err(e) => panic!("Failed to enqueue sample before fork: {e}"), + } + + // Give the manager time to process + std::thread::sleep(std::time::Duration::from_millis(50)); + + // Pause the profiler manager before forking + println!("[test] Pausing profiler manager before fork"); + let pause_result = unsafe { ddog_prof_ProfilerManager_pause() }; + match pause_result { + VoidResult::Ok => println!("[test] Profiler manager paused successfully"), + VoidResult::Err(e) => panic!("Failed to pause profiler manager: {e}"), + } + + // Fork the process + println!("[test] Forking process"); + match unsafe { libc::fork() } { + -1 => panic!("Failed to fork"), + 0 => { + // Child process - just exit immediately + println!("[child] Child process started, exiting immediately"); + std::process::exit(0); + } + child_pid => { + // Parent process - test restart_in_parent + println!( + "[parent] Parent process continuing, child PID: {}", + child_pid + ); + + // Parent should restart preserving profile data + println!("[parent] Restarting profiler manager in parent"); + let restart_result = unsafe { ddog_prof_ProfilerManager_restart_in_parent() }; + let mut parent_client_handle = match restart_result { + ddcommon_ffi::Result::Ok(handle) => { + println!("[parent] Profiler manager restarted successfully in parent"); + handle + } + ddcommon_ffi::Result::Err(e) => { + panic!("[parent] Failed to restart profiler manager in parent: {e}") + } + }; + + // Send another sample in parent process + println!("[parent] Sending sample in parent process"); + let parent_sample = create_test_sample(200); + let parent_sample_ptr = Box::into_raw(Box::new(parent_sample)) as *mut c_void; + + let parent_enqueue_result = unsafe { + ddog_prof_ProfilerManager_enqueue_sample( + &mut parent_client_handle, + parent_sample_ptr, + ) + }; + match parent_enqueue_result { + VoidResult::Ok => println!("[parent] Sample enqueued successfully in parent"), + VoidResult::Err(e) => panic!("[parent] Failed to enqueue sample in parent: {e}"), + } + + // Give the manager time to process + std::thread::sleep(std::time::Duration::from_millis(50)); + + // Wait for child to complete + println!("[parent] Waiting for child process to complete"); + let mut status = 0; + let wait_result = unsafe { libc::waitpid(child_pid, &mut status, 0) }; + if wait_result == -1 { + panic!("[parent] Failed to wait for child process"); + } + + if libc::WIFEXITED(status) { + let exit_code = libc::WEXITSTATUS(status); + println!("[parent] Child process exited with code: {}", exit_code); + assert_eq!(exit_code, 0, "Child process should exit successfully"); + } else { + println!("[parent] Child process terminated by signal"); + } + + // Terminate the profiler manager in parent + println!("[parent] Terminating profiler manager in parent"); + let terminate_result = unsafe { ddog_prof_ProfilerManager_terminate() }; + let mut final_profile_handle = match terminate_result { + ddcommon_ffi::Result::Ok(handle) => { + println!("[parent] Profiler manager terminated successfully in parent"); + handle + } + ddcommon_ffi::Result::Err(e) => { + panic!("[parent] Failed to terminate profiler manager in parent: {e}") + } + }; + + // Check that the expected sample is present in the final profile + let profile_result = unsafe { final_profile_handle.take() }; + assert_profile_has_sample_values(profile_result, &[42, 200]); + + // Drop the client handles + let drop_result = unsafe { ddog_prof_ProfilerClient_drop(&mut client_handle) }; + match drop_result { + VoidResult::Ok => println!("[parent] Original client handle dropped successfully"), + VoidResult::Err(e) => { + println!("[parent] Warning: failed to drop original client handle: {e}") + } + } + + let drop_result2 = unsafe { ddog_prof_ProfilerClient_drop(&mut parent_client_handle) }; + match drop_result2 { + VoidResult::Ok => println!("[parent] Parent client handle dropped successfully"), + VoidResult::Err(e) => { + println!("[parent] Warning: failed to drop parent client handle: {e}") + } + } + + println!("[parent] Parent process completed successfully"); + } + } +} diff --git a/datadog-profiling-ffi/tests/test_ffi_lifecycle_basic.rs b/datadog-profiling-ffi/tests/test_ffi_lifecycle_basic.rs index 202eadd256..450aeeab79 100644 --- a/datadog-profiling-ffi/tests/test_ffi_lifecycle_basic.rs +++ b/datadog-profiling-ffi/tests/test_ffi_lifecycle_basic.rs @@ -5,6 +5,7 @@ use datadog_profiling_ffi::{ ddog_prof_ProfilerManager_start, ddog_prof_ProfilerManager_terminate, ManagedSampleCallbacks, ProfileNewResult, ProfilerManagerConfig, Slice, ValueType, VoidResult, }; +use ddcommon_ffi::ToInner; use test_utils::*; #[test] @@ -83,8 +84,7 @@ fn test_ffi_lifecycle_basic() { // Terminate the profiler manager immediately println!("[test] Calling ddog_prof_ProfilerManager_terminate"); let terminate_result = unsafe { ddog_prof_ProfilerManager_terminate() }; - println!("[test] ddog_prof_ProfilerManager_terminate returned"); - let _final_profile_handle = match terminate_result { + let mut final_profile_handle = match terminate_result { ddcommon_ffi::Result::Ok(handle) => { println!("[test] Profiler manager terminated successfully"); handle @@ -92,6 +92,11 @@ fn test_ffi_lifecycle_basic() { ddcommon_ffi::Result::Err(e) => panic!("Failed to terminate profiler manager: {e}"), }; + // Check that the profile is empty (no samples) + let profile_result = unsafe { final_profile_handle.take() }; + let pprof = roundtrip_to_pprof(profile_result); + assert_eq!(pprof.samples.len(), 0, "Profile should have no samples"); + // Drop the client handles println!("[test] Dropping first client handle"); let drop_result = unsafe { ddog_prof_ProfilerClient_drop(&mut client_handle) }; diff --git a/datadog-profiling-ffi/tests/test_ffi_lifecycle_with_data.rs b/datadog-profiling-ffi/tests/test_ffi_lifecycle_with_data.rs index c82551aa0e..22e80a20f3 100644 --- a/datadog-profiling-ffi/tests/test_ffi_lifecycle_with_data.rs +++ b/datadog-profiling-ffi/tests/test_ffi_lifecycle_with_data.rs @@ -6,6 +6,7 @@ use datadog_profiling_ffi::{ ddog_prof_ProfilerManager_terminate, ManagedSampleCallbacks, ProfileNewResult, ProfilerManagerConfig, Slice, ValueType, VoidResult, }; +use ddcommon_ffi::ToInner; use std::ffi::c_void; use test_utils::*; @@ -104,8 +105,7 @@ fn test_ffi_lifecycle_with_data() { // Terminate the profiler manager immediately println!("[test] Calling ddog_prof_ProfilerManager_terminate"); let terminate_result = unsafe { ddog_prof_ProfilerManager_terminate() }; - println!("[test] ddog_prof_ProfilerManager_terminate returned"); - let _final_profile_handle = match terminate_result { + let mut final_profile_handle = match terminate_result { ddcommon_ffi::Result::Ok(handle) => { println!("[test] Profiler manager terminated successfully"); handle @@ -113,6 +113,10 @@ fn test_ffi_lifecycle_with_data() { ddcommon_ffi::Result::Err(e) => panic!("Failed to terminate profiler manager: {e}"), }; + // Check that the expected sample is present in the final profile + let profile_result = unsafe { final_profile_handle.take() }; + assert_profile_has_sample_values(profile_result, &[42]); + // Drop the client handles println!("[test] Dropping first client handle"); let drop_result = unsafe { ddog_prof_ProfilerClient_drop(&mut client_handle) }; diff --git a/datadog-profiling-ffi/tests/test_utils.rs b/datadog-profiling-ffi/tests/test_utils.rs index 9196d4a967..5ae4c187d0 100644 --- a/datadog-profiling-ffi/tests/test_utils.rs +++ b/datadog-profiling-ffi/tests/test_utils.rs @@ -2,7 +2,9 @@ use std::sync::atomic::{AtomicUsize, Ordering}; use datadog_profiling::internal::Profile; use datadog_profiling_ffi::*; +use datadog_profiling_protobuf::prost_impls::Profile as ProstProfile; use ddcommon_ffi::{Handle, Slice}; +use prost::Message; use tokio_util::sync::CancellationToken; pub extern "C" fn test_cpu_sampler_callback(_profile: *mut Profile) {} @@ -85,3 +87,44 @@ pub extern "C" fn test_drop(sample: SendSample) { let _test_sample = unsafe { Box::from_raw(sample.as_ptr() as *mut TestSample) }; // Box will be dropped here, freeing the memory } + +// --- Helpers for profile sample checking --- + +pub fn decode_pprof(encoded: &[u8]) -> ProstProfile { + let mut decoder = lz4_flex::frame::FrameDecoder::new(encoded); + let mut buf = std::vec::Vec::new(); + use std::io::Read; + decoder.read_to_end(&mut buf).unwrap(); + ProstProfile::decode(buf.as_slice()).unwrap() +} + +pub fn roundtrip_to_pprof( + profile: std::result::Result, anyhow::Error>, +) -> ProstProfile { + let encoded = (*profile.expect("Failed to extract profile")) + .serialize_into_compressed_pprof(None, None) + .unwrap(); + decode_pprof(&encoded.buffer) +} + +pub fn assert_profile_has_sample_values( + profile: std::result::Result, anyhow::Error>, + expected_values: &[i64], +) { + let pprof = roundtrip_to_pprof(profile); + let mut found = vec![false; expected_values.len()]; + for sample in &pprof.samples { + for (i, &expected) in expected_values.iter().enumerate() { + if sample.values.contains(&expected) { + found[i] = true; + } + } + } + for (i, &was_found) in found.iter().enumerate() { + assert!( + was_found, + "Expected sample value {} not found in profile", + expected_values[i] + ); + } +} From 742cdd8111b4013d643df9882ebdf927ac92a83f Mon Sep 17 00:00:00 2001 From: Daniel Schwartz-Narbonne Date: Thu, 3 Jul 2025 13:44:47 -0400 Subject: [PATCH 46/52] clippy clean --- .../tests/test_ffi_fork_child_only.rs | 7 ++---- .../tests/test_ffi_fork_data_preservation.rs | 24 +++++++------------ .../tests/test_ffi_fork_parent_child.rs | 7 ++---- .../tests/test_ffi_fork_parent_only.rs | 7 ++---- datadog-profiling-ffi/tests/test_utils.rs | 2 ++ 5 files changed, 17 insertions(+), 30 deletions(-) diff --git a/datadog-profiling-ffi/tests/test_ffi_fork_child_only.rs b/datadog-profiling-ffi/tests/test_ffi_fork_child_only.rs index f9ffdb96dc..16aea0c587 100644 --- a/datadog-profiling-ffi/tests/test_ffi_fork_child_only.rs +++ b/datadog-profiling-ffi/tests/test_ffi_fork_child_only.rs @@ -148,10 +148,7 @@ fn test_ffi_fork_child_only() { } child_pid => { // Parent process - just wait for child and clean up - println!( - "[parent] Parent process continuing, child PID: {}", - child_pid - ); + println!("[parent] Parent process continuing, child PID: {child_pid}"); // Wait for child to complete println!("[parent] Waiting for child process to complete"); @@ -163,7 +160,7 @@ fn test_ffi_fork_child_only() { if libc::WIFEXITED(status) { let exit_code = libc::WEXITSTATUS(status); - println!("[parent] Child process exited with code: {}", exit_code); + println!("[parent] Child process exited with code: {exit_code}"); assert_eq!(exit_code, 0, "Child process should exit successfully"); } else { println!("[parent] Child process terminated by signal"); diff --git a/datadog-profiling-ffi/tests/test_ffi_fork_data_preservation.rs b/datadog-profiling-ffi/tests/test_ffi_fork_data_preservation.rs index ba43c30b8f..4cb859218f 100644 --- a/datadog-profiling-ffi/tests/test_ffi_fork_data_preservation.rs +++ b/datadog-profiling-ffi/tests/test_ffi_fork_data_preservation.rs @@ -89,8 +89,8 @@ fn test_ffi_fork_data_preservation() { let enqueue_result = unsafe { ddog_prof_ProfilerManager_enqueue_sample(&mut client_handle, sample_ptr) }; match enqueue_result { - VoidResult::Ok => println!("[test] Sample {} enqueued successfully before fork", i), - VoidResult::Err(e) => panic!("Failed to enqueue sample {} before fork: {e}", i), + VoidResult::Ok => println!("[test] Sample {i} enqueued successfully before fork"), + VoidResult::Err(e) => panic!("Failed to enqueue sample {i} before fork: {e}"), } } @@ -141,10 +141,10 @@ fn test_ffi_fork_data_preservation() { }; match child_enqueue_result { VoidResult::Ok => { - println!("[child] Sample {} enqueued successfully in child", i) + println!("[child] Sample {i} enqueued successfully in child") } VoidResult::Err(e) => { - panic!("[child] Failed to enqueue sample {} in child: {e}", i) + panic!("[child] Failed to enqueue sample {i} in child: {e}") } } } @@ -184,10 +184,7 @@ fn test_ffi_fork_data_preservation() { } child_pid => { // Parent process - should restart preserving profile data - println!( - "[parent] Parent process continuing, child PID: {}", - child_pid - ); + println!("[parent] Parent process continuing, child PID: {child_pid}"); // Parent should restart preserving profile data println!("[parent] Restarting profiler manager in parent"); @@ -216,10 +213,10 @@ fn test_ffi_fork_data_preservation() { }; match parent_enqueue_result { VoidResult::Ok => { - println!("[parent] Sample {} enqueued successfully in parent", i) + println!("[parent] Sample {i} enqueued successfully in parent") } VoidResult::Err(e) => { - panic!("[parent] Failed to enqueue sample {} in parent: {e}", i) + panic!("[parent] Failed to enqueue sample {i} in parent: {e}") } } } @@ -238,7 +235,7 @@ fn test_ffi_fork_data_preservation() { if libc::WIFEXITED(status) { let exit_code = libc::WEXITSTATUS(status); - println!("[parent] Child process exited with code: {}", exit_code); + println!("[parent] Child process exited with code: {exit_code}"); assert_eq!(exit_code, 0, "Child process should exit successfully"); } else { println!("[parent] Child process terminated by signal"); @@ -280,10 +277,7 @@ fn test_ffi_fork_data_preservation() { // Verify that we had uploads (indicating data was processed) let total_uploads = UPLOAD_COUNT.load(Ordering::SeqCst); - println!( - "[parent] Total uploads across all processes: {}", - total_uploads - ); + println!("[parent] Total uploads across all processes: {total_uploads}"); // The exact number of uploads may vary due to timing, but we should have some assert!(total_uploads > 0, "Should have had at least some uploads"); diff --git a/datadog-profiling-ffi/tests/test_ffi_fork_parent_child.rs b/datadog-profiling-ffi/tests/test_ffi_fork_parent_child.rs index 9c1ec7a4f6..08a2b8bf14 100644 --- a/datadog-profiling-ffi/tests/test_ffi_fork_parent_child.rs +++ b/datadog-profiling-ffi/tests/test_ffi_fork_parent_child.rs @@ -149,10 +149,7 @@ fn test_ffi_fork_parent_child() { } child_pid => { // Parent process - println!( - "[parent] Parent process continuing, child PID: {}", - child_pid - ); + println!("[parent] Parent process continuing, child PID: {child_pid}"); // Parent should restart preserving profile data println!("[parent] Restarting profiler manager in parent"); @@ -196,7 +193,7 @@ fn test_ffi_fork_parent_child() { if libc::WIFEXITED(status) { let exit_code = libc::WEXITSTATUS(status); - println!("[parent] Child process exited with code: {}", exit_code); + println!("[parent] Child process exited with code: {exit_code}"); assert_eq!(exit_code, 0, "Child process should exit successfully"); } else { println!("[parent] Child process terminated by signal"); diff --git a/datadog-profiling-ffi/tests/test_ffi_fork_parent_only.rs b/datadog-profiling-ffi/tests/test_ffi_fork_parent_only.rs index 7258388d5e..673a148d98 100644 --- a/datadog-profiling-ffi/tests/test_ffi_fork_parent_only.rs +++ b/datadog-profiling-ffi/tests/test_ffi_fork_parent_only.rs @@ -91,10 +91,7 @@ fn test_ffi_fork_parent_only() { } child_pid => { // Parent process - test restart_in_parent - println!( - "[parent] Parent process continuing, child PID: {}", - child_pid - ); + println!("[parent] Parent process continuing, child PID: {child_pid}"); // Parent should restart preserving profile data println!("[parent] Restarting profiler manager in parent"); @@ -138,7 +135,7 @@ fn test_ffi_fork_parent_only() { if libc::WIFEXITED(status) { let exit_code = libc::WEXITSTATUS(status); - println!("[parent] Child process exited with code: {}", exit_code); + println!("[parent] Child process exited with code: {exit_code}"); assert_eq!(exit_code, 0, "Child process should exit successfully"); } else { println!("[parent] Child process terminated by signal"); diff --git a/datadog-profiling-ffi/tests/test_utils.rs b/datadog-profiling-ffi/tests/test_utils.rs index 5ae4c187d0..02f61913f1 100644 --- a/datadog-profiling-ffi/tests/test_utils.rs +++ b/datadog-profiling-ffi/tests/test_utils.rs @@ -1,3 +1,5 @@ +#![allow(dead_code)] + use std::sync::atomic::{AtomicUsize, Ordering}; use datadog_profiling::internal::Profile; From 6e3c6f7b57baba80329a3d19d9ead36208c63aa6 Mon Sep 17 00:00:00 2001 From: Daniel Schwartz-Narbonne Date: Tue, 8 Jul 2025 14:44:44 -0400 Subject: [PATCH 47/52] fix the tests --- datadog-profiling-ffi/src/manager/samples.rs | 2 +- .../tests/test_ffi_fork_data_preservation.rs | 89 +++++++++++++------ .../tests/test_ffi_fork_parent_child.rs | 4 +- .../tests/test_ffi_fork_parent_only.rs | 2 +- datadog-profiling-ffi/tests/test_utils.rs | 5 ++ ddcommon-ffi/src/handle.rs | 1 + 6 files changed, 74 insertions(+), 29 deletions(-) diff --git a/datadog-profiling-ffi/src/manager/samples.rs b/datadog-profiling-ffi/src/manager/samples.rs index dd38d484af..9151480cff 100644 --- a/datadog-profiling-ffi/src/manager/samples.rs +++ b/datadog-profiling-ffi/src/manager/samples.rs @@ -80,7 +80,7 @@ impl ClientSampleChannels { /// 2. The caller transfers ownership of the sample to this function /// - The sample is not being used by any other thread /// - The sample must not be accessed by the caller after this call - /// - The manager will either free the sample or recycle it back + /// - The sample will be properly cleaned up if it cannot be sent /// 3. The sample will be properly cleaned up if it cannot be sent pub unsafe fn send_sample(&self, sample: *mut c_void) -> Result<(), SendError> { self.samples_sender.send(SendSample::new(sample)) diff --git a/datadog-profiling-ffi/tests/test_ffi_fork_data_preservation.rs b/datadog-profiling-ffi/tests/test_ffi_fork_data_preservation.rs index 4cb859218f..65929f1240 100644 --- a/datadog-profiling-ffi/tests/test_ffi_fork_data_preservation.rs +++ b/datadog-profiling-ffi/tests/test_ffi_fork_data_preservation.rs @@ -60,7 +60,7 @@ fn test_ffi_fork_data_preservation() { let config = ProfilerManagerConfig { channel_depth: 10, cpu_sampling_interval_ms: 10, // 10ms for very fast testing - upload_interval_ms: 20, // 20ms for very fast testing + upload_interval_ms: 100_000, // 100 seconds - prevent uploads }; // Start the profiler manager @@ -153,22 +153,10 @@ fn test_ffi_fork_data_preservation() { println!("[child] Waiting for processing in child"); std::thread::sleep(std::time::Duration::from_millis(100)); - // Terminate the profiler manager in child - println!("[child] Terminating profiler manager in child"); - let terminate_result = unsafe { ddog_prof_ProfilerManager_terminate() }; - let mut final_profile_handle = match terminate_result { - ddcommon_ffi::Result::Ok(handle) => { - println!("[child] Profiler manager terminated successfully in child"); - handle - } - ddcommon_ffi::Result::Err(e) => { - panic!("[child] Failed to terminate profiler manager in child: {e}") - } - }; - - // Check that the expected sample is present in the final profile - let profile_result = unsafe { final_profile_handle.take() }; - assert_profile_has_sample_values(profile_result, &[100, 101, 102]); + // Don't terminate the profiler manager in child - this causes conflicts with parent + // Instead, just exit cleanly and let the parent handle the termination + println!("[child] Skipping terminate in child to avoid conflicts with parent"); + println!("[child] Child process completed successfully"); // Drop the child client handle let drop_result = unsafe { ddog_prof_ProfilerClient_drop(&mut child_client_handle) }; @@ -213,7 +201,10 @@ fn test_ffi_fork_data_preservation() { }; match parent_enqueue_result { VoidResult::Ok => { - println!("[parent] Sample {i} enqueued successfully in parent") + println!("[parent] Sample {i} enqueued successfully in parent"); + // Add debugging: try to print the profile contents after enqueue + // (This would require an FFI call to extract the profile, which we don't have, so just print a marker) + println!("[parent] (debug) Enqueued sample value {}", 200 + i); } VoidResult::Err(e) => { panic!("[parent] Failed to enqueue sample {i} in parent: {e}") @@ -225,6 +216,9 @@ fn test_ffi_fork_data_preservation() { println!("[parent] Waiting for processing in parent"); std::thread::sleep(std::time::Duration::from_millis(100)); + // Print a marker before terminate + println!("[parent] (debug) About to terminate, expecting to see post-fork samples in profile"); + // Wait for child to complete println!("[parent] Waiting for child process to complete"); let mut status = 0; @@ -238,11 +232,13 @@ fn test_ffi_fork_data_preservation() { println!("[parent] Child process exited with code: {exit_code}"); assert_eq!(exit_code, 0, "Child process should exit successfully"); } else { - println!("[parent] Child process terminated by signal"); + println!("[parent] Child process terminated by signal {status}"); } + + println!("[parent] Child process completed, parent profile state should still be intact"); // Terminate the profiler manager in parent - println!("[parent] Terminating profiler manager in parent"); + println!("[parent] About to terminate profiler manager in parent"); let terminate_result = unsafe { ddog_prof_ProfilerManager_terminate() }; let mut final_profile_handle = match terminate_result { ddcommon_ffi::Result::Ok(handle) => { @@ -256,7 +252,51 @@ fn test_ffi_fork_data_preservation() { // Check that the expected sample is present in the final profile let profile_result = unsafe { final_profile_handle.take() }; - assert_profile_has_sample_values(profile_result, &[42, 43, 44, 45, 46, 200, 201, 202]); + let pprof = roundtrip_to_pprof(profile_result); + println!("[debug] Profile contains {} samples", pprof.samples.len()); + for (i, sample) in pprof.samples.iter().enumerate() { + println!("[debug] Sample {}: values = {:?}", i, sample.values); + } + + // Check pre-fork values + let mut found = [false; 5]; + for sample in &pprof.samples { + for (i, &expected) in [42, 43, 44, 45, 46].iter().enumerate() { + if sample.values.contains(&expected) { + found[i] = true; + } + } + } + for (i, &was_found) in found.iter().enumerate() { + assert!(was_found, "Expected pre-fork sample value {} not found in profile", [42, 43, 44, 45, 46][i]); + } + + // Check for merged post-fork sample + let mut found_merged = false; + for sample in &pprof.samples { + if sample.values.contains(&603) { + // Check function name + if let Some(loc_id) = sample.location_ids.get(0) { + let loc_obj = pprof.locations.iter().find(|l| l.id == *loc_id); + if let Some(loc_obj) = loc_obj { + let fn_id = loc_obj.lines.get(0).map(|l| l.function_id); + if let Some(fn_id) = fn_id { + let fn_obj = pprof.functions.iter().find(|f| f.id == fn_id); + if let Some(fn_obj) = fn_obj { + // fn_obj.name is an index into the string table + let name_idx = fn_obj.name as usize; + if let Some(name) = pprof.string_table.get(name_idx) { + if name == "unknown_function" { + found_merged = true; + } + } + } + } + } + } + } + } + assert!(found_merged, "Expected merged post-fork sample value 603 with function name 'unknown_function' not found in profile"); // Drop the client handles let drop_result = unsafe { ddog_prof_ProfilerClient_drop(&mut client_handle) }; @@ -275,13 +315,12 @@ fn test_ffi_fork_data_preservation() { } } - // Verify that we had uploads (indicating data was processed) + // Note: We don't require uploads in this test since we're using a long upload interval + // to prevent premature uploads that could interfere with data preservation testing. + // The test has already verified that samples are correctly preserved across fork boundaries. let total_uploads = UPLOAD_COUNT.load(Ordering::SeqCst); println!("[parent] Total uploads across all processes: {total_uploads}"); - // The exact number of uploads may vary due to timing, but we should have some - assert!(total_uploads > 0, "Should have had at least some uploads"); - println!("[parent] Parent process completed successfully"); } } diff --git a/datadog-profiling-ffi/tests/test_ffi_fork_parent_child.rs b/datadog-profiling-ffi/tests/test_ffi_fork_parent_child.rs index 08a2b8bf14..b6fe219c81 100644 --- a/datadog-profiling-ffi/tests/test_ffi_fork_parent_child.rs +++ b/datadog-profiling-ffi/tests/test_ffi_fork_parent_child.rs @@ -34,11 +34,11 @@ fn test_ffi_fork_parent_child() { // Create sample callbacks let sample_callbacks = ManagedSampleCallbacks::new(test_converter, test_reset, test_drop); - // Create config with short intervals for testing + // Create config with long upload interval to prevent premature uploads during test let config = ProfilerManagerConfig { channel_depth: 10, cpu_sampling_interval_ms: 50, // 50ms for faster testing - upload_interval_ms: 100, // 100ms for faster testing + upload_interval_ms: 10000, // 10 seconds - very long to prevent uploads during test }; // Start the profiler manager diff --git a/datadog-profiling-ffi/tests/test_ffi_fork_parent_only.rs b/datadog-profiling-ffi/tests/test_ffi_fork_parent_only.rs index 673a148d98..f0d17e44e9 100644 --- a/datadog-profiling-ffi/tests/test_ffi_fork_parent_only.rs +++ b/datadog-profiling-ffi/tests/test_ffi_fork_parent_only.rs @@ -138,7 +138,7 @@ fn test_ffi_fork_parent_only() { println!("[parent] Child process exited with code: {exit_code}"); assert_eq!(exit_code, 0, "Child process should exit successfully"); } else { - println!("[parent] Child process terminated by signal"); + println!("[parent] Child process terminated by signal {status}"); } // Terminate the profiler manager in parent diff --git a/datadog-profiling-ffi/tests/test_utils.rs b/datadog-profiling-ffi/tests/test_utils.rs index 02f61913f1..909505abb6 100644 --- a/datadog-profiling-ffi/tests/test_utils.rs +++ b/datadog-profiling-ffi/tests/test_utils.rs @@ -114,6 +114,11 @@ pub fn assert_profile_has_sample_values( expected_values: &[i64], ) { let pprof = roundtrip_to_pprof(profile); + println!("[debug] Profile contains {} samples", pprof.samples.len()); + for (i, sample) in pprof.samples.iter().enumerate() { + println!("[debug] Sample {}: values = {:?}", i, sample.values); + } + let mut found = vec![false; expected_values.len()]; for sample in &pprof.samples { for (i, &expected) in expected_values.iter().enumerate() { diff --git a/ddcommon-ffi/src/handle.rs b/ddcommon-ffi/src/handle.rs index 25a19b59b0..2740fa889c 100644 --- a/ddcommon-ffi/src/handle.rs +++ b/ddcommon-ffi/src/handle.rs @@ -8,6 +8,7 @@ use anyhow::Context; /// Represents an object that should only be referred to by its handle. /// Do not access its member for any reason, only use the C API functions on this struct. #[repr(C)] +#[derive(Debug)] pub struct Handle { // This may be null, but if not it will point to a valid . inner: *mut T, From 8fd752c6ffb0c3148dc3df3e7cccd1e4c69b1d22 Mon Sep 17 00:00:00 2001 From: Daniel Schwartz-Narbonne Date: Tue, 8 Jul 2025 14:48:14 -0400 Subject: [PATCH 48/52] clippy --- .../tests/test_ffi_fork_data_preservation.rs | 22 +++++++++++++------ datadog-profiling-ffi/tests/test_utils.rs | 2 +- 2 files changed, 16 insertions(+), 8 deletions(-) diff --git a/datadog-profiling-ffi/tests/test_ffi_fork_data_preservation.rs b/datadog-profiling-ffi/tests/test_ffi_fork_data_preservation.rs index 65929f1240..028e7fed39 100644 --- a/datadog-profiling-ffi/tests/test_ffi_fork_data_preservation.rs +++ b/datadog-profiling-ffi/tests/test_ffi_fork_data_preservation.rs @@ -203,7 +203,8 @@ fn test_ffi_fork_data_preservation() { VoidResult::Ok => { println!("[parent] Sample {i} enqueued successfully in parent"); // Add debugging: try to print the profile contents after enqueue - // (This would require an FFI call to extract the profile, which we don't have, so just print a marker) + // (This would require an FFI call to extract the profile, which we don't + // have, so just print a marker) println!("[parent] (debug) Enqueued sample value {}", 200 + i); } VoidResult::Err(e) => { @@ -234,8 +235,10 @@ fn test_ffi_fork_data_preservation() { } else { println!("[parent] Child process terminated by signal {status}"); } - - println!("[parent] Child process completed, parent profile state should still be intact"); + + println!( + "[parent] Child process completed, parent profile state should still be intact" + ); // Terminate the profiler manager in parent println!("[parent] About to terminate profiler manager in parent"); @@ -268,7 +271,11 @@ fn test_ffi_fork_data_preservation() { } } for (i, &was_found) in found.iter().enumerate() { - assert!(was_found, "Expected pre-fork sample value {} not found in profile", [42, 43, 44, 45, 46][i]); + assert!( + was_found, + "Expected pre-fork sample value {} not found in profile", + [42, 43, 44, 45, 46][i] + ); } // Check for merged post-fork sample @@ -276,10 +283,10 @@ fn test_ffi_fork_data_preservation() { for sample in &pprof.samples { if sample.values.contains(&603) { // Check function name - if let Some(loc_id) = sample.location_ids.get(0) { + if let Some(loc_id) = sample.location_ids.first() { let loc_obj = pprof.locations.iter().find(|l| l.id == *loc_id); if let Some(loc_obj) = loc_obj { - let fn_id = loc_obj.lines.get(0).map(|l| l.function_id); + let fn_id = loc_obj.lines.first().map(|l| l.function_id); if let Some(fn_id) = fn_id { let fn_obj = pprof.functions.iter().find(|f| f.id == fn_id); if let Some(fn_obj) = fn_obj { @@ -317,7 +324,8 @@ fn test_ffi_fork_data_preservation() { // Note: We don't require uploads in this test since we're using a long upload interval // to prevent premature uploads that could interfere with data preservation testing. - // The test has already verified that samples are correctly preserved across fork boundaries. + // The test has already verified that samples are correctly preserved across fork + // boundaries. let total_uploads = UPLOAD_COUNT.load(Ordering::SeqCst); println!("[parent] Total uploads across all processes: {total_uploads}"); diff --git a/datadog-profiling-ffi/tests/test_utils.rs b/datadog-profiling-ffi/tests/test_utils.rs index 909505abb6..b474ec8f76 100644 --- a/datadog-profiling-ffi/tests/test_utils.rs +++ b/datadog-profiling-ffi/tests/test_utils.rs @@ -118,7 +118,7 @@ pub fn assert_profile_has_sample_values( for (i, sample) in pprof.samples.iter().enumerate() { println!("[debug] Sample {}: values = {:?}", i, sample.values); } - + let mut found = vec![false; expected_values.len()]; for sample in &pprof.samples { for (i, &expected) in expected_values.iter().enumerate() { From 7dc3a40961cca6f501c0a83909d891ce41b8fedb Mon Sep 17 00:00:00 2001 From: Daniel Schwartz-Narbonne Date: Tue, 8 Jul 2025 14:57:41 -0400 Subject: [PATCH 49/52] terminate back in --- .../tests/test_ffi_fork_data_preservation.rs | 19 +- .../tests/test_ffi_fork_lifecycle.rs | 239 ++++++++++++++++++ 2 files changed, 254 insertions(+), 4 deletions(-) create mode 100644 datadog-profiling-ffi/tests/test_ffi_fork_lifecycle.rs diff --git a/datadog-profiling-ffi/tests/test_ffi_fork_data_preservation.rs b/datadog-profiling-ffi/tests/test_ffi_fork_data_preservation.rs index 028e7fed39..e75c7edbef 100644 --- a/datadog-profiling-ffi/tests/test_ffi_fork_data_preservation.rs +++ b/datadog-profiling-ffi/tests/test_ffi_fork_data_preservation.rs @@ -153,10 +153,21 @@ fn test_ffi_fork_data_preservation() { println!("[child] Waiting for processing in child"); std::thread::sleep(std::time::Duration::from_millis(100)); - // Don't terminate the profiler manager in child - this causes conflicts with parent - // Instead, just exit cleanly and let the parent handle the termination - println!("[child] Skipping terminate in child to avoid conflicts with parent"); - println!("[child] Child process completed successfully"); + // Terminate the profiler manager in child (added back) + println!("[child] Terminating profiler manager in child"); + let terminate_result = unsafe { ddog_prof_ProfilerManager_terminate() }; + let mut _final_profile_handle = match terminate_result { + ddcommon_ffi::Result::Ok(handle) => { + println!("[child] Profiler manager terminated successfully in child"); + handle + } + ddcommon_ffi::Result::Err(e) => { + panic!("[child] Failed to terminate profiler manager in child: {e}") + } + }; + // Extract the profile and assert expected values + let profile_result = unsafe { _final_profile_handle.take() }; + assert_profile_has_sample_values(profile_result, &[100, 101, 102]); // Drop the child client handle let drop_result = unsafe { ddog_prof_ProfilerClient_drop(&mut child_client_handle) }; diff --git a/datadog-profiling-ffi/tests/test_ffi_fork_lifecycle.rs b/datadog-profiling-ffi/tests/test_ffi_fork_lifecycle.rs new file mode 100644 index 0000000000..bf2dde52d1 --- /dev/null +++ b/datadog-profiling-ffi/tests/test_ffi_fork_lifecycle.rs @@ -0,0 +1,239 @@ +mod test_utils; +use datadog_profiling_ffi::{ + ddog_prof_Profile_new, ddog_prof_ProfilerClient_drop, ddog_prof_ProfilerManager_enqueue_sample, + ddog_prof_ProfilerManager_pause, ddog_prof_ProfilerManager_reset_for_testing, + ddog_prof_ProfilerManager_restart_in_child, ddog_prof_ProfilerManager_restart_in_parent, + ddog_prof_ProfilerManager_start, ddog_prof_ProfilerManager_terminate, ManagedSampleCallbacks, + ProfileNewResult, ProfilerManagerConfig, Slice, ValueType, VoidResult, +}; +use ddcommon_ffi::ToInner; +use std::ffi::c_void; +use test_utils::*; + +#[test] +fn test_ffi_fork_lifecycle() { + println!("[test] Starting fork lifecycle test"); + + // Reset global state for this test + unsafe { ddog_prof_ProfilerManager_reset_for_testing() }.unwrap(); + // Reset upload count for this test + UPLOAD_COUNT.store(0, std::sync::atomic::Ordering::SeqCst); + + // Create a profile + println!("[test] Creating profile"); + let sample_types = [ValueType::new("samples", "count")]; + let profile_result = unsafe { ddog_prof_Profile_new(Slice::from(&sample_types[..]), None) }; + let mut profile = match profile_result { + ProfileNewResult::Ok(p) => p, + ProfileNewResult::Err(e) => { + panic!("Failed to create profile: {e}") + } + }; + + // Create sample callbacks + let sample_callbacks = ManagedSampleCallbacks::new(test_converter, test_reset, test_drop); + + // Create config with long intervals to prevent uploads during test + let config = ProfilerManagerConfig { + channel_depth: 10, + cpu_sampling_interval_ms: 50, // 50ms for faster testing + upload_interval_ms: 10000, // 10 seconds - prevent uploads during test + }; + + // Start the profiler manager + println!("[test] Starting profiler manager"); + let client_result = unsafe { + ddog_prof_ProfilerManager_start( + &mut profile, + test_cpu_sampler_callback, + test_upload_callback, + sample_callbacks, + config, + ) + }; + + let mut client_handle = match client_result { + ddcommon_ffi::Result::Ok(handle) => handle, + ddcommon_ffi::Result::Err(e) => panic!("Failed to start profiler manager: {e}"), + }; + + // Send a sample before forking + println!("[test] Sending sample before fork"); + let test_sample = create_test_sample(42); + let sample_ptr = Box::into_raw(Box::new(test_sample)) as *mut c_void; + + let enqueue_result = + unsafe { ddog_prof_ProfilerManager_enqueue_sample(&mut client_handle, sample_ptr) }; + match enqueue_result { + VoidResult::Ok => println!("[test] Sample enqueued successfully before fork"), + VoidResult::Err(e) => panic!("Failed to enqueue sample before fork: {e}"), + } + + // Give the manager time to process + std::thread::sleep(std::time::Duration::from_millis(50)); + + // Pause the profiler manager before forking + println!("[test] Pausing profiler manager before fork"); + let pause_result = unsafe { ddog_prof_ProfilerManager_pause() }; + match pause_result { + VoidResult::Ok => println!("[test] Profiler manager paused successfully"), + VoidResult::Err(e) => panic!("Failed to pause profiler manager: {e}"), + } + + // Fork the process + println!("[test] Forking process"); + match unsafe { libc::fork() } { + -1 => panic!("Failed to fork"), + 0 => { + // Child process + println!("[child] Child process started"); + + // Child should restart with fresh profile (discards previous data) + println!("[child] Restarting profiler manager in child"); + let restart_result = unsafe { ddog_prof_ProfilerManager_restart_in_child() }; + let mut child_client_handle = match restart_result { + ddcommon_ffi::Result::Ok(handle) => { + println!("[child] Profiler manager restarted successfully in child"); + handle + } + ddcommon_ffi::Result::Err(e) => { + panic!("[child] Failed to restart profiler manager in child: {e}") + } + }; + + // Send a sample in child process + println!("[child] Sending sample in child process"); + let child_sample = create_test_sample(100); + let child_sample_ptr = Box::into_raw(Box::new(child_sample)) as *mut c_void; + + let child_enqueue_result = unsafe { + ddog_prof_ProfilerManager_enqueue_sample(&mut child_client_handle, child_sample_ptr) + }; + match child_enqueue_result { + VoidResult::Ok => println!("[child] Sample enqueued successfully in child"), + VoidResult::Err(e) => panic!("[child] Failed to enqueue sample in child: {e}"), + } + + // Give the manager time to process + std::thread::sleep(std::time::Duration::from_millis(50)); + + // Terminate the profiler manager in child + println!("[child] Terminating profiler manager in child"); + let terminate_result = unsafe { ddog_prof_ProfilerManager_terminate() }; + let mut final_profile_handle = match terminate_result { + ddcommon_ffi::Result::Ok(handle) => { + println!("[child] Profiler manager terminated successfully in child"); + handle + } + ddcommon_ffi::Result::Err(e) => { + panic!("[child] Failed to terminate profiler manager in child: {e}") + } + }; + + // Check that the expected sample is present in the final profile + let profile_result = unsafe { final_profile_handle.take() }; + assert_profile_has_sample_values(profile_result, &[100]); + + // Drop the child client handle + let drop_result = unsafe { ddog_prof_ProfilerClient_drop(&mut child_client_handle) }; + match drop_result { + VoidResult::Ok => println!("[child] Child client handle dropped successfully"), + VoidResult::Err(e) => { + println!("[child] Warning: failed to drop child client handle: {e}") + } + } + + println!("[child] Child process completed successfully"); + std::process::exit(0); + } + child_pid => { + // Parent process + println!("[parent] Parent process continuing, child PID: {child_pid}"); + + // Parent should restart preserving profile data + println!("[parent] Restarting profiler manager in parent"); + let restart_result = unsafe { ddog_prof_ProfilerManager_restart_in_parent() }; + let mut parent_client_handle = match restart_result { + ddcommon_ffi::Result::Ok(handle) => { + println!("[parent] Profiler manager restarted successfully in parent"); + handle + } + ddcommon_ffi::Result::Err(e) => { + panic!("[parent] Failed to restart profiler manager in parent: {e}") + } + }; + + // Send a sample in parent process + println!("[parent] Sending sample in parent process"); + let parent_sample = create_test_sample(200); + let parent_sample_ptr = Box::into_raw(Box::new(parent_sample)) as *mut c_void; + + let parent_enqueue_result = unsafe { + ddog_prof_ProfilerManager_enqueue_sample( + &mut parent_client_handle, + parent_sample_ptr, + ) + }; + match parent_enqueue_result { + VoidResult::Ok => println!("[parent] Sample enqueued successfully in parent"), + VoidResult::Err(e) => panic!("[parent] Failed to enqueue sample in parent: {e}"), + } + + // Give the manager time to process + std::thread::sleep(std::time::Duration::from_millis(50)); + + // Wait for child to complete + println!("[parent] Waiting for child process to complete"); + let mut status = 0; + let wait_result = unsafe { libc::waitpid(child_pid, &mut status, 0) }; + if wait_result == -1 { + panic!("[parent] Failed to wait for child process"); + } + + if libc::WIFEXITED(status) { + let exit_code = libc::WEXITSTATUS(status); + println!("[parent] Child process exited with code: {exit_code}"); + assert_eq!(exit_code, 0, "Child process should exit successfully"); + } else { + println!("[parent] Child process terminated by signal"); + } + + // Terminate the profiler manager in parent + println!("[parent] Terminating profiler manager in parent"); + let terminate_result = unsafe { ddog_prof_ProfilerManager_terminate() }; + let mut final_profile_handle = match terminate_result { + ddcommon_ffi::Result::Ok(handle) => { + println!("[parent] Profiler manager terminated successfully in parent"); + handle + } + ddcommon_ffi::Result::Err(e) => { + panic!("[parent] Failed to terminate profiler manager in parent: {e}") + } + }; + + // Check that the expected samples are present in the final profile + // Parent should have both pre-fork (42) and post-fork (200) samples + let profile_result = unsafe { final_profile_handle.take() }; + assert_profile_has_sample_values(profile_result, &[42, 200]); + + // Drop the client handles + let drop_result = unsafe { ddog_prof_ProfilerClient_drop(&mut client_handle) }; + match drop_result { + VoidResult::Ok => println!("[parent] Original client handle dropped successfully"), + VoidResult::Err(e) => { + println!("[parent] Warning: failed to drop original client handle: {e}") + } + } + + let drop_result2 = unsafe { ddog_prof_ProfilerClient_drop(&mut parent_client_handle) }; + match drop_result2 { + VoidResult::Ok => println!("[parent] Parent client handle dropped successfully"), + VoidResult::Err(e) => { + println!("[parent] Warning: failed to drop parent client handle: {e}") + } + } + + println!("[parent] Parent process completed successfully"); + } + } +} From 9bae67478053d05b7b3ed726675cae61c7f080f5 Mon Sep 17 00:00:00 2001 From: Daniel Schwartz-Narbonne Date: Tue, 8 Jul 2025 15:27:32 -0400 Subject: [PATCH 50/52] copyright headers --- datadog-profiling-ffi/src/manager/client.rs | 2 ++ datadog-profiling-ffi/src/manager/ffi_utils.rs | 2 ++ datadog-profiling-ffi/src/manager/profiler_manager.rs | 2 ++ datadog-profiling-ffi/src/manager/samples.rs | 2 ++ datadog-profiling-ffi/src/manager/tests.rs | 2 ++ datadog-profiling-ffi/tests/test_ffi_fork_child_only.rs | 2 ++ datadog-profiling-ffi/tests/test_ffi_fork_data_preservation.rs | 2 ++ datadog-profiling-ffi/tests/test_ffi_fork_lifecycle.rs | 2 ++ datadog-profiling-ffi/tests/test_ffi_fork_parent_child.rs | 2 ++ datadog-profiling-ffi/tests/test_ffi_fork_parent_only.rs | 2 ++ datadog-profiling-ffi/tests/test_ffi_lifecycle_basic.rs | 2 ++ datadog-profiling-ffi/tests/test_ffi_lifecycle_with_data.rs | 2 ++ datadog-profiling-ffi/tests/test_ffi_start_pause_terminate.rs | 2 ++ datadog-profiling-ffi/tests/test_ffi_start_terminate.rs | 2 ++ datadog-profiling-ffi/tests/test_utils.rs | 2 ++ 15 files changed, 30 insertions(+) diff --git a/datadog-profiling-ffi/src/manager/client.rs b/datadog-profiling-ffi/src/manager/client.rs index d0db6229eb..4c5a690bc7 100644 --- a/datadog-profiling-ffi/src/manager/client.rs +++ b/datadog-profiling-ffi/src/manager/client.rs @@ -1,3 +1,5 @@ +// Copyright 2025-Present Datadog, Inc. https://www.datadoghq.com/ + use std::sync::Arc; use std::{ffi::c_void, sync::atomic::AtomicBool}; diff --git a/datadog-profiling-ffi/src/manager/ffi_utils.rs b/datadog-profiling-ffi/src/manager/ffi_utils.rs index 297a3c1819..ae08fb2a8c 100644 --- a/datadog-profiling-ffi/src/manager/ffi_utils.rs +++ b/datadog-profiling-ffi/src/manager/ffi_utils.rs @@ -1,3 +1,5 @@ +// Copyright 2025-Present Datadog, Inc. https://www.datadoghq.com/ + use std::num::NonZeroI64; use crate::profiles::datatypes::{ProfileResult, Sample}; diff --git a/datadog-profiling-ffi/src/manager/profiler_manager.rs b/datadog-profiling-ffi/src/manager/profiler_manager.rs index 0dfa4ae321..920649456d 100644 --- a/datadog-profiling-ffi/src/manager/profiler_manager.rs +++ b/datadog-profiling-ffi/src/manager/profiler_manager.rs @@ -1,3 +1,5 @@ +// Copyright 2025-Present Datadog, Inc. https://www.datadoghq.com/ + #![allow(clippy::unwrap_used)] use std::sync::atomic::AtomicBool; diff --git a/datadog-profiling-ffi/src/manager/samples.rs b/datadog-profiling-ffi/src/manager/samples.rs index 9151480cff..89fd2a9f72 100644 --- a/datadog-profiling-ffi/src/manager/samples.rs +++ b/datadog-profiling-ffi/src/manager/samples.rs @@ -1,3 +1,5 @@ +// Copyright 2025-Present Datadog, Inc. https://www.datadoghq.com/ + use std::ffi::c_void; use crossbeam_channel::{Receiver, SendError, Sender, TryRecvError}; diff --git a/datadog-profiling-ffi/src/manager/tests.rs b/datadog-profiling-ffi/src/manager/tests.rs index d924117163..233a0b919c 100644 --- a/datadog-profiling-ffi/src/manager/tests.rs +++ b/datadog-profiling-ffi/src/manager/tests.rs @@ -1,3 +1,5 @@ +// Copyright 2025-Present Datadog, Inc. https://www.datadoghq.com/ + use std::ffi::c_void; use std::sync::atomic::{AtomicUsize, Ordering}; diff --git a/datadog-profiling-ffi/tests/test_ffi_fork_child_only.rs b/datadog-profiling-ffi/tests/test_ffi_fork_child_only.rs index 16aea0c587..695becc1ce 100644 --- a/datadog-profiling-ffi/tests/test_ffi_fork_child_only.rs +++ b/datadog-profiling-ffi/tests/test_ffi_fork_child_only.rs @@ -1,3 +1,5 @@ +// Copyright 2025-Present Datadog, Inc. https://www.datadoghq.com/ + mod test_utils; use datadog_profiling_ffi::{ ddog_prof_Profile_new, ddog_prof_ProfilerClient_drop, ddog_prof_ProfilerManager_enqueue_sample, diff --git a/datadog-profiling-ffi/tests/test_ffi_fork_data_preservation.rs b/datadog-profiling-ffi/tests/test_ffi_fork_data_preservation.rs index e75c7edbef..34b457a7c7 100644 --- a/datadog-profiling-ffi/tests/test_ffi_fork_data_preservation.rs +++ b/datadog-profiling-ffi/tests/test_ffi_fork_data_preservation.rs @@ -1,3 +1,5 @@ +// Copyright 2025-Present Datadog, Inc. https://www.datadoghq.com/ + mod test_utils; use datadog_profiling_ffi::{ ddog_prof_Profile_new, ddog_prof_ProfilerClient_drop, ddog_prof_ProfilerManager_enqueue_sample, diff --git a/datadog-profiling-ffi/tests/test_ffi_fork_lifecycle.rs b/datadog-profiling-ffi/tests/test_ffi_fork_lifecycle.rs index bf2dde52d1..c0acb4ed09 100644 --- a/datadog-profiling-ffi/tests/test_ffi_fork_lifecycle.rs +++ b/datadog-profiling-ffi/tests/test_ffi_fork_lifecycle.rs @@ -1,3 +1,5 @@ +// Copyright 2025-Present Datadog, Inc. https://www.datadoghq.com/ + mod test_utils; use datadog_profiling_ffi::{ ddog_prof_Profile_new, ddog_prof_ProfilerClient_drop, ddog_prof_ProfilerManager_enqueue_sample, diff --git a/datadog-profiling-ffi/tests/test_ffi_fork_parent_child.rs b/datadog-profiling-ffi/tests/test_ffi_fork_parent_child.rs index b6fe219c81..22f5c045ed 100644 --- a/datadog-profiling-ffi/tests/test_ffi_fork_parent_child.rs +++ b/datadog-profiling-ffi/tests/test_ffi_fork_parent_child.rs @@ -1,3 +1,5 @@ +// Copyright 2025-Present Datadog, Inc. https://www.datadoghq.com/ + mod test_utils; use datadog_profiling_ffi::{ ddog_prof_Profile_new, ddog_prof_ProfilerClient_drop, ddog_prof_ProfilerManager_enqueue_sample, diff --git a/datadog-profiling-ffi/tests/test_ffi_fork_parent_only.rs b/datadog-profiling-ffi/tests/test_ffi_fork_parent_only.rs index f0d17e44e9..53e1d62224 100644 --- a/datadog-profiling-ffi/tests/test_ffi_fork_parent_only.rs +++ b/datadog-profiling-ffi/tests/test_ffi_fork_parent_only.rs @@ -1,3 +1,5 @@ +// Copyright 2025-Present Datadog, Inc. https://www.datadoghq.com/ + mod test_utils; use datadog_profiling_ffi::{ ddog_prof_Profile_new, ddog_prof_ProfilerClient_drop, ddog_prof_ProfilerManager_enqueue_sample, diff --git a/datadog-profiling-ffi/tests/test_ffi_lifecycle_basic.rs b/datadog-profiling-ffi/tests/test_ffi_lifecycle_basic.rs index 450aeeab79..073dc0ba60 100644 --- a/datadog-profiling-ffi/tests/test_ffi_lifecycle_basic.rs +++ b/datadog-profiling-ffi/tests/test_ffi_lifecycle_basic.rs @@ -1,3 +1,5 @@ +// Copyright 2025-Present Datadog, Inc. https://www.datadoghq.com/ + mod test_utils; use datadog_profiling_ffi::{ ddog_prof_Profile_new, ddog_prof_ProfilerClient_drop, ddog_prof_ProfilerManager_pause, diff --git a/datadog-profiling-ffi/tests/test_ffi_lifecycle_with_data.rs b/datadog-profiling-ffi/tests/test_ffi_lifecycle_with_data.rs index 22e80a20f3..fc72d5216d 100644 --- a/datadog-profiling-ffi/tests/test_ffi_lifecycle_with_data.rs +++ b/datadog-profiling-ffi/tests/test_ffi_lifecycle_with_data.rs @@ -1,3 +1,5 @@ +// Copyright 2025-Present Datadog, Inc. https://www.datadoghq.com/ + mod test_utils; use datadog_profiling_ffi::{ ddog_prof_Profile_new, ddog_prof_ProfilerClient_drop, ddog_prof_ProfilerManager_enqueue_sample, diff --git a/datadog-profiling-ffi/tests/test_ffi_start_pause_terminate.rs b/datadog-profiling-ffi/tests/test_ffi_start_pause_terminate.rs index 71cd212cbd..dba8fa9941 100644 --- a/datadog-profiling-ffi/tests/test_ffi_start_pause_terminate.rs +++ b/datadog-profiling-ffi/tests/test_ffi_start_pause_terminate.rs @@ -1,3 +1,5 @@ +// Copyright 2025-Present Datadog, Inc. https://www.datadoghq.com/ + mod test_utils; use datadog_profiling_ffi::{ ddog_prof_Profile_new, ddog_prof_ProfilerClient_drop, ddog_prof_ProfilerManager_enqueue_sample, diff --git a/datadog-profiling-ffi/tests/test_ffi_start_terminate.rs b/datadog-profiling-ffi/tests/test_ffi_start_terminate.rs index 4201fa03ac..2c84538789 100644 --- a/datadog-profiling-ffi/tests/test_ffi_start_terminate.rs +++ b/datadog-profiling-ffi/tests/test_ffi_start_terminate.rs @@ -1,3 +1,5 @@ +// Copyright 2025-Present Datadog, Inc. https://www.datadoghq.com/ + mod test_utils; use datadog_profiling_ffi::{ ddog_prof_Profile_new, ddog_prof_ProfilerClient_drop, ddog_prof_ProfilerManager_enqueue_sample, diff --git a/datadog-profiling-ffi/tests/test_utils.rs b/datadog-profiling-ffi/tests/test_utils.rs index b474ec8f76..9524377955 100644 --- a/datadog-profiling-ffi/tests/test_utils.rs +++ b/datadog-profiling-ffi/tests/test_utils.rs @@ -1,3 +1,5 @@ +// Copyright 2025-Present Datadog, Inc. https://www.datadoghq.com/ + #![allow(dead_code)] use std::sync::atomic::{AtomicUsize, Ordering}; From f07a0fcf9bc460fe63af0aa4a359c91d10e23eac Mon Sep 17 00:00:00 2001 From: Daniel Schwartz-Narbonne Date: Tue, 8 Jul 2025 15:30:34 -0400 Subject: [PATCH 51/52] SPDX --- datadog-profiling-ffi/src/manager/client.rs | 1 + datadog-profiling-ffi/src/manager/ffi_utils.rs | 1 + datadog-profiling-ffi/src/manager/profiler_manager.rs | 1 + datadog-profiling-ffi/src/manager/samples.rs | 1 + datadog-profiling-ffi/src/manager/tests.rs | 1 + datadog-profiling-ffi/tests/test_ffi_fork_child_only.rs | 1 + datadog-profiling-ffi/tests/test_ffi_fork_data_preservation.rs | 1 + datadog-profiling-ffi/tests/test_ffi_fork_lifecycle.rs | 1 + datadog-profiling-ffi/tests/test_ffi_fork_parent_child.rs | 1 + datadog-profiling-ffi/tests/test_ffi_fork_parent_only.rs | 1 + datadog-profiling-ffi/tests/test_ffi_lifecycle_basic.rs | 1 + datadog-profiling-ffi/tests/test_ffi_lifecycle_with_data.rs | 1 + datadog-profiling-ffi/tests/test_ffi_start_pause_terminate.rs | 1 + datadog-profiling-ffi/tests/test_ffi_start_terminate.rs | 1 + datadog-profiling-ffi/tests/test_utils.rs | 1 + 15 files changed, 15 insertions(+) diff --git a/datadog-profiling-ffi/src/manager/client.rs b/datadog-profiling-ffi/src/manager/client.rs index 4c5a690bc7..465ac61f7c 100644 --- a/datadog-profiling-ffi/src/manager/client.rs +++ b/datadog-profiling-ffi/src/manager/client.rs @@ -1,4 +1,5 @@ // Copyright 2025-Present Datadog, Inc. https://www.datadoghq.com/ +// SPDX-License-Identifier: Apache-2.0 use std::sync::Arc; use std::{ffi::c_void, sync::atomic::AtomicBool}; diff --git a/datadog-profiling-ffi/src/manager/ffi_utils.rs b/datadog-profiling-ffi/src/manager/ffi_utils.rs index ae08fb2a8c..d07b3c5a53 100644 --- a/datadog-profiling-ffi/src/manager/ffi_utils.rs +++ b/datadog-profiling-ffi/src/manager/ffi_utils.rs @@ -1,4 +1,5 @@ // Copyright 2025-Present Datadog, Inc. https://www.datadoghq.com/ +// SPDX-License-Identifier: Apache-2.0 use std::num::NonZeroI64; diff --git a/datadog-profiling-ffi/src/manager/profiler_manager.rs b/datadog-profiling-ffi/src/manager/profiler_manager.rs index 920649456d..70f90f73d3 100644 --- a/datadog-profiling-ffi/src/manager/profiler_manager.rs +++ b/datadog-profiling-ffi/src/manager/profiler_manager.rs @@ -1,4 +1,5 @@ // Copyright 2025-Present Datadog, Inc. https://www.datadoghq.com/ +// SPDX-License-Identifier: Apache-2.0 #![allow(clippy::unwrap_used)] diff --git a/datadog-profiling-ffi/src/manager/samples.rs b/datadog-profiling-ffi/src/manager/samples.rs index 89fd2a9f72..44003ff33e 100644 --- a/datadog-profiling-ffi/src/manager/samples.rs +++ b/datadog-profiling-ffi/src/manager/samples.rs @@ -1,4 +1,5 @@ // Copyright 2025-Present Datadog, Inc. https://www.datadoghq.com/ +// SPDX-License-Identifier: Apache-2.0 use std::ffi::c_void; diff --git a/datadog-profiling-ffi/src/manager/tests.rs b/datadog-profiling-ffi/src/manager/tests.rs index 233a0b919c..6ff8c75df7 100644 --- a/datadog-profiling-ffi/src/manager/tests.rs +++ b/datadog-profiling-ffi/src/manager/tests.rs @@ -1,4 +1,5 @@ // Copyright 2025-Present Datadog, Inc. https://www.datadoghq.com/ +// SPDX-License-Identifier: Apache-2.0 use std::ffi::c_void; use std::sync::atomic::{AtomicUsize, Ordering}; diff --git a/datadog-profiling-ffi/tests/test_ffi_fork_child_only.rs b/datadog-profiling-ffi/tests/test_ffi_fork_child_only.rs index 695becc1ce..cf35d4218f 100644 --- a/datadog-profiling-ffi/tests/test_ffi_fork_child_only.rs +++ b/datadog-profiling-ffi/tests/test_ffi_fork_child_only.rs @@ -1,4 +1,5 @@ // Copyright 2025-Present Datadog, Inc. https://www.datadoghq.com/ +// SPDX-License-Identifier: Apache-2.0 mod test_utils; use datadog_profiling_ffi::{ diff --git a/datadog-profiling-ffi/tests/test_ffi_fork_data_preservation.rs b/datadog-profiling-ffi/tests/test_ffi_fork_data_preservation.rs index 34b457a7c7..2872647a93 100644 --- a/datadog-profiling-ffi/tests/test_ffi_fork_data_preservation.rs +++ b/datadog-profiling-ffi/tests/test_ffi_fork_data_preservation.rs @@ -1,4 +1,5 @@ // Copyright 2025-Present Datadog, Inc. https://www.datadoghq.com/ +// SPDX-License-Identifier: Apache-2.0 mod test_utils; use datadog_profiling_ffi::{ diff --git a/datadog-profiling-ffi/tests/test_ffi_fork_lifecycle.rs b/datadog-profiling-ffi/tests/test_ffi_fork_lifecycle.rs index c0acb4ed09..b68cbf6f7a 100644 --- a/datadog-profiling-ffi/tests/test_ffi_fork_lifecycle.rs +++ b/datadog-profiling-ffi/tests/test_ffi_fork_lifecycle.rs @@ -1,4 +1,5 @@ // Copyright 2025-Present Datadog, Inc. https://www.datadoghq.com/ +// SPDX-License-Identifier: Apache-2.0 mod test_utils; use datadog_profiling_ffi::{ diff --git a/datadog-profiling-ffi/tests/test_ffi_fork_parent_child.rs b/datadog-profiling-ffi/tests/test_ffi_fork_parent_child.rs index 22f5c045ed..e0be8d45c2 100644 --- a/datadog-profiling-ffi/tests/test_ffi_fork_parent_child.rs +++ b/datadog-profiling-ffi/tests/test_ffi_fork_parent_child.rs @@ -1,4 +1,5 @@ // Copyright 2025-Present Datadog, Inc. https://www.datadoghq.com/ +// SPDX-License-Identifier: Apache-2.0 mod test_utils; use datadog_profiling_ffi::{ diff --git a/datadog-profiling-ffi/tests/test_ffi_fork_parent_only.rs b/datadog-profiling-ffi/tests/test_ffi_fork_parent_only.rs index 53e1d62224..9e71d74847 100644 --- a/datadog-profiling-ffi/tests/test_ffi_fork_parent_only.rs +++ b/datadog-profiling-ffi/tests/test_ffi_fork_parent_only.rs @@ -1,4 +1,5 @@ // Copyright 2025-Present Datadog, Inc. https://www.datadoghq.com/ +// SPDX-License-Identifier: Apache-2.0 mod test_utils; use datadog_profiling_ffi::{ diff --git a/datadog-profiling-ffi/tests/test_ffi_lifecycle_basic.rs b/datadog-profiling-ffi/tests/test_ffi_lifecycle_basic.rs index 073dc0ba60..f1bb7360f9 100644 --- a/datadog-profiling-ffi/tests/test_ffi_lifecycle_basic.rs +++ b/datadog-profiling-ffi/tests/test_ffi_lifecycle_basic.rs @@ -1,4 +1,5 @@ // Copyright 2025-Present Datadog, Inc. https://www.datadoghq.com/ +// SPDX-License-Identifier: Apache-2.0 mod test_utils; use datadog_profiling_ffi::{ diff --git a/datadog-profiling-ffi/tests/test_ffi_lifecycle_with_data.rs b/datadog-profiling-ffi/tests/test_ffi_lifecycle_with_data.rs index fc72d5216d..25017ae353 100644 --- a/datadog-profiling-ffi/tests/test_ffi_lifecycle_with_data.rs +++ b/datadog-profiling-ffi/tests/test_ffi_lifecycle_with_data.rs @@ -1,4 +1,5 @@ // Copyright 2025-Present Datadog, Inc. https://www.datadoghq.com/ +// SPDX-License-Identifier: Apache-2.0 mod test_utils; use datadog_profiling_ffi::{ diff --git a/datadog-profiling-ffi/tests/test_ffi_start_pause_terminate.rs b/datadog-profiling-ffi/tests/test_ffi_start_pause_terminate.rs index dba8fa9941..003dd38404 100644 --- a/datadog-profiling-ffi/tests/test_ffi_start_pause_terminate.rs +++ b/datadog-profiling-ffi/tests/test_ffi_start_pause_terminate.rs @@ -1,4 +1,5 @@ // Copyright 2025-Present Datadog, Inc. https://www.datadoghq.com/ +// SPDX-License-Identifier: Apache-2.0 mod test_utils; use datadog_profiling_ffi::{ diff --git a/datadog-profiling-ffi/tests/test_ffi_start_terminate.rs b/datadog-profiling-ffi/tests/test_ffi_start_terminate.rs index 2c84538789..eed2b73d2e 100644 --- a/datadog-profiling-ffi/tests/test_ffi_start_terminate.rs +++ b/datadog-profiling-ffi/tests/test_ffi_start_terminate.rs @@ -1,4 +1,5 @@ // Copyright 2025-Present Datadog, Inc. https://www.datadoghq.com/ +// SPDX-License-Identifier: Apache-2.0 mod test_utils; use datadog_profiling_ffi::{ diff --git a/datadog-profiling-ffi/tests/test_utils.rs b/datadog-profiling-ffi/tests/test_utils.rs index 9524377955..4a8907016e 100644 --- a/datadog-profiling-ffi/tests/test_utils.rs +++ b/datadog-profiling-ffi/tests/test_utils.rs @@ -1,4 +1,5 @@ // Copyright 2025-Present Datadog, Inc. https://www.datadoghq.com/ +// SPDX-License-Identifier: Apache-2.0 #![allow(dead_code)] From 05e60f6a3f433dd4f6bc9c14ca5cbabf2b7e6982 Mon Sep 17 00:00:00 2001 From: Daniel Schwartz-Narbonne Date: Tue, 8 Jul 2025 15:48:34 -0400 Subject: [PATCH 52/52] unix only tests --- datadog-profiling-ffi/tests/test_ffi_fork_child_only.rs | 1 + datadog-profiling-ffi/tests/test_ffi_fork_data_preservation.rs | 1 + datadog-profiling-ffi/tests/test_ffi_fork_lifecycle.rs | 1 + datadog-profiling-ffi/tests/test_ffi_fork_parent_child.rs | 1 + datadog-profiling-ffi/tests/test_ffi_fork_parent_only.rs | 1 + datadog-profiling-ffi/tests/test_ffi_lifecycle_basic.rs | 1 + datadog-profiling-ffi/tests/test_ffi_lifecycle_with_data.rs | 1 + datadog-profiling-ffi/tests/test_ffi_start_pause_terminate.rs | 1 + datadog-profiling-ffi/tests/test_ffi_start_terminate.rs | 1 + 9 files changed, 9 insertions(+) diff --git a/datadog-profiling-ffi/tests/test_ffi_fork_child_only.rs b/datadog-profiling-ffi/tests/test_ffi_fork_child_only.rs index cf35d4218f..4f27a0cba0 100644 --- a/datadog-profiling-ffi/tests/test_ffi_fork_child_only.rs +++ b/datadog-profiling-ffi/tests/test_ffi_fork_child_only.rs @@ -1,6 +1,7 @@ // Copyright 2025-Present Datadog, Inc. https://www.datadoghq.com/ // SPDX-License-Identifier: Apache-2.0 +#[cfg(unix)] mod test_utils; use datadog_profiling_ffi::{ ddog_prof_Profile_new, ddog_prof_ProfilerClient_drop, ddog_prof_ProfilerManager_enqueue_sample, diff --git a/datadog-profiling-ffi/tests/test_ffi_fork_data_preservation.rs b/datadog-profiling-ffi/tests/test_ffi_fork_data_preservation.rs index 2872647a93..8091ca2e64 100644 --- a/datadog-profiling-ffi/tests/test_ffi_fork_data_preservation.rs +++ b/datadog-profiling-ffi/tests/test_ffi_fork_data_preservation.rs @@ -1,6 +1,7 @@ // Copyright 2025-Present Datadog, Inc. https://www.datadoghq.com/ // SPDX-License-Identifier: Apache-2.0 +#[cfg(unix)] mod test_utils; use datadog_profiling_ffi::{ ddog_prof_Profile_new, ddog_prof_ProfilerClient_drop, ddog_prof_ProfilerManager_enqueue_sample, diff --git a/datadog-profiling-ffi/tests/test_ffi_fork_lifecycle.rs b/datadog-profiling-ffi/tests/test_ffi_fork_lifecycle.rs index b68cbf6f7a..6f4a98d793 100644 --- a/datadog-profiling-ffi/tests/test_ffi_fork_lifecycle.rs +++ b/datadog-profiling-ffi/tests/test_ffi_fork_lifecycle.rs @@ -1,6 +1,7 @@ // Copyright 2025-Present Datadog, Inc. https://www.datadoghq.com/ // SPDX-License-Identifier: Apache-2.0 +#[cfg(unix)] mod test_utils; use datadog_profiling_ffi::{ ddog_prof_Profile_new, ddog_prof_ProfilerClient_drop, ddog_prof_ProfilerManager_enqueue_sample, diff --git a/datadog-profiling-ffi/tests/test_ffi_fork_parent_child.rs b/datadog-profiling-ffi/tests/test_ffi_fork_parent_child.rs index e0be8d45c2..2ecf38c6ac 100644 --- a/datadog-profiling-ffi/tests/test_ffi_fork_parent_child.rs +++ b/datadog-profiling-ffi/tests/test_ffi_fork_parent_child.rs @@ -1,6 +1,7 @@ // Copyright 2025-Present Datadog, Inc. https://www.datadoghq.com/ // SPDX-License-Identifier: Apache-2.0 +#[cfg(unix)] mod test_utils; use datadog_profiling_ffi::{ ddog_prof_Profile_new, ddog_prof_ProfilerClient_drop, ddog_prof_ProfilerManager_enqueue_sample, diff --git a/datadog-profiling-ffi/tests/test_ffi_fork_parent_only.rs b/datadog-profiling-ffi/tests/test_ffi_fork_parent_only.rs index 9e71d74847..5f5ec2936d 100644 --- a/datadog-profiling-ffi/tests/test_ffi_fork_parent_only.rs +++ b/datadog-profiling-ffi/tests/test_ffi_fork_parent_only.rs @@ -1,6 +1,7 @@ // Copyright 2025-Present Datadog, Inc. https://www.datadoghq.com/ // SPDX-License-Identifier: Apache-2.0 +#[cfg(unix)] mod test_utils; use datadog_profiling_ffi::{ ddog_prof_Profile_new, ddog_prof_ProfilerClient_drop, ddog_prof_ProfilerManager_enqueue_sample, diff --git a/datadog-profiling-ffi/tests/test_ffi_lifecycle_basic.rs b/datadog-profiling-ffi/tests/test_ffi_lifecycle_basic.rs index f1bb7360f9..34ba000843 100644 --- a/datadog-profiling-ffi/tests/test_ffi_lifecycle_basic.rs +++ b/datadog-profiling-ffi/tests/test_ffi_lifecycle_basic.rs @@ -1,6 +1,7 @@ // Copyright 2025-Present Datadog, Inc. https://www.datadoghq.com/ // SPDX-License-Identifier: Apache-2.0 +#[cfg(unix)] mod test_utils; use datadog_profiling_ffi::{ ddog_prof_Profile_new, ddog_prof_ProfilerClient_drop, ddog_prof_ProfilerManager_pause, diff --git a/datadog-profiling-ffi/tests/test_ffi_lifecycle_with_data.rs b/datadog-profiling-ffi/tests/test_ffi_lifecycle_with_data.rs index 25017ae353..1fa82e7c29 100644 --- a/datadog-profiling-ffi/tests/test_ffi_lifecycle_with_data.rs +++ b/datadog-profiling-ffi/tests/test_ffi_lifecycle_with_data.rs @@ -1,6 +1,7 @@ // Copyright 2025-Present Datadog, Inc. https://www.datadoghq.com/ // SPDX-License-Identifier: Apache-2.0 +#[cfg(unix)] mod test_utils; use datadog_profiling_ffi::{ ddog_prof_Profile_new, ddog_prof_ProfilerClient_drop, ddog_prof_ProfilerManager_enqueue_sample, diff --git a/datadog-profiling-ffi/tests/test_ffi_start_pause_terminate.rs b/datadog-profiling-ffi/tests/test_ffi_start_pause_terminate.rs index 003dd38404..9c8d45cb80 100644 --- a/datadog-profiling-ffi/tests/test_ffi_start_pause_terminate.rs +++ b/datadog-profiling-ffi/tests/test_ffi_start_pause_terminate.rs @@ -1,6 +1,7 @@ // Copyright 2025-Present Datadog, Inc. https://www.datadoghq.com/ // SPDX-License-Identifier: Apache-2.0 +#[cfg(unix)] mod test_utils; use datadog_profiling_ffi::{ ddog_prof_Profile_new, ddog_prof_ProfilerClient_drop, ddog_prof_ProfilerManager_enqueue_sample, diff --git a/datadog-profiling-ffi/tests/test_ffi_start_terminate.rs b/datadog-profiling-ffi/tests/test_ffi_start_terminate.rs index eed2b73d2e..0827346458 100644 --- a/datadog-profiling-ffi/tests/test_ffi_start_terminate.rs +++ b/datadog-profiling-ffi/tests/test_ffi_start_terminate.rs @@ -1,6 +1,7 @@ // Copyright 2025-Present Datadog, Inc. https://www.datadoghq.com/ // SPDX-License-Identifier: Apache-2.0 +#[cfg(unix)] mod test_utils; use datadog_profiling_ffi::{ ddog_prof_Profile_new, ddog_prof_ProfilerClient_drop, ddog_prof_ProfilerManager_enqueue_sample,