Merge pull request #1826 from tursodatabase/lucio/add-sys-wal-fn

libsql: add wal frame_count and get_frame safe fn
tursodatabase · Nov 19, 2024 · b80f611 · b80f611
2 parents fd44afa + 682ac9d
commit b80f611
Show file tree

Hide file tree

Showing 2 changed files with 40 additions and 18 deletions.
diff --git a/libsql/src/local/connection.rs b/libsql/src/local/connection.rs
@@ -445,6 +445,43 @@ impl Connection {
             }
         }
     }
+
+    pub(crate) fn wal_frame_count(&self) -> u32 {
+        let mut max_frame_no: std::os::raw::c_uint = 0;
+        unsafe { libsql_sys::ffi::libsql_wal_frame_count(self.handle(), &mut max_frame_no) };
+
+        max_frame_no
+    }
+
+    pub(crate) fn wal_get_frame(&self, frame_no: u32, page_size: u32) -> Result<bytes::BytesMut> {
+        use bytes::BufMut;
+
+        let frame_size: usize = 24 + page_size as usize;
+
+        // Use a BytesMut to provide cheaper clones of frame data (think retries)
+        // and more efficient buffer usage for extracting wal frames and spliting them off.
+        let mut buf = bytes::BytesMut::with_capacity(frame_size);
+
+        let rc = unsafe {
+            libsql_sys::ffi::libsql_wal_get_frame(
+                self.handle(),
+                frame_no,
+                buf.chunk_mut().as_mut_ptr() as *mut _,
+                frame_size as u32,
+            )
+        };
+
+        if rc != 0 {
+            return Err(crate::errors::Error::SqliteFailure(
+                rc as std::ffi::c_int,
+                format!("Failed to get frame: {}", frame_no),
+            ));
+        }
+
+        unsafe { buf.advance_mut(frame_size) };
+
+        Ok(buf)
+    }
 }
 
 impl fmt::Debug for Connection {

diff --git a/libsql/src/local/database.rs b/libsql/src/local/database.rs
@@ -395,8 +395,7 @@ impl Database {
             page_size
         };
 
-        let mut max_frame_no: std::os::raw::c_uint = 0;
-        unsafe { libsql_sys::ffi::libsql_wal_frame_count(conn.handle(), &mut max_frame_no) };
+        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;
@@ -432,22 +431,8 @@ impl Database {
         frame_no: u32,
         page_size: u32,
     ) -> Result<u32> {
-        let frame_size: usize = 24 + page_size as usize;
-        let frame = vec![0; frame_size];
-        let rc = unsafe {
-            libsql_sys::ffi::libsql_wal_get_frame(
-                conn.handle(),
-                frame_no,
-                frame.as_ptr() as *mut _,
-                frame_size as u32,
-            )
-        };
-        if rc != 0 {
-            return Err(crate::errors::Error::SqliteFailure(
-                rc as std::ffi::c_int,
-                format!("Failed to get frame: {}", frame_no),
-            ));
-        }
+        let frame = conn.wal_get_frame(frame_no, page_size)?;
+
         let uri = format!(
             "{}/sync/{}/{}/{}",
             sync_ctx.sync_url,