feat: fetch logs instead of listening for events via ws

Author: MarcosNicolau

aggregation-mode/src/backend/fetcher.rs

@@ -1,22 +1,16 @@
-use std::{str::FromStr, sync::Arc};
+use std::str::FromStr;
 
 use super::{
     config::Config,
     queue::ProofsQueue,
-    s3::S3Client,
     types::{AlignedLayerServiceManager, AlignedLayerServiceManagerContract},
 };
-use crate::zk::{
-    backends::sp1::{vk_from_elf, SP1Proof},
-    Proof,
+use crate::{
+    backend::s3::get_aligned_batch_from_s3,
+    zk::{backends::sp1::SP1Proof, Proof},
 };
 use aligned_sdk::core::types::ProvingSystemId;
-use alloy::{
-    primitives::Address,
-    providers::{ProviderBuilder, WsConnect},
-};
-use futures_util::stream::StreamExt;
-use tokio::sync::Mutex;
+use alloy::{primitives::Address, providers::ProviderBuilder};
 use tracing::{error, info};
 
 /// This services is in charge of:
@@ -26,79 +20,77 @@ use tracing::{error, info};
 /// 4. Push the proofs to the queue
 pub struct ProofsFetcher {
     aligned_service_manager: AlignedLayerServiceManagerContract,
-    s3: S3Client,
-    queue: Arc<Mutex<ProofsQueue>>,
 }
 
 impl ProofsFetcher {
-    pub async fn new(config: &Config, queue: Arc<Mutex<ProofsQueue>>) -> Self {
-        let ws = WsConnect::new(&config.eth_ws_url);
-        let provider = ProviderBuilder::new()
-            .on_ws(ws)
-            .await
-            .expect("Successful connection");
-
+    pub fn new(config: &Config) -> Self {
+        let rpc_url = config.eth_rpc_url.parse().expect("correct url");
+        let provider = ProviderBuilder::new().on_http(rpc_url);
         let aligned_service_manager = AlignedLayerServiceManager::new(
             Address::from_str(&config.aligned_service_manager_address)
                 .expect("Address to be correct"),
             provider,
         );
 
-        let s3 = S3Client::new(config.bucket_name.clone(), None).await;
-
         Self {
             aligned_service_manager,
-            s3,
-            queue,
         }
     }
 
-    pub async fn start(&self) {
+    // TODO: remove panic and return proofs instead of taking queue
+    pub async fn fetch(&self, queue: &mut ProofsQueue) {
+        info!("Fetching proofs from batch logs");
         // Subscribe to NewBatch event from AlignedServiceManager
-        let event_sub = self
+        let logs = self
             .aligned_service_manager
-            .NewBatch_filter()
-            .subscribe()
+            .NewBatchV3_filter()
+            .from_block(0)
+            .query()
             .await
-            .expect("To subscribe to event");
-        let mut stream = event_sub.into_stream();
+            .expect("to get logs");
 
-        while let Some(log) = stream.next().await {
-            let Ok(log) = log else {
-                continue;
-            };
+        info!("Logs collected {}", logs.len());
+
+        for (batch, _) in logs {
+            info!(
+                "New batch submitted, about to process. Batch merkle root {}...",
+                batch.batchMerkleRoot
+            );
 
             // Download batch proofs from s3
-            let Ok(data) = self.s3.get_aligned_batch(log.0.batchDataPointer).await else {
-                error!("Error while downloading proofs from s3");
-                continue;
+            let data = match get_aligned_batch_from_s3(batch.batchDataPointer).await {
+                Ok(data) => data,
+                Err(err) => {
+                    error!("Error while downloading proofs from s3. Err {:?}", err);
+                    continue;
+                }
             };
 
+            info!("Data downloaded from S3, number of proofs {}", data.len());
+
             // Filter SP1 compressed proofs to and push to queue to be aggregated
-            let proofs: Vec<(Proof, Vec<u8>)> = data
+            let proofs: Vec<Proof> = data
                 .into_iter()
                 .filter_map(|p| match p.proving_system {
                     ProvingSystemId::SP1 => {
                         let elf = p.vm_program_code?;
                         let proof = bincode::deserialize(&p.proof).ok()?;
-                        let sp1_proof = SP1Proof {
-                            proof,
-                            vk: vk_from_elf(&elf),
-                        };
+                        let sp1_proof = SP1Proof { proof, elf };
 
-                        Some((Proof::SP1(sp1_proof), elf))
+                        Some(Proof::SP1(sp1_proof))
                     }
                     _ => None,
                 })
                 .collect();
 
+            info!("SP1 proofs filtered, total proofs to add {}", proofs.len());
+
             // try to add them to the queue
-            let mut queue_lock = self.queue.lock().await;
-            for (proof, elf) in proofs {
-                match queue_lock.add_proof(proof, &elf) {
+            for proof in proofs {
+                match queue.add_proof(proof) {
                     Ok(_) => info!(
                         "New proof added to queue, current length {}",
-                        queue_lock.proofs().len()
+                        queue.proofs().len()
                     ),
                     Err(e) => error!("Could not add proof, reason: {:?}", e),
                 };

aggregation-mode/src/backend/mod.rs

@@ -18,11 +18,11 @@ use alloy::{
     signers::local::LocalSigner,
 };
 use config::Config;
+use fetcher::ProofsFetcher;
 use merkle_tree::compute_proofs_merkle_root;
 use queue::ProofsQueue;
 use sp1_sdk::HashableKey;
-use std::{str::FromStr, sync::Arc, time::Duration};
-use tokio::sync::Mutex;
+use std::{str::FromStr, time::Duration};
 use tracing::{error, info, warn};
 use types::{AlignedProofAggregationService, AlignedProofAggregationServiceContract};
 
@@ -38,11 +38,12 @@ pub struct ProofAggregator {
     engine: ZKVMEngine,
     submit_proof_every_secs: u64,
     proof_aggregation_service: AlignedProofAggregationServiceContract,
-    queue: Arc<Mutex<ProofsQueue>>,
+    fetcher: ProofsFetcher,
+    queue: ProofsQueue,
 }
 
 impl ProofAggregator {
-    pub async fn new(config: &Config, queue: Arc<Mutex<ProofsQueue>>) -> Self {
+    pub fn new(config: &Config) -> Self {
         let rpc_url = config.eth_rpc_url.parse().expect("correct url");
         let signer = LocalSigner::decrypt_keystore(
             config.ecdsa.private_key_store_path.clone(),
@@ -51,17 +52,19 @@ impl ProofAggregator {
         .expect("Correct keystore signer");
         let wallet = EthereumWallet::from(signer);
         let provider = ProviderBuilder::new().wallet(wallet).on_http(rpc_url);
-        let proof_aggregation_service = AlignedProofAggregationService::new(
+        let proof_aggregation_service: AlignedProofAggregationService::AlignedProofAggregationServiceInstance<(), alloy::providers::fillers::FillProvider<alloy::providers::fillers::JoinFill<alloy::providers::fillers::JoinFill<alloy::providers::Identity, alloy::providers::fillers::JoinFill<alloy::providers::fillers::GasFiller, alloy::providers::fillers::JoinFill<alloy::providers::fillers::BlobGasFiller, alloy::providers::fillers::JoinFill<alloy::providers::fillers::NonceFiller, alloy::providers::fillers::ChainIdFiller>>>>, alloy::providers::fillers::WalletFiller<EthereumWallet>>, alloy::providers::RootProvider>> = AlignedProofAggregationService::new(
             Address::from_str(&config.proof_aggregation_service_address)
                 .expect("Address to be correct"),
             provider,
         );
+        let fetcher = ProofsFetcher::new(config);
 
         Self {
             engine: ZKVMEngine::SP1,
             submit_proof_every_secs: config.submit_proofs_every_secs,
             proof_aggregation_service,
-            queue,
+            queue: ProofsQueue::new(config.max_proofs_in_queue),
+            fetcher,
         }
     }
 
@@ -96,7 +99,9 @@ impl ProofAggregator {
     async fn aggregate_and_submit_proofs_on_chain(
         &mut self,
     ) -> Result<(), AggregatedProofSubmissionError> {
-        let proofs = self.queue.lock().await.clear();
+        self.fetcher.fetch(&mut self.queue).await;
+        let proofs = self.queue.clear();
+
         if proofs.len() == 0 {
             warn!("No proofs in queue, skipping iteration...");
             return Ok(());
@@ -141,7 +146,7 @@ impl ProofAggregator {
                     .proof_aggregation_service
                     .verify(
                         blob_tx_hash.into(),
-                        proof.vk.bytes32_raw().into(),
+                        proof.vk().bytes32_raw().into(),
                         proof.proof.public_values.to_vec().into(),
                         proof.proof.bytes().into(),
                     )

aggregation-mode/src/backend/s3.rs

@@ -1,63 +1,32 @@
 use aligned_sdk::core::types::VerificationData;
-use aws_config::{meta::region::RegionProviderChain, BehaviorVersion};
-use aws_sdk_s3::Client;
-use tracing::info;
-
-pub struct S3Client {
-    client: Client,
-    bucket_name: String,
-}
 
+#[derive(Debug)]
 pub enum GetBatchProofsError {
     Fetching,
     Deserialization,
     EmptyBody,
+    StatusFailed,
 }
 
-impl S3Client {
-    pub async fn new(bucket_name: String, endpoint_url: Option<String>) -> Self {
-        let region_provider = RegionProviderChain::default_provider().or_else("us-east-2");
-        let mut config = aws_config::defaults(BehaviorVersion::latest()).region(region_provider);
-        if let Some(endpoint_url) = &endpoint_url {
-            info!("Using custom endpoint: {}", endpoint_url);
-            config = config.endpoint_url(endpoint_url);
-        }
-        let config = config.load().await;
-
-        let mut s3_config_builder = aws_sdk_s3::config::Builder::from(&config);
-        if endpoint_url.is_some() {
-            info!("Forcing path style for custom endpoint");
-            s3_config_builder = s3_config_builder.force_path_style(true);
-        }
+pub async fn get_aligned_batch_from_s3(
+    url: String,
+) -> Result<Vec<VerificationData>, GetBatchProofsError> {
+    let response = reqwest::get(url)
+        .await
+        .map_err(|_| GetBatchProofsError::Fetching)?;
 
-        let client = Client::from_conf(s3_config_builder.build());
-
-        Self {
-            client,
-            bucket_name,
-        }
+    if !response.status().is_success() {
+        return Err(GetBatchProofsError::StatusFailed);
     }
 
-    pub async fn get_aligned_batch(
-        &self,
-        key: String,
-    ) -> Result<Vec<VerificationData>, GetBatchProofsError> {
-        let result = self
-            .client
-            .get_object()
-            .bucket(self.bucket_name.clone())
-            .key(key)
-            .send()
-            .await
-            .map_err(|_| GetBatchProofsError::Fetching)?;
-
-        let Some(bytes) = result.body.bytes() else {
-            return Err(GetBatchProofsError::EmptyBody);
-        };
+    let bytes = response
+        .bytes()
+        .await
+        .map_err(|_| GetBatchProofsError::EmptyBody)?;
+    let bytes: &[u8] = bytes.iter().as_slice();
 
-        let data: Vec<VerificationData> =
-            ciborium::from_reader(bytes).map_err(|_| GetBatchProofsError::Deserialization)?;
+    let data: Vec<VerificationData> =
+        ciborium::from_reader(bytes).map_err(|_| GetBatchProofsError::Deserialization)?;
 
-        Ok(data)
-    }
+    Ok(data)
 }

aggregation-mode/src/main.rs

@@ -1,9 +1,6 @@
-use std::{env, sync::Arc};
+use std::env;
 
-use proof_aggregator::backend::{
-    config::Config, fetcher::ProofsFetcher, queue::ProofsQueue, ProofAggregator,
-};
-use tokio::sync::Mutex;
+use proof_aggregator::backend::{config::Config, ProofAggregator};
 use tracing_subscriber::FmtSubscriber;
 
 fn read_config_filepath_from_args() -> String {
@@ -29,13 +26,9 @@ async fn main() {
     let config = Config::from_file(&config_file_path).expect("Config is valid");
     tracing::info!("Config loaded");
 
-    let queue = Arc::new(Mutex::new(ProofsQueue::new(config.max_proofs_in_queue)));
-    let mut proof_aggregator = ProofAggregator::new(&config, queue.clone()).await;
-    let proofs_fetcher = ProofsFetcher::new(&config, queue).await;
+    let mut proof_aggregator = ProofAggregator::new(&config);
 
-    // start tasks -> Proof aggregator + Proofs fetcher
     let proof_aggregator_handle = tokio::spawn(async move { proof_aggregator.start().await });
-    let proofs_fetcher_handle = tokio::spawn(async move { proofs_fetcher.start().await });
 
-    let _ = tokio::join!(proof_aggregator_handle, proofs_fetcher_handle);
+    let _ = tokio::join!(proof_aggregator_handle);
 }