Merge branch 'main' into unicheat

crates/cli/src/commands/types.rs → crates/cli/src/commands/contender_subcommand.rs

@@ -47,13 +47,14 @@ Requires --priv-key to be set for each 'from' address in the given testfile.",
         )]
         duration: Option<usize>,
 
-        /// The seed to use for generating spam transactions. If not provided, one is generated.
+        /// The seed to use for generating spam transactions & accounts.
         #[arg(
             short,
             long,
-            long_help = "The seed to use for generating spam transactions"
+            long_help = "The seed to use for generating spam transactions",
+            default_value = "0xffffffffffffffffffffffffffffffff13131313131313131313131313131313"
         )]
-        seed: Option<String>,
+        seed: String,
 
         /// The private keys to use for blockwise spamming.
         /// Required if `txs_per_block` is set.
@@ -109,6 +110,23 @@ May be specified multiple times."
             default_value = "1.0"
         )]
         min_balance: String,
+
+        /// The seed used to generate pool accounts.
+        #[arg(
+            short,
+            long,
+            long_help = "The seed used to generate pool accounts.",
+            default_value = "0xffffffffffffffffffffffffffffffff13131313131313131313131313131313"
+        )]
+        seed: String,
+
+        /// The number of signers to generate for each pool.
+        #[arg(
+            short,
+            long = "signers-per-pool",
+            long_help = "The number of signers to generate for each pool."
+        )]
+        num_signers_per_pool: Option<usize>,
     },
 
     #[command(

crates/cli/src/commands/mod.rs

@@ -1,16 +1,16 @@
+mod contender_subcommand;
 mod report;
 mod run;
 mod setup;
 mod spam;
-mod types;
 
 use clap::Parser;
 
+pub use contender_subcommand::ContenderSubcommand;
 pub use report::report;
 pub use run::run;
 pub use setup::setup;
 pub use spam::{spam, SpamCommandArgs};
-pub use types::ContenderSubcommand;
 
 #[derive(Parser, Debug)]
 pub struct ContenderCli {

crates/cli/src/commands/setup.rs

@@ -2,12 +2,17 @@ use alloy::{
     network::AnyNetwork, primitives::utils::parse_ether, providers::ProviderBuilder,
     signers::local::PrivateKeySigner, transports::http::reqwest::Url,
 };
-use contender_core::{generator::RandSeed, test_scenario::TestScenario};
+use contender_core::{
+    agent_controller::{AgentStore, SignerStore},
+    generator::RandSeed,
+    test_scenario::TestScenario,
+};
 use contender_testfile::TestConfig;
 use std::str::FromStr;
 
 use crate::util::{
-    check_private_keys_fns, find_insufficient_balance_addrs, get_signers_with_defaults,
+    check_private_keys_fns, find_insufficient_balance_addrs, fund_accounts, get_setup_pools,
+    get_signers_with_defaults,
 };
 
 pub async fn setup(
@@ -16,11 +21,14 @@ pub async fn setup(
     rpc_url: impl AsRef<str>,
     private_keys: Option<Vec<String>>,
     min_balance: String,
+    seed: RandSeed,
+    signers_per_period: usize,
 ) -> Result<(), Box<dyn std::error::Error>> {
     let url = Url::parse(rpc_url.as_ref()).expect("Invalid RPC URL");
     let rpc_client = ProviderBuilder::new()
         .network::<AnyNetwork>()
         .on_http(url.to_owned());
+    let eth_client = ProviderBuilder::new().on_http(url.to_owned());
     let testconfig: TestConfig = TestConfig::from_file(testfile.as_ref())?;
     let min_balance = parse_ether(&min_balance)?;
 
@@ -30,10 +38,10 @@ pub async fn setup(
         .iter()
         .map(|key| PrivateKeySigner::from_str(key).expect("invalid private key"))
         .collect::<Vec<PrivateKeySigner>>();
-    let signers = get_signers_with_defaults(private_keys);
+    let default_signers = get_signers_with_defaults(private_keys);
     check_private_keys_fns(
         &testconfig.setup.to_owned().unwrap_or_default(),
-        signers.as_slice(),
+        &default_signers,
     );
     let broke_accounts = find_insufficient_balance_addrs(
         &user_signers.iter().map(|s| s.address()).collect::<Vec<_>>(),
@@ -42,17 +50,41 @@ pub async fn setup(
     )
     .await?;
     if !broke_accounts.is_empty() {
-        panic!("Some accounts do not have sufficient balance");
+        panic!("Insufficient balance in provided user account(s)");
+    }
+
+    let mut agents = AgentStore::new();
+    let from_pools = get_setup_pools(&testconfig);
+    let mut all_signers = vec![];
+    all_signers.extend_from_slice(&user_signers);
+
+    for from_pool in &from_pools {
+        if agents.has_agent(from_pool) {
+            continue;
+        }
+
+        let agent = SignerStore::new_random(signers_per_period, &seed, from_pool);
+        all_signers.extend_from_slice(&agent.signers);
+        agents.add_agent(from_pool, agent);
     }
 
+    fund_accounts(
+        &rpc_client,
+        &eth_client,
+        min_balance,
+        &all_signers,
+        &default_signers[0],
+    )
+    .await?;
+
     let mut scenario = TestScenario::new(
         testconfig.to_owned(),
         db.clone().into(),
         url,
         None,
-        RandSeed::new(),
-        &signers,
-        Default::default(),
+        seed,
+        &default_signers,
+        agents,
     )
     .await?;
 

crates/cli/src/commands/spam.rs

@@ -14,7 +14,7 @@ use contender_core::{
 use contender_testfile::TestConfig;
 
 use crate::util::{
-    check_private_keys, fund_accounts, get_from_pools, get_signers_with_defaults,
+    check_private_keys, fund_accounts, get_signers_with_defaults, get_spam_pools,
     spam_callback_default, SpamCallbackType,
 };
 
@@ -26,7 +26,7 @@ pub struct SpamCommandArgs {
     pub txs_per_block: Option<usize>,
     pub txs_per_second: Option<usize>,
     pub duration: Option<usize>,
-    pub seed: Option<String>,
+    pub seed: String,
     pub private_keys: Option<Vec<String>>,
     pub disable_reports: bool,
     pub min_balance: String,
@@ -37,10 +37,7 @@ pub async fn spam(
     args: SpamCommandArgs,
 ) -> Result<(), Box<dyn std::error::Error>> {
     let testconfig = TestConfig::from_file(&args.testfile)?;
-    let rand_seed = args
-        .seed
-        .map(|s| RandSeed::seed_from_str(s.as_ref()))
-        .unwrap_or_default();
+    let rand_seed = RandSeed::seed_from_str(&args.seed);
     let url = Url::parse(&args.rpc_url).expect("Invalid RPC URL");
     let rpc_client = ProviderBuilder::new()
         .network::<AnyNetwork>()
@@ -57,7 +54,7 @@ pub async fn spam(
         .expect("No spam function calls found in testfile");
 
     // distill all from_pool arguments from the spam requests
-    let from_pools = get_from_pools(&testconfig);
+    let from_pools = get_spam_pools(&testconfig);
 
     let mut agents = AgentStore::new();
     let signers_per_period = args

crates/cli/src/main.rs

@@ -5,7 +5,7 @@ mod util;
 use std::sync::LazyLock;
 
 use commands::{ContenderCli, ContenderSubcommand, SpamCommandArgs};
-use contender_core::db::DbOps;
+use contender_core::{db::DbOps, generator::RandSeed};
 use contender_sqlite::SqliteDb;
 
 static DB: LazyLock<SqliteDb> = std::sync::LazyLock::new(|| {
@@ -16,14 +16,28 @@ static DB: LazyLock<SqliteDb> = std::sync::LazyLock::new(|| {
 async fn main() -> Result<(), Box<dyn std::error::Error>> {
     let args = ContenderCli::parse_args();
     let _ = DB.create_tables(); // ignore error; tables already exist
+    let db = DB.clone();
 
     match args.command {
         ContenderSubcommand::Setup {
             testfile,
             rpc_url,
             private_keys,
             min_balance,
-        } => commands::setup(&DB.clone(), testfile, rpc_url, private_keys, min_balance).await?,
+            seed,
+            num_signers_per_pool,
+        } => {
+            commands::setup(
+                &db,
+                testfile,
+                rpc_url,
+                private_keys,
+                min_balance,
+                RandSeed::seed_from_str(&seed),
+                num_signers_per_pool.unwrap_or(1),
+            )
+            .await?
+        }
 
         ContenderSubcommand::Spam {
             testfile,
@@ -38,7 +52,7 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
             min_balance,
         } => {
             commands::spam(
-                &DB.clone(),
+                &db,
                 SpamCommandArgs {
                     testfile,
                     rpc_url,
@@ -55,9 +69,7 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
             .await?
         }
 
-        ContenderSubcommand::Report { id, out_file } => {
-            commands::report(&DB.clone(), id, out_file)?
-        }
+        ContenderSubcommand::Report { id, out_file } => commands::report(&db, id, out_file)?,
 
         ContenderSubcommand::Run {
             scenario,
@@ -68,7 +80,7 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
             txs_per_duration,
         } => {
             commands::run(
-                &DB.clone(),
+                &db,
                 scenario,
                 rpc_url,
                 private_key,

crates/cli/src/util.rs

@@ -33,7 +33,17 @@ pub const DEFAULT_PRV_KEYS: [&str; 10] = [
     "0x2a871d0798f97d79848a013d4936a73bf4cc922c825d33c1cf7073dff6d409c6",
 ];
 
-pub fn get_from_pools(testconfig: &TestConfig) -> Vec<String> {
+pub fn get_setup_pools(testconfig: &TestConfig) -> Vec<String> {
+    testconfig
+        .setup
+        .to_owned()
+        .unwrap_or_default()
+        .into_iter()
+        .filter_map(|s| s.from_pool)
+        .collect()
+}
+
+pub fn get_spam_pools(testconfig: &TestConfig) -> Vec<String> {
     let mut from_pools = vec![];
     let spam = testconfig
         .spam

crates/core/src/generator/mod.rs

@@ -203,27 +203,41 @@ where
             }
             PlanType::Setup(on_setup_step) => {
                 let setup_steps = conf.get_setup_steps()?;
-                for step in setup_steps.iter() {
-                    // lookup placeholders in DB & update map before templating
-                    templater.find_fncall_placeholders(step, db, &mut placeholder_map)?;
-
-                    // setup tx with template values
-                    let tx = NamedTxRequest::new(
-                        templater.template_function_call(
-                            &self.make_strict_call(step, 0)?,
-                            &placeholder_map,
-                        )?,
-                        None,
-                        step.kind.to_owned(),
-                    );
 
-                    let handle = on_setup_step(tx.to_owned())?;
-                    if let Some(handle) = handle {
-                        handle.await.map_err(|e| {
-                            ContenderError::with_err(e, "join error; callback crashed")
-                        })?;
+                let agents = self.get_agent_store();
+                let num_accts = agents
+                    .all_agents()
+                    .next()
+                    .map(|(_, store)| store.signers.len())
+                    .unwrap_or(1);
+
+                for i in 0..(num_accts) {
+                    for step in setup_steps.iter() {
+                        if i > 0 && step.from_pool.is_none() {
+                            // only loop on from_pool steps; single-account steps can't be repeated
+                            continue;
+                        }
+                        // lookup placeholders in DB & update map before templating
+                        templater.find_fncall_placeholders(step, db, &mut placeholder_map)?;
+
+                        // setup tx with template values
+                        let tx = NamedTxRequest::new(
+                            templater.template_function_call(
+                                &self.make_strict_call(step, i)?, // 'from' address injected here
+                                &placeholder_map,
+                            )?,
+                            None,
+                            step.kind.to_owned(),
+                        );
+
+                        let handle = on_setup_step(tx.to_owned())?;
+                        if let Some(handle) = handle {
+                            handle.await.map_err(|e| {
+                                ContenderError::with_err(e, "join error; callback crashed")
+                            })?;
+                        }
+                        txs.push(tx.into());
                     }
-                    txs.push(tx.into());
                 }
             }
             PlanType::Spam(num_txs, on_spam_setup) => {
@@ -281,7 +295,7 @@ where
                     for step in spam_steps.iter() {
                         // converts a FunctionCallDefinition to a NamedTxRequest (filling in fuzzable args),
                         // returns a callback handle and the processed tx request
-                        let process_tx = |req| {
+                        let prepare_tx = |req| {
                             let args = get_fuzzed_args(req, &canonical_fuzz_map, i);
                             let fuzz_tx_value = get_fuzzed_tx_value(req, &canonical_fuzz_map, i);
                             let mut req = req.to_owned();
@@ -304,7 +318,7 @@ where
 
                         match step {
                             SpamRequest::Tx(req) => {
-                                let (handle, tx) = process_tx(req)?;
+                                let (handle, tx) = prepare_tx(req)?;
                                 if let Some(handle) = handle {
                                     handle.await.map_err(|e| {
                                         ContenderError::with_err(e, "error from callback")
@@ -315,7 +329,7 @@ where
                             SpamRequest::Bundle(req) => {
                                 let mut bundle_txs = vec![];
                                 for tx in req.txs.iter() {
-                                    let (handle, txr) = process_tx(tx)?;
+                                    let (handle, txr) = prepare_tx(tx)?;
                                     if let Some(handle) = handle {
                                         handle.await.map_err(|e| {
                                             ContenderError::with_err(e, "error from callback")