fixup! Abstract compression function, so we can instrument input and output

Author: kim

crates/commitlog/src/repo/fs.rs

@@ -4,12 +4,11 @@ use std::io;
 use std::sync::Arc;
 
 use log::{debug, warn};
-use spacetimedb_fs_utils::compression::{CompressReader, CompressionStats};
+use spacetimedb_fs_utils::compression::{compress_with_zstd, 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";
@@ -261,15 +260,18 @@ impl Repo for Fs {
         fs::remove_file(self.segment_path(offset))
     }
 
-    fn compress_segment_with(&self, offset: u64, f: impl CompressOnce) -> io::Result<CompressionStats> {
+    fn compress_segment(&self, offset: u64) -> 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)?;
-        let stats = f.compress(&mut src, &mut dst)?;
+        // 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))?;
         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_with(&self, _: u64, _: impl super::CompressOnce) -> io::Result<CompressionStats> {
+    fn compress_segment(&self, _offset: u64) -> io::Result<CompressionStats> {
         Ok(<_>::default())
     }
 

crates/commitlog/src/repo/mod.rs

@@ -4,8 +4,8 @@ use std::{
 };
 
 use log::{debug, warn};
-use spacetimedb_fs_utils::compression::Zstd;
-pub use spacetimedb_fs_utils::compression::{CompressOnce, CompressionStats};
+
+pub use spacetimedb_fs_utils::compression::CompressionStats;
 
 use crate::{
     commit::Commit,
@@ -123,13 +123,7 @@ 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> {
-        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>;
+    fn compress_segment(&self, offset: u64) -> io::Result<CompressionStats>;
 
     /// Traverse all segments in this repository and return list of their
     /// offsets, sorted in ascending order.
@@ -172,8 +166,8 @@ impl<T: Repo> Repo for &T {
         T::remove_segment(self, offset)
     }
 
-    fn compress_segment_with(&self, offset: u64, f: impl CompressOnce) -> io::Result<CompressionStats> {
-        T::compress_segment_with(self, offset, f)
+    fn compress_segment(&self, offset: u64) -> io::Result<CompressionStats> {
+        T::compress_segment(self, offset)
     }
 
     fn existing_offsets(&self) -> io::Result<Vec<u64>> {
@@ -362,17 +356,6 @@ 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_with(&self, offset: u64, f: impl repo::CompressOnce) -> io::Result<CompressionStats> {
-        self.inner.compress_segment_with(offset, f)
+    fn compress_segment(&self, offset: u64) -> io::Result<CompressionStats> {
+        self.inner.compress_segment(offset)
     }
 
     fn existing_offsets(&self) -> io::Result<Vec<u64>> {

crates/fs-utils/src/compression.rs

@@ -150,27 +150,6 @@ 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,