Merge #301

301: Snapshot merging r=MarinPostma a=MarinPostma

This PR implements snapshot merging. When the size taken by snapshots grow larger than twice the size of the database, or when there are more than 32 snapshot files, then all snapshots are merged together.


Co-authored-by: ad hoc <[email protected]>

Author: bors[bot]

sqld/src/replication/client.rs

@@ -529,7 +529,7 @@ impl ReadReplicationHook {
                         }
                         Err(e) => return Err(e.into()),
                     };
-                    let frame = WalFrame::decode(raw_frame.data)?;
+                    let frame = WalFrame::try_from_bytes(raw_frame.data)?;
                     debug_assert_eq!(
                         Some(frame.header.frame_id),
                         current_offset

sqld/src/replication/log_compaction.rs

@@ -10,9 +10,7 @@ use anyhow::{bail, ensure};
 use bytemuck::{bytes_of, pod_read_unaligned, try_from_bytes, Pod, Zeroable};
 use bytes::{Bytes, BytesMut};
 use crc::Crc;
-use once_cell::sync::Lazy;
 use parking_lot::RwLock;
-use regex::Regex;
 use rusqlite::ffi::SQLITE_ERROR;
 use uuid::Uuid;
 
@@ -22,8 +20,7 @@ use crate::libsql::ffi::{
 };
 use crate::libsql::{ffi::PageHdrIter, wal_hook::WalHook};
 
-use super::log_compaction::LogCompactor;
-use super::snapshot::SnapshotFile;
+use super::snapshot::{find_snapshot_file, LogCompactor, SnapshotFile};
 
 pub const WAL_PAGE_SIZE: i32 = 4096;
 pub const WAL_MAGIC: u64 = u64::from_le_bytes(*b"SQLDWAL\0");
@@ -99,7 +96,7 @@ unsafe impl WalHook for ReplicationLoggerHook {
                 .maybe_compact(
                     self.logger.compactor.clone(),
                     ntruncate,
-                    &self.logger.path,
+                    &self.logger.db_path,
                     self.logger.current_checksum.load(Ordering::Relaxed),
                 )
                 .unwrap();
@@ -136,6 +133,24 @@ pub struct WalFrame {
     pub data: Bytes,
 }
 
+impl WalFrame {
+    /// size of a single frame
+    pub const SIZE: usize = size_of::<FrameHeader>() + WAL_PAGE_SIZE as usize;
+
+    pub fn try_from_bytes(mut data: Bytes) -> anyhow::Result<Self> {
+        let header_bytes = data.split_to(size_of::<FrameHeader>());
+        ensure!(
+            data.len() == WAL_PAGE_SIZE as usize,
+            "invalid frame size, expected: {}, found: {}",
+            WAL_PAGE_SIZE,
+            data.len()
+        );
+        let header = FrameHeader::decode(&header_bytes)?;
+
+        Ok(Self { header, data })
+    }
+}
+
 pub struct WalFrameBorrowed<'a> {
     pub header: FrameHeader,
     pub data: &'a [u8],
@@ -150,20 +165,6 @@ impl<'a> From<&'a WalFrame> for WalFrameBorrowed<'a> {
     }
 }
 
-impl WalFrame {
-    pub fn decode(mut data: Bytes) -> anyhow::Result<Self> {
-        let header_bytes = data.split_to(size_of::<FrameHeader>());
-        ensure!(
-            data.len() == WAL_PAGE_SIZE as usize,
-            "invalid frame size, expected: {}, found: {}",
-            WAL_PAGE_SIZE,
-            data.len()
-        );
-        let header = FrameHeader::decode(&header_bytes)?;
-        Ok(Self { header, data })
-    }
-}
-
 impl ReplicationLoggerHook {
     pub fn new(logger: Arc<ReplicationLogger>) -> Self {
         Self {
@@ -519,12 +520,12 @@ pub struct ReplicationLogger {
     pub generation: Generation,
     pub log_file: RwLock<LogFile>,
     compactor: LogCompactor,
-    path: PathBuf,
+    db_path: PathBuf,
 }
 
 impl ReplicationLogger {
-    pub fn open(path: impl AsRef<Path>, max_log_size: u64) -> anyhow::Result<Self> {
-        let log_path = path.as_ref().join("wallog");
+    pub fn open(db_path: &Path, max_log_size: u64) -> anyhow::Result<Self> {
+        let log_path = db_path.join("wallog");
         let file = OpenOptions::new()
             .create(true)
             .write(true)
@@ -544,9 +545,9 @@ impl ReplicationLogger {
         Ok(Self {
             current_checksum,
             generation: Generation::new(generation_start_frame_id),
+            compactor: LogCompactor::new(db_path, log_file.header.db_id)?,
             log_file: RwLock::new(log_file),
-            compactor: LogCompactor::new(path.as_ref().to_owned()),
-            path: path.as_ref().to_owned(),
+            db_path: db_path.to_owned(),
         })
     }
 
@@ -616,43 +617,8 @@ impl ReplicationLogger {
             .store(new_current_checksum, Ordering::Relaxed);
     }
 
-    pub async fn locate_snapshot(
-        &self,
-        start_idx: FrameId,
-    ) -> anyhow::Result<Option<SnapshotFile>> {
-        static SNAPSHOT_FILE_MATCHER: Lazy<Regex> = Lazy::new(|| {
-            Regex::new(
-                r#"(?x)
-                # match database id
-                (\w{8}-\w{4}-\w{4}-\w{4}-\w{12})-
-                # match start frame id
-                (\d*)-
-                # match end frame-id
-                (\d*).snap"#,
-            )
-            .unwrap()
-        });
-
-        let snapshots_dir = self.path.join("snapshots");
-        let mut entries = tokio::fs::read_dir(snapshots_dir).await?;
-
-        while let Some(entry) = entries.next_entry().await? {
-            let path = entry.path();
-            let Some(name) = path.file_name() else {continue;};
-            let Some(name_str) = name.to_str() else { continue;};
-            let Some(captures) = SNAPSHOT_FILE_MATCHER.captures(name_str) else {continue;};
-            let _db_id = captures.get(1).unwrap();
-            let start_index: u64 = captures.get(2).unwrap().as_str().parse().unwrap();
-            let end_index: u64 = captures.get(3).unwrap().as_str().parse().unwrap();
-            // we're looking for the frame right after the last applied frame on the replica
-            if (start_index..=end_index).contains(&(start_idx + 1)) {
-                tracing::debug!("found snapshot for frame {start_idx} at {path:?}");
-                let snapshot_file = SnapshotFile::open(&path)?;
-                return Ok(Some(snapshot_file));
-            }
-        }
-
-        Ok(None)
+    pub fn get_snapshot_file(&self, from: FrameId) -> anyhow::Result<Option<SnapshotFile>> {
+        find_snapshot_file(&self.db_path, from)
     }
 }
 
@@ -677,7 +643,7 @@ mod test {
 
         let log_file = logger.log_file.write();
         for i in 0..10 {
-            let frame = WalFrame::decode(log_file.frame_bytes(i).unwrap()).unwrap();
+            let frame = WalFrame::try_from_bytes(log_file.frame_bytes(i).unwrap()).unwrap();
             assert_eq!(frame.header.page_no, i as u32);
             assert!(frame.data.iter().all(|x| i as u8 == *x));
         }