add proper handling for deffered transactions in shadow state

Author: pedrocarlo

simulator/model/mod.rs

@@ -18,6 +18,7 @@ use sql_generation::model::{
 };
 use turso_parser::ast::Distinctness;
 
+use crate::runner::env::TransactionMode;
 use crate::{generation::Shadow, runner::env::ShadowTablesMut};
 
 pub mod interactions;
@@ -146,6 +147,16 @@ impl Query {
     pub fn is_dml(&self) -> bool {
         matches!(self, Self::Insert(..) | Self::Update(..) | Self::Delete(..))
     }
+
+    #[inline]
+    pub fn is_write(&self) -> bool {
+        self.is_ddl() || self.is_dml()
+    }
+
+    #[inline]
+    pub fn is_select(&self) -> bool {
+        matches!(self, Self::Select(_))
+    }
 }
 
 impl Display for Query {
@@ -173,6 +184,9 @@ impl Shadow for Query {
     type Result = anyhow::Result<Vec<Vec<SimValue>>>;
 
     fn shadow(&self, env: &mut ShadowTablesMut) -> Self::Result {
+        // First check if we are not in a deffered transaction, if we are create a snapshot
+        env.upgrade_transaction(self);
+
         match self {
             Query::Create(create) => create.shadow(env),
             Query::Insert(insert) => insert.shadow(env),
@@ -567,9 +581,18 @@ impl Shadow for Select {
 impl Shadow for Begin {
     type Result = Vec<Vec<SimValue>>;
     fn shadow(&self, tables: &mut ShadowTablesMut) -> Self::Result {
-        // FIXME: currently the snapshot is taken eagerly
-        // this is wrong for Deffered transactions
-        tables.create_snapshot();
+        match self {
+            Begin::Deferred => {
+                tables.create_deffered_snapshot();
+            }
+            Begin::Immediate => {
+                tables.create_snapshot(TransactionMode::Write);
+            }
+            Begin::Concurrent => {
+                tables.create_snapshot(TransactionMode::Concurrent);
+            }
+        };
+
         vec![]
     }
 }

simulator/runner/env.rs

@@ -37,6 +37,13 @@ pub(crate) enum SimulationPhase {
     Shrink,
 }
 
+#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
+pub enum TransactionMode {
+    Read = 0,
+    Concurrent = 1,
+    Write = 2,
+}
+
 /// Represents a single operation during a transaction, applied in order.
 /// TODO: encode table/index creations etc. as operations too (although
 /// in both WAL mode and MVCC DDL requires an exclusive transaction)
@@ -59,38 +66,93 @@ pub enum RowOperation {
     },
 }
 
+/// Database snapshot
 #[derive(Debug, Clone)]
-pub struct TransactionTables {
+pub struct Snapshot {
     /// The current state after applying transaction's changes (used for reads within the transaction)
     current_tables: Vec<Table>,
     /// Operations recorded during this transaction, in order
     operations: Vec<RowOperation>,
+
+    transaction_mode: TransactionMode,
+}
+
+impl Snapshot {
+    #[inline]
+    fn set_transaction_mode(&mut self, transaction_mode: TransactionMode) {
+        self.transaction_mode = transaction_mode;
+    }
+}
+
+#[derive(Debug, Clone)]
+pub enum TransactionTables {
+    /// Signifies a Deffered transaction
+    Deffered,
+    Snapshot(Snapshot),
 }
 
 impl TransactionTables {
-    pub fn new(tables: Vec<Table>) -> Self {
-        Self {
-            current_tables: tables,
-            operations: Vec::new(),
+    #[inline]
+    fn into_snapshot(self) -> Option<Snapshot> {
+        match self {
+            TransactionTables::Deffered => None,
+            TransactionTables::Snapshot(snapshot) => Some(snapshot),
+        }
+    }
+
+    #[inline]
+    fn as_snaphot_opt(&self) -> Option<&Snapshot> {
+        match self {
+            TransactionTables::Deffered => None,
+            TransactionTables::Snapshot(snapshot) => Some(snapshot),
+        }
+    }
+
+    #[inline]
+    fn as_snaphot_mut_opt(&mut self) -> Option<&mut Snapshot> {
+        match self {
+            TransactionTables::Deffered => None,
+            TransactionTables::Snapshot(snapshot) => Some(snapshot),
         }
     }
 
+    #[inline]
+    fn as_snaphot(&self) -> &Snapshot {
+        self.as_snaphot_opt()
+            .expect("called as_snaphot on deffered transaction")
+    }
+
+    #[inline]
+    fn as_snaphot_mut(&mut self) -> &mut Snapshot {
+        self.as_snaphot_mut_opt()
+            .expect("called as_snaphot_mut on deffered transaction")
+    }
+
+    #[inline]
     pub fn record_insert(&mut self, table_name: String, row: Vec<SimValue>) {
-        self.operations
+        self.as_snaphot_mut()
+            .operations
             .push(RowOperation::Insert { table_name, row });
     }
 
+    #[inline]
     pub fn record_delete(&mut self, table_name: String, row: Vec<SimValue>) {
-        self.operations
+        self.as_snaphot_mut()
+            .operations
             .push(RowOperation::Delete { table_name, row });
     }
 
+    #[inline]
     pub fn record_drop_table(&mut self, table_name: String) {
-        self.operations.push(RowOperation::DropTable { table_name });
+        self.as_snaphot_mut()
+            .operations
+            .push(RowOperation::DropTable { table_name });
     }
 
+    #[inline]
     pub fn record_rename_table(&mut self, old_name: String, new_name: String) {
-        self.operations
+        self.as_snaphot_mut()
+            .operations
             .push(RowOperation::RenameTable { old_name, new_name });
     }
 }
@@ -110,6 +172,7 @@ pub struct ShadowTablesMut<'a> {
 impl<'a> ShadowTables<'a> {
     fn tables(&self) -> &'a Vec<Table> {
         self.transaction_tables
+            .and_then(|v| v.as_snaphot_opt())
             .map_or(self.commited_tables, |v| &v.current_tables)
     }
 }
@@ -129,15 +192,21 @@ where
     fn tables(&'a self) -> &'a Vec<Table> {
         self.transaction_tables
             .as_ref()
-            .map(|t| &t.current_tables)
-            .unwrap_or(self.commited_tables)
+            .map_or(self.commited_tables, |v| {
+                // Want to panic here with `as_snapshot_mut` because
+                // we don't want to accidently return a the commited tables on deffered transactions
+                &v.as_snaphot().current_tables
+            })
     }
 
     fn tables_mut(&'b mut self) -> &'b mut Vec<Table> {
         self.transaction_tables
             .as_mut()
-            .map(|t| &mut t.current_tables)
-            .unwrap_or(self.commited_tables)
+            .map_or(self.commited_tables, |v| {
+                // Want to panic here with `as_snapshot_mut` because
+                // we don't want to accidently return a the commited tables on deffered transactions
+                &mut v.as_snaphot_mut().current_tables
+            })
     }
 
     /// Record that a row was inserted during the current transaction
@@ -168,12 +237,60 @@ where
         }
     }
 
-    pub fn create_snapshot(&mut self) {
-        *self.transaction_tables = Some(TransactionTables::new(self.commited_tables.clone()));
+    /// Tries to upgrade the Transaction Mode
+    #[inline]
+    pub fn upgrade_transaction(&mut self, query: &Query) {
+        let transaction_mode = if query.is_write() {
+            TransactionMode::Write
+        } else if query.is_select() {
+            TransactionMode::Read
+        } else {
+            return;
+        };
+        if let Some(txn) = self.transaction_tables.as_mut() {
+            match txn {
+                TransactionTables::Deffered => self.create_snapshot(transaction_mode),
+                TransactionTables::Snapshot(snapshot) => {
+                    match (snapshot.transaction_mode, transaction_mode) {
+                        (_, TransactionMode::Concurrent) => {
+                            unreachable!();
+                        }
+                        (TransactionMode::Read, TransactionMode::Write) => {
+                            snapshot.set_transaction_mode(transaction_mode)
+                        }
+                        (TransactionMode::Concurrent, TransactionMode::Write) => {
+                            if query.is_ddl() {
+                                // Only upgrade on DDL for MVCC
+                                snapshot.set_transaction_mode(transaction_mode)
+                            }
+                        }
+                        _ => {}
+                    };
+                }
+            }
+        }
+    }
+
+    #[inline]
+    pub fn create_deffered_snapshot(&mut self) {
+        *self.transaction_tables = Some(TransactionTables::Deffered);
+    }
+
+    #[inline]
+    pub fn create_snapshot(&mut self, transaction_mode: TransactionMode) {
+        *self.transaction_tables = Some(TransactionTables::Snapshot(Snapshot {
+            current_tables: self.commited_tables.clone(),
+            operations: Vec::new(),
+            transaction_mode,
+        }));
     }
 
     pub fn apply_snapshot(&mut self) {
-        if let Some(transaction_tables) = self.transaction_tables.take() {
+        if let Some(transaction_tables) = self
+            .transaction_tables
+            .take()
+            .and_then(|transaction_tables| transaction_tables.into_snapshot())
+        {
             // Apply all operations in recorded order, including renames.
             // This ensures operations like INSERT foo, RENAME foo->bar, INSERT bar
             // are applied correctly.