feat: upload compression config

Author: morrisonlevi

datadog-profiling-ffi/src/lib.rs

@@ -14,6 +14,8 @@ mod exporter;
 mod profiles;
 mod string_storage;
 
+pub use profiles::*;
+
 // re-export crashtracker ffi
 #[cfg(feature = "crashtracker-ffi")]
 pub use datadog_crashtracker_ffi::*;

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

@@ -70,6 +70,15 @@ impl From<anyhow::Result<()>> for ProfileResult {
     }
 }
 
+impl From<ProfileResult> for Result<(), Error> {
+    fn from(result: ProfileResult) -> Self {
+        match result {
+            ProfileResult::Ok(_) => Ok(()),
+            ProfileResult::Err(err) => Err(err),
+        }
+    }
+}
+
 /// Returned by [ddog_prof_Profile_new].
 #[allow(dead_code)]
 #[repr(C)]
@@ -79,6 +88,15 @@ pub enum ProfileNewResult {
     Err(Error),
 }
 
+impl From<ProfileNewResult> for Result<Profile, Error> {
+    fn from(value: ProfileNewResult) -> Self {
+        match value {
+            ProfileNewResult::Ok(v) => Ok(v),
+            ProfileNewResult::Err(e) => Err(e),
+        }
+    }
+}
+
 #[allow(dead_code)]
 #[repr(C)]
 pub enum SerializeResult {
@@ -390,7 +408,7 @@ impl<'a> From<Sample<'a>> for api::StringIdSample<'a> {
 ///   time.
 ///
 /// # Safety
-/// All slices must be have pointers that are suitably aligned for their type
+/// All slices must have pointers that are suitably aligned for their type
 /// and must have the correct number of elements for the slice.
 #[no_mangle]
 #[must_use]
@@ -401,7 +419,11 @@ pub unsafe extern "C" fn ddog_prof_Profile_new(
     profile_new(sample_types, period, None)
 }
 
-/// Same as `ddog_profile_new` but also configures a `string_storage` for the profile.
+/// Same as [ddog_prof_Profile_new] but also configures a `string_storage` for
+/// the profile.
+///
+/// # Safety
+/// Has all same safety conditions as [ddog_prof_Profile_new].
 #[no_mangle]
 #[must_use]
 /// TODO: @ivoanjo Should this take a `*mut ManagedStringStorage` like Profile APIs do?
@@ -447,26 +469,6 @@ pub unsafe extern "C" fn ddog_prof_Profile_drop(profile: *mut Profile) {
     }
 }
 
-#[cfg(test)]
-impl From<ProfileResult> for Result<(), Error> {
-    fn from(result: ProfileResult) -> Self {
-        match result {
-            ProfileResult::Ok(_) => Ok(()),
-            ProfileResult::Err(err) => Err(err),
-        }
-    }
-}
-
-#[cfg(test)]
-impl From<ProfileNewResult> for Result<Profile, Error> {
-    fn from(result: ProfileNewResult) -> Self {
-        match result {
-            ProfileNewResult::Ok(p) => Ok(p),
-            ProfileNewResult::Err(err) => Err(err),
-        }
-    }
-}
-
 /// # Safety
 /// The `profile` ptr must point to a valid Profile object created by this
 /// module. All pointers inside the `sample` need to be valid for the duration
@@ -547,6 +549,58 @@ pub unsafe extern "C" fn ddog_prof_Profile_set_endpoint(
     .into()
 }
 
+/// Sets the profile's upload compression algorithm. Useful only for testing.
+///
+/// # Errors
+///
+/// Returns an error if either the pointer or the inner profiler ptr is null.
+///
+/// # Safety
+///
+/// The `profile` ptr must point to a valid Profile object.
+///
+/// # Examples
+///
+/// ```
+/// # use datadog_profiling_ffi::{Error, ddog_prof_Profile_set_upload_compression, ddog_prof_Profile_new, ValueType, Slice, CharSlice};
+/// # use std::ptr::addr_of_mut;
+/// # fn main() -> Result<(), Error> { unsafe {
+/// use datadog_profiling::UploadCompression;
+///
+/// let mut profile = Result::from(ddog_prof_Profile_new(
+///     Slice::from([ValueType {
+///         type_: CharSlice::from("sample"),
+///         unit: CharSlice::from("count"),
+///     }].as_slice()),
+///     None,
+/// ))?;
+///
+/// // Set compression off (only for testing).
+/// Result::from(ddog_prof_Profile_set_upload_compression(
+///     addr_of_mut!(profile),
+///     UploadCompression::Off
+/// ))?;
+/// # } Ok(()) }
+/// ```
+#[no_mangle]
+#[must_use]
+#[named]
+pub unsafe extern "C" fn ddog_prof_Profile_set_upload_compression(
+    profile: *mut Profile,
+    upload_compression: datadog_profiling::UploadCompression,
+) -> ProfileResult {
+    match profile_ptr_to_inner(profile) {
+        Ok(profile) => {
+            profile.set_upload_compression(upload_compression);
+            ProfileResult::Ok(true)
+        }
+        Err(err) => {
+            let e = err.context(concat!(function_name!(), " failed"));
+            ProfileResult::Err(e.into())
+        }
+    }
+}
+
 /// Count the number of times an endpoint has been seen.
 ///
 /// # Arguments
@@ -759,6 +813,14 @@ pub unsafe extern "C" fn ddog_prof_Profile_serialize(
     .into()
 }
 
+/// Borrows FFI Vec as an FFI Slice.
+///
+/// # Safety
+/// The input needs to be a valid reference to an FFI Vec, and the slice needs
+/// to be used in a way consistent with the lifetime and safety rules. Some
+/// things to avoid:
+///  - Do not modify the Vec at all while the Slice is alive.
+///  - Do not drop the Vec while the Slice is alive.
 #[must_use]
 #[no_mangle]
 pub unsafe extern "C" fn ddog_Vec_U8_as_slice(vec: &ddcommon_ffi::Vec<u8>) -> Slice<u8> {

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

@@ -407,14 +407,9 @@ pub unsafe extern "C" fn ddog_prof_Profile_get_generation(
 }
 
 /// This functions returns whether the given generations are equal.
-///
-/// # Safety: No safety requirements
 #[must_use]
 #[no_mangle]
-pub unsafe extern "C" fn ddog_prof_Profile_generations_are_equal(
-    a: Generation,
-    b: Generation,
-) -> bool {
+pub extern "C" fn ddog_prof_Profile_generations_are_equal(a: Generation, b: Generation) -> bool {
     a == b
 }
 

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

@@ -3,3 +3,6 @@
 
 mod datatypes;
 mod interning_api;
+
+pub use datatypes::*;
+pub use interning_api::*;

datadog-profiling/src/internal/profile/mod.rs

@@ -6,18 +6,16 @@ mod fuzz_tests;
 
 pub mod interning_api;
 
-use self::api::UpscalingInfo;
 use super::*;
-use crate::api;
-use crate::api::ManagedStringId;
+use crate::api::{self, ManagedStringId, UpscalingInfo};
 use crate::collections::identifiable::*;
 use crate::collections::string_storage::{CachedProfileId, ManagedStringStorage};
 use crate::collections::string_table::StringTable;
 use crate::iter::{IntoLendingIterator, LendingIterator};
+use crate::UploadCompression;
 use anyhow::Context;
 use datadog_profiling_protobuf::{self as protobuf, Field, NO_OPT_ZERO, OPT_ZERO};
 use interning_api::Generation;
-use lz4_flex::frame::FrameEncoder;
 use std::borrow::Cow;
 use std::collections::HashMap;
 use std::io;
@@ -51,6 +49,7 @@ pub struct Profile {
     string_storage: Option<Arc<Mutex<ManagedStringStorage>>>,
     string_storage_cached_profile_id: Option<CachedProfileId>,
     timestamp_key: StringId,
+    upload_compression: UploadCompression,
     upscaling_rules: UpscalingRules,
 }
 
@@ -313,6 +312,13 @@ impl Profile {
         Ok(profile)
     }
 
+    /// Sets the upload compression algorithm. The default level is
+    /// [UploadCompression::On]. The exact algorithm that on uses may change
+    /// over time.
+    pub fn set_upload_compression(&mut self, compression: UploadCompression) {
+        self.upload_compression = compression;
+    }
+
     /// Serialize the aggregated profile, adding the end time and duration.
     /// # Arguments
     /// * `end_time` - Optional end time of the profile. Passing None will use the current time.
@@ -336,11 +342,25 @@ impl Profile {
         // size of 32KiB should definitely outperform starting at zero for
         // time consumed, allocator pressure, and allocator fragmentation.
         const INITIAL_PPROF_BUFFER_SIZE: usize = 32 * 1024;
-        let mut compressor = FrameEncoder::new(Vec::with_capacity(INITIAL_PPROF_BUFFER_SIZE));
-
-        let mut encoded_profile = self.encode(&mut compressor, end_time, duration)?;
-        encoded_profile.buffer = compressor.finish()?;
-        Ok(encoded_profile)
+        let mut buffer = Vec::with_capacity(INITIAL_PPROF_BUFFER_SIZE);
+
+        Ok(match self.upload_compression {
+            UploadCompression::Off => {
+                let encoded_profile = self.encode(&mut buffer, end_time, duration)?;
+                EncodedProfile {
+                    buffer,
+                    ..encoded_profile
+                }
+            }
+            UploadCompression::On | UploadCompression::Lz4 => {
+                let mut compressor = lz4_flex::frame::FrameEncoder::new(buffer);
+                let encoded_profile = self.encode(&mut compressor, end_time, duration)?;
+                EncodedProfile {
+                    buffer: compressor.finish()?,
+                    ..encoded_profile
+                }
+            }
+        })
     }
 
     /// Encodes the profile. Note that the buffer will be empty. The caller
@@ -735,6 +755,7 @@ impl Profile {
             string_storage_cached_profile_id: None, /* Never reuse an id! See comments on
                                                      * CachedProfileId for why. */
             timestamp_key: Default::default(),
+            upload_compression: Default::default(),
             upscaling_rules: Default::default(),
         };
 

datadog-profiling/src/lib.rs

@@ -12,3 +12,13 @@ pub mod exporter;
 pub mod internal;
 pub mod iter;
 pub mod pprof;
+
+#[repr(C)]
+#[derive(Copy, Clone, Debug, Default)]
+pub enum UploadCompression {
+    Off,
+    /// On is the default. The compression algorithm used can change over time.
+    #[default]
+    On,
+    Lz4,
+}

ddcommon-ffi/src/result.rs

@@ -46,6 +46,15 @@ pub enum Result<T> {
     Err(Error),
 }
 
+impl From<VoidResult> for std::result::Result<(), Error> {
+    fn from(result: VoidResult) -> Self {
+        match result {
+            VoidResult::Ok(_) => Ok(()),
+            VoidResult::Err(err) => Err(err),
+        }
+    }
+}
+
 impl<T> Result<T> {
     pub fn unwrap(self) -> T {
         match self {