feat: Modularize code for different build; Fallback to zeekstd for x86

Author: SersemPeca

runtime_tracing/Cargo.toml

@@ -22,6 +22,11 @@ serde_repr = "0.1"
 capnp = "0.21.1"
 cbor4ii = { version = "1.0.0", features = ["serde1", "use_std"] }
 fscommon = "0.1.1"
+
+[target.'cfg(not(target_arch = "x86_64"))'.dependencies]
+zeekstd = "0.6.0"
+
+[target.'cfg(target_arch = "wasm32")'.dependencies]
 ruzstd = "0.8.1"
 
 [build-dependencies]

runtime_tracing/src/cbor_zstd_reader.rs

@@ -2,7 +2,7 @@ use std::io::{self, BufRead, BufReader, Read, Seek, Write};
 
 use fscommon::StreamSlice;
 
-use ruzstd::decoding::StreamingDecoder;
+use zeekstd::Decoder;
 
 use crate::{TraceLowLevelEvent, cbor_zstd_writer::HEADERV1};
 
@@ -11,6 +11,7 @@ fn is_at_eof<R: BufRead>(reader: &mut R) -> io::Result<bool> {
     Ok(buffer.is_empty())
 }
 
+#[cfg(target_arch = "x86_64")]
 pub fn read_trace(input: &mut (impl Read + Write + Seek)) -> Result<Vec<TraceLowLevelEvent>, Box<dyn std::error::Error>> {
     let end_pos = input.seek(io::SeekFrom::End(0))?;
     input.seek(io::SeekFrom::Start(0))?;
@@ -25,7 +26,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 = StreamingDecoder::new(input2)?;
+    let decoder = Decoder::new(input2)?;
     let mut buf_reader = BufReader::new(decoder);
 
     let mut result: Vec<TraceLowLevelEvent> = vec![];

runtime_tracing/src/cbor_zstd_reader_wasm.rs

@@ -0,0 +1,39 @@
+use std::io::{self, BufRead, BufReader, Read, Seek, Write};
+
+use fscommon::StreamSlice;
+
+use ruzstd::decoding::StreamingDecoder;
+
+use crate::{TraceLowLevelEvent, cbor_zstd_writer::HEADERV1};
+
+fn is_at_eof<R: BufRead>(reader: &mut R) -> io::Result<bool> {
+    let buffer = reader.fill_buf()?;
+    Ok(buffer.is_empty())
+}
+
+pub fn read_trace(input: &mut (impl Read + Write + Seek)) -> Result<Vec<TraceLowLevelEvent>, Box<dyn std::error::Error>> {
+    let end_pos = input.seek(io::SeekFrom::End(0))?;
+    input.seek(io::SeekFrom::Start(0))?;
+
+    let mut header_buf = [0; 8];
+    let mut buf_reader = BufReader::new(&mut *input);
+    buf_reader.read_exact(&mut header_buf)?;
+    if header_buf != HEADERV1 {
+        panic!("Invalid file header (wrong file format or incompatible version)");
+    }
+
+    input.seek(io::SeekFrom::Start(0))?;
+    let input2 = StreamSlice::new(&mut *input, 8, end_pos)?;
+
+    let decoder = StreamingDecoder::new(input2)?;
+    let mut buf_reader = BufReader::new(decoder);
+
+    let mut result: Vec<TraceLowLevelEvent> = vec![];
+
+    while !is_at_eof(&mut buf_reader)? {
+        let obj = cbor4ii::serde::from_reader::<TraceLowLevelEvent, _>(&mut buf_reader)?;
+        result.push(obj);
+    }
+
+    Ok(result)
+}

runtime_tracing/src/cbor_zstd_writer.rs

@@ -4,7 +4,8 @@ use std::{
     path::PathBuf,
 };
 
-use ruzstd::encoding::{CompressionLevel, compress};
+#[cfg(target_arch = "x86_64")]
+use zeekstd::Encoder;
 
 use crate::{
     TraceLowLevelEvent,
@@ -20,28 +21,26 @@ pub const HEADERV1: &[u8] = &[
     0x00, 0x00,
 ]; // Reserved, must be zero in this version.
 
-pub struct CborZstdTraceWriter {
+pub struct CborZstdTraceWriter<'a> {
     base: AbstractTraceWriterData,
 
     trace_events_path: Option<PathBuf>,
-    trace_events_file: Option<File>,
-    uncompressed_buf: Vec<u8>,
+    trace_events_file_zstd_encoder: Option<Encoder<'a, File>>,
 }
 
-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: None,
-            uncompressed_buf: Vec::new(),
+            trace_events_file_zstd_encoder: None,
         }
     }
 }
 
-impl AbstractTraceWriter for CborZstdTraceWriter {
+impl AbstractTraceWriter for CborZstdTraceWriter<'_> {
     fn get_data(&self) -> &AbstractTraceWriterData {
         &self.base
     }
@@ -52,37 +51,35 @@ 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).expect("CBOR encode failed");
-        self.uncompressed_buf.extend_from_slice(&q);
+        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();
+        }
     }
 
     fn append_events(&mut self, events: &mut Vec<TraceLowLevelEvent>) {
-        for e in events.drain(..) {
-            <Self as AbstractTraceWriter>::add_event(self, e);
+        for e in events {
+            AbstractTraceWriter::add_event(self, e.clone());
         }
     }
 }
 
-impl TraceWriter for CborZstdTraceWriter {
+#[cfg(target_arch = "x86_64")]
+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 = Some(file_output);
+        self.trace_events_file_zstd_encoder = Some(Encoder::new(file_output)?);
 
         Ok(())
     }
 
     fn finish_writing_trace_events(&mut self) -> Result<(), Box<dyn std::error::Error>> {
-        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()?;
+        if let Some(enc) = self.trace_events_file_zstd_encoder.take() {
+            enc.finish()?;
 
-            self.uncompressed_buf.clear();
             Ok(())
         } else {
             panic!("finish_writing_trace_events() called without previous call to begin_writing_trace_events()");

runtime_tracing/src/cbor_zstd_writer_wasm.rs

@@ -0,0 +1,91 @@
+use std::{
+    fs::File,
+    io::{Cursor, Write},
+    path::PathBuf,
+};
+
+use ruzstd::encoding::{CompressionLevel, compress};
+
+use crate::{
+    TraceLowLevelEvent,
+    abstract_trace_writer::{AbstractTraceWriter, AbstractTraceWriterData},
+    trace_writer::TraceWriter,
+};
+
+/// The next 3 bytes are reserved/version info.
+/// The header is 8 bytes in size, ensuring 64-bit alignment for the rest of the file.
+pub const HEADERV1: &[u8] = &[
+    0xC0, 0xDE, 0x72, 0xAC, 0xE2, // The first 5 bytes identify the file as a CodeTracer file (hex l33tsp33k - C0DE72ACE2 for "CodeTracer").
+    0x01, // Indicates version 1 of the file format
+    0x00, 0x00,
+]; // Reserved, must be zero in this version.
+
+pub struct CborZstdTraceWriter {
+    base: AbstractTraceWriterData,
+
+    trace_events_path: Option<PathBuf>,
+    trace_events_file: Option<File>,
+    uncompressed_buf: Vec<u8>,
+}
+
+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: None,
+            uncompressed_buf: vec![],
+        }
+    }
+}
+
+impl AbstractTraceWriter for CborZstdTraceWriter {
+    fn get_data(&self) -> &AbstractTraceWriterData {
+        &self.base
+    }
+
+    fn get_mut_data(&mut self) -> &mut AbstractTraceWriterData {
+        &mut self.base
+    }
+
+    fn add_event(&mut self, event: TraceLowLevelEvent) {
+        let buf: Vec<u8> = Vec::new();
+        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.drain(..) {
+            <Self as AbstractTraceWriter>::add_event(self, e);
+        }
+    }
+}
+
+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 = Some(file_output);
+
+        Ok(())
+    }
+
+    fn finish_writing_trace_events(&mut self) -> Result<(), Box<dyn std::error::Error>> {
+        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()");
+        }
+    }
+}

runtime_tracing/src/lib.rs

@@ -11,8 +11,19 @@
 mod abstract_trace_writer;
 mod base64;
 mod capnptrace;
+
+#[cfg(target_arch = "wasm32")]
+#[path = "./cbor_zstd_reader_wasm.rs"]
 mod cbor_zstd_reader;
+#[cfg(target_arch = "wasm32")]
+#[path = "./cbor_zstd_writer_wasm.rs"]
 mod cbor_zstd_writer;
+
+#[cfg(not(target_arch = "wasm32"))]
+mod cbor_zstd_reader;
+#[cfg(not(target_arch = "wasm32"))]
+mod cbor_zstd_writer;
+
 mod non_streaming_trace_writer;
 mod trace_readers;
 mod trace_writer;