Namespace forking (#646)

* add namespace fork endpoint

* add fork method to namespace store

* implement general forking mecanism

* refactor namespace module

* fix bottomless did_recover

* fix clippy

* handle fork replica error

* fix rebase quirks

* fix rebase issues

* add docstring to restore method

Author: MarinPostma

bottomless/src/lib.rs

@@ -415,7 +415,7 @@ pub extern "C" fn xGetPathname(buf: *mut c_char, orig: *const c_char, orig_len:
 
 async fn try_restore(replicator: &mut replicator::Replicator) -> i32 {
     match replicator.restore(None, None).await {
-        Ok(replicator::RestoreAction::SnapshotMainDbFile) => {
+        Ok((replicator::RestoreAction::SnapshotMainDbFile, _)) => {
             replicator.new_generation();
             match replicator.snapshot_main_db_file(None).await {
                 Ok(Some(h)) => {
@@ -437,7 +437,7 @@ async fn try_restore(replicator: &mut replicator::Replicator) -> i32 {
                 return ffi::SQLITE_CANTOPEN;
             }
         }
-        Ok(replicator::RestoreAction::ReuseGeneration(gen)) => {
+        Ok((replicator::RestoreAction::ReuseGeneration(gen), _)) => {
             replicator.set_generation(gen);
         }
         Err(e) => {

bottomless/src/replicator.rs

@@ -962,12 +962,13 @@ impl Replicator {
         Some((first_frame_no, last_frame_no, timestamp, compression_kind))
     }
 
-    // Restores the database state from given remote generation
+    /// Restores the database state from given remote generation
+    /// On success, returns the RestoreAction, and whether the database was recovered from backup.
     async fn restore_from(
         &mut self,
         generation: Uuid,
         utc_time: Option<NaiveDateTime>,
-    ) -> Result<RestoreAction> {
+    ) -> Result<(RestoreAction, bool)> {
         if let Some(tombstone) = self.get_tombstone().await? {
             if let Some(timestamp) = Self::generation_to_timestamp(&generation) {
                 if tombstone.timestamp() as u64 >= timestamp.to_unix().0 {
@@ -989,7 +990,7 @@ impl Replicator {
         let last_frame = self.get_last_consistent_frame(&generation).await?;
         tracing::debug!("Last consistent remote frame in generation {generation}: {last_frame}.");
         if let Some(action) = self.compare_with_local(generation, last_frame).await? {
-            return Ok(action);
+            return Ok((action, false));
         }
 
         // at this point we know, we should do a full restore
@@ -1071,11 +1072,11 @@ impl Replicator {
 
         if applied_wal_frame {
             tracing::info!("WAL file has been applied onto database file in generation {}. Requesting snapshot.", generation);
-            Ok::<_, anyhow::Error>(RestoreAction::SnapshotMainDbFile)
+            Ok::<_, anyhow::Error>((RestoreAction::SnapshotMainDbFile, true))
         } else {
             tracing::info!("Reusing generation {}.", generation);
             // since WAL was not applied, we can reuse the latest generation
-            Ok::<_, anyhow::Error>(RestoreAction::ReuseGeneration(generation))
+            Ok::<_, anyhow::Error>((RestoreAction::ReuseGeneration(generation), true))
         }
     }
 
@@ -1299,19 +1300,20 @@ impl Replicator {
         Ok(())
     }
 
-    // Restores the database state from newest remote generation
+    /// Restores the database state from newest remote generation
+    /// On success, returns the RestoreAction, and whether the database was recovered from backup.
     pub async fn restore(
         &mut self,
         generation: Option<Uuid>,
         timestamp: Option<NaiveDateTime>,
-    ) -> Result<RestoreAction> {
+    ) -> Result<(RestoreAction, bool)> {
         let generation = match generation {
             Some(gen) => gen,
             None => match self.latest_generation_before(timestamp.as_ref()).await {
                 Some(gen) => gen,
                 None => {
                     tracing::debug!("No generation found, nothing to restore");
-                    return Ok(RestoreAction::SnapshotMainDbFile);
+                    return Ok((RestoreAction::SnapshotMainDbFile, false));
                 }
             },
         };

sqld/src/admin_api.rs

@@ -15,21 +15,25 @@ use crate::connection::config::{DatabaseConfig, DatabaseConfigStore};
 use crate::error::LoadDumpError;
 use crate::namespace::{DumpStream, MakeNamespace, NamespaceStore, RestoreOption};
 
-struct AppState<F: MakeNamespace> {
+struct AppState<M: MakeNamespace> {
     db_config_store: Arc<DatabaseConfigStore>,
-    namespaces: Arc<NamespaceStore<F>>,
+    namespaces: Arc<NamespaceStore<M>>,
 }
 
-pub async fn run_admin_api<F: MakeNamespace>(
+pub async fn run_admin_api<M: MakeNamespace>(
     addr: SocketAddr,
     db_config_store: Arc<DatabaseConfigStore>,
-    namespaces: Arc<NamespaceStore<F>>,
+    namespaces: Arc<NamespaceStore<M>>,
 ) -> anyhow::Result<()> {
     use axum::routing::{get, post};
     let router = axum::Router::new()
         .route("/", get(handle_get_index))
         .route("/v1/config", get(handle_get_config))
         .route("/v1/block", post(handle_post_block))
+        .route(
+            "/v1/namespaces/:namespace/fork/:to",
+            post(handle_fork_namespace),
+        )
         .route(
             "/v1/namespaces/:namespace/create",
             post(handle_create_namespace),
@@ -60,8 +64,8 @@ async fn handle_get_index() -> &'static str {
     "Welcome to the sqld admin API"
 }
 
-async fn handle_get_config<F: MakeNamespace>(
-    State(app_state): State<Arc<AppState<F>>>,
+async fn handle_get_config<M: MakeNamespace>(
+    State(app_state): State<Arc<AppState<M>>>,
 ) -> Json<Arc<DatabaseConfig>> {
     Json(app_state.db_config_store.get())
 }
@@ -79,8 +83,8 @@ struct CreateNamespaceReq {
     dump_url: Option<Url>,
 }
 
-async fn handle_post_block<F: MakeNamespace>(
-    State(app_state): State<Arc<AppState<F>>>,
+async fn handle_post_block<M: MakeNamespace>(
+    State(app_state): State<Arc<AppState<M>>>,
     Json(req): Json<BlockReq>,
 ) -> (axum::http::StatusCode, &'static str) {
     let mut config = (*app_state.db_config_store.get()).clone();
@@ -97,8 +101,8 @@ async fn handle_post_block<F: MakeNamespace>(
     }
 }
 
-async fn handle_create_namespace<F: MakeNamespace>(
-    State(app_state): State<Arc<AppState<F>>>,
+async fn handle_create_namespace<M: MakeNamespace>(
+    State(app_state): State<Arc<AppState<M>>>,
     Path(namespace): Path<String>,
     Json(req): Json<CreateNamespaceReq>,
 ) -> crate::Result<()> {
@@ -111,6 +115,14 @@ async fn handle_create_namespace<F: MakeNamespace>(
     Ok(())
 }
 
+async fn handle_fork_namespace<M: MakeNamespace>(
+    State(app_state): State<Arc<AppState<M>>>,
+    Path((from, to)): Path<(String, String)>,
+) -> crate::Result<()> {
+    app_state.namespaces.fork(from.into(), to.into()).await?;
+    Ok(())
+}
+
 async fn dump_stream_from_url(url: &Url) -> Result<DumpStream, LoadDumpError> {
     match url.scheme() {
         "http" => {

sqld/src/error.rs

@@ -3,7 +3,7 @@ use hyper::StatusCode;
 use tonic::metadata::errors::InvalidMetadataValueBytes;
 
 use crate::{
-    auth::AuthError, query_result_builder::QueryResultBuilderError,
+    auth::AuthError, namespace::ForkError, query_result_builder::QueryResultBuilderError,
     replication::replica::error::ReplicationError,
 };
 
@@ -77,6 +77,8 @@ pub enum Error {
     LoadDumpExistingDb,
     #[error("cannot restore database when conflicting params were provided")]
     ConflictingRestoreParameters,
+    #[error("failed to fork database: {0}")]
+    Fork(#[from] ForkError),
 }
 
 trait ResponseError: std::error::Error {
@@ -126,6 +128,7 @@ impl IntoResponse for Error {
             ReplicaRestoreError => self.format_err(StatusCode::BAD_REQUEST),
             LoadDumpExistingDb => self.format_err(StatusCode::BAD_REQUEST),
             ConflictingRestoreParameters => self.format_err(StatusCode::BAD_REQUEST),
+            Fork(e) => e.into_response(),
         }
     }
 }
@@ -181,3 +184,17 @@ impl IntoResponse for LoadDumpError {
         }
     }
 }
+
+impl ResponseError for ForkError {}
+
+impl IntoResponse for ForkError {
+    fn into_response(self) -> axum::response::Response {
+        match self {
+            ForkError::Internal(_)
+            | ForkError::Io(_)
+            | ForkError::LogRead(_)
+            | ForkError::CreateNamespace(_) => self.format_err(StatusCode::INTERNAL_SERVER_ERROR),
+            ForkError::ForkReplica => self.format_err(StatusCode::BAD_REQUEST),
+        }
+    }
+}

sqld/src/lib.rs

@@ -443,21 +443,18 @@ pub async fn init_bottomless_replicator(
         .ok_or_else(|| anyhow::anyhow!("Invalid db path"))?
         .to_owned();
     let mut replicator = bottomless::replicator::Replicator::with_options(path, options).await?;
-    let mut did_recover = false;
 
     let (generation, timestamp) = match restore_option {
         RestoreOption::Latest | RestoreOption::Dump(_) => (None, None),
         RestoreOption::Generation(generation) => (Some(*generation), None),
         RestoreOption::PointInTime(timestamp) => (None, Some(*timestamp)),
     };
 
-    match replicator.restore(generation, timestamp).await? {
+    let (action, did_recover) = replicator.restore(generation, timestamp).await?;
+    match action {
         bottomless::replicator::RestoreAction::SnapshotMainDbFile => {
             replicator.new_generation();
-            if let Some(handle) = replicator.snapshot_main_db_file(None).await? {
-                handle.await?;
-                did_recover = true;
-            }
+            replicator.snapshot_main_db_file(None).await?;
             // Restoration process only leaves the local WAL file if it was
             // detected to be newer than its remote counterpart.
             replicator.maybe_replicate_wal().await?

sqld/src/namespace/fork.rs

@@ -0,0 +1,133 @@
+use std::io::SeekFrom;
+use std::path::PathBuf;
+use std::sync::Arc;
+use std::time::Duration;
+
+use bytes::Bytes;
+use tokio::io::{AsyncReadExt, AsyncSeekExt, AsyncWriteExt};
+use tokio_stream::StreamExt;
+
+use crate::database::PrimaryDatabase;
+use crate::replication::frame::Frame;
+use crate::replication::primary::frame_stream::FrameStream;
+use crate::replication::{LogReadError, ReplicationLogger};
+
+use super::{MakeNamespace, RestoreOption};
+
+// FIXME: get this const from somewhere else (crate wide)
+const PAGE_SIZE: usize = 4096;
+
+type Result<T> = crate::Result<T, ForkError>;
+
+#[derive(Debug, thiserror::Error)]
+pub enum ForkError {
+    #[error("internal error: {0}")]
+    Internal(anyhow::Error),
+    #[error("io error: {0}")]
+    Io(#[from] std::io::Error),
+    #[error("failed to read frame from replication log: {0}")]
+    LogRead(anyhow::Error),
+    #[error("an error occured creating the namespace: {0}")]
+    CreateNamespace(Box<crate::error::Error>),
+    #[error("cannot fork a replica, try again with the primary.")]
+    ForkReplica,
+}
+
+impl From<tokio::task::JoinError> for ForkError {
+    fn from(e: tokio::task::JoinError) -> Self {
+        Self::Internal(e.into())
+    }
+}
+
+async fn write_frame(frame: Frame, temp_file: &mut tokio::fs::File) -> Result<()> {
+    let page_no = frame.header().page_no;
+    let page_pos = (page_no - 1) as usize * PAGE_SIZE;
+    temp_file.seek(SeekFrom::Start(page_pos as u64)).await?;
+    temp_file.write_all(frame.page()).await?;
+
+    Ok(())
+}
+
+pub struct ForkTask<'a> {
+    pub base_path: PathBuf,
+    pub logger: Arc<ReplicationLogger>,
+    pub dest_namespace: Bytes,
+    pub make_namespace: &'a dyn MakeNamespace<Database = PrimaryDatabase>,
+}
+
+impl ForkTask<'_> {
+    pub async fn fork(self) -> Result<super::Namespace<PrimaryDatabase>> {
+        match self.try_fork().await {
+            Err(e) => {
+                let _ = tokio::fs::remove_dir_all(
+                    self.base_path
+                        .join("dbs")
+                        .join(std::str::from_utf8(&self.dest_namespace).unwrap()),
+                )
+                .await;
+                Err(e)
+            }
+            Ok(ns) => Ok(ns),
+        }
+    }
+
+    async fn try_fork(&self) -> Result<super::Namespace<PrimaryDatabase>> {
+        // until what index to replicate
+        let base_path = self.base_path.clone();
+        let temp_dir =
+            tokio::task::spawn_blocking(move || tempfile::tempdir_in(base_path)).await??;
+        let mut data_file = tokio::fs::File::create(temp_dir.path().join("data")).await?;
+
+        let logger = self.logger.clone();
+        let end_frame_no = *logger.new_frame_notifier.borrow();
+        let mut next_frame_no = 0;
+        while next_frame_no < end_frame_no {
+            let mut streamer = FrameStream::new(logger.clone(), next_frame_no, false, None)
+                .map_err(|e| ForkError::LogRead(e.into()))?;
+            while let Some(res) = streamer.next().await {
+                match res {
+                    Ok(frame) => {
+                        next_frame_no = next_frame_no.max(frame.header().frame_no + 1);
+                        write_frame(frame, &mut data_file).await?;
+                    }
+                    Err(LogReadError::SnapshotRequired) => {
+                        let snapshot = loop {
+                            if let Some(snap) = logger
+                                .get_snapshot_file(next_frame_no)
+                                .map_err(ForkError::Internal)?
+                            {
+                                break snap;
+                            }
+
+                            // the snapshot must exist, it is just not yet available.
+                            tokio::time::sleep(Duration::from_millis(100)).await;
+                        };
+
+                        let iter = snapshot.frames_iter_from(next_frame_no);
+                        for frame in iter {
+                            let frame = frame.map_err(ForkError::LogRead)?;
+                            next_frame_no = next_frame_no.max(frame.header().frame_no + 1);
+                            write_frame(frame, &mut data_file).await?;
+                        }
+                    }
+                    Err(LogReadError::Ahead) => {
+                        unreachable!("trying to fork ahead of the forked database!")
+                    }
+                    Err(LogReadError::Error(e)) => return Err(ForkError::LogRead(e)),
+                }
+            }
+        }
+
+        let dest_path = self
+            .base_path
+            .join("dbs")
+            .join(std::str::from_utf8(&self.dest_namespace).unwrap());
+        tokio::fs::rename(temp_dir.path(), dest_path).await?;
+
+        tokio::io::stdin().read_i8().await.unwrap();
+        self.make_namespace
+            .create(self.dest_namespace.clone(), RestoreOption::Latest, true)
+            .await
+            .map_err(|e| ForkError::CreateNamespace(Box::new(e)))
+    }
+}