Merge pull request #269 from akshaylg0314/containerinfo

feat(api)-Store ContainerInfo into ETCD at MonitoringServer[#265]

Author: leeeunkoo

src/agent/nodeagent/src/config.rs

@@ -2,8 +2,8 @@ use serde::Deserialize;
 use std::fs::File;
 use std::io::Read;
 use std::path::Path;
-use thiserror::Error;
 use std::sync::OnceLock;
+use thiserror::Error;
 
 // Global config instance
 static NODEAGENT_CONFIG: OnceLock<Config> = OnceLock::new();
@@ -12,7 +12,7 @@ static NODEAGENT_CONFIG: OnceLock<Config> = OnceLock::new();
 pub enum ConfigError {
     #[error("Failed to read config file: {0}")]
     IoError(#[from] std::io::Error),
-    
+
     #[error("Failed to parse YAML: {0}")]
     YamlError(#[from] serde_yaml::Error),
 }
@@ -62,7 +62,7 @@ impl Config {
         let mut file = File::open(path)?;
         let mut contents = String::new();
         file.read_to_string(&mut contents)?;
-        
+
         let config = serde_yaml::from_str(&contents)?;
         Ok(config)
     }

src/agent/nodeagent/src/main.rs

@@ -3,14 +3,14 @@
 //! This file sets up the asynchronous runtime, initializes the manager and gRPC server,
 //! and launches both concurrently. It also provides unit tests for initialization.
 
-use common::nodeagent::HandleYamlRequest;
 use clap::Parser;
+use common::nodeagent::HandleYamlRequest;
 use std::path::PathBuf;
 mod bluechi;
+pub mod config;
 pub mod grpc;
 pub mod manager;
 pub mod resource;
-pub mod config;
 
 use common::nodeagent::node_agent_connection_server::NodeAgentConnectionServer;
 use tokio::sync::mpsc::{channel, Receiver, Sender};
@@ -19,15 +19,19 @@ use tokio::sync::mpsc::{channel, Receiver, Sender};
 ///
 /// This function creates the manager, initializes it, and then runs it.
 /// If initialization or running fails, errors are printed to stderr.
-async fn launch_manager(rx_grpc: Receiver<HandleYamlRequest>, hostname: String, config: config::Config) {
+async fn launch_manager(
+    rx_grpc: Receiver<HandleYamlRequest>,
+    hostname: String,
+    config: config::Config,
+) {
     let mut manager = manager::NodeAgentManager::new(rx_grpc, hostname.clone()).await;
 
     match manager.initialize().await {
         Ok(_) => {
             println!("NodeAgentManager successfully initialized");
             // Add registration with API server
             let mut sender = grpc::sender::NodeAgentSender::default();
-            
+
             // Use IP address from config file
             let host_ip = config.get_host_ip();
             let node_id = format!("{}-{}", hostname, host_ip);
@@ -141,12 +145,16 @@ async fn main() {
             config
         }
         Err(err) => {
-            eprintln!("Error loading configuration from {}: {}", args.config.display(), err);
+            eprintln!(
+                "Error loading configuration from {}: {}",
+                args.config.display(),
+                err
+            );
             eprintln!("Falling back to default configuration");
             config::Config::default()
         }
     };
-    
+
     // Set global config for other parts of the application
     config::Config::set_global(app_config.clone());
 
@@ -169,13 +177,13 @@ async fn main() {
 
 #[cfg(test)]
 mod tests {
-    use crate::launch_manager;
     use crate::config::Config;
+    use crate::launch_manager;
     use common::nodeagent::HandleYamlRequest;
+    use std::path::PathBuf;
     use tokio::sync::mpsc::{channel, Receiver, Sender};
     use tokio::task::LocalSet;
     use tokio::time::{sleep, Duration};
-    use std::path::PathBuf;
 
     #[tokio::test]
     async fn test_main_initializes_channels() {

src/server/monitoringserver/src/data_structures.rs

@@ -3,6 +3,7 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
+use common::monitoringserver::ContainerInfo;
 use common::monitoringserver::NodeInfo;
 use serde::{Deserialize, Serialize};
 use std::collections::HashMap;
@@ -54,6 +55,8 @@ pub struct DataStore {
     pub nodes: HashMap<String, NodeInfo>,
     pub socs: HashMap<String, SocInfo>,
     pub boards: HashMap<String, BoardInfo>,
+    pub containers: HashMap<String, ContainerInfo>,
+    pub container_node_mapping: HashMap<String, String>, // ADD THIS LINE
 }
 
 impl DataStore {
@@ -62,6 +65,8 @@ impl DataStore {
             nodes: HashMap::new(),
             socs: HashMap::new(),
             boards: HashMap::new(),
+            containers: HashMap::new(),
+            container_node_mapping: HashMap::new(), // ADD THIS LINE
         }
     }
 
@@ -208,6 +213,133 @@ impl DataStore {
         Ok(())
     }
 
+    /// Stores ContainerInfo to memory and etcd
+    pub async fn store_container_info(
+        &mut self,
+        container_info: ContainerInfo,
+    ) -> Result<(), String> {
+        let container_id = container_info.id.clone();
+
+        // Store container in memory
+        self.containers
+            .insert(container_id.clone(), container_info.clone());
+
+        // Store to etcd with error handling
+        if let Err(e) = crate::etcd_storage::store_container_info(&container_info).await {
+            eprintln!(
+                "[ETCD] Warning: Failed to store ContainerInfo to etcd: {}",
+                e
+            );
+            // Don't fail the entire operation, just log the warning
+        }
+
+        println!(
+            "[DataStore] Stored container info: {} from node {}",
+            container_id,
+            container_info
+                .names
+                .first()
+                .unwrap_or(&"unnamed".to_string())
+        );
+        Ok(())
+    }
+
+    /// Stores ContainerInfo to memory and etcd, with explicit node association
+    pub async fn store_container_info_with_node(
+        &mut self,
+        container_info: ContainerInfo,
+        node_name: String,
+    ) -> Result<(), String> {
+        let container_id = container_info.id.clone();
+
+        // Store container in memory
+        self.containers
+            .insert(container_id.clone(), container_info.clone());
+
+        // Store the node association
+        self.container_node_mapping
+            .insert(container_id.clone(), node_name.clone());
+
+        // Store to etcd
+        if let Err(e) = crate::etcd_storage::store_container_info(&container_info).await {
+            eprintln!(
+                "[ETCD] Warning: Failed to store ContainerInfo to etcd: {}",
+                e
+            );
+        }
+
+        println!(
+            "[DataStore] Stored container {} on node {}",
+            container_id, node_name
+        );
+        Ok(())
+    }
+
+    /// Retrieves ContainerInfo from memory, fallback to etcd
+    pub async fn get_container_info(&self, container_id: &str) -> Result<ContainerInfo, String> {
+        // Try memory first
+        if let Some(container_info) = self.containers.get(container_id) {
+            return Ok(container_info.clone());
+        }
+
+        // Fallback to etcd
+        match crate::etcd_storage::get_container_info(container_id).await {
+            Ok(container_info) => Ok(container_info),
+            Err(e) => Err(format!("Container not found in memory or etcd: {}", e)),
+        }
+    }
+
+    /// Gets all containers from memory
+    pub fn get_all_containers(&self) -> &HashMap<String, ContainerInfo> {
+        &self.containers
+    }
+
+    /// Gets all containers for a specific node
+    pub fn get_containers_by_node(&self, node_name: &str) -> Vec<&ContainerInfo> {
+        self.container_node_mapping
+            .iter()
+            .filter(|(_, mapped_node)| *mapped_node == node_name)
+            .filter_map(|(container_id, _)| self.containers.get(container_id))
+            .collect()
+    }
+
+    /// Removes container from memory and etcd
+    pub async fn remove_container_info(&mut self, container_id: &str) -> Result<(), String> {
+        // Remove from memory
+        self.containers.remove(container_id);
+
+        // Remove from etcd
+        if let Err(e) = crate::etcd_storage::delete_container_info(container_id).await {
+            eprintln!(
+                "[ETCD] Warning: Failed to delete ContainerInfo from etcd: {}",
+                e
+            );
+        }
+
+        println!("[DataStore] Removed container info: {}", container_id);
+        Ok(())
+    }
+
+    /// Load all containers from etcd into memory (useful for initialization)
+    pub async fn load_containers_from_etcd(&mut self) -> Result<(), String> {
+        match crate::etcd_storage::get_all_containers().await {
+            Ok(containers) => {
+                for container in containers {
+                    self.containers.insert(container.id.clone(), container);
+                }
+                println!(
+                    "[DataStore] Loaded {} containers from etcd",
+                    self.containers.len()
+                );
+                Ok(())
+            }
+            Err(e) => {
+                eprintln!("[ETCD] Warning: Failed to load containers from etcd: {}", e);
+                Ok(()) // Don't fail initialization
+            }
+        }
+    }
+
     pub fn get_node_info(&self, node_name: &str) -> Option<&NodeInfo> {
         self.nodes.get(node_name)
     }
@@ -231,6 +363,40 @@ impl DataStore {
     pub fn get_all_boards(&self) -> &HashMap<String, BoardInfo> {
         &self.boards
     }
+
+    /// ADD THIS METHOD for cleanup with etcd deletion
+    pub async fn cleanup_node_containers(
+        &mut self,
+        node_name: &str,
+        current_containers: &[String],
+    ) {
+        let containers_to_remove: Vec<String> = self
+            .container_node_mapping
+            .iter()
+            .filter(|(_, mapped_node)| *mapped_node == node_name)
+            .filter(|(container_id, _)| !current_containers.contains(container_id))
+            .map(|(container_id, _)| container_id.clone())
+            .collect();
+
+        for container_id in containers_to_remove {
+            // Remove from memory
+            self.containers.remove(&container_id);
+            self.container_node_mapping.remove(&container_id);
+
+            // Remove from etcd
+            if let Err(e) = crate::etcd_storage::delete_container_info(&container_id).await {
+                eprintln!(
+                    "[ETCD] Warning: Failed to delete container {} from etcd: {}",
+                    container_id, e
+                );
+            }
+
+            println!(
+                "[DataStore] Removed obsolete container {} from node {}",
+                container_id, node_name
+            );
+        }
+    }
 }
 
 impl SocInfo {