[FEATURE]: use JoinSet to run transaction checks concurrently (#33)

src/actions/shoot.rs

@@ -7,12 +7,12 @@ use crate::utils::{
 use color_eyre::eyre::Context;
 use color_eyre::{eyre::eyre, Report as EyreReport, Result};
 
-use log::{debug, error, info, warn};
+use log::{debug, info, warn};
 use starknet::core::types::contract::SierraClass;
 
-use futures::stream::StreamExt;
 use std::collections::HashMap;
 use std::path::Path;
+use tokio::task::JoinSet;
 
 use crate::metrics::BenchmarkReport;
 
@@ -41,7 +41,7 @@ use url::Url;
 // Used to bypass validation
 pub static MAX_FEE: FieldElement = felt!("0xffffffff");
 pub static CHECK_INTERVAL: Duration = Duration::from_millis(500);
-pub static TX_VALIDATION_CHUNKS: usize = 10;
+pub static TX_VALIDATION_CONCURRENCY: usize = 10;
 
 type StarknetAccount = SingleOwnerAccount<Arc<JsonRpcClient<HttpTransport>>, LocalWallet>;
 
@@ -218,40 +218,42 @@ impl GatlingShooter {
         let mut accepted_txs = Vec::new();
         let mut errors = Vec::new();
 
-        // Verify transactions in parallel
-        let chunk_size = transactions.len() / TX_VALIDATION_CHUNKS;
-        let transactions_sets = transactions.chunks(chunk_size).map(|chunk| chunk.to_vec());
-
-        let fetches = futures::stream::iter(transactions_sets.map(|transactions_set| {
-            // Should we clone the rpc client or is it ok to share it between threads ?
-            let starknet_rpc = self.starknet_rpc.clone();
-            tokio::spawn(async move {
-                let mut results = Vec::new();
-                for tx in transactions_set {
-                    let res = wait_for_tx(&starknet_rpc, tx, CHECK_INTERVAL)
+        let mut set = JoinSet::new();
+
+        let mut transactions = transactions.into_iter();
+
+        for _ in 0..TX_VALIDATION_CONCURRENCY {
+            if let Some(transaction) = transactions.next() {
+                let starknet_rpc = Arc::clone(&self.starknet_rpc);
+                set.spawn(async move {
+                    wait_for_tx(&starknet_rpc, transaction, CHECK_INTERVAL)
+                        .await
+                        .map(|_| transaction)
+                        .map_err(|err| (err, transaction))
+                });
+            }
+        }
+
+        while let Some(res) = set.join_next().await {
+            if let Some(transaction) = transactions.next() {
+                let starknet_rpc = Arc::clone(&self.starknet_rpc);
+                set.spawn(async move {
+                    wait_for_tx(&starknet_rpc, transaction, CHECK_INTERVAL)
                         .await
-                        .map(|_| tx);
-                    debug!("Transaction {:#064x} result: {:?}", tx, res);
-                    results.push(res);
+                        .map(|_| transaction)
+                        .map_err(|err| (err, transaction))
+                });
+            }
+
+            match res.unwrap() {
+                Ok(transaction) => {
+                    accepted_txs.push(transaction);
+                    debug!("Transaction {:#064x} accepted", transaction)
                 }
-                results
-            })
-        }))
-        .buffer_unordered(TX_VALIDATION_CHUNKS) // Adjust the concurrency level to the number of connections
-        .collect::<Vec<_>>();
-
-        // fetches.await will resolve to a Vec<Result<Vec<Result<Transaction, Error>>, JoinError>>
-        for fetch_result in fetches.await {
-            match fetch_result {
-                Ok(results) => {
-                    for res in results {
-                        match res {
-                            Ok(tx) => accepted_txs.push(tx),
-                            Err(e) => errors.push(e),
-                        }
-                    }
+                Err((err, transaction)) => {
+                    errors.push(err);
+                    debug!("Transaction {:#064x} rejected", transaction)
                 }
-                Err(e) => error!("JoinError {:?}", e),
             }
         }