feat(su): reupload client

Author: VinceJuliano

servers/su/cli

@@ -3,18 +3,19 @@ use std::io;
 use su::domain::migrate_to_disk;
 use su::domain::migrate_to_local;
 use su::domain::sync_local_drives;
+use su::domain::reupload_bundles;
 
 #[tokio::main]
 async fn main() -> io::Result<()> {
     let args: Vec<String> = env::args().collect();
 
     if args.len() < 2 {
         eprintln!("Usage: {} <function_name>", args[0]);
-        eprintln!("Available functions: migrate_to_disk, migrate_to_local, sync_local_drives");
+        eprintln!("Available functions: migrate_to_disk, migrate_to_local, sync_local_drives, reupload_bundles");
         return Ok(());
     }
 
-    let interval = if args.len() >= 3 {
+    let interval = if args.len() >= 3 && args[1] == "sync_local_drives" {
         match args[2].parse::<u64>() {
             Ok(val) => val,
             Err(_) => {
@@ -26,6 +27,20 @@ async fn main() -> io::Result<()> {
         5
     };
 
+    let pids = if args.len() > 3 && args[1] == "reupload_bundles" {
+        args[2].clone()
+    } else {
+        eprintln!("Must provid pids, and since value");
+        return Ok(());
+    };
+
+    let since = if args.len() > 3 && args[1] == "reupload_bundles" {
+        args[3].clone()
+    } else {
+        eprintln!("Must provid since");
+        return Ok(());
+    };
+
     match args[1].as_str() {
         "migrate_to_disk" => {
             migrate_to_disk().await.unwrap();
@@ -36,6 +51,9 @@ async fn main() -> io::Result<()> {
         "sync_local_drives" => {
             sync_local_drives(interval).await.unwrap();
         }
+        "reupload_bundles" => {
+            reupload_bundles(pids, since).await.unwrap();
+        }
         _ => {
             eprintln!("Invalid function name: {}", args[1]);
             eprintln!("Available functions: migrate_to_disk, migrate_to_local, sync_local_drives");

servers/su/src/bin/cli.rs

@@ -494,6 +494,14 @@ impl DataStore for LocalStoreClient {
         Err(StoreErrorType::NotFound("Process not found".to_string()))
     }
 
+    fn get_bundle_by_assignment(&self, tx_id: &str) -> Result<Vec<u8>, StoreErrorType> {
+        let assignment_key = self.msg_assignment_key(tx_id);
+        if let Some(message_bundle) = self.file_db.get(assignment_key.as_bytes())? {
+            return Ok(message_bundle);
+        }
+        Err(StoreErrorType::DatabaseError("Failed to get bundle".to_string()))
+    }
+
     fn get_message(&self, tx_id: &str) -> Result<Message, StoreErrorType> {
         let assignment_key = self.msg_assignment_key(tx_id);
         if let Some(message_bundle) = self.file_db.get(assignment_key.as_bytes())? {
@@ -819,6 +827,22 @@ impl DataStore for LocalStoreClient {
         )?)
     }
 
+    async fn assignments_since(
+        &self,
+        process_id: &String,
+        since: &String,
+        limit: i64
+    ) -> Result<Vec<String>, StoreErrorType> {
+        let (keys, _has_next_page) = self.fetch_message_range(
+            process_id,
+            &Some(since.clone()),
+            &None,
+            &Some(limit.try_into().unwrap())
+        ).await?;
+
+        Ok(keys.into_iter().map(|(_, assignment_id)| assignment_id).collect())
+    }
+
     /*
       This is a stripped down version of get_messages
       used for retrieving bundles

servers/su/src/domain/clients/local_store/store.rs

@@ -1309,6 +1309,25 @@ impl DataStore for StoreClient {
         ))
     }
 
+    fn get_bundle_by_assignment(&self, tx_id: &str) -> Result<Vec<u8>, StoreErrorType> {
+        use super::schema::messages::dsl::*;
+        let conn = &mut self.get_read_conn()?;
+
+        let db_message_result: Result<Option<DbMessage>, DieselError> = messages
+            .filter(assignment_id.eq(tx_id))
+            .order(timestamp.asc())
+            .first(conn)
+            .optional();
+
+        match db_message_result {
+            Ok(Some(db_message)) => {
+                Ok(db_message.bundle)
+            }
+            Ok(None) => Err(StoreErrorType::NotFound("Bundle not found".to_string())),
+            Err(e) => Err(StoreErrorType::from(e)),
+        }
+    }
+
     fn get_message(&self, tx_id: &str) -> Result<Message, StoreErrorType> {
         use super::schema::messages::dsl::*;
         let conn = &mut self.get_read_conn()?;
@@ -1381,6 +1400,43 @@ impl DataStore for StoreClient {
             Err(e) => Err(StoreErrorType::from(e)),
         }
     }
+
+    async fn assignments_since(
+        &self,
+        pid: &String,
+        since: &String,
+        limit: i64
+    ) -> Result<Vec<String>, StoreErrorType> {
+        use super::schema::messages::dsl::*;
+        let conn = &mut self.get_read_conn()?;
+        let mut query = messages
+            .filter(process_id.eq(pid.clone()))
+            .into_boxed();
+        let from_timestamp = since
+            .parse::<i64>()
+            .map_err(StoreErrorType::from)?;
+        query = query.filter(timestamp.gt(from_timestamp));
+
+        let db_assignment_ids_result: Result<Vec<Option<String>>, DieselError> = query
+            .select(assignment_id)
+            .order(timestamp.asc())
+            .limit(limit)
+            .load(conn);
+
+        let assignment_ids_o = match db_assignment_ids_result {
+            Ok(d) => d,
+            Err(_) => return Err(
+                StoreErrorType::DatabaseError("Failed to read assignments".to_string())
+            )
+        };
+
+        let assignment_ids: Vec<String> = assignment_ids_o
+            .into_iter()
+            .filter_map(|id| id)
+            .collect();
+
+        Ok(assignment_ids)
+    }
 }
 
 impl RouterDataStore for StoreClient {

servers/su/src/domain/clients/store.rs

@@ -1,4 +1,5 @@
 use std::sync::Arc;
+use std::io;
 
 use reqwest::{Client, Url};
 
@@ -12,6 +13,13 @@ use crate::domain::core::dal::{Uploader, UploaderErrorType};
 use crate::domain::bytes;
 use crate::domain::Log;
 
+use crate::domain::clients::store::StoreClient;
+use crate::domain::clients::local_store::store::LocalStoreClient;
+use crate::domain::config::AoConfig;
+use crate::domain::core::dal::{DataStore, Gateway};
+use crate::domain::clients::gateway::ArweaveGateway;
+use crate::domain::SuLog;
+
 pub struct UploaderClient {
     node_url: Url,
     cache_url: Url,
@@ -234,3 +242,51 @@ impl Uploader for UploaderClient {
         Ok(())
     }
 }
+
+pub async fn reupload_bundles(pids: String, since: String) -> io::Result<()> {
+    let config = AoConfig::new(None).expect("Failed to read configuration");
+    let gateway: Arc<dyn Gateway> = Arc::new(
+        ArweaveGateway::new().await.expect("Failed to init the gateway")
+    );
+    let logger = SuLog::init();
+    let uploader = UploaderClient::new(
+        &config.upload_node_url, &config.cache_url, logger.clone()
+    ).unwrap();
+
+    let data_store: Arc<dyn DataStore> = if config.use_local_store == false {
+        Arc::new(StoreClient::new_single_connection()
+            .expect("Failed to create StoreClient"))
+    } else {
+        Arc::new(LocalStoreClient::new_read_only(
+            &config.su_file_db_dir, 
+            &config.su_index_db_dir
+        ).expect("Failed to create LocalStoreClient"))
+    };
+
+    for pid in pids.split(',') {
+      let p = pid.to_string();
+      let assignments = data_store.assignments_since(&p, &since, 5000000).await.unwrap();
+      logger.log(format!("processing pid {}", pid));
+
+      for assignment in assignments {
+        sleep(Duration::from_millis(300)).await;
+        let gateway_tx = gateway.gql_tx(&assignment).await;
+        match gateway_tx {
+            Ok(_) => {
+                logger.log(format!("Assignment found, skipping {}", assignment));
+            },
+            Err(e) => {
+                if e == "Transaction not found" {
+                    logger.log(format!("Assignment not found, reuploading {}", assignment));
+                    let bundle = data_store.get_bundle_by_assignment(&assignment).unwrap();
+                    uploader.upload(bundle).unwrap();
+                } else {
+                    logger.log(format!("Error fetching tx {}: {}", assignment, e));
+                }
+            }
+        }
+      }
+    }
+
+    return Ok(())
+}

servers/su/src/domain/clients/uploader.rs

@@ -181,6 +181,13 @@ pub trait DataStore: Send + Sync {
         process_id: &String,
         deep_hash: &String,
     ) -> Result<(), StoreErrorType>;
+    async fn assignments_since(
+        &self,
+        process_id: &String,
+        since: &String,
+        limit: i64
+    ) -> Result<Vec<String>, StoreErrorType>;
+    fn get_bundle_by_assignment(&self, tx_id: &str) -> Result<Vec<u8>, StoreErrorType>;
 }
 
 #[async_trait]

servers/su/src/domain/core/dal.rs

@@ -16,9 +16,10 @@ use clients::{
 };
 use config::AoConfig;
 use core::dal::{Config, DataStore, Gateway, Log, MockRouterDataStore, ExtRouter};
-use logger::SuLog;
+pub use logger::SuLog;
 
 pub use clients::metrics::PromMetrics;
+pub use clients::uploader::reupload_bundles;
 pub use core::flows;
 pub use core::router;
 pub use core::bytes;