From cd0e521232759c14797f6425685a80144a2a6119 Mon Sep 17 00:00:00 2001
From: Xun Li <lxfind@gmail.com>
Date: Thu, 31 Oct 2024 09:53:56 -0700
Subject: [PATCH] Make number of tasks configurable

---
 src/benchmarks/kms_stress.rs |  6 +++---
 src/bin/tool.rs              | 10 ++++++++--
 2 files changed, 11 insertions(+), 5 deletions(-)

diff --git a/src/benchmarks/kms_stress.rs b/src/benchmarks/kms_stress.rs
index d1e974a..28d8705 100644
--- a/src/benchmarks/kms_stress.rs
+++ b/src/benchmarks/kms_stress.rs
@@ -8,7 +8,7 @@ use std::time::{Duration, Instant};
 use sui_types::base_types::{random_object_ref, SuiAddress};
 use sui_types::transaction::{ProgrammableTransaction, TransactionData, TransactionKind};
 
-pub async fn run_kms_stress_test(kms_url: String) {
+pub async fn run_kms_stress_test(kms_url: String, num_tasks: usize) {
     let signer = SidecarTxSigner::new(kms_url).await;
     let test_tx_data = TransactionData::new(
         TransactionKind::ProgrammableTransaction(ProgrammableTransaction {
@@ -28,9 +28,9 @@ pub async fn run_kms_stress_test(kms_url: String) {
     let stop_flag = Arc::new(AtomicBool::new(false));
 
     // Vector to hold worker task handles
-    let mut handles = Vec::with_capacity(200);
+    let mut handles = Vec::with_capacity(num_tasks);
 
-    for _ in 0..200 {
+    for _ in 0..num_tasks {
         let success_counter = Arc::clone(&success_counter);
         let failure_counter = Arc::clone(&failure_counter);
         let stop_flag = Arc::clone(&stop_flag);
diff --git a/src/bin/tool.rs b/src/bin/tool.rs
index dee73d3..55d2b4f 100644
--- a/src/bin/tool.rs
+++ b/src/bin/tool.rs
@@ -42,6 +42,12 @@ pub enum ToolCommand {
     StressKMS {
         #[arg(long, help = "Full URL to the KMS signer")]
         kms_url: String,
+        #[arg(
+            long,
+            default_value_t = 300,
+            help = "Number of tasks to spawn to send requests to servers."
+        )]
+        num_tasks: usize,
     },
     /// Generate a sample config file and put it in the specified path.
     #[clap(name = "generate-sample-config")]
@@ -94,8 +100,8 @@ impl ToolCommand {
                     .run_benchmark(gas_station_url, reserve_duration_sec, num_clients)
                     .await
             }
-            ToolCommand::StressKMS { kms_url } => {
-                run_kms_stress_test(kms_url).await;
+            ToolCommand::StressKMS { kms_url, num_tasks } => {
+                run_kms_stress_test(kms_url, num_tasks).await;
             }
             ToolCommand::GenerateSampleConfig {
                 config_path,