diff --git a/crates/spammer/src/cli.rs b/crates/spammer/src/cli.rs index e55fa2a..635f53f 100644 --- a/crates/spammer/src/cli.rs +++ b/crates/spammer/src/cli.rs @@ -21,6 +21,7 @@ use clap::Args; use crate::{ accounts::PartitionMode, config::{Erc20FnWeights, GuzzlerFnWeights, TxTypeMix}, + rate_limiter::MAX_TPS, Config, }; @@ -83,7 +84,7 @@ pub struct SpammerArgs { #[clap(short = 'n', long, default_value_t = defaults::NUM_TXS, global = true)] pub num_txs: u64, /// Number of transactions to send per second (all generators combined) - #[clap(short = 'r', long, default_value_t = defaults::RATE, global = true, value_parser = clap::value_parser!(u64).range(1..))] + #[clap(short = 'r', long, default_value_t = defaults::RATE, global = true, value_parser = clap::value_parser!(u64).range(1..=MAX_TPS))] pub rate: u64, /// Maximum time in seconds to send transactions (applies to all generators) (0 for no limit) #[clap(short = 't', long, default_value_t = defaults::TIME, global = true)] diff --git a/crates/spammer/src/config.rs b/crates/spammer/src/config.rs index 1cab34b..13f1af1 100644 --- a/crates/spammer/src/config.rs +++ b/crates/spammer/src/config.rs @@ -16,7 +16,7 @@ use std::path::PathBuf; -use crate::accounts::PartitionMode; +use crate::{accounts::PartitionMode, rate_limiter::MAX_TPS}; use color_eyre::eyre::{self, Result}; use serde::{Deserialize, Serialize}; use std::str::FromStr; @@ -356,6 +356,12 @@ impl Config { if self.num_generators == 0 { eyre::bail!("num_generators must be greater than 0"); } + if self.max_rate > MAX_TPS { + eyre::bail!( + "--rate ({}) exceeds the maximum supported rate ({MAX_TPS})", + self.max_rate + ); + } if !self.max_num_accounts.is_multiple_of(self.num_generators) { eyre::bail!( "Expected max_num_accounts ({}) to be divisible by num_generators ({})", @@ -461,6 +467,28 @@ mod tests { } } + #[test] + fn config_accepts_max_supported_rate() { + Config { + max_rate: MAX_TPS, + ..default_config() + } + .validate() + .expect("maximum representable rate should be supported"); + } + + #[test] + fn config_rejects_rate_above_supported_max() { + let err = Config { + max_rate: MAX_TPS + 1, + ..default_config() + } + .validate() + .expect_err("oversized rate must be rejected"); + assert!(err.to_string().contains("--rate")); + assert!(err.to_string().contains(&MAX_TPS.to_string())); + } + #[test] fn tx_type_mix_parses_full_spec() { let mix = TxTypeMix::from_str("transfer=70,erc20=20,guzzler=10").expect("should parse"); diff --git a/crates/spammer/src/main.rs b/crates/spammer/src/main.rs index 64bc3ae..452311e 100644 --- a/crates/spammer/src/main.rs +++ b/crates/spammer/src/main.rs @@ -476,6 +476,20 @@ mod tests { assert_eq!(cli.args.rate, 42); } + #[test] + fn cli_rejects_rate_above_supported_max() { + let max_rate = 1_000_000_000_u64.to_string(); + Cli::try_parse_from(["spammer", "--rate", max_rate.as_str(), "ws"]) + .expect("maximum representable rate should be accepted"); + + let oversized_rate = 1_000_000_001_u64.to_string(); + let err = match Cli::try_parse_from(["spammer", "--rate", oversized_rate.as_str(), "ws"]) { + Ok(_) => panic!("rate above the supported maximum must be rejected"), + Err(err) => err, + }; + assert_eq!(err.kind(), ErrorKind::ValueValidation); + } + #[test] fn cli_parses_fire_and_forget_flag() { // Default: fire_and_forget is false (backpressure is the default) diff --git a/crates/spammer/src/rate_limiter.rs b/crates/spammer/src/rate_limiter.rs index 26e428e..e72ae9a 100644 --- a/crates/spammer/src/rate_limiter.rs +++ b/crates/spammer/src/rate_limiter.rs @@ -17,8 +17,13 @@ use std::num::NonZeroU32; use std::sync::atomic::{AtomicU64, Ordering}; +use color_eyre::eyre::{self, Result, WrapErr}; use governor::{Jitter, Quota}; +/// Highest rate `governor` can represent without rounding its per-token +/// replenish interval down to zero nanoseconds. +pub(crate) const MAX_TPS: u64 = 1_000_000_000; + /// Token-bucket rate limiter for transaction sending. /// /// Spaces sends evenly across each second using `governor` instead of @@ -32,22 +37,28 @@ pub(crate) struct RateLimiter { } impl RateLimiter { - pub fn new(tps: u64, max_num_txs: u64, num_senders: usize) -> Self { - let tps_u32 = u32::try_from(tps).expect("TPS must fit in u32"); - let tps_nz = NonZeroU32::new(tps_u32).expect("TPS must be > 0"); + pub fn new(tps: u64, max_num_txs: u64, num_senders: usize) -> Result { + if tps > MAX_TPS { + eyre::bail!("rate {tps} exceeds the maximum supported rate {MAX_TPS}"); + } + let tps_u32 = u32::try_from(tps) + .wrap_err_with(|| format!("rate {tps} exceeds the maximum supported rate {MAX_TPS}"))?; + let tps_nz = + NonZeroU32::new(tps_u32).ok_or_else(|| eyre::eyre!("rate must be greater than 0"))?; let burst = (tps / num_senders.max(1) as u64).max(1); - let burst_nz = NonZeroU32::new(u32::try_from(burst).expect("burst must fit in u32")) - .expect("burst must be > 0"); + let burst_u32 = u32::try_from(burst).wrap_err("rate limiter burst must fit in u32")?; + let burst_nz = NonZeroU32::new(burst_u32) + .ok_or_else(|| eyre::eyre!("rate limiter burst must be greater than 0"))?; let quota = Quota::per_second(tps_nz).allow_burst(burst_nz); let limiter = governor::RateLimiter::direct(quota); // Uniformly random jitter up to half the interval let jitter = Jitter::up_to(quota.replenish_interval() / 2); - Self { + Ok(Self { limiter, jitter, max_num_txs, total_counter: AtomicU64::new(0), - } + }) } /// Wait until the rate limiter permits the next send. @@ -62,3 +73,23 @@ impl RateLimiter { prev < self.max_num_txs } } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn new_accepts_max_supported_rate() { + assert!(RateLimiter::new(MAX_TPS, 1, 1).is_ok()); + } + + #[test] + fn new_rejects_rate_above_supported_max() { + let err = RateLimiter::new(MAX_TPS + 1, 1, 1) + .err() + .expect("oversized rate must be rejected"); + assert!(err + .to_string() + .contains("exceeds the maximum supported rate")); + } +} diff --git a/crates/spammer/src/spammer.rs b/crates/spammer/src/spammer.rs index d9f636d..c9d427e 100644 --- a/crates/spammer/src/spammer.rs +++ b/crates/spammer/src/spammer.rs @@ -267,7 +267,7 @@ impl Spammer { config.max_rate, config.max_num_txs, config.num_generators, - )); + )?); // Create transaction generators and senders let (tx_generators, tx_senders, tx_ack_receivers) = if config.fire_and_forget { @@ -689,7 +689,7 @@ impl Spammer { config.max_rate, config.max_num_txs, num_generators, - )); + )?); let mut tx_generators = Vec::new(); let mut tx_senders = Vec::new();