Abstract compression function, so we can instrument input and output

Author: kim

crates/commitlog/src/repo/fs.rs

@@ -4,11 +4,12 @@ use std::io;
 use std::sync::Arc;
 
 use log::{debug, warn};
-use spacetimedb_fs_utils::compression::{compress_with_zstd, CompressReader, CompressionStats};
+use spacetimedb_fs_utils::compression::{CompressReader, CompressionStats};
 use spacetimedb_paths::server::{CommitLogDir, SegmentFile};
 use tempfile::NamedTempFile;
 
 use super::{Repo, SegmentLen, SegmentReader, TxOffset, TxOffsetIndex, TxOffsetIndexMut};
+use crate::repo::CompressOnce;
 use crate::segment::FileLike;
 
 const SEGMENT_FILE_EXT: &str = ".stdb.log";
@@ -260,18 +261,15 @@ impl Repo for Fs {
         fs::remove_file(self.segment_path(offset))
     }
 
-    fn compress_segment(&self, offset: u64) -> io::Result<CompressionStats> {
+    fn compress_segment_with(&self, offset: u64, f: impl CompressOnce) -> io::Result<CompressionStats> {
         let src = self.open_segment_reader(offset)?;
         // if it's already compressed, leave it be
         let CompressReader::None(mut src) = src.inner else {
             return Ok(<_>::default());
         };
 
         let mut dst = NamedTempFile::new_in(&self.root)?;
-        // bytes per frame. in the future, it might be worth looking into putting
-        // every commit into its own frame, to make seeking more efficient.
-        let max_frame_size = 0x1000;
-        let stats = compress_with_zstd(&mut src, &mut dst, Some(max_frame_size))?;
+        let stats = f.compress(&mut src, &mut dst)?;
         dst.persist(self.segment_path(offset))?;
 
         Ok(stats)

crates/commitlog/src/repo/mem.rs

@@ -104,7 +104,7 @@ impl Repo for Memory {
         Ok(())
     }
 
-    fn compress_segment(&self, _offset: u64) -> io::Result<CompressionStats> {
+    fn compress_segment_with(&self, _: u64, _: impl super::CompressOnce) -> io::Result<CompressionStats> {
         Ok(<_>::default())
     }
 

crates/commitlog/src/repo/mod.rs

@@ -4,8 +4,8 @@ use std::{
 };
 
 use log::{debug, warn};
-
-pub use spacetimedb_fs_utils::compression::CompressionStats;
+use spacetimedb_fs_utils::compression::Zstd;
+pub use spacetimedb_fs_utils::compression::{CompressOnce, CompressionStats};
 
 use crate::{
     commit::Commit,
@@ -123,7 +123,13 @@ pub trait Repo: Clone + fmt::Display {
     fn remove_segment(&self, offset: u64) -> io::Result<()>;
 
     /// Compress a segment in storage, marking it as immutable.
-    fn compress_segment(&self, offset: u64) -> io::Result<CompressionStats>;
+    fn compress_segment(&self, offset: u64) -> io::Result<CompressionStats> {
+        self.compress_segment_with(offset, segment_compressor())
+    }
+
+    /// Compress a segment using a supplied [CompressOnce], marking it as
+    /// immutable.
+    fn compress_segment_with(&self, offset: u64, f: impl CompressOnce) -> io::Result<CompressionStats>;
 
     /// Traverse all segments in this repository and return list of their
     /// offsets, sorted in ascending order.
@@ -166,8 +172,8 @@ impl<T: Repo> Repo for &T {
         T::remove_segment(self, offset)
     }
 
-    fn compress_segment(&self, offset: u64) -> io::Result<CompressionStats> {
-        T::compress_segment(self, offset)
+    fn compress_segment_with(&self, offset: u64, f: impl CompressOnce) -> io::Result<CompressionStats> {
+        T::compress_segment_with(self, offset, f)
     }
 
     fn existing_offsets(&self) -> io::Result<Vec<u64>> {
@@ -356,6 +362,17 @@ pub fn open_segment_reader<R: Repo>(
         .map_err(|source| with_segment_context("reading segment header", repo, offset, source))
 }
 
+/// Obtain the canonical [CompressOnce] compressor for segments.
+///
+/// The compressor will create seekable [Zstd] archives with a max frame size
+/// of 4KiB. That is, seeking to an arbitrary byte offset (of the uncompressed
+/// segment) within the archive will decompress 4KiB of data on average.
+pub fn segment_compressor() -> Zstd {
+    Zstd {
+        max_frame_size: Some(0x1000),
+    }
+}
+
 fn segment_label<R: Repo>(repo: &R, offset: u64) -> String {
     repo.segment_file_path(offset)
         .unwrap_or_else(|| format!("offset {offset}"))

crates/commitlog/src/tests/partial.rs

@@ -296,8 +296,8 @@ impl Repo for ShortMem {
         self.inner.remove_segment(offset)
     }
 
-    fn compress_segment(&self, offset: u64) -> io::Result<CompressionStats> {
-        self.inner.compress_segment(offset)
+    fn compress_segment_with(&self, offset: u64, f: impl repo::CompressOnce) -> io::Result<CompressionStats> {
+        self.inner.compress_segment_with(offset, f)
     }
 
     fn existing_offsets(&self) -> io::Result<Vec<u64>> {

crates/fs-utils/src/compression.rs

@@ -150,6 +150,27 @@ pub fn new_zstd_writer<'a, W: io::Write>(inner: W, max_frame_size: Option<u32>)
     .build()
 }
 
+/// Compress `src` to `dst` in one go.
+///
+/// Like a `FnOnce` closure, but polymorphic over the arguments.
+pub trait CompressOnce {
+    fn compress(self, src: impl io::Read, dst: impl io::Write) -> io::Result<CompressionStats>;
+}
+
+/// Implements [CompressOnce] for the zstd compression algorithm.
+pub struct Zstd {
+    /// If `Some`, add a seek table with `max_frame_size` to the compressed output.
+    ///
+    /// See [zstd_framed::writer::ZstdWriterBuilder::with_seek_table].
+    pub max_frame_size: Option<u32>,
+}
+
+impl CompressOnce for Zstd {
+    fn compress(self, src: impl io::Read, dst: impl io::Write) -> io::Result<CompressionStats> {
+        compress_with_zstd(src, dst, self.max_frame_size)
+    }
+}
+
 #[derive(Clone, Copy, Debug, Default)]
 pub struct CompressionStats {
     pub bytes_read: u64,