feat: rewrite the compression logic to use ruzstd instead of zeekstd

Author: SersemPeca

runtime_tracing/Cargo.toml

@@ -21,9 +21,9 @@ serde_json = "1.0"
 serde_repr = "0.1"
 schemars = "0.8.2"
 capnp = "0.21.1"
-zeekstd = "0.5.0"
 cbor4ii = { version = "1.0.0", features = ["serde1", "use_std"] }
 fscommon = "0.1.1"
+ruzstd = "0.8.1"
 
 [build-dependencies]
 capnpc = "0.21.0"

runtime_tracing/src/cbor_zstd_reader.rs

@@ -1,7 +1,8 @@
 use std::io::{self, BufRead, BufReader, Read, Seek, Write};
 
 use fscommon::StreamSlice;
-use zeekstd::Decoder;
+
+use ruzstd::decoding::StreamingDecoder;
 
 use crate::{TraceLowLevelEvent, cbor_zstd_writer::HEADERV1};
 
@@ -24,7 +25,7 @@ pub fn read_trace(input: &mut (impl Read + Write + Seek)) -> Result<Vec<TraceLow
     input.seek(io::SeekFrom::Start(0))?;
     let input2 = StreamSlice::new(&mut *input, 8, end_pos)?;
 
-    let decoder = Decoder::new(input2)?;
+    let decoder = StreamingDecoder::new(input2)?;
     let mut buf_reader = BufReader::new(decoder);
 
     let mut result: Vec<TraceLowLevelEvent> = vec![];

runtime_tracing/src/cbor_zstd_writer.rs

@@ -1,5 +1,10 @@
-use std::{fs::File, io::Write, path::PathBuf};
-use zeekstd::Encoder;
+use std::{
+    fs::File,
+    io::{Cursor, Write},
+    path::PathBuf,
+};
+
+use ruzstd::encoding::{CompressionLevel, compress};
 
 use crate::{
     TraceLowLevelEvent,
@@ -15,26 +20,28 @@ pub const HEADERV1: &[u8] = &[
     0x00, 0x00,
 ]; // Reserved, must be zero in this version.
 
-pub struct CborZstdTraceWriter<'a> {
+pub struct CborZstdTraceWriter {
     base: AbstractTraceWriterData,
 
     trace_events_path: Option<PathBuf>,
-    trace_events_file_zstd_encoder: Option<Encoder<'a, File>>,
+    trace_events_file: Option<File>,
+    uncompressed_buf: Vec<u8>,
 }
 
-impl CborZstdTraceWriter<'_> {
+impl CborZstdTraceWriter {
     /// Create a new tracer instance for the given program and arguments.
     pub fn new(program: &str, args: &[String]) -> Self {
         CborZstdTraceWriter {
             base: AbstractTraceWriterData::new(program, args),
 
             trace_events_path: None,
-            trace_events_file_zstd_encoder: None,
+            trace_events_file: None,
+            uncompressed_buf: Vec::new(),
         }
     }
 }
 
-impl AbstractTraceWriter for CborZstdTraceWriter<'_> {
+impl AbstractTraceWriter for CborZstdTraceWriter {
     fn get_data(&self) -> &AbstractTraceWriterData {
         &self.base
     }
@@ -45,34 +52,37 @@ impl AbstractTraceWriter for CborZstdTraceWriter<'_> {
 
     fn add_event(&mut self, event: TraceLowLevelEvent) {
         let buf: Vec<u8> = Vec::new();
-        let q = cbor4ii::serde::to_vec(buf, &event).unwrap();
-        if let Some(enc) = &mut self.trace_events_file_zstd_encoder {
-            enc.write_all(&q).unwrap();
-        }
+        let q = cbor4ii::serde::to_vec(buf, &event).expect("CBOR encode failed");
+        self.uncompressed_buf.extend_from_slice(&q);
     }
 
     fn append_events(&mut self, events: &mut Vec<TraceLowLevelEvent>) {
-        for e in events {
-            AbstractTraceWriter::add_event(self, e.clone());
+        for e in events.drain(..) {
+            <Self as AbstractTraceWriter>::add_event(self, e);
         }
     }
 }
 
-impl TraceWriter for CborZstdTraceWriter<'_> {
+impl TraceWriter for CborZstdTraceWriter {
     fn begin_writing_trace_events(&mut self, path: &std::path::Path) -> Result<(), Box<dyn std::error::Error>> {
         let pb = path.to_path_buf();
         self.trace_events_path = Some(pb.clone());
+
         let mut file_output = std::fs::File::create(pb)?;
         file_output.write_all(HEADERV1)?;
-        self.trace_events_file_zstd_encoder = Some(Encoder::new(file_output)?);
+        self.trace_events_file = Some(file_output);
 
         Ok(())
     }
 
     fn finish_writing_trace_events(&mut self) -> Result<(), Box<dyn std::error::Error>> {
-        if let Some(enc) = self.trace_events_file_zstd_encoder.take() {
-            enc.finish()?;
+        if let Some(mut file) = self.trace_events_file.take() {
+            let mut cursor = Cursor::new(&self.uncompressed_buf);
+            compress(&mut cursor, &mut file, CompressionLevel::Fastest);
+
+            file.flush()?;
 
+            self.uncompressed_buf.clear();
             Ok(())
         } else {
             panic!("finish_writing_trace_events() called without previous call to begin_writing_trace_events()");