From e2aafacbb4f3e3dc5ee1bb647a40ec26f0b574ae Mon Sep 17 00:00:00 2001
From: pedrocarlo <pedro.muniz.carlo@gmail.com>
Date: Wed, 18 Jun 2025 00:52:06 -0300
Subject: [PATCH] add cli option to adjust latency probability

---
 simulator/runner/cli.rs  | 13 +++++++++++++
 simulator/runner/env.rs  |  3 ++-
 simulator/runner/file.rs | 11 ++++++++---
 simulator/runner/io.rs   |  5 ++++-
 4 files changed, 27 insertions(+), 5 deletions(-)

diff --git a/simulator/runner/cli.rs b/simulator/runner/cli.rs
index 3b3bd6452..6e1bb5c9a 100644
--- a/simulator/runner/cli.rs
+++ b/simulator/runner/cli.rs
@@ -84,6 +84,12 @@ pub struct SimulatorCLI {
     pub disable_select_optimizer: bool,
     #[clap(long, help = "disable Reopen-Database fault", default_value_t = false)]
     pub disable_reopen_database: bool,
+    #[clap(
+        long = "latency_prob",
+        help = "added IO latency probability",
+        default_value_t = 0
+    )]
+    pub latency_probability: usize,
 }
 
 #[derive(Parser, Debug, Clone, Serialize, Deserialize, PartialEq, PartialOrd, Eq, Ord)]
@@ -141,6 +147,13 @@ impl SimulatorCLI {
             anyhow::bail!("Cannot set seed and load plan at the same time");
         }
 
+        if self.latency_probability > 100 {
+            anyhow::bail!(
+                "latency probability must be a number between 0 and 100. Got `{}`",
+                self.latency_probability
+            );
+        }
+
         Ok(())
     }
 }
diff --git a/simulator/runner/env.rs b/simulator/runner/env.rs
index 49eeb15de..606aec768 100644
--- a/simulator/runner/env.rs
+++ b/simulator/runner/env.rs
@@ -123,7 +123,8 @@ impl SimulatorEnv {
             disable_reopen_database: cli_opts.disable_reopen_database,
         };
 
-        let io = Arc::new(SimulatorIO::new(seed, opts.page_size).unwrap());
+        let io =
+            Arc::new(SimulatorIO::new(seed, opts.page_size, cli_opts.latency_probability).unwrap());
 
         // Remove existing database file if it exists
         if db_path.exists() {
diff --git a/simulator/runner/file.rs b/simulator/runner/file.rs
index 4ef3ef5b8..435a60997 100644
--- a/simulator/runner/file.rs
+++ b/simulator/runner/file.rs
@@ -6,6 +6,7 @@ use std::{
 use limbo_core::{File, Result};
 use rand::Rng as _;
 use rand_chacha::ChaCha8Rng;
+use tracing::{instrument, Level};
 pub(crate) struct SimulatorFile {
     pub(crate) inner: Arc<dyn File>,
     pub(crate) fault: Cell<bool>,
@@ -28,6 +29,8 @@ pub(crate) struct SimulatorFile {
     pub(crate) page_size: usize,
 
     pub(crate) rng: RefCell<ChaCha8Rng>,
+
+    pub latency_probability: usize,
 }
 
 unsafe impl Send for SimulatorFile {}
@@ -67,12 +70,13 @@ impl SimulatorFile {
         stats_table.join("\n")
     }
 
+    #[instrument(skip_all, level = Level::TRACE)]
     fn generate_latency(&self) {
         let mut rng = self.rng.borrow_mut();
-        // 5% chance to introduce some latency
-        if rng.gen_bool(0.05) {
+        // Chance to introduce some latency
+        if rng.gen_bool(self.latency_probability as f64 / 100.0) {
             let latency = std::time::Duration::from_millis(rng.gen_range(20..50));
-            tracing::trace!(?latency, "latency");
+            tracing::trace!(?latency);
             std::thread::sleep(latency);
         }
     }
@@ -85,6 +89,7 @@ impl File for SimulatorFile {
                 "Injected fault".into(),
             ));
         }
+        self.generate_latency();
         self.inner.lock_file(exclusive)
     }
 
diff --git a/simulator/runner/io.rs b/simulator/runner/io.rs
index 660ddc605..bd5710c40 100644
--- a/simulator/runner/io.rs
+++ b/simulator/runner/io.rs
@@ -17,13 +17,14 @@ pub(crate) struct SimulatorIO {
     pub(crate) nr_run_once_faults: Cell<usize>,
     pub(crate) page_size: usize,
     seed: u64,
+    latency_probability: usize,
 }
 
 unsafe impl Send for SimulatorIO {}
 unsafe impl Sync for SimulatorIO {}
 
 impl SimulatorIO {
-    pub(crate) fn new(seed: u64, page_size: usize) -> Result<Self> {
+    pub(crate) fn new(seed: u64, page_size: usize, latency_probability: usize) -> Result<Self> {
         let inner = Box::new(PlatformIO::new()?);
         let fault = Cell::new(false);
         let files = RefCell::new(Vec::new());
@@ -37,6 +38,7 @@ impl SimulatorIO {
             nr_run_once_faults,
             page_size,
             seed,
+            latency_probability,
         })
     }
 
@@ -82,6 +84,7 @@ impl IO for SimulatorIO {
             nr_sync_calls: Cell::new(0),
             page_size: self.page_size,
             rng: RefCell::new(ChaCha8Rng::seed_from_u64(self.seed)),
+            latency_probability: self.latency_probability,
         });
         self.files.borrow_mut().push(file.clone());
         Ok(file)