Remove erc20 setup from shooter setup

keep-starknet-strange · May 16, 2024 · b356969 · b356969
1 parent fff7e6c
commit b356969
Show file tree

Hide file tree

Showing 3 changed files with 65 additions and 46 deletions.
diff --git a/src/actions/mod.rs b/src/actions/mod.rs
@@ -20,10 +20,7 @@ pub async fn shoot(config: GatlingConfig) -> color_eyre::Result<()> {
     let total_txs = config.run.num_erc20_transfers + config.run.num_erc721_mints;
 
     let mut shooter_setup = GatlingSetup::from_config(config).await?;
-    let transfer_shooter = TransferShooter::setup(&mut shooter_setup).await?;
-    shooter_setup
-        .setup_accounts(transfer_shooter.erc20_address)
-        .await?;
+    shooter_setup.setup_accounts().await?;
 
     let mut global_report = GlobalReport {
         users: shooter_setup.config().run.concurrency,
@@ -35,6 +32,7 @@ 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);

diff --git a/src/actions/setup.rs b/src/actions/setup.rs
@@ -87,7 +87,7 @@ impl GatlingSetup {
     }
 
     /// Setup the simulation.
-    pub async fn setup_accounts(&mut self, erc20_address: FieldElement) -> Result<()> {
+    pub async fn setup_accounts(&mut self) -> Result<()> {
         let account_contract = self.config.setup.account_contract.clone();
 
         let account_class_hash = self.declare_contract(&account_contract).await?;
@@ -102,7 +102,6 @@ impl GatlingSetup {
                 account_class_hash,
                 self.config.run.concurrency as usize,
                 execution_encoding,
-                erc20_address,
             )
             .await?;
 
@@ -111,37 +110,15 @@ impl GatlingSetup {
         Ok(())
     }
 
-    async fn transfer(
-        &mut self,
+    pub async fn transfer(
+        &self,
         contract_address: FieldElement,
         account: StarknetAccount,
         recipient: FieldElement,
         amount: FieldElement,
         nonce: FieldElement,
     ) -> Result<FieldElement> {
-        let from_address = account.address();
-
-        debug!(
-            "Transferring {amount} of {contract_address:#064x} from address {from_address:#064x} to address {recipient:#064x} with nonce={}",
-            nonce,
-        );
-
-        let (amount_low, amount_high) = (amount, felt!("0"));
-
-        let call = Call {
-            to: contract_address,
-            selector: selector!("transfer"),
-            calldata: vec![recipient, amount_low, amount_high],
-        };
-
-        let result = account
-            .execute(vec![call])
-            .max_fee(MAX_FEE)
-            .nonce(nonce)
-            .send()
-            .await?;
-
-        Ok(result.transaction_hash)
+        transfer(account, nonce, amount, contract_address, recipient).await
     }
 
     /// Create accounts.
@@ -161,7 +138,6 @@ impl GatlingSetup {
         class_hash: FieldElement,
         num_accounts: usize,
         execution_encoding: ExecutionEncoding,
-        erc20_address: FieldElement,
     ) -> Result<Vec<StarknetAccount>> {
         info!("Creating {} accounts", num_accounts);
 
@@ -212,18 +188,6 @@ impl GatlingSetup {
                 }
             }
 
-            info!("Funding account {i} at address {address:#064x}");
-            let tx_hash = self
-                .transfer(
-                    erc20_address,
-                    self.account.clone(),
-                    address,
-                    felt!("0xFFFFF"),
-                    nonce,
-                )
-                .await?;
-            nonce += FieldElement::ONE;
-            wait_for_tx(&self.starknet_rpc, tx_hash, CHECK_INTERVAL).await?;
             let tx_hash = self
                 .transfer(
                     fee_token_address,
@@ -370,6 +334,38 @@ impl GatlingSetup {
     }
 }
 
+pub async fn transfer(
+    account: StarknetAccount,
+    nonce: FieldElement,
+    amount: FieldElement,
+    contract_address: FieldElement,
+    recipient: FieldElement,
+) -> color_eyre::Result<FieldElement> {
+    let from_address = account.address();
+
+    debug!(
+        "Transferring {amount} of {contract_address:#064x} from address {from_address:#064x} to address {recipient:#064x} with nonce={}",
+        nonce,
+    );
+
+    let (amount_low, amount_high) = (amount, felt!("0"));
+
+    let call = Call {
+        to: contract_address,
+        selector: selector!("transfer"),
+        calldata: vec![recipient, amount_low, amount_high],
+    };
+
+    let result = account
+        .execute(vec![call])
+        .max_fee(MAX_FEE)
+        .nonce(nonce)
+        .send()
+        .await?;
+
+    Ok(result.transaction_hash)
+}
+
 /// Create a StarkNet RPC provider from a URL.
 /// # Arguments
 /// * `rpc` - The URL of the StarkNet RPC provider.

diff --git a/src/actions/shooters/transfer.rs b/src/actions/shooters/transfer.rs
@@ -7,9 +7,10 @@ use starknet::{
     macros::{felt, selector},
     providers::Provider,
 };
+use tokio::task::JoinSet;
 
 use crate::{
-    actions::setup::{GatlingSetup, StarknetAccount, CHECK_INTERVAL, MAX_FEE},
+    actions::setup::{self, GatlingSetup, StarknetAccount, CHECK_INTERVAL, MAX_FEE},
     config::GatlingConfig,
     utils::{compute_contract_address, wait_for_tx},
 };
@@ -33,7 +34,7 @@ impl Shooter for TransferShooter {
             .await?;
 
         let contract_factory = ContractFactory::new(class_hash, setup.deployer_account().clone());
-        let nonce = setup.deployer_account().get_nonce().await?;
+        let mut nonce = setup.deployer_account().get_nonce().await?;
 
         let name = selector!("TestToken");
         let symbol = selector!("TT");
@@ -80,6 +81,7 @@ impl Shooter for TransferShooter {
         );
 
         let result = deploy.nonce(nonce).max_fee(MAX_FEE).send().await?;
+        nonce += FieldElement::ONE;
         wait_for_tx(setup.rpc_client(), result.transaction_hash, CHECK_INTERVAL).await?;
 
         debug!(
@@ -89,6 +91,29 @@ impl Shooter for TransferShooter {
 
         info!("ERC20 contract deployed at address {:#064x}", address);
 
+        let mut joinset = JoinSet::new();
+
+        for account in setup.accounts() {
+            info!("Funding account at address {address:#064x}");
+
+            let tx_hash = setup::transfer(
+                setup.deployer_account().clone(),
+                nonce,
+                felt!("0xFFFFF"),
+                address,
+                account.address(),
+            )
+            .await?;
+
+            nonce += FieldElement::ONE;
+            let rpc_client = setup.rpc_client().clone();
+            joinset.spawn(async move { wait_for_tx(&rpc_client, tx_hash, CHECK_INTERVAL).await });
+        }
+
+        while let Some(result) = joinset.join_next().await {
+            result??;
+        }
+
         Ok(TransferShooter {
             erc20_address: address,
             account: setup.deployer_account().clone(),