bench: Implement uniform_v1_std workload

Additional changes:
- Simplified setup of benchmark
- Updated help menu

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

sky-bench/help_text/help

@@ -21,9 +21,14 @@ OPTIONS:
     --keysize     Set the default primary key size. defaults to 7
     --rowcount    Set the number of rows to be manipulated for the benchmark
                   Defaults to 1,000,000 rows.
-    --engine      Set the engine for benchmarking. `rookie` is the stable engine
+    --workload    Set the workload to run (see below).
+    --engine      (DEPRECATED) Set the engine for benchmarking. `rookie` is the stable engine
                   and `fury` is the new experimental engine. Defaults to `fury`
 
+WORKLOADS:
+- 'std_uniform_v1': (current default) This real-world workload creates and manipulates unique
+                    rows with an uniform distribution of executed queries
+
 NOTES:
     - If no password is supplied, we look for the `{password_env_var}`
       environment variable

sky-bench/src/args.rs

@@ -28,6 +28,7 @@ use {
     crate::{
         error::{BenchError, BenchResult},
         setup,
+        workload::{workloads, Workload},
     },
     libsky::{env_vars, CliAction},
     std::{collections::hash_map::HashMap, env},
@@ -180,15 +181,21 @@ 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),
+            workloads::UniformV1Std::ID => 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 '{}'",
+                    workloads::UniformV1Std::ID
+                );
+                BenchType::Workload(BenchWorkload::UniformV1)
+            }
             Some(engine) => BenchType::Legacy(match engine.as_str() {
                 "rookie" => BenchEngine::Rookie,
                 "fury" => BenchEngine::Fury,

sky-bench/src/bench.rs

@@ -24,21 +24,21 @@
  *
 */
 
-use skytable::response::Value;
-
-use crate::{
-    args::{BenchEngine, BenchType, BenchWorkload},
-    setup, workload,
-};
-
 use {
     crate::{
-        args::BenchConfig,
+        args::{BenchConfig, BenchEngine, BenchType, BenchWorkload},
         error::{self, BenchResult},
         legacy::runtime::{fury, rookie},
+        setup,
         stats::RuntimeStats,
+        workload::{self, workloads},
+    },
+    skytable::{
+        error::Error,
+        query,
+        response::{Response, Value},
+        Config, Connection, Query,
     },
-    skytable::{error::Error, query, response::Response, Config, Connection, Query},
     std::{fmt, time::Instant},
 };
 
@@ -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(workload) => match workload {
+            BenchWorkload::UniformV1 => workload::run_bench(workloads::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(())
 }
 
@@ -205,24 +210,18 @@ fn print_table(data: Vec<(&'static str, RuntimeStats)>) {
 pub struct BenchmarkTask {
     gen_query: fn(&Self, u64) -> Query,
     check_resp: fn(&Self, u64, Response) -> bool,
-    pk_len: usize,
 }
 
 impl BenchmarkTask {
     fn new(
-        pk_len: usize,
         gen_query: fn(&Self, u64) -> Query,
         check_resp: fn(&Self, u64, Response) -> bool,
     ) -> Self {
         Self {
             gen_query,
             check_resp,
-            pk_len,
         }
     }
-    fn fmt_pk(&self, current: u64) -> Vec<u8> {
-        format!("{:0>width$}", current, width = self.pk_len).into_bytes()
-    }
     pub fn generate_query(&self, current: u64) -> Query {
         (self.gen_query)(self, current)
     }
@@ -266,20 +265,33 @@ fn prepare_bench_spec() -> Vec<BenchItem> {
         BenchItem::new(
             "INSERT",
             BenchmarkTask::new(
-                config_instance.object_size(),
-                |me, current| query!("insert into bench(?, ?)", me.fmt_pk(current), 0u64),
+                |_, current| {
+                    query!(
+                        "insert into bench(?, ?)",
+                        unsafe { setup::instance() }.fmt_pk(current),
+                        0u64
+                    )
+                },
                 |_, _, actual_resp| actual_resp == Response::Empty,
             ),
             config_instance.object_count(),
         ),
         BenchItem::new(
             "SELECT",
             BenchmarkTask::new(
-                config_instance.object_size(),
-                |me, current| query!("select * from bench where un = ?", me.fmt_pk(current)),
-                |me, current, resp| match resp {
+                |_, current| {
+                    query!(
+                        "select * from bench where un = ?",
+                        unsafe { setup::instance() }.fmt_pk(current)
+                    )
+                },
+                |_, current, resp| match resp {
                     Response::Row(r) => {
-                        r.into_values() == vec![Value::Binary(me.fmt_pk(current)), Value::UInt8(0)]
+                        r.into_values()
+                            == vec![
+                                Value::Binary(unsafe { setup::instance() }.fmt_pk(current)),
+                                Value::UInt8(0),
+                            ]
                     }
                     _ => false,
                 },
@@ -289,12 +301,11 @@ fn prepare_bench_spec() -> Vec<BenchItem> {
         BenchItem::new(
             "UPDATE",
             BenchmarkTask::new(
-                config_instance.object_size(),
-                |me, current| {
+                |_, current| {
                     query!(
                         "update bench set pw += ? where un = ?",
                         1u64,
-                        me.fmt_pk(current)
+                        unsafe { setup::instance() }.fmt_pk(current)
                     )
                 },
                 |_, _, resp| resp == Response::Empty,
@@ -304,8 +315,12 @@ fn prepare_bench_spec() -> Vec<BenchItem> {
         BenchItem::new(
             "DELETE",
             BenchmarkTask::new(
-                config_instance.object_size(),
-                |me, current| query!("delete from bench where un = ?", me.fmt_pk(current)),
+                |_, current| {
+                    query!(
+                        "delete from bench where un = ?",
+                        unsafe { setup::instance() }.fmt_pk(current)
+                    )
+                },
                 |_, _, resp| resp == Response::Empty,
             ),
             config_instance.object_count(),

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)
     }
 }

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,
@@ -73,6 +74,12 @@ impl RunnerSetup {
     pub fn object_count(&self) -> usize {
         self.object_count
     }
+    pub fn fmt_pk(&self, current: u64) -> Vec<u8> {
+        Self::_fmt_pk(current, self.object_size())
+    }
+    fn _fmt_pk(current: u64, width: usize) -> Vec<u8> {
+        format!("{current:0>width$}",).into_bytes()
+    }
 }
 
 pub unsafe fn instance() -> &'static RunnerSetup {
@@ -98,3 +105,11 @@ pub unsafe fn configure(
     SETUP.object_size = object_size;
     SETUP.object_count = object_count;
 }
+
+#[test]
+fn fmt_pk() {
+    assert_eq!(
+        RunnerSetup::_fmt_pk(123456789, 18),
+        "000000000123456789".as_bytes()
+    );
+}