Merge pull request #49 from flashbots/sign-timed-spamming

generalize spamming

crates/bundle_provider/src/bundle_provider.rs

@@ -14,7 +14,7 @@ pub struct EthSendBundleResponse {
 }
 
 impl BundleClient {
-    pub fn new(url: String) -> Self {
+    pub fn new(url: impl AsRef<str>) -> Self {
         let client = HttpClient::builder()
             .build(url)
             .expect("failed to connect to RPC provider");

crates/cli/src/main.rs

@@ -19,7 +19,9 @@ use contender_core::{
         types::{AnyProvider, EthProvider, FunctionCallDefinition, SpamRequest},
         RandSeed,
     },
-    spammer::{BlockwiseSpammer, LogCallback, NilCallback, TimedSpammer},
+    spammer::{
+        blockwise::BlockwiseSpammer, timed::TimedSpammer, LogCallback, NilCallback, Spammer,
+    },
     test_scenario::TestScenario,
 };
 use contender_sqlite::SqliteDb;
@@ -81,7 +83,8 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
                 RandSeed::new(),
                 &signers,
                 Default::default(),
-            );
+            )
+            .await?;
 
             scenario.deploy_contracts().await?;
             scenario.run_setup().await?;
@@ -141,7 +144,8 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
             }
 
             let mut agents = AgentStore::new();
-            let signers_per_block = txs_per_block.unwrap_or(spam.len()) / spam.len();
+            let signers_per_period =
+                txs_per_block.unwrap_or(txs_per_second.unwrap_or(spam.len())) / spam.len();
 
             let mut all_signers = vec![];
             all_signers.extend_from_slice(&user_signers);
@@ -151,7 +155,7 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
                     continue;
                 }
 
-                let agent = SignerStore::new_random(signers_per_block, &rand_seed, &from_pool);
+                let agent = SignerStore::new_random(signers_per_period, &rand_seed, &from_pool);
                 all_signers.extend_from_slice(&agent.signers);
                 agents.add_agent(from_pool, agent);
             }
@@ -214,16 +218,18 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
                 panic!("Must set either --txs-per-block (--tpb) or --txs-per-second (--tps)");
             }
 
+            let mut scenario = TestScenario::new(
+                testconfig,
+                DB.clone().into(),
+                url,
+                builder_url.map(|url| Url::parse(&url).expect("Invalid builder URL")),
+                rand_seed,
+                &user_signers,
+                agents,
+            )
+            .await?;
+
             if let Some(txs_per_block) = txs_per_block {
-                let scenario = TestScenario::new(
-                    testconfig,
-                    DB.clone().into(),
-                    url,
-                    builder_url.map(|url| Url::parse(&url).expect("Invalid builder URL")),
-                    rand_seed,
-                    &user_signers,
-                    agents,
-                );
                 println!("Blockwise spamming with {} txs per block", txs_per_block);
                 match spam_callback_default(!disable_reports, Arc::new(rpc_client).into()).await {
                     SpamCallbackType::Log(cback) => {
@@ -232,35 +238,45 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
                             .expect("Time went backwards")
                             .as_millis();
                         let run_id = DB.insert_run(timestamp as u64, txs_per_block * duration)?;
-                        let mut spammer = BlockwiseSpammer::new(scenario, cback).await;
+                        let spammer = BlockwiseSpammer::new::<SqliteDb>(cback);
                         spammer
-                            .spam_rpc(txs_per_block, duration, Some(run_id.into()))
+                            .spam_rpc(&mut scenario, txs_per_block, duration, Some(run_id))
                             .await?;
                         println!("Saved run. run_id = {}", run_id);
                     }
                     SpamCallbackType::Nil(cback) => {
-                        let mut spammer = BlockwiseSpammer::new(scenario, cback).await;
-                        spammer.spam_rpc(txs_per_block, duration, None).await?;
+                        let spammer = BlockwiseSpammer::new::<SqliteDb>(cback);
+                        spammer
+                            .spam_rpc(&mut scenario, txs_per_block, duration, None)
+                            .await?;
                     }
                 };
                 return Ok(());
             }
 
-            // NOTE: private keys are not currently used for timed spamming.
-            // Timed spamming only works with unlocked accounts, because it uses the `eth_sendTransaction` RPC method.
-            let scenario = TestScenario::new(
-                testconfig,
-                DB.clone().into(),
-                url,
-                None,
-                rand_seed,
-                &user_signers,
-                agents,
-            );
             let tps = txs_per_second.unwrap_or(10);
             println!("Timed spamming with {} txs per second", tps);
-            let spammer = TimedSpammer::new(scenario, NilCallback::new());
-            spammer.spam_rpc(tps, duration).await?;
+
+            let interval = std::time::Duration::from_secs(1);
+
+            match spam_callback_default(!disable_reports, Arc::new(rpc_client).into()).await {
+                SpamCallbackType::Log(cback) => {
+                    let timestamp = std::time::SystemTime::now()
+                        .duration_since(std::time::UNIX_EPOCH)
+                        .expect("Time went backwards")
+                        .as_millis();
+                    let run_id = DB.insert_run(timestamp as u64, tps * duration)?;
+                    let spammer = TimedSpammer::new::<SqliteDb>(cback, interval);
+                    spammer
+                        .spam_rpc(&mut scenario, tps, duration, Some(run_id))
+                        .await?;
+                    println!("Saved run. run_id = {}", run_id);
+                }
+                SpamCallbackType::Nil(cback) => {
+                    let spammer = TimedSpammer::new::<SqliteDb>(cback, interval);
+                    spammer.spam_rpc(&mut scenario, tps, duration, None).await?;
+                }
+            };
         }
         ContenderSubcommand::Report { id, out_file } => {
             let num_runs = DB.clone().num_runs()?;

crates/core/src/agent_controller.rs

@@ -24,6 +24,7 @@ pub struct SignerStore {
     pub signers: Vec<PrivateKeySigner>,
 }
 
+#[derive(Debug)]
 pub struct AgentStore {
     agents: HashMap<String, SignerStore>,
 }