Performance: optimise offline sync

Currently, the sync code pulls all the generations and frames one by one.
However, this is very inefficient, for databases with many generations
it could take hours.

This patch adds a optimisation that it pulls the latest generation
by calling `export` endpoint and boostraps the db file quickly.

My testing shows insane improvements. For a db with 30 generations,
and about 100 frames in each gen, it has brought down the initial sync
from 1hr to 1 minute.

Author: avinassh

libsql/src/sync.rs

@@ -60,6 +60,10 @@ pub enum SyncError {
     RedirectHeader(http::header::ToStrError),
     #[error("redirect response with no location header")]
     NoRedirectLocationHeader,
+    #[error("failed to pull db export: status={0}, error={1}")]
+    PullDb(StatusCode, String),
+    #[error("server returned a lower generation than local: local={0}, remote={1}")]
+    InvalidLocalGeneration(u32, u32),
 }
 
 impl SyncError {
@@ -86,6 +90,11 @@ pub enum PullResult {
     EndOfGeneration { max_generation: u32 },
 }
 
+#[derive(serde::Deserialize)]
+struct InfoResult {
+    current_generation: u32,
+}
+
 pub struct SyncContext {
     db_path: String,
     client: hyper::Client<ConnectorService, Body>,
@@ -97,6 +106,9 @@ pub struct SyncContext {
     durable_generation: u32,
     /// Represents the max_frame_no from the server.
     durable_frame_num: u32,
+    /// whenever sync is called very first time, we will call the remote server
+    /// to get the generation information and sync the db file if needed
+    initial_server_sync: bool,
 }
 
 impl SyncContext {
@@ -125,6 +137,7 @@ impl SyncContext {
             client,
             durable_generation: 1,
             durable_frame_num: 0,
+            initial_server_sync: false,
         };
 
         if let Err(e) = me.read_metadata().await {
@@ -458,6 +471,114 @@ impl SyncContext {
 
         Ok(())
     }
+
+    /// get_remote_info calls the remote server to get the current generation information.
+    async fn get_remote_info(&self) -> Result<InfoResult> {
+        let uri = format!("{}/info", self.sync_url);
+        let mut req = http::Request::builder().method("GET").uri(&uri);
+
+        if let Some(auth_token) = &self.auth_token {
+            req = req.header("Authorization", auth_token);
+        }
+
+        let req = req.body(Body::empty()).expect("valid request");
+
+        let res = self
+            .client
+            .request(req)
+            .await
+            .map_err(SyncError::HttpDispatch)?;
+
+        if !res.status().is_success() {
+            let status = res.status();
+            let body = hyper::body::to_bytes(res.into_body())
+                .await
+                .map_err(SyncError::HttpBody)?;
+            return Err(
+                SyncError::PullDb(status, String::from_utf8_lossy(&body).to_string()).into(),
+            );
+        }
+
+        let body = hyper::body::to_bytes(res.into_body())
+            .await
+            .map_err(SyncError::HttpBody)?;
+
+        let info = serde_json::from_slice(&body).map_err(SyncError::JsonDecode)?;
+
+        Ok(info)
+    }
+
+    async fn sync_db_if_needed(&mut self, generation: u32) -> Result<()> {
+        // we will get the export file only if the remote generation is different from the one we have
+        if generation == self.durable_generation {
+            return Ok(());
+        }
+        // somehow we are ahead of the remote in generations. following should not happen because
+        // we checkpoint only if the remote server tells us to do so.
+        if self.durable_generation > generation {
+            tracing::error!(
+                "server returned a lower generation than what we have: sent={}, got={}",
+                self.durable_generation,
+                generation
+            );
+            return Err(
+                SyncError::InvalidLocalGeneration(self.durable_generation, generation).into(),
+            );
+        }
+        self.sync_db(generation).await
+    }
+
+    /// sync_db will download the db file from the remote server and replace the local file.
+    async fn sync_db(&mut self, generation: u32) -> Result<()> {
+        let uri = format!("{}/export/{}", self.sync_url, generation);
+        let mut req = http::Request::builder().method("GET").uri(&uri);
+
+        if let Some(auth_token) = &self.auth_token {
+            req = req.header("Authorization", auth_token);
+        }
+
+        let req = req.body(Body::empty()).expect("valid request");
+
+        let res = self
+            .client
+            .request(req)
+            .await
+            .map_err(SyncError::HttpDispatch)?;
+
+        if !res.status().is_success() {
+            let status = res.status();
+            let body = hyper::body::to_bytes(res.into_body())
+                .await
+                .map_err(SyncError::HttpBody)?;
+            return Err(
+                SyncError::PullFrame(status, String::from_utf8_lossy(&body).to_string()).into(),
+            );
+        }
+
+        // todo: do streaming write to the disk
+        let bytes = hyper::body::to_bytes(res.into_body())
+            .await
+            .map_err(SyncError::HttpBody)?;
+
+
+        // since we are starting from new generation, we need to clean up and remove the old files
+        // we will remove `.db-wal` and `.db-shm` manually. The `.db` file will be replaced by the
+        // atomic write.
+        let db_path_wal = format!("{}-wal", self.db_path);
+        let db_path_shm = format!("{}-shm", self.db_path);
+        tokio::fs::remove_file(db_path_wal)
+            .await
+            .map_err(SyncError::io("remove wal file"))?;
+        tokio::fs::remove_file(db_path_shm)
+            .await
+            .map_err(SyncError::io("remove shm file"))?;
+
+        atomic_write(&self.db_path, &bytes).await?;
+        self.durable_generation = generation;
+        self.durable_frame_num = 0;
+        self.write_metadata().await?;
+        Ok(())
+    }
 }
 
 #[derive(serde::Serialize, serde::Deserialize, Debug)]
@@ -555,6 +676,19 @@ pub async fn sync_offline(
             Err(e) => Err(e),
         }
     } else {
+        // todo: we are checking with the remote server only during initialisation. ideally,
+        // we should check everytime we try to sync with the remote server. However, we need to close
+        // all the ongoing connections since we replace `.db` file and remove the `.db-wal` file
+        if !sync_ctx.initial_server_sync {
+            // sync is being called first time. so we will call remote, get the generation information
+            // if we are lagging behind, then we will call the export API and get to the latest
+            // generation directly.
+            let info = sync_ctx.get_remote_info().await?;
+            sync_ctx
+                .sync_db_if_needed(info.current_generation)
+                .await?;
+            sync_ctx.initial_server_sync = true;
+        }
         try_pull(sync_ctx, conn).await
     }
     .or_else(|err| {