Fix: multiple group assignments (#327)

* improve redis tests, fix node group assignments

* add groups to api endpoint

Author: JannikSt

crates/orchestrator/src/api/routes/nodes.rs

@@ -11,7 +11,6 @@ use std::str::FromStr;
 use std::time::Duration;
 // Timeout for node operations in seconds
 const NODE_REQUEST_TIMEOUT: u64 = 30;
-
 async fn get_nodes(app_state: Data<AppState>) -> HttpResponse {
     let nodes = app_state.store_context.node_store.get_nodes();
 
@@ -29,11 +28,34 @@ async fn get_nodes(app_state: Data<AppState>) -> HttpResponse {
         }
     }
 
-    HttpResponse::Ok().json(json!({
+    let mut response = json!({
         "success": true,
         "nodes": nodes,
         "counts": status_counts
-    }))
+    });
+
+    // If node groups plugin exists, add group information to each node
+    if let Some(node_groups_plugin) = &app_state.node_groups_plugin {
+        let mut nodes_with_groups = Vec::new();
+
+        for node in &nodes {
+            let mut node_json = json!(node);
+
+            if let Ok(Some(group)) = node_groups_plugin.get_node_group(&node.address.to_string()) {
+                node_json["group"] = json!({
+                    "id": group.id,
+                    "size": group.nodes.len(),
+                    "created_at": group.created_at
+                });
+            }
+
+            nodes_with_groups.push(node_json);
+        }
+
+        response["nodes"] = json!(nodes_with_groups);
+    }
+
+    HttpResponse::Ok().json(response)
 }
 
 async fn restart_node_task(node_id: web::Path<String>, app_state: Data<AppState>) -> HttpResponse {

crates/orchestrator/src/api/server.rs

@@ -4,6 +4,7 @@ use crate::api::routes::task::tasks_routes;
 use crate::api::routes::{heartbeat::heartbeat_routes, metrics::metrics_routes};
 use crate::models::node::NodeStatus;
 use crate::scheduler::Scheduler;
+use crate::status_update::plugins::node_groups::NodeGroupsPlugin;
 use crate::store::core::{RedisStore, StoreContext};
 use crate::utils::loop_heartbeats::LoopHeartbeats;
 use crate::ServerMode;
@@ -30,6 +31,7 @@ pub struct AppState {
     pub contracts: Option<Arc<Contracts>>,
     pub pool_id: u32,
     pub scheduler: Scheduler,
+    pub node_groups_plugin: Option<Arc<NodeGroupsPlugin>>,
 }
 
 #[allow(clippy::too_many_arguments)]
@@ -48,6 +50,7 @@ pub async fn start_server(
     pool_id: u32,
     server_mode: ServerMode,
     scheduler: Scheduler,
+    node_groups_plugin: Option<Arc<NodeGroupsPlugin>>,
 ) -> Result<(), Error> {
     info!("Starting server at http://{}:{}", host, port);
     let app_state = Data::new(AppState {
@@ -61,6 +64,7 @@ pub async fn start_server(
         contracts,
         pool_id,
         scheduler,
+        node_groups_plugin,
     });
     let node_store = app_state.store_context.node_store.clone();
     let node_store_clone = node_store.clone();

crates/orchestrator/src/api/tests/helper.rs

@@ -52,6 +52,7 @@ pub async fn create_test_app_state() -> Data<AppState> {
         hourly_upload_limit: 12,
         redis_store: store.clone(),
         scheduler,
+        node_groups_plugin: None,
     })
 }
 

crates/orchestrator/src/main.rs

@@ -201,13 +201,16 @@ async fn main() -> Result<()> {
     let group_store_context = store_context.clone();
     let mut scheduler_plugins: Vec<Box<dyn SchedulerPlugin>> = Vec::new();
     let mut status_update_plugins: Vec<Box<dyn StatusUpdatePlugin>> = vec![];
+    let mut node_groups_plugin: Option<Arc<NodeGroupsPlugin>> = None;
 
     // Add group plugin if enabled
     if args.with_basic_group_plugin {
         let group_size: usize = args.group_size as usize;
         let group_plugin =
             NodeGroupsPlugin::new(group_size, group_size, store.clone(), group_store_context);
         let status_group_plugin = group_plugin.clone();
+        let group_plugin_for_server = group_plugin.clone();
+        node_groups_plugin = Some(Arc::new(group_plugin_for_server));
         scheduler_plugins.push(Box::new(group_plugin));
         status_update_plugins.push(Box::new(status_group_plugin));
     }
@@ -295,6 +298,7 @@ async fn main() -> Result<()> {
             compute_pool_id,
             server_mode,
             scheduler,
+            node_groups_plugin,
         ) => {
             if let Err(e) = res {
                 error!("Server error: {}", e);