mecha10-node-runner 0.8.0

//! Node supervision with restart policies and health monitoring
//!
//! This module provides process supervision for node binaries, including:
//! - Restart policies (never, on-failure, always)
//! - Exponential backoff for restarts
//! - Process health monitoring
//! - Resource limit enforcement
//! - Detailed logging of lifecycle events

use crate::limits::{ResourceLimits, ResourceMonitor, ResourceStatus};
use anyhow::{Context, Result};
use std::collections::HashMap;
use std::process::Stdio;
use std::time::Duration;
use tokio::process::{Child, Command};
use tokio::time::sleep;
use tracing::{error, info, warn};

#[derive(Debug, Clone)]
pub struct RestartPolicy {
    pub policy: PolicyType,
    pub max_retries: u32,
    pub backoff: Duration,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum PolicyType {
    Never,
    OnFailure,
    Always,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum HealthStatus {
    Running,
    Stopped,
}

#[derive(Debug)]
enum MonitorResult {
    Exited(std::process::ExitStatus),
    IntentionalShutdown,
    ResourceExceeded(String),
}

pub struct Supervisor {
    restart_policy: RestartPolicy,
    health_interval: Option<Duration>,
    resource_limits: ResourceLimits,
}

impl Supervisor {
    mecha10_core::new_with_defaults!(
        /// Create a new supervisor with the given restart policy and health check interval
        pub fn new(restart_policy: RestartPolicy, health_interval: Option<Duration>, resource_limits: ResourceLimits) -> Self
    );

    /// Supervise node execution with restart and health monitoring
    pub async fn supervise(
        &self,
        node_name: &str,
        binary_path: &str,
        args: &[String],
        env: HashMap<String, String>,
    ) -> Result<()> {
        info!(
            "Starting supervision for '{}' with policy: {:?}",
            node_name, self.restart_policy.policy
        );

        let mut attempt = 0;

        loop {
            if attempt > 0 {
                info!("Restart attempt {} for '{}'", attempt, node_name);
            }

            let child = self.spawn_node(binary_path, args, &env).await?;
            let pid = child.id();

            info!("Node '{}' started (PID: {:?})", node_name, pid);

            let monitor_result = self
                .monitor_with_health_checks(node_name, child)
                .await
                .context("Failed to monitor node")?;

            match monitor_result {
                MonitorResult::IntentionalShutdown => {
                    info!(
                        "Node '{}' received shutdown signal - stopping without restart",
                        node_name
                    );
                    break;
                }
                MonitorResult::ResourceExceeded(resource) => {
                    error!("Node '{}' terminated due to {} limit exceeded", node_name, resource);
                    if !should_restart_on_failure(&self.restart_policy, attempt) {
                        anyhow::bail!("Node '{}' terminated due to {} limit exceeded", node_name, resource);
                    }
                }
                MonitorResult::Exited(exit_status) => {
                    info!(
                        "Node '{}' exited with status: {}",
                        node_name,
                        exit_status
                            .code()
                            .map(|c| c.to_string())
                            .unwrap_or_else(|| "signal".to_string())
                    );

                    if !should_restart(&self.restart_policy, attempt, &exit_status) {
                        if attempt >= self.restart_policy.max_retries
                            && self.restart_policy.policy == PolicyType::OnFailure
                        {
                            error!(
                                "Node '{}' exceeded max retries ({}), giving up",
                                node_name, self.restart_policy.max_retries
                            );
                            anyhow::bail!("Node '{}' failed after {} restart attempts", node_name, attempt);
                        }
                        info!("Node '{}' supervision ended", node_name);
                        break;
                    }
                }
            }

            let backoff_duration = calculate_backoff(&self.restart_policy, attempt);
            warn!("Node '{}' will restart in {:?}", node_name, backoff_duration);
            sleep(backoff_duration).await;

            attempt += 1;
        }

        Ok(())
    }

    /// Spawn node process
    ///
    /// On Unix, spawns in a new process group to ensure cleanup of child processes
    /// when the node-runner is terminated.
    async fn spawn_node(&self, binary_path: &str, args: &[String], env: &HashMap<String, String>) -> Result<Child> {
        let mut cmd = Command::new(binary_path);
        cmd.args(args)
            .envs(env)
            .stdin(Stdio::null())
            .stdout(Stdio::inherit())
            .stderr(Stdio::inherit());

        #[cfg(unix)]
        {
            cmd.process_group(0);
        }

        let child = cmd
            .spawn()
            .with_context(|| format!("Failed to spawn node binary '{}'", binary_path))?;

        Ok(child)
    }

    /// Monitor node with periodic health checks and resource limits
    ///
    /// Returns:
    /// - `MonitorResult::Exited(status)` - process exited normally or crashed
    /// - `MonitorResult::IntentionalShutdown` - received SIGTERM/SIGINT (don't restart)
    /// - `MonitorResult::ResourceExceeded(resource)` - killed due to resource limit
    async fn monitor_with_health_checks(&self, node_name: &str, child: Child) -> Result<MonitorResult> {
        let resource_monitor = if self.resource_limits.has_limits() {
            let pid = child.id().ok_or_else(|| anyhow::anyhow!("Failed to get child PID"))?;
            Some(ResourceMonitor::new(pid, self.resource_limits.clone()))
        } else {
            None
        };

        #[cfg(unix)]
        return self.monitor_with_signals_unix(node_name, child, resource_monitor).await;

        #[cfg(not(unix))]
        return self.monitor_without_signals(node_name, child, resource_monitor).await;
    }

    #[cfg(unix)]
    async fn monitor_with_signals_unix(
        &self,
        node_name: &str,
        mut child: Child,
        mut resource_monitor: Option<ResourceMonitor>,
    ) -> Result<MonitorResult> {
        use tokio::signal::unix::{signal, SignalKind};

        let mut sigterm = signal(SignalKind::terminate())?;
        let mut sigint = signal(SignalKind::interrupt())?;

        if let Some(interval) = self.health_interval {
            let pid = child.id();
            let mut last_status = HealthStatus::Running;

            loop {
                tokio::select! {
                    result = child.wait() => {
                        let status = result.context("Failed to wait for node process")?;
                        return Ok(MonitorResult::Exited(status));
                    }

                    _ = sigterm.recv() => {
                        info!("Received SIGTERM, forwarding to child process");
                        let _ = child.kill().await;
                        let _ = child.wait().await;
                        return Ok(MonitorResult::IntentionalShutdown);
                    }

                    _ = sigint.recv() => {
                        info!("Received SIGINT, forwarding to child process");
                        let _ = child.kill().await;
                        let _ = child.wait().await;
                        return Ok(MonitorResult::IntentionalShutdown);
                    }

                    _ = sleep(interval) => {
                        let current_status = check_process_health(pid);

                        if current_status != last_status {
                            match current_status {
                                HealthStatus::Running => {
                                    info!("Node '{}' health: Running", node_name);
                                }
                                HealthStatus::Stopped => {
                                    warn!("Node '{}' health: Stopped", node_name);
                                }
                            }
                            last_status = current_status;
                        }

                        if let Some(ref mut monitor) = resource_monitor {
                            match monitor.check_resources()? {
                                ResourceStatus::Exceeded { resource, current, limit } => {
                                    error!(
                                        "Resource limit exceeded for '{}': {} = {} > {}",
                                        node_name, resource, current, limit
                                    );
                                    let _ = child.kill().await;
                                    let _ = child.wait().await;
                                    return Ok(MonitorResult::ResourceExceeded(resource.to_string()));
                                }
                                ResourceStatus::ProcessNotFound => {
                                }
                                ResourceStatus::Ok => {
                                }
                            }
                        }

                    }
                }
            }
        } else {
            if let Some(mut monitor) = resource_monitor {
                let check_interval = Duration::from_secs(5);

                loop {
                    tokio::select! {
                        result = child.wait() => {
                            let status = result.context("Failed to wait for node process")?;
                            return Ok(MonitorResult::Exited(status));
                        }

                        _ = sigterm.recv() => {
                            info!("Received SIGTERM, forwarding to child process");
                            let _ = child.kill().await;
                            let _ = child.wait().await;
                            return Ok(MonitorResult::IntentionalShutdown);
                        }

                        _ = sigint.recv() => {
                            info!("Received SIGINT, forwarding to child process");
                            let _ = child.kill().await;
                            let _ = child.wait().await;
                            return Ok(MonitorResult::IntentionalShutdown);
                        }

                        _ = sleep(check_interval) => {
                            match monitor.check_resources()? {
                                ResourceStatus::Exceeded { resource, current, limit } => {
                                    error!(
                                        "Resource limit exceeded for '{}': {} = {} > {}",
                                        node_name, resource, current, limit
                                    );
                                    let _ = child.kill().await;
                                    let _ = child.wait().await;
                                    return Ok(MonitorResult::ResourceExceeded(resource.to_string()));
                                }
                                ResourceStatus::ProcessNotFound => {
                                }
                                ResourceStatus::Ok => {
                                }
                            }
                        }
                    }
                }
            } else {
                tokio::select! {
                    result = child.wait() => {
                        let status = result.context("Failed to wait for node process")?;
                        Ok(MonitorResult::Exited(status))
                    }

                    _ = sigterm.recv() => {
                        info!("Received SIGTERM, forwarding to child process");
                        let _ = child.kill().await;
                        let _ = child.wait().await;
                        Ok(MonitorResult::IntentionalShutdown)
                    }

                    _ = sigint.recv() => {
                        info!("Received SIGINT, forwarding to child process");
                        let _ = child.kill().await;
                        let _ = child.wait().await;
                        Ok(MonitorResult::IntentionalShutdown)
                    }
                }
            }
        }
    }

    #[cfg(not(unix))]
    async fn monitor_without_signals(
        &self,
        node_name: &str,
        mut child: Child,
        mut resource_monitor: Option<ResourceMonitor>,
    ) -> Result<MonitorResult> {
        if let Some(interval) = self.health_interval {
            let pid = child.id();
            let mut last_status = HealthStatus::Running;

            loop {
                tokio::select! {
                    result = child.wait() => {
                        let status = result.context("Failed to wait for node process")?;
                        return Ok(MonitorResult::Exited(status));
                    }

                    _ = sleep(interval) => {
                        let current_status = check_process_health(pid);

                        if current_status != last_status {
                            match current_status {
                                HealthStatus::Running => {
                                    info!("Node '{}' is running", node_name);
                                }
                                HealthStatus::Stopped => {
                                    warn!("Node '{}' stopped", node_name);
                                }
                            }
                            last_status = current_status;
                        }

                        if let Some(ref mut monitor) = resource_monitor {
                            match monitor.check_resources()? {
                                ResourceStatus::Exceeded { resource, current, limit } => {
                                    error!(
                                        "Resource limit exceeded for '{}': {} = {} > {}",
                                        node_name, resource, current, limit
                                    );
                                    let _ = child.kill().await;
                                    let _ = child.wait().await;
                                    return Ok(MonitorResult::ResourceExceeded(resource.to_string()));
                                }
                                ResourceStatus::ProcessNotFound => {
                                }
                                ResourceStatus::Ok => {
                                }
                            }
                        }

                    }
                }
            }
        } else {
            if let Some(mut monitor) = resource_monitor {
                let check_interval = Duration::from_secs(5);

                loop {
                    tokio::select! {
                        result = child.wait() => {
                            let status = result.context("Failed to wait for node process")?;
                            return Ok(MonitorResult::Exited(status));
                        }

                        _ = sleep(check_interval) => {
                            match monitor.check_resources()? {
                                ResourceStatus::Exceeded { resource, current, limit } => {
                                    error!(
                                        "Resource limit exceeded for '{}': {} = {} > {}",
                                        node_name, resource, current, limit
                                    );
                                    let _ = child.kill().await;
                                    let _ = child.wait().await;
                                    return Ok(MonitorResult::ResourceExceeded(resource.to_string()));
                                }
                                ResourceStatus::ProcessNotFound => {
                                }
                                ResourceStatus::Ok => {
                                }
                            }
                        }
                    }
                }
            } else {
                let status = child.wait().await.context("Failed to wait for node process")?;
                Ok(MonitorResult::Exited(status))
            }
        }
    }
}

/// Check if process should restart based on policy and exit status
fn should_restart(policy: &RestartPolicy, attempt: u32, exit_status: &std::process::ExitStatus) -> bool {
    match policy.policy {
        PolicyType::Never => false,
        PolicyType::OnFailure => !exit_status.success() && attempt < policy.max_retries,
        PolicyType::Always => true,
    }
}

/// Check if process should restart on failure (for resource exceeded cases)
///
/// This is similar to should_restart but doesn't need an exit status -
/// used when we know it's a failure condition (resource exceeded).
fn should_restart_on_failure(policy: &RestartPolicy, attempt: u32) -> bool {
    match policy.policy {
        PolicyType::Never => false,
        PolicyType::OnFailure => attempt < policy.max_retries,
        PolicyType::Always => true,
    }
}

/// Calculate exponential backoff duration
fn calculate_backoff(policy: &RestartPolicy, attempt: u32) -> Duration {
    if policy.policy == PolicyType::Always {
        policy.backoff
    } else {
        let multiplier = 2_u64.saturating_pow(attempt);
        policy.backoff.saturating_mul(multiplier.try_into().unwrap_or(1))
    }
}

/// Check if process is still running
fn check_process_health(pid: Option<u32>) -> HealthStatus {
    let Some(pid) = pid else {
        return HealthStatus::Stopped;
    };

    #[cfg(unix)]
    {
        use nix::sys::signal::{kill, Signal};
        use nix::unistd::Pid;

        match kill(Pid::from_raw(pid as i32), Some(Signal::SIGCONT)) {
            Ok(_) => HealthStatus::Running,
            Err(_) => HealthStatus::Stopped,
        }
    }

    #[cfg(not(unix))]
    {
        use sysinfo::{Pid as SysInfoPid, System};

        let mut system = System::new();
        system.refresh_processes();

        if system.process(SysInfoPid::from(pid as usize)).is_some() {
            HealthStatus::Running
        } else {
            HealthStatus::Stopped
        }
    }
}