Merge pull request #1711 from tursodatabase/libsql-wal-db-restore

libsql-wal: Sync dbs from remote storage on startup

Author: MarinPostma

bottomless/src/replicator.rs

@@ -1652,6 +1652,7 @@ impl Replicator {
                     let frame = RpcFrame {
                         data: frame_to_inject.bytes(),
                         timestamp: None,
+                        durable_frame_no: None,
                     };
                     injector.inject_frame(frame).await?;
                     applied_wal_frame = true;

libsql-replication/proto/replication_log.proto

@@ -36,6 +36,7 @@ message Frame {
     // if this frames is a commit frame, then this can be set
     // to the time when the transaction was commited
     optional int64 timestamp = 2;
+    optional uint64 durable_frame_no = 3;
 }
 
 message Frames {

libsql-replication/src/generated/wal_log.rs

@@ -83,6 +83,8 @@ pub struct Frame {
     /// to the time when the transaction was commited
     #[prost(int64, optional, tag = "2")]
     pub timestamp: ::core::option::Option<i64>,
+    #[prost(uint64, optional, tag = "3")]
+    pub durable_frame_no: ::core::option::Option<u64>,
 }
 #[allow(clippy::derive_partial_eq_without_eq)]
 #[derive(Clone, PartialEq, ::prost::Message)]

libsql-replication/src/injector/libsql_injector.rs

@@ -49,4 +49,8 @@ impl super::Injector for LibsqlInjector {
             .map_err(|e| Error::FatalInjectError(e.into()))?;
         Ok(None)
     }
+
+    fn durable_frame_no(&mut self, frame_no: u64) {
+        self.injector.set_durable(frame_no);
+    }
 }

libsql-replication/src/injector/mod.rs

@@ -29,4 +29,6 @@ pub trait Injector {
     /// Trigger a dummy write, and flush the cache to trigger a call to xFrame. The buffer's frame
     /// are then injected into the wal.
     fn flush(&mut self) -> impl Future<Output = Result<Option<FrameNo>>> + Send;
+
+    fn durable_frame_no(&mut self, frame_no: u64);
 }

libsql-replication/src/injector/sqlite_injector/mod.rs

@@ -46,6 +46,9 @@ impl Injector for SqliteInjector {
         let inner = self.inner.clone();
         spawn_blocking(move || inner.lock().flush()).await.unwrap()
     }
+
+    #[inline]
+    fn durable_frame_no(&mut self, _frame_no: u64) {}
 }
 
 impl SqliteInjector {

libsql-replication/src/replicator.rs

@@ -328,6 +328,10 @@ where
     async fn inject_frame(&mut self, frame: RpcFrame) -> Result<(), Error> {
         self.frames_synced += 1;
 
+        if let Some(frame_no) = frame.durable_frame_no {
+            self.injector.durable_frame_no(frame_no);
+        }
+
         match self.injector.inject_frame(frame).await? {
             Some(commit_fno) => {
                 self.client.commit_frame_no(commit_fno).await?;
@@ -772,6 +776,7 @@ mod test {
                         .map(|f| RpcFrame {
                             data: f.bytes(),
                             timestamp: None,
+                            durable_frame_no: None,
                         })
                         .take(2)
                         .map(Ok)
@@ -785,6 +790,7 @@ mod test {
                         .map(|f| RpcFrame {
                             data: f.bytes(),
                             timestamp: None,
+                            durable_frame_no: None,
                         })
                         .map(Ok)
                         .collect::<Vec<_>>();

libsql-server/src/bottomless_migrate.rs

@@ -162,13 +162,7 @@ async fn migrate_one(
     .await
     .unwrap()?;
 
-    let mut tx = shared.begin_read(0).into();
-    shared.upgrade(&mut tx).unwrap();
-    let guard = tx
-        .into_write()
-        .unwrap_or_else(|_| panic!("should be a write txn"))
-        .into_lock_owned();
-    let mut injector = Injector::new(shared.clone(), guard, 10)?;
+    let mut injector = Injector::new(shared.clone(), 10)?;
     let orig_db_path = base_path
         .join("dbs")
         .join(config.namespace().as_str())

libsql-server/src/lib.rs

@@ -177,6 +177,9 @@ pub struct Server<C = HttpConnector, A = AddrIncoming, D = HttpsConnector<HttpCo
     pub connector: Option<D>,
     pub migrate_bottomless: bool,
     pub enable_deadlock_monitor: bool,
+    pub should_sync_from_storage: bool,
+    pub force_load_wals: bool,
+    pub sync_conccurency: usize,
 }
 
 impl<C, A, D> Default for Server<C, A, D> {
@@ -203,6 +206,9 @@ impl<C, A, D> Default for Server<C, A, D> {
             connector: None,
             migrate_bottomless: false,
             enable_deadlock_monitor: false,
+            should_sync_from_storage: false,
+            force_load_wals: false,
+            sync_conccurency: 8,
         }
     }
 }
@@ -781,10 +787,10 @@ where
         tokio::select! {
             _ = shutdown.notified() => {
                 let shutdown = async {
+                    namespace_store.shutdown().await?;
                     task_manager.shutdown().await?;
                     // join_set.shutdown().await;
                     service_shutdown.notify_waiters();
-                    namespace_store.shutdown().await?;
 
                     Ok::<_, crate::Error>(())
                 };
@@ -958,19 +964,30 @@ where
             Ok(())
         });
 
-        // If we have performed the migration, load all shared wals to force flush to storage with
-        // the new registry
-        if did_migrate {
+        // If we performed a migration from bottomless to libsql-wal earlier, then we need to
+        // forecefully load all the wals, to trigger segment storage with the actual storage. This
+        // is because migration didn't actually send anything to storage, but just created the
+        // segments.
+        if did_migrate || self.should_sync_from_storage || self.force_load_wals {
+            // eagerly load all namespaces, then call sync_all on the registry
+            // TODO: do conccurently
             let dbs_path = base_config.base_path.join("dbs");
             let stream = meta_store.namespaces();
             tokio::pin!(stream);
             while let Some(conf) = stream.next().await {
                 let registry = registry.clone();
                 let namespace = conf.namespace().clone();
-                let path = dbs_path.join(namespace.as_str()).join("data");
-                tokio::task::spawn_blocking(move || registry.open(&path, &namespace.into()))
-                    .await
-                    .unwrap()?;
+                let path = dbs_path.join(namespace.as_str());
+                tokio::fs::create_dir_all(&path).await?;
+                tokio::task::spawn_blocking(move || {
+                    registry.open(&path.join("data"), &namespace.into())
+                })
+                .await
+                .unwrap()?;
+            }
+
+            if self.should_sync_from_storage {
+                registry.sync_all(self.sync_conccurency).await?;
             }
         }
 
@@ -1236,31 +1253,31 @@ where
         base_config: &BaseNamespaceConfig,
         primary_config: &PrimaryConfig,
     ) -> anyhow::Result<bool> {
-        let is_previous_migration_successful = self.check_previous_migration_success()?;
-        let is_libsql_wal = matches!(self.use_custom_wal, Some(CustomWAL::LibsqlWal));
-        let is_bottomless_enabled = self.db_config.bottomless_replication.is_some();
         let is_primary = self.rpc_client_config.is_none();
-        let should_attempt_migration = self.migrate_bottomless
-            && is_primary
-            && is_bottomless_enabled
-            && !is_previous_migration_successful
-            && is_libsql_wal;
-
-        if should_attempt_migration {
-            bottomless_migrate(meta_store, base_config.clone(), primary_config.clone()).await?;
-            Ok(true)
-        } else {
-            // the wals directory is present and so is the _dbs. This means that a crash occured
-            // before we could remove it. clean it up now. see code in `bottomless_migrate.rs`
-            let tmp_dbs_path = base_config.base_path.join("_dbs");
-            if tmp_dbs_path.try_exists()? {
-                tracing::info!("removed dangling `_dbs` folder");
-                tokio::fs::remove_dir_all(&tmp_dbs_path).await?;
-            }
+        if self.migrate_bottomless && is_primary {
+            let is_previous_migration_successful = self.check_previous_migration_success()?;
+            let is_libsql_wal = matches!(self.use_custom_wal, Some(CustomWAL::LibsqlWal));
+            let is_bottomless_enabled = self.db_config.bottomless_replication.is_some();
+            let should_attempt_migration =
+                is_bottomless_enabled && !is_previous_migration_successful && is_libsql_wal;
+
+            if should_attempt_migration {
+                bottomless_migrate(meta_store, base_config.clone(), primary_config.clone()).await?;
+                return Ok(true);
+            } else {
+                // the wals directory is present and so is the _dbs. This means that a crash occured
+                // before we could remove it. clean it up now. see code in `bottomless_migrate.rs`
+                let tmp_dbs_path = base_config.base_path.join("_dbs");
+                if tmp_dbs_path.try_exists()? {
+                    tracing::info!("removed dangling `_dbs` folder");
+                    tokio::fs::remove_dir_all(&tmp_dbs_path).await?;
+                }
 
-            tracing::info!("bottomless already migrated, skipping...");
-            Ok(false)
+                tracing::info!("bottomless already migrated, skipping...");
+            }
         }
+
+        Ok(false)
     }
 
     fn check_previous_migration_success(&self) -> anyhow::Result<bool> {

libsql-server/src/main.rs

@@ -280,6 +280,27 @@ struct Cli {
     /// Auth key for the admin API
     #[clap(long, env = "LIBSQL_ADMIN_AUTH_KEY", requires = "admin_listen_addr")]
     admin_auth_key: Option<String>,
+
+    /// Whether to perform a sync of all namespaces with remote on startup
+    #[clap(
+        long,
+        env = "LIBSQL_SYNC_FROM_STORAGE",
+        requires = "enable_bottomless_replication"
+    )]
+    sync_from_storage: bool,
+    /// Whether to force loading all WAL at startup, with libsql-wal
+    /// By default, WALs are loaded lazily, as the databases are openned.
+    /// Whether to force loading all wal at startup
+    #[clap(long)]
+    force_load_wals: bool,
+    /// Sync conccurency
+    #[clap(
+        long,
+        env = "LIBSQL_SYNC_CONCCURENCY",
+        requires = "sync_from_storage",
+        default_value = "8"
+    )]
+    sync_conccurency: usize,
 }
 
 #[derive(clap::Subcommand, Debug)]
@@ -681,6 +702,9 @@ async fn build_server(config: &Cli) -> anyhow::Result<Server> {
         connector: Some(https),
         migrate_bottomless: config.migrate_bottomless,
         enable_deadlock_monitor: config.enable_deadlock_monitor,
+        should_sync_from_storage: config.sync_from_storage,
+        force_load_wals: config.force_load_wals,
+        sync_conccurency: config.sync_conccurency,
     })
 }