add BufWriter to limit calls to ZSTD_compressStream

Author: realFlowControl

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

@@ -12,12 +12,12 @@ use crate::collections::identifiable::Id;
 use crate::internal::Timestamp;
 use crate::profiles::{DefaultObservationCodec as DefaultCodec, ObservationCodec};
 use byteorder::{NativeEndian, ReadBytesExt};
-use std::io::{self, Write};
+use std::io::{self, BufWriter, Write};
 
 pub type TimestampedObservations = TimestampedObservationsImpl<DefaultCodec>;
 
 pub struct TimestampedObservationsImpl<C: ObservationCodec> {
-    compressed_timestamped_data: C::Encoder,
+    compressed_timestamped_data: BufWriter<C::Encoder>,
     sample_types_len: usize,
 }
 
@@ -40,10 +40,10 @@ impl<C: ObservationCodec> TimestampedObservationsImpl<C> {
 
     pub fn try_new(sample_types_len: usize) -> io::Result<Self> {
         Ok(Self {
-            compressed_timestamped_data: C::new_encoder(
-                Self::DEFAULT_BUFFER_SIZE,
-                Self::MAX_CAPACITY,
-            )?,
+            compressed_timestamped_data: BufWriter::with_capacity(
+                C::recommended_input_buf_size(),
+                C::new_encoder(Self::DEFAULT_BUFFER_SIZE, Self::MAX_CAPACITY)?,
+            ),
             sample_types_len,
         })
     }
@@ -74,8 +74,12 @@ impl<C: ObservationCodec> TimestampedObservationsImpl<C> {
     }
 
     pub fn try_into_iter(self) -> io::Result<TimestampedObservationsIterImpl<C>> {
+        let encoder = self
+            .compressed_timestamped_data
+            .into_inner()
+            .map_err(|e| e.into_error())?;
         Ok(TimestampedObservationsIterImpl {
-            decoder: C::encoder_into_decoder(self.compressed_timestamped_data)?,
+            decoder: C::encoder_into_decoder(encoder)?,
             sample_types_len: self.sample_types_len,
         })
     }

libdd-profiling/src/profiles/compressor.rs

@@ -145,6 +145,10 @@ pub trait ObservationCodec {
 
     fn new_encoder(size_hint: usize, max_capacity: usize) -> io::Result<Self::Encoder>;
     fn encoder_into_decoder(encoder: Self::Encoder) -> io::Result<Self::Decoder>;
+
+    /// Returns the recommended input buffer size for the encoder.
+    /// Used to size the `BufWriter` that wraps the encoder.
+    fn recommended_input_buf_size() -> usize;
 }
 
 #[allow(unused)]
@@ -161,6 +165,10 @@ impl ObservationCodec for NoopObservationCodec {
     fn encoder_into_decoder(encoder: Self::Encoder) -> io::Result<Self::Decoder> {
         Ok(io::Cursor::new(encoder))
     }
+
+    fn recommended_input_buf_size() -> usize {
+        0
+    }
 }
 
 #[allow(unused)]
@@ -181,6 +189,10 @@ impl ObservationCodec for ZstdObservationCodec {
             Err((_enc, error)) => Err(error),
         }
     }
+
+    fn recommended_input_buf_size() -> usize {
+        zstd::Encoder::<SizeRestrictedBuffer>::recommended_input_size()
+    }
 }
 
 #[cfg(not(miri))]