feat(profiling)!: Add upload compression config

For now, the options are Off, On, and Lz4:
 - Off: no compression.
 - On: default, and which algorithm it chooses can change.
 - Lz4: lz4 frame encoded. Currently what On uses.

This is being done in preparation for having Zstd as an option, and
possibly the new default.

Author: morrisonlevi

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

@@ -12,6 +12,7 @@ use function_name::named;
 use std::num::NonZeroI64;
 use std::str::Utf8Error;
 use std::time::SystemTime;
+use datadog_profiling::serializer::UploadCompression;
 
 /// Represents a profile. Do not access its member for any reason, only use
 /// the C API functions on this struct.
@@ -743,6 +744,7 @@ pub unsafe extern "C" fn ddog_prof_Profile_serialize(
     profile: *mut Profile,
     start_time: Option<&Timespec>,
     end_time: Option<&Timespec>,
+    compressor: UploadCompression,
 ) -> SerializeResult {
     (|| {
         let profile = profile_ptr_to_inner(profile)?;
@@ -753,7 +755,7 @@ pub unsafe extern "C" fn ddog_prof_Profile_serialize(
         }
 
         let end_time = end_time.map(SystemTime::from);
-        old_profile.serialize_into_compressed_pprof(end_time, None)
+        old_profile.serialize_into_compressed_pprof(end_time, None, compressor)
     })()
     .context("ddog_prof_Profile_serialize failed")
     .into()

datadog-profiling-replayer/src/main.rs

@@ -225,8 +225,11 @@ fn main() -> anyhow::Result<()> {
 
     if let Some(file) = output {
         println!("Writing out pprof to file {file}");
-        let encoded = outprof
-            .serialize_into_compressed_pprof(Some(replayer.start_time), Some(replayer.duration))?;
+        let encoded = outprof.serialize_into_compressed_pprof(
+            Some(replayer.start_time),
+            Some(replayer.duration),
+            Default::default(),
+        )?;
         if let Some(s) = &mut sysinfo {
             s.measure_memory("After serializing");
         }

datadog-profiling/examples/profiles.rs

@@ -54,7 +54,7 @@ fn main() {
         Err(_) => exit(1),
     }
 
-    match profile.serialize_into_compressed_pprof(None, None) {
+    match profile.serialize_into_compressed_pprof(None, None, Default::default()) {
         Ok(encoded_profile) => {
             let buffer = &encoded_profile.buffer;
             assert!(buffer.len() > 100);

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

@@ -15,7 +15,7 @@ use crate::collections::string_storage::{CachedProfileId, ManagedStringStorage};
 use crate::collections::string_table::StringTable;
 use crate::iter::{IntoLendingIterator, LendingIterator};
 use crate::pprof::sliced_proto::*;
-use crate::serializer::CompressedProtobufSerializer;
+use crate::serializer::{CompressedProtobufSerializer, UploadCompression};
 use anyhow::Context;
 use interning_api::Generation;
 use std::borrow::Cow;
@@ -323,6 +323,7 @@ impl Profile {
         mut self,
         end_time: Option<SystemTime>,
         duration: Option<Duration>,
+        upload_compression: UploadCompression,
     ) -> anyhow::Result<EncodedProfile> {
         let end = end_time.unwrap_or_else(SystemTime::now);
         let start = self.start_time;
@@ -353,7 +354,10 @@ impl Profile {
         // size of 32KiB should definitely out-perform starting at zero for
         // time consumed, allocator pressure, and allocator fragmentation.
         const INITIAL_PPROF_BUFFER_SIZE: usize = 32 * 1024;
-        let mut encoder = CompressedProtobufSerializer::with_capacity(INITIAL_PPROF_BUFFER_SIZE);
+        let mut encoder = CompressedProtobufSerializer::with_config_and_capacity(
+            upload_compression,
+            INITIAL_PPROF_BUFFER_SIZE,
+        );
 
         for (sample, timestamp, mut values) in std::mem::take(&mut self.observations).into_iter() {
             let labels = self.enrich_sample_labels(sample, timestamp)?;
@@ -1295,7 +1299,7 @@ mod api_tests {
         let profile: Profile = Profile::new(&sample_types, None);
 
         let encoded_profile = profile
-            .serialize_into_compressed_pprof(None, None)
+            .serialize_into_compressed_pprof(None, None, Default::default())
             .expect("Unable to encode/serialize the profile");
 
         let endpoints_stats = encoded_profile.endpoints_stats;
@@ -1319,7 +1323,7 @@ mod api_tests {
         profile.add_endpoint_count(Cow::from(second_endpoint), 1)?;
 
         let encoded_profile = profile
-            .serialize_into_compressed_pprof(None, None)
+            .serialize_into_compressed_pprof(None, None, Default::default())
             .expect("Unable to encode/serialize the profile");
 
         let endpoints_stats = encoded_profile.endpoints_stats;

datadog-profiling/src/pprof/test_utils.rs

@@ -1,6 +1,8 @@
 // Copyright 2023-Present Datadog, Inc. https://www.datadoghq.com/
 // SPDX-License-Identifier: Apache-2.0
 
+use crate::serializer::UploadCompression;
+
 pub fn deserialize_compressed_pprof(encoded: &[u8]) -> anyhow::Result<super::Profile> {
     use prost::Message;
     use std::io::Read;
@@ -12,7 +14,9 @@ pub fn deserialize_compressed_pprof(encoded: &[u8]) -> anyhow::Result<super::Pro
     Ok(profile)
 }
 
+// todo: allow choosing what compressor to use, because in some case we may
+//       want that, and other cases None will execute faster.
 pub fn roundtrip_to_pprof(profile: crate::internal::Profile) -> anyhow::Result<super::Profile> {
-    let encoded = profile.serialize_into_compressed_pprof(None, None)?;
+    let encoded = profile.serialize_into_compressed_pprof(None, None, UploadCompression::Lz4)?;
     deserialize_compressed_pprof(&encoded.buffer)
 }

datadog-profiling/src/serializer/compressed_streaming_encoder.rs

@@ -1,14 +1,36 @@
 // Copyright 2023-Present Datadog, Inc. https://www.datadoghq.com/
 // SPDX-License-Identifier: Apache-2.0
 
+use super::UploadCompression;
 use bytes::BufMut;
 use lz4_flex::frame::FrameEncoder;
 use prost::encoding::{encode_key, encode_varint, encoded_len_varint, key_len, WireType};
-use std::io::Write;
+use std::io::{self, Write};
+
+enum Compressor {
+    None,
+    Lz4 { zipper: FrameEncoder<Vec<u8>> },
+}
+
+impl Compressor {
+    #[inline]
+    fn compress(&mut self, buffer: &mut Vec<u8>) -> io::Result<()> {
+        match self {
+            Compressor::None => Ok(()),
+            Compressor::Lz4 { zipper } => {
+                zipper.write_all(buffer)?;
+                buffer.clear();
+                Ok(())
+            }
+        }
+    }
+}
 
 pub struct CompressedProtobufSerializer {
+    /// Buffer that protobuf is encoded into. Lz4 uses this as a temporary
+    /// buffer, while None uses this as the final output buffer.
     buffer: Vec<u8>,
-    zipper: FrameEncoder<Vec<u8>>,
+    compressor: Compressor,
 }
 
 // I've opened a PR for a generic version of this upstream:
@@ -20,44 +42,54 @@ fn encode_str(tag: u32, value: &str, buf: &mut Vec<u8>) {
 }
 
 impl CompressedProtobufSerializer {
-    pub fn encode(&mut self, item: impl prost::Message) -> anyhow::Result<()> {
-        item.encode(&mut self.buffer)?;
-        self.zipper.write_all(&self.buffer)?;
-        self.buffer.clear();
-        Ok(())
+    pub fn encode(&mut self, item: impl prost::Message) -> io::Result<()> {
+        let buffer = &mut self.buffer;
+        item.encode(buffer)?;
+        self.compressor.compress(buffer)
     }
 
     /// Only meant for string table strings. This is essentially an
     /// implementation of [prost::Message::encode] but for any `AsRef<str>`,
     /// and specialized for handling the unlikely OOM conditions of writing
     /// into a `Vec<u8>`.
-    pub(crate) fn encode_string_table_entry(
-        &mut self,
-        item: impl AsRef<str>,
-    ) -> anyhow::Result<()> {
+    pub(crate) fn encode_string_table_entry(&mut self, item: impl AsRef<str>) -> io::Result<()> {
+        let buffer = &mut self.buffer;
         // In pprof, string tables are tag 6 on the Profile message.
         let tag = 6u32;
         let str = item.as_ref();
         let encoded_len = encoded_len_varint(str.len() as u64);
         let required = key_len(tag) + encoded_len + str.len();
-        if let Err(err) = self.buffer.try_reserve(required) {
-            return Err(anyhow::Error::from(err)
-                .context("failed to encode Protobuf str; insufficient buffer capacity"));
-        }
+        buffer.try_reserve(required)?;
 
-        encode_str(tag, str, &mut self.buffer);
-        self.zipper.write_all(&self.buffer)?;
-        self.buffer.clear();
-        Ok(())
+        encode_str(tag, str, buffer);
+        self.compressor.compress(buffer)
     }
 
-    pub fn finish(self) -> anyhow::Result<Vec<u8>> {
-        Ok(self.zipper.finish()?)
+    pub fn finish(self) -> io::Result<Vec<u8>> {
+        match self.compressor {
+            Compressor::None => Ok(self.buffer),
+            Compressor::Lz4 { zipper } => {
+                debug_assert!(self.buffer.is_empty());
+                Ok(zipper.finish()?)
+            }
+        }
     }
 
-    pub fn with_capacity(capacity: usize) -> Self {
+    pub fn with_config_and_capacity(config: UploadCompression, capacity: usize) -> Self {
+        // Final output buffer.
         let buffer = Vec::with_capacity(capacity);
-        let zipper = FrameEncoder::new(Vec::with_capacity(capacity));
-        Self { buffer, zipper }
+        match config {
+            UploadCompression::Off => Self {
+                buffer,
+                compressor: Compressor::None,
+            },
+            UploadCompression::On | UploadCompression::Lz4 => Self {
+                // Temporary input buffer.
+                buffer: Vec::with_capacity(256),
+                compressor: Compressor::Lz4 {
+                    zipper: FrameEncoder::new(buffer),
+                },
+            },
+        }
     }
 }

datadog-profiling/src/serializer/mod.rs

@@ -4,3 +4,19 @@
 mod compressed_streaming_encoder;
 
 pub use compressed_streaming_encoder::*;
+
+#[repr(C)]
+#[derive(Debug)]
+pub enum UploadCompression {
+    Off,
+    /// On is the default, with the exact compression algorithm being
+    /// unspecified, and free to change. For example, we're testing zstd.
+    On,
+    Lz4,
+}
+
+impl Default for UploadCompression {
+    fn default() -> Self {
+        UploadCompression::On
+    }
+}