Implement zstd streaming log compression

Cargo.toml

@@ -30,3 +30,4 @@ tabled = "0.18"
 tempfile = "3.27.0"
 tracing = "0.1"
 tracing-subscriber = "0.3"
+zstd = { version = "0.13.3", default-features = false }

candumpr/Cargo.toml

@@ -21,6 +21,7 @@ libc.workspace = true
 neli.workspace = true
 tracing.workspace = true
 tracing-subscriber.workspace = true
+zstd.workspace = true
 
 [dev-dependencies]
 ctor.workspace = true

candumpr/src/main.rs

@@ -109,10 +109,12 @@ enum Format {
 
 impl Format {
     /// Log file extension for each output format
-    fn ext(&self) -> &'static str {
-        match self {
-            Format::CandumpFile => "log",
-            Format::CandumpConsole => "txt",
+    fn ext(&self, compress: bool) -> &'static str {
+        match (self, compress) {
+            (Format::CandumpFile, false) => "log",
+            (Format::CandumpFile, true) => "log.zst",
+            (Format::CandumpConsole, false) => "txt",
+            (Format::CandumpConsole, true) => "txt.zst",
         }
     }
 }
@@ -137,6 +139,10 @@ struct Cli {
     #[arg(long, value_enum, default_value = "candump-file")]
     format: Format,
 
+    /// Compress output with zstd. Requires --log or --output.
+    #[arg(long, short = 'c')]
+    compress: bool,
+
     /// Timestamp rendering mode. Only applies to the candump formats.
     #[arg(long, value_enum, default_value = "absolute")]
     timestamp: TimestampMode,
@@ -170,6 +176,23 @@ fn main() -> ExitCode {
         return ExitCode::FAILURE;
     }
 
+    // Compressing stdout would hand a terminal a binary stream, and the pipe case is served just as
+    // well by `candumpr can0 | zstd`.
+    if cli.compress && !cli.log && cli.output.is_none() {
+        tracing::error!("--compress requires --log or --output");
+        return ExitCode::FAILURE;
+    }
+    if let Some(path) = &cli.output
+        && cli.compress
+        && path.extension().is_none_or(|ext| ext != "zst")
+    {
+        // File extensions are useful, but not required. Give a QoL warning.
+        tracing::warn!(
+            path = %path.display(),
+            "output is zstd-compressed but the path does not end in .zst"
+        );
+    }
+
     // The sockets vector defines the canonical interface ordering. The orderings of:
     //
     // 1. cli.interfaces
@@ -253,7 +276,7 @@ fn main() -> ExitCode {
             .map(|interface| Output::Template {
                 dir: ".".into(),
                 interface: interface.clone(),
-                ext: cli.format.ext().to_string(),
+                ext: cli.format.ext(cli.compress).to_string(),
                 next_index: 0,
             })
             .collect()
@@ -268,6 +291,7 @@ fn main() -> ExitCode {
             let formatter = make_formatter();
             let mut config = SinkConfig::new(output);
             config.header = formatter.header().map(|h| h.to_vec());
+            config.compress = cli.compress;
             (formatter, Sink::new(config))
         })
         .collect();

candumpr/src/sink.rs

@@ -3,7 +3,7 @@ use std::path::PathBuf;
 use std::time::{Duration, Instant};
 
 use crate::recv::Timestamp;
-use crate::writer::{FileWriter, StdoutWriter, Writer};
+use crate::writer::{FileWriter, StdoutWriter, Writer, ZstdWriter};
 
 /// Where a [Sink] writes.
 pub enum Output {
@@ -21,42 +21,6 @@ pub enum Output {
     },
 }
 
-impl Output {
-    /// Construct the writer for this output type.
-    fn open(
-        &self,
-        timestamp: Timestamp,
-        flush_threshold_bytes: usize,
-    ) -> std::io::Result<Box<dyn Writer>> {
-        let path = match self {
-            Output::Stdout => return Ok(Box::new(StdoutWriter::new())),
-            Output::Path(path) => path.clone(),
-            Output::Template {
-                dir,
-                interface,
-                ext,
-                next_index,
-            } => dir.join(template_filename(
-                *next_index,
-                interface,
-                timestamp.sec,
-                ext,
-            )),
-        };
-        if let Some(parent) = path.parent()
-            && !parent.as_os_str().is_empty()
-        {
-            std::fs::create_dir_all(parent)?;
-        }
-        let file = std::fs::File::create(&path)?;
-        tracing::info!(path = %path.display(), "created log file");
-        Ok(Box::new(std::io::BufWriter::with_capacity(
-            flush_threshold_bytes,
-            FileWriter::new(file),
-        )))
-    }
-}
-
 /// Configuration for a [Sink].
 pub struct SinkConfig {
     pub output: Output,
@@ -70,6 +34,8 @@ pub struct SinkConfig {
     pub sync_interval: Option<Duration>,
     /// Whether activation failures that waiting could heal are retried, or should be fatal.
     pub retry_activation_failures: bool,
+    /// Compress file output with zstd. Ignored for [Output::Stdout].
+    pub compress: bool,
 }
 
 impl SinkConfig {
@@ -81,6 +47,42 @@ impl SinkConfig {
             flush_interval: Some(Duration::from_secs(5)),
             sync_interval: Some(Duration::from_secs(5 * 60)),
             retry_activation_failures: false,
+            compress: false,
+        }
+    }
+
+    /// Construct the writer stack for this config's output.
+    fn open_writer(&self, timestamp: Timestamp) -> std::io::Result<Box<dyn Writer>> {
+        let path = match &self.output {
+            // Unreachable with compression: main rejects --compress against stdout.
+            Output::Stdout => return Ok(Box::new(StdoutWriter::new())),
+            Output::Path(path) => path.clone(),
+            Output::Template {
+                dir,
+                interface,
+                ext,
+                next_index,
+            } => dir.join(template_filename(
+                *next_index,
+                interface,
+                timestamp.sec,
+                ext,
+            )),
+        };
+        if let Some(parent) = path.parent()
+            && !parent.as_os_str().is_empty()
+        {
+            std::fs::create_dir_all(parent)?;
+        }
+        let file = FileWriter::new(std::fs::File::create(&path)?);
+        tracing::info!(path = %path.display(), "created log file");
+        if self.compress {
+            Ok(Box::new(ZstdWriter::new(file)?))
+        } else {
+            Ok(Box::new(std::io::BufWriter::with_capacity(
+                self.flush_threshold_bytes,
+                file,
+            )))
         }
     }
 }
@@ -151,11 +153,7 @@ impl Sink {
             self.state = SinkState::Pending {
                 last_attempt: Some(Instant::now()),
             };
-            let mut writer = match self
-                .config
-                .output
-                .open(timestamp, self.config.flush_threshold_bytes)
-            {
+            let mut writer = match self.config.open_writer(timestamp) {
                 Ok(writer) => writer,
                 Err(e) => {
                     if classify(e.kind(), self.config.retry_activation_failures)

candumpr/src/writer/mod.rs

@@ -1,10 +1,12 @@
 mod file;
 mod stdout;
+mod zstd;
 
 use std::io::Write;
 
 pub use file::FileWriter;
 pub use stdout::StdoutWriter;
+pub use zstd::ZstdWriter;
 
 /// Writes formatted frame data to an output destination.
 pub trait Writer: std::io::Write {

candumpr/src/writer/zstd.rs

@@ -0,0 +1,184 @@
+use std::io::Write;
+
+use zstd::stream::raw::CParameter;
+
+use super::Writer;
+
+// TODO: Tune these given real sample data
+const LEVEL: i32 = 1;
+const WINDOW_LOG: u32 = 15;
+
+/// Compresses formatted frames with streaming zstd, one zstd frame per file.
+pub struct ZstdWriter<W: Writer> {
+    // Write to a Vec<u8> instead of the inner Writer so that we can always write chunks of
+    // compressed data evenly divisible by CanFrames (no partial frames are written). We do this by
+    // consuming from the Vec only upon flush(), which is only ever called at a Frame boundary.
+    encoder: zstd::stream::write::Encoder<'static, Vec<u8>>,
+    inner: W,
+}
+
+impl<W: Writer> ZstdWriter<W> {
+    pub fn new(inner: W) -> std::io::Result<Self> {
+        Self::with_params(inner, LEVEL, WINDOW_LOG)
+    }
+
+    pub fn with_params(inner: W, level: i32, window_log: u32) -> std::io::Result<Self> {
+        let mut encoder = zstd::stream::raw::Encoder::new(level)?;
+        encoder.set_parameter(CParameter::WindowLog(window_log))?;
+        Ok(Self {
+            encoder: zstd::stream::write::Encoder::with_encoder(Vec::new(), encoder),
+            inner,
+        })
+    }
+}
+
+impl<W: Writer> std::io::Write for ZstdWriter<W> {
+    fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
+        self.encoder.write(buf)
+    }
+
+    /// Close the current block and flush
+    ///
+    /// The [ZstdWriter] does not write to disk outside of [ZstdWriter::flush] in order to guarantee
+    /// that a CAN frame is never partially written across two compressed chunks.
+    fn flush(&mut self) -> std::io::Result<()> {
+        self.encoder.flush()?;
+        let buf = self.encoder.get_mut();
+        // We can't use Write::write_all, because in the case of a partial write it does not
+        // indicate how much was written. So we mirror the implementation of BufWriter::flush_buf
+        // which lets us retain unwritten bytes and retry. If we dropped bytes on write failure,
+        // even if we started writing bytes again afterwards, everything after that point would be
+        // corrupted.
+        let mut written = 0;
+        let result = loop {
+            if written == buf.len() {
+                break Ok(());
+            }
+            match self.inner.write(&buf[written..]) {
+                Ok(0) => {
+                    break Err(std::io::Error::new(
+                        std::io::ErrorKind::WriteZero,
+                        "failed to write compressed data",
+                    ));
+                }
+                Ok(n) => written += n,
+                Err(e) if e.kind() == std::io::ErrorKind::Interrupted => {}
+                Err(e) => break Err(e),
+            }
+        };
+        buf.drain(..written);
+        result?;
+        self.inner.flush()
+    }
+}
+
+impl<W: Writer> Writer for ZstdWriter<W> {
+    fn sync(&mut self) -> std::io::Result<()> {
+        self.flush()?;
+        self.inner.sync()
+    }
+
+    /// Finish writing.
+    ///
+    /// Flushes internal buffers, and writes the zstd frame epilogue. Do not write additional data
+    /// after calling finish().
+    fn finish(&mut self) -> std::io::Result<()> {
+        self.encoder.do_finish()?;
+        self.flush()?;
+        self.inner.finish()
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::recv::receiver::BATCH_CAPACITY;
+    use crate::writer::FileWriter;
+
+    /// Formatted candump frames for `range`, grouped the way a [Sink](crate::sink::Sink) writes
+    /// them: one `Vec` per receive batch.
+    fn batches(range: std::ops::Range<usize>) -> Vec<Vec<u8>> {
+        range
+            .map(|i| format!("(1732117385.{:06}) vcan0 123#{i:08X}\n", i % 1_000_000))
+            .collect::<Vec<_>>()
+            .chunks(BATCH_CAPACITY)
+            .map(|batch| batch.concat().into_bytes())
+            .collect()
+    }
+
+    /// The inner [Writer] under a [ZstdWriter] under test.
+    ///
+    /// Keeps what it accepts, and counts `write` calls, which is how many times compressed bytes
+    /// left the [ZstdWriter].
+    #[derive(Default)]
+    struct WriteCountingWriter {
+        bytes: Vec<u8>,
+        writes: usize,
+        /// When set, the first `write` accepts only this many bytes and the second fails with EIO
+        flaky: Option<usize>,
+    }
+
+    impl std::io::Write for WriteCountingWriter {
+        fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
+            self.writes += 1;
+            let accepted = match (self.flaky, self.writes) {
+                (Some(n), 1) => n.min(buf.len()),
+                (Some(_), 2) => return Err(std::io::Error::from_raw_os_error(libc::EIO)),
+                _ => buf.len(),
+            };
+            self.bytes.extend_from_slice(&buf[..accepted]);
+            Ok(accepted)
+        }
+
+        fn flush(&mut self) -> std::io::Result<()> {
+            Ok(())
+        }
+    }
+
+    impl Writer for WriteCountingWriter {
+        fn sync(&mut self) -> std::io::Result<()> {
+            Ok(())
+        }
+
+        fn finish(&mut self) -> std::io::Result<()> {
+            Ok(())
+        }
+    }
+
+    // happy path
+    #[test]
+    fn round_trip_decodes_to_input() {
+        let file = tempfile::NamedTempFile::new().unwrap();
+        let batches = batches(0..1000);
+
+        let mut writer = ZstdWriter::new(FileWriter::new(file.reopen().unwrap())).unwrap();
+        for batch in &batches {
+            writer.write_all(batch).unwrap();
+        }
+        writer.finish().unwrap();
+
+        let raw = std::fs::read(file.path()).unwrap();
+        assert_eq!(zstd::decode_all(&raw[..]).unwrap(), batches.concat());
+    }
+
+    // recovery after write errors shouldn't result in something that can't decompress
+    #[test]
+    fn partial_write_retains_the_remainder() {
+        let mut writer = ZstdWriter::new(WriteCountingWriter {
+            flaky: Some(100),
+            ..WriteCountingWriter::default()
+        })
+        .unwrap();
+        let input = batches(0..1000).concat();
+        writer.write_all(&input).unwrap();
+
+        let err = writer.flush().unwrap_err();
+        assert_eq!(err.raw_os_error(), Some(libc::EIO));
+        assert_eq!(writer.inner.bytes.len(), 100, "kept only what was accepted");
+
+        // Retrying recovers everything: the bytes the inner writer refused are still buffered.
+        // Dropping them instead would leave a hole that costs every frame after it.
+        writer.finish().unwrap();
+        assert_eq!(zstd::decode_all(&writer.inner.bytes[..]).unwrap(), input);
+    }
+}