Single now takes batch size argument

benchmark/src/main.rs

@@ -84,7 +84,7 @@ mod single;
 mod tenkrandom;
 mod zipf;
 
-#[derive(Debug, Subcommand)]
+#[derive(Debug, Subcommand, PartialEq)]
 enum TestName {
     Create,
     TenKRandom,
@@ -122,6 +122,10 @@ static GLOBAL: tikv_jemallocator::Jemalloc = tikv_jemallocator::Jemalloc;
 async fn main() -> Result<(), Box<dyn Error>> {
     let args = Args::parse();
 
+    if args.test_name == TestName::Single && args.batch_size > 1000 {
+        panic!("Single test is not designed to handle batch sizes > 1000");
+    }
+
     env_logger::Builder::new()
         .filter_level(match args.global_opts.log_level.as_str() {
             "debug" => LevelFilter::Debug,

benchmark/src/single.rs

@@ -16,14 +16,19 @@ pub struct Single;
 impl TestRunner for Single {
     async fn run(&self, db: &Db, args: &crate::Args) -> Result<(), Box<dyn Error>> {
         let start = Instant::now();
-        let inner_key = Sha256::digest(0u64.to_ne_bytes()).to_vec();
+        let inner_keys: Vec<_> = (0..args.batch_size)
+            .map(|i| Sha256::digest(i.to_ne_bytes()))
+            .collect();
         let mut batch_id = 0;
 
         while start.elapsed().as_secs() / 60 < args.global_opts.duration_minutes {
-            let batch = vec![BatchOp::Put {
-                key: inner_key.clone(),
-                value: vec![batch_id as u8],
-            }];
+            let batch = inner_keys
+                .iter()
+                .map(|key| BatchOp::Put {
+                    key,
+                    value: vec![batch_id as u8],
+                })
+                .collect();
             let proposal = db.propose(batch).await.expect("proposal should succeed");
             proposal.commit().await?;
 

benchmark/src/zipf.rs

@@ -13,7 +13,7 @@ use std::collections::HashSet;
 use std::error::Error;
 use std::time::Instant;
 
-#[derive(clap::Args, Debug)]
+#[derive(clap::Args, Debug, PartialEq)]
 pub struct Args {
     #[arg(short, long, help = "zipf exponent", default_value_t = 1.2)]
     exponent: f64,