Merge pull request #90 from flashbots/bugfix/too-many-agent-accounts

fix "too many accounts" bug

README.md

@@ -88,7 +88,8 @@ Run the spammer with a custom scenario (10 tx/sec for 3 seconds):
 contender spam ./scenarios/stress.toml $RPC_URL --tps 10 -d 3
 ```
 
-Setting `--tps` defines the number of "agent accounts" (generated EOAs used to send txs).
+Setting `--tps` defines the number of "agent accounts" (generated EOAs used to send txs). The number of accounts each agent has is determined by `txs_per_period / num_agents`, where `num_agents` is defined by the scenario. For example, if the `stress.toml` scenario has 4 agents (defined by `from_pool` declarations), passing `--tps` 10 will generate `10 / 4 = 2.5` accounts, rounded down.
+
 Pass a private key with `-p` to fund agent accounts from your account:
 
 ```bash

crates/cli/src/commands/contender_subcommand.rs

@@ -123,14 +123,6 @@ May be specified multiple times."
         /// The seed used to generate pool accounts.
         #[arg(short, long, long_help = "The seed used to generate pool accounts.")]
         seed: Option<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/setup.rs

@@ -22,7 +22,6 @@ pub async fn setup(
     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()
@@ -60,27 +59,39 @@ pub async fn setup(
         );
     }
 
-    // collect all signers
-    let mut all_signers = vec![];
-    all_signers.extend_from_slice(&user_signers);
-
     // load agents from setup and create pools
-    let mut from_pools = get_setup_pools(&testconfig);
-    from_pools.extend(get_create_pools(&testconfig));
+    let from_pool_declarations =
+        [get_setup_pools(&testconfig), get_create_pools(&testconfig)].concat();
+
+    // create agents for each from_pool declaration
     let mut agents = AgentStore::new();
-    for from_pool in &from_pools {
+    for from_pool in &from_pool_declarations {
         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);
+        let agent = SignerStore::new_random(1, &seed, from_pool);
         agents.add_agent(from_pool, agent);
     }
 
+    let all_signer_addrs = [
+        // don't include default accounts (`user_signers_with_defaults`) here because if you're using them, they should already be funded
+        user_signers
+            .iter()
+            .map(|signer| signer.address())
+            .collect::<Vec<_>>(),
+        agents
+            .all_agents()
+            .flat_map(|(_, agent)| agent.signers.iter().map(|signer| signer.address()))
+            .collect::<Vec<_>>(),
+    ]
+    .concat();
+
+    let admin_signer = &user_signers_with_defaults[0];
+
     fund_accounts(
-        &all_signers,
-        &user_signers_with_defaults[0],
+        &all_signer_addrs,
+        admin_signer,
         &rpc_client,
         &eth_client,
         min_balance,

crates/cli/src/commands/spam.rs

@@ -55,30 +55,42 @@ 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_spam_pools(&testconfig);
+    let from_pool_declarations = get_spam_pools(&testconfig);
 
     let mut agents = AgentStore::new();
     let signers_per_period = args
         .txs_per_block
         .unwrap_or(args.txs_per_second.unwrap_or(spam.len()));
 
-    let mut all_signers = vec![];
-    all_signers.extend_from_slice(&user_signers);
-
-    for from_pool in &from_pools {
+    for from_pool in &from_pool_declarations {
         if agents.has_agent(from_pool) {
             continue;
         }
 
-        let agent = SignerStore::new_random(signers_per_period, &rand_seed, from_pool);
-        all_signers.extend_from_slice(&agent.signers);
+        let agent = SignerStore::new_random(
+            signers_per_period / from_pool_declarations.len(),
+            &rand_seed,
+            from_pool,
+        );
         agents.add_agent(from_pool, agent);
     }
 
-    check_private_keys(&testconfig, &all_signers);
+    let all_signer_addrs = [
+        user_signers
+            .iter()
+            .map(|signer| signer.address())
+            .collect::<Vec<_>>(),
+        agents
+            .all_agents()
+            .flat_map(|(_, agent)| agent.signers.iter().map(|signer| signer.address()))
+            .collect::<Vec<_>>(),
+    ]
+    .concat();
+
+    check_private_keys(&testconfig, &user_signers);
 
     fund_accounts(
-        &all_signers,
+        &all_signer_addrs,
         &user_signers[0],
         &rpc_client,
         &eth_client,

crates/cli/src/main.rs

@@ -55,7 +55,6 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
             private_keys,
             min_balance,
             seed,
-            num_signers_per_pool,
         } => {
             let seed = seed.unwrap_or(stored_seed);
             commands::setup(
@@ -65,7 +64,6 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
                 private_keys,
                 min_balance,
                 RandSeed::seed_from_str(&seed),
-                num_signers_per_pool.unwrap_or(1),
             )
             .await?
         }

crates/cli/src/util.rs

@@ -153,18 +153,14 @@ async fn is_balance_sufficient(
 }
 
 pub async fn fund_accounts(
-    accounts: &[PrivateKeySigner],
+    recipient_addresses: &[Address],
     fund_with: &PrivateKeySigner,
     rpc_client: &AnyProvider,
     eth_client: &EthProvider,
     min_balance: U256,
 ) -> Result<(), Box<dyn std::error::Error>> {
-    let insufficient_balance_addrs = find_insufficient_balance_addrs(
-        &accounts.iter().map(|s| s.address()).collect::<Vec<_>>(),
-        min_balance,
-        rpc_client,
-    )
-    .await?;
+    let insufficient_balance_addrs =
+        find_insufficient_balance_addrs(recipient_addresses, min_balance, rpc_client).await?;
 
     let mut pending_fund_txs = vec![];
     let admin_nonce = rpc_client

crates/core/src/spammer/blockwise.rs

@@ -50,40 +50,113 @@ where
 
 #[cfg(test)]
 mod tests {
+    use alloy::{consensus::constants::ETH_TO_WEI, primitives::U256, providers::ProviderBuilder};
+
     use crate::{
-        agent_controller::AgentStore,
+        agent_controller::{AgentStore, SignerStore},
         db::MockDb,
         generator::util::test::spawn_anvil,
-        spammer::util::test::{get_test_signers, MockCallback},
+        spammer::util::test::{fund_account, get_test_signers, MockCallback},
         test_scenario::tests::MockConfig,
     };
+    use std::collections::HashSet;
     use std::sync::Arc;
 
     use super::*;
 
     #[tokio::test]
     async fn watches_blocks_and_spams_them() {
         let anvil = spawn_anvil();
+        let provider = ProviderBuilder::new().on_http(anvil.endpoint_url().to_owned());
         println!("anvil url: {}", anvil.endpoint_url());
         let seed = crate::generator::RandSeed::seed_from_str("444444444444");
+        let mut agents = AgentStore::new();
+        let txs_per_period = 10;
+        let periods = 3;
+        agents.add_agent(
+            "pool1",
+            SignerStore::new_random(txs_per_period / periods, &seed, "eeeeeeee"),
+        );
+        agents.add_agent(
+            "pool2",
+            SignerStore::new_random(txs_per_period / periods, &seed, "11111111"),
+        );
+
+        let user_signers = get_test_signers();
+        let mut nonce = provider
+            .get_transaction_count(user_signers[0].address())
+            .await
+            .unwrap();
+
+        for (_pool_name, agent) in agents.all_agents() {
+            for signer in &agent.signers {
+                let res = fund_account(
+                    &user_signers[0],
+                    signer.address(),
+                    U256::from(ETH_TO_WEI),
+                    &provider,
+                    Some(nonce),
+                )
+                .await
+                .unwrap();
+                println!("funded signer: {:?}", res);
+                provider.watch_pending_transaction(res).await.unwrap();
+                nonce += 1;
+            }
+        }
+
         let mut scenario = TestScenario::new(
             MockConfig,
             MockDb.into(),
             anvil.endpoint_url(),
             None,
             seed,
-            get_test_signers().as_slice(),
-            AgentStore::new(),
+            &user_signers,
+            agents,
         )
         .await
         .unwrap();
         let callback_handler = MockCallback;
         let spammer = BlockwiseSpammer {};
 
+        let start_block = provider.get_block_number().await.unwrap();
+
         let result = spammer
-            .spam_rpc(&mut scenario, 10, 3, None, Arc::new(callback_handler))
+            .spam_rpc(
+                &mut scenario,
+                txs_per_period,
+                periods,
+                None,
+                Arc::new(callback_handler),
+            )
             .await;
-        println!("{:?}", result);
         assert!(result.is_ok());
+
+        let mut unique_addresses = HashSet::new();
+        let mut n_block = start_block;
+        let current_block = provider.get_block_number().await.unwrap();
+
+        while n_block <= current_block {
+            let block = provider
+                .get_block(
+                    n_block.into(),
+                    alloy::rpc::types::BlockTransactionsKind::Full,
+                )
+                .await
+                .unwrap();
+            if let Some(block) = block {
+                for tx in block.transactions.into_transactions() {
+                    unique_addresses.insert(tx.from);
+                }
+            }
+            n_block += 1;
+        }
+
+        for addr in unique_addresses.iter() {
+            println!("unique address: {}", addr);
+        }
+
+        assert!(unique_addresses.len() >= (txs_per_period / periods));
+        assert!(unique_addresses.len() <= txs_per_period);
     }
 }

crates/core/src/spammer/util.rs

@@ -2,11 +2,18 @@
 pub mod test {
     use std::{collections::HashMap, str::FromStr, sync::Arc};
 
-    use alloy::{providers::PendingTransactionConfig, signers::local::PrivateKeySigner};
+    use alloy::{
+        network::{EthereumWallet, TransactionBuilder},
+        primitives::{Address, U256},
+        providers::PendingTransactionConfig,
+        providers::Provider,
+        rpc::types::TransactionRequest,
+        signers::local::PrivateKeySigner,
+    };
     use tokio::task::JoinHandle;
 
     use crate::{
-        generator::NamedTxRequest,
+        generator::{types::EthProvider, NamedTxRequest},
         spammer::{tx_actor::TxActorHandle, OnTxSent},
     };
 
@@ -34,4 +41,37 @@ pub mod test {
         .map(|s| PrivateKeySigner::from_str(s).unwrap())
         .collect::<Vec<PrivateKeySigner>>()
     }
+
+    pub async fn fund_account(
+        sender: &PrivateKeySigner,
+        recipient: Address,
+        amount: U256,
+        rpc_client: &EthProvider,
+        nonce: Option<u64>,
+    ) -> Result<PendingTransactionConfig, Box<dyn std::error::Error>> {
+        println!(
+            "funding account {} with user account {}",
+            recipient,
+            sender.address()
+        );
+
+        let gas_price = rpc_client.get_gas_price().await?;
+        let nonce = nonce.unwrap_or(rpc_client.get_transaction_count(sender.address()).await?);
+        let chain_id = rpc_client.get_chain_id().await?;
+        let tx_req = TransactionRequest {
+            from: Some(sender.address()),
+            to: Some(alloy::primitives::TxKind::Call(recipient)),
+            value: Some(amount),
+            gas: Some(21000),
+            gas_price: Some(gas_price + 4_200_000_000),
+            nonce: Some(nonce),
+            chain_id: Some(chain_id),
+            ..Default::default()
+        };
+        let eth_wallet = EthereumWallet::from(sender.to_owned());
+        let tx = tx_req.build(&eth_wallet).await?;
+        let res = rpc_client.send_tx_envelope(tx).await?;
+
+        Ok(res.into_inner())
+    }
 }