From e57a0ad932987491e58b2ba176a7c1b307657533 Mon Sep 17 00:00:00 2001
From: Ron Kuris <ron.kuris@avalabs.org>
Date: Mon, 18 Nov 2024 09:18:54 -0800
Subject: [PATCH] Single now takes batch size argument (#753)

---
 benchmark/src/main.rs   |  6 +++++-
 benchmark/src/single.rs | 15 ++++++++++-----
 benchmark/src/zipf.rs   |  2 +-
 3 files changed, 16 insertions(+), 7 deletions(-)

diff --git a/benchmark/src/main.rs b/benchmark/src/main.rs
index 9c8b2e49..30d4aff6 100644
--- a/benchmark/src/main.rs
+++ b/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,
diff --git a/benchmark/src/single.rs b/benchmark/src/single.rs
index b43684b9..e90ae354 100644
--- a/benchmark/src/single.rs
+++ b/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?;
 
diff --git a/benchmark/src/zipf.rs b/benchmark/src/zipf.rs
index 4e6dfba6..a054abdf 100644
--- a/benchmark/src/zipf.rs
+++ b/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,