feat: AWS support

Author: Adrian Nagy

cli/src/commands/node/mod.rs

@@ -138,10 +138,9 @@ pub struct Node {
     #[arg(short = 'c', long, env)]
     pub config: Option<PathBuf>,
 
-    /// Enable archive mode (seding blocks to the archive process).
-    #[arg(long, env)]
-    pub archive_address: Option<Url>,
-
+    // /// Enable archive mode (seding blocks to the archive process).
+    // #[arg(long, env)]
+    // pub archive_address: Option<Url>,
     /// Enable local precomputed storage.
     #[arg(long, env)]
     pub archive_local_storage: bool,
@@ -157,6 +156,13 @@ pub struct Node {
     pub archive_gcp_storage: bool,
 
     /// Enable AWS precomputed storage.
+    /// 
+    /// This requires the following environment variables to be set:
+    /// - AWS_ACCESS_KEY_ID
+    /// - AWS_SECRET_ACCESS_KEY
+    /// - AWS_SESSION_TOKEN
+    /// - AWS_DEFAULT_REGION
+    /// - OPENMINA_AWS_BUCKET_NAME
     #[arg(long, env)]
     pub archive_aws_storage: bool,
 }
@@ -314,17 +320,21 @@ impl Node {
             .map(|(_, option)| option.clone()),
         );
 
-        if let Some(address) = self.archive_address {
+        if archive_storage_options.is_enabled() {
             node::core::info!(
                 summary = "Archive mode enabled",
-                address = address.to_string()
+                local_storage = archive_storage_options.uses_local_precomputed_storage(),
+                archiver_process = archive_storage_options.uses_archiver_process(),
+                gcp_storage = archive_storage_options.uses_gcp_precomputed_storage(),
+                aws_storage = archive_storage_options.uses_aws_precomputed_storage(),
             );
-            // Convert URL to SocketAddr
-            let socket_addrs = address.socket_addrs(|| None).expect("Invalid URL");
 
-            let socket_addr = socket_addrs.first().expect("No socket address found");
+            archive_storage_options
+                .validate_env_vars()
+                .map_err(|e| anyhow::anyhow!(e))?;
+
             // TODO(adonagy): add options
-            node_builder.archive(*socket_addr, archive_storage_options);
+            node_builder.archive(archive_storage_options, work_dir.clone());
         }
 
         if let Some(sec_key) = self.run_snarker {

node/common/src/service/archive.rs

@@ -1,9 +1,10 @@
 use bitflags::bitflags;
-use mina_p2p_messages::v2::{self, ArchiveTransitionFronntierDiff};
+use mina_p2p_messages::v2::{self};
 use node::core::{channels::mpsc, thread};
 use node::ledger::write::BlockApplyResult;
 use std::env;
-use std::net::SocketAddr;
+use std::io::Write;
+use std::path::PathBuf;
 
 use super::NodeService;
 
@@ -23,7 +24,7 @@ bitflags! {
 
 impl ArchiveStorageOptions {
     pub fn is_enabled(&self) -> bool {
-        self.is_empty()
+        !self.is_empty()
     }
 
     pub fn requires_precomputed_block(&self) -> bool {
@@ -60,6 +61,27 @@ impl ArchiveStorageOptions {
                         .to_string(),
                 );
             }
+
+            if env::var("AWS_DEFAULT_REGION").is_err() {
+                return Err(
+                    "AWS_DEFAULT_REGION is required when AWS_PRECOMPUTED_STORAGE is enabled"
+                        .to_string(),
+                );
+            }
+
+            if env::var("OPENMINA_AWS_BUCKET_NAME").is_err() {
+                return Err(
+                    "OPENMINA_AWS_BUCKET_NAME is required when AWS_PRECOMPUTED_STORAGE is enabled"
+                        .to_string(),
+                );
+            }
+
+            // if env::var("OPENMINA_AWS_BUCKET_PATH").is_err() {
+            //     return Err(
+            //         "OPENMINA_AWS_BUCKET_PATH is required when AWS_PRECOMPUTED_STORAGE is enabled"
+            //             .to_string(),
+            //     );
+            // }
         }
 
         // TODO(adonagy): Add GCP precomputed storage validation
@@ -89,31 +111,58 @@ pub struct ArchiveService {
 }
 
 struct ArchiveServiceClients {
-    aws_client: Option<aws_sdk_s3::Client>,
+    aws_client: Option<ArchiveAWSClient>,
     gcp_client: Option<()>,
+    local_path: Option<String>,
+}
+
+struct ArchiveAWSClient {
+    client: aws_sdk_s3::Client,
+    bucket_name: String,
+    bucket_path: String,
 }
 
 impl ArchiveServiceClients {
-    async fn new(options: &ArchiveStorageOptions) -> Self {
+    async fn new(options: &ArchiveStorageOptions, work_dir: String) -> Self {
         let aws_client = if options.uses_aws_precomputed_storage() {
             let config = aws_config::load_from_env().await;
-            Some(aws_sdk_s3::Client::new(&config))
+            let bucket_name = env::var("OPENMINA_AWS_BUCKET_NAME").unwrap_or_default();
+            let bucket_path = env::var("OPENMINA_AWS_BUCKET_PATH").unwrap_or_default();
+            Some(ArchiveAWSClient {
+                client: aws_sdk_s3::Client::new(&config),
+                bucket_name,
+                bucket_path,
+            })
         } else {
             None
         };
+
+        let local_path = if options.uses_local_precomputed_storage() {
+            let env_path = env::var("OPENMINA_LOCAL_PRECOMPUTED_STORAGE_PATH");
+            let default = format!("{}/archive-precomputed", work_dir);
+            Some(env_path.unwrap_or(default))
+        } else {
+            None
+        };
+
         Self {
             aws_client,
             gcp_client: None,
+            local_path,
         }
     }
 
-    pub fn aws_client(&self) -> Option<&aws_sdk_s3::Client> {
+    pub fn aws_client(&self) -> Option<&ArchiveAWSClient> {
         self.aws_client.as_ref()
     }
 
     pub fn gcp_client(&self) -> Option<&()> {
         self.gcp_client.as_ref()
     }
+
+    pub fn local_path(&self) -> Option<&str> {
+        self.local_path.as_deref()
+    }
 }
 
 impl ArchiveService {
@@ -124,23 +173,35 @@ impl ArchiveService {
     #[cfg(not(target_arch = "wasm32"))]
     async fn run(
         mut archive_receiver: mpsc::UnboundedReceiver<BlockApplyResult>,
-        address: SocketAddr,
         options: ArchiveStorageOptions,
+        work_dir: String,
     ) {
+        use std::{fs::File, path::Path};
+
         use mina_p2p_messages::v2::PrecomputedBlock;
+        use openmina_core::NetworkConfig;
+        use reqwest::Url;
 
-        let clients = ArchiveServiceClients::new(&options).await;
+        let clients = ArchiveServiceClients::new(&options, work_dir).await;
 
-        while let Some(breadcrumb) = archive_receiver.blocking_recv() {
+        while let Some(breadcrumb) = archive_receiver.recv().await {
             if options.uses_archiver_process() {
+                let address = std::env::var("OPENMINA_ARCHIVE_ADDRESS")
+                    .expect("OPENMINA_ARCHIVE_ADDRESS is not set");
+                let address = Url::parse(&address).expect("Invalid URL");
+
+                // Convert URL to SocketAddr
+                let socket_addrs = address.socket_addrs(|| None).expect("Invalid URL");
+
+                let socket_addr = socket_addrs.first().expect("No socket address found");
                 let mut retries = ARCHIVE_SEND_RETRIES;
 
                 let archive_transition_frontier_diff: v2::ArchiveTransitionFronntierDiff =
                     breadcrumb.clone().try_into().unwrap();
 
                 while retries > 0 {
                     match rpc::send_diff(
-                        address,
+                        *socket_addr,
                         v2::ArchiveRpc::SendDiff(archive_transition_frontier_diff.clone()),
                     ) {
                         Ok(result) => {
@@ -149,9 +210,10 @@ impl ArchiveService {
                                     summary = "Archive suddenly closed connection, retrying..."
                                 );
                                 retries -= 1;
-                                std::thread::sleep(std::time::Duration::from_millis(
+                                tokio::time::sleep(std::time::Duration::from_millis(
                                     RETRY_INTERVAL_MS,
-                                ));
+                                ))
+                                .await;
                             } else {
                                 node::core::warn!(summary = "Successfully sent diff to archive");
                                 break;
@@ -164,15 +226,25 @@ impl ArchiveService {
                                 retries = retries
                             );
                             retries -= 1;
-                            std::thread::sleep(std::time::Duration::from_millis(RETRY_INTERVAL_MS));
+                            tokio::time::sleep(std::time::Duration::from_millis(RETRY_INTERVAL_MS))
+                                .await;
                         }
                     }
                 }
             }
 
             if options.requires_precomputed_block() {
-                // let key =
-                // TODO(adonagy)
+                let network_name = NetworkConfig::global().name;
+                let height = breadcrumb.block.height();
+                let state_hash = breadcrumb.block.hash();
+
+                let key = format!("{network_name}-{height}-{state_hash}.json");
+
+                node::core::info!(
+                    summary = "Uploading precomputed block to archive",
+                    key = key.clone()
+                );
+
                 let precomputed_block: PrecomputedBlock =
                     if let Ok(precomputed_block) = breadcrumb.try_into() {
                         precomputed_block
@@ -184,16 +256,46 @@ impl ArchiveService {
                     };
 
                 if options.uses_local_precomputed_storage() {
-                    // TODO(adonagy): Implement local precomputed storage
+                    // TODO(adonagy): Cleanup the unwraps (fn that returns a Result + log the error)
+                    if let Some(path) = clients.local_path() {
+                        let file_path = Path::new(path).join(key.clone());
+                        let mut file = File::create(file_path).unwrap();
+                        let json = serde_json::to_vec(&precomputed_block).unwrap();
+                        file.write_all(&json).unwrap();
+                    } else {
+                        node::core::warn!(summary = "Local precomputed storage path not set");
+                    }
                 }
 
                 if options.uses_gcp_precomputed_storage() {
                     // TODO(adonagy): Implement GCP precomputed storage
                 }
 
                 if options.uses_aws_precomputed_storage() {
-                    let client = clients.aws_client().unwrap();
-                    // put
+                    if let Some(client) = clients.aws_client() {
+                        let json = serde_json::to_string(&precomputed_block).unwrap();
+                        let res = client
+                            .client
+                            .put_object()
+                            .bucket(client.bucket_name.clone())
+                            .key(format!("{}/{}", client.bucket_path, key))
+                            .body(json.as_bytes().to_vec().into())
+                            .send()
+                            .await;
+
+                        if let Err(e) = res {
+                            node::core::warn!(
+                                summary = "Failed to upload precomputed block to AWS",
+                                error = e.to_string()
+                            );
+                        } else {
+                            node::core::warn!(
+                                summary = "Successfully uploaded precomputed block to AWS"
+                            );
+                        }
+                    } else {
+                        node::core::warn!(summary = "AWS client not initialized");
+                    }
                 }
             }
         }
@@ -209,7 +311,7 @@ impl ArchiveService {
         unimplemented!()
     }
 
-    pub fn start(address: SocketAddr, options: ArchiveStorageOptions) -> Self {
+    pub fn start(options: ArchiveStorageOptions, work_dir: String) -> Self {
         let (archive_sender, archive_receiver) = mpsc::unbounded_channel::<BlockApplyResult>();
 
         let runtime = tokio::runtime::Builder::new_current_thread()
@@ -220,7 +322,7 @@ impl ArchiveService {
         thread::Builder::new()
             .name("openmina_archive".to_owned())
             .spawn(move || {
-                runtime.block_on(Self::run(archive_receiver, address, options));
+                runtime.block_on(Self::run(archive_receiver, options, work_dir));
             })
             .unwrap();
 

node/common/src/service/builder.rs

@@ -1,4 +1,4 @@
-use std::net::SocketAddr;
+use std::{net::SocketAddr, path::PathBuf};
 
 use ledger::proofs::provers::BlockProver;
 use node::{
@@ -101,12 +101,8 @@ impl NodeServiceCommonBuilder {
         self
     }
 
-    pub fn archive_init(
-        &mut self,
-        address: SocketAddr,
-        options: ArchiveStorageOptions,
-    ) -> &mut Self {
-        self.archive = Some(ArchiveService::start(address, options));
+    pub fn archive_init(&mut self, options: ArchiveStorageOptions, work_dir: String) -> &mut Self {
+        self.archive = Some(ArchiveService::start(options, work_dir));
         self
     }
 

node/native/src/node/builder.rs

@@ -1,7 +1,7 @@
 use std::{
     fs::File,
     io::{BufRead, BufReader, Read},
-    net::{IpAddr, SocketAddr},
+    net::IpAddr,
     path::Path,
     sync::Arc,
     time::Duration,
@@ -218,9 +218,9 @@ impl NodeBuilder {
         Ok(self.block_producer(key, provers))
     }
 
-    pub fn archive(&mut self, address: SocketAddr, options: ArchiveStorageOptions) -> &mut Self {
-        self.archive = Some(ArchiveConfig::new(&address.to_string()));
-        self.service.archive_init(address, options);
+    pub fn archive(&mut self, options: ArchiveStorageOptions, work_dir: String) -> &mut Self {
+        self.archive = Some(ArchiveConfig::new(work_dir.clone()));
+        self.service.archive_init(options, work_dir.clone());
         self
     }
 

node/native/src/service/builder.rs

@@ -1,4 +1,4 @@
-use std::net::SocketAddr;
+use std::{net::SocketAddr, path::PathBuf};
 
 use ledger::proofs::provers::BlockProver;
 use node::{
@@ -56,12 +56,8 @@ impl NodeServiceBuilder {
         self
     }
 
-    pub fn archive_init(
-        &mut self,
-        address: SocketAddr,
-        options: ArchiveStorageOptions,
-    ) -> &mut Self {
-        self.common.archive_init(address, options);
+    pub fn archive_init(&mut self, options: ArchiveStorageOptions, work_dir: String) -> &mut Self {
+        self.common.archive_init(options, work_dir);
         self
     }