feat(memtrack): add zstd compression support

Author: not-matthias

Cargo.lock

@@ -3,10 +3,11 @@ use clap::Parser;
 use ipc_channel::ipc::{self};
 use log::{debug, info};
 use memtrack::{MemtrackIpcMessage, Tracker, handle_ipc_message};
-use runner_shared::artifacts::{ArtifactExt, MemtrackArtifact, MemtrackEvent};
-use std::path::PathBuf;
+use runner_shared::artifacts::{ArtifactExt, MemtrackArtifact, MemtrackEvent, MemtrackWriter};
+use std::path::{Path, PathBuf};
 use std::process::Command;
 use std::sync::atomic::{AtomicBool, Ordering};
+use std::sync::mpsc::channel;
 use std::sync::{Arc, Mutex};
 use std::thread;
 
@@ -51,10 +52,8 @@ fn main() -> Result<()> {
         } => {
             debug!("Starting memtrack for command: {command}");
 
-            let (root_pid, events, status) =
-                track_command(&command, ipc_server).context("Failed to track command")?;
-            let result = MemtrackArtifact { events };
-            result.save_with_pid_to(&out_dir, root_pid as libc::pid_t)?;
+            let status =
+                track_command(&command, ipc_server, &out_dir).context("Failed to track command")?;
 
             std::process::exit(status.code().unwrap_or(1));
         }
@@ -64,7 +63,8 @@ fn main() -> Result<()> {
 fn track_command(
     cmd_string: &str,
     ipc_server_name: Option<String>,
-) -> anyhow::Result<(u32, Vec<MemtrackEvent>, std::process::ExitStatus)> {
+    out_dir: &Path,
+) -> anyhow::Result<std::process::ExitStatus> {
     let tracker = Tracker::new()?;
 
     let tracker_arc = Arc::new(Mutex::new(tracker));
@@ -95,37 +95,67 @@ fn track_command(
     let event_rx = { tracker_arc.lock().unwrap().track(root_pid)? };
     info!("Spawned child with pid {root_pid}");
 
-    // Spawn event processing thread
-    let process_events = Arc::new(AtomicBool::new(true));
-    let process_events_clone = process_events.clone();
-    let processing_thread = thread::spawn(move || {
-        let mut events = Vec::new();
+    // Create the artifact, then drain and write the events to disk
+    let file_name = MemtrackArtifact::file_name(Some(root_pid));
+    let out_file = std::fs::File::create(out_dir.join(file_name))?;
+
+    let (write_tx, write_rx) = channel::<MemtrackEvent>();
+
+    // Stage A: Fast drain thread - This is required so that we immediately clear the ring buffer
+    // because it only has a limited size.
+    static DRAIN_EVENTS: AtomicBool = AtomicBool::new(true);
+    let write_tx_clone = write_tx.clone();
+    let drain_thread = thread::spawn(move || {
         loop {
-            if !process_events_clone.load(Ordering::Relaxed) {
+            if !DRAIN_EVENTS.load(Ordering::Relaxed) {
                 break;
             }
-
             let Ok(event) = event_rx.try_recv() else {
                 continue;
             };
+            let _ = write_tx_clone.send(event.into());
+        }
+    });
+
+    // Stage B: Writer thread - Immediately writes the events to disk
+    let writer_thread = thread::spawn(move || -> anyhow::Result<()> {
+        let mut writer = MemtrackWriter::new(out_file)?;
 
-            events.push(event.into());
+        while let Ok(first) = write_rx.recv() {
+            writer.write_event(&first)?;
+
+            // Drain any backlog in a tight loop (batching)
+            while let Ok(ev) = write_rx.try_recv() {
+                writer.write_event(&ev)?;
+            }
         }
-        events
+        writer.finish()?;
+
+        Ok(())
     });
 
     // Wait for the command to complete
     let status = child.wait().context("Failed to wait for command")?;
     info!("Command exited with status: {status}");
 
-    info!("Waiting for the event processing thread to finish");
-    process_events.store(false, Ordering::Relaxed);
-    let events = processing_thread
+    // Close the write channel to signal threads to finish
+    DRAIN_EVENTS.store(false, Ordering::Relaxed);
+    drop(write_tx);
+
+    // Wait for drain thread to finish
+    info!("Waiting for the drain thread to finish");
+    drain_thread
+        .join()
+        .map_err(|_| anyhow::anyhow!("Failed to join drain thread"))?;
+
+    // Wait for writer thread to finish and propagate errors
+    info!("Waiting for the writer thread to finish");
+    writer_thread
         .join()
-        .map_err(|_| anyhow::anyhow!("Failed to join event thread"))?;
+        .map_err(|_| anyhow::anyhow!("Failed to join writer thread"))??;
 
     // IPC thread will exit when channel closes
     drop(ipc_handle);
 
-    Ok((root_pid as u32, events, status))
+    Ok(status)
 }

crates/memtrack/src/main.rs

@@ -13,3 +13,4 @@ log = { workspace = true }
 rmp = "0.8.14"
 rmp-serde = "1.3.0"
 libc = { workspace = true }
+zstd = "0.13"

crates/runner-shared/Cargo.toml

@@ -8,20 +8,22 @@ pub struct MemtrackArtifact {
 }
 impl super::ArtifactExt for MemtrackArtifact {
     fn encode_to_writer<W: Write>(&self, writer: W) -> anyhow::Result<()> {
-        // This is required for `decode_streamed`: We can't stream the deserialization of
-        // the whole artifact, so we have to encode them one by one.
-        let mut serializer = rmp_serde::Serializer::new(writer);
+        let mut writer = MemtrackWriter::new(writer)?;
         for event in &self.events {
-            event.serialize(&mut serializer)?;
+            writer.write_event(event)?;
         }
+        writer.finish()?;
         Ok(())
     }
 }
 
 impl MemtrackArtifact {
-    pub fn decode_streamed<R: std::io::Read>(reader: R) -> anyhow::Result<MemtrackEventStream<R>> {
+    pub fn decode_streamed<R: std::io::Read>(
+        reader: R,
+    ) -> anyhow::Result<MemtrackEventStream<zstd::Decoder<'static, std::io::BufReader<R>>>> {
+        let decoder = zstd::Decoder::new(reader)?;
         Ok(MemtrackEventStream {
-            deserializer: rmp_serde::Deserializer::new(reader),
+            deserializer: rmp_serde::Deserializer::new(decoder),
         })
     }
 }
@@ -61,6 +63,33 @@ impl<R: Read> Iterator for MemtrackEventStream<R> {
     }
 }
 
+/// Streaming writer for memtrack events with compression
+pub struct MemtrackWriter<W: Write> {
+    serializer: rmp_serde::Serializer<zstd::Encoder<'static, W>>,
+}
+
+impl<W: Write> MemtrackWriter<W> {
+    pub fn new(writer: W) -> anyhow::Result<Self> {
+        let encoder = zstd::Encoder::new(writer, 1)?;
+        Ok(Self {
+            serializer: rmp_serde::Serializer::new(encoder),
+        })
+    }
+
+    /// Write a single event to the stream
+    pub fn write_event(&mut self, event: &MemtrackEvent) -> anyhow::Result<()> {
+        event.serialize(&mut self.serializer)?;
+        Ok(())
+    }
+
+    /// Finish writing and flush the compression stream
+    pub fn finish(self) -> anyhow::Result<()> {
+        let encoder = self.serializer.into_inner();
+        encoder.finish()?;
+        Ok(())
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use crate::artifacts::ArtifactExt;
@@ -92,6 +121,7 @@ mod tests {
         };
         let mut buf = Vec::new();
         artifact.encode_to_writer(&mut buf)?;
+        std::fs::write("test.bin", &buf)?;
 
         let stream = MemtrackArtifact::decode_streamed(Cursor::new(buf))?;
         let collected: Vec<_> = stream.collect();

crates/runner-shared/src/artifacts/memtrack.rs

@@ -17,6 +17,14 @@ where
         std::any::type_name::<Self>().rsplit("::").next().unwrap()
     }
 
+    fn file_name(pid: Option<pid_t>) -> String {
+        if let Some(pid) = pid {
+            format!("{pid}.{}.msgpack", Self::name())
+        } else {
+            format!("{}.msgpack", Self::name())
+        }
+    }
+
     fn encode_to_writer<W: std::io::Write>(&self, mut writer: W) -> anyhow::Result<()> {
         let encoded = rmp_serde::to_vec_named(self)?;
         writer.write_all(&encoded)?;
@@ -37,14 +45,14 @@ where
     }
 
     fn save_to<P: AsRef<std::path::Path>>(&self, folder: P) -> anyhow::Result<()> {
-        self.save_file_to(folder, &format!("{}.msgpack", Self::name()))
+        self.save_file_to(folder, &Self::file_name(None))
     }
 
     fn save_with_pid_to<P: AsRef<std::path::Path>>(
         &self,
         folder: P,
         pid: pid_t,
     ) -> anyhow::Result<()> {
-        self.save_file_to(folder, &format!("{pid}.{}.msgpack", Self::name()))
+        self.save_file_to(folder, &Self::file_name(Some(pid)))
     }
 }