refactor: use enum for noop vs zstd compressor

Author: morrisonlevi

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

@@ -12,6 +12,7 @@ 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::*;

datadog-profiling/src/internal/observation/observations.rs

@@ -4,10 +4,11 @@
 //! See the mod.rs file comment for why this module and file exists.
 
 use super::super::Sample;
-use super::timestamped_observations::TimestampedObservations;
+use super::timestamped_observations::{EncodingType, TimestampedObservations};
 use super::trimmed_observation::{ObservationLength, TrimmedObservation};
 use crate::internal::Timestamp;
 use std::collections::HashMap;
+use std::io;
 
 struct NonEmptyObservations {
     // Samples with no timestamps are aggregated in-place as each observation is added
@@ -27,14 +28,28 @@ pub struct Observations {
 /// Public API
 impl Observations {
     pub fn new(observations_len: usize) -> Self {
-        Observations {
+        // 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")
+    }
+
+    pub fn try_new(encoding_type: EncodingType, observations_len: usize) -> io::Result<Self> {
+        Ok(Observations {
             inner: Some(NonEmptyObservations {
                 aggregated_data: AggregatedObservations::new(observations_len),
-                timestamped_data: TimestampedObservations::new(observations_len),
+                timestamped_data: TimestampedObservations::try_new(
+                    encoding_type,
+                    observations_len,
+                )?,
                 obs_len: ObservationLength::new(observations_len),
                 timestamped_samples_count: 0,
             }),
-        }
+        })
     }
 
     pub fn add(

datadog-profiling/src/internal/observation/timestamped_observations.rs

@@ -10,40 +10,120 @@ use super::super::Sample;
 use super::super::StackTraceId;
 use crate::collections::identifiable::Id;
 use crate::internal::Timestamp;
+use crate::profiles::SizeRestrictedBuffer;
 use byteorder::{NativeEndian, ReadBytesExt};
-use std::io::Write;
-use std::io::{BufReader, Cursor};
+use std::io::{self, Cursor, Read, Write};
+
+// todo: document
+#[derive(Clone, Copy, Debug, Default)]
+pub enum EncodingType {
+    #[default]
+    None,
+    Zstd,
+}
+
+enum ObservationEncoder {
+    Noop(SizeRestrictedBuffer),
+    Zstd(zstd::Encoder<'static, SizeRestrictedBuffer>),
+}
 
 pub struct TimestampedObservations {
-    compressed_timestamped_data: zstd::Encoder<'static, Vec<u8>>,
+    compressed_timestamped_data: ObservationEncoder,
+    sample_types_len: usize,
+}
+
+enum ObservationDecoder {
+    Noop(Cursor<SizeRestrictedBuffer>),
+    Zstd(zstd::Decoder<'static, Cursor<SizeRestrictedBuffer>>),
+}
+
+pub struct TimestampedObservationsIter {
+    decoder: ObservationDecoder,
     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 {
     // 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.
     // To avoid this, we've picked a default buffer size of 1MB that very rarely needs to grow, and
     // when it does, is expected to grow in larger steps.
-    const DEFAULT_BUFFER_SIZE: usize = 1_048_576;
-
-    pub fn new(sample_types_len: usize) -> Self {
-        #[allow(clippy::expect_used)] // previous API panic'd implicitly
-        Self {
-            compressed_timestamped_data: zstd::Encoder::new(
-                Vec::with_capacity(Self::DEFAULT_BUFFER_SIZE),
-                1,
-            )
-            .expect("failed to create zstd encoder"),
+    const DEFAULT_BUFFER_SIZE: usize = 1024 * 1024;
+
+    // Protobufs can't exceed 2 GiB, if our observations grow this large, then
+    // 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> {
+        Ok(Self {
+            compressed_timestamped_data: ObservationEncoder::try_new(
+                encoding_type,
+                Self::DEFAULT_BUFFER_SIZE,
+                Self::MAX_CAPACITY,
+            )?,
             sample_types_len,
-        }
+        })
     }
 
-    pub fn with_no_backing_store() -> Self {
-        #[allow(clippy::expect_used)] // previous API panic'd implicitly
+    pub const fn with_no_backing_store() -> Self {
         Self {
-            compressed_timestamped_data: zstd::Encoder::new(vec![], 1)
-                .expect("failed to create zstd encoder"),
+            compressed_timestamped_data: ObservationEncoder::Noop(
+                SizeRestrictedBuffer::zero_capacity(),
+            ),
             sample_types_len: 0,
         }
     }
@@ -74,22 +154,17 @@ impl TimestampedObservations {
     }
 
     pub fn into_iter(self) -> TimestampedObservationsIter {
-        #[allow(clippy::expect_used, clippy::unwrap_used)]
+        #[allow(clippy::expect_used)]
         TimestampedObservationsIter {
-            decoder: zstd::Decoder::new(Cursor::new(
-                self.compressed_timestamped_data.finish().unwrap(),
-            ))
-            .expect("failed to create zstd decoder"),
+            decoder: self
+                .compressed_timestamped_data
+                .try_into_decoder()
+                .expect("failed to initialize timestamped observation decoder"),
             sample_types_len: self.sample_types_len,
         }
     }
 }
 
-pub struct TimestampedObservationsIter {
-    decoder: zstd::Decoder<'static, BufReader<Cursor<Vec<u8>>>>,
-    sample_types_len: usize,
-}
-
 impl Iterator for TimestampedObservationsIter {
     type Item = (Sample, Timestamp, Vec<i64>);
 

datadog-profiling/src/profiles/compressor.rs

@@ -13,6 +13,14 @@ pub struct SizeRestrictedBuffer {
 }
 
 impl SizeRestrictedBuffer {
+    /// Returns a buffer which never holds any data.
+    pub const fn zero_capacity() -> Self {
+        Self {
+            vec: Vec::new(),
+            max_capacity: 0,
+        }
+    }
+
     pub fn try_new(size_hint: usize, max_capacity: usize) -> io::Result<Self> {
         if size_hint > max_capacity {
             return Err(io::Error::new(
@@ -27,7 +35,7 @@ impl SizeRestrictedBuffer {
         Ok(SizeRestrictedBuffer { vec, max_capacity })
     }
 
-    pub fn as_slice(&self) -> &[u8] {
+    pub const fn as_slice(&self) -> &[u8] {
         self.vec.as_slice()
     }
 }
@@ -38,6 +46,12 @@ impl From<SizeRestrictedBuffer> for Vec<u8> {
     }
 }
 
+impl AsRef<[u8]> for SizeRestrictedBuffer {
+    fn as_ref(&self) -> &[u8] {
+        self.as_slice()
+    }
+}
+
 impl Write for SizeRestrictedBuffer {
     #[inline]
     fn write(&mut self, buf: &[u8]) -> io::Result<usize> {