consolidate SimulatorEnv rollback code

simulator/generation/plan.rs

@@ -216,10 +216,8 @@ impl InteractionPlan {
         let num_interactions = env.opts.max_interactions as usize;
         if self.len() < num_interactions {
             let conn_index = env.choose_conn(rng);
-            dbg!(self.mvcc, env.conn_in_transaction(conn_index));
             let interactions = if self.mvcc && !env.conn_in_transaction(conn_index) {
                 let query = Query::Begin(Begin::Concurrent);
-                env.update_conn_last_interaction(conn_index, Some(&query));
                 Interactions::new(conn_index, InteractionsType::Query(query))
             } else if self.mvcc
                 && env.conn_in_transaction(conn_index)
@@ -227,10 +225,8 @@ impl InteractionPlan {
                 && rng.random_bool(0.4)
             {
                 let query = Query::Commit(Commit);
-                env.update_conn_last_interaction(conn_index, Some(&query));
                 Interactions::new(conn_index, InteractionsType::Query(query))
             } else {
-                env.update_conn_last_interaction(conn_index, None);
                 let conn_ctx = &env.connection_context(conn_index);
                 Interactions::arbitrary_from(rng, conn_ctx, (env, self.stats(), conn_index))
             };
@@ -786,13 +782,15 @@ impl InteractionType {
                 match fault {
                     Fault::Disconnect => {
                         if env.connections[conn_index].is_connected() {
+                            if env.conn_in_transaction(conn_index) {
+                                env.rollback_conn(conn_index);
+                            }
                             env.connections[conn_index].disconnect();
                         } else {
                             return Err(turso_core::LimboError::InternalError(
                                 "connection already disconnected".into(),
                             ));
                         }
-                        env.connections[conn_index] = SimConnection::Disconnected;
                     }
                     Fault::ReopenDatabase => {
                         reopen_database(env);

simulator/generation/property.rs

@@ -454,7 +454,7 @@ impl Property {
                                     Err(e) => {
                                         if e.to_string().to_lowercase().contains(&format!("table {table_name} already exists")) {
                                              // On error we rollback the transaction if there is any active here
-                                            Rollback.shadow(&mut env.get_conn_tables_mut(connection_index));
+                                            env.rollback_conn(connection_index);
                                             Ok(Ok(()))
                                         } else {
                                             Ok(Err(format!("expected table already exists error, got: {e}")))
@@ -808,7 +808,7 @@ impl Property {
                                 tracing::error!("Fault injection produced error: {err}");
 
                                 // On error we rollback the transaction if there is any active here
-                                Rollback.shadow(&mut env.get_conn_tables_mut(connection_index));
+                                env.rollback_conn(connection_index);
                                 Ok(Ok(()))
                             }
                         }

simulator/runner/env.rs

@@ -10,9 +10,11 @@ use garde::Validate;
 use rand::{Rng, SeedableRng};
 use rand_chacha::ChaCha8Rng;
 use sql_generation::generation::GenerationContext;
+use sql_generation::model::query::transaction::Rollback;
 use sql_generation::model::table::Table;
 use turso_core::Database;
 
+use crate::generation::Shadow;
 use crate::model::Query;
 use crate::profiles::Profile;
 use crate::runner::SimIO;
@@ -495,6 +497,11 @@ impl SimulatorEnv {
         self.connection_last_query.set(conn_index, value);
     }
 
+    pub fn rollback_conn(&mut self, conn_index: usize) {
+        Rollback.shadow(&mut self.get_conn_tables_mut(conn_index));
+        self.update_conn_last_interaction(conn_index, Some(&Query::Rollback(Rollback)));
+    }
+
     pub fn get_conn_tables<'a>(&'a self, conn_index: usize) -> ShadowTables<'a> {
         ShadowTables {
             transaction_tables: self.connection_tables.get(conn_index).unwrap().as_ref(),

simulator/runner/execution.rs

@@ -160,7 +160,7 @@ pub fn execute_interaction(
     }
 }
 
-#[instrument(skip(env, interaction, stack), fields(seed = %env.opts.seed, interaction = %interaction))]
+#[instrument(skip(env, interaction, stack), fields(conn_index = interaction.connection_index, interaction = %interaction))]
 pub fn execute_interaction_turso(
     env: &mut SimulatorEnv,
     interaction: &Interaction,
@@ -173,7 +173,7 @@ pub fn execute_interaction_turso(
     // Leave this empty info! here to print the span of the execution
     tracing::info!("");
     match &interaction.interaction {
-        InteractionType::Query(_) => {
+        InteractionType::Query(query) => {
             tracing::debug!(?interaction);
             let results = interaction
                 .execute_query(conn)
@@ -189,6 +189,7 @@ pub fn execute_interaction_turso(
             if !env.profile.experimental_mvcc {
                 limbo_integrity_check(conn)?;
             }
+            env.update_conn_last_interaction(interaction.connection_index, Some(query));
         }
         InteractionType::FsyncQuery(query) => {
             let results = interaction
@@ -308,6 +309,7 @@ fn execute_interaction_rusqlite(
             }
             tracing::debug!("{:?}", results);
             stack.push(results);
+            env.update_conn_last_interaction(interaction.connection_index, Some(query));
         }
         InteractionType::FsyncQuery(..) => {
             unimplemented!("cannot implement fsync query in rusqlite, as we do not control IO");