Remove chain groups (#4267)

## Motivation

Since we're using the full flow, that will take care of synchronizing
things for us, so we don't need the chain groups anymore. They also
complicate things a bit more than necessary.

## Proposal

Remove the chain groups

## Test Plan

Run locally, see everything working as expected

## Release Plan

- Nothing to do / These changes follow the usual release cycle.

Author: ndr-ds

Cargo.lock

@@ -154,7 +154,6 @@ mini-moka = "0.10.3"
 num-bigint = "0.4.3"
 num-format = "0.4.4"
 num-traits = "0.2.18"
-num_cpus = "1.17.0"
 octocrab = "0.42.1"
 oneshot = "0.1.6"
 pathdiff = "0.2.1"

Cargo.toml

@@ -81,22 +81,22 @@ pub struct Benchmark<Env: Environment> {
 impl<Env: Environment> Benchmark<Env> {
     #[expect(clippy::too_many_arguments)]
     pub async fn run_benchmark<C: ClientContext<Environment = Env> + 'static>(
-        num_chain_groups: usize,
+        num_chains: usize,
         transactions_per_block: usize,
         bps: usize,
-        chain_clients: Vec<Vec<ChainClient<Env>>>,
-        blocks_infos: Vec<Vec<Vec<Operation>>>,
+        chain_clients: Vec<ChainClient<Env>>,
+        blocks_infos: Vec<Vec<Operation>>,
         health_check_endpoints: Option<String>,
         runtime_in_seconds: Option<u64>,
-        delay_between_chain_groups_ms: Option<u64>,
+        delay_between_chains_ms: Option<u64>,
         chain_listener: ChainListener<C>,
         shutdown_notifier: &CancellationToken,
     ) -> Result<(), BenchmarkError> {
-        let bps_counts = (0..num_chain_groups)
+        let bps_counts = (0..num_chains)
             .map(|_| Arc::new(AtomicUsize::new(0)))
             .collect::<Vec<_>>();
         let notifier = Arc::new(Notify::new());
-        let barrier = Arc::new(Barrier::new(num_chain_groups + 1));
+        let barrier = Arc::new(Barrier::new(num_chains + 1));
 
         let chain_listener_handle = tokio::spawn(
             async move {
@@ -117,19 +117,19 @@ impl<Env: Environment> Benchmark<Env> {
         );
 
         let (runtime_control_task, runtime_control_sender) =
-            Self::runtime_control_task(shutdown_notifier, runtime_in_seconds, num_chain_groups);
+            Self::runtime_control_task(shutdown_notifier, runtime_in_seconds, num_chains);
 
-        let bps_initial_share = bps / num_chain_groups;
-        let mut bps_remainder = bps % num_chain_groups;
+        let bps_initial_share = bps / num_chains;
+        let mut bps_remainder = bps % num_chains;
         let mut join_set = task::JoinSet::<Result<(), BenchmarkError>>::new();
-        for (chain_group_index, (chain_group, chain_clients)) in blocks_infos
+        for (chain_idx, (block_info, chain_client)) in blocks_infos
             .into_iter()
             .zip(chain_clients.into_iter())
             .enumerate()
         {
             let shutdown_notifier_clone = shutdown_notifier.clone();
             let barrier_clone = barrier.clone();
-            let bps_count_clone = bps_counts[chain_group_index].clone();
+            let bps_count_clone = bps_counts[chain_idx].clone();
             let notifier_clone = notifier.clone();
             let runtime_control_sender_clone = runtime_control_sender.clone();
             let bps_share = if bps_remainder > 0 {
@@ -138,38 +138,44 @@ impl<Env: Environment> Benchmark<Env> {
             } else {
                 bps_initial_share
             };
+            let chain_id = chain_client.chain_id();
             join_set.spawn(
                 async move {
                     Box::pin(Self::run_benchmark_internal(
-                        chain_group_index,
+                        chain_idx,
+                        chain_id,
                         bps_share,
-                        chain_group,
-                        chain_clients,
+                        block_info,
+                        chain_client,
                         shutdown_notifier_clone,
                         bps_count_clone,
                         barrier_clone,
                         notifier_clone,
                         runtime_control_sender_clone,
-                        delay_between_chain_groups_ms,
+                        delay_between_chains_ms,
                     ))
                     .await?;
 
                     Ok(())
                 }
-                .instrument(
-                    tracing::info_span!("chain_group", chain_group_index = ?chain_group_index),
-                ),
+                .instrument(tracing::info_span!("chain_id", chain_id = ?chain_id)),
             );
         }
 
+        // Wait for tasks and fail immediately if any task returns an error or panics
+        while let Some(result) = join_set.join_next().await {
+            let inner_result = result?;
+            if let Err(e) = inner_result {
+                error!("Benchmark task failed: {}", e);
+                shutdown_notifier.cancel();
+                join_set.abort_all();
+                return Err(e);
+            }
+        }
+
         let metrics_watcher =
             Self::metrics_watcher(health_check_endpoints, shutdown_notifier).await?;
 
-        join_set
-            .join_all()
-            .await
-            .into_iter()
-            .collect::<Result<Vec<_>, _>>()?;
         info!("All benchmark tasks completed");
         bps_control_task.await?;
         if let Some(metrics_watcher) = metrics_watcher {
@@ -517,42 +523,39 @@ impl<Env: Environment> Benchmark<Env> {
 
     #[expect(clippy::too_many_arguments)]
     async fn run_benchmark_internal(
-        chain_group_index: usize,
+        chain_idx: usize,
+        chain_id: ChainId,
         bps: usize,
-        chain_group: Vec<Vec<Operation>>,
-        chain_clients: Vec<ChainClient<Env>>,
+        operations: Vec<Operation>,
+        chain_client: ChainClient<Env>,
         shutdown_notifier: CancellationToken,
         bps_count: Arc<AtomicUsize>,
         barrier: Arc<Barrier>,
         notifier: Arc<Notify>,
         runtime_control_sender: Option<mpsc::Sender<()>>,
-        delay_between_chain_groups_ms: Option<u64>,
+        delay_between_chains_ms: Option<u64>,
     ) -> Result<(), BenchmarkError> {
         barrier.wait().await;
-        if let Some(delay_between_chain_groups_ms) = delay_between_chain_groups_ms {
+        if let Some(delay_between_chains_ms) = delay_between_chains_ms {
             time::sleep(time::Duration::from_millis(
-                (chain_group_index as u64) * delay_between_chain_groups_ms,
+                (chain_idx as u64) * delay_between_chains_ms,
             ))
             .await;
         }
-        info!("Starting benchmark for chain group {:?}", chain_group_index);
+        info!("Starting benchmark for chain {:?}", chain_id);
 
         if let Some(runtime_control_sender) = runtime_control_sender {
             runtime_control_sender.send(()).await?;
         }
 
-        for (operations, chain_client) in chain_group
-            .into_iter()
-            .zip(chain_clients.into_iter())
-            .cycle()
-        {
+        loop {
             if shutdown_notifier.is_cancelled() {
                 info!("Shutdown signal received, stopping benchmark");
                 break;
             }
 
             chain_client
-                .execute_operations(operations, vec![])
+                .execute_operations(operations.clone(), vec![])
                 .await
                 .map_err(BenchmarkError::ChainClient)?
                 .expect("should execute block with operations");
@@ -588,60 +591,56 @@ impl<Env: Environment> Benchmark<Env> {
 
     /// Generates information related to one block per chain.
     pub fn make_benchmark_block_info(
-        benchmark_chains: Vec<Vec<(ChainId, AccountOwner)>>,
+        benchmark_chains: Vec<(ChainId, AccountOwner)>,
         transactions_per_block: usize,
         fungible_application_id: Option<ApplicationId>,
-    ) -> Vec<Vec<Vec<Operation>>> {
+    ) -> Vec<Vec<Operation>> {
         let mut blocks_infos = Vec::new();
-        for chains in benchmark_chains {
-            let mut infos = Vec::new();
-            let chains_len = chains.len();
+        for (i, (_, owner)) in benchmark_chains.iter().enumerate() {
             let amount = Amount::from(1);
-            for i in 0..chains_len {
-                let owner = chains[i].1;
-                let mut operations = Vec::new();
+            let mut operations = Vec::new();
+
+            let mut other_chains: Vec<_> = if benchmark_chains.len() == 1 {
+                // If there's only one chain, just have it send to itself.
+                benchmark_chains
+                    .iter()
+                    .map(|(chain_id, _)| *chain_id)
+                    .collect()
+            } else {
+                // If there's more than one chain, have it send to all other chains, and don't
+                // send to self.
+                benchmark_chains
+                    .iter()
+                    .enumerate()
+                    .filter(|(j, _)| i != *j)
+                    .map(|(_, (chain_id, _))| *chain_id)
+                    .collect()
+            };
 
-                let mut other_chains: Vec<_> = if chains_len == 1 {
-                    // If there's only one chain, just have it send to itself.
-                    chains.iter().map(|(chain_id, _)| *chain_id).collect()
-                } else {
-                    // If there's more than one chain, have it send to all other chains, and don't
-                    // send to self.
-                    chains
-                        .iter()
-                        .enumerate()
-                        .filter(|(j, _)| i != *j)
-                        .map(|(_, (chain_id, _))| *chain_id)
-                        .collect()
+            other_chains.shuffle(&mut thread_rng());
+            for recipient_chain_id in other_chains {
+                let operation = match fungible_application_id {
+                    Some(application_id) => Self::fungible_transfer(
+                        application_id,
+                        recipient_chain_id,
+                        *owner,
+                        *owner,
+                        amount,
+                    ),
+                    None => Operation::system(SystemOperation::Transfer {
+                        owner: AccountOwner::CHAIN,
+                        recipient: Recipient::chain(recipient_chain_id),
+                        amount,
+                    }),
                 };
-
-                other_chains.shuffle(&mut thread_rng());
-
-                for recipient_chain_id in other_chains {
-                    let operation = match fungible_application_id {
-                        Some(application_id) => Self::fungible_transfer(
-                            application_id,
-                            recipient_chain_id,
-                            owner,
-                            owner,
-                            amount,
-                        ),
-                        None => Operation::system(SystemOperation::Transfer {
-                            owner: AccountOwner::CHAIN,
-                            recipient: Recipient::chain(recipient_chain_id),
-                            amount,
-                        }),
-                    };
-                    operations.push(operation);
-                }
-                let operations = operations
-                    .into_iter()
-                    .cycle()
-                    .take(transactions_per_block)
-                    .collect();
-                infos.push(operations);
+                operations.push(operation);
             }
-            blocks_infos.push(infos);
+            let operations = operations
+                .into_iter()
+                .cycle()
+                .take(transactions_per_block)
+                .collect();
+            blocks_infos.push(operations);
         }
         blocks_infos
     }