bench: Implement uniform_v1_std workload

skytable · Apr 28, 2024 · 54b0b0d · 54b0b0d
1 parent 2c81da4
commit 54b0b0d
Show file tree

Hide file tree

Showing 9 changed files with 481 additions and 522 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -2,6 +2,18 @@
 
 All changes in this project will be noted in this file.
 
+## Version 0.8.3
+
+### Additions
+
+- Benchmark tool:
+  - Switch to using `workload` based benchmarks
+
+### Fixes
+
+- Benchmark tool:
+  - Running a SIGINT now gracefully terminates the workload
+
 ## Version 0.8.2
 
 ### Additions

diff --git a/sky-bench/src/args.rs b/sky-bench/src/args.rs
@@ -180,15 +180,18 @@ pub fn parse_and_setup() -> BenchResult<Task> {
     };
     let workload = match args.remove("--workload") {
         Some(workload) => match workload.as_str() {
-            "uniform_v1" => BenchType::Workload(BenchWorkload::UniformV1),
+            "uniform_v1_std" => BenchType::Workload(BenchWorkload::UniformV1),
             _ => {
                 return Err(BenchError::ArgsErr(format!(
                     "unknown workload choice {workload}"
                 )))
             }
         },
         None => match args.remove("--engine") {
-            None => BenchType::Workload(BenchWorkload::UniformV1),
+            None => {
+                warn!("workload not specified. choosing default workload 'uniform_v1_std'");
+                BenchType::Workload(BenchWorkload::UniformV1)
+            }
             Some(engine) => BenchType::Legacy(match engine.as_str() {
                 "rookie" => BenchEngine::Rookie,
                 "fury" => BenchEngine::Fury,

diff --git a/sky-bench/src/bench.rs b/sky-bench/src/bench.rs
@@ -125,17 +125,20 @@ pub fn run(bench: BenchConfig) -> error::BenchResult<()> {
         config_instance.username(),
         config_instance.password(),
     ));
-    let mut main_thread_db;
+    let mut main_thread_db = None;
     let stats = match bench.workload {
-        BenchType::Workload(BenchWorkload::UniformV1) => return workload::run_bench(),
+        BenchType::Workload(BenchWorkload::UniformV1) => {
+            workload::run_bench(workload::UniformV1Std::new())
+        }
         BenchType::Legacy(l) => {
             warn!("using `--engine` is now deprecated. please consider switching to `--workload`");
             info!("running preliminary checks and creating model `bench.bench` with definition: `{{un: binary, pw: uint8}}`");
-            main_thread_db = bench_config.config.connect()?;
-            main_thread_db.query_parse::<()>(&query!("create space bench"))?;
-            main_thread_db.query_parse::<()>(&query!(format!(
+            let mut mt_db = bench_config.config.connect()?;
+            mt_db.query_parse::<()>(&query!("create space bench"))?;
+            mt_db.query_parse::<()>(&query!(format!(
                 "create model {BENCHMARK_SPACE_ID}.{BENCHMARK_MODEL_ID}(un: binary, pw: uint8)"
             )))?;
+            main_thread_db = Some(mt_db);
             match l {
                 BenchEngine::Rookie => bench_rookie(bench_config),
                 BenchEngine::Fury => bench_fury(),
@@ -170,9 +173,11 @@ pub fn run(bench: BenchConfig) -> error::BenchResult<()> {
     util
 */
 
-fn cleanup(mut main_thread_db: Connection) -> Result<(), error::BenchError> {
+fn cleanup(main_thread_db: Option<Connection>) -> Result<(), error::BenchError> {
     trace!("dropping space and table");
-    main_thread_db.query_parse::<()>(&query!("drop space allow not empty bench"))?;
+    if let Some(mut db) = main_thread_db {
+        db.query_parse::<()>(&query!("drop space allow not empty bench"))?;
+    }
     Ok(())
 }
 

diff --git a/sky-bench/src/error.rs b/sky-bench/src/error.rs
@@ -28,7 +28,7 @@ use {
     crate::{
         bench::BombardTask,
         legacy::runtime::{fury, rookie::BombardError},
-        workload::error::WorkloadDriverError,
+        workload::error::WorkloadError,
     },
     core::fmt,
     skytable::error::Error,
@@ -40,7 +40,7 @@ pub type BenchResult<T> = Result<T, BenchError>;
 pub enum BenchError {
     ArgsErr(String),
     DirectDbError(Error),
-    WorkloadDriverError(WorkloadDriverError),
+    WorkloadDriverError(WorkloadError),
     // legacy
     LegacyRookieEngineError(BombardError<BombardTask>),
     LegacyFuryEngineError(fury::FuryError),
@@ -52,8 +52,8 @@ impl From<fury::FuryError> for BenchError {
     }
 }
 
-impl From<WorkloadDriverError> for BenchError {
-    fn from(e: WorkloadDriverError) -> Self {
+impl From<WorkloadError> for BenchError {
+    fn from(e: WorkloadError) -> Self {
         Self::WorkloadDriverError(e)
     }
 }

diff --git a/sky-bench/src/setup.rs b/sky-bench/src/setup.rs
@@ -37,6 +37,7 @@ static mut SETUP: RunnerSetup = RunnerSetup {
     object_count: 0,
 };
 
+#[derive(Debug)]
 pub struct RunnerSetup {
     username: String,
     password: String,