Optimise sync_pull to pull frames in batches

Previously, we pulled frames one by one. This patch changes it pull
frames in batches. Currently, the batch size is set to 128 (the
maximum supported by the server)

Author: avinassh

libsql/src/sync.rs

@@ -20,6 +20,7 @@ const METADATA_VERSION: u32 = 0;
 
 const DEFAULT_MAX_RETRIES: usize = 5;
 const DEFAULT_PUSH_BATCH_SIZE: u32 = 128;
+const DEFAULT_PULL_BATCH_SIZE: u32 = 128;
 
 #[derive(thiserror::Error, Debug)]
 #[non_exhaustive]
@@ -66,6 +67,8 @@ pub enum SyncError {
     InvalidLocalGeneration(u32, u32),
     #[error("invalid local state: {0}")]
     InvalidLocalState(String),
+    #[error("server returned invalid length of frames: {0}")]
+    InvalidPullFrameBytes(usize),
 }
 
 impl SyncError {
@@ -98,8 +101,8 @@ pub enum PushStatus {
 }
 
 pub enum PullResult {
-    /// A frame was successfully pulled.
-    Frame(Bytes),
+    /// Frames were successfully pulled.
+    Frames(Bytes),
     /// We've reached the end of the generation.
     EndOfGeneration { max_generation: u32 },
 }
@@ -122,6 +125,7 @@ pub struct SyncContext {
     auth_token: Option<HeaderValue>,
     max_retries: usize,
     push_batch_size: u32,
+    pull_batch_size: u32,
     /// The current durable generation.
     durable_generation: u32,
     /// Represents the max_frame_no from the server.
@@ -154,6 +158,7 @@ impl SyncContext {
             auth_token,
             max_retries: DEFAULT_MAX_RETRIES,
             push_batch_size: DEFAULT_PUSH_BATCH_SIZE,
+            pull_batch_size: DEFAULT_PULL_BATCH_SIZE,
             client,
             durable_generation: 0,
             durable_frame_num: 0,
@@ -175,7 +180,7 @@ impl SyncContext {
     }
 
     #[tracing::instrument(skip(self))]
-    pub(crate) async fn pull_one_frame(
+    pub(crate) async fn pull_frames(
         &mut self,
         generation: u32,
         frame_no: u32,
@@ -185,9 +190,10 @@ impl SyncContext {
             self.sync_url,
             generation,
             frame_no,
-            frame_no + 1
+            // the server expects the range of [start, end) frames, i.e. end is exclusive
+            frame_no + self.pull_batch_size
         );
-        tracing::debug!("pulling frame");
+        tracing::debug!("pulling frame (uri={})", uri);
         self.pull_with_retry(uri, self.max_retries).await
     }
 
@@ -417,10 +423,24 @@ impl SyncContext {
                 .map_err(SyncError::HttpDispatch)?;
 
             if res.status().is_success() {
-                let frame = hyper::body::to_bytes(res.into_body())
+                let frames = hyper::body::to_bytes(res.into_body())
                     .await
                     .map_err(SyncError::HttpBody)?;
-                return Ok(PullResult::Frame(frame));
+                // a success result should always return some frames
+                if frames.is_empty() {
+                    tracing::error!("server returned empty frames in pull response");
+                    return Err(SyncError::InvalidPullFrameBytes(0).into());
+                }
+                // the minimum payload size cannot be less than a single frame
+                if frames.len() < FRAME_SIZE {
+                    tracing::error!(
+                        "server returned frames with invalid length: {} < {}",
+                        frames.len(),
+                        FRAME_SIZE
+                    );
+                    return Err(SyncError::InvalidPullFrameBytes(frames.len()).into());
+                }
+                return Ok(PullResult::Frames(frames));
             }
             // BUG ALERT: The server returns a 500 error if the remote database is empty.
             // This is a bug and should be fixed.
@@ -887,6 +907,11 @@ async fn try_push(
     })
 }
 
+/// PAGE_SIZE used by the sync / diskless server
+const PAGE_SIZE: usize = 4096;
+const FRAME_HEADER_SIZE: usize = 24;
+const FRAME_SIZE: usize = PAGE_SIZE + FRAME_HEADER_SIZE;
+
 pub async fn try_pull(
     sync_ctx: &mut SyncContext,
     conn: &Connection,
@@ -898,10 +923,39 @@ pub async fn try_pull(
     loop {
         let generation = sync_ctx.durable_generation();
         let frame_no = sync_ctx.durable_frame_num() + 1;
-        match sync_ctx.pull_one_frame(generation, frame_no).await {
-            Ok(PullResult::Frame(frame)) => {
-                insert_handle.insert(&frame)?;
-                sync_ctx.durable_frame_num = frame_no;
+        match sync_ctx.pull_frames(generation, frame_no).await {
+            Ok(PullResult::Frames(frames)) => {
+                tracing::debug!(
+                    "pull_frames: generation={}, start_frame_no={} (batch_size={}), frame_size={}",
+                    generation,
+                    frame_no,
+                    sync_ctx.pull_batch_size,
+                    frames.len(),
+                );
+                if frames.len() % FRAME_SIZE != 0 {
+                    tracing::error!(
+                        "frame size {} is not a multiple of the expected size {}",
+                        frames.len(),
+                        FRAME_SIZE,
+                    );
+                    return Err(SyncError::InvalidPullFrameBytes(frames.len()).into());
+                }
+                for chunk in frames.chunks(FRAME_SIZE) {
+                    tracing::debug!(
+                        "inserting frame (frame_no={})",
+                        sync_ctx.durable_frame_num + 1
+                    );
+                    let r = insert_handle.insert(&chunk);
+                    if let Err(e) = r {
+                        tracing::debug!(
+                            "insert error (frame= {}) : {:?}",
+                            sync_ctx.durable_frame_num + 1,
+                            e
+                        );
+                        return Err(e);
+                    }
+                    sync_ctx.durable_frame_num += 1;
+                }
             }
             Ok(PullResult::EndOfGeneration { max_generation }) => {
                 // If there are no more generations to pull, we're done.
@@ -920,7 +974,7 @@ pub async fn try_pull(
                 insert_handle.begin()?;
             }
             Err(e) => {
-                tracing::debug!("pull_one_frame error: {:?}", e);
+                tracing::debug!("pull_frames error: {:?}", e);
                 err.replace(e);
                 break;
             }