Merge origin/master into joshua/remove-old-query-code

Resolve import conflict in relational_db.rs (take master's new test utilities).

Author: clockwork-labs-bot

crates/commitlog/src/repo/fs.rs

@@ -10,7 +10,7 @@ use tempfile::NamedTempFile;
 
 use crate::segment::FileLike;
 
-use super::{Repo, SegmentLen, TxOffset, TxOffsetIndex, TxOffsetIndexMut};
+use super::{Repo, SegmentLen, SegmentReader, TxOffset, TxOffsetIndex, TxOffsetIndexMut};
 
 const SEGMENT_FILE_EXT: &str = ".stdb.log";
 
@@ -154,9 +154,52 @@ impl FileLike for NamedTempFile {
     }
 }
 
+/// A file-backed, read-only segment.
+///
+/// Transparently handles reading compressed segments.
+/// [Self::sealed] returns `true` if the segment is compressed.
+pub struct ReadOnlySegment {
+    inner: CompressReader,
+}
+
+impl SegmentReader for ReadOnlySegment {
+    #[inline]
+    fn sealed(&self) -> bool {
+        self.inner.is_compressed()
+    }
+}
+
+impl io::Read for ReadOnlySegment {
+    #[inline]
+    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
+        self.inner.read(buf)
+    }
+}
+
+impl io::BufRead for ReadOnlySegment {
+    #[inline]
+    fn fill_buf(&mut self) -> io::Result<&[u8]> {
+        self.inner.fill_buf()
+    }
+
+    #[inline]
+    fn consume(&mut self, amount: usize) {
+        self.inner.consume(amount);
+    }
+}
+
+impl io::Seek for ReadOnlySegment {
+    #[inline]
+    fn seek(&mut self, pos: io::SeekFrom) -> io::Result<u64> {
+        self.inner.seek(pos)
+    }
+}
+
+impl SegmentLen for ReadOnlySegment {}
+
 impl Repo for Fs {
     type SegmentWriter = File;
-    type SegmentReader = CompressReader;
+    type SegmentReader = ReadOnlySegment;
 
     fn create_segment(&self, offset: u64) -> io::Result<Self::SegmentWriter> {
         File::options()
@@ -198,8 +241,9 @@ impl Repo for Fs {
     }
 
     fn open_segment_reader(&self, offset: u64) -> io::Result<Self::SegmentReader> {
+        debug!("fs: open segment at {}", self.segment_path(offset).display());
         let file = File::open(self.segment_path(offset))?;
-        CompressReader::new(file)
+        CompressReader::new(file).map(|inner| ReadOnlySegment { inner })
     }
 
     fn remove_segment(&self, offset: u64) -> io::Result<()> {
@@ -215,7 +259,7 @@ impl Repo for Fs {
     fn compress_segment(&self, offset: u64) -> io::Result<()> {
         let src = self.open_segment_reader(offset)?;
         // if it's already compressed, leave it be
-        let CompressReader::None(mut src) = src else {
+        let CompressReader::None(mut src) = src.inner else {
             return Ok(());
         };
 

crates/commitlog/src/repo/mem.rs

@@ -10,7 +10,7 @@ use crate::repo::{
 };
 
 mod segment;
-pub use segment::Segment;
+pub use segment::{ReadOnlySegment, Segment};
 
 pub const PAGE_SIZE: usize = 4096;
 
@@ -52,7 +52,7 @@ impl fmt::Display for Memory {
 
 impl Repo for Memory {
     type SegmentWriter = Segment;
-    type SegmentReader = io::BufReader<Segment>;
+    type SegmentReader = ReadOnlySegment;
 
     fn create_segment(&self, offset: u64) -> io::Result<Self::SegmentWriter> {
         let mut inner = self.segments.write().unwrap();
@@ -88,7 +88,7 @@ impl Repo for Memory {
     }
 
     fn open_segment_reader(&self, offset: u64) -> io::Result<Self::SegmentReader> {
-        self.open_segment_writer(offset).map(io::BufReader::new)
+        self.open_segment_writer(offset).map(Into::into)
     }
 
     fn remove_segment(&self, offset: u64) -> io::Result<()> {

crates/commitlog/src/repo/mem/segment.rs

@@ -7,7 +7,7 @@ use std::{
 use crate::{
     repo::{
         mem::{SpaceOnDevice, PAGE_SIZE},
-        SegmentLen,
+        SegmentLen, SegmentReader,
     },
     segment::FileLike,
 };
@@ -318,3 +318,50 @@ mod async_impls {
         }
     }
 }
+
+pub struct ReadOnlySegment {
+    inner: io::BufReader<Segment>,
+}
+
+impl From<Segment> for ReadOnlySegment {
+    fn from(inner: Segment) -> Self {
+        Self {
+            inner: io::BufReader::new(inner),
+        }
+    }
+}
+
+impl SegmentReader for ReadOnlySegment {
+    /// Memory segments dont' support compression, so are never sealed.
+    fn sealed(&self) -> bool {
+        false
+    }
+}
+
+impl io::Read for ReadOnlySegment {
+    #[inline]
+    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
+        self.inner.read(buf)
+    }
+}
+
+impl io::BufRead for ReadOnlySegment {
+    #[inline]
+    fn fill_buf(&mut self) -> io::Result<&[u8]> {
+        self.inner.fill_buf()
+    }
+
+    #[inline]
+    fn consume(&mut self, amount: usize) {
+        self.inner.consume(amount);
+    }
+}
+
+impl io::Seek for ReadOnlySegment {
+    #[inline]
+    fn seek(&mut self, pos: io::SeekFrom) -> io::Result<u64> {
+        self.inner.seek(pos)
+    }
+}
+
+impl SegmentLen for ReadOnlySegment {}

crates/commitlog/src/repo/mod.rs

@@ -1,4 +1,7 @@
-use std::{fmt, io};
+use std::{
+    fmt,
+    io::{self, Seek},
+};
 
 use log::{debug, warn};
 
@@ -52,8 +55,14 @@ pub trait SegmentLen: io::Seek {
     }
 }
 
-pub trait SegmentReader: io::BufRead + SegmentLen + Send + Sync {}
-impl<T: io::BufRead + SegmentLen + Send + Sync> SegmentReader for T {}
+pub trait SegmentReader: io::BufRead + SegmentLen + Send + Sync {
+    /// Whether the segment is considered immutable.
+    ///
+    /// Currently, this is true when the segment is compressed.
+    /// [resume_segment_writer] uses this method to indicate that a new segment
+    /// should be created when opening a commitlog.
+    fn sealed(&self) -> bool;
+}
 
 pub trait SegmentWriter: FileLike + io::Read + io::Write + SegmentLen + Send + Sync {}
 impl<T: FileLike + io::Read + io::Write + SegmentLen + Send + Sync> SegmentWriter for T {}
@@ -243,21 +252,9 @@ pub fn resume_segment_writer<R: Repo>(
     opts: Options,
     offset: u64,
 ) -> io::Result<Result<Writer<R::SegmentWriter>, Metadata>> {
-    let mut storage = repo.open_segment_writer(offset)?;
-    // Ensure we have enough space for this segment.
-    // The segment could have been created without the `fallocate` feature
-    // enabled, so we call this here again to ensure writes can't fail due to
-    // ENOSPC.
-    fallocate(&mut storage, &opts)?;
+    let mut reader = repo.open_segment_reader(offset)?;
     let offset_index = repo.get_offset_index(offset).ok();
-    let Metadata {
-        header,
-        tx_range,
-        size_in_bytes,
-        max_epoch,
-        max_commit_offset: _,
-        max_commit: _,
-    } = match Metadata::extract(offset, &mut storage, offset_index.as_ref()) {
+    let meta = match Metadata::extract(offset, &mut reader, offset_index.as_ref()) {
         Err(error::SegmentMetadata::InvalidCommit { sofar, source }) => {
             warn!("invalid commit in segment {offset}: {source}");
             debug!("sofar={sofar:?}");
@@ -266,34 +263,55 @@ pub fn resume_segment_writer<R: Repo>(
         Err(error::SegmentMetadata::Io(e)) => return Err(e),
         Ok(meta) => meta,
     };
-    header
+    meta.header
         .ensure_compatible(opts.log_format_version, Commit::CHECKSUM_ALGORITHM)
         .map_err(|msg| io::Error::new(io::ErrorKind::InvalidData, msg))?;
     // When resuming, the log format version must be equal.
-    if header.log_format_version != opts.log_format_version {
+    if meta.header.log_format_version != opts.log_format_version {
         return Err(io::Error::new(
             io::ErrorKind::InvalidData,
             format!(
                 "log format version mismatch: current={} segment={}",
-                opts.log_format_version, header.log_format_version
+                opts.log_format_version, meta.header.log_format_version
             ),
         ));
     }
 
-    Ok(Ok(Writer {
-        commit: Commit {
-            min_tx_offset: tx_range.end,
-            n: 0,
-            records: Vec::new(),
-            epoch: max_epoch,
-        },
-        inner: io::BufWriter::new(storage),
-
-        min_tx_offset: tx_range.start,
-        bytes_written: size_in_bytes,
-
-        offset_index_head: create_offset_index_writer(repo, offset, opts),
-    }))
+    if reader.sealed() {
+        Ok(Err(meta))
+    } else {
+        let Metadata {
+            header: _,
+            tx_range,
+            size_in_bytes,
+            max_epoch,
+            max_commit_offset: _,
+            max_commit: _,
+        } = meta;
+        let mut writer = repo.open_segment_writer(offset)?;
+        // Ensure we have enough space for this segment.
+        // The segment could have been created without the `fallocate` feature
+        // enabled, so we call this here again to ensure writes can't fail due
+        // to ENOSPC.
+        fallocate(&mut writer, &opts)?;
+        // We use `O_APPEND`, but make the file offset consistent regardless.
+        writer.seek(io::SeekFrom::End(0))?;
+
+        Ok(Ok(Writer {
+            commit: Commit {
+                min_tx_offset: tx_range.end,
+                n: 0,
+                records: Vec::new(),
+                epoch: max_epoch,
+            },
+            inner: io::BufWriter::new(writer),
+
+            min_tx_offset: tx_range.start,
+            bytes_written: size_in_bytes,
+
+            offset_index_head: create_offset_index_writer(repo, offset, opts),
+        }))
+    }
 }
 
 /// Open the existing segment at `offset` for reading.

crates/commitlog/src/tests/partial.rs

@@ -271,7 +271,7 @@ impl fmt::Display for ShortMem {
 
 impl Repo for ShortMem {
     type SegmentWriter = ShortSegment;
-    type SegmentReader = io::BufReader<repo::mem::Segment>;
+    type SegmentReader = repo::mem::ReadOnlySegment;
 
     fn create_segment(&self, offset: u64) -> io::Result<Self::SegmentWriter> {
         self.inner.create_segment(offset).map(|inner| ShortSegment {

crates/commitlog/tests/random_payload/mod.rs

@@ -1,6 +1,7 @@
 use log::info;
+use spacetimedb_commitlog::repo::Repo;
 use spacetimedb_commitlog::tests::helpers::enable_logging;
-use spacetimedb_commitlog::{payload, Commitlog, Options};
+use spacetimedb_commitlog::{commitlog, payload, repo, Commitlog, Options};
 use spacetimedb_paths::server::CommitLogDir;
 use spacetimedb_paths::FromPathUnchecked;
 use tempfile::tempdir;
@@ -75,6 +76,12 @@ fn resets() {
     }
 }
 
+/// Try to generate commitlogs that will be amenable to compression -
+/// random data doesn't compress well, so try and have there be repetition
+fn compressible_payloads() -> impl Iterator<Item = [u8; 256]> {
+    (0..4).map(|_| gen_payload()).cycle()
+}
+
 #[test]
 fn compression() {
     enable_logging();
@@ -90,9 +97,7 @@ fn compression() {
     )
     .unwrap();
 
-    // try to generate commitlogs that will be amenable to compression -
-    // random data doesn't compress well, so try and have there be repetition
-    let payloads = (0..4).map(|_| gen_payload()).cycle().take(1024).collect::<Vec<_>>();
+    let payloads = compressible_payloads().take(1024).collect::<Vec<_>>();
     for (i, payload) in payloads.iter().enumerate() {
         clog.commit([(i as u64, *payload)]).unwrap();
     }
@@ -114,3 +119,45 @@ fn compression() {
         .enumerate()
         .all(|(i, x)| x.offset == i as u64 && x.txdata == payloads[i]));
 }
+
+/// When restoring an archived commitlog, all segments are compressed and should
+/// remain immutable.
+///
+/// Tests that this is upheld, i.e. a fresh segment is created when resuming
+/// writes.
+#[test]
+fn all_segments_sealed() {
+    enable_logging();
+
+    let root = tempdir().unwrap();
+    let path = CommitLogDir::from_path_unchecked(root.path());
+    let opts = Options {
+        max_segment_size: 64 * 1024,
+        ..<_>::default()
+    };
+    let num_commits = 1024;
+    let repo = repo::Fs::new(path, None).unwrap();
+    {
+        let mut clog = commitlog::Generic::open(&repo, opts).unwrap();
+        for (i, payload) in compressible_payloads().take(num_commits).enumerate() {
+            clog.commit([(i as u64, payload)]).unwrap();
+        }
+        clog.flush().unwrap();
+        clog.sync();
+    }
+
+    let segments = repo.existing_offsets().unwrap();
+    let num_segments = segments.len();
+
+    // Compress all segments via the `repo`,
+    // to not trigger the assert that the head segment cannot be compressed.
+    for segment in segments {
+        repo.compress_segment(segment).unwrap();
+    }
+
+    // Re-opening the commitlog should create a fresh segment at offset `num_commits`.
+    let _ = commitlog::Generic::<_, [u8; 256]>::open(&repo, opts).unwrap();
+    let segments = repo.existing_offsets().unwrap();
+    assert_eq!(num_segments + 1, segments.len());
+    assert_eq!(segments.last().copied(), Some(num_commits as u64));
+}