remove txn state from execute return

It is now passed to the result builder

Author: MarinPostma

Cargo.lock

@@ -4,7 +4,7 @@ use std::sync::atomic::{AtomicBool, Ordering};
 use std::sync::Arc;
 
 use parking_lot::{Mutex, RwLock};
-use rusqlite::{DatabaseName, ErrorCode, OpenFlags, StatementStatus};
+use rusqlite::{DatabaseName, ErrorCode, OpenFlags, StatementStatus, TransactionState};
 use sqld_libsql_bindings::wal_hook::{TransparentMethods, WalMethodsHook};
 use tokio::sync::{watch, Notify};
 use tokio::time::{Duration, Instant};
@@ -144,7 +144,6 @@ where
     }
 }
 
-#[derive(Clone)]
 pub struct LibSqlConnection<W: WalHook> {
     inner: Arc<Mutex<Connection<W>>>,
 }
@@ -160,6 +159,12 @@ impl<W: WalHook> std::fmt::Debug for LibSqlConnection<W> {
     }
 }
 
+impl<W: WalHook> Clone for LibSqlConnection<W> {
+    fn clone(&self) -> Self {
+        Self { inner: self.inner.clone() }
+    }
+}
+
 pub fn open_conn<W>(
     path: &Path,
     wal_methods: &'static WalMethodsHook<W>,
@@ -219,6 +224,15 @@ where
             inner: Arc::new(Mutex::new(conn)),
         })
     }
+
+    pub fn txn_status(&self) -> crate::Result<TxnStatus> {
+        Ok(self
+            .inner
+            .lock()
+            .conn
+            .transaction_state(Some(DatabaseName::Main))?
+            .into())
+    }
 }
 
 struct Connection<W: WalHook = TransparentMethods> {
@@ -351,6 +365,16 @@ unsafe extern "C" fn busy_handler<W: WalHook>(state: *mut c_void, _retries: c_in
     })
 }
 
+impl From<TransactionState> for TxnStatus {
+    fn from(value: TransactionState) -> Self {
+        use TransactionState as Tx;
+        match value {
+            Tx::None => TxnStatus::Init,
+            Tx::Read | Tx::Write => TxnStatus::Txn,
+            _ => unreachable!(),
+        }
+    }
+}
 impl<W: WalHook> Connection<W> {
     fn new(
         path: &Path,
@@ -405,7 +429,7 @@ impl<W: WalHook> Connection<W> {
         this: Arc<Mutex<Self>>,
         pgm: Program,
         mut builder: B,
-    ) -> Result<(B, TxnStatus)> {
+    ) -> Result<B> {
         use rusqlite::TransactionState as Tx;
 
         let state = this.lock().state.clone();
@@ -469,23 +493,18 @@ impl<W: WalHook> Connection<W> {
             results.push(res);
         }
 
-        let status = if matches!(
-            this.lock()
-                .conn
-                .transaction_state(Some(DatabaseName::Main))?,
-            Tx::Read | Tx::Write
-        ) {
-            TxnStatus::Txn
-        } else {
-            TxnStatus::Init
-        };
+        let status = this
+            .lock()
+            .conn
+            .transaction_state(Some(DatabaseName::Main))?
+            .into();
 
         builder.finish(
             *this.lock().current_frame_no_receiver.borrow_and_update(),
             status,
         )?;
 
-        Ok((builder, status))
+        Ok(builder)
     }
 
     fn execute_step(
@@ -736,7 +755,7 @@ where
         auth: Authenticated,
         builder: B,
         _replication_index: Option<FrameNo>,
-    ) -> Result<(B, TxnStatus)> {
+    ) -> Result<B> {
         check_program_auth(auth, &pgm)?;
         let conn = self.inner.clone();
         tokio::task::spawn_blocking(move || Connection::run(conn, pgm, builder))
@@ -828,7 +847,7 @@ mod test {
     fn test_libsql_conn_builder_driver() {
         test_driver(1000, |b| {
             let conn = setup_test_conn();
-            Connection::run(conn, Program::seq(&["select * from test"]), b).map(|x| x.0)
+            Connection::run(conn, Program::seq(&["select * from test"]), b)
         })
     }
 
@@ -852,23 +871,23 @@ mod test {
 
         tokio::time::pause();
         let conn = make_conn.make_connection().await.unwrap();
-        let (_builder, state) = Connection::run(
+        let _builder = Connection::run(
             conn.inner.clone(),
             Program::seq(&["BEGIN IMMEDIATE"]),
             TestBuilder::default(),
         )
         .unwrap();
-        assert_eq!(state, TxnStatus::Txn);
+        assert_eq!(conn.txn_status().unwrap(), TxnStatus::Txn);
 
         tokio::time::advance(TXN_TIMEOUT * 2).await;
 
-        let (builder, state) = Connection::run(
+        let builder = Connection::run(
             conn.inner.clone(),
             Program::seq(&["BEGIN IMMEDIATE"]),
             TestBuilder::default(),
         )
         .unwrap();
-        assert_eq!(state, TxnStatus::Init);
+        assert_eq!(conn.txn_status().unwrap(), TxnStatus::Init);
         assert!(matches!(builder.into_ret()[0], Err(Error::LibSqlTxTimeout)));
     }
 
@@ -896,13 +915,13 @@ mod test {
         for _ in 0..10 {
             let conn = make_conn.make_connection().await.unwrap();
             set.spawn_blocking(move || {
-                let (builder, state) = Connection::run(
-                    conn.inner,
+                let builder = Connection::run(
+                    conn.inner.clone(),
                     Program::seq(&["BEGIN IMMEDIATE"]),
                     TestBuilder::default(),
                 )
                 .unwrap();
-                assert_eq!(state, TxnStatus::Txn);
+                assert_eq!(conn.txn_status().unwrap(), TxnStatus::Txn);
                 assert!(builder.into_ret()[0].is_ok());
             });
         }
@@ -937,15 +956,15 @@ mod test {
 
         let conn1 = make_conn.make_connection().await.unwrap();
         tokio::task::spawn_blocking({
-            let conn = conn1.inner.clone();
+            let conn = conn1.clone();
             move || {
-                let (builder, state) = Connection::run(
-                    conn,
+                let builder = Connection::run(
+                    conn.inner.clone(),
                     Program::seq(&["BEGIN IMMEDIATE"]),
                     TestBuilder::default(),
                 )
                 .unwrap();
-                assert_eq!(state, TxnStatus::Txn);
+                assert_eq!(conn.txn_status().unwrap(), TxnStatus::Txn);
                 assert!(builder.into_ret()[0].is_ok());
             }
         })
@@ -954,16 +973,16 @@ mod test {
 
         let conn2 = make_conn.make_connection().await.unwrap();
         let handle = tokio::task::spawn_blocking({
-            let conn = conn2.inner.clone();
+            let conn = conn2.clone();
             move || {
                 let before = Instant::now();
-                let (builder, state) = Connection::run(
-                    conn,
+                let builder = Connection::run(
+                    conn.inner.clone(),
                     Program::seq(&["BEGIN IMMEDIATE"]),
                     TestBuilder::default(),
                 )
                 .unwrap();
-                assert_eq!(state, TxnStatus::Txn);
+                assert_eq!(conn.txn_status().unwrap(), TxnStatus::Txn);
                 assert!(builder.into_ret()[0].is_ok());
                 before.elapsed()
             }
@@ -973,12 +992,12 @@ mod test {
         tokio::time::sleep(wait_time).await;
 
         tokio::task::spawn_blocking({
-            let conn = conn1.inner.clone();
+            let conn = conn1.clone();
             move || {
-                let (builder, state) =
-                    Connection::run(conn, Program::seq(&["COMMIT"]), TestBuilder::default())
+                let builder =
+                    Connection::run(conn.inner.clone(), Program::seq(&["COMMIT"]), TestBuilder::default())
                         .unwrap();
-                assert_eq!(state, TxnStatus::Init);
+                assert_eq!(conn.txn_status().unwrap(), TxnStatus::Init);
                 assert!(builder.into_ret()[0].is_ok());
             }
         })

sqld/src/connection/libsql.rs

@@ -8,7 +8,7 @@ use tokio::{sync::Semaphore, time::timeout};
 use crate::auth::Authenticated;
 use crate::error::Error;
 use crate::query::{Params, Query};
-use crate::query_analysis::{Statement, TxnStatus};
+use crate::query_analysis::Statement;
 use crate::query_result_builder::{IgnoreResult, QueryResultBuilder};
 use crate::replication::FrameNo;
 use crate::Result;
@@ -32,7 +32,7 @@ pub trait Connection: Send + Sync + 'static {
         auth: Authenticated,
         response_builder: B,
         replication_index: Option<FrameNo>,
-    ) -> Result<(B, TxnStatus)>;
+    ) -> Result<B>;
 
     /// Execute all the queries in the batch sequentially.
     /// If an query in the batch fails, the remaining queries are ignores, and the batch current
@@ -43,7 +43,7 @@ pub trait Connection: Send + Sync + 'static {
         auth: Authenticated,
         result_builder: B,
         replication_index: Option<FrameNo>,
-    ) -> Result<(B, TxnStatus)> {
+    ) -> Result<B> {
         let batch_len = batch.len();
         let mut steps = make_batch_program(batch);
 
@@ -67,11 +67,11 @@ pub trait Connection: Send + Sync + 'static {
 
         // ignore the rollback result
         let builder = result_builder.take(batch_len);
-        let (builder, state) = self
+        let builder = self
             .execute_program(pgm, auth, builder, replication_index)
             .await?;
 
-        Ok((builder.into_inner(), state))
+        Ok(builder.into_inner())
     }
 
     /// Execute all the queries in the batch sequentially.
@@ -82,7 +82,7 @@ pub trait Connection: Send + Sync + 'static {
         auth: Authenticated,
         result_builder: B,
         replication_index: Option<FrameNo>,
-    ) -> Result<(B, TxnStatus)> {
+    ) -> Result<B> {
         let steps = make_batch_program(batch);
         let pgm = Program::new(steps);
         self.execute_program(pgm, auth, result_builder, replication_index)
@@ -312,7 +312,7 @@ impl<DB: Connection> Connection for TrackedConnection<DB> {
         auth: Authenticated,
         builder: B,
         replication_index: Option<FrameNo>,
-    ) -> crate::Result<(B, TxnStatus)> {
+    ) -> crate::Result<B> {
         self.atime.store(now_millis(), Ordering::Relaxed);
         self.inner
             .execute_program(pgm, auth, builder, replication_index)
@@ -367,7 +367,7 @@ mod test {
             _auth: Authenticated,
             _builder: B,
             _replication_index: Option<FrameNo>,
-        ) -> crate::Result<(B, TxnStatus)> {
+        ) -> crate::Result<B> {
             unreachable!()
         }
 

sqld/src/connection/mod.rs

@@ -184,7 +184,7 @@ impl WriteProxyConnection {
         status: &mut TxnStatus,
         auth: Authenticated,
         builder: B,
-    ) -> Result<(B, TxnStatus)> {
+    ) -> Result<B> {
         self.stats.inc_write_requests_delegated();
         *status = TxnStatus::Invalid;
         let res = self
@@ -208,7 +208,7 @@ impl WriteProxyConnection {
             self.update_last_write_frame_no(current_frame_no);
         }
 
-        Ok((builder, new_status))
+        Ok(builder)
     }
 
     fn update_last_write_frame_no(&self, new_frame_no: FrameNo) {
@@ -339,8 +339,8 @@ impl RemoteConnection {
                 }
                 Err(e) => {
                     tracing::error!("received error from connection stream: {e}");
-                    return Err(Error::StreamDisconnect)
-                },
+                    return Err(Error::StreamDisconnect);
+                }
             }
         }
 
@@ -356,26 +356,30 @@ impl Connection for WriteProxyConnection {
         auth: Authenticated,
         builder: B,
         replication_index: Option<FrameNo>,
-    ) -> Result<(B, TxnStatus)> {
+    ) -> Result<B> {
         let mut state = self.state.lock().await;
 
         // This is a fresh namespace, and it is not replicated yet, proxy the first request.
         if self.applied_frame_no_receiver.borrow().is_none() {
             self.execute_remote(pgm, &mut state, auth, builder).await
         } else if *state == TxnStatus::Init && pgm.is_read_only() {
+            // set the state to invalid before doing anything, and set it to a valid state after.
+            *state = TxnStatus::Invalid;
             self.wait_replication_sync(replication_index).await?;
             // We know that this program won't perform any writes. We attempt to run it on the
             // replica. If it leaves an open transaction, then this program is an interactive
             // transaction, so we rollback the replica, and execute again on the primary.
-            let (builder, new_state) = self
+            let builder = self
                 .read_conn
                 .execute_program(pgm.clone(), auth.clone(), builder, replication_index)
                 .await?;
+            let new_state = self.read_conn.txn_status()?;
             if new_state != TxnStatus::Init {
                 self.read_conn.rollback(auth.clone()).await?;
                 self.execute_remote(pgm, &mut state, auth, builder).await
             } else {
-                Ok((builder, new_state))
+                *state = new_state;
+                Ok(builder)
             }
         } else {
             self.execute_remote(pgm, &mut state, auth, builder).await
@@ -433,7 +437,10 @@ pub mod test {
     use rand::Fill;
 
     use super::*;
-    use crate::{query_result_builder::test::test_driver, rpc::proxy::rpc::{ExecuteResults, query_result::RowResult}};
+    use crate::{
+        query_result_builder::test::test_driver,
+        rpc::proxy::rpc::{query_result::RowResult, ExecuteResults},
+    };
 
     /// generate an arbitraty rpc value. see build.rs for usage.
     pub fn arbitrary_rpc_value(u: &mut Unstructured) -> arbitrary::Result<Vec<u8>> {
@@ -497,12 +504,15 @@ pub mod test {
     /// In this test, we generate random ExecuteResults, and ensures that the `execute_results_to_builder` drives the builder FSM correctly.
     #[test]
     fn test_execute_results_to_builder() {
-        test_driver(1000, |b| -> std::result::Result<crate::query_result_builder::test::FsmQueryBuilder, Error> {
-            let mut data = [0; 10_000];
-            data.try_fill(&mut rand::thread_rng()).unwrap();
-            let mut un = Unstructured::new(&data);
-            let res = ExecuteResults::arbitrary(&mut un).unwrap();
-            execute_results_to_builder(res, b, &QueryBuilderConfig::default())
-        });
+        test_driver(
+            1000,
+            |b| -> std::result::Result<crate::query_result_builder::test::FsmQueryBuilder, Error> {
+                let mut data = [0; 10_000];
+                data.try_fill(&mut rand::thread_rng()).unwrap();
+                let mut un = Unstructured::new(&data);
+                let res = ExecuteResults::arbitrary(&mut un).unwrap();
+                execute_results_to_builder(res, b, &QueryBuilderConfig::default())
+            },
+        );
     }
 }