fix/core: use CAS to protect against concurrent TX state changes

Concurrent use a of a connection across multiple threads is not supported,
but we should not crash even if a client attempts to do so.

We do crash when this happens. See: #3911

The proper behavior is to return Busy if the same connection is already doing
transactional work on another thread.

To ensure the above, this PR implements Connection::atomic_swap_tx_state() which
uses CAS instead of simply setting the state.

This PR also adds a regression test modeled after the reproduction in #3911 that
only accepts `"database is locked"` errors and panics on anything else.

Author: jussisaurio

Cargo.lock

@@ -22,6 +22,7 @@ thiserror = { workspace = true }
 tracing-subscriber.workspace = true
 tracing.workspace = true
 mimalloc = { workspace = true, optional = true }
+rusqlite.workspace = true
 
 [dev-dependencies]
 tempfile = { workspace = true }

bindings/rust/Cargo.toml

@@ -1,3 +1,5 @@
+use std::time::Duration;
+
 use tokio::fs;
 use turso::{Builder, EncryptionOpts, Error, Value};
 
@@ -531,3 +533,74 @@ async fn test_connection_clone() {
     let id: i64 = row.get(0).unwrap();
     assert_eq!(id, 1);
 }
+
+use tokio::task::JoinSet;
+
+#[tokio::test(flavor = "multi_thread", worker_threads = 8)]
+/// Sharing a connection across threads is not supported, but we must not crash.
+/// The proper behavior is to return Busy to the caller if another statement is already running a transaction.
+/// This test verifies that we end up inserting at least some of the rows and that integrity check passes.
+async fn concurrent_inserts_on_shared_connection() {
+    let temp_file = tempfile::NamedTempFile::new().expect("create temp file");
+    let db_path = temp_file.path().to_str().expect("path to string");
+
+    let db = Builder::new_local(db_path)
+        .build()
+        .await
+        .expect("temp file db");
+    let conn = db.connect().expect("connect");
+    conn.execute("CREATE TABLE IF NOT EXISTS t (value INTEGER)", ())
+        .await
+        .expect("create table");
+
+    conn.busy_timeout(Duration::from_millis(1000)).unwrap();
+
+    let attempts = 200usize;
+    let mut join_set = JoinSet::new();
+
+    for i in 0..attempts {
+        let conn = conn.clone();
+        join_set.spawn(async move {
+            let res = conn
+                .execute("INSERT INTO t (value) VALUES (?1)", [i as i64])
+                .await;
+            match res {
+                Ok(_) => {}
+                Err(Error::SqlExecutionFailure(e)) if e.contains("database is locked") => {}
+                Err(e) => panic!("unexpected error: {e:?}"),
+            }
+        });
+    }
+
+    while let Some(res) = join_set.join_next().await {
+        res.expect("task panicked");
+    }
+
+    let mut rows = conn
+        .query("SELECT COUNT(*) FROM t", ())
+        .await
+        .expect("count rows");
+    let count_value = rows
+        .next()
+        .await
+        .expect("step result")
+        .expect("row")
+        .get_value(0)
+        .expect("count value");
+    match count_value {
+        Value::Integer(count) => assert!(count > 0),
+        other => panic!("expected integer count, got {other:?}"),
+    }
+    assert!(rows.next().await.expect("consume rows").is_none());
+
+    // Close turso connection before opening with rusqlite
+    drop(conn);
+    drop(db);
+
+    // Verify database integrity with rusqlite
+    let rusqlite_conn = rusqlite::Connection::open(db_path).expect("open with rusqlite");
+    let integrity_result: String = rusqlite_conn
+        .query_row("PRAGMA integrity_check", [], |row| row.get(0))
+        .expect("integrity check");
+    assert_eq!(integrity_result, "ok", "database integrity check failed");
+}

bindings/rust/tests/integration_tests.rs

@@ -2623,6 +2623,16 @@ impl Connection {
         self.transaction_state.get()
     }
 
+    /// Atomically compare and exchange transaction state.
+    /// Returns Ok(current) on success, or Err(actual_value) if the state was modified concurrently.
+    fn atomic_swap_tx_state(
+        &self,
+        current: TransactionState,
+        new: TransactionState,
+    ) -> Result<TransactionState, TransactionState> {
+        self.transaction_state.compare_exchange(current, new)
+    }
+
     pub(crate) fn get_mv_tx_id(&self) -> Option<u64> {
         self.mv_tx.read().map(|(tx_id, _)| tx_id)
     }

core/lib.rs

@@ -2234,6 +2234,22 @@ pub fn op_transaction_inner(
 
                 // 1. We try to upgrade current version
                 let current_state = conn.get_tx_state();
+                let auto_commit = conn.auto_commit.load(Ordering::SeqCst);
+
+                // In autocommit mode, each statement is its own transaction. If we see
+                // a transaction already in progress (state != None), AND we haven't started
+                // any transaction work in this execution yet, it means another concurrent
+                // execution owns that transaction. Return Busy to avoid using their transaction.
+                //
+                // We check auto_txn_cleanup to distinguish:
+                // - RollbackTxn: We started a transaction, so state != None is expected (e.g., Read -> Write upgrade)
+                // - None: We haven't started anything, so state != None means someone else has the transaction
+                let connection_already_has_tx = !matches!(current_state, TransactionState::None)
+                    && matches!(state.auto_txn_cleanup, TxnCleanup::None);
+                if auto_commit && !conn.is_nested_stmt() && connection_already_has_tx {
+                    return Err(LimboError::Busy);
+                }
+
                 let (new_transaction_state, updated) = if conn.is_nested_stmt() {
                     (current_state, false)
                 } else {
@@ -2275,7 +2291,22 @@ pub fn op_transaction_inner(
                     }
                 };
 
-                // 2. Start transaction if needed
+                // 2. Atomically claim the new transaction state BEFORE doing pager operations.
+                // This prevents races where two threads both think they can upgrade the state.
+                // If compare_exchange fails, another thread modified the state concurrently -
+                // return Busy to signal the caller should retry. We can't just retry here because
+                // pager transaction state is per-connection, and another thread's transaction
+                // is not ours to use.
+                if updated {
+                    if conn
+                        .atomic_swap_tx_state(current_state, new_transaction_state)
+                        .is_err()
+                    {
+                        return Err(LimboError::Busy);
+                    }
+                }
+
+                // 3. Start transaction on the pager. If this fails, restore the old state.
                 if let Some(mv_store) = mv_store.as_ref() {
                     // In MVCC we don't have write exclusivity, therefore we just need to start a transaction if needed.
                     // Programs can run Transaction twice, first with read flag and then with write flag. So a single txid is enough
@@ -2286,6 +2317,10 @@ pub fn op_transaction_inner(
                     let conn_has_executed_begin_deferred = !has_existing_mv_tx
                         && !program.connection.auto_commit.load(Ordering::SeqCst);
                     if conn_has_executed_begin_deferred && *tx_mode == TransactionMode::Concurrent {
+                        // Restore state before returning error
+                        if updated {
+                            conn.set_tx_state(current_state);
+                        }
                         return Err(LimboError::TxError(
                             "Cannot start CONCURRENT transaction after BEGIN DEFERRED".to_string(),
                         ));
@@ -2318,6 +2353,10 @@ pub fn op_transaction_inner(
                     }
                 } else {
                     if matches!(tx_mode, TransactionMode::Concurrent) {
+                        // Restore state before returning error
+                        if updated {
+                            conn.set_tx_state(current_state);
+                        }
                         return Err(LimboError::TxError(
                             "Concurrent transaction mode is only supported when MVCC is enabled"
                                 .to_string(),
@@ -2344,27 +2383,20 @@ pub fn op_transaction_inner(
                             // start a new one.
                             if matches!(current_state, TransactionState::None) {
                                 pager.end_read_tx();
-                                conn.set_tx_state(TransactionState::None);
                                 state.auto_txn_cleanup = TxnCleanup::None;
                             }
-                            assert_eq!(conn.get_tx_state(), current_state);
+                            // Restore the transaction state since pager ops failed
+                            conn.set_tx_state(current_state);
                             return Err(LimboError::Busy);
                         }
                         if let IOResult::IO(io) = begin_w_tx_res? {
                             // set the transaction state to pending so we don't have to
                             // end the read transaction.
-                            program
-                                .connection
-                                .set_tx_state(TransactionState::PendingUpgrade);
+                            conn.set_tx_state(TransactionState::PendingUpgrade);
                             return Ok(InsnFunctionStepResult::IO(io));
                         }
                     }
                 }
-
-                // 3. Transaction state should be updated before checking for Schema cookie so that the tx is ended properly on error
-                if updated {
-                    conn.set_tx_state(new_transaction_state);
-                }
                 state.op_transaction_state = OpTransactionState::CheckSchemaCookie;
                 continue;
             }

core/vdbe/execute.rs

@@ -263,6 +263,23 @@ pub(crate) fn derive_atomic_enum_inner(input: TokenStream) -> TokenStream {
                 let prev = self.0.swap(Self::to_storage(&val), ::std::sync::atomic::Ordering::SeqCst);
                 Self::from_storage(prev)
             }
+
+            #[inline]
+            /// Compare and exchange: if the current value equals `current`, replace it with `new`.
+            /// Returns Ok(current) on success, or Err(actual_value) on failure.
+            pub fn compare_exchange(&self, current: #name, new: #name) -> Result<#name, #name> {
+                let current_storage = Self::to_storage(&current);
+                let new_storage = Self::to_storage(&new);
+                match self.0.compare_exchange(
+                    current_storage,
+                    new_storage,
+                    ::std::sync::atomic::Ordering::SeqCst,
+                    ::std::sync::atomic::Ordering::SeqCst,
+                ) {
+                    Ok(prev) => Ok(Self::from_storage(prev)),
+                    Err(actual) => Err(Self::from_storage(actual)),
+                }
+            }
         }
 
         impl From<#name> for #atomic_name {