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 mecha10_core::rwlock_collections::RwLockMapExt;
use std::sync::Arc;
use tokio::task::JoinHandle;

#[derive(Clone, Debug, PartialEq)]
pub enum NodeStatus {
    Starting,

    Running,

    ShuttingDown,

    Stopped,

    Failed { reason: String },

    Restarting { attempt: usize },
}

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(())
    }
}

pub struct Supervisor {
    health_checker: Arc<HealthChecker>,
    shutdown: ShutdownHandle,
    restart_policy: RestartPolicy,
    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;

            health_checker.register(node_name.clone(), HealthStatus::Healthy).await;

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

                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;

                                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;
                    }
                }
            }

            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)]
    pub async fn start_node(&self, name: &str) -> Result<()> {
        tracing::info!("Starting node: {}", name);

        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)]
    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) {
            handle.abort();

            self.health_checker.unregister(name).await;

            tracing::info!("Node {} stopped", name);
            Ok(())
        } else {
            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)]
    pub async fn get_running_nodes(&self) -> Vec<String> {
        self.nodes.cloned_keys().await
    }
}

#[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> {
        vec![]
    }
}