refactor: convert enum to traits

morrisonlevi · morrisonlevi · commit f28efa6a7b0f · 2025-10-09T16:57:50.000-06:00
diff --git a/datadog-profiling/src/exporter/mod.rs b/datadog-profiling/src/exporter/mod.rs
@@ -29,7 +29,7 @@ pub use connector::uds::{socket_path_from_uri, socket_path_to_uri};
 pub use connector::named_pipe::{named_pipe_path_from_uri, named_pipe_path_to_uri};
 
 use crate::internal::EncodedProfile;
-use crate::profiles::Compressor;
+use crate::profiles::{Compressor, DefaultProfileCodec};
 
 const DURATION_ZERO: std::time::Duration = std::time::Duration::from_millis(0);
 
@@ -282,8 +282,12 @@ impl ProfileExporter {
             let capacity = (file.bytes.len() / 10).next_power_of_two();
             let max_capacity = 50 * 1024 * 1024;
             let compression_level = 1;
-            let mut encoder = Compressor::try_new(capacity, max_capacity, compression_level)
-                .context("failed to create compressor")?;
+            let mut encoder = Compressor::<DefaultProfileCodec>::try_new(
+                capacity,
+                max_capacity,
+                compression_level,
+            )
+            .context("failed to create compressor")?;
             encoder.write_all(file.bytes)?;
             let encoded = encoder.finish()?;
             /* The Datadog RFC examples strip off the file extension, but the exact behavior
diff --git a/datadog-profiling/src/internal/observation/mod.rs b/datadog-profiling/src/internal/observation/mod.rs
@@ -12,7 +12,6 @@ mod observations;
 mod timestamped_observations;
 mod trimmed_observation;
 
-pub use timestamped_observations::EncodingType;
 // We keep trimmed_observation private, to ensure that only maps can make and
 // operate on trimmed objects, which helps ensure safety.
 pub use observations::*;
diff --git a/datadog-profiling/src/internal/observation/observations.rs b/datadog-profiling/src/internal/observation/observations.rs
@@ -4,7 +4,7 @@
 //! See the mod.rs file comment for why this module and file exists.
 
 use super::super::Sample;
-use super::timestamped_observations::{EncodingType, TimestampedObservations};
+use super::timestamped_observations::TimestampedObservations;
 use super::trimmed_observation::{ObservationLength, TrimmedObservation};
 use crate::internal::Timestamp;
 use std::collections::HashMap;
@@ -15,7 +15,7 @@ struct NonEmptyObservations {
     aggregated_data: AggregatedObservations,
     // Samples with timestamps are all separately kept (so we can know the exact values at the
     // given timestamp)
-    timestamped_data: TimestampedObservations,
+    timestamped_data: Option<TimestampedObservations>,
     obs_len: ObservationLength,
     timestamped_samples_count: usize,
 }
@@ -28,24 +28,15 @@ pub struct Observations {
 /// Public API
 impl Observations {
     pub fn new(observations_len: usize) -> Self {
-        // zstd does FFI calls which miri cannot handle
-        let encoding_type = if cfg!(not(miri)) {
-            EncodingType::Zstd
-        } else {
-            EncodingType::None
-        };
         #[allow(clippy::expect_used)]
-        Self::try_new(encoding_type, observations_len).expect("failed to initialize observations")
+        Self::try_new(observations_len).expect("failed to initialize observations")
     }
 
-    pub fn try_new(encoding_type: EncodingType, observations_len: usize) -> io::Result<Self> {
+    pub fn try_new(observations_len: usize) -> io::Result<Self> {
         Ok(Observations {
             inner: Some(NonEmptyObservations {
                 aggregated_data: AggregatedObservations::new(observations_len),
-                timestamped_data: TimestampedObservations::try_new(
-                    encoding_type,
-                    observations_len,
-                )?,
+                timestamped_data: Some(TimestampedObservations::try_new(observations_len)?),
                 obs_len: ObservationLength::new(observations_len),
                 timestamped_samples_count: 0,
             }),
@@ -73,8 +64,13 @@ impl Observations {
             values.len()
         );
 
+        #[allow(clippy::expect_used)]
         if let Some(ts) = timestamp {
-            observations.timestamped_data.add(sample, ts, values)?;
+            observations
+                .timestamped_data
+                .as_mut()
+                .expect("timestamped_data should be present")
+                .add(sample, ts, values)?;
             observations.timestamped_samples_count += 1;
         } else {
             observations.aggregated_data.add(sample, values)?;
@@ -188,12 +184,13 @@ impl IntoIterator for Observations {
 
     fn into_iter(self) -> Self::IntoIter {
         let it = self.inner.into_iter().flat_map(|mut observations| {
-            let timestamped_data_it = std::mem::replace(
-                &mut observations.timestamped_data,
-                TimestampedObservations::with_no_backing_store(),
-            )
-            .into_iter()
-            .map(|(s, t, o)| (s, Some(t), o));
+            #[allow(clippy::expect_used)]
+            let timestamped_data_it = observations
+                .timestamped_data
+                .take()
+                .expect("timestamped_data should be present")
+                .into_iter()
+                .map(|(s, t, o)| (s, Some(t), o));
             let aggregated_data_it = std::mem::take(&mut observations.aggregated_data.data)
                 .into_iter()
                 .map(|(s, o)| (s, None, o))
diff --git a/datadog-profiling/src/internal/observation/timestamped_observations.rs b/datadog-profiling/src/internal/observation/timestamped_observations.rs
@@ -10,93 +10,23 @@ use super::super::Sample;
 use super::super::StackTraceId;
 use crate::collections::identifiable::Id;
 use crate::internal::Timestamp;
-use crate::profiles::SizeRestrictedBuffer;
+use crate::profiles::{DefaultObservationCodec as DefaultCodec, ObservationCodec};
 use byteorder::{NativeEndian, ReadBytesExt};
-use std::io::{self, Cursor, Read, Write};
-
-// todo: document
-#[derive(Clone, Copy, Debug, Default)]
-pub enum EncodingType {
-    #[default]
-    None,
-    Zstd,
-}
+use std::io::{self, Write};
 
-enum ObservationEncoder {
-    Noop(SizeRestrictedBuffer),
-    Zstd(zstd::Encoder<'static, SizeRestrictedBuffer>),
-}
+pub type TimestampedObservations = TimestampedObservationsImpl<DefaultCodec>;
 
-pub struct TimestampedObservations {
-    compressed_timestamped_data: ObservationEncoder,
+pub struct TimestampedObservationsImpl<C: ObservationCodec> {
+    compressed_timestamped_data: C::Encoder,
     sample_types_len: usize,
 }
 
-enum ObservationDecoder {
-    Noop(Cursor<SizeRestrictedBuffer>),
-    Zstd(zstd::Decoder<'static, Cursor<SizeRestrictedBuffer>>),
-}
-
-pub struct TimestampedObservationsIter {
-    decoder: ObservationDecoder,
+pub struct TimestampedObservationsIterImpl<C: ObservationCodec> {
+    decoder: C::Decoder,
     sample_types_len: usize,
 }
 
-impl ObservationEncoder {
-    pub fn try_new(
-        encoding_type: EncodingType,
-        size_hint: usize,
-        max_capacity: usize,
-    ) -> io::Result<Self> {
-        let output_buffer = SizeRestrictedBuffer::try_new(size_hint, max_capacity)?;
-        match encoding_type {
-            EncodingType::None => Ok(ObservationEncoder::Noop(output_buffer)),
-            EncodingType::Zstd => Ok(ObservationEncoder::Zstd(zstd::Encoder::new(
-                output_buffer,
-                1,
-            )?)),
-        }
-    }
-
-    pub fn try_into_decoder(self) -> io::Result<ObservationDecoder> {
-        match self {
-            ObservationEncoder::Noop(buffer) => Ok(ObservationDecoder::Noop(Cursor::new(buffer))),
-            ObservationEncoder::Zstd(encoder) => match encoder.try_finish() {
-                Ok(buffer) => Ok(ObservationDecoder::Zstd(zstd::Decoder::with_buffer(
-                    Cursor::new(buffer),
-                )?)),
-                Err((_encoder, error)) => Err(error),
-            },
-        }
-    }
-}
-
-impl Write for ObservationEncoder {
-    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
-        match self {
-            ObservationEncoder::Noop(encoder) => encoder.write(buf),
-            ObservationEncoder::Zstd(encoder) => encoder.write(buf),
-        }
-    }
-
-    fn flush(&mut self) -> io::Result<()> {
-        match self {
-            ObservationEncoder::Noop(encoder) => encoder.flush(),
-            ObservationEncoder::Zstd(encoder) => encoder.flush(),
-        }
-    }
-}
-
-impl Read for ObservationDecoder {
-    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
-        match self {
-            ObservationDecoder::Noop(decoder) => decoder.read(buf),
-            ObservationDecoder::Zstd(decoder) => decoder.read(buf),
-        }
-    }
-}
-
-impl TimestampedObservations {
+impl<C: ObservationCodec> TimestampedObservationsImpl<C> {
     // As documented in the internal Datadog doc "Ruby timeline memory fragmentation impact
     // investigation", allowing the timeline storage vec to slowly expand creates A LOT of
     // memory fragmentation for apps that employ multiple threads.
@@ -108,26 +38,16 @@ impl TimestampedObservations {
     // the profile as a whole would defintely exceed this.
     const MAX_CAPACITY: usize = i32::MAX as usize;
 
-    pub fn try_new(encoding_type: EncodingType, sample_types_len: usize) -> io::Result<Self> {
+    pub fn try_new(sample_types_len: usize) -> io::Result<Self> {
         Ok(Self {
-            compressed_timestamped_data: ObservationEncoder::try_new(
-                encoding_type,
+            compressed_timestamped_data: C::new_encoder(
                 Self::DEFAULT_BUFFER_SIZE,
                 Self::MAX_CAPACITY,
             )?,
             sample_types_len,
         })
     }
 
-    pub const fn with_no_backing_store() -> Self {
-        Self {
-            compressed_timestamped_data: ObservationEncoder::Noop(
-                SizeRestrictedBuffer::zero_capacity(),
-            ),
-            sample_types_len: 0,
-        }
-    }
-
     pub fn add(&mut self, sample: Sample, ts: Timestamp, values: &[i64]) -> anyhow::Result<()> {
         // We explicitly turn the data into a stream of bytes, feeding it to the compressor.
         // @ivoanjo: I played with introducing a structure to serialize it all-at-once, but it seems
@@ -153,19 +73,17 @@ impl TimestampedObservations {
         Ok(())
     }
 
-    pub fn into_iter(self) -> TimestampedObservationsIter {
+    pub fn into_iter(self) -> TimestampedObservationsIterImpl<C> {
         #[allow(clippy::expect_used)]
-        TimestampedObservationsIter {
-            decoder: self
-                .compressed_timestamped_data
-                .try_into_decoder()
+        TimestampedObservationsIterImpl {
+            decoder: C::encoder_into_decoder(self.compressed_timestamped_data)
                 .expect("failed to initialize timestamped observation decoder"),
             sample_types_len: self.sample_types_len,
         }
     }
 }
 
-impl Iterator for TimestampedObservationsIter {
+impl<C: ObservationCodec> Iterator for TimestampedObservationsIterImpl<C> {
     type Item = (Sample, Timestamp, Vec<i64>);
 
     fn next(&mut self) -> Option<Self::Item> {
diff --git a/datadog-profiling/src/internal/profile/mod.rs b/datadog-profiling/src/internal/profile/mod.rs
@@ -14,7 +14,7 @@ use crate::collections::identifiable::*;
 use crate::collections::string_storage::{CachedProfileId, ManagedStringStorage};
 use crate::collections::string_table::{self, StringTable};
 use crate::iter::{IntoLendingIterator, LendingIterator};
-use crate::profiles::Compressor;
+use crate::profiles::{Compressor, DefaultProfileCodec};
 use anyhow::Context;
 use datadog_profiling_protobuf::{self as protobuf, Record, Value, NO_OPT_ZERO, OPT_ZERO};
 use interning_api::Generation;
@@ -348,7 +348,7 @@ impl Profile {
         // level 1 provided better compressed files while taking less time
         // compared to lz4.
         const COMPRESSION_LEVEL: i32 = 1;
-        let mut compressor = Compressor::try_new(
+        let mut compressor = Compressor::<DefaultProfileCodec>::try_new(
             INITIAL_PPROF_BUFFER_SIZE,
             MAX_PROFILE_SIZE,
             COMPRESSION_LEVEL,
diff --git a/datadog-profiling/src/pprof/test_utils.rs b/datadog-profiling/src/pprof/test_utils.rs
@@ -5,14 +5,20 @@ use anyhow::Context;
 use datadog_profiling_protobuf::prost_impls::{Profile, Sample};
 use std::io::Cursor;
 
-pub fn deserialize_compressed_pprof(encoded: &[u8]) -> anyhow::Result<Profile> {
+fn deserialize_compressed_pprof(encoded: &[u8]) -> anyhow::Result<Profile> {
     use prost::Message;
     use std::io::Read;
 
-    let mut decoder =
-        zstd::Decoder::new(Cursor::new(encoded)).context("failed to create zstd decoder")?;
-    let mut buf = Vec::new();
-    decoder.read_to_end(&mut buf)?;
+    #[cfg(miri)]
+    let buf = encoded.to_vec();
+    #[cfg(not(miri))]
+    let buf = {
+        let mut decoder =
+            zstd::Decoder::new(Cursor::new(encoded)).context("failed to create zstd decoder")?;
+        let mut out = Vec::new();
+        decoder.read_to_end(&mut out)?;
+        out
+    };
     let profile = Profile::decode(buf.as_slice())?;
     Ok(profile)
 }
diff --git a/datadog-profiling/src/profiles/compressor.rs b/datadog-profiling/src/profiles/compressor.rs
