Merge pull request #1608 from tursodatabase/replicate-from-stored-segments

replicate from stored segments

Author: MarinPostma

Cargo.lock

@@ -41,6 +41,7 @@ use libsql_sys::wal::either::Either as EitherWAL;
 #[cfg(feature = "durable-wal")]
 use libsql_sys::wal::either::Either3 as EitherWAL;
 use libsql_sys::wal::Sqlite3WalManager;
+use libsql_wal::checkpointer::LibsqlCheckpointer;
 use libsql_wal::registry::WalRegistry;
 use libsql_wal::storage::NoStorage;
 use libsql_wal::wal::LibsqlWalManager;
@@ -456,7 +457,7 @@ where
         let (stats_sender, stats_receiver) = mpsc::channel(1024);
 
         // chose the wal backend
-        let (make_wal_manager, registry_shutdown) = self.configure_wal_manager()?;
+        let (make_wal_manager, registry_shutdown) = self.configure_wal_manager(&mut join_set)?;
 
         let ns_config = NamespaceConfig {
             db_kind,
@@ -660,6 +661,7 @@ where
 
     fn configure_wal_manager(
         &self,
+        join_set: &mut JoinSet<anyhow::Result<()>>,
     ) -> anyhow::Result<(
         Arc<dyn Fn() -> InnerWalManager + Sync + Send + 'static>,
         Pin<Box<dyn Future<Output = anyhow::Result<()>> + Send + Sync + 'static>>,
@@ -703,15 +705,19 @@ where
 
         match self.use_custom_wal {
             Some(CustomWAL::LibsqlWal) => {
-                let registry = Arc::new(WalRegistry::new(wal_path, NoStorage)?);
+                let (sender, receiver) = tokio::sync::mpsc::channel(64);
+                let registry = Arc::new(WalRegistry::new(wal_path, NoStorage, sender)?);
+                let checkpointer = LibsqlCheckpointer::new(registry.clone(), receiver, 8);
+                join_set.spawn(async move {
+                    checkpointer.run().await;
+                    Ok(())
+                });
 
                 let wal = LibsqlWalManager::new(registry.clone(), Arc::new(namespace_resolver));
                 let shutdown_notify = self.shutdown.clone();
                 let shutdown_fut = Box::pin(async move {
                     shutdown_notify.notified().await;
-                    tokio::task::spawn_blocking(move || registry.shutdown())
-                        .await
-                        .unwrap()?;
+                    registry.shutdown().await?;
                     Ok(())
                 });
 

libsql-server/src/lib.rs

@@ -41,6 +41,7 @@ clap = { version = "4.5.9", optional = true, features = ["derive"] }
 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"
 
 [dev-dependencies]
 criterion = "0.5.1"

libsql-wal/Cargo.toml

@@ -59,7 +59,8 @@ fn with_libsql_conn(f: impl FnOnce(&mut Connection<LibsqlWal<StdIO>>)) {
     let tmp = tempdir().unwrap();
     let resolver = |_: &Path| NamespaceName::from_string("test".into());
 
-    let registry = Arc::new(WalRegistry::new(tmp.path().join("wals"), NoStorage).unwrap());
+    let (sender, _) = tokio::sync::mpsc::channel(12);
+    let registry = Arc::new(WalRegistry::new(tmp.path().join("wals"), NoStorage, sender).unwrap());
     let wal_manager = LibsqlWalManager::new(registry.clone(), Arc::new(resolver));
 
     let mut conn = libsql_sys::Connection::open(

libsql-wal/benches/benchmarks.rs

@@ -6,6 +6,7 @@ use aws_config::{BehaviorVersion, Region};
 use aws_credential_types::Credentials;
 use aws_sdk_s3::config::SharedCredentialsProvider;
 use clap::{Parser, ValueEnum};
+use libsql_wal::checkpointer::LibsqlCheckpointer;
 use tokio::task::{block_in_place, JoinSet};
 
 use libsql_sys::name::NamespaceName;
@@ -83,21 +84,24 @@ async fn main() {
 
     if cli.s3_args.enable_s3 {
         let storage = setup_s3_storage(&cli, &mut join_set).await;
-        handle(&cli, storage).await;
+        handle(&cli, storage, &mut join_set).await;
     } else {
         todo!()
     }
 
     while join_set.join_next().await.is_some() {}
 }
 
-async fn handle<S>(cli: &Cli, storage: S)
+async fn handle<S>(cli: &Cli, storage: S, join_set: &mut JoinSet<()>)
 where
     S: Storage<Segment = SealedSegment<std::fs::File>>,
 {
     match &cli.subcommand {
         Subcommand::Shell { db_path } => {
-            let registry = WalRegistry::new(db_path.clone(), storage).unwrap();
+            let (sender, receiver) = tokio::sync::mpsc::channel(64);
+            let registry = Arc::new(WalRegistry::new(db_path.clone(), storage, sender).unwrap());
+            let checkpointer = LibsqlCheckpointer::new(registry.clone(), receiver, 64);
+            join_set.spawn(checkpointer.run());
             run_shell(
                 registry,
                 &db_path,
@@ -143,13 +147,15 @@ where
     println!("max durable frame: {durable}");
 }
 
-async fn run_shell<S>(registry: WalRegistry<StdIO, S>, db_path: &Path, namespace: NamespaceName)
-where
+async fn run_shell<S>(
+    registry: Arc<WalRegistry<StdIO, S>>,
+    db_path: &Path,
+    namespace: NamespaceName,
+) where
     S: Storage<Segment = SealedSegment<std::fs::File>>,
 {
     let db_path = db_path.join("dbs").join(namespace.as_str());
     tokio::fs::create_dir_all(&db_path).await.unwrap();
-    let registry = Arc::new(registry);
     let resolver = move |path: &Path| {
         NamespaceName::from_string(
             path.parent()
@@ -216,7 +222,7 @@ where
 
     drop(conn);
 
-    registry.shutdown().unwrap();
+    registry.shutdown().await.unwrap();
 }
 
 async fn handle_builtin<S>(

libsql-wal/src/bins/shell/main.rs

@@ -11,16 +11,29 @@ use crate::registry::WalRegistry;
 
 pub(crate) type NotifyCheckpointer = mpsc::Sender<NamespaceName>;
 
-type LibsqlCheckpointer<IO, S> = Checkpointer<WalRegistry<IO, S>>;
+pub enum CheckpointMessage {
+    /// notify that a namespace may be checkpointable
+    Namespace(NamespaceName),
+    /// shutdown initiated
+    Shutdown,
+}
+
+impl From<NamespaceName> for CheckpointMessage {
+    fn from(value: NamespaceName) -> Self {
+        Self::Namespace(value)
+    }
+}
+
+pub type LibsqlCheckpointer<IO, S> = Checkpointer<WalRegistry<IO, S>>;
 
 impl<IO, S> LibsqlCheckpointer<IO, S>
 where
     IO: Io,
     S: Sync + Send + 'static,
 {
     pub fn new(
-        registry: WalRegistry<IO, S>,
-        notifier: mpsc::Receiver<NamespaceName>,
+        registry: Arc<WalRegistry<IO, S>>,
+        notifier: mpsc::Receiver<CheckpointMessage>,
         max_checkpointing_conccurency: usize,
     ) -> Self {
         Self::new_with_performer(registry, notifier, max_checkpointing_conccurency)
@@ -70,7 +83,7 @@ pub struct Checkpointer<P> {
     checkpointing: HashSet<NamespaceName>,
     /// the checkpointer is notifier whenever there is a change to a namespage that could trigger a
     /// checkpoint
-    recv: mpsc::Receiver<NamespaceName>,
+    recv: mpsc::Receiver<CheckpointMessage>,
     max_checkpointing_conccurency: usize,
     shutting_down: bool,
     join_set: JoinSet<(NamespaceName, crate::error::Result<()>)>,
@@ -84,12 +97,12 @@ where
     P: PerformCheckpoint + Send + Sync + 'static,
 {
     fn new_with_performer(
-        perform_checkpoint: P,
-        notifier: mpsc::Receiver<NamespaceName>,
+        perform_checkpoint: Arc<P>,
+        notifier: mpsc::Receiver<CheckpointMessage>,
         max_checkpointing_conccurency: usize,
     ) -> Self {
         Self {
-            perform_checkpoint: Arc::new(perform_checkpoint),
+            perform_checkpoint,
             scheduled: Default::default(),
             checkpointing: Default::default(),
             recv: notifier,
@@ -141,10 +154,10 @@ where
             }
             notified = self.recv.recv(), if !self.shutting_down => {
                 match notified {
-                    Some(namespace) => {
+                    Some(CheckpointMessage::Namespace(namespace)) => {
                         self.scheduled.insert(namespace);
                     }
-                    None => {
+                    None | Some(CheckpointMessage::Shutdown) => {
                         self.shutting_down = true;
                     }
                 }
@@ -201,10 +214,11 @@ mod test {
         }
 
         let (sender, receiver) = mpsc::channel(8);
-        let mut checkpointer = Checkpointer::new_with_performer(TestPerformCheckoint, receiver, 5);
+        let mut checkpointer =
+            Checkpointer::new_with_performer(TestPerformCheckoint.into(), receiver, 5);
         let ns = NamespaceName::from("test");
 
-        sender.send(ns.clone()).await.unwrap();
+        sender.send(ns.clone().into()).await.unwrap();
 
         checkpointer.step().await;
 
@@ -233,10 +247,11 @@ mod test {
         }
 
         let (sender, receiver) = mpsc::channel(8);
-        let mut checkpointer = Checkpointer::new_with_performer(TestPerformCheckoint, receiver, 5);
+        let mut checkpointer =
+            Checkpointer::new_with_performer(TestPerformCheckoint.into(), receiver, 5);
         let ns = NamespaceName::from("test");
 
-        sender.send(ns.clone()).await.unwrap();
+        sender.send(ns.clone().into()).await.unwrap();
 
         checkpointer.step().await;
         assert_eq!(checkpointer.errors, 0);
@@ -264,7 +279,8 @@ mod test {
         }
 
         let (sender, receiver) = mpsc::channel(8);
-        let mut checkpointer = Checkpointer::new_with_performer(TestPerformCheckoint, receiver, 5);
+        let mut checkpointer =
+            Checkpointer::new_with_performer(TestPerformCheckoint.into(), receiver, 5);
 
         drop(sender);
 
@@ -290,7 +306,8 @@ mod test {
         }
 
         let (sender, receiver) = mpsc::channel(8);
-        let mut checkpointer = Checkpointer::new_with_performer(TestPerformCheckoint, receiver, 5);
+        let mut checkpointer =
+            Checkpointer::new_with_performer(TestPerformCheckoint.into(), receiver, 5);
 
         drop(sender);
 
@@ -323,12 +340,13 @@ mod test {
         }
 
         let (sender, receiver) = mpsc::channel(8);
-        let mut checkpointer = Checkpointer::new_with_performer(TestPerformCheckoint, receiver, 5);
+        let mut checkpointer =
+            Checkpointer::new_with_performer(TestPerformCheckoint.into(), receiver, 5);
 
         let ns: NamespaceName = "test".into();
 
-        sender.send(ns.clone()).await.unwrap();
-        sender.send(ns.clone()).await.unwrap();
+        sender.send(ns.clone().into()).await.unwrap();
+        sender.send(ns.clone().into()).await.unwrap();
 
         checkpointer.step().await;
 
@@ -355,13 +373,14 @@ mod test {
         }
 
         let (sender, receiver) = mpsc::channel(8);
-        let mut checkpointer = Checkpointer::new_with_performer(TestPerformCheckoint, receiver, 5);
+        let mut checkpointer =
+            Checkpointer::new_with_performer(TestPerformCheckoint.into(), receiver, 5);
 
         let ns1: NamespaceName = "test1".into();
         let ns2: NamespaceName = "test2".into();
 
-        sender.send(ns1.clone()).await.unwrap();
-        sender.send(ns2.clone()).await.unwrap();
+        sender.send(ns1.clone().into()).await.unwrap();
+        sender.send(ns2.clone().into()).await.unwrap();
 
         checkpointer.step().await;
 
@@ -390,15 +409,16 @@ mod test {
         }
 
         let (sender, receiver) = mpsc::channel(8);
-        let mut checkpointer = Checkpointer::new_with_performer(TestPerformCheckoint, receiver, 2);
+        let mut checkpointer =
+            Checkpointer::new_with_performer(TestPerformCheckoint.into(), receiver, 2);
 
         let ns1: NamespaceName = "test1".into();
         let ns2: NamespaceName = "test2".into();
         let ns3: NamespaceName = "test3".into();
 
-        sender.send(ns1.clone()).await.unwrap();
-        sender.send(ns2.clone()).await.unwrap();
-        sender.send(ns3.clone()).await.unwrap();
+        sender.send(ns1.clone().into()).await.unwrap();
+        sender.send(ns2.clone().into()).await.unwrap();
+        sender.send(ns3.clone().into()).await.unwrap();
 
         checkpointer.step().await;
         checkpointer.step().await;