[profiling] Rust driver for python profiler

datadog-profiling-ffi/Cargo.toml

@@ -35,6 +35,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 }
@@ -50,4 +51,10 @@ 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"
+
+[dev-dependencies]
+lz4_flex = "0.11.3"
+prost = "0.13.5"
+datadog-profiling-protobuf = { path = "../datadog-profiling-protobuf" }

datadog-profiling-ffi/src/lib.rs

@@ -11,7 +11,8 @@
 pub use symbolizer_ffi::*;
 
 mod exporter;
-mod profiles;
+mod manager;
+pub mod profiles;
 mod string_storage;
 
 // re-export crashtracker ffi
@@ -37,3 +38,17 @@ pub use datadog_log_ffi::*;
 // re-export tracer metadata functions
 #[cfg(feature = "ddcommon-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,
+};

datadog-profiling-ffi/src/manager/client.rs

@@ -0,0 +1,49 @@
+// 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};
+
+use crossbeam_channel::SendError;
+
+use super::ClientSampleChannels;
+use super::SendSample;
+
+#[derive(Debug, Clone)]
+pub struct ManagedProfilerClient {
+    channels: ClientSampleChannels,
+    is_shutdown: Arc<AtomicBool>,
+}
+
+impl ManagedProfilerClient {
+    pub fn new(channels: ClientSampleChannels, is_shutdown: Arc<AtomicBool>) -> Self {
+        Self {
+            channels,
+            is_shutdown,
+        }
+    }
+
+    /// # 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<SendSample>> {
+        if self.is_shutdown.load(std::sync::atomic::Ordering::SeqCst) {
+            return Err(SendError(unsafe { SendSample::new(sample) }));
+        }
+        self.channels.send_sample(sample)
+    }
+
+    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()
+    }
+}

datadog-profiling-ffi/src/manager/ffi_api.rs

@@ -0,0 +1,193 @@
+// Copyright 2025-Present Datadog, Inc. https://www.datadoghq.com/
+// SPDX-License-Identifier: Apache-2.0
+
+use crate::manager::{
+    profiler_manager::{
+        ManagedSampleCallbacks, ManagerCallbacks, ProfilerManager, ProfilerManagerConfig,
+    },
+    ManagedProfilerClient,
+};
+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, ToInner,
+    VoidResult,
+};
+use function_name::named;
+use std::ffi::c_void;
+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.
+/// - 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(
+    profile: *mut Profile,
+    cpu_sampler_callback: extern "C" fn(*mut internal::Profile),
+    upload_callback: extern "C" fn(*mut Handle<internal::Profile>, &mut Option<CancellationToken>),
+    sample_callbacks: ManagedSampleCallbacks,
+    config: ProfilerManagerConfig,
+) -> FFIResult<Handle<ManagedProfilerClient>> {
+    wrap_with_ffi_result!({
+        let internal_profile = *profile.take()?;
+        let callbacks = ManagerCallbacks {
+            cpu_sampler_callback,
+            upload_callback,
+            sample_callbacks,
+        };
+        let client = ProfilerManager::start(internal_profile, callbacks, config)?;
+        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(
+    mut handle: *mut Handle<ManagedProfilerClient>,
+    sample_ptr: *mut c_void,
+) -> VoidResult {
+    wrap_with_void_ffi_result!({
+        handle
+            .to_inner_mut()?
+            .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(
+    mut handle: *mut Handle<ManagedProfilerClient>,
+) -> FFIResult<*mut c_void> {
+    wrap_with_ffi_result!({
+        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")),
+        }
+    })
+}
+
+/// Pauses the global profiler manager, shutting down the current instance and storing the profile.
+/// The manager can be restarted later using restart functions.
+///
+/// # Safety
+/// - 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_pause() -> VoidResult {
+    wrap_with_void_ffi_result!({
+        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<Handle<ManagedProfilerClient>> {
+    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<Handle<ManagedProfilerClient>> {
+    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<Handle<internal::Profile>> {
+    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]
+#[named]
+pub unsafe extern "C" fn ddog_prof_ProfilerClient_drop(
+    mut handle: *mut Handle<ManagedProfilerClient>,
+) -> VoidResult {
+    wrap_with_void_ffi_result!({
+        handle
+            .take()
+            .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))?;
+    })
+}

datadog-profiling-ffi/src/manager/ffi_utils.rs

@@ -0,0 +1,38 @@
+// Copyright 2025-Present Datadog, Inc. https://www.datadoghq.com/
+// SPDX-License-Identifier: Apache-2.0
+
+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<NonZeroI64>,
+) -> 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()
+}