mecha10_runtime/

process.rs

//! External process management for node lifecycle
//!
//! This module provides process management for external node processes (spawned as child processes),
//! complementing the Supervisor which manages in-process NodeRunner implementations.

use crate::shutdown::ShutdownHandle;
use anyhow::Result;
use std::collections::HashMap;
use std::process::Child;
use std::time::Duration;
use tracing::{debug, error, info, warn};

/// Information about a tracked process
#[derive(Debug, Clone)]
pub struct ProcessInfo {
    /// Process name
    pub name: String,
    /// Process ID
    pub pid: u32,
    /// Current status
    pub status: ProcessStatus,
}

/// Process status
#[derive(Debug, Clone, PartialEq)]
pub enum ProcessStatus {
    /// Process is running
    Running,
    /// Process exited with code
    Exited(i32),
    /// Process encountered an error
    Error,
}

/// Manages external child processes with graceful shutdown and dependency ordering
///
/// ProcessManager tracks and manages child processes spawned as external binaries.
/// It provides:
/// - Process tracking with PIDs
/// - Dependency-based shutdown ordering
/// - Graceful shutdown with configurable timeout
/// - Force kill fallback
/// - Integration with runtime's ShutdownHandle
pub struct ProcessManager {
    /// Tracked child processes
    processes: HashMap<String, Child>,
    /// Dependency graph: node -> dependencies
    dependency_graph: HashMap<String, Vec<String>>,
    /// Shutdown handle for coordinating graceful shutdown
    shutdown_handle: Option<ShutdownHandle>,
    /// Default shutdown timeout
    shutdown_timeout: Duration,
}

impl ProcessManager {
    /// Create a new process manager
    pub fn new() -> Self {
        Self {
            processes: HashMap::new(),
            dependency_graph: HashMap::new(),
            shutdown_handle: None,
            shutdown_timeout: Duration::from_secs(10),
        }
    }

    /// Create a new process manager with shutdown handle
    pub fn with_shutdown_handle(shutdown_handle: ShutdownHandle) -> Self {
        Self {
            processes: HashMap::new(),
            dependency_graph: HashMap::new(),
            shutdown_handle: Some(shutdown_handle),
            shutdown_timeout: Duration::from_secs(10),
        }
    }

    /// Set the default shutdown timeout
    pub fn set_shutdown_timeout(&mut self, timeout: Duration) {
        self.shutdown_timeout = timeout;
    }

    /// Add a process to track
    ///
    /// # Arguments
    ///
    /// * `name` - Process name/identifier
    /// * `child` - Child process handle
    pub fn track(&mut self, name: String, child: Child) {
        let pid = child.id();
        debug!("Tracking process: {} (PID: {})", name, pid);
        self.processes.insert(name, child);
    }

    /// Remove and return a process by name
    pub fn remove(&mut self, name: &str) -> Option<Child> {
        self.processes.remove(name)
    }

    /// Add a dependency relationship
    ///
    /// Specifies that `node` depends on `dependency` and should be stopped before it.
    ///
    /// # Arguments
    ///
    /// * `node` - The node that depends on another
    /// * `dependency` - The node that is depended upon
    pub fn add_dependency(&mut self, node: String, dependency: String) {
        self.dependency_graph.entry(node).or_default().push(dependency);
    }

    /// Get list of all tracked process names
    pub fn list(&self) -> Vec<String> {
        self.processes.keys().cloned().collect()
    }

    /// Get the number of tracked processes
    pub fn len(&self) -> usize {
        self.processes.len()
    }

    /// Check if any processes are being tracked
    pub fn is_empty(&self) -> bool {
        self.processes.is_empty()
    }

    /// Get status of all processes
    pub fn status_all(&mut self) -> HashMap<String, ProcessStatus> {
        let mut status_map = HashMap::new();
        let mut to_remove = Vec::new();

        for (name, child) in self.processes.iter_mut() {
            let status = match child.try_wait() {
                Ok(Some(exit_status)) => {
                    to_remove.push(name.clone());
                    ProcessStatus::Exited(exit_status.code().unwrap_or(-1))
                }
                Ok(None) => ProcessStatus::Running,
                Err(_) => {
                    to_remove.push(name.clone());
                    ProcessStatus::Error
                }
            };
            status_map.insert(name.clone(), status);
        }

        // Remove exited/errored processes
        for name in to_remove {
            self.processes.remove(&name);
        }

        status_map
    }

    /// Get information about all tracked processes
    pub fn processes(&mut self) -> Vec<ProcessInfo> {
        let mut processes = Vec::new();
        let mut to_remove = Vec::new();

        for (name, child) in self.processes.iter_mut() {
            let pid = child.id();
            let status = match child.try_wait() {
                Ok(Some(exit_status)) => {
                    to_remove.push(name.clone());
                    ProcessStatus::Exited(exit_status.code().unwrap_or(-1))
                }
                Ok(None) => ProcessStatus::Running,
                Err(_) => {
                    to_remove.push(name.clone());
                    ProcessStatus::Error
                }
            };

            processes.push(ProcessInfo {
                name: name.clone(),
                pid,
                status,
            });
        }

        // Remove exited/errored processes
        for name in to_remove {
            self.processes.remove(&name);
        }

        processes
    }

    /// Calculate shutdown order based on dependency graph
    ///
    /// Returns nodes in the order they should be stopped (dependents before dependencies).
    pub fn shutdown_order(&self) -> Vec<String> {
        let mut order = Vec::new();
        let mut visited = std::collections::HashSet::new();

        // Build reverse graph (dependents)
        let mut dependents: HashMap<String, Vec<String>> = HashMap::new();
        for (node, deps) in &self.dependency_graph {
            for dep in deps {
                dependents.entry(dep.clone()).or_default().push(node.clone());
            }
        }

        // Get all tracked nodes
        let all_nodes: Vec<String> = self.list();

        // Visit each node (prioritize nodes with dependents)
        for node in &all_nodes {
            if !visited.contains(node) {
                visit_for_shutdown(node, &dependents, &mut visited, &mut order);
            }
        }

        order
    }

    /// Stop a specific process gracefully with timeout
    ///
    /// # Arguments
    ///
    /// * `name` - Name of the process to stop
    /// * `timeout` - How long to wait for graceful shutdown before force killing
    ///
    /// # Returns
    ///
    /// Ok(()) if stopped successfully (graceful or force), Err if not found or error
    pub fn stop_graceful(&mut self, name: &str, timeout: Duration) -> Result<()> {
        use std::thread;
        use std::time::Instant;

        if let Some(mut child) = self.remove(name) {
            let pid = child.id();
            debug!("Stopping process: {} (PID: {})", name, pid);

            // Try graceful shutdown (SIGTERM on Unix)
            #[cfg(unix)]
            {
                use std::process::Command;
                // Send SIGTERM to the process group (negative PID kills the group)
                // This ensures we kill the node-runner AND all its child nodes
                let _ = Command::new("kill").arg("--").arg(format!("-{}", pid)).output();

                // Also send to the process itself as fallback
                let _ = Command::new("kill").arg(pid.to_string()).output();
            }

            #[cfg(not(unix))]
            {
                // Windows: terminate immediately
                child.kill()?;
            }

            // Wait with timeout
            let start = Instant::now();
            loop {
                match child.try_wait() {
                    Ok(Some(_)) => {
                        // Process exited gracefully
                        info!("Process {} stopped gracefully", name);
                        return Ok(());
                    }
                    Ok(None) => {
                        // Still running
                        if start.elapsed() >= timeout {
                            // Timeout - force kill the process group
                            warn!("Process {} did not stop within {:?}, force killing", name, timeout);

                            #[cfg(unix)]
                            {
                                use std::process::Command;
                                // Force kill the entire process group
                                let _ = Command::new("kill")
                                    .arg("-9")
                                    .arg("--")
                                    .arg(format!("-{}", pid))
                                    .output();
                            }

                            child.kill()?;
                            let _ = child.wait();
                            return Ok(());
                        }
                        thread::sleep(Duration::from_millis(100));
                    }
                    Err(e) => {
                        // Error checking status - try to kill anyway
                        error!("Error checking process {} status: {}", name, e);

                        #[cfg(unix)]
                        {
                            use std::process::Command;
                            let _ = Command::new("kill")
                                .arg("-9")
                                .arg("--")
                                .arg(format!("-{}", pid))
                                .output();
                        }

                        child.kill()?;
                        let _ = child.wait();
                        return Ok(());
                    }
                }
            }
        } else {
            Err(anyhow::anyhow!("Process '{}' not found", name))
        }
    }

    /// Force kill a process immediately
    pub fn force_kill(&mut self, name: &str) -> Result<()> {
        if let Some(mut child) = self.remove(name) {
            let pid = child.id();

            #[cfg(unix)]
            {
                use std::process::Command;
                // Force kill the entire process group
                let _ = Command::new("kill")
                    .arg("-9")
                    .arg("--")
                    .arg(format!("-{}", pid))
                    .output();
            }

            child.kill()?;
            let _ = child.wait();
            info!("Process {} force killed", name);
            Ok(())
        } else {
            Err(anyhow::anyhow!("Process '{}' not found", name))
        }
    }

    /// Stop all processes in dependency order with graceful shutdown
    pub fn shutdown_all(&mut self) {
        if self.is_empty() {
            debug!("No processes to shut down");
            return;
        }

        info!("Shutting down all processes in dependency order");

        // Trigger shutdown signal if available
        if let Some(ref shutdown_handle) = self.shutdown_handle {
            shutdown_handle.shutdown();
        }

        let shutdown_order = self.shutdown_order();

        for node in shutdown_order {
            info!("Stopping {} with {:?} timeout", node, self.shutdown_timeout);

            if let Err(e) = self.stop_graceful(&node, self.shutdown_timeout) {
                warn!("Failed to stop {}: {}, attempting force kill", node, e);
                if let Err(kill_err) = self.force_kill(&node) {
                    error!("Failed to force kill {}: {}", node, kill_err);
                }
            }
        }

        info!("All processes stopped");
    }

    /// Check if a specific process is running
    pub fn is_running(&mut self, name: &str) -> bool {
        if let Some(child) = self.processes.get_mut(name) {
            matches!(child.try_wait(), Ok(None))
        } else {
            false
        }
    }
}

impl Default for ProcessManager {
    fn default() -> Self {
        Self::new()
    }
}

impl Drop for ProcessManager {
    fn drop(&mut self) {
        // Ensure all processes are cleaned up
        self.shutdown_all();
    }
}

/// Recursively visit nodes for shutdown ordering
///
/// Helper function for topological sort (DFS-based)
fn visit_for_shutdown(
    node: &str,
    dependents: &HashMap<String, Vec<String>>,
    visited: &mut std::collections::HashSet<String>,
    order: &mut Vec<String>,
) {
    visited.insert(node.to_string());

    // Visit dependents first (they must be stopped before this node)
    if let Some(deps) = dependents.get(node) {
        for dep in deps {
            if !visited.contains(dep) {
                visit_for_shutdown(dep, dependents, visited, order);
            }
        }
    }

    // Add this node to the order
    order.push(node.to_string());
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_dependency_ordering() {
        let mut pm = ProcessManager::new();

        // Setup: planner -> motor, camera -> planner
        pm.add_dependency("planner".to_string(), "motor".to_string());
        pm.add_dependency("camera".to_string(), "planner".to_string());

        // Simulate tracking processes
        pm.processes.insert("motor".to_string(), mock_child());
        pm.processes.insert("planner".to_string(), mock_child());
        pm.processes.insert("camera".to_string(), mock_child());

        let order = pm.shutdown_order();

        // Camera should stop before planner, planner before motor
        let camera_idx = order.iter().position(|n| n == "camera").unwrap();
        let planner_idx = order.iter().position(|n| n == "planner").unwrap();
        let motor_idx = order.iter().position(|n| n == "motor").unwrap();

        assert!(camera_idx < planner_idx);
        assert!(planner_idx < motor_idx);
    }

    #[cfg(unix)]
    fn mock_child() -> Child {
        use std::process::{Command, Stdio};
        Command::new("sleep")
            .arg("1000")
            .stdin(Stdio::null())
            .stdout(Stdio::null())
            .stderr(Stdio::null())
            .spawn()
            .unwrap()
    }

    #[cfg(not(unix))]
    fn mock_child() -> Child {
        use std::process::{Command, Stdio};
        Command::new("timeout")
            .arg("/t")
            .arg("1000")
            .stdin(Stdio::null())
            .stdout(Stdio::null())
            .stderr(Stdio::null())
            .spawn()
            .unwrap()
    }
}