Merge pull request #1827 from tursodatabase/lucio/refactor-sync-ctx

libsql: refactor push and sync code

Author: LucioFranco

libsql/src/local/database.rs

@@ -33,7 +33,7 @@ pub struct Database {
     #[cfg(feature = "replication")]
     pub replication_ctx: Option<ReplicationContext>,
     #[cfg(feature = "sync")]
-    pub sync_ctx: Option<SyncContext>,
+    pub sync_ctx: Option<tokio::sync::Mutex<SyncContext>>,
 }
 
 impl Database {
@@ -143,7 +143,10 @@ impl Database {
             endpoint
         };
         let mut db = Database::open(&db_path, flags)?;
-        db.sync_ctx = Some(SyncContext::new(endpoint, Some(auth_token)));
+        db.sync_ctx = Some(tokio::sync::Mutex::new(SyncContext::new(
+            endpoint,
+            Some(auth_token),
+        )));
         Ok(db)
     }
 
@@ -383,7 +386,7 @@ impl Database {
     #[cfg(feature = "sync")]
     /// Push WAL frames to remote.
     pub async fn push(&self) -> Result<crate::database::Replicated> {
-        let sync_ctx = self.sync_ctx.as_ref().unwrap();
+        let sync_ctx = self.sync_ctx.as_ref().unwrap().lock().await;
         let conn = self.connect()?;
 
         let page_size = {
@@ -398,17 +401,20 @@ impl Database {
         let max_frame_no = conn.wal_frame_count();
 
         let generation = 1; // TODO: Probe from WAL.
-        let start_frame_no = sync_ctx.durable_frame_num + 1;
+        let start_frame_no = sync_ctx.durable_frame_num() + 1;
         let end_frame_no = max_frame_no;
 
         let mut frame_no = start_frame_no;
         while frame_no <= end_frame_no {
+            let frame = conn.wal_get_frame(frame_no, page_size)?;
+
             // The server returns its maximum frame number. To avoid resending
             // frames the server already knows about, we need to update the
             // frame number to the one returned by the server.
-            let max_frame_no = self
-                .push_one_frame(&conn, &sync_ctx, generation, frame_no, page_size)
+            let max_frame_no = sync_ctx
+                .push_one_frame(frame.to_vec(), generation, frame_no)
                 .await?;
+
             if max_frame_no > frame_no {
                 frame_no = max_frame_no;
             }
@@ -422,72 +428,6 @@ impl Database {
         })
     }
 
-    #[cfg(feature = "sync")]
-    async fn push_one_frame(
-        &self,
-        conn: &Connection,
-        sync_ctx: &SyncContext,
-        generation: u32,
-        frame_no: u32,
-        page_size: u32,
-    ) -> Result<u32> {
-        let frame = conn.wal_get_frame(frame_no, page_size)?;
-
-        let uri = format!(
-            "{}/sync/{}/{}/{}",
-            sync_ctx.sync_url,
-            generation,
-            frame_no,
-            frame_no + 1
-        );
-        let max_frame_no = self
-            .push_with_retry(
-                uri,
-                &sync_ctx.auth_token,
-                frame.to_vec(),
-                sync_ctx.max_retries,
-            )
-            .await?;
-        Ok(max_frame_no)
-    }
-
-    #[cfg(feature = "sync")]
-    async fn push_with_retry(
-        &self,
-        uri: String,
-        auth_token: &Option<String>,
-        frame: Vec<u8>,
-        max_retries: usize,
-    ) -> Result<u32> {
-        let mut nr_retries = 0;
-        loop {
-            let client = reqwest::Client::new();
-            let mut builder = client.post(uri.to_owned());
-            match auth_token {
-                Some(ref auth_token) => {
-                    builder = builder
-                        .header("Authorization", format!("Bearer {}", auth_token.to_owned()));
-                }
-                None => {}
-            }
-            let res = builder.body(frame.to_vec()).send().await.unwrap();
-            if res.status().is_success() {
-                let resp = res.json::<serde_json::Value>().await.unwrap();
-                let max_frame_no = resp.get("max_frame_no").unwrap().as_u64().unwrap();
-                return Ok(max_frame_no as u32);
-            }
-            if nr_retries > max_retries {
-                return Err(crate::errors::Error::ConnectionFailed(format!(
-                    "Failed to push frame: {}",
-                    res.status()
-                )));
-            }
-            let delay = std::time::Duration::from_millis(100 * (1 << nr_retries));
-            tokio::time::sleep(delay).await;
-            nr_retries += 1;
-        }
-    }
-
     pub(crate) fn path(&self) -> &str {
         &self.db_path
     }

libsql/src/sync.rs

@@ -1,9 +1,11 @@
+use crate::Result;
+
 const DEFAULT_MAX_RETRIES: usize = 5;
 pub struct SyncContext {
-    pub sync_url: String,
-    pub auth_token: Option<String>,
-    pub max_retries: usize,
-    pub durable_frame_num: u32,
+    sync_url: String,
+    auth_token: Option<String>,
+    max_retries: usize,
+    durable_frame_num: u32,
 }
 
 impl SyncContext {
@@ -15,4 +17,62 @@ impl SyncContext {
             max_retries: DEFAULT_MAX_RETRIES,
         }
     }
+
+    pub(crate) async fn push_one_frame(
+        &self,
+        frame: Vec<u8>,
+        generation: u32,
+        frame_no: u32,
+    ) -> Result<u32> {
+        let uri = format!(
+            "{}/sync/{}/{}/{}",
+            self.sync_url,
+            generation,
+            frame_no,
+            frame_no + 1
+        );
+        let max_frame_no = self
+            .push_with_retry(uri, frame.to_vec(), self.max_retries)
+            .await?;
+
+        Ok(max_frame_no)
+    }
+
+    async fn push_with_retry(
+        &self,
+        uri: String,
+        frame: Vec<u8>,
+        max_retries: usize,
+    ) -> Result<u32> {
+        let mut nr_retries = 0;
+        loop {
+            let client = reqwest::Client::new();
+            let mut builder = client.post(uri.to_owned());
+            match &self.auth_token {
+                Some(ref auth_token) => {
+                    builder = builder.header("Authorization", format!("Bearer {}", auth_token));
+                }
+                None => {}
+            }
+            let res = builder.body(frame.to_vec()).send().await.unwrap();
+            if res.status().is_success() {
+                let resp = res.json::<serde_json::Value>().await.unwrap();
+                let max_frame_no = resp.get("max_frame_no").unwrap().as_u64().unwrap();
+                return Ok(max_frame_no as u32);
+            }
+            if nr_retries > max_retries {
+                return Err(crate::errors::Error::ConnectionFailed(format!(
+                    "Failed to push frame: {}",
+                    res.status()
+                )));
+            }
+            let delay = std::time::Duration::from_millis(100 * (1 << nr_retries));
+            tokio::time::sleep(delay).await;
+            nr_retries += 1;
+        }
+    }
+
+    pub(crate) fn durable_frame_num(&self) -> u32 {
+        self.durable_frame_num
+    }
 }