Use read lock after truncate checkpoint to prevent races

The previous implementation attempted to hold checkpoint locks during file copy by adding a keep_lock param to the checkpoint function.
However, checkpoint can have multiple early exit paths like empty WAL where no lock is acquired. Also, it clutters the checkpoint function.

This implementation executes TRUNCATE checkpoint and then acquires read lock on WAL before copying database file. This ensures new data is written
to WAL and new Checkpoint can not be taken before this read lock is released.

Signed-off-by: Prateek Singh Rathore <[email protected]>

Author: psrvere

core/lib.rs

@@ -2025,11 +2025,13 @@ impl Connection {
 
     /// Create a point-in-time snapshot of the database by copying the database file.
     ///
-    /// It checkpoints and keeps the lock after checkpointing to avoid race window between
-    /// checkpoint completion and file copy operation where concurrent writers could modify
-    /// the database.
+    /// This function:
+    /// 1. Performs a TRUNCATE checkpoint to flush all WAL data to the database file
+    /// 2. Acquires a read lock to prevent other checkpoints during copy
+    /// 3. Copies the database file (new writes go to WAL, not DB file)
+    /// 4. Releases the read lock
     #[cfg(feature = "fs")]
-    pub fn _snapshot_copy_file(self: &Arc<Self>, output_path: &str) -> Result<()> {
+    fn _snapshot_copy_file(self: &Arc<Self>, output_path: &str) -> Result<()> {
         use crate::util::MEMORY_PATH;
         use std::path::Path;
 
@@ -2064,23 +2066,36 @@ impl Connection {
         }
 
         let pager = self.pager.load();
-        let result = (|| -> Result<()> {
-            // Checkpoint and keep the lock
-            let _res = pager.blocking_checkpoint_keep_lock(
-                CheckpointMode::Truncate {
-                    upper_bound_inclusive: None,
-                },
-                self.get_sync_mode(),
-            )?;
+        let Some(wal) = pager.wal.as_ref() else {
+            return Err(LimboError::InternalError(
+                "Cannot snapshot database without WAL".to_string(),
+            ));
+        };
 
+        // TRUNCATE checkpoint - flushes all WAL data to DB file and empties WAL
+        let _ = pager.blocking_checkpoint(
+            CheckpointMode::Truncate {
+                upper_bound_inclusive: None,
+            },
+            self.get_sync_mode(),
+        )?;
+
+        // Acquire read lock - prevents other checkpoints during copy
+        // After TRUNCATE, DB file is complete. New writes go to WAL only, not DB file.
+        pager.begin_read_tx()?;
+
+        // Copy the database file
+        let result = (|| -> Result<()> {
             let source_path = Path::new(&self.db.path);
             std::fs::copy(source_path, path).map_err(|e| {
                 LimboError::InternalError(format!("Failed to copy database file: {e}"))
             })?;
-
-            Ok(()) // lock is dropped here when _res is dropped here
+            Ok(())
         })();
 
+        // Release read lock
+        wal.end_read_tx();
+
         result
     }
 

core/storage/pager.rs

@@ -2856,40 +2856,17 @@ impl Pager {
         }
     }
 
-    #[instrument(skip_all, level = Level::DEBUG)]
-    /// checkpoint_keep_lock checkpoints the WAL to the database file (if needed) and keeps the lock
-    pub fn checkpoint_keep_lock(
-        &self,
-        mode: CheckpointMode,
-        sync_mode: crate::SyncMode,
-        clear_page_cache: bool,
-    ) -> Result<IOResult<CheckpointResult>> {
-        self.checkpoint_inner(mode, sync_mode, clear_page_cache, true)
-    }
-
-    #[instrument(skip_all, level = Level::DEBUG)]
-    /// checkpoint checkpoints the WAL to the database file (if needed) and drops the lock
-    pub fn checkpoint(
-        &self,
-        mode: CheckpointMode,
-        sync_mode: crate::SyncMode,
-        clear_page_cache: bool,
-    ) -> Result<IOResult<CheckpointResult>> {
-        self.checkpoint_inner(mode, sync_mode, clear_page_cache, false)
-    }
-
     #[instrument(skip_all, level = Level::DEBUG, name = "pager_checkpoint",)]
     /// Checkpoint the WAL to the database file (if needed).
     /// Args:
     /// - mode: The checkpoint mode to use (PASSIVE, FULL, RESTART, TRUNCATE)
     /// - sync_mode: The fsync mode to use (OFF, NORMAL, FULL)
     /// - clear_page_cache: Whether to clear the page cache after checkpointing
-    fn checkpoint_inner(
+    pub fn checkpoint(
         &self,
         mode: CheckpointMode,
         sync_mode: crate::SyncMode,
         clear_page_cache: bool,
-        keep_lock: bool,
     ) -> Result<IOResult<CheckpointResult>> {
         let Some(wal) = self.wal.as_ref() else {
             return Err(LimboError::InternalError(
@@ -3017,10 +2994,8 @@ impl Pager {
                     let mut state = self.checkpoint_state.write();
                     let mut res = state.result.take().expect("result should be set");
 
-                    // Release checkpoint guard if lock is not to be kept
-                    if !keep_lock {
-                        res.release_guard();
-                    }
+                    // Release checkpoint guard
+                    res.release_guard();
 
                     // Clear page cache only if requested (explicit checkpoints do this, auto-checkpoint does not)
                     if clear_page_cache {
@@ -3117,16 +3092,6 @@ impl Pager {
         self.io.block(|| self.checkpoint(mode, sync_mode, true))
     }
 
-    #[instrument(skip_all, level = Level::DEBUG)]
-    pub fn blocking_checkpoint_keep_lock(
-        &self,
-        mode: CheckpointMode,
-        sync_mode: crate::SyncMode,
-    ) -> Result<CheckpointResult> {
-        self.io
-            .block(|| self.checkpoint_keep_lock(mode, sync_mode, true))
-    }
-
     pub fn freepage_list(&self) -> u32 {
         self.io
             .block(|| HeaderRefMut::from_pager(self))