Repository: tursodatabase/limbo (rust)

Issue:
Transaction commit in program halt is not re-entrant
If the `Program::halt()` function returns I/O status code, we attempt to commit a transaction that has already been committed as we retry the instruction. Part of the problem is that we perform checkpointing for every transaction, which is not what we want.

Pull request description:
Transaction finish lock fixes
While adding logs to our locks in `Wal` I noticed we weren't cleaning up connection's transaction state. This PR set to `TransactionState::None` once commited and calls `end_read_txn` and `end_tx` in case of write.

Fixes #1004


Patch:
diff --git a/core/storage/wal.rs b/core/storage/wal.rs
index 9df9121f..e9e61025 100644
--- a/core/storage/wal.rs
+++ b/core/storage/wal.rs
@@ -62,7 +62,7 @@ impl LimboRwLock {
     /// Shared lock. Returns true if it was successful, false if it couldn't lock it
     pub fn read(&mut self) -> bool {
         let lock = self.lock.load(Ordering::SeqCst);
-        match lock {
+        let ok = match lock {
             NO_LOCK => {
                 let res = self.lock.compare_exchange(
                     lock,
@@ -82,13 +82,15 @@ impl LimboRwLock {
             }
             WRITE_LOCK => false,
             _ => unreachable!(),
-        }
+        };
+        tracing::trace!("read_lock({})", ok);
+        ok
     }
 
     /// Locks exclusively. Returns true if it was successful, false if it couldn't lock it
     pub fn write(&mut self) -> bool {
         let lock = self.lock.load(Ordering::SeqCst);
-        match lock {
+        let ok = match lock {
             NO_LOCK => {
                 let res = self.lock.compare_exchange(
                     lock,
@@ -104,12 +106,15 @@ impl LimboRwLock {
             }
             WRITE_LOCK => true,
             _ => unreachable!(),
-        }
+        };
+        tracing::trace!("write_lock({})", ok);
+        ok
     }
 
     /// Unlock the current held lock.
     pub fn unlock(&mut self) {
         let lock = self.lock.load(Ordering::SeqCst);
+        tracing::trace!("unlock(lock={})", lock);
         match lock {
             NO_LOCK => {}
             SHARED_LOCK => {
@@ -355,7 +360,7 @@ impl Wal for WalFile {
         self.max_frame_read_lock_index = max_read_mark_index as usize;
         self.max_frame = max_read_mark as u64;
         self.min_frame = shared.nbackfills + 1;
-        trace!(
+        tracing::debug!(
             "begin_read_tx(min_frame={}, max_frame={}, lock={})",
             self.min_frame,
             self.max_frame,
@@ -366,6 +371,7 @@ impl Wal for WalFile {
 
     /// End a read transaction.
     fn end_read_tx(&self) -> Result<LimboResult> {
+        tracing::debug!("end_read_tx");
         let mut shared = self.shared.write();
         let read_lock = &mut shared.read_locks[self.max_frame_read_lock_index];
         read_lock.unlock();
@@ -376,6 +382,7 @@ impl Wal for WalFile {
     fn begin_write_tx(&mut self) -> Result<LimboResult> {
         let mut shared = self.shared.write();
         let busy = !shared.write_lock.write();
+        tracing::debug!("begin_write_transaction(busy={})", busy);
         if busy {
             return Ok(LimboResult::Busy);
         }
@@ -384,6 +391,7 @@ impl Wal for WalFile {
 
     /// End a write transaction
     fn end_write_tx(&self) -> Result<LimboResult> {
+        tracing::debug!("end_write_txn");
         let mut shared = self.shared.write();
         shared.write_lock.unlock();
         Ok(LimboResult::Ok)
diff --git a/core/vdbe/mod.rs b/core/vdbe/mod.rs
index c0067443..3f16d525 100644
--- a/core/vdbe/mod.rs
+++ b/core/vdbe/mod.rs
@@ -30,7 +30,6 @@ use crate::functions::datetime::{
     exec_date, exec_datetime_full, exec_julianday, exec_strftime, exec_time, exec_unixepoch,
 };
 use crate::functions::printf::exec_printf;
-use crate::info;
 use crate::pseudo::PseudoCursor;
 use crate::result::LimboResult;
 use crate::schema::{affinity, Affinity};
@@ -55,6 +54,7 @@ use crate::{
     json::json_error_position, json::json_extract, json::json_object, json::json_patch,
     json::json_quote, json::json_remove, json::json_set, json::json_type,
 };
+use crate::{info, CheckpointStatus};
 use crate::{resolve_ext_path, Connection, Result, TransactionState, DATABASE_VERSION};
 use insn::{
     exec_add, exec_and, exec_bit_and, exec_bit_not, exec_bit_or, exec_boolean_not, exec_concat,
@@ -71,6 +71,7 @@ use std::cell::{Cell, RefCell};
 use std::collections::HashMap;
 use std::ffi::c_void;
 use std::num::NonZero;
+use std::ops::Deref;
 use std::rc::{Rc, Weak};
 use std::sync::{Arc, Mutex};
 use tracing::instrument;
@@ -219,6 +220,11 @@ impl VTabOpaqueCursor {
     }
 }
 
+#[derive(Copy, Clone)]
+enum HaltState {
+    Checkpointing,
+}
+
 /// The program state describes the environment in which the program executes.
 pub struct ProgramState {
     pub pc: InsnReference,
@@ -231,6 +237,7 @@ pub struct ProgramState {
     regex_cache: RegexCache,
     interrupted: bool,
     parameters: HashMap<NonZero<usize>, OwnedValue>,
+    halt_state: Option<HaltState>,
 }
 
 impl ProgramState {
@@ -249,6 +256,7 @@ impl ProgramState {
             regex_cache: RegexCache::new(),
             interrupted: false,
             parameters: HashMap::new(),
+            halt_state: None,
         }
     }
 
@@ -1196,7 +1204,7 @@ impl Program {
                             )));
                         }
                     }
-                    return self.halt(pager);
+                    return self.halt(pager, state);
                 }
                 Insn::Transaction { write } => {
                     let connection = self.connection.upgrade().unwrap();
@@ -1212,7 +1220,6 @@ impl Program {
 
                     if updated && matches!(current_state, TransactionState::None) {
                         if let LimboResult::Busy = pager.begin_read_tx()? {
-                            tracing::trace!("begin_read_tx busy");
                             return Ok(StepResult::Busy);
                         }
                     }
@@ -1254,7 +1261,7 @@ impl Program {
                             "cannot commit - no transaction is active".to_string(),
                         ));
                     }
-                    return self.halt(pager);
+                    return self.halt(pager, state);
                 }
                 Insn::Goto { target_pc } => {
                     assert!(target_pc.is_offset());
@@ -3077,39 +3084,68 @@ impl Program {
         }
     }
 
-    fn halt(&self, pager: Rc<Pager>) -> Result<StepResult> {
+    fn halt(&self, pager: Rc<Pager>, program_state: &mut ProgramState) -> Result<StepResult> {
         let connection = self
             .connection
             .upgrade()
             .expect("only weak ref to connection?");
         let auto_commit = *connection.auto_commit.borrow();
         tracing::trace!("Halt auto_commit {}", auto_commit);
-        if auto_commit {
-            let current_state = connection.transaction_state.borrow().clone();
-            if current_state == TransactionState::Read {
-                pager.end_read_tx()?;
-                return Ok(StepResult::Done);
-            }
-            match pager.end_tx() {
-                Ok(crate::storage::wal::CheckpointStatus::IO) => Ok(StepResult::IO),
-                Ok(crate::storage::wal::CheckpointStatus::Done(_)) => {
-                    if self.change_cnt_on {
-                        if let Some(conn) = self.connection.upgrade() {
-                            conn.set_changes(self.n_change.get());
-                        }
+        assert!(
+            program_state.halt_state.is_none()
+                || (matches!(program_state.halt_state.unwrap(), HaltState::Checkpointing))
+        );
+        if program_state.halt_state.is_some() {
+            self.step_end_write_txn(&pager, &mut program_state.halt_state, connection.deref())
+        } else {
+            if auto_commit {
+                let current_state = connection.transaction_state.borrow().clone();
+                match current_state {
+                    TransactionState::Write => self.step_end_write_txn(
+                        &pager,
+                        &mut program_state.halt_state,
+                        connection.deref(),
+                    ),
+                    TransactionState::Read => {
+                        pager.end_read_tx()?;
+                        Ok(StepResult::Done)
                     }
-                    Ok(StepResult::Done)
+                    TransactionState::None => Ok(StepResult::Done),
                 }
-                Err(e) => Err(e),
-            }
-        } else {
-            if self.change_cnt_on {
-                if let Some(conn) = self.connection.upgrade() {
-                    conn.set_changes(self.n_change.get());
+            } else {
+                if self.change_cnt_on {
+                    if let Some(conn) = self.connection.upgrade() {
+                        conn.set_changes(self.n_change.get());
+                    }
                 }
+                Ok(StepResult::Done)
+            }
+        }
+    }
+
+    fn step_end_write_txn(
+        &self,
+        pager: &Rc<Pager>,
+        halt_state: &mut Option<HaltState>,
+        connection: &Connection,
+    ) -> Result<StepResult> {
+        let checkpoint_status = pager.end_tx()?;
+        match checkpoint_status {
+            CheckpointStatus::Done(_) => {
+                connection.transaction_state.replace(TransactionState::None);
+                let _ = halt_state.take();
+            }
+            CheckpointStatus::IO => {
+                *halt_state = Some(HaltState::Checkpointing);
+                return Ok(StepResult::IO);
+            }
+        }
+        if self.change_cnt_on {
+            if let Some(conn) = self.connection.upgrade() {
+                conn.set_changes(self.n_change.get());
             }
-            Ok(StepResult::Done)
         }
+        Ok(StepResult::Done)
     }
 }
 


Test patch:
diff --git a/tests/integration/wal/test_wal.rs b/tests/integration/wal/test_wal.rs
index a0173a2b..5a92e01c 100644
--- a/tests/integration/wal/test_wal.rs
+++ b/tests/integration/wal/test_wal.rs
@@ -1,7 +1,9 @@
 use crate::common::{do_flush, TempDatabase};
 use limbo_core::{Connection, LimboError, Result, StepResult};
 use std::cell::RefCell;
+use std::ops::Deref;
 use std::rc::Rc;
+use std::sync::{Arc, Mutex};
 use tracing_subscriber::layer::SubscriberExt;
 use tracing_subscriber::util::SubscriberInitExt;
 use tracing_subscriber::EnvFilter;
@@ -33,6 +35,87 @@ fn test_wal_checkpoint_result() -> Result<()> {
     Ok(())
 }
 
+#[test]
+#[ignore = "ignored for now because it's flaky"]
+fn test_wal_1_writer_1_reader() -> Result<()> {
+    maybe_setup_tracing();
+    let tmp_db = Arc::new(Mutex::new(TempDatabase::new("test_wal.db")));
+    let db = tmp_db.lock().unwrap().limbo_database();
+
+    {
+        let conn = db.connect().unwrap();
+        match conn.query("CREATE TABLE t (id)")? {
+            Some(ref mut rows) => loop {
+                match rows.step().unwrap() {
+                    StepResult::Row => {}
+                    StepResult::IO => {
+                        tmp_db.lock().unwrap().io.run_once().unwrap();
+                    }
+                    StepResult::Interrupt => break,
+                    StepResult::Done => break,
+                    StepResult::Busy => unreachable!(),
+                }
+            },
+            None => todo!(),
+        }
+        do_flush(&conn, tmp_db.lock().unwrap().deref()).unwrap();
+    }
+    let rows = Arc::new(std::sync::Mutex::new(0));
+    let rows_ = rows.clone();
+    const ROWS_WRITE: usize = 1000;
+    let tmp_db_w = db.clone();
+    let writer_thread = std::thread::spawn(move || {
+        let conn = tmp_db_w.connect().unwrap();
+        for i in 0..ROWS_WRITE {
+            conn.execute(format!("INSERT INTO t values({})", i).as_str())
+                .unwrap();
+            let mut rows = rows_.lock().unwrap();
+            *rows += 1;
+        }
+    });
+    let rows_ = rows.clone();
+    let reader_thread = std::thread::spawn(move || {
+        let conn = db.connect().unwrap();
+        loop {
+            let rows = *rows_.lock().unwrap();
+            let mut i = 0;
+            match conn.query("SELECT * FROM t") {
+                Ok(Some(ref mut rows)) => loop {
+                    match rows.step().unwrap() {
+                        StepResult::Row => {
+                            let row = rows.row().unwrap();
+                            let first_value = row.get_value(0);
+                            let id = match first_value {
+                                limbo_core::OwnedValue::Integer(i) => *i as i32,
+                                _ => unreachable!(),
+                            };
+                            assert_eq!(id, i);
+                            i += 1;
+                        }
+                        StepResult::IO => {
+                            tmp_db.lock().unwrap().io.run_once().unwrap();
+                        }
+                        StepResult::Interrupt => break,
+                        StepResult::Done => break,
+                        StepResult::Busy => unreachable!(),
+                    }
+                },
+                Ok(None) => {}
+                Err(err) => {
+                    eprintln!("{}", err);
+                }
+            }
+            if rows == ROWS_WRITE {
+                break;
+            }
+        }
+    });
+
+    writer_thread.join().unwrap();
+    reader_thread.join().unwrap();
+    Ok(())
+}
+
 fn maybe_setup_tracing() {
     let _ = tracing_subscriber::registry()
         .with(