mecha10-node-runner 0.6.3

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

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

/// Restart policy type
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum PolicyType {
    /// Never restart on failure
    Never,
    /// Restart on failure with exponential backoff up to max_retries
    OnFailure,
    /// Restart indefinitely with constant backoff
    Always,
}

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

/// Result of monitoring a node process
#[derive(Debug)]
enum MonitorResult {
    /// Process exited normally or crashed - check restart policy
    Exited(std::process::ExitStatus),
    /// Intentional shutdown via signal (SIGTERM/SIGINT) - never restart
    IntentionalShutdown,
    /// Resource limit exceeded - may restart based on policy
    ResourceExceeded(String),
}

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

impl Supervisor {
    /// 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 {
        Self {
            restart_policy,
            health_interval,
            resource_limits,
        }
    }

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

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

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

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

            // Handle the monitoring result
            match monitor_result {
                MonitorResult::IntentionalShutdown => {
                    // User requested shutdown (SIGTERM/SIGINT) - never restart
                    info!(
                        "Node '{}' received shutdown signal - stopping without restart",
                        node_name
                    );
                    break;
                }
                MonitorResult::ResourceExceeded(resource) => {
                    // Resource limit exceeded - treat as failure for restart policy
                    error!("Node '{}' terminated due to {} limit exceeded", node_name, resource);
                    // Create a fake "failure" status for restart policy check
                    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())
                    );

                    // Check if should restart based on exit status
                    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;
                    }
                }
            }

            // Calculate exponential backoff
            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());

        // On Unix: Create new process group so killing node-runner kills all children
        #[cfg(unix)]
        {
            cmd.process_group(0); // 0 = create new process group with same ID as child PID
        }

        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> {
        // Setup resource monitoring if limits are configured
        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
        };

        // On Unix: Setup signal handlers to forward termination signals to child
        // On Windows: Signals work differently, child.kill() is sufficient
        #[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 {
            // Health monitoring enabled
            let pid = child.id();
            let mut last_status = HealthStatus::Running;

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

                    // Forward SIGTERM to child - intentional shutdown, don't restart
                    _ = sigterm.recv() => {
                        info!("Received SIGTERM, forwarding to child process");
                        let _ = child.kill().await;
                        let _ = child.wait().await; // Wait for child to exit
                        return Ok(MonitorResult::IntentionalShutdown);
                    }

                    // Forward SIGINT to child - intentional shutdown, don't restart
                    _ = sigint.recv() => {
                        info!("Received SIGINT, forwarding to child process");
                        let _ = child.kill().await;
                        let _ = child.wait().await; // Wait for child to exit
                        return Ok(MonitorResult::IntentionalShutdown);
                    }

                    // Periodic health check
                    _ = sleep(interval) => {
                        // Check process health
                        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;
                        }

                        // Check resource limits
                        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 => {
                                    // Process already exited, wait() will handle it
                                }
                                ResourceStatus::Ok => {
                                    // Resources within limits
                                }
                            }
                        }

                        // If health check failed, the wait() will return soon
                    }
                }
            }
        } else {
            // No health monitoring, but check resource limits periodically if configured
            if let Some(mut monitor) = resource_monitor {
                let check_interval = Duration::from_secs(5); // Check every 5 seconds without health monitoring

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

                        // Forward SIGTERM to child - intentional shutdown, don't restart
                        _ = sigterm.recv() => {
                            info!("Received SIGTERM, forwarding to child process");
                            let _ = child.kill().await;
                            let _ = child.wait().await;
                            return Ok(MonitorResult::IntentionalShutdown);
                        }

                        // Forward SIGINT to child - intentional shutdown, don't restart
                        _ = sigint.recv() => {
                            info!("Received SIGINT, forwarding to child process");
                            let _ = child.kill().await;
                            let _ = child.wait().await;
                            return Ok(MonitorResult::IntentionalShutdown);
                        }

                        // Periodic resource check
                        _ = 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 => {
                                    // Process already exited
                                }
                                ResourceStatus::Ok => {
                                    // Resources within limits
                                }
                            }
                        }
                    }
                }
            } else {
                // No monitoring at all - wait for process or termination signal
                tokio::select! {
                    // Wait for process to exit
                    result = child.wait() => {
                        let status = result.context("Failed to wait for node process")?;
                        Ok(MonitorResult::Exited(status))
                    }

                    // Forward SIGTERM to child - intentional shutdown, don't restart
                    _ = sigterm.recv() => {
                        info!("Received SIGTERM, forwarding to child process");
                        let _ = child.kill().await;
                        let _ = child.wait().await;
                        Ok(MonitorResult::IntentionalShutdown)
                    }

                    // Forward SIGINT to child - intentional shutdown, don't restart
                    _ = 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> {
        // Windows: No signal forwarding needed, child.kill() is sufficient
        // Note: Windows doesn't have SIGTERM/SIGINT handling like Unix
        // Process termination via Task Manager or similar will just kill the process
        if let Some(interval) = self.health_interval {
            // Health monitoring enabled
            let pid = child.id();
            let mut last_status = HealthStatus::Running;

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

                    // Periodic health check
                    _ = sleep(interval) => {
                        // Check process health
                        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;
                        }

                        // Check resource limits if configured
                        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 => {
                                    // Process already exited, wait() will handle it
                                }
                                ResourceStatus::Ok => {
                                    // Resources within limits
                                }
                            }
                        }

                        // If health check failed, the wait() will return soon
                    }
                }
            }
        } else {
            // No health monitoring, but check resource limits periodically if configured
            if let Some(mut monitor) = resource_monitor {
                let check_interval = Duration::from_secs(5);

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

                        // Periodic resource check
                        _ = 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 => {
                                    // Process already exited
                                }
                                ResourceStatus::Ok => {
                                    // Resources within limits
                                }
                            }
                        }
                    }
                }
            } else {
                // No monitoring at all, just wait for process
                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 => {
            // Restart only on failure and within retry limit
            !exit_status.success() && attempt < policy.max_retries
        }
        PolicyType::Always => {
            // Always restart regardless of exit status
            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 => {
            // Restart only within retry limit
            attempt < policy.max_retries
        }
        PolicyType::Always => {
            // Always restart
            true
        }
    }
}

/// Calculate exponential backoff duration
fn calculate_backoff(policy: &RestartPolicy, attempt: u32) -> Duration {
    if policy.policy == PolicyType::Always {
        // Constant backoff for "always" policy
        policy.backoff
    } else {
        // Exponential backoff: backoff * 2^attempt
        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;

        // Use kill(pid, 0) to check if process exists without sending a signal
        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};

        // Windows: use sysinfo to check if process exists
        let mut system = System::new();
        system.refresh_processes();

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