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/replication/remote_client.rs

@@ -18,7 +18,7 @@ use tonic::metadata::AsciiMetadataValue;
 use tonic::{Response, Status};
 use zerocopy::FromBytes;
 
-async fn time<O>(fut: impl Future<Output = O>) -> (O, Duration) {
+pub async fn time<O>(fut: impl Future<Output = O>) -> (O, Duration) {
     let before = Instant::now();
     let out = fut.await;
     (out, before.elapsed())

libsql/src/sync.rs

@@ -2,6 +2,7 @@ use crate::{local::Connection, util::ConnectorService, Error, Result};
 
 use std::path::Path;
 
+use crate::replication::remote_client::time;
 use bytes::Bytes;
 use chrono::Utc;
 use http::{HeaderValue, StatusCode};
@@ -20,6 +21,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 +68,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 +102,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 +126,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 +159,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 +181,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 +191,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,20 +424,39 @@ 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.
             if res.status() == StatusCode::BAD_REQUEST
                 || res.status() == StatusCode::INTERNAL_SERVER_ERROR
             {
+                let status = res.status();
                 let res_body = hyper::body::to_bytes(res.into_body())
                     .await
                     .map_err(SyncError::HttpBody)?;
-
+                tracing::trace!(
+                    "server returned: {} body: {}",
+                    status,
+                    String::from_utf8_lossy(&res_body[..])
+                );
                 let resp = serde_json::from_slice::<serde_json::Value>(&res_body[..])
                     .map_err(SyncError::JsonDecode)?;
 
@@ -650,22 +676,33 @@ impl SyncContext {
 
         let req = req.body(Body::empty()).expect("valid request");
 
-        let res = self
-            .client
-            .request(req)
-            .await
-            .map_err(SyncError::HttpDispatch)?;
+        let (res, http_duration) = time(self.client.request(req)).await;
+        let res = res.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)?;
+            tracing::error!(
+                "failed to pull db file from remote server, status={}, body={}, url={}, duration={:?}",
+                status,
+                String::from_utf8_lossy(&body),
+                uri,
+                http_duration
+            );
             return Err(
                 SyncError::PullFrame(status, String::from_utf8_lossy(&body).to_string()).into(),
             );
         }
 
+        tracing::debug!(
+            "pulled db file from remote server, status={}, url={}, duration={:?}",
+            res.status(),
+            uri,
+            http_duration
+        );
+
         // todo: do streaming write to the disk
         let bytes = hyper::body::to_bytes(res.into_body())
             .await
@@ -887,6 +924,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 +940,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 +991,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;
             }