diff --git a/src/benchmarks/kms_stress.rs b/src/benchmarks/kms_stress.rs new file mode 100644 index 0000000..d1e974a --- /dev/null +++ b/src/benchmarks/kms_stress.rs @@ -0,0 +1,89 @@ +// Copyright (c) Mysten Labs, Inc. +// SPDX-License-Identifier: Apache-2.0 + +use crate::tx_signer::{SidecarTxSigner, TxSigner}; +use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; +use std::sync::Arc; +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) { + let signer = SidecarTxSigner::new(kms_url).await; + let test_tx_data = TransactionData::new( + TransactionKind::ProgrammableTransaction(ProgrammableTransaction { + inputs: vec![], + commands: vec![], + }), + SuiAddress::ZERO, + random_object_ref(), + 1000, + 1000, + ); + // Shared atomic counters for successes and failures + let success_counter = Arc::new(AtomicUsize::new(0)); + let failure_counter = Arc::new(AtomicUsize::new(0)); + + // Shared atomic flag to signal workers to stop + let stop_flag = Arc::new(AtomicBool::new(false)); + + // Vector to hold worker task handles + let mut handles = Vec::with_capacity(200); + + for _ in 0..200 { + let success_counter = Arc::clone(&success_counter); + let failure_counter = Arc::clone(&failure_counter); + let stop_flag = Arc::clone(&stop_flag); + let signer = signer.clone(); // Ensure Signer is Clone + let test_tx_data = test_tx_data.clone(); // Clone the test data + + // Spawn a worker task + let handle = tokio::spawn(async move { + while !stop_flag.load(Ordering::Relaxed) { + match signer.sign_transaction(&test_tx_data).await { + Ok(_) => { + // Increment the success counter on successful call + success_counter.fetch_add(1, Ordering::Relaxed); + } + Err(_) => { + // Increment the failure counter on failed call + failure_counter.fetch_add(1, Ordering::Relaxed); + } + } + } + }); + + handles.push(handle); + } + + // Start the timer + let start = Instant::now(); + + // Let the test run for the specified duration + tokio::time::sleep(Duration::from_secs(30)).await; + + // Signal all workers to stop + stop_flag.store(true, Ordering::Relaxed); + + // Wait for all workers to finish + for handle in handles { + let _ = handle.await; + } + + // Calculate elapsed time + let elapsed = start.elapsed().as_secs_f64(); + + // Get the total number of successful and failed calls + let total_success = success_counter.load(Ordering::Relaxed); + let total_failures = failure_counter.load(Ordering::Relaxed); + let total_calls = total_success + total_failures; + + // Calculate throughput + let throughput = total_success as f64 / elapsed; + + println!("Total calls completed: {}", total_calls); + println!(" - Successful: {}", total_success); + println!(" - Failed: {}", total_failures); + println!("Elapsed time: {:.2} seconds", elapsed); + println!("Throughput: {:.2} successful calls/second", throughput); +} diff --git a/src/benchmarks/mod.rs b/src/benchmarks/mod.rs index 16fa8f9..fa7d15b 100644 --- a/src/benchmarks/mod.rs +++ b/src/benchmarks/mod.rs @@ -1,6 +1,8 @@ // Copyright (c) Mysten Labs, Inc. // SPDX-License-Identifier: Apache-2.0 +pub mod kms_stress; + use crate::rpc::client::GasPoolRpcClient; use clap::ValueEnum; use parking_lot::RwLock; diff --git a/src/bin/tool.rs b/src/bin/tool.rs index 720e0e5..dee73d3 100644 --- a/src/bin/tool.rs +++ b/src/bin/tool.rs @@ -4,6 +4,7 @@ use clap::*; use std::path::PathBuf; use sui_config::Config; +use sui_gas_station::benchmarks::kms_stress::run_kms_stress_test; use sui_gas_station::benchmarks::BenchmarkMode; use sui_gas_station::config::{GasPoolStorageConfig, GasStationConfig, TxSignerConfig}; use sui_gas_station::rpc::client::GasPoolRpcClient; @@ -37,6 +38,11 @@ pub enum ToolCommand { #[arg(long, help = "Benchmark mode.", default_value = "reserve-only")] benchmark_mode: BenchmarkMode, }, + #[clap(name = "stress-kms")] + StressKMS { + #[arg(long, help = "Full URL to the KMS signer")] + kms_url: String, + }, /// Generate a sample config file and put it in the specified path. #[clap(name = "generate-sample-config")] GenerateSampleConfig { @@ -88,6 +94,9 @@ 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::GenerateSampleConfig { config_path, with_sidecar_signer, diff --git a/src/sui_client.rs b/src/sui_client.rs index 994acef..1aab955 100644 --- a/src/sui_client.rs +++ b/src/sui_client.rs @@ -233,15 +233,13 @@ impl SuiClient { .read_api() .get_object_with_options(obj_ref.0, SuiObjectDataOptions::default()) .await; - match response { - Ok(SuiObjectResponse { - data: Some(data), .. - }) => { - if data.version == obj_ref.1 { - break; - } + if let Ok(SuiObjectResponse { + data: Some(data), .. + }) = response + { + if data.version == obj_ref.1 { + break; } - _ => (), } tokio::time::sleep(Duration::from_millis(200)).await; }