fix migration

Author: MarinPostma

libsql-server/src/bottomless_migrate.rs

@@ -1,17 +1,17 @@
 use std::path::{Path, PathBuf};
 use std::sync::Arc;
 
-use libsql_replication::LIBSQL_PAGE_SIZE;
+use libsql_sys::ffi::Sqlite3DbHeader;
 use libsql_sys::wal::Sqlite3WalManager;
 use libsql_wal::io::StdIO;
 use libsql_wal::replication::injector::Injector;
 use libsql_wal::segment::{Frame, FrameHeader};
-use libsql_wal::storage::Storage;
-use libsql_wal::{registry::WalRegistry, segment::sealed::SealedSegment};
+use libsql_wal::storage::NoStorage;
+use libsql_wal::registry::WalRegistry;
 use tempfile::TempDir;
 use tokio::io::AsyncReadExt;
 use tokio_stream::StreamExt;
-use zerocopy::FromZeroes;
+use zerocopy::{FromZeroes, FromBytes};
 
 #[cfg(not(feature = "durable-wal"))]
 use libsql_sys::wal::either::Either as EitherWAL;
@@ -27,14 +27,11 @@ use crate::namespace::NamespaceStore;
 
 /// The process for migrating from bottomless to libsql wal is simple:
 /// 1) iteratate over all namespaces, and make sure that they   
-pub async fn bottomless_migrate<S>(
+pub async fn bottomless_migrate(
     meta_store: MetaStore,
-    storage: Arc<S>,
     base_config: BaseNamespaceConfig,
     primary_config: PrimaryConfig,
 ) -> anyhow::Result<()>
-where
-    S: Storage<Segment = SealedSegment<std::fs::File>>,
 {
     tracing::info!("attempting bottomless migration to libsql-wal");
 
@@ -46,8 +43,21 @@ where
     let configs_stream = meta_store.namespaces();
     tokio::pin!(configs_stream);
 
-    let (sender, _) = tokio::sync::mpsc::channel(1);
-    let tmp_registry = Arc::new(WalRegistry::new(tmp.path().join("wals"), storage, sender)?);
+    let (sender, mut rcv) = tokio::sync::mpsc::channel(1);
+
+    tokio::spawn(async move {
+        loop {
+            match rcv.recv().await {
+                Some(libsql_wal::checkpointer::CheckpointMessage::Shutdown) => {
+                    break
+                }
+                Some(_) => (),
+                None => break,
+            }
+        }
+    });
+
+    let tmp_registry = Arc::new(WalRegistry::new(tmp.path().join("wals"), NoStorage.into(), sender)?);
 
     let mut configurators = NamespaceConfigurators::default();
 
@@ -78,6 +88,14 @@ where
         .await?;
     }
 
+    tmp_registry.shutdown().await?;
+
+    // FIXME: this is not atomic!!
+    tokio::fs::remove_dir_all(base_config.base_path.join("dbs")).await?;
+    tokio::fs::rename(tmp.path().join("dbs"), base_config.base_path.join("dbs")).await?;
+    tokio::fs::rename(tmp.path().join("wals"), base_config.base_path.join("wals")).await?;
+
+
     Ok(())
 }
 
@@ -88,20 +106,18 @@ where
 /// - either the migration didn't happen
 /// - a crash happened before we could swap the directories
 #[tracing::instrument(skip_all, fields(namespace = config.namespace().as_str()))]
-async fn migrate_one<S>(
+async fn migrate_one(
     configurators: &NamespaceConfigurators,
     config: MetaStoreHandle,
     dummy_store: NamespaceStore,
     tmp: &Path,
-    tmp_registry: Arc<WalRegistry<StdIO, S>>,
+    tmp_registry: Arc<WalRegistry<StdIO, NoStorage>>,
     base_path: &Path,
 ) -> anyhow::Result<()>
-where
-    S: Storage<Segment = SealedSegment<std::fs::File>>,
 {
     let broadcasters = BroadcasterRegistry::default();
     // TODO: check if we already have a backup for this db from storage
-    tracing::info!("started db migrating");
+    tracing::info!("started namespace migration");
     // we load the namespace ensuring it's restored to the latest version
     configurators
         .configure_primary()?
@@ -118,13 +134,16 @@ where
         )
         .await?;
 
-    let db_path = tmp
+    let db_dir = tmp
         .join("dbs")
-        .join(config.namespace().as_str())
-        .join("data");
+        .join(config.namespace().as_str());
+    tokio::fs::create_dir_all(&db_dir).await?;
+    let db_path = db_dir.join("data");
     let registry = tmp_registry.clone();
     let namespace = config.namespace().clone();
-    let shared = tokio::task::spawn_blocking(move || registry.open(&db_path, &namespace.into()))
+    let shared = tokio::task::spawn_blocking({
+        let registry = registry.clone();
+        move || registry.open(&db_path, &namespace.into()) })
         .await
         .unwrap()
         .unwrap();
@@ -141,18 +160,35 @@ where
         .join(config.namespace().as_str())
         .join("data");
     let mut orig_db_file = tokio::fs::File::open(orig_db_path).await?;
-    let orig_db_file_len = orig_db_file.metadata().await?.len();
-    for i in 0..(orig_db_file_len / LIBSQL_PAGE_SIZE as u64) {
+    let mut db_size = usize::MAX;
+    let mut current = 0;
+    while current < db_size {
         let mut frame: Box<Frame> = Frame::new_box_zeroed();
+        orig_db_file.read_exact(frame.data_mut()).await?;
+        if current == 0 {
+            let header: Sqlite3DbHeader = Sqlite3DbHeader::read_from_prefix(frame.data()).unwrap();
+            db_size = header.db_size.get() as usize;
+        }
+        let size_after = if current == db_size - 1 {
+            db_size as u32
+        } else {
+            0
+        };
         *frame.header_mut() = FrameHeader {
-            page_no: (i as u32 + 1).into(),
-            size_after: 0.into(),
-            frame_no: (i + 1).into(),
+            page_no: (current as u32 + 1).into(),
+            size_after: size_after.into(),
+            frame_no: (current as u64 + 1).into(),
         };
-        orig_db_file.read_exact(frame.data_mut()).await?;
         injector.insert_frame(frame).await?;
+        current += 1;
     }
 
+    drop(injector);
+
+    tokio::task::spawn_blocking(move || {
+        shared.seal_current()
+    }).await.unwrap()?;
+
     tracing::info!("sucessfull migration");
 
     Ok(())

libsql-server/src/lib.rs

@@ -54,7 +54,7 @@ use libsql_wal::registry::WalRegistry;
 use libsql_wal::segment::sealed::SealedSegment;
 use libsql_wal::storage::async_storage::{AsyncStorage, AsyncStorageInitConfig};
 use libsql_wal::storage::backend::s3::S3Backend;
-use libsql_wal::storage::{NoStorage, Storage};
+use libsql_wal::storage::NoStorage;
 use namespace::meta_store::MetaStoreHandle;
 use namespace::NamespaceName;
 use net::Connector;
@@ -65,6 +65,7 @@ use tokio::runtime::Runtime;
 use tokio::sync::{mpsc, Notify, Semaphore};
 use tokio::task::JoinSet;
 use tokio::time::Duration;
+use tokio_stream::StreamExt as _;
 use tonic::transport::Channel;
 use url::Url;
 use utils::services::idle_shutdown::IdleShutdownKicker;
@@ -878,24 +879,42 @@ where
         };
 
         // perform migration before creating the actual registry creation
-        self.maybe_migrate_bottomless(
-            meta_store,
-            storage.clone(),
+        let did_migrate = self.maybe_migrate_bottomless(
+            meta_store.clone(),
             &base_config,
             &primary_config,
         ).await?;
 
+
         if self.rpc_server_config.is_some() && matches!(*storage, Either::B(_)) {
             anyhow::bail!("replication without bottomless not supported yet");
         }
 
+
         let registry = Arc::new(WalRegistry::new(wal_path, storage, sender)?);
         let checkpointer = LibsqlCheckpointer::new(registry.clone(), receiver, 8);
         task_manager.spawn_with_shutdown_notify(|_| async move {
             checkpointer.run().await;
             Ok(())
         });
 
+        // If we have performed the migration, load all shared wals to force flush to storage with
+        // the new registry
+        if did_migrate {
+            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 namespace_resolver = Arc::new(|path: &Path| {
             NamespaceName::from_string(
                 path.parent()
@@ -1139,15 +1158,13 @@ where
     /// - migrate_bottomless flag is raised
     /// - there hasn't been a previous successfull migration (wals directory is either absent,
     /// or emtpy)
-    async fn maybe_migrate_bottomless<S>(
+    /// returns whether the migration was performed
+    async fn maybe_migrate_bottomless(
         &self,
         meta_store: MetaStore,
-        storage: Arc<S>,
         base_config: &BaseNamespaceConfig,
         primary_config: &PrimaryConfig,
-        ) -> anyhow::Result<()>
-        where S: Storage<Segment = SealedSegment<std::fs::File>>,
-    {
+        ) -> 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();
@@ -1159,18 +1176,24 @@ where
             && is_libsql_wal;
 
         if should_attempt_migration {
-            bottomless_migrate(meta_store, storage, base_config.clone(), primary_config.clone()).await?;
+            bottomless_migrate(meta_store, base_config.clone(), primary_config.clone()).await?;
+            Ok(true)
+        } else {
+            tracing::info!("bottomless already migrated, skipping...");
+            Ok(false)
         }
-
-        Ok(())
     }
 
     fn check_previous_migration_success(&self) -> anyhow::Result<bool> {
         let wals_path = self.path.join("wals");
+        if !wals_path.try_exists()? {
+            return Ok(false)
+        }
+
         let dir = std::fs::read_dir(&wals_path)?;
 
         // wals dir exist and is not empty
-        Ok(wals_path.try_exists()? && dir.count() != 0)
+        Ok(dir.count() != 0)
     }
 }