diff --git a/Cargo.lock b/Cargo.lock index 985ad877..d6e93744 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2195,6 +2195,7 @@ dependencies = [ "thiserror 2.0.17", "tokio", "tokio-util", + "toml", "tower 0.5.2", "tracing", "tracing-subscriber 0.3.20", diff --git a/crates/cli/src/commands/campaign.rs b/crates/cli/src/commands/campaign.rs index f00fbd08..00392f8c 100644 --- a/crates/cli/src/commands/campaign.rs +++ b/crates/cli/src/commands/campaign.rs @@ -51,15 +51,6 @@ pub struct CampaignCliArgs { )] pub pending_timeout: u64, - /// The number of accounts to generate for each agent (`from_pool` in scenario files) - #[arg( - short, - long, - visible_aliases = ["na", "accounts"], - default_value_t = 10 - )] - pub accounts_per_agent: u64, - /// Max number of txs to send in a single json-rpc batch request. #[arg( long = "rpc-batch-size", @@ -159,29 +150,33 @@ pub async fn run_campaign( .or_else(|| campaign.spam.seed.map(|s| s.to_string())) .unwrap_or(load_seedfile()?); - // Setup phase. Skip builtin scenarios since they do their own setup at spam time. - let provider = args.eth_json_rpc_args.new_rpc_provider()?; + // Setup phase: run setup for each (stage, mix) with the same derived seed that spam will use. + // This ensures setup creates accounts matching what spam expects. + // Skip builtin scenarios since they do their own setup at spam time. if !args.skip_setup { - for scenario_label in campaign.setup_scenarios() { - let scenario = match parse_builtin_reference(&scenario_label) { - Some(builtin) => SpamScenario::Builtin( - builtin - .to_builtin_scenario( - &provider, - &create_spam_cli_args(None, &args, CampaignMode::Tps, 1, 1), - /* TODO: KLUDGE: - - I don't think a `BuiltinScenarioCli` *needs* `rate` or `duration` -- that's for the spammer. - - we should use a different interface for `to_builtin_scenario` (replace `SpamCliArgs`) - */ - ) - .await?, - ), - None => SpamScenario::Testfile(scenario_label.to_owned()), - }; - let mut setup_args = args.eth_json_rpc_args.clone(); - setup_args.seed = Some(base_seed.clone()); - let setup_cmd = SetupCommandArgs::new(scenario, setup_args)?; - commands::setup(db, setup_cmd).await?; + for stage in &stages { + let stage_seed = bump_seed(&base_seed, &stage.name); + for (mix_idx, mix) in stage.mix.iter().enumerate() { + if mix.rate == 0 { + continue; + } + // Skip builtins - they do their own setup during spam + if parse_builtin_reference(&mix.scenario).is_some() { + continue; + } + + let scenario_seed = bump_seed(&stage_seed, &mix_idx.to_string()); + let scenario = SpamScenario::Testfile(mix.scenario.clone()); + + let mut setup_args = args.eth_json_rpc_args.clone(); + setup_args.seed = Some(scenario_seed); + // Ensure accounts_per_agent uses campaign default (10) if not explicitly set + if setup_args.accounts_per_agent.is_none() { + setup_args.accounts_per_agent = Some(10); + } + let setup_cmd = SetupCommandArgs::new(scenario, setup_args)?; + commands::setup(db, setup_cmd).await?; + } } } @@ -316,6 +311,8 @@ fn create_spam_cli_args( spam_mode: CampaignMode, spam_rate: u64, spam_duration: u64, + skip_setup: bool, + redeploy: bool, ) -> SpamCliArgs { SpamCliArgs { eth_json_rpc_args: ScenarioSendTxsCliArgs { @@ -337,13 +334,12 @@ fn create_spam_cli_args( duration: spam_duration, pending_timeout: args.pending_timeout, run_forever: false, - accounts_per_agent: args.accounts_per_agent, }, ignore_receipts: args.ignore_receipts, optimistic_nonces: args.optimistic_nonces, gen_report: false, - redeploy: args.redeploy, - skip_setup: true, + redeploy, + skip_setup, rpc_batch_size: args.rpc_batch_size, spam_timeout: args.spam_timeout, } @@ -383,12 +379,21 @@ async fn execute_stage( args.eth_json_rpc_args.seed = Some(scenario_seed.clone()); debug!("mix {mix_idx} seed: {}", scenario_seed); + // Check if this is a builtin scenario to determine skip_setup/redeploy behavior: + // - Builtins: respect campaign's flags (they do their own setup during spam) + // - Toml scenarios: always skip setup (ran in Phase 1), redeploy not applicable + let is_builtin = parse_builtin_reference(&mix.scenario).is_some(); + let skip_setup = if is_builtin { args.skip_setup } else { true }; + let redeploy = if is_builtin { args.redeploy } else { false }; + let spam_cli_args = create_spam_cli_args( Some(mix.scenario.clone()), &args, campaign.spam.mode, mix.rate, stage.duration, + skip_setup, + redeploy, ); let spam_scenario = if let Some(builtin_cli) = parse_builtin_reference(&mix.scenario) { diff --git a/crates/cli/src/commands/common.rs b/crates/cli/src/commands/common.rs index fa4230b3..aee5a234 100644 --- a/crates/cli/src/commands/common.rs +++ b/crates/cli/src/commands/common.rs @@ -128,9 +128,23 @@ Requires --auth-rpc-url and --jwt-secret to be set.", value_parser = parse_value, )] pub gas_price: Option, + + /// The number of accounts to generate for each agent (`from_pool` in scenario files). + /// Defaults to 1 for standalone setup, 10 for spam and campaign. + #[arg( + short = 'a', + long, + visible_aliases = ["na", "accounts"], + )] + pub accounts_per_agent: Option, } impl SendTxsCliArgsInner { + /// Returns the accounts_per_agent value, or the provided default if not set. + pub fn accounts_per_agent_or(&self, default: u64) -> u64 { + self.accounts_per_agent.unwrap_or(default) + } + pub fn new_rpc_provider(&self) -> Result, ArgsError> { info!("connecting to {}", self.rpc_url); Ok(DynProvider::new( @@ -301,15 +315,6 @@ Requires --priv-key to be set for each 'from' address in the given testfile.", visible_aliases = ["indefinite", "indefinitely", "infinite"] )] pub run_forever: bool, - - /// The number of accounts to generate for each agent (`from_pool` in scenario files) - #[arg( - short, - long, - visible_aliases = ["na", "accounts"], - default_value_t = 10 - )] - pub accounts_per_agent: u64, } #[derive(Copy, Debug, Clone, clap::ValueEnum)] diff --git a/crates/cli/src/commands/setup.rs b/crates/cli/src/commands/setup.rs index 7c1bd94e..6c52e32d 100644 --- a/crates/cli/src/commands/setup.rs +++ b/crates/cli/src/commands/setup.rs @@ -36,8 +36,10 @@ pub async fn setup( env, bundle_type, override_senders, + accounts_per_agent, .. } = args.eth_json_rpc_args.clone(); + let accounts_per_agent = accounts_per_agent.unwrap_or(1) as usize; let engine_params = args.engine_params().await?; let rpc_client = args.eth_json_rpc_args.new_rpc_provider()?; let user_signers_with_defaults = args.eth_json_rpc_args.user_signers_with_defaults(); @@ -86,7 +88,7 @@ pub async fn setup( continue; } - let agent = SignerStore::new(1, &args.seed, from_pool); + let agent = SignerStore::new(accounts_per_agent, &args.seed, from_pool); agents.add_agent(from_pool, agent); } } diff --git a/crates/cli/src/commands/spam.rs b/crates/cli/src/commands/spam.rs index 5553e97d..902cc225 100644 --- a/crates/cli/src/commands/spam.rs +++ b/crates/cli/src/commands/spam.rs @@ -219,7 +219,6 @@ impl SpamCommandArgs { duration, pending_timeout, run_forever, - accounts_per_agent, } = self.spam_args.spam_args.clone(); let SendTxsCliArgsInner { min_balance, @@ -227,8 +226,10 @@ impl SpamCommandArgs { bundle_type, env, override_senders, + accounts_per_agent, .. } = self.spam_args.eth_json_rpc_args.rpc_args.clone(); + let accounts_per_agent = accounts_per_agent.unwrap_or(10); let mut testconfig = self.testconfig().await?; let spam_len = testconfig.spam.as_ref().map(|s| s.len()).unwrap_or(0); diff --git a/crates/cli/src/default_scenarios/builtin.rs b/crates/cli/src/default_scenarios/builtin.rs index df3bfced..d4e342f5 100644 --- a/crates/cli/src/default_scenarios/builtin.rs +++ b/crates/cli/src/default_scenarios/builtin.rs @@ -98,7 +98,10 @@ impl BuiltinScenarioCli { let mut agents = AgentStore::new(); agents.init( &["spammers"], - spam_args.spam_args.accounts_per_agent as usize, + spam_args + .eth_json_rpc_args + .rpc_args + .accounts_per_agent_or(10) as usize, &seed, ); let spammers = agents diff --git a/crates/core/Cargo.toml b/crates/core/Cargo.toml index 34b52d6c..26e24c4e 100644 --- a/crates/core/Cargo.toml +++ b/crates/core/Cargo.toml @@ -32,6 +32,7 @@ thiserror = { workspace = true } [dev-dependencies] contender_testfile = { workspace = true } +toml = { workspace = true } [package.metadata.cargo-udeps.ignore] development = ["contender_testfile"] diff --git a/crates/core/src/generator/error.rs b/crates/core/src/generator/error.rs index 3bb57363..821ab9e6 100644 --- a/crates/core/src/generator/error.rs +++ b/crates/core/src/generator/error.rs @@ -37,6 +37,9 @@ pub enum GeneratorError { #[error("fuzz invalid")] FuzzInvalid, + #[error("for_all_accounts requires from_pool to be set")] + ForAllAccountsRequiresFromPool, + #[error("must specify from or from_pool in scenario config")] InvalidSender, diff --git a/crates/core/src/generator/function_def.rs b/crates/core/src/generator/function_def.rs index 003228f0..fdb5391d 100644 --- a/crates/core/src/generator/function_def.rs +++ b/crates/core/src/generator/function_def.rs @@ -32,6 +32,10 @@ pub struct FunctionCallDefinition { pub blob_data: Option, /// Optional setCode data; tx type must be set to EIP7702 by spammer pub authorization_address: Option, + /// If true and `from_pool` is set, run this setup transaction for all accounts in the pool. + /// Defaults to false (only runs for the first account). + #[serde(default)] + pub for_all_accounts: bool, } /// User-facing definition of a function call to be executed. @@ -55,6 +59,7 @@ impl FunctionCallDefinition { gas_limit: None, blob_data: None, authorization_address: None, + for_all_accounts: false, } } @@ -103,6 +108,10 @@ impl FunctionCallDefinition { self.authorization_address = Some(auth_addr.as_ref().to_owned()); self } + pub fn with_for_all_accounts(mut self, for_all_accounts: bool) -> Self { + self.for_all_accounts = for_all_accounts; + self + } pub fn sidecar_data(&self) -> Result, GeneratorError> { let sidecar_data = if let Some(data) = self.blob_data.as_ref() { @@ -150,3 +159,56 @@ pub struct FuzzParam { /// Maximum value fuzzer will use. pub max: Option, } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn for_all_accounts_defaults_to_false() { + let toml = r#" + to = "0x1234567890123456789012345678901234567890" + from_pool = "test_pool" + signature = "test()" + "#; + let def: FunctionCallDefinition = toml::from_str(toml).unwrap(); + assert!(!def.for_all_accounts); + } + + #[test] + fn for_all_accounts_parses_true() { + let toml = r#" + to = "0x1234567890123456789012345678901234567890" + from_pool = "test_pool" + signature = "test()" + for_all_accounts = true + "#; + let def: FunctionCallDefinition = toml::from_str(toml).unwrap(); + assert!(def.for_all_accounts); + } + + #[test] + fn for_all_accounts_parses_false() { + let toml = r#" + to = "0x1234567890123456789012345678901234567890" + from_pool = "test_pool" + signature = "test()" + for_all_accounts = false + "#; + let def: FunctionCallDefinition = toml::from_str(toml).unwrap(); + assert!(!def.for_all_accounts); + } + + #[test] + fn with_for_all_accounts_builder() { + let def = FunctionCallDefinition::new("0x1234") + .with_from_pool("test_pool") + .with_for_all_accounts(true); + assert!(def.for_all_accounts); + + let def = FunctionCallDefinition::new("0x1234") + .with_from_pool("test_pool") + .with_for_all_accounts(false); + assert!(!def.for_all_accounts); + } +} diff --git a/crates/core/src/generator/trait.rs b/crates/core/src/generator/trait.rs index 1ee2e2cd..f26bc47e 100644 --- a/crates/core/src/generator/trait.rs +++ b/crates/core/src/generator/trait.rs @@ -393,33 +393,60 @@ where self.get_genesis_hash(), )?; - // setup tx with template values - let mut tx = NamedTxRequest::new( - templater.template_function_call( - &self.make_strict_call(step, 0)?, // 'from' address injected here - &placeholder_map, - )?, - None, - step.kind.to_owned(), - ); - - // assign a unique nonce to each tx (tracker per sender) - // - we need to block on the future to ensure it's correct before sending the tx) - let from = tx.tx.from.expect("from address"); - if let std::collections::hash_map::Entry::Vacant(e) = next_nonce.entry(from) { - let nonce = self.get_rpc_provider().get_transaction_count(from).await?; - e.insert(nonce); - } - let nonce = next_nonce.get_mut(&from).expect("nonce"); - tx.tx.nonce = Some(*nonce); - *nonce += 1; + // Determine which account indices to use for this setup step + let account_indices: Vec = if step.for_all_accounts { + if let Some(ref from_pool) = step.from_pool { + // Get all account indices for this pool + let agents = self.get_agent_store(); + if let Some(agent) = agents.get_agent(from_pool) { + (0..agent.signers.len()).collect() + } else { + return Err(crate::Error::Generator( + GeneratorError::from_pool_not_found(from_pool), + )); + } + } else { + // for_all_accounts requires from_pool to be set + return Err(crate::Error::Generator( + GeneratorError::ForAllAccountsRequiresFromPool, + )); + } + } else { + // Default behavior: only use first account (idx 0) + vec![0] + }; - // spawn and store handle (will await all txs later) - let handle = on_setup_step(tx.to_owned())?; - if let Some(handle) = handle { - handles.push(handle); + // Generate a setup transaction for each account index + for account_idx in account_indices { + // setup tx with template values + let mut tx = NamedTxRequest::new( + templater.template_function_call( + &self.make_strict_call(step, account_idx)?, // 'from' address injected here + &placeholder_map, + )?, + None, + step.kind.to_owned(), + ); + + // assign a unique nonce to each tx (tracker per sender) + // - we need to block on the future to ensure it's correct before sending the tx) + let from = tx.tx.from.expect("from address"); + if let std::collections::hash_map::Entry::Vacant(e) = next_nonce.entry(from) + { + let nonce = self.get_rpc_provider().get_transaction_count(from).await?; + e.insert(nonce); + } + let nonce = next_nonce.get_mut(&from).expect("nonce"); + tx.tx.nonce = Some(*nonce); + *nonce += 1; + + // spawn and store handle (will await all txs later) + let handle = on_setup_step(tx.to_owned())?; + if let Some(handle) = handle { + handles.push(handle); + } + txs.push(tx.into()); } - txs.push(tx.into()); } for handle in handles {