Make adding shooters easier

README.md

@@ -84,8 +84,11 @@ The configuration is defined by the following spec
 - `run`
 
   - `concurrency`: How many transactions to do simultaneously
-  - `num_erc20_transfers`: Number of ERC20 `transfer` transactions
-  - `num_erc721_mints`: Number of ERC721 `mint` transactions
+  - `shooter`: A list of write shooter benchmarks to run
+
+    - `name`: The name of the shooter, must be either `transfer` or `mint`
+    - `shoot`: How many transactions to do
+
   - `read_benches`: A list of read benchmarks to run
 
     - `name`: The name to write on the output report

config/default.yaml

@@ -17,8 +17,11 @@ setup:
 
 run:
   concurrency: 5
-  num_erc20_transfers: 5000
-  num_erc721_mints: 5000
+  shooters:
+    - name: "transfer"
+      shoot: 5000
+    - name: "mint"
+      shoot: 5000
   read_benches:
     - name: "Get Events"
       num_requests: 100

config/katana.yaml

@@ -17,8 +17,11 @@ setup:
 
 run:
   concurrency: 5
-  num_erc20_transfers: 300
-  num_erc721_mints: 300
+  shooters:
+    - name: "transfer"
+      shoot: 300
+    - name: "mint"
+      shoot: 300
   read_benches:
     - name: "Get Events"
       num_requests: 100

config/sharingan.yaml

@@ -17,8 +17,11 @@ setup:
 
 run:
   concurrency: 5
-  num_erc20_transfers: 1200
-  num_erc721_mints: 1200
+  shooters:
+    - name: "transfer"
+      shoot: 1200
+    - name: "mint"
+      shoot: 1200
   read_benches:
     - name: "Get Events"
       num_requests: 100

config/v2.1.0.yaml

@@ -23,8 +23,11 @@ setup:
 
 run:
   concurrency: 5
-  num_erc20_transfers: 100
-  num_erc721_mints: 100
+  shooters:
+    - name: "transfer"
+      shoot: 100
+    - name: "mint"
+      shoot: 100
   read_benches:
     - name: "Get Events"
       num_requests: 100

src/actions/mod.rs

@@ -1,5 +1,6 @@
-use std::{fs::File, sync::Arc};
+use std::{fs::File, mem, sync::Arc};
 
+use color_eyre::eyre::bail;
 use log::info;
 
 use crate::{
@@ -16,8 +17,9 @@ mod goose;
 mod setup;
 mod shooters;
 
-pub async fn shoot(config: GatlingConfig) -> color_eyre::Result<()> {
-    let total_txs = config.run.num_erc20_transfers + config.run.num_erc721_mints;
+pub async fn shoot(mut config: GatlingConfig) -> color_eyre::Result<()> {
+    let shooters = mem::take(&mut config.run.shooters);
+    let total_txs: u64 = shooters.iter().map(|s| s.shoot).sum();
 
     let mut shooter_setup = GatlingSetup::from_config(config).await?;
     shooter_setup.setup_accounts().await?;
@@ -31,25 +33,25 @@ pub async fn shoot(config: GatlingConfig) -> color_eyre::Result<()> {
 
     let mut blocks = Option::<(u64, u64)>::None;
 
-    if shooter_setup.config().run.num_erc20_transfers != 0 {
-        let transfer_shooter = TransferShooter::setup(&mut shooter_setup).await?;
-        let report = make_report_over_shooter(transfer_shooter, &shooter_setup).await?;
-
-        global_report.benches.push(report.0);
-        blocks.get_or_insert((report.1, report.2)).1 = report.2;
-    } else {
-        log::info!("Skipping erc20 transfers")
-    }
-
-    if shooter_setup.config().run.num_erc721_mints != 0 {
-        let shooter = MintShooter::setup(&mut shooter_setup).await?;
+    for shooter in shooters {
+        if shooter.shoot == 0 {
+            log::info!("Skipping {} transfers", shooter.name);
+            continue;
+        }
 
-        let report = make_report_over_shooter(shooter, &shooter_setup).await?;
+        let (report, first_block, last_block) = match shooter.name.as_str() {
+            "transfer" => {
+                make_report_over_shooter::<TransferShooter>(&mut shooter_setup, shooter.shoot)
+                    .await?
+            }
+            "mint" => {
+                make_report_over_shooter::<MintShooter>(&mut shooter_setup, shooter.shoot).await?
+            }
+            name => bail!("Shooter `{name}` not found!"),
+        };
 
-        global_report.benches.push(report.0);
-        blocks.get_or_insert((report.1, report.2)).1 = report.2;
-    } else {
-        log::info!("Skipping erc721 mints")
+        global_report.benches.push(report);
+        blocks.get_or_insert((first_block, last_block)).1 = last_block;
     }
 
     let mut all_bench_report = BenchmarkReport::new("".into(), total_txs as usize);
@@ -80,10 +82,11 @@ pub async fn shoot(config: GatlingConfig) -> color_eyre::Result<()> {
 }
 
 async fn make_report_over_shooter<S: Shooter + Send + Sync + 'static>(
-    shooter: S,
-    setup: &GatlingSetup,
+    setup: &mut GatlingSetup,
+    amount: u64,
 ) -> color_eyre::Result<(BenchmarkReport, u64, u64)> {
-    let goose_config = S::get_goose_config(setup.config())?;
+    let shooter = S::setup(setup).await?;
+    let goose_config = S::get_goose_config(setup.config(), amount)?;
 
     let ShooterAttack {
         goose_metrics,

src/actions/shooters.rs

@@ -42,10 +42,11 @@ pub trait Shooter {
     where
         Self: Sized;
 
-    fn get_amount(config: &GatlingConfig) -> u64;
-
-    fn get_goose_config(config: &GatlingConfig) -> color_eyre::Result<GooseConfiguration> {
-        make_goose_config(config, Self::get_amount(config), Self::NAME)
+    fn get_goose_config(
+        config: &GatlingConfig,
+        amount: u64,
+    ) -> color_eyre::Result<GooseConfiguration> {
+        make_goose_config(config, amount, Self::NAME)
     }
 
     async fn goose_attack(

src/actions/shooters/mint.rs

@@ -13,7 +13,6 @@ use tokio::task::JoinSet;
 
 use crate::{
     actions::setup::{GatlingSetup, StarknetAccount, CHECK_INTERVAL, MAX_FEE},
-    config::GatlingConfig,
     generators::get_rng,
     utils::{compute_contract_address, wait_for_tx},
 };
@@ -61,10 +60,6 @@ impl Shooter for MintShooter {
         })
     }
 
-    fn get_amount(config: &GatlingConfig) -> u64 {
-        config.run.num_erc721_mints
-    }
-
     fn get_execution_data(&self, account: &StarknetAccount) -> Call {
         let recipient = account.address();
 

src/actions/shooters/transfer.rs

@@ -11,7 +11,6 @@ use tokio::task::JoinSet;
 
 use crate::{
     actions::setup::{self, GatlingSetup, StarknetAccount, CHECK_INTERVAL, MAX_FEE},
-    config::GatlingConfig,
     utils::{compute_contract_address, wait_for_tx},
 };
 
@@ -120,10 +119,6 @@ impl Shooter for TransferShooter {
         })
     }
 
-    fn get_amount(config: &GatlingConfig) -> u64 {
-        config.run.num_erc20_transfers
-    }
-
     fn get_execution_data(&self, _account: &StarknetAccount) -> Call {
         let (amount_low, amount_high) = (felt!("1"), felt!("0"));
 

src/config.rs

@@ -110,11 +110,16 @@ pub struct DeployerConfig {
 #[derive(Debug, Deserialize, Clone)]
 pub struct RunConfig {
     pub concurrency: u64,
-    pub num_erc20_transfers: u64,
-    pub num_erc721_mints: u64,
+    pub shooters: Vec<Shooters>,
     pub read_benches: Vec<ReadBenchConfig>,
 }
 
+#[derive(Debug, Deserialize, Clone)]
+pub struct Shooters {
+    pub name: String,
+    pub shoot: u64,
+}
+
 #[derive(Debug, Deserialize, Clone)]
 pub struct ReadBenchConfig {
     pub name: String,