imp(orchestrator): cleanup mutex handling to proper async (#530)

* cleanup mutex handling to proper async

Author: JannikSt

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

@@ -335,6 +335,7 @@ fn generate_file_name(template: &str, original_name: &str) -> String {
 mod tests {
 
     use std::collections::HashMap;
+    use std::sync::Arc;
 
     use super::*;
     use crate::plugins::StatusUpdatePlugin;
@@ -580,13 +581,14 @@ mod tests {
             compute_requirements: None,
         };
 
-        let plugin = NodeGroupsPlugin::new(
+        let plugin = Arc::new(NodeGroupsPlugin::new(
             vec![config],
             app_state.redis_store.clone(),
             app_state.store_context.clone(),
             None,
             None,
-        );
+        ));
+        let _ = plugin.clone().register_observer().await;
 
         let _ = plugin
             .handle_status_change(&node, &NodeStatus::Healthy)

crates/orchestrator/src/main.rs

@@ -259,7 +259,14 @@ async fn main() -> Result<()> {
 
                 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));
+                let group_plugin_arc = Arc::new(group_plugin_for_server);
+
+                // Register the plugin as a task observer
+                if let Err(e) = group_plugin_arc.clone().register_observer().await {
+                    error!("Failed to register node groups plugin as observer: {}", e);
+                }
+
+                node_groups_plugin = Some(group_plugin_arc);
                 scheduler_plugins.push(Box::new(group_plugin));
                 status_update_plugins.push(Box::new(status_group_plugin));
                 info!("Plugin: Node group plugin initialized");

crates/orchestrator/src/plugins/node_groups/mod.rs

@@ -166,22 +166,24 @@ impl NodeGroupsPlugin {
             }
         });
 
-        let plugin = Self {
+        Self {
             configuration_templates: sorted_configs,
             store,
             store_context,
             node_groups_heartbeats,
             webhook_plugins,
             task_switching_policy,
             proximity_optimization_policy,
-        };
+        }
+    }
 
-        plugin
-            .store_context
+    /// Register this plugin as a task observer (async)
+    pub async fn register_observer(self: Arc<Self>) -> Result<()> {
+        self.store_context
             .task_store
-            .add_observer(Arc::new(plugin.clone()));
-
-        plugin
+            .add_observer(self.clone())
+            .await;
+        Ok(())
     }
 
     /// Check if a node is compatible with a configuration's compute requirements