Update

Author: AndreiEres

cumulus/zombienet/zombienet-sdk/tests/zombie_ci/statement_store_bench.rs

@@ -16,7 +16,7 @@ use zombienet_sdk::{
 };
 
 const GROUP_SIZE: u32 = 6;
-const PARTICIPANT_SIZE: u32 = GROUP_SIZE * 8333;
+const PARTICIPANT_SIZE: u32 = GROUP_SIZE * 8333; // Target ~50,000 total
 const MESSAGE_SIZE: usize = 5 * 1024; // 5KiB
 const MESSAGE_COUNT: usize = 2;
 const MAX_RETRIES: u32 = 100;
@@ -33,7 +33,7 @@ async fn statement_store_one_node_bench() -> Result<(), anyhow::Error> {
 
 	info!("Starting statement store benchmark with {} participants", PARTICIPANT_SIZE);
 
-	let target_node = "charlie";
+	let target_node = "alice";
 	let node = network.get_node(target_node)?;
 	let rpc_client = node.rpc().await?;
 	info!("Created single RPC client for target node: {}", target_node);
@@ -45,13 +45,54 @@ async fn statement_store_one_node_bench() -> Result<(), anyhow::Error> {
 
 	let handles: Vec<_> = participants
 		.into_iter()
-		.map(|mut participant| {
-			tokio::spawn(async move {
-				// Add staggered start delay to reduce initial network congestion
-				tokio::time::sleep(Duration::from_millis(participant.idx as u64 * 1)).await;
-				participant.run().await
-			})
-		})
+		.map(|mut p| tokio::spawn(async move { p.run().await }))
+		.collect();
+
+	let mut all_stats = Vec::new();
+	for handle in handles {
+		let stats = handle.await??;
+		all_stats.push(stats);
+	}
+
+	let aggregated_stats = ParticipantStats::aggregate(all_stats);
+	aggregated_stats.log_summary();
+
+	Ok(())
+}
+
+#[tokio::test(flavor = "multi_thread")]
+async fn statement_store_many_nodes_bench() -> Result<(), anyhow::Error> {
+	let _ = env_logger::try_init_from_env(
+		env_logger::Env::default().filter_or(env_logger::DEFAULT_FILTER_ENV, "info"),
+	);
+
+	let network = spawn_network().await?;
+
+	info!("Starting statement store benchmark with {} participants", PARTICIPANT_SIZE);
+
+	let collator_names = ["alice", "bob", "charlie", "dave", "eve", "ferdie"];
+	let mut rpc_clients = Vec::new();
+	for &name in &collator_names {
+		let node = network.get_node(name)?;
+		let rpc_client = node.rpc().await?;
+		rpc_clients.push(rpc_client);
+	}
+	info!("Created RPC clients for {} collator nodes", rpc_clients.len());
+
+	let mut participants = Vec::with_capacity(PARTICIPANT_SIZE as usize);
+	for i in 0..(PARTICIPANT_SIZE) as usize {
+		let client_idx = i % GROUP_SIZE as usize;
+		participants.push(Participant::new(i as u32, rpc_clients[client_idx].clone()));
+	}
+	info!(
+		"Participants distributed across nodes: {} participants per group of {} nodes",
+		PARTICIPANT_SIZE / GROUP_SIZE,
+		GROUP_SIZE
+	);
+
+	let handles: Vec<_> = participants
+		.into_iter()
+		.map(|mut participant| tokio::spawn(async move { participant.run().await }))
 		.collect();
 
 	let mut all_stats = Vec::new();
@@ -232,9 +273,9 @@ impl AggregatedStats {
 }
 
 struct Participant {
+	idx: u32,
 	keyring: sr25519::Pair,
 	session_key: sr25519::Pair,
-	idx: u32,
 	group_members: Vec<u32>,
 	session_keys: HashMap<u32, sr25519::Public>,
 	symmetric_keys: HashMap<u32, sr25519::Public>,
@@ -303,9 +344,7 @@ impl Participant {
 		if self.retry_count % 10 == 0 {
 			debug!(target: &self.log_target(), "[{}] Retry attempt {}", self.idx, self.retry_count);
 		}
-		let delay_ms =
-			std::cmp::min(RETRY_DELAY_MS * (1 << std::cmp::min(self.retry_count / 5, 4)), 5000);
-		tokio::time::sleep(tokio::time::Duration::from_millis(delay_ms)).await;
+		tokio::time::sleep(tokio::time::Duration::from_millis(RETRY_DELAY_MS)).await;
 
 		Ok(())
 	}
@@ -314,11 +353,6 @@ impl Participant {
 		tokio::time::sleep(tokio::time::Duration::from_millis(RECEIVE_DELAY_MS)).await;
 	}
 
-	async fn send_session_key(&mut self) -> Result<(), anyhow::Error> {
-		let statement = self.public_key_statement();
-		self.statement_submit(statement).await
-	}
-
 	async fn receive_statements_with_retry<T, F, R>(
 		&mut self,
 		mut pending: Vec<T>,
@@ -335,7 +369,7 @@ impl Participant {
 
 			for item in &pending {
 				match timeout(
-					Duration::from_millis(RECEIVE_DELAY_MS * 3), // Increase timeout 3x
+					Duration::from_millis(RETRY_DELAY_MS),
 					self.statement_broadcasts_statement(topic_generator(item)),
 				)
 				.await
@@ -351,11 +385,9 @@ impl Participant {
 					},
 					Ok(Ok(statements)) if statements.is_empty() => {
 						debug!(target: &self.log_target(), "[{}] No statements received for item {:?}", self.idx, item);
-						self.receive_sleep().await;
 					},
 					err => {
 						debug!(target: &self.log_target(), "[{}] Cannot receive statements for item {:?}, err: {:?}", self.idx, item, err);
-						self.receive_sleep().await;
 					},
 				}
 			}
@@ -372,6 +404,11 @@ impl Participant {
 		Ok(())
 	}
 
+	async fn send_session_key(&mut self) -> Result<(), anyhow::Error> {
+		let statement = self.public_key_statement();
+		self.statement_submit(statement).await
+	}
+
 	async fn receive_session_keys(&mut self) -> Result<(), anyhow::Error> {
 		let group_members = self.group_members.clone();
 
@@ -956,27 +993,30 @@ impl Participant {
 		self.receive_sleep().await;
 		self.receive_symmetric_keys().await?;
 
-		debug!(target: &self.log_target(), "[{}] Symmetric key acknowledgments", self.idx);
+		debug!(target: &self.log_target(), "[{}] Symmetric key acknowledgments exchange", self.idx);
 		self.send_symmetric_key_acknowledgments().await?;
 		self.receive_sleep().await;
 		self.receive_symmetric_key_acknowledgments().await?;
 
 		debug!(target: &self.log_target(), "[{}] Preparation finished", self.idx);
-		for round in 0..MESSAGE_COUNT {
-			debug!(target: &self.log_target(), "[{}] Req/res exchange round {}", self.idx, round + 1);
 
+		for round in 0..MESSAGE_COUNT {
+			debug!(target: &self.log_target(), "[{}] Requests exchange, round {}", self.idx, round + 1);
 			self.send_requests(round).await?;
 			self.receive_sleep().await;
 			self.receive_requests(round).await?;
 
+			debug!(target: &self.log_target(), "[{}] Request acknowledgments exchange, round {}", self.idx, round + 1);
 			self.send_request_acknowledgments().await?;
 			self.receive_sleep().await;
 			self.receive_request_acknowledgments().await?;
 
+			debug!(target: &self.log_target(), "[{}] Responses exchange, round {}", self.idx, round + 1);
 			self.send_responses(round).await?;
 			self.receive_sleep().await;
 			self.receive_responses(round).await?;
 
+			debug!(target: &self.log_target(), "[{}] Response acknowledgments exchange, round {}", self.idx, round + 1);
 			self.send_response_acknowledgments().await?;
 			self.receive_sleep().await;
 			self.receive_response_acknowledgments().await?;