mecha10_runtime/

supervisor.rs

//! Node supervision and lifecycle management

use crate::config::RestartPolicy;
use crate::health::{HealthChecker, HealthStatus};
use crate::node::NodeRunner;
use crate::shutdown::ShutdownHandle;
use anyhow::Result;
use std::sync::Arc;
use tokio::task::JoinHandle;

/// Status of a supervised node
#[derive(Clone, Debug, PartialEq)]
pub enum NodeStatus {
    /// Node is starting up
    Starting,

    /// Node is running normally
    Running,

    /// Node is shutting down
    ShuttingDown,

    /// Node has stopped
    Stopped,

    /// Node failed with error
    Failed { reason: String },

    /// Node is restarting
    Restarting { attempt: usize },
}

/// Handle to a supervised node
pub struct NodeHandle {
    name: String,
    task: JoinHandle<()>,
}

impl NodeHandle {
    /// Get the node name
    pub fn name(&self) -> &str {
        &self.name
    }

    /// Check if the node task is finished
    pub fn is_finished(&self) -> bool {
        self.task.is_finished()
    }

    /// Abort the node task
    pub fn abort(&self) {
        self.task.abort();
    }

    /// Wait for the node to finish
    pub async fn wait(self) -> Result<()> {
        self.task.await?;
        Ok(())
    }
}

/// Supervisor for managing node lifecycle
pub struct Supervisor {
    health_checker: Arc<HealthChecker>,
    shutdown: ShutdownHandle,
    restart_policy: RestartPolicy,
    /// Tracked node handles by name
    nodes: Arc<tokio::sync::RwLock<std::collections::HashMap<String, NodeHandle>>>,
}

impl Supervisor {
    /// Create a new supervisor
    pub fn new(health_checker: Arc<HealthChecker>, shutdown: ShutdownHandle, restart_policy: RestartPolicy) -> Self {
        Self {
            health_checker,
            shutdown,
            restart_policy,
            nodes: Arc::new(tokio::sync::RwLock::new(std::collections::HashMap::new())),
        }
    }

    /// Launch and supervise a node
    pub fn launch_node(&self, mut node: Box<dyn NodeRunner>) -> NodeHandle {
        let name = node.name().to_string();
        let health_checker = self.health_checker.clone();
        let shutdown = self.shutdown.clone();
        let restart_policy = self.restart_policy.clone();

        tracing::info!("Launching node: {}", name);

        let task = tokio::spawn(async move {
            let node_name = node.name().to_string();
            let mut attempt = 0;

            // Register with health checker
            health_checker.register(node_name.clone(), HealthStatus::Healthy).await;

            loop {
                tracing::debug!("Starting node: {} (attempt {})", node_name, attempt);

                // Run the node with shutdown signal
                let mut shutdown_rx = shutdown.subscribe();
                let run_future = node.run();

                tokio::select! {
                    result = run_future => {
                        match result {
                            Ok(()) => {
                                tracing::info!("Node {} completed successfully", node_name);
                                health_checker
                                    .update(&node_name, HealthStatus::Healthy)
                                    .await;
                                break;
                            }
                            Err(e) => {
                                tracing::error!("Node {} failed: {}", node_name, e);
                                health_checker
                                    .update(
                                        &node_name,
                                        HealthStatus::Unhealthy {
                                            reason: e.to_string(),
                                        },
                                    )
                                    .await;

                                // Check if should restart
                                if restart_policy.should_restart(attempt) {
                                    let backoff = restart_policy.backoff_duration(attempt);
                                    tracing::info!(
                                        "Restarting node {} in {:?} (attempt {})",
                                        node_name,
                                        backoff,
                                        attempt + 1
                                    );
                                    tokio::time::sleep(backoff).await;
                                    attempt += 1;
                                    continue;
                                } else {
                                    tracing::error!(
                                        "Node {} exhausted restart attempts, giving up",
                                        node_name
                                    );
                                    break;
                                }
                            }
                        }
                    }
                    _ = shutdown_rx.recv() => {
                        tracing::info!("Shutdown signal received, stopping node: {}", node_name);
                        if let Err(e) = node.shutdown().await {
                            tracing::error!("Error during node {} shutdown: {}", node_name, e);
                        }
                        health_checker
                            .update(&node_name, HealthStatus::Healthy)
                            .await;
                        break;
                    }
                }
            }

            // Unregister from health checker
            health_checker.unregister(&node_name).await;
            tracing::debug!("Node {} supervision ended", node_name);
        });

        NodeHandle { name, task }
    }

    /// Launch multiple nodes
    pub fn launch_nodes(&self, nodes: Vec<Box<dyn NodeRunner>>) -> Vec<NodeHandle> {
        nodes.into_iter().map(|node| self.launch_node(node)).collect()
    }

    /// Get the health checker
    pub fn health_checker(&self) -> &Arc<HealthChecker> {
        &self.health_checker
    }

    /// Get the shutdown handle
    pub fn shutdown(&self) -> &ShutdownHandle {
        &self.shutdown
    }

    /// Start a node by name (for lifecycle management)
    ///
    /// Note: This is a stub implementation for the lifecycle system.
    /// Full implementation requires node registry or factory pattern.
    ///
    /// # Arguments
    ///
    /// * `name` - Name of the node to start
    ///
    /// # Errors
    ///
    /// Returns error if node cannot be found or started
    #[allow(dead_code)] // Will be used in Phase 3
    pub async fn start_node(&self, name: &str) -> Result<()> {
        tracing::info!("Starting node: {}", name);

        // TODO: Implement node registry/factory to create nodes by name
        // For now, this is a placeholder that will be implemented when
        // we integrate with the full runtime system

        // Register with health checker as a placeholder
        self.health_checker
            .register(name.to_string(), HealthStatus::Healthy)
            .await;

        tracing::debug!("Node {} registered (stub)", name);
        Ok(())
    }

    /// Stop a node by name (for lifecycle management)
    ///
    /// # Arguments
    ///
    /// * `name` - Name of the node to stop
    ///
    /// # Errors
    ///
    /// Returns error if node cannot be found or stopped
    #[allow(dead_code)] // Will be used in Phase 3
    pub async fn stop_node(&self, name: &str) -> Result<()> {
        tracing::info!("Stopping node: {}", name);

        let mut nodes = self.nodes.write().await;

        if let Some(handle) = nodes.remove(name) {
            // Abort the node task
            handle.abort();

            // Unregister from health checker
            self.health_checker.unregister(name).await;

            tracing::info!("Node {} stopped", name);
            Ok(())
        } else {
            // Node not found, but unregister anyway in case it was tracked elsewhere
            self.health_checker.unregister(name).await;
            tracing::warn!(
                "Node {} not found in tracked nodes, unregistered from health checker",
                name
            );
            Ok(())
        }
    }

    /// Get list of currently running nodes
    ///
    /// Returns the names of all nodes being supervised
    #[allow(dead_code)] // Will be used in Phase 3
    pub async fn get_running_nodes(&self) -> Vec<String> {
        let nodes = self.nodes.read().await;
        nodes.keys().cloned().collect()
    }
}

/// Implement the SupervisorTrait for Supervisor
#[async_trait::async_trait]
impl crate::lifecycle::SupervisorTrait for Supervisor {
    async fn start_node(&self, name: &str) -> Result<()> {
        self.start_node(name).await
    }

    async fn stop_node(&self, name: &str) -> Result<()> {
        self.stop_node(name).await
    }

    fn get_running_nodes(&self) -> Vec<String> {
        // Note: This is synchronous but SupervisorTrait requires sync
        // In real implementation, we'd need to use a different approach
        // For now, return empty vec as placeholder
        vec![]
    }
}