From 2810b5d8a21012dbe1033ef47fb96c2d8d03e9fe Mon Sep 17 00:00:00 2001
From: PThorpe92 <preston@turso.tech>
Date: Fri, 15 Aug 2025 17:54:32 -0400
Subject: [PATCH] Add --keep-files flag and force Drop of database in simulator
 to prevent issues with the Registry

---
 simulator/generation/plan.rs  | 8 +++++---
 simulator/main.rs             | 2 +-
 simulator/runner/cli.rs       | 6 ++++++
 simulator/runner/env.rs       | 9 ++++++---
 simulator/runner/execution.rs | 5 +++--
 simulator/runner/watch.rs     | 5 +++--
 6 files changed, 24 insertions(+), 11 deletions(-)

diff --git a/simulator/generation/plan.rs b/simulator/generation/plan.rs
index c3b992cb1..2a58a3b82 100644
--- a/simulator/generation/plan.rs
+++ b/simulator/generation/plan.rs
@@ -717,6 +717,7 @@ fn reopen_database(env: &mut SimulatorEnv) {
             }
         }
         SimulationType::Default | SimulationType::Doublecheck => {
+            env.db = None;
             let db = match turso_core::Database::open_file(
                 env.io.clone(),
                 env.get_db_path().to_str().expect("path should be 'to_str'"),
@@ -734,11 +735,12 @@ fn reopen_database(env: &mut SimulatorEnv) {
                 }
             };
 
-            env.db = db;
+            env.db = Some(db);
 
             for _ in 0..num_conns {
-                env.connections
-                    .push(SimConnection::LimboConnection(env.db.connect().unwrap()));
+                env.connections.push(SimConnection::LimboConnection(
+                    env.db.as_ref().expect("db to be Some").connect().unwrap(),
+                ));
             }
         }
     };
diff --git a/simulator/main.rs b/simulator/main.rs
index 23450a3cf..d2a31f099 100644
--- a/simulator/main.rs
+++ b/simulator/main.rs
@@ -136,7 +136,7 @@ fn testing_main(cli_opts: &SimulatorCLI) -> anyhow::Result<()> {
     println!("seed: {seed}");
     println!("path: {}", paths.base.display());
 
-    if result.is_ok() {
+    if !cli_opts.keep_files && result.is_ok() {
         paths.delete_all_files();
     }
 
diff --git a/simulator/runner/cli.rs b/simulator/runner/cli.rs
index 3bfae4349..b67b6fef3 100644
--- a/simulator/runner/cli.rs
+++ b/simulator/runner/cli.rs
@@ -129,6 +129,12 @@ pub struct SimulatorCLI {
     pub experimental_mvcc: bool,
     #[clap(long, help = "Disable experimental indexing feature")]
     pub disable_experimental_indexes: bool,
+    #[clap(
+        long,
+        help = "Keep all database and plan files",
+        default_value_t = false
+    )]
+    pub keep_files: bool,
 }
 
 #[derive(Parser, Debug, Clone, Serialize, Deserialize, PartialEq, PartialOrd, Eq, Ord)]
diff --git a/simulator/runner/env.rs b/simulator/runner/env.rs
index ec0b389f9..c061cfc52 100644
--- a/simulator/runner/env.rs
+++ b/simulator/runner/env.rs
@@ -62,7 +62,7 @@ pub(crate) struct SimulatorEnv {
     pub(crate) opts: SimulatorOpts,
     pub(crate) connections: Vec<SimConnection>,
     pub(crate) io: Arc<SimulatorIO>,
-    pub(crate) db: Arc<Database>,
+    pub(crate) db: Option<Arc<Database>>,
     pub(crate) rng: ChaCha8Rng,
     pub(crate) paths: Paths,
     pub(crate) type_: SimulationType,
@@ -115,6 +115,7 @@ impl SimulatorEnv {
         if wal_path.exists() {
             std::fs::remove_file(&wal_path).unwrap();
         }
+        self.db = None;
 
         let db = match Database::open_file(
             io.clone(),
@@ -129,7 +130,7 @@ impl SimulatorEnv {
             }
         };
         self.io = io;
-        self.db = db;
+        self.db = Some(db);
     }
 
     pub(crate) fn get_db_path(&self) -> PathBuf {
@@ -298,7 +299,7 @@ impl SimulatorEnv {
             paths,
             rng,
             io,
-            db,
+            db: Some(db),
             type_: simulation_type,
             phase: SimulationPhase::Test,
         }
@@ -320,6 +321,8 @@ impl SimulatorEnv {
             SimulationType::Default | SimulationType::Doublecheck => {
                 self.connections[connection_index] = SimConnection::LimboConnection(
                     self.db
+                        .as_ref()
+                        .expect("db to be Some")
                         .connect()
                         .expect("Failed to connect to Limbo database"),
                 );
diff --git a/simulator/runner/execution.rs b/simulator/runner/execution.rs
index d943ebc85..206423e4e 100644
--- a/simulator/runner/execution.rs
+++ b/simulator/runner/execution.rs
@@ -123,8 +123,9 @@ fn execute_plan(
 
     if let SimConnection::Disconnected = connection {
         tracing::debug!("connecting {}", connection_index);
-        env.connections[connection_index] =
-            SimConnection::LimboConnection(env.db.connect().unwrap());
+        env.connections[connection_index] = SimConnection::LimboConnection(
+            env.db.as_ref().expect("db to be Some").connect().unwrap(),
+        );
     } else {
         tracing::debug!("connection {} already connected", connection_index);
         match execute_interaction(env, connection_index, interaction, &mut state.stack) {
diff --git a/simulator/runner/watch.rs b/simulator/runner/watch.rs
index 7c3795f01..90e8edc68 100644
--- a/simulator/runner/watch.rs
+++ b/simulator/runner/watch.rs
@@ -98,8 +98,9 @@ fn execute_plan(
 
     if let SimConnection::Disconnected = connection {
         tracing::debug!("connecting {}", connection_index);
-        env.connections[connection_index] =
-            SimConnection::LimboConnection(env.db.connect().unwrap());
+        env.connections[connection_index] = SimConnection::LimboConnection(
+            env.db.as_ref().expect("db to be Some").connect().unwrap(),
+        );
     } else {
         match execute_interaction(env, connection_index, interaction, &mut state.stack) {
             Ok(next_execution) => {