fix(su): add max read memory #980

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

@@ -1092,18 +1092,32 @@ mod bytestore {
             &self,
             ids: Vec<(String, Option<String>, String, String)>,
         ) -> Result<DashMap<(String, Option<String>, String, String), Vec<u8>>, String> {
+            let max_memory_usage = self.config.max_read_memory;
             let binaries = Arc::new(DashMap::new());
             let db = match self.db.read() {
                 Ok(r) => r,
                 Err(_) => return Err("Failed to acquire read lock".into()),
             };
 
             if let Some(ref db) = *db {
+                let mut total_memory_usage: usize = 0;
+
                 for id in ids {
                     let binaries = binaries.clone();
-
                     let key = ByteStore::create_key(&id.0, &id.1, &id.2, &id.3);
                     if let Ok(Some(value)) = db.get(&key) {
+                        /*
+                          This is added here because really large message lists
+                          with large messages are filling up the machines memory
+                          and freezing it.
+                        */
+                        total_memory_usage += value.len();
+                        if total_memory_usage > max_memory_usage {
+                            return Err(format!(
+                                "Memory usage exceeded the limit: {} bytes",
+                                max_memory_usage
+                            ));
+                        }
                         binaries.insert(id.clone(), value);
                     }
                 }

servers/su/src/domain/config.rs

@@ -20,6 +20,7 @@ pub struct AoConfig {
     pub db_write_connections: u32,
     pub db_read_connections: u32,
     pub enable_metrics: bool,
+    pub max_read_memory: usize,
 }
 
 impl AoConfig {
@@ -65,6 +66,10 @@ impl AoConfig {
             Ok(val) => val == "true",
             Err(_e) => false,
         };
+        let max_read_memory = match env::var("MAX_READ_MEMORY") {
+            Ok(val) => val.parse().unwrap(),
+            Err(_e) => 1_073_741_824,
+        };
         Ok(AoConfig {
             database_url: env::var("DATABASE_URL")?,
             database_read_url,
@@ -80,6 +85,7 @@ impl AoConfig {
             db_write_connections,
             db_read_connections,
             enable_metrics,
+            max_read_memory,
         })
     }
 }
@@ -91,4 +97,4 @@ impl Config for AoConfig {
     fn scheduler_list_path(&self) -> String {
         self.scheduler_list_path.clone()
     }
-}
+}