fix(batcher): send correct responses to each client (#244)

Co-authored-by: jotabulacios <[email protected]>
Co-authored-by: Mauro Toscano <[email protected]>

Author: 3 people

Makefile

@@ -48,7 +48,7 @@ anvil_start:
 
 anvil_start_with_block_time:
 	@echo "Starting Anvil..."
-	anvil --load-state contracts/scripts/anvil/state/alignedlayer-deployed-anvil-state.json --block-time 2
+	anvil --load-state contracts/scripts/anvil/state/alignedlayer-deployed-anvil-state.json --block-time 3
 
 # TODO: Allow enviroment variables / different configuration files
 aggregator_start:
@@ -154,13 +154,13 @@ batcher_send_sp1_task:
 		--vm_program test_files/sp1/sp1_fibonacci-elf \
 		--proof_generator_addr 0x66f9664f97F2b50F62D13eA064982f936dE76657
 
-batcher_send_sp1_burst_5:
+batcher_send_sp1_burst:
 	@echo "Sending SP1 fibonacci task to Batcher..."
 	@cd batcher/client/ && cargo run --release -- \
 		--proving_system SP1 \
 		--proof test_files/sp1/sp1_fibonacci.proof \
 		--vm_program test_files/sp1/sp1_fibonacci-elf \
-		--repetitions 5 \
+		--repetitions 15 \
 		--proof_generator_addr 0x66f9664f97F2b50F62D13eA064982f936dE76657
 
 batcher_send_infinite_sp1:
@@ -176,6 +176,16 @@ batcher_send_plonk_bn254_task: batcher/client/target/release/batcher-client
 		--vk test_files/plonk_bn254/plonk.vk \
 		--proof_generator_addr 0x66f9664f97F2b50F62D13eA064982f936dE76657
 
+batcher_send_plonk_bn254_burst: batcher/client/target/release/batcher-client
+	@echo "Sending Groth16Bn254 1!=0 task to Batcher..."
+	@cd batcher/client/ && cargo run --release -- \
+		--proving_system GnarkPlonkBn254 \
+		--proof test_files/plonk_bn254/plonk.proof \
+		--public_input test_files/plonk_bn254/plonk_pub_input.pub \
+		--vk test_files/plonk_bn254/plonk.vk \
+		--proof_generator_addr 0x66f9664f97F2b50F62D13eA064982f936dE76657 \
+		--repetitions 15
+
 batcher_send_plonk_bls12_381_task: batcher/client/target/release/batcher-client
 	@echo "Sending Groth16 BLS12-381 1!=0 task to Batcher..."
 	@cd batcher/client/ && cargo run --release -- \
@@ -185,6 +195,16 @@ batcher_send_plonk_bls12_381_task: batcher/client/target/release/batcher-client
 		--vk test_files/plonk_bls12_381/plonk.vk \
 		--proof_generator_addr 0x66f9664f97F2b50F62D13eA064982f936dE76657
 
+batcher_send_plonk_bls12_381_burst: batcher/client/target/release/batcher-client
+	@echo "Sending Groth16 BLS12-381 1!=0 task to Batcher..."
+	@cd batcher/client/ && cargo run --release -- \
+		--proving_system GnarkPlonkBls12_381 \
+		--proof test_files/plonk_bls12_381/plonk.proof \
+		--public_input test_files/plonk_bls12_381/plonk_pub_input.pub \
+		--vk test_files/plonk_bls12_381/plonk.vk \
+		--proof_generator_addr 0x66f9664f97F2b50F62D13eA064982f936dE76657 \
+		--repetitions 15
+
 
 batcher_send_groth16_bn254_task: batcher/client/target/release/batcher-client
 	@echo "Sending Groth16Bn254 1!=0 task to Batcher..."
@@ -195,14 +215,14 @@ batcher_send_groth16_bn254_task: batcher/client/target/release/batcher-client
 		--vk test_files/groth16/ineq_1_groth16.vk \
 		--proof_generator_addr 0x66f9664f97F2b50F62D13eA064982f936dE76657
 
-batcher_send_groth16_burst_5: batcher/client/target/release/batcher-client
+batcher_send_groth16_burst: batcher/client/target/release/batcher-client
 	@echo "Sending Groth16Bn254 1!=0 task to Batcher..."
 	@cd batcher/client/ && cargo run --release -- \
 		--proving_system Groth16Bn254 \
 		--proof test_files/groth16/ineq_1_groth16.proof \
 		--public_input test_files/groth16/ineq_1_groth16.pub \
 		--vk test_files/groth16/ineq_1_groth16.vk \
-		--repetitions 5 \
+		--repetitions 15 \
 		--proof_generator_addr 0x66f9664f97F2b50F62D13eA064982f936dE76657
 
 batcher_send_infinite_groth16: ./batcher/client/target/release/batcher-client ## Send a different Groth16 BN254 proof using the task sender every 3 seconds

batcher/client/src/main.rs

@@ -1,14 +1,20 @@
-use std::path::PathBuf;
+use std::{path::PathBuf, sync::Arc, time::Duration};
 
 use alloy_primitives::Address;
 use env_logger::Env;
-use futures_util::{future, SinkExt, StreamExt, TryStreamExt};
-use log::{info};
-use tokio_tungstenite::connect_async;
+use futures_util::{
+    future,
+    stream::{SplitSink, SplitStream},
+    SinkExt, StreamExt, TryStreamExt,
+};
+use log::info;
+use tokio::{net::TcpStream, sync::Mutex};
+use tokio_tungstenite::{connect_async, MaybeTlsStream, WebSocketStream};
 
-use batcher::types::{parse_proving_system, VerificationData};
+use batcher::types::{parse_proving_system, BatchInclusionData, VerificationData};
 
 use clap::Parser;
+use tungstenite::Message;
 
 #[derive(Parser, Debug)]
 #[command(version, about, long_about = None)]
@@ -48,15 +54,14 @@ struct Args {
         long = "repetitions",
         default_value = "1"
     )]
-    repetitions: u32,
+    repetitions: usize,
 
     #[arg(
         name = "Proof generator address",
         long = "proof_generator_addr",
         default_value = "0xf39Fd6e51aad88F6F4ce6aB8827279cffFb92266"
     )] // defaults to anvil address 1
     proof_generator_addr: String,
-
 }
 
 #[tokio::main]
@@ -99,7 +104,8 @@ async fn main() {
         info!("No VM program code file provided, continuing without VM program code...");
     }
 
-    let proof_generator_addr: Address = Address::parse_checksummed(&args.proof_generator_addr, None).unwrap();
+    let proof_generator_addr: Address =
+        Address::parse_checksummed(&args.proof_generator_addr, None).unwrap();
 
     let verification_data = VerificationData {
         proving_system,
@@ -112,23 +118,41 @@ async fn main() {
 
     let json_data = serde_json::to_string(&verification_data).expect("Failed to serialize task");
     for _ in 0..args.repetitions {
+        // NOTE(marian): This sleep is only for ease of testing interactions between client and batcher,
+        // it can be removed.
+        std::thread::sleep(Duration::from_millis(500));
         ws_write
             .send(tungstenite::Message::Text(json_data.to_string()))
             .await
             .unwrap();
+        info!("Message sent...")
     }
 
+    let num_responses = Arc::new(Mutex::new(0));
+    let ws_write = Arc::new(Mutex::new(ws_write));
+
+    receive(ws_read, ws_write, args.repetitions, num_responses).await;
+}
+
+async fn receive(
+    ws_read: SplitStream<WebSocketStream<MaybeTlsStream<TcpStream>>>,
+    ws_write: Arc<Mutex<SplitSink<WebSocketStream<MaybeTlsStream<TcpStream>>, Message>>>,
+    total_messages: usize,
+    num_responses: Arc<Mutex<usize>>,
+) {
     ws_read
-        .try_filter(|msg| future::ready(msg.is_text()))
-        .for_each(|msg| async move {
-            let data = msg.unwrap().into_text().unwrap();
-            info!("Batch merkle root received: {}", data);
+        .try_filter(|msg| future::ready(msg.is_text() || msg.is_binary()))
+        .for_each(|msg| async {
+            let mut num_responses_lock = num_responses.lock().await;
+            *num_responses_lock += 1;
+            let data = msg.unwrap().into_data();
+            let deserialized_data: BatchInclusionData = serde_json::from_slice(&data).unwrap();
+            info!("Batcher response received: {}", deserialized_data);
+
+            if *num_responses_lock == total_messages {
+                info!("All messages responded. Closing connection...");
+                ws_write.lock().await.close().await.unwrap();
+            }
         })
         .await;
-
-    info!("Closing connection...");
-    ws_write
-        .close()
-        .await
-        .expect("Failed to close WebSocket connection");
 }

batcher/src/eth/mod.rs

@@ -41,7 +41,7 @@ pub async fn get_contract(
 }
 
 pub async fn create_new_task(
-    service_manager: AlignedLayerServiceManager,
+    service_manager: &AlignedLayerServiceManager,
     batch_merkle_root: [u8; 32],
     batch_data_pointer: String,
 ) -> Result<TransactionReceipt, anyhow::Error> {