feat: upload batches to multiple storage services

Author: JuArce

.gitignore

@@ -12,6 +12,7 @@ crates/cli/batch_inclusion_responses/*
 **/aligned_verification_data
 **/broadcast
 volume
+volume2
 config-files/*.last_processed_batch.json
 config-files/*.last_aggregated_block.json
 

aggregation_mode/src/backend/fetcher.rs

@@ -9,7 +9,7 @@ use crate::{
         risc0_aggregator::Risc0ProofReceiptAndImageId, sp1_aggregator::SP1ProofWithPubValuesAndElf,
         AlignedProof, ZKVMEngine,
     },
-    backend::s3::get_aligned_batch_from_s3,
+    backend::s3::get_aligned_batch_from_s3_with_multiple_urls,
 };
 use aligned_sdk::common::types::ProvingSystemId;
 use alloy::{
@@ -97,7 +97,7 @@ impl ProofsFetcher {
             );
 
             // Download batch proofs from s3
-            let data = match get_aligned_batch_from_s3(batch.batchDataPointer).await {
+            let data = match get_aligned_batch_from_s3_with_multiple_urls(batch.batchDataPointer).await {
                 Ok(data) => data,
                 Err(err) => {
                     error!("Error while downloading proofs from s3. Err {:?}", err);

aggregation_mode/src/backend/s3.rs

@@ -1,3 +1,4 @@
+use tracing::{info, warn};
 use aligned_sdk::common::types::VerificationData;
 
 #[derive(Debug)]
@@ -12,10 +13,57 @@ pub enum GetBatchProofsError {
 
 // needed to make S3 bucket work
 const DEFAULT_USER_AGENT: &str = "proof-aggregator/aligned-layer";
+const MAX_BATCH_URLS: usize = 5;
+
+// get_aligned_batch_from_s3_with_multiple_urls tries multiple comma-separated URLs until first successful response
+pub async fn get_aligned_batch_from_s3_with_multiple_urls(
+    urls: String,
+) -> Result<Vec<VerificationData>, GetBatchProofsError> {
+    // Parse comma-separated URLs and limit to max 5
+    let parsed_urls = parse_batch_urls(&urls);
+    info!("Getting batch from data service with {} URLs: {:?}", parsed_urls.len(), parsed_urls);
+    
+    let mut errors = Vec::new();
+
+    // Try each URL until first successful response
+    for (_i, url) in parsed_urls.iter().enumerate() {
+        match get_aligned_batch_from_s3(url.clone()).await {
+            Ok(data) => {
+                return Ok(data);
+            }
+            Err(err) => {
+                warn!("Failed to fetch batch from URL {}: {:?}", url, err);
+                errors.push(format!("URL {}: {:?}", url, err));
+            }
+        }
+    }
+
+    // All URLs failed
+    Err(GetBatchProofsError::FetchingS3Batch(format!(
+        "Failed to get batch from all URLs, errors: {}",
+        errors.join("; ")
+    )))
+}
+
+// parse_batch_urls parses comma-separated URLs and limits to max 5
+fn parse_batch_urls(batch_urls: &str) -> Vec<String> {
+    let mut urls = Vec::new();
+    for url in batch_urls.split(',') {
+        let trimmed_url = url.trim();
+        if !trimmed_url.is_empty() {
+            urls.push(trimmed_url.to_string());
+            if urls.len() > MAX_BATCH_URLS {
+                break;
+            }
+        }
+    }
+    urls
+}
 
 pub async fn get_aligned_batch_from_s3(
     url: String,
 ) -> Result<Vec<VerificationData>, GetBatchProofsError> {
+    info!("Fetching batch from S3 URL: {}", url);
     let client = reqwest::Client::builder()
         .user_agent(DEFAULT_USER_AGENT)
         .build()

crates/batcher/.env.dev

@@ -1,8 +1,18 @@
+# Main S3 bucket configuration
 AWS_SECRET_ACCESS_KEY=test
 AWS_REGION=us-east-2
 AWS_ACCESS_KEY_ID=test
 AWS_BUCKET_NAME=aligned.storage
 UPLOAD_ENDPOINT=http://localhost:4566
 DOWNLOAD_ENDPOINT=http://localhost:4566/aligned.storage
+
+# Secondary S3 bucket configuration
+AWS_SECRET_ACCESS_KEY_SECONDARY=test2
+AWS_REGION_SECONDARY=us-west-1
+AWS_ACCESS_KEY_ID_SECONDARY=test2
+AWS_BUCKET_NAME_SECONDARY=aligned.storage
+UPLOAD_ENDPOINT_SECONDARY=http://localhost:4567
+DOWNLOAD_ENDPOINT_SECONDARY=http://localhost:4567/aligned.storage
+
 RUST_LOG=info
 RUST_BACKTRACE=1

crates/batcher/.env.docker

@@ -1,8 +1,18 @@
+# Main S3 bucket configuration
 AWS_SECRET_ACCESS_KEY=test
 AWS_REGION=us-east-2
 AWS_ACCESS_KEY_ID=test
 AWS_BUCKET_NAME=aligned.storage
 UPLOAD_ENDPOINT=http://localstack:4566
 DOWNLOAD_ENDPOINT=http://localstack:4566/aligned.storage
+
+# Secondary S3 bucket configuration
+AWS_SECRET_ACCESS_KEY_SECONDARY=test2
+AWS_REGION_SECONDARY=us-west-1
+AWS_ACCESS_KEY_ID_SECONDARY=test2
+AWS_BUCKET_NAME_SECONDARY=aligned.storage
+UPLOAD_ENDPOINT_SECONDARY=http://localhost:4567
+DOWNLOAD_ENDPOINT_SECONDARY=http://localhost:4567/aligned.storage
+
 RUST_LOG=info
 RUST_BACKTRACE=1

crates/batcher/.env.example

@@ -1,7 +1,17 @@
+# Main S3 bucket configuration
 AWS_SECRET_ACCESS_KEY=<secret_access_key>
 AWS_REGION=<region>
 AWS_ACCESS_KEY_ID=<access_key_id>
 AWS_BUCKET_NAME=<bucket_name>
 UPLOAD_ENDPOINT=<upload_endpoint. For non-aws storage>
 DOWNLOAD_ENDPOINT=<download_endpoint>
+
+# Secondary S3 bucket configuration
+AWS_SECRET_ACCESS_KEY_SECONDARY=<secondary_secret_access_key>
+AWS_REGION_SECONDARY=<secondary_region>
+AWS_ACCESS_KEY_ID_SECONDARY=<secondary_access_key_id>
+AWS_BUCKET_NAME_SECONDARY=<secondary_bucket_name>
+UPLOAD_ENDPOINT_SECONDARY=<secondary_upload_endpoint. For non-aws storage>
+DOWNLOAD_ENDPOINT_SECONDARY=<secondary_download_endpoint>
+
 RUST_LOG=<log_level> # info, debug, error, warn, etc.

crates/batcher/src/lib.rs

@@ -75,6 +75,9 @@ pub struct Batcher {
     s3_client: S3Client,
     s3_bucket_name: String,
     download_endpoint: String,
+    s3_client_secondary: S3Client,
+    s3_bucket_name_secondary: String,
+    download_endpoint_secondary: String,
     eth_ws_url: String,
     eth_ws_url_fallback: String,
     batcher_signer: Arc<SignerMiddlewareT>,
@@ -106,15 +109,36 @@ impl Batcher {
         dotenv().ok();
 
         // https://docs.aws.amazon.com/sdk-for-rust/latest/dg/localstack.html
-        let upload_endpoint = env::var("UPLOAD_ENDPOINT").ok();
+        // Primary S3 configuration
+        let s3_config_primary = s3::S3Config {
+            access_key_id: env::var("AWS_ACCESS_KEY_ID").ok(),
+            secret_access_key: env::var("AWS_SECRET_ACCESS_KEY").ok(),
+            region: env::var("AWS_REGION").ok(),
+            endpoint_url: env::var("UPLOAD_ENDPOINT").ok(),
+        };
 
         let s3_bucket_name =
             env::var("AWS_BUCKET_NAME").expect("AWS_BUCKET_NAME not found in environment");
 
         let download_endpoint =
             env::var("DOWNLOAD_ENDPOINT").expect("DOWNLOAD_ENDPOINT not found in environment");
 
-        let s3_client = s3::create_client(upload_endpoint).await;
+        let s3_client = s3::create_client(s3_config_primary).await;
+
+        // Secondary S3 configuration
+        let s3_config_secondary = s3::S3Config {
+            access_key_id: env::var("AWS_ACCESS_KEY_ID_SECONDARY").ok(),
+            secret_access_key: env::var("AWS_SECRET_ACCESS_KEY_SECONDARY").ok(),
+            region: env::var("AWS_REGION_SECONDARY").ok(),
+            endpoint_url: env::var("UPLOAD_ENDPOINT_SECONDARY").ok(),
+        };
+
+        let s3_bucket_name_secondary = env::var("AWS_BUCKET_NAME_SECONDARY")
+            .expect("AWS_BUCKET_NAME_SECONDARY not found in environment");
+        let download_endpoint_secondary = env::var("DOWNLOAD_ENDPOINT_SECONDARY")
+            .expect("DOWNLOAD_ENDPOINT_SECONDARY not found in environment");
+
+        let s3_client_secondary = s3::create_client(s3_config_secondary).await;
 
         let config = ConfigFromYaml::new(config_file);
         // Ensure max_batch_bytes_size can at least hold one proof of max_proof_size,
@@ -252,6 +276,9 @@ impl Batcher {
             s3_client,
             s3_bucket_name,
             download_endpoint,
+            s3_client_secondary,
+            s3_bucket_name_secondary,
+            download_endpoint_secondary,
             eth_ws_url: config.eth_ws_url,
             eth_ws_url_fallback: config.eth_ws_url_fallback,
             batcher_signer,
@@ -1541,7 +1568,16 @@ impl Batcher {
         let batch_merkle_root_hex = hex::encode(batch_merkle_root);
         info!("Batch merkle root: 0x{}", batch_merkle_root_hex);
         let file_name = batch_merkle_root_hex.clone() + ".json";
-        let batch_data_pointer: String = "".to_owned() + &self.download_endpoint + "/" + &file_name;
+        
+        let batch_data_pointer = self.upload_batch_to_multiple_s3(batch_bytes, &file_name).await?;
+        if let Err(e) = self
+            .telemetry
+            .task_uploaded_to_s3(&batch_merkle_root_hex)
+            .await
+        {
+            warn!("Failed to send task status to telemetry: {:?}", e);
+        };
+        info!("Batch upload to: {}", batch_data_pointer);
 
         let num_proofs_in_batch = leaves.len();
         let gas_per_proof = (self.constant_gas_cost()
@@ -1577,16 +1613,6 @@ impl Batcher {
             .gas_price_used_on_latest_batch
             .set(gas_price.as_u64() as i64);
 
-        info!("Uploading batch to S3...");
-        self.upload_batch_to_s3(batch_bytes, &file_name).await?;
-        if let Err(e) = self
-            .telemetry
-            .task_uploaded_to_s3(&batch_merkle_root_hex)
-            .await
-        {
-            warn!("Failed to send task status to telemetry: {:?}", e);
-        };
-        info!("Batch sent to S3 with name: {}", file_name);
         if let Err(e) = self
             .telemetry
             .task_created(
@@ -1857,22 +1883,61 @@ impl Batcher {
         unlocked
     }
 
+    /// Uploads the batch to both S3 buckets and returns the comma-separated URLs of successful uploads.
+    /// Returns an error only if all uploads fail.
+    async fn upload_batch_to_multiple_s3(
+        &self,
+        batch_bytes: &[u8],
+        file_name: &str,
+    ) -> Result<String, BatcherError> {
+        // Upload to both S3 buckets and collect successful URLs
+        let mut successful_urls = Vec::new();
+        
+        // Try primary S3 upload
+        if let Ok(_) = self.upload_batch_to_s3(&self.s3_client, batch_bytes, file_name, &self.s3_bucket_name).await {
+            let primary_url = format!("{}/{}", self.download_endpoint, file_name);
+            successful_urls.push(primary_url.clone());
+            info!("Successfully uploaded batch to primary S3: {}", primary_url);
+        } else {
+            warn!("Failed to upload batch to primary S3");
+        }
+        
+        // Try secondary S3 upload
+        if let Ok(_) = self.upload_batch_to_s3(&self.s3_client_secondary, batch_bytes, file_name, &self.s3_bucket_name_secondary).await {
+            let secondary_url = format!("{}/{}", self.download_endpoint_secondary, file_name);
+            successful_urls.push(secondary_url.clone());
+            info!("Successfully uploaded batch to secondary S3: {}", secondary_url);
+        } else {
+            warn!("Failed to upload batch to secondary S3");
+        }
+        
+        // If no uploads succeeded, return error
+        if successful_urls.is_empty() {
+            error!("Failed to upload batch to both S3 buckets");
+            return Err(BatcherError::BatchUploadError("Failed to upload to any S3 bucket".to_string()));
+        }
+        
+        Ok(successful_urls.join(","))
+    }
+
     /// Uploads the batch to s3.
     /// Retries on recoverable errors using exponential backoff up to `ETHEREUM_CALL_MAX_RETRIES` times:
     /// (0,5 secs - 1 secs - 2 secs - 4 secs - 8 secs).
     async fn upload_batch_to_s3(
         &self,
+        s3_client: &S3Client,
         batch_bytes: &[u8],
         file_name: &str,
+        bucket_name: &str,
     ) -> Result<(), BatcherError> {
         let start = Instant::now();
         let result = retry_function(
             || {
                 Self::upload_batch_to_s3_retryable(
                     batch_bytes,
                     file_name,
-                    self.s3_client.clone(),
-                    &self.s3_bucket_name,
+                    s3_client.clone(),
+                    bucket_name,
                 )
             },
             ETHEREUM_CALL_MIN_RETRY_DELAY,

crates/batcher/src/s3/mod.rs

@@ -4,19 +4,47 @@ use aws_sdk_s3::error::SdkError;
 use aws_sdk_s3::operation::put_object::{PutObjectError, PutObjectOutput};
 use aws_sdk_s3::primitives::ByteStream;
 use aws_sdk_s3::Client;
+use aws_sdk_s3::config::Region;
 use log::info;
 
-pub async fn create_client(endpoint_url: Option<String>) -> Client {
-    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 {
+pub struct S3Config {
+    pub access_key_id: Option<String>,
+    pub secret_access_key: Option<String>,
+    pub region: Option<String>,
+    pub endpoint_url: Option<String>,
+}
+
+pub async fn create_client(s3_config: S3Config) -> Client {
+    let mut config = aws_config::defaults(BehaviorVersion::latest());
+    
+    if let Some(region) = s3_config.region {
+        let region_provider = RegionProviderChain::first_try(Region::new(region)).or_else("us-east-2");
+        config = config.region(region_provider);
+    } else {
+        let region_provider = RegionProviderChain::default_provider().or_else("us-east-2");
+        config = config.region(region_provider);
+    }
+
+    if let (Some(access_key_id), Some(secret_access_key)) = (s3_config.access_key_id, s3_config.secret_access_key) {
+        let credentials = aws_sdk_s3::config::Credentials::new(
+            access_key_id,
+            secret_access_key,
+            None,
+            None,
+            "custom",
+        );
+        config = config.credentials_provider(credentials);
+    }
+
+    if let Some(endpoint_url) = &s3_config.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() {
+    if s3_config.endpoint_url.is_some() {
         info!("Forcing path style for custom endpoint");
         s3_config_builder = s3_config_builder.force_path_style(true);
     }