diff --git a/bench-tps/src/bench.rs b/bench-tps/src/bench.rs index 9c5240412..69312be87 100644 --- a/bench-tps/src/bench.rs +++ b/bench-tps/src/bench.rs @@ -17,7 +17,6 @@ use solana_sdk::transaction::Transaction; use std::cmp; use std::collections::VecDeque; use std::net::SocketAddr; -use std::process::exit; use std::sync::atomic::{AtomicBool, AtomicIsize, AtomicUsize, Ordering}; use std::sync::{Arc, RwLock}; use std::thread::sleep; @@ -28,6 +27,13 @@ use std::time::Instant; pub const MAX_SPENDS_PER_TX: u64 = 4; pub const NUM_LAMPORTS_PER_ACCOUNT: u64 = 128; +#[derive(Debug)] +pub enum BenchTpsError { + AirdropFailure, +} + +pub type Result = std::result::Result; + pub type SharedTransactions = Arc>>>; pub struct Config { @@ -447,7 +453,7 @@ pub fn airdrop_lamports( drone_addr: &SocketAddr, id: &Keypair, tx_count: u64, -) { +) -> Result<()> { let starting_balance = client.get_balance(&id.pubkey()).unwrap_or(0); metrics_submit_lamport_balance(starting_balance); println!("starting balance {}", starting_balance); @@ -496,9 +502,10 @@ pub fn airdrop_lamports( current_balance, starting_balance ); - exit(1); + return Err(BenchTpsError::AirdropFailure); } } + Ok(()) } fn compute_and_report_stats( @@ -596,7 +603,7 @@ pub fn generate_and_fund_keypairs( funding_pubkey: &Keypair, tx_count: usize, lamports_per_account: u64, -) -> (Vec, u64) { +) -> Result<(Vec, u64)> { info!("Creating {} keypairs...", tx_count * 2); let mut keypairs = generate_keypairs(funding_pubkey, tx_count as u64 * 2); @@ -614,7 +621,7 @@ pub fn generate_and_fund_keypairs( lamports_per_account - last_keypair_balance + fee_calculator.lamports_per_signature; let total = extra * (keypairs.len() as u64); if client.get_balance(&funding_pubkey.pubkey()).unwrap_or(0) < total { - airdrop_lamports(client, &drone_addr.unwrap(), funding_pubkey, total); + airdrop_lamports(client, &drone_addr.unwrap(), funding_pubkey, total)?; } info!("adding more lamports {}", extra); fund_keys( @@ -629,7 +636,7 @@ pub fn generate_and_fund_keypairs( // 'generate_keypairs' generates extra keys to be able to have size-aligned funding batches for fund_keys. keypairs.truncate(2 * tx_count); - (keypairs, last_keypair_balance) + Ok((keypairs, last_keypair_balance)) } #[cfg(test)] @@ -696,7 +703,8 @@ mod tests { &config.id, config.tx_count, lamports_per_account, - ); + ) + .unwrap(); let total = do_bench_tps(vec![client], config, keypairs, 0); assert!(total > 100); @@ -714,7 +722,7 @@ mod tests { config.duration = Duration::from_secs(5); let (keypairs, _keypair_balance) = - generate_and_fund_keypairs(&clients[0], None, &config.id, config.tx_count, 20); + generate_and_fund_keypairs(&clients[0], None, &config.id, config.tx_count, 20).unwrap(); do_bench_tps(clients, config, keypairs, 0); } @@ -728,7 +736,7 @@ mod tests { let lamports = 20; let (keypairs, _keypair_balance) = - generate_and_fund_keypairs(&client, None, &id, tx_count, lamports); + generate_and_fund_keypairs(&client, None, &id, tx_count, lamports).unwrap(); for kp in &keypairs { // TODO: This should be >= lamports, but fails at the moment diff --git a/bench-tps/src/main.rs b/bench-tps/src/main.rs index a50dd0ebb..a804db4c0 100644 --- a/bench-tps/src/main.rs +++ b/bench-tps/src/main.rs @@ -90,6 +90,10 @@ fn main() { tx_count, NUM_LAMPORTS_PER_ACCOUNT, ) + .unwrap_or_else(|e| { + eprintln!("Error could not fund keys: {:?}", e); + exit(1); + }) }; let config = Config {