libsql-{sqlite3,ffi}: Add xFrameCount() to the virtual WAL API

penberg · penberg · commit 3ac4aa1689e8 · 2024-10-29T13:30:27.000+02:00
diff --git a/libsql-ffi/bundled/bindings/bindgen.rs b/libsql-ffi/bundled/bindings/bindgen.rs
@@ -3316,6 +3316,13 @@ pub struct libsql_wal_methods {
             aWalData: *mut ::std::os::raw::c_uint,
         ) -> ::std::os::raw::c_int,
     >,
+    pub xFrameCount: ::std::option::Option<
+        unsafe extern "C" fn(
+            pWal: *mut wal_impl,
+            arg1: ::std::os::raw::c_int,
+            arg2: *mut ::std::os::raw::c_uint,
+        ) -> ::std::os::raw::c_int,
+    >,
     pub xFrames: ::std::option::Option<
         unsafe extern "C" fn(
             pWal: *mut wal_impl,
diff --git a/libsql-ffi/bundled/bindings/session_bindgen.rs b/libsql-ffi/bundled/bindings/session_bindgen.rs
@@ -3822,6 +3822,13 @@ pub struct libsql_wal_methods {
             aWalData: *mut ::std::os::raw::c_uint,
         ) -> ::std::os::raw::c_int,
     >,
+    pub xFrameCount: ::std::option::Option<
+        unsafe extern "C" fn(
+            pWal: *mut wal_impl,
+            arg1: ::std::os::raw::c_int,
+            arg2: *mut ::std::os::raw::c_uint,
+        ) -> ::std::os::raw::c_int,
+    >,
     pub xFrames: ::std::option::Option<
         unsafe extern "C" fn(
             pWal: *mut wal_impl,
diff --git a/libsql-replication/src/injector/sqlite_injector/injector_wal.rs b/libsql-replication/src/injector/sqlite_injector/injector_wal.rs
@@ -134,6 +134,10 @@ impl Wal for InjectorWal {
         self.inner.savepoint_undo(rollback_data)
     }
 
+    fn frame_count(&self, locked: i32) -> Result<u32> {
+        self.inner.frame_count(locked)
+    }
+
     fn insert_frames(
         &mut self,
         page_size: c_int,
diff --git a/libsql-sqlite3/src/main.c b/libsql-sqlite3/src/main.c
@@ -2449,6 +2449,32 @@ int libsql_wal_disable_checkpoint(sqlite3 *db) {
   return SQLITE_OK;
 }
 
+/*
+** Return the number of frames in the WAL of the given database.
+*/
+int libsql_wal_frame_count(
+  sqlite3* db,
+  unsigned int *pnFrame
+){
+  int rc = SQLITE_OK;
+  Pager *pPager;
+
+#ifdef SQLITE_OMIT_WAL
+  *pnFrame = 0;
+  return SQLITE_OK;
+#else
+#ifdef SQLITE_ENABLE_API_ARMOR
+  if( !sqlite3SafetyCheckOk(db) ) return SQLITE_MISUSE_BKPT;
+#endif
+
+  sqlite3_mutex_enter(db->mutex);
+  pPager = sqlite3BtreePager(db->aDb[0].pBt);
+  rc = sqlite3PagerWalFrameCount(pPager, pnFrame);
+  sqlite3_mutex_leave(db->mutex);
+  return rc;
+#endif
+}
+
 /*
 ** Register a function to be invoked prior to each autovacuum that
 ** determines the number of pages to vacuum.
diff --git a/libsql-sqlite3/src/pager.c b/libsql-sqlite3/src/pager.c
@@ -7771,6 +7771,19 @@ int sqlite3PagerCloseWal(Pager *pPager, sqlite3 *db){
   return rc;
 }
 
+/**
+** Return the number of frames in the WAL file.
+**
+** If the pager is not in WAL mode or we failed to obtain an exclusive write lock, returns -1.
+**/
+int sqlite3PagerWalFrameCount(Pager *pPager, unsigned int *pnFrames){
+  if( pagerUseWal(pPager) ){
+    return pPager->wal->methods.xFrameCount(pPager->wal->pData, 0, pnFrames);
+  }else{
+    return SQLITE_ERROR;
+  }
+}
+
 #ifdef SQLITE_ENABLE_SETLK_TIMEOUT
 /*
 ** If pager pPager is a wal-mode database not in exclusive locking mode,
diff --git a/libsql-sqlite3/src/pager.h b/libsql-sqlite3/src/pager.h
@@ -133,6 +133,8 @@ int sqlite3PagerReadFileheader(Pager*, int, unsigned char*);
 void sqlite3PagerSetBusyHandler(Pager*, int(*)(void *), void *);
 int sqlite3PagerSetPagesize(Pager*, u32*, int);
 Pgno sqlite3PagerMaxPageCount(Pager*, Pgno);
+int sqlite3PagerWalFrameCount(Pager *, unsigned int *);
+
 void sqlite3PagerSetCachesize(Pager*, int);
 int sqlite3PagerSetSpillsize(Pager*, int);
 void sqlite3PagerSetMmapLimit(Pager *, sqlite3_int64);
diff --git a/libsql-sqlite3/src/sqlite.h.in b/libsql-sqlite3/src/sqlite.h.in
@@ -10564,6 +10564,15 @@ void *libsql_close_hook(sqlite3 *db, void (*xClose)(void *pCtx, sqlite3 *db), vo
 */
 int libsql_wal_disable_checkpoint(sqlite3 *db);
 
+/*
+** CAPI3REF: Get the number of frames in the WAL file
+** METHOD: sqlite3
+**
+** ^The [libsql_wal_frame_count(D,P)] interface returns the number of frames
+** in the WAL file for [database connection] D into *P.
+*/
+int libsql_wal_frame_count(sqlite3*, unsigned int*);
+
 /*
 ** CAPI3REF: Low-level system error code
 ** METHOD: sqlite3
diff --git a/libsql-sqlite3/src/wal.c b/libsql-sqlite3/src/wal.c
@@ -4001,6 +4001,19 @@ static int walFrames(
   return rc;
 }
 
+int sqlite3WalFrameCount(Wal *pWal, int locked, unsigned int *pnFrames){
+  int rc = SQLITE_OK;
+  if( locked==0 ) {
+    rc = walLockExclusive(pWal, WAL_WRITE_LOCK, 1);
+    if (rc != SQLITE_OK) return rc;
+  }
+  *pnFrames = pWal->hdr.mxFrame;
+  if( locked==0 ) {
+    walUnlockExclusive(pWal, WAL_WRITE_LOCK, 1);
+  }
+  return SQLITE_OK;
+}
+
 /* 
 ** Write a set of frames to the log. The caller must hold the write-lock
 ** on the log file (obtained using sqlite3WalBeginWriteTransaction()).
@@ -4506,6 +4519,7 @@ static int sqlite3WalOpen(
     out->methods.xUndo = (int (*)(wal_impl *, int (*)(void *, unsigned int), void *))sqlite3WalUndo;
     out->methods.xSavepoint = (void (*)(wal_impl *, unsigned int *))sqlite3WalSavepoint;
     out->methods.xSavepointUndo = (int (*)(wal_impl *, unsigned int *))sqlite3WalSavepointUndo;
+    out->methods.xFrameCount = (int (*)(wal_impl *, int, unsigned int *))sqlite3WalFrameCount;
     out->methods.xFrames = (int (*)(wal_impl *, int, libsql_pghdr *, unsigned int, int, int, int *))sqlite3WalFrames;
     out->methods.xCheckpoint = (int (*)(wal_impl *, sqlite3 *, int, int (*)(void *), void *, int, int, unsigned char *, int *, int *, int (*)(void*, int, const unsigned char*, int, int, int), void*))sqlite3WalCheckpoint;
     out->methods.xCallback = (int (*)(wal_impl *))sqlite3WalCallback;
diff --git a/libsql-sqlite3/src/wal.h b/libsql-sqlite3/src/wal.h
@@ -75,6 +75,9 @@ typedef struct libsql_wal_methods {
   ** response to a ROLLBACK TO command. */
   int (*xSavepointUndo)(wal_impl* pWal, unsigned int *aWalData);
 
+  /* Return the number of frames in the WAL */
+  int (*xFrameCount)(wal_impl* pWal, int, unsigned int *);
+
   /* Write a frame or frames to the log. */
   int (*xFrames)(wal_impl* pWal, int, libsql_pghdr *, unsigned int, int, int, int*);
 
diff --git a/libsql-sqlite3/test/rust_suite/src/virtual_wal.rs b/libsql-sqlite3/test/rust_suite/src/virtual_wal.rs
@@ -98,6 +98,10 @@ mod tests {
             self.0.read_frame(frame_no, buffer)
         }
 
+        fn frame_count(&self, locked: i32) -> libsql_sys::wal::Result<u32> {
+            self.0.frame_count(locked)
+        }
+
         fn db_size(&self) -> u32 {
             self.0.db_size()
         }
diff --git a/libsql-storage/src/lib.rs b/libsql-storage/src/lib.rs
@@ -266,6 +266,10 @@ impl Wal for DurableWal {
         Ok(())
     }
 
+    fn frame_count(&self, _locked: i32) -> Result<u32> {
+        todo!()
+    }
+
     fn db_size(&self) -> u32 {
         let size = self.local_cache.get_max_frame_num().unwrap();
         trace!("DurableWal::db_size() => {}", size);
diff --git a/libsql-sys/src/wal/either.rs b/libsql-sys/src/wal/either.rs
@@ -75,6 +75,12 @@ macro_rules! create_either {
                 }
             }
 
+            fn frame_count(&self, locked: i32) -> super::Result<u32> {
+                match self {
+                    $( $name::$t(inner) => inner.frame_count(locked) ),*
+                }
+            }
+
             fn insert_frames(
                 &mut self,
                 page_size: std::ffi::c_int,
diff --git a/libsql-sys/src/wal/ffi.rs b/libsql-sys/src/wal/ffi.rs
@@ -1,4 +1,4 @@
-use std::ffi::{c_char, c_int, c_longlong, c_void, CStr};
+use std::ffi::{c_char, c_int, c_longlong, c_uint, c_void, CStr};
 use std::num::NonZeroU32;
 use std::ptr::null;
 
@@ -29,6 +29,7 @@ pub(crate) fn construct_libsql_wal<W: Wal>(wal: *mut W) -> libsql_wal {
             xUndo: Some(undo::<W>),
             xSavepoint: Some(savepoint::<W>),
             xSavepointUndo: Some(savepoint_undo::<W>),
+            xFrameCount: Some(frame_count::<W>),
             xFrames: Some(frames::<W>),
             xCheckpoint: Some(checkpoint::<W>),
             xCallback: Some(callback::<W>),
@@ -276,6 +277,25 @@ pub unsafe extern "C" fn savepoint_undo<T: Wal>(wal: *mut wal_impl, wal_data: *m
     }
 }
 
+pub unsafe extern "C" fn frame_count<T: Wal>(
+    wal: *mut wal_impl,
+    locked: i32,
+    out: *mut c_uint,
+) -> c_int {
+    let this = &mut (*(wal as *mut T));
+    match this.frame_count(locked) {
+        Ok(n) => {
+            if !out.is_null() {
+                unsafe {
+                    *out = n as _;
+                }
+            }
+            SQLITE_OK
+        }
+        Err(code) => code.extended_code,
+    }
+}
+
 pub unsafe extern "C" fn frames<T: Wal>(
     wal: *mut wal_impl,
     page_size: c_int,
diff --git a/libsql-sys/src/wal/mod.rs b/libsql-sys/src/wal/mod.rs
@@ -197,6 +197,8 @@ pub trait Wal {
     fn savepoint(&mut self, rollback_data: &mut [u32]);
     fn savepoint_undo(&mut self, rollback_data: &mut [u32]) -> Result<()>;
 
+    fn frame_count(&self, locked: i32) -> Result<u32>;
+
     /// Insert frames in the wal. On commit, returns the number of inserted frames for that
     /// transaction, or 0 for non-commit calls.
     fn insert_frames(
diff --git a/libsql-sys/src/wal/sqlite3_wal.rs b/libsql-sys/src/wal/sqlite3_wal.rs
@@ -272,6 +272,18 @@ impl Wal for Sqlite3Wal {
         }
     }
 
+    fn frame_count(&self, locked: i32) -> Result<u32> {
+        let mut out: u32 = 0;
+        let rc = unsafe {
+            (self.inner.methods.xFrameCount.unwrap())(self.inner.pData, locked, &mut out)
+        };
+        if rc != 0 {
+            Err(Error::new(rc))
+        } else {
+            Ok(out)
+        }
+    }
+
     fn insert_frames(
         &mut self,
         page_size: c_int,
diff --git a/libsql-sys/src/wal/wrapper.rs b/libsql-sys/src/wal/wrapper.rs
@@ -85,6 +85,10 @@ impl<T: WrapWal<W>, W: Wal> Wal for WalRef<T, W> {
         unsafe { (*self.wrapper).savepoint_undo(&mut *self.wrapped, rollback_data) }
     }
 
+    fn frame_count(&self, locked: i32) -> super::Result<u32> {
+        unsafe { (*self.wrapper).frame_count(&*self.wrapped, locked) }
+    }
+
     fn insert_frames(
         &mut self,
         page_size: c_int,
@@ -259,6 +263,10 @@ where
             .savepoint_undo(&mut self.wrapped, rollback_data)
     }
 
+    fn frame_count(&self, locked: i32) -> super::Result<u32> {
+        self.wrapper.frame_count(&self.wrapped, locked)
+    }
+
     fn insert_frames(
         &mut self,
         page_size: std::ffi::c_int,
@@ -383,6 +391,10 @@ pub trait WrapWal<W: Wal> {
         wrapped.savepoint_undo(rollback_data)
     }
 
+    fn frame_count(&self, wrapped: &W, locked: i32) -> super::Result<u32> {
+        wrapped.frame_count(locked)
+    }
+
     fn insert_frames(
         &mut self,
         wrapped: &mut W,
diff --git a/libsql-wal/src/wal.rs b/libsql-wal/src/wal.rs
@@ -266,6 +266,11 @@ where
         }
     }
 
+    #[tracing::instrument(skip_all, fields(id = self.conn_id))]
+    fn frame_count(&self, _locked: i32) -> libsql_sys::wal::Result<u32> {
+        Err(libsql_sys::wal::Error::new(10)) // SQLITE_IOERR
+    }
+
     #[tracing::instrument(skip_all, fields(id = self.conn_id))]
     fn insert_frames(
         &mut self,

Original file line number	Diff line number	Diff line change
`@@ -98,6 +98,10 @@ mod tests {`
`98`	`98`	`self.0.read_frame(frame_no, buffer)`
`99`	`99`	`}`
`100`	`100`
	`101`	`+ fn frame_count(&self, locked: i32) -> libsql_sys::wal::Result<u32> {`
	`102`	`+ self.0.frame_count(locked)`
	`103`	`+ }`
	`104`	`+`
`101`	`105`	`fn db_size(&self) -> u32 {`
`102`	`106`	`self.0.db_size()`
`103`	`107`	`}`