Merge pull request #1613 from tursodatabase/current-segment-recovery

current segment recovery

Author: MarinPostma

libsql-wal/Cargo.toml

@@ -42,6 +42,7 @@ inquire = { version = "0.7.5", optional = true }
 tracing-subscriber = { version = "0.3.18", optional = true }
 aws-credential-types = { version = "1.2.0", optional = true }
 dashmap = "6.0.1"
+rand = "0.8.5"
 
 [dev-dependencies]
 criterion = "0.5.1"

libsql-wal/src/checkpointer.rs

@@ -89,6 +89,8 @@ pub struct Checkpointer<P> {
     join_set: JoinSet<(NamespaceName, crate::error::Result<()>)>,
     processing: Vec<NamespaceName>,
     errors: usize,
+    /// previous iteration of the loop resulted in no work being enqueued
+    no_work: bool,
 }
 
 #[allow(private_bounds)]
@@ -111,6 +113,7 @@ where
             join_set: JoinSet::new(),
             processing: Vec::new(),
             errors: 0,
+            no_work: false,
         }
     }
 
@@ -164,17 +167,19 @@ where
             }
             // don't wait if there is stuff to enqueue
             _ = std::future::ready(()), if !self.scheduled.is_empty()
-                && self.join_set.len() < self.max_checkpointing_conccurency => (),
+                && self.join_set.len() < self.max_checkpointing_conccurency && !self.no_work => (),
         }
 
         let n_available = self.max_checkpointing_conccurency - self.join_set.len();
         if n_available > 0 {
+            self.no_work = true;
             for namespace in self
                 .scheduled
                 .difference(&self.checkpointing)
                 .take(n_available)
                 .cloned()
             {
+                self.no_work = false;
                 self.processing.push(namespace.clone());
                 let perform_checkpoint = self.perform_checkpoint.clone();
                 self.join_set.spawn(async move {

libsql-wal/src/error.rs

@@ -11,6 +11,10 @@ pub enum Error {
     BusySnapshot,
     #[error("invalid segment header checksum")]
     InvalidHeaderChecksum,
+    #[error("invalid segment header magic")]
+    InvalidHeaderMagic,
+    #[error("invalid segment header version")]
+    InvalidHeaderVersion,
 }
 
 impl Into<libsql_sys::ffi::Error> for Error {

libsql-wal/src/io/file.rs

@@ -214,6 +214,10 @@ impl<W> BufCopy<W> {
         let Self { w, buf } = self;
         (w, buf)
     }
+
+    pub fn get_ref(&self) -> &W {
+        &self.w
+    }
 }
 
 impl<W: Write> Write for BufCopy<W> {

libsql-wal/src/io/mod.rs

@@ -3,6 +3,7 @@ use std::path::Path;
 use std::sync::Arc;
 
 use chrono::{DateTime, Utc};
+use rand::{rngs::ThreadRng, thread_rng, Rng};
 use uuid::Uuid;
 
 pub use self::file::FileExt;
@@ -14,6 +15,7 @@ pub mod file;
 pub trait Io: Send + Sync + 'static {
     type File: FileExt;
     type TempFile: FileExt;
+    type Rng: Rng;
 
     fn create_dir_all(&self, path: &Path) -> io::Result<()>;
     /// TODO: when adding an async variant make sure all places where async is needed are replaced
@@ -30,6 +32,9 @@ pub trait Io: Send + Sync + 'static {
     fn now(&self) -> DateTime<Utc>;
     fn uuid(&self) -> Uuid;
     fn hard_link(&self, src: &Path, dst: &Path) -> io::Result<()>;
+    fn with_rng<F, R>(&self, f: F) -> R
+    where
+        F: FnOnce(&mut Self::Rng) -> R;
 }
 
 #[derive(Default, Debug, Clone, Copy)]
@@ -38,6 +43,7 @@ pub struct StdIO(pub(crate) ());
 impl Io for StdIO {
     type File = std::fs::File;
     type TempFile = std::fs::File;
+    type Rng = ThreadRng;
 
     fn create_dir_all(&self, path: &Path) -> io::Result<()> {
         std::fs::create_dir_all(path)
@@ -72,11 +78,19 @@ impl Io for StdIO {
     fn hard_link(&self, src: &Path, dst: &Path) -> io::Result<()> {
         std::fs::hard_link(src, dst)
     }
+
+    fn with_rng<F, R>(&self, f: F) -> R
+    where
+        F: FnOnce(&mut Self::Rng) -> R,
+    {
+        f(&mut thread_rng())
+    }
 }
 
 impl<T: Io> Io for Arc<T> {
     type File = T::File;
     type TempFile = T::TempFile;
+    type Rng = T::Rng;
 
     fn create_dir_all(&self, path: &Path) -> io::Result<()> {
         self.as_ref().create_dir_all(path)
@@ -107,4 +121,41 @@ impl<T: Io> Io for Arc<T> {
     fn hard_link(&self, src: &Path, dst: &Path) -> io::Result<()> {
         self.as_ref().hard_link(src, dst)
     }
+
+    fn with_rng<F, R>(&self, f: F) -> R
+    where
+        F: FnOnce(&mut Self::Rng) -> R,
+    {
+        self.as_ref().with_rng(f)
+    }
+}
+
+pub struct Inspect<W, F> {
+    inner: W,
+    f: F,
+}
+
+impl<W, F> Inspect<W, F> {
+    pub fn new(inner: W, f: F) -> Self {
+        Self { inner, f }
+    }
+
+    pub(crate) fn into_inner(self) -> W {
+        self.inner
+    }
+}
+
+impl<W, F> io::Write for Inspect<W, F>
+where
+    W: io::Write,
+    for<'a> F: FnMut(&'a [u8]),
+{
+    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
+        (self.f)(buf);
+        self.inner.write(buf)
+    }
+
+    fn flush(&mut self) -> io::Result<()> {
+        self.inner.flush()
+    }
 }

libsql-wal/src/lib.rs

@@ -16,24 +16,27 @@ const LIBSQL_MAGIC: u64 = u64::from_be_bytes(*b"LIBSQL\0\0");
 #[cfg(any(debug_assertions, test))]
 pub mod test {
     use std::fs::OpenOptions;
+    use std::path::Path;
     use std::path::PathBuf;
-    use std::{path::Path, sync::Arc};
+    use std::sync::Arc;
 
-    use libsql_sys::{name::NamespaceName, rusqlite::OpenFlags};
+    use libsql_sys::name::NamespaceName;
+    use libsql_sys::rusqlite::OpenFlags;
     use tempfile::{tempdir, TempDir};
     use tokio::sync::mpsc;
 
     use crate::checkpointer::LibsqlCheckpointer;
+    use crate::io::Io;
     use crate::io::StdIO;
     use crate::registry::WalRegistry;
     use crate::shared_wal::SharedWal;
     use crate::storage::TestStorage;
     use crate::wal::{LibsqlWal, LibsqlWalManager};
 
-    pub struct TestEnv {
-        pub tmp: TempDir,
-        pub registry: Arc<WalRegistry<StdIO, TestStorage>>,
-        pub wal: LibsqlWalManager<StdIO, TestStorage>,
+    pub struct TestEnv<IO: Io = StdIO> {
+        pub tmp: Arc<TempDir>,
+        pub registry: Arc<WalRegistry<IO, TestStorage<IO>>>,
+        pub wal: LibsqlWalManager<IO, TestStorage<IO>>,
     }
 
     impl TestEnv {
@@ -42,7 +45,17 @@ pub mod test {
         }
 
         pub fn new_store(store: bool) -> Self {
+            TestEnv::new_io(StdIO(()), store)
+        }
+    }
+
+    impl<IO: Io + Clone> TestEnv<IO> {
+        pub fn new_io(io: IO, store: bool) -> Self {
             let tmp = tempdir().unwrap();
+            Self::new_io_and_tmp(io, tmp.into(), store)
+        }
+
+        pub fn new_io_and_tmp(io: IO, tmp: Arc<TempDir>, store: bool) -> Self {
             let resolver = |path: &Path| {
                 let name = path
                     .parent()
@@ -56,23 +69,26 @@ pub mod test {
 
             let (sender, receiver) = mpsc::channel(128);
             let registry = Arc::new(
-                WalRegistry::new(
+                WalRegistry::new_with_io(
+                    io.clone(),
                     tmp.path().join("test/wals"),
-                    TestStorage::new_io(store, StdIO(())),
+                    TestStorage::new_io(store, io),
                     sender,
                 )
                 .unwrap(),
             );
+
             if store {
                 let checkpointer = LibsqlCheckpointer::new(registry.clone(), receiver, 5);
                 tokio::spawn(checkpointer.run());
             }
+
             let wal = LibsqlWalManager::new(registry.clone(), Arc::new(resolver));
 
             Self { tmp, registry, wal }
         }
 
-        pub fn shared(&self, namespace: &str) -> Arc<SharedWal<StdIO>> {
+        pub fn shared(&self, namespace: &str) -> Arc<SharedWal<IO>> {
             let path = self.tmp.path().join(namespace).join("data");
             let registry = self.registry.clone();
             let namespace = NamespaceName::from_string(namespace.into());
@@ -83,10 +99,7 @@ pub mod test {
             self.tmp.path().join(namespace)
         }
 
-        pub fn open_conn(
-            &self,
-            namespace: &'static str,
-        ) -> libsql_sys::Connection<LibsqlWal<StdIO>> {
+        pub fn open_conn(&self, namespace: &'static str) -> libsql_sys::Connection<LibsqlWal<IO>> {
             let path = self.db_path(namespace);
             let wal = self.wal.clone();
             std::fs::create_dir_all(&path).unwrap();
@@ -110,7 +123,7 @@ pub mod test {
         }
     }
 
-    pub fn seal_current_segment(shared: &SharedWal<StdIO>) {
+    pub fn seal_current_segment<IO: Io>(shared: &SharedWal<IO>) {
         let mut tx = shared.begin_read(99999).into();
         shared.upgrade(&mut tx).unwrap();
         {

libsql-wal/src/registry.rs

@@ -6,6 +6,7 @@ use std::sync::Arc;
 use dashmap::DashMap;
 use libsql_sys::ffi::Sqlite3DbHeader;
 use parking_lot::{Condvar, Mutex};
+use rand::Rng;
 use tokio::sync::{mpsc, Notify, Semaphore};
 use tokio::task::JoinSet;
 use zerocopy::{AsBytes, FromZeroes};
@@ -32,7 +33,7 @@ enum Slot<IO: Io> {
 
 /// Wal Registry maintains a set of shared Wal, and their respective set of files.
 pub struct WalRegistry<IO: Io, S> {
-    fs: IO,
+    io: IO,
     path: PathBuf,
     shutdown: AtomicBool,
     opened: DashMap<NamespaceName, Slot<IO>>,
@@ -59,7 +60,7 @@ impl<IO: Io, S> WalRegistry<IO, S> {
     ) -> Result<Self> {
         io.create_dir_all(&path)?;
         let registry = Self {
-            fs: io,
+            io,
             path,
             opened: Default::default(),
             shutdown: Default::default(),
@@ -105,14 +106,15 @@ where
             .join(shared.namespace().as_str())
             .join(format!("{}:{start_frame_no:020}.seg", shared.namespace()));
 
-        let segment_file = self.fs.open(true, true, true, &path)?;
-
+        let segment_file = self.io.open(true, true, true, &path)?;
+        let salt = self.io.with_rng(|rng| rng.gen());
         let new = CurrentSegment::create(
             segment_file,
             path,
             start_frame_no,
             current.db_size(),
             current.tail().clone(),
+            salt,
         )?;
         // sealing must the last fallible operation, because we don't want to end up in a situation
         // where the current log is sealed and it wasn't swapped.
@@ -126,9 +128,8 @@ where
                     update_durable(fno, notifier, durable_frame_no, namespace).await;
                 })
             });
-            self.storage
-                .store(&shared.namespace, sealed.clone(), None, cb);
-            new.tail().push(sealed);
+            new.tail().push(sealed.clone());
+            self.storage.store(&shared.namespace, sealed, None, cb);
         }
 
         shared.current.swap(Arc::new(new));
@@ -226,7 +227,7 @@ where
         db_path: &Path,
     ) -> Result<Arc<SharedWal<IO>>> {
         let path = self.path.join(namespace.as_str());
-        self.fs.create_dir_all(&path)?;
+        self.io.create_dir_all(&path)?;
         // TODO: handle that with abstract io
         let dir = walkdir::WalkDir::new(&path).sort_by_file_name().into_iter();
 
@@ -246,7 +247,7 @@ where
                 continue;
             }
 
-            let file = self.fs.open(false, true, true, entry.path())?;
+            let file = self.io.open(false, true, true, entry.path())?;
 
             if let Some(sealed) =
                 SealedSegment::open(file.into(), entry.path().to_path_buf(), Default::default())?
@@ -265,7 +266,7 @@ where
             }
         }
 
-        let db_file = self.fs.open(false, true, true, db_path)?;
+        let db_file = self.io.open(false, true, true, db_path)?;
 
         let mut header: Sqlite3DbHeader = Sqlite3DbHeader::new_zeroed();
         db_file.read_exact_at(header.as_bytes_mut(), 0)?;
@@ -283,14 +284,16 @@ where
 
         let current_path = path.join(format!("{namespace}:{next_frame_no:020}.seg"));
 
-        let segment_file = self.fs.open(true, true, true, &current_path)?;
+        let segment_file = self.io.open(true, true, true, &current_path)?;
+        let salt = self.io.with_rng(|rng| rng.gen());
 
         let current = arc_swap::ArcSwap::new(Arc::new(CurrentSegment::create(
             segment_file,
             current_path,
             next_frame_no,
             db_size,
             tail.into(),
+            salt,
         )?));
 
         let (new_frame_notifier, _) = tokio::sync::watch::channel(next_frame_no.get() - 1);
@@ -309,6 +312,7 @@ where
                 namespace.clone(),
             )),
             shutdown: false.into(),
+            checkpoint_notifier: self.checkpoint_notifier.clone(),
         });
 
         self.opened

libsql-wal/src/replication/injector.rs

@@ -54,7 +54,7 @@ impl<'a, IO: Io> Injector<'a, IO> {
             self.max_tx_frame_no = 0;
         }
         let buffer = current
-            .insert_frames(buffer, commit_data, &mut self.tx)
+            .inject_frames(buffer, commit_data, &mut self.tx)
             .await?;
         self.buffer = buffer;
 

libsql-wal/src/replication/replicator.rs

@@ -299,7 +299,6 @@ mod test {
         let mut replica_content = vec![0u8; db_content.len()];
         while let Some(f) = stream.next().await {
             let frame = f.unwrap();
-            dbg!(frame.header().page_no());
             let offset = (frame.header().page_no() as usize - 1) * 4096;
             tmp.as_file()
                 .write_all_at(frame.data(), offset as u64)