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

    /// 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. Once the directly-tracked child (e.g. node-runner) has
            // exited, `child.try_wait()` will keep reporting `Ok(Some(_))` since the child is
            // already reaped - so from that point on we can't use it to tell whether we're
            // still waiting on anything. Instead, once the tracked child has exited, switch to
            // polling whether the whole process *group* is empty: node-runner may exit near-
            // instantly while a grandchild it was supervising (e.g. simulator's Python MuJoCo
            // subprocess) is still alive in the same group with nothing left to signal it.
            let start = Instant::now();
            let mut tracked_child_exited = false;

            loop {
                if !tracked_child_exited {
                    match child.try_wait() {
                        Ok(Some(_)) => {
                            tracked_child_exited = true;
                        }
                        Ok(None) => {
                            // Still running.
                        }
                        Err(e) => {
                            // Error checking status - try to kill anyway
                            error!("Error checking process {} status: {}", name, e);
                            return kill_group_and_reap(pid, &mut child);
                        }
                    }
                }

                let done = if tracked_child_exited {
                    #[cfg(unix)]
                    {
                        process_group_is_empty(pid)
                    }
                    #[cfg(not(unix))]
                    {
                        true
                    }
                } else {
                    false
                };

                if done {
                    return Ok(());
                }

                if start.elapsed() >= timeout {
                    if tracked_child_exited {
                        warn!(
                            "Process {} exited but its process group did not stop within {:?}, force killing",
                            name, timeout
                        );
                    } else {
                        warn!("Process {} did not stop within {:?}, force killing", name, timeout);
                    }
                    return kill_group_and_reap(pid, &mut child);
                }

                thread::sleep(Duration::from_millis(100));
            }
        } 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();
    }
}

/// Check whether a process group is empty (no processes remain in it).
///
/// Sends signal 0 to the process group (`kill -0 -- -{pid}`), which performs no actual
/// signaling but reports via its exit status whether any process in that group still exists.
/// A successful exit status means at least one process in the group is still alive; any
/// failure (notably ESRCH - "no such process") means the group is empty.
#[cfg(unix)]
fn process_group_is_empty(pid: u32) -> bool {
    use std::process::Command;
    match Command::new("kill")
        .arg("-0")
        .arg("--")
        .arg(format!("-{}", pid))
        .output()
    {
        Ok(output) => !output.status.success(),
        // If we can't even run `kill` to check, don't block shutdown on it - treat as empty.
        Err(_) => true,
    }
}

/// Force kill the process group and reap the directly-tracked child, returning `Ok(())`.
///
/// Used by `stop_graceful`'s timeout/error paths - mirrors `force_kill`'s unconditional
/// "kill -9 the group, then kill+wait the tracked child" behavior.
#[allow(unused_variables)]
fn kill_group_and_reap(pid: u32, child: &mut Child) -> Result<()> {
    #[cfg(unix)]
    {
        use std::process::Command;
        let _ = Command::new("kill")
            .arg("-9")
            .arg("--")
            .arg(format!("-{}", pid))
            .output();
    }

    child.kill()?;
    let _ = child.wait();
    Ok(())
}

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

    /// Regression test for the bug where `stop_graceful` declared success the instant the
    /// directly-tracked process exited, without checking whether the rest of its process
    /// group (e.g. a supervised grandchild subprocess, like simulator's Python MuJoCo
    /// process) was actually gone yet.
    ///
    /// Spawns a "parent" process as the leader of its own process group (mirroring how
    /// node-runner is spawned with `process_group(0)`), which itself forks a background
    /// child that ignores SIGTERM and stays alive for a bit before exiting on its own -
    /// simulating a supervised subprocess that's slower to shut down than its supervisor.
    /// `stop_graceful` must not return until *both* have exited, not just the tracked parent.
    #[cfg(unix)]
    #[test]
    fn stop_graceful_waits_for_entire_process_group() {
        use std::os::unix::process::CommandExt;
        use std::process::{Command, Stdio};
        use std::time::Instant;

        let mut pm = ProcessManager::new();

        // The background grandchild traps SIGTERM (ignores it) and sleeps for ~300ms before
        // exiting on its own - standing in for a subprocess whose own shutdown (e.g. Redis
        // disconnect, GL context teardown) takes longer than its supervisor's exit. The
        // parent shell backgrounds it and then exits immediately, so the tracked process
        // (the parent) exits almost instantly while the grandchild lingers in the same group.
        // The outer shell also ignores SIGTERM so that it reliably forks the background
        // grandchild and exits normally, regardless of exactly when `stop_graceful` delivers
        // the signal (avoiding a race where SIGTERM kills the outer shell before it gets a
        // chance to spawn the grandchild at all).
        let child = Command::new("sh")
            .arg("-c")
            .arg(r#"trap "" TERM; sh -c 'trap "" TERM; sleep 0.3' & exit 0"#)
            .process_group(0)
            .stdin(Stdio::null())
            .stdout(Stdio::null())
            .stderr(Stdio::null())
            .spawn()
            .expect("failed to spawn test process group leader");

        pm.track("grandchild-test".to_string(), child);

        let start = Instant::now();
        let result = pm.stop_graceful("grandchild-test", Duration::from_secs(5));
        let elapsed = start.elapsed();

        assert!(result.is_ok(), "stop_graceful should succeed: {:?}", result);
        assert!(
            elapsed >= Duration::from_millis(250),
            "stop_graceful returned after {:?}, before the grandchild (~300ms) could have \
             exited - it must have declared success as soon as the tracked parent exited \
             without checking the rest of the process group",
            elapsed
        );
        // Should still be comfortably under the 5s timeout - i.e. it noticed the group
        // emptied on its own rather than waiting for a force-kill timeout.
        assert!(
            elapsed < Duration::from_secs(2),
            "stop_graceful took {:?}, suggesting it fell through to the timeout/force-kill \
             path instead of detecting the group became empty",
            elapsed
        );
    }
}