diff --git a/stress/main.rs b/stress/main.rs index 6bd7a6bef..6180faf30 100644 --- a/stress/main.rs +++ b/stress/main.rs @@ -9,7 +9,9 @@ use limbo::Builder; use opts::Opts; use serde_json::json; use std::collections::HashSet; -use std::sync::Arc; +use std::fs::File; +use std::io::Write; +use std::sync::{Arc, Mutex}; /// Represents a column in a SQLite table #[derive(Debug, Clone)] @@ -297,11 +299,19 @@ async fn main() -> Result<(), Box> { let schema = gen_schema(); let ddl_statements = schema.to_sql(); - for stmt in &ddl_statements { - println!("{}", stmt); - } let opts = Opts::parse(); + let log_file = File::create(&opts.log_file)?; + let log_file = Arc::new(Mutex::new(log_file)); + + // Write DDL statements to log file + { + let mut file = log_file.lock().unwrap(); + for stmt in &ddl_statements { + writeln!(file, "{}", stmt)?; + } + } + let mut handles = Vec::with_capacity(opts.nr_threads); for _ in 0..opts.nr_threads { @@ -318,12 +328,16 @@ async fn main() -> Result<(), Box> { let nr_iterations = opts.nr_iterations; let db = db.clone(); let schema = schema.clone(); + let log_file = log_file.clone(); let handle = tokio::spawn(async move { let conn = db.connect()?; for _ in 0..nr_iterations { let sql = generate_random_statement(&schema); - println!("{}", sql); + { + let mut file = log_file.lock().unwrap(); + writeln!(file, "{}", sql)?; + } if let Err(e) = conn.execute(&sql, ()).await { println!("Error: {}", e); } @@ -336,6 +350,6 @@ async fn main() -> Result<(), Box> { for handle in handles { handle.await??; } - println!("Done."); + println!("Done. SQL statements written to {}", opts.log_file); Ok(()) } diff --git a/stress/opts.rs b/stress/opts.rs index da00e1a00..e7799d29a 100644 --- a/stress/opts.rs +++ b/stress/opts.rs @@ -4,8 +4,11 @@ use clap::{command, Parser}; #[command(name = "limbo_stress")] #[command(author, version, about, long_about = None)] pub struct Opts { + /// Number of threads to run #[clap(short = 't', long, help = "the number of threads", default_value_t = 8)] pub nr_threads: usize, + + /// Number of iterations per thread #[clap( short = 'i', long, @@ -13,4 +16,13 @@ pub struct Opts { default_value_t = 100000 )] pub nr_iterations: usize, + + /// Log file for SQL statements + #[clap( + short = 'l', + long, + help = "log file for SQL statements", + default_value = "limbostress.log" + )] + pub log_file: String, }