fix: upscaling based on Christophe's review

Author: morrisonlevi

datadog-profiling-ffi/src/profiles/pprof_builder/mod.rs

@@ -3,14 +3,15 @@
 
 mod upscaling;
 
-use function_name::named;
-pub use upscaling::UpscalingRule;
+pub use upscaling::ProportionalUpscalingRule;
 
 use crate::profile_handle::ProfileHandle;
-use crate::profiles::{ensure_non_null_out_parameter, Utf8Option};
+use crate::profiles::pprof_builder::upscaling::PoissonUpscalingRule;
+use crate::profiles::{
+    ensure_non_null_out_parameter, Utf8ConversionError, Utf8Option,
+};
 use crate::{ArcHandle, ProfileStatus};
 use datadog_profiling::exporter::EncodedProfile;
-use datadog_profiling::profiles;
 use datadog_profiling::profiles::datatypes::{
     Profile, ProfilesDictionary, ScratchPad,
 };
@@ -48,49 +49,127 @@ pub unsafe extern "C" fn ddog_prof_PprofBuilder_new(
     }())
 }
 
-/// Adds a profile to the builder with the attached upscaling rules.
+/// Adds a profile to the builder without upscaling rules.
+///
+/// # Safety
+///
+/// - `handle` must refer to a live builder, and no other mutable
+///   references to that builder may be active for the duration of the call.
+/// - `profile` must be non-null and point to a valid `Profile` that
+///   remains alive until the pprof builder is done.
+/// TODO: finish safety
+#[must_use]
+#[no_mangle]
+pub unsafe extern "C" fn ddog_prof_PprofBuilder_add_profile(
+    mut handle: ProfileHandle<PprofBuilder>,
+    profile: *const Profile,
+) -> ProfileStatus {
+    crate::profiles::ensure_non_null_insert!(profile);
+    let result = || -> Result<(), ProfileStatus> {
+        let builder = unsafe {
+            handle
+                .as_inner_mut()
+                .map_err(|e| ProfileStatus::from_ffi_safe_error_message(e))?
+        };
+        let prof_ref = unsafe { &*profile };
+        builder
+            .try_add_profile(prof_ref)
+            .map_err(ProfileStatus::from_ffi_safe_error_message)
+    }();
+    match result {
+        Ok(_) => ProfileStatus::OK,
+        Err(err) => err,
+    }
+}
+
+/// Adds a profile to the builder with the attached poisson upscaling rule.
+///
+/// # Safety
+///
+/// - `handle` must refer to a live builder, and no other mutable
+///   references to that builder may be active for the duration of the call.
+/// - `profile` must be non-null and point to a valid `Profile` that
+///   remains alive until the pprof builder is done.
+/// TODO: finish safety
+#[must_use]
+#[no_mangle]
+pub unsafe extern "C" fn ddog_prof_PprofBuilder_add_profile_with_poisson_upscaling(
+    mut handle: ProfileHandle<PprofBuilder>,
+    profile: *const Profile,
+    upscaling_rule: PoissonUpscalingRule,
+) -> ProfileStatus {
+    crate::profiles::ensure_non_null_insert!(profile);
+    let result = || -> Result<(), ProfileStatus> {
+        let builder = unsafe {
+            handle
+                .as_inner_mut()
+                .map_err(|e| ProfileStatus::from_ffi_safe_error_message(e))?
+        };
+        let prof_ref = unsafe { &*profile };
+
+        let upscaling_rule = upscaling_rule
+            .try_into()
+            .map_err(ProfileStatus::from_ffi_safe_error_message)?;
+        builder
+            .try_add_profile_with_poisson_upscaling(prof_ref, upscaling_rule)
+            .map_err(ProfileStatus::from_ffi_safe_error_message)
+    }();
+    match result {
+        Ok(_) => ProfileStatus::OK,
+        Err(status) => status,
+    }
+}
+
+/// Adds a profile to the builder with the attached proportional rule.
 ///
 /// # Safety
 ///
 /// - `handle` must refer to a live builder, and no other mutable
 ///   references to that builder may be active for the duration of the call.
 /// - `profile` must be non-null and point to a valid `Profile` that
-///   remains alive for the duration of the call.
+///   remains alive until the pprof builder is done.
+/// TODO: finish safety
 #[must_use]
-#[named]
 #[no_mangle]
-pub unsafe extern "C" fn ddog_prof_PprofBuilder_add_profile<'a>(
+pub unsafe extern "C" fn ddog_prof_PprofBuilder_add_profile_with_proportional_upscaling<
+    'a,
+>(
     mut handle: ProfileHandle<PprofBuilder<'a>>,
     profile: *const Profile,
-    upscaling_rules: Slice<UpscalingRule<'a>>,
+    upscaling_rules: Slice<ProportionalUpscalingRule<'a>>,
     utf8_option: Utf8Option,
 ) -> ProfileStatus {
     crate::profiles::ensure_non_null_insert!(profile);
-    ProfileStatus::from(|| -> Result<(), ProfileError> {
-        let builder = unsafe { handle.as_inner_mut()? };
+    let result = || -> Result<(), ProfileStatus> {
+        let builder = unsafe { handle.as_inner_mut() }
+            .map_err(ProfileStatus::from_ffi_safe_error_message)?;
         let prof_ref = unsafe { &*profile };
 
-        let upscaling_rules = upscaling_rules.try_as_slice().map_err(|e| {
-            ProfileError::fmt(format_args!(
-                "{} failed: upscaling rules couldn't be converted to a Rust slice: {e}",
-                function_name!(),
-            ))
-        })?;
+        let upscaling_rules = upscaling_rules
+            .try_as_slice()
+            .map_err(ProfileStatus::from_ffi_safe_error_message)?;
 
-        builder.try_add_profile(
-            prof_ref,
-            upscaling_rules.iter().map(|rule| {
-                let key = utf8_option
-                    .try_as_bytes_convert(rule.group_by_label_key)?;
-                let value = utf8_option
-                    .try_as_bytes_convert(rule.group_by_label_value)?;
-                let group_by_label = (key, value);
-                let upscaling_info =
-                    profiles::UpscalingInfo::try_from(rule.upscaling_info)?;
-                Ok(profiles::UpscalingRule { group_by_label, upscaling_info })
-            }),
-        )
-    }())
+        builder
+            .try_add_profile_with_proportional_upscaling(
+                prof_ref,
+                upscaling_rules.iter().map(
+                    |rule| -> Result<_, Utf8ConversionError> {
+                        let key = rule.group_by_label.key;
+                        let value = utf8_option
+                            .try_as_bytes_convert(rule.group_by_label.value)?;
+                        Ok((
+                            (key, value),
+                            rule.sampled as f64 / rule.real as f64,
+                        ))
+                    },
+                ),
+            )
+            .map_err(ProfileStatus::from_ffi_safe_error_message)
+    }();
+    match result {
+        Ok(_) => ProfileStatus::OK,
+        Err(status) => status,
+    }
 }
 
 /// Builds and returns a compressed `EncodedProfile` via `out_profile`.

datadog-profiling-ffi/src/profiles/pprof_builder/upscaling.rs

@@ -1,73 +1,87 @@
 // Copyright 2023-Present Datadog, Inc. https://www.datadoghq.com/
 // SPDX-License-Identifier: Apache-2.0
 
-use datadog_profiling::profiles;
-use datadog_profiling::profiles::ProfileError;
+use datadog_profiling::profiles::collections::StringId;
+use datadog_profiling::profiles::datatypes::MAX_SAMPLE_TYPES;
+use ddcommon::error::FfiSafeErrorMessage;
 use ddcommon_ffi::slice::CharSlice;
+use std::ffi::CStr;
+use std::num::NonZeroU64;
 
 #[repr(C)]
 #[derive(Clone, Copy, Debug)]
-pub struct UpscalingProportional {
-    scale: f64,
+pub struct GroupByLabel<'a> {
+    pub key: StringId,
+    pub value: CharSlice<'a>,
 }
 
 #[repr(C)]
 #[derive(Clone, Copy, Debug)]
-pub struct UpscalingPoisson {
-    sampling_distance: u64,
+pub struct ProportionalUpscalingRule<'a> {
+    /// The labels to group the sample values by. If it should apply to all
+    /// samples and not group by label, then use the empty StringId and empty
+    /// CharSlice.
+    pub group_by_label: GroupByLabel<'a>,
+    pub sampled: u64,
+    pub real: u64,
 }
 
 #[repr(C)]
 #[derive(Clone, Copy, Debug)]
-pub struct UpscalingPoissonNonSampleTypeCount {
-    count_value: u64,
+pub struct PoissonUpscalingRule {
+    /// Which offset in the profile's sample is the sum. Must be disjoint from
+    /// `count_offset`.
+    sum_offset: u32,
+    /// Which offset in the profile's sample is the count. Must be disjoint
+    /// from `sum_offset`.
+    count_offset: u32,
     sampling_distance: u64,
 }
 
-/// cbindgen:prefix-with-name=false
-/// cbindgen:rename-all=ScreamingSnakeCase
-#[repr(C)]
-#[derive(Clone, Copy, Debug)]
-#[allow(clippy::enum_variant_names)] // doing this for FFI purposes
-pub enum UpscalingInfo {
-    DdogPprofUpscalingProportional(UpscalingProportional),
-    DdogPprofUpscalingPoisson(UpscalingPoisson),
-    DdogPprofUpscalingPoissonNonSampleTypeCount(
-        UpscalingPoissonNonSampleTypeCount,
-    ),
+#[derive(Debug)]
+pub enum PoissonUpscalingConversionError {
+    SamplingDistance,
+    SumOffset,
+    CountOffset,
 }
 
-impl TryFrom<UpscalingInfo> for profiles::UpscalingInfo {
-    type Error = ProfileError;
+// SAFETY: all cases use Rust c-str literals.
+unsafe impl FfiSafeErrorMessage for PoissonUpscalingConversionError {
+    fn as_ffi_str(&self) -> &'static CStr {
+        match self {
+            PoissonUpscalingConversionError::SamplingDistance => c"PoissonUpscalingRule.sampling_distance cannot be zero",
+            PoissonUpscalingConversionError::SumOffset => c"PoissonUpscalingRule.sum_offset must be less than MAX_SAMPLE_TYPES",
+            PoissonUpscalingConversionError::CountOffset => c"PoissonUpscalingRule.count_offset must be less than MAX_SAMPLE_TYPES",
+        }
+    }
+}
 
-    fn try_from(info: UpscalingInfo) -> Result<Self, Self::Error> {
-        Ok(match info {
-            UpscalingInfo::DdogPprofUpscalingProportional(UpscalingProportional { scale }) => {
-                profiles::UpscalingInfo::Proportional { scale }
-            }
-            UpscalingInfo::DdogPprofUpscalingPoisson(UpscalingPoisson { sampling_distance }) => {
-                profiles::UpscalingInfo::Poisson {
-                    sampling_distance: sampling_distance
-                        .try_into()
-                        .map_err(|_| ProfileError::other("invalid input: upscaling Poisson sampling distance was zero"))?,
-                }
-            }
-            UpscalingInfo::DdogPprofUpscalingPoissonNonSampleTypeCount(UpscalingPoissonNonSampleTypeCount{
-                count_value,
-                sampling_distance,
-            }) => profiles::UpscalingInfo::PoissonNonSampleTypeCount {
-                count_value: count_value.try_into()
-                    .map_err(|_| ProfileError::other("invalid input: upscaling PoissonNonSampleTypeCount count value was zero"))?,
-                sampling_distance: sampling_distance.try_into()
-                    .map_err(|_| ProfileError::other("invalid input: upscaling PoissonNonSampleTypeCount sampling distance was zero"))?,
-            },
-        })
+impl core::fmt::Display for PoissonUpscalingConversionError {
+    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
+        self.as_rust_str().fmt(f)
     }
 }
 
-#[repr(C)]
-pub struct UpscalingRule<'a> {
-    pub group_by_label_key: CharSlice<'a>, // todo: this one is interned already
-    pub group_by_label_value: CharSlice<'a>, // "OutOfBoundsException", "0 - 9ms"
-    pub upscaling_info: UpscalingInfo,
+impl core::error::Error for PoissonUpscalingConversionError {}
+
+impl TryFrom<PoissonUpscalingRule>
+    for datadog_profiling::profiles::PoissonUpscalingRule
+{
+    type Error = PoissonUpscalingConversionError;
+
+    fn try_from(value: PoissonUpscalingRule) -> Result<Self, Self::Error> {
+        let Some(sampling_distance) = NonZeroU64::new(value.sampling_distance)
+        else {
+            return Err(PoissonUpscalingConversionError::SamplingDistance);
+        };
+        let sum_offset = value.count_offset as usize;
+        let count_offset = value.count_offset as usize;
+        if sum_offset >= MAX_SAMPLE_TYPES {
+            return Err(PoissonUpscalingConversionError::SumOffset);
+        }
+        if count_offset >= MAX_SAMPLE_TYPES {
+            return Err(PoissonUpscalingConversionError::CountOffset);
+        }
+        Ok(Self { sum_offset, count_offset, sampling_distance })
+    }
 }

datadog-profiling-ffi/src/profiles/scratchpad.rs

@@ -336,7 +336,7 @@ pub unsafe extern "C" fn ddog_prof_ScratchPad_get_trace_endpoint_str(
 ) -> ProfileStatus {
     ensure_non_null_out_parameter!(result);
     let Ok(pad) = handle.as_inner() else {
-        return ProfileStatus::from_thin_error(EmptyHandleError);
+        return ProfileStatus::from_ffi_safe_error_message(EmptyHandleError);
     };
     if let Some(s) =
         pad.endpoint_tracker().get_trace_endpoint_str(local_root_span_id)