warpdrive_proxy/process/

supervisor.rs

//! Process supervisor for managing upstream application servers
//!
//! This module provides process supervision capabilities to spawn and monitor
//! application servers (e.g., Rails, Node.js). It handles:
//! - Process spawning with environment setup
//! - Signal forwarding (SIGTERM/SIGINT)
//! - Graceful shutdown with timeout
//! - Automatic restart on crash (optional)
//!
//! # Example
//!
//! ```no_run
//! use warpdrive::process::ProcessSupervisor;
//! use anyhow::Result;
//!
//! #[tokio::main]
//! async fn main() -> Result<()> {
//!     let mut supervisor = ProcessSupervisor::new(
//!         "bundle".to_string(),
//!         vec!["exec".to_string(), "puma".to_string()],
//!     );
//!
//!     // Start the process with PORT env var
//!     supervisor.start(3000).await?;
//!
//!     // Handle signals in background
//!     let signal_handle = tokio::spawn({
//!         let supervisor = supervisor.clone();
//!         async move {
//!             supervisor.handle_signals().await;
//!         }
//!     });
//!
//!     // Wait for process to exit
//!     let exit_code = supervisor.wait().await?;
//!
//!     // Clean up signal handler
//!     signal_handle.abort();
//!
//!     std::process::exit(exit_code);
//! }
//! ```

use anyhow::{Context, Result, anyhow};
use std::sync::Arc;
use std::time::Duration;
use tokio::process::{Child, Command};
use tokio::signal::unix::{SignalKind, signal};
use tokio::sync::{Mutex, Notify};
use tracing::{debug, error, info, warn};

/// Graceful shutdown timeout before sending SIGKILL
const SHUTDOWN_TIMEOUT: Duration = Duration::from_secs(30);

/// Process supervisor for managing upstream application servers
///
/// The ProcessSupervisor spawns and monitors a child process, forwarding signals
/// and handling graceful shutdown. It's designed to work with application servers
/// that listen on a specific port.
#[derive(Clone)]
pub struct ProcessSupervisor {
    /// Command to execute
    command: String,

    /// Arguments for the command
    args: Vec<String>,

    /// Child process handle (protected by mutex for async access)
    child: Arc<Mutex<Option<Child>>>,

    /// Notification channel for process startup
    started: Arc<Notify>,

    /// Notification channel for shutdown signal
    shutdown: Arc<Notify>,
}

impl ProcessSupervisor {
    /// Create a new process supervisor
    ///
    /// # Arguments
    ///
    /// * `command` - The command to execute (e.g., "bundle", "node")
    /// * `args` - Command-line arguments
    ///
    /// # Example
    ///
    /// ```
    /// use warpdrive::process::ProcessSupervisor;
    ///
    /// let supervisor = ProcessSupervisor::new(
    ///     "bundle".to_string(),
    ///     vec!["exec".to_string(), "puma".to_string()],
    /// );
    /// ```
    pub fn new(command: String, args: Vec<String>) -> Self {
        Self {
            command,
            args,
            child: Arc::new(Mutex::new(None)),
            started: Arc::new(Notify::new()),
            shutdown: Arc::new(Notify::new()),
        }
    }

    /// Start the child process with the specified port
    ///
    /// Spawns the child process with the PORT environment variable set.
    /// The process inherits stdin, stdout, and stderr from the parent.
    ///
    /// # Arguments
    ///
    /// * `port` - Port number to pass via PORT environment variable
    ///
    /// # Errors
    ///
    /// Returns an error if:
    /// - The command cannot be found or executed
    /// - The process is already running
    ///
    /// # Example
    ///
    /// ```no_run
    /// # use warpdrive::process::ProcessSupervisor;
    /// # use anyhow::Result;
    /// # async fn example() -> Result<()> {
    /// let mut supervisor = ProcessSupervisor::new("rails".to_string(), vec!["server".to_string()]);
    /// supervisor.start(3000).await?;
    /// # Ok(())
    /// # }
    /// ```
    pub async fn start(&self, port: u16) -> Result<()> {
        let mut child_guard = self.child.lock().await;

        // Check if process is already running
        if child_guard.is_some() {
            return Err(anyhow!("Process is already running"));
        }

        info!(
            "Starting upstream process: {} {}",
            self.command,
            self.args.join(" ")
        );

        // Spawn the child process with PORT environment variable
        let child = Command::new(&self.command)
            .args(&self.args)
            .env("PORT", port.to_string())
            .stdin(std::process::Stdio::inherit())
            .stdout(std::process::Stdio::inherit())
            .stderr(std::process::Stdio::inherit())
            .spawn()
            .with_context(|| format!("Failed to spawn command: {}", self.command))?;

        let pid = child
            .id()
            .ok_or_else(|| anyhow!("Failed to get child process ID"))?;

        info!("Upstream process started with PID: {}", pid);

        *child_guard = Some(child);
        drop(child_guard);

        // Notify that process has started
        self.started.notify_waiters();

        Ok(())
    }

    /// Wait for the child process to exit
    ///
    /// This method blocks until the child process terminates and returns
    /// the exit code. A non-zero exit code indicates the process was
    /// terminated by a signal or crashed.
    ///
    /// # Errors
    ///
    /// Returns an error if:
    /// - No process is running
    /// - Failed to wait for the process
    ///
    /// # Returns
    ///
    /// The exit code of the process (0 for success, non-zero for error/signal)
    ///
    /// # Example
    ///
    /// ```no_run
    /// # use warpdrive::process::ProcessSupervisor;
    /// # use anyhow::Result;
    /// # async fn example() -> Result<()> {
    /// # let mut supervisor = ProcessSupervisor::new("sleep".to_string(), vec!["1".to_string()]);
    /// # supervisor.start(3000).await?;
    /// let exit_code = supervisor.wait().await?;
    /// println!("Process exited with code: {}", exit_code);
    /// # Ok(())
    /// # }
    /// ```
    pub async fn wait(&self) -> Result<i32> {
        let mut child_guard = self.child.lock().await;

        let child = child_guard
            .as_mut()
            .ok_or_else(|| anyhow!("No process is running"))?;

        let status = child
            .wait()
            .await
            .context("Failed to wait for child process")?;

        // Clear the child handle
        *child_guard = None;
        drop(child_guard);

        let exit_code = status.code().unwrap_or(-1);

        if status.success() {
            info!("Upstream process exited successfully");
        } else {
            warn!("Upstream process exited with code: {}", exit_code);
        }

        Ok(exit_code)
    }

    /// Stop the child process gracefully
    ///
    /// Sends SIGTERM to the process and waits for it to exit within the
    /// timeout period. If the process doesn't exit within the timeout,
    /// sends SIGKILL to force termination.
    ///
    /// # Errors
    ///
    /// Returns an error if:
    /// - No process is running
    /// - Failed to send signals to the process
    ///
    /// # Example
    ///
    /// ```no_run
    /// # use warpdrive::process::ProcessSupervisor;
    /// # use anyhow::Result;
    /// # async fn example() -> Result<()> {
    /// # let mut supervisor = ProcessSupervisor::new("sleep".to_string(), vec!["100".to_string()]);
    /// # supervisor.start(3000).await?;
    /// supervisor.stop().await?;
    /// # Ok(())
    /// # }
    /// ```
    pub async fn stop(&self) -> Result<()> {
        let mut child_guard = self.child.lock().await;

        let child = match child_guard.as_mut() {
            Some(c) => c,
            None => return Ok(()), // No process running, nothing to stop
        };

        let pid = child
            .id()
            .ok_or_else(|| anyhow!("Failed to get child process ID"))?;

        info!("Stopping upstream process (PID: {})", pid);

        // Send SIGTERM for graceful shutdown
        debug!("Sending SIGTERM to PID {}", pid);
        send_signal(pid, SignalKind::terminate())?;

        // Wait for process to exit with timeout
        let wait_result = tokio::select! {
            status = child.wait() => {
                match status {
                    Ok(s) => {
                        info!("Process exited gracefully with code: {}", s.code().unwrap_or(-1));
                        Ok(())
                    }
                    Err(e) => Err(anyhow!("Failed to wait for process: {}", e))
                }
            }
            _ = tokio::time::sleep(SHUTDOWN_TIMEOUT) => {
                warn!("Process did not exit within timeout, sending SIGKILL");
                Err(anyhow!("Shutdown timeout"))
            }
        };

        // If graceful shutdown failed, force kill
        if wait_result.is_err() {
            debug!("Sending SIGKILL to PID {}", pid);
            if let Err(e) = child.kill().await {
                error!("Failed to kill process: {}", e);
            }

            // Wait for process to be reaped
            if let Err(e) = child.wait().await {
                error!("Failed to wait for killed process: {}", e);
            }
        }

        // Clear the child handle
        *child_guard = None;

        info!("Upstream process stopped");
        Ok(())
    }

    /// Handle signals and forward them to the child process
    ///
    /// This method listens for SIGTERM and SIGINT signals and forwards them
    /// to the child process. It should be spawned in a background task.
    ///
    /// The method will block until a signal is received or the shutdown
    /// notification is triggered.
    ///
    /// # Example
    ///
    /// ```no_run
    /// # use warpdrive::process::ProcessSupervisor;
    /// # use anyhow::Result;
    /// # async fn example() -> Result<()> {
    /// # let supervisor = ProcessSupervisor::new("sleep".to_string(), vec!["100".to_string()]);
    /// # supervisor.start(3000).await?;
    /// // Spawn signal handler in background
    /// let signal_handle = tokio::spawn({
    ///     let supervisor = supervisor.clone();
    ///     async move {
    ///         supervisor.handle_signals().await;
    ///     }
    /// });
    ///
    /// // Later, clean up
    /// signal_handle.abort();
    /// # Ok(())
    /// # }
    /// ```
    pub async fn handle_signals(&self) {
        let mut sigterm = match signal(SignalKind::terminate()) {
            Ok(s) => s,
            Err(e) => {
                error!("Failed to setup SIGTERM handler: {}", e);
                return;
            }
        };

        let mut sigint = match signal(SignalKind::interrupt()) {
            Ok(s) => s,
            Err(e) => {
                error!("Failed to setup SIGINT handler: {}", e);
                return;
            }
        };

        tokio::select! {
            _ = sigterm.recv() => {
                info!("Received SIGTERM, forwarding to child process");
                self.forward_signal(SignalKind::terminate()).await;
            }
            _ = sigint.recv() => {
                info!("Received SIGINT, forwarding to child process");
                self.forward_signal(SignalKind::interrupt()).await;
            }
            _ = self.shutdown.notified() => {
                debug!("Shutdown notification received");
            }
        }
    }

    /// Wait for the process to start
    ///
    /// This method blocks until the `start()` method has been called and
    /// the process has been spawned.
    ///
    /// # Example
    ///
    /// ```no_run
    /// # use warpdrive::process::ProcessSupervisor;
    /// # use anyhow::Result;
    /// # async fn example() -> Result<()> {
    /// # let supervisor = ProcessSupervisor::new("sleep".to_string(), vec!["1".to_string()]);
    /// // Start in background
    /// let supervisor_clone = supervisor.clone();
    /// tokio::spawn(async move {
    ///     supervisor_clone.start(3000).await.unwrap();
    /// });
    ///
    /// // Wait for startup
    /// supervisor.wait_for_start().await;
    /// println!("Process has started");
    /// # Ok(())
    /// # }
    /// ```
    pub async fn wait_for_start(&self) {
        self.started.notified().await;
    }

    /// Forward a signal to the child process
    async fn forward_signal(&self, kind: SignalKind) {
        let child_guard = self.child.lock().await;

        if let Some(child) = child_guard.as_ref() {
            if let Some(pid) = child.id() {
                if let Err(e) = send_signal(pid, kind) {
                    error!("Failed to forward signal to child: {}", e);
                }
            }
        }
    }

    /// Trigger shutdown notification
    ///
    /// This notifies the signal handler to stop listening for signals.
    pub fn trigger_shutdown(&self) {
        self.shutdown.notify_waiters();
    }

    /// Check if the process is currently running
    ///
    /// Returns `true` if a child process is running, `false` otherwise.
    ///
    /// # Example
    ///
    /// ```no_run
    /// # use warpdrive::process::ProcessSupervisor;
    /// # use anyhow::Result;
    /// # async fn example() -> Result<()> {
    /// # let supervisor = ProcessSupervisor::new("sleep".to_string(), vec!["1".to_string()]);
    /// # supervisor.start(3000).await?;
    /// if supervisor.is_running().await {
    ///     println!("Process is running");
    /// }
    /// # Ok(())
    /// # }
    /// ```
    pub async fn is_running(&self) -> bool {
        self.child.lock().await.is_some()
    }
}

/// Send a Unix signal to a process
///
/// # Arguments
///
/// * `pid` - Process ID to send signal to
/// * `kind` - Type of signal to send
///
/// # Errors
///
/// Returns an error if the signal cannot be sent or the process doesn't exist
fn send_signal(pid: u32, kind: SignalKind) -> Result<()> {
    use nix::sys::signal::{Signal, kill};
    use nix::unistd::Pid;

    let signal = match kind {
        k if k == SignalKind::terminate() => Signal::SIGTERM,
        k if k == SignalKind::interrupt() => Signal::SIGINT,
        k if k == SignalKind::quit() => Signal::SIGQUIT,
        _ => return Err(anyhow!("Unsupported signal type")),
    };

    kill(Pid::from_raw(pid as i32), signal).context("Failed to send signal to process")?;

    Ok(())
}

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

    #[tokio::test]
    async fn test_process_supervisor_new() {
        let supervisor = ProcessSupervisor::new("echo".to_string(), vec!["test".to_string()]);

        // Verify initial state
        assert!(supervisor.child.lock().await.is_none());
    }

    #[tokio::test]
    async fn test_spawn_and_wait() {
        let supervisor = ProcessSupervisor::new("sleep".to_string(), vec!["0.1".to_string()]);

        supervisor.start(3000).await.unwrap();
        let exit_code = supervisor.wait().await.unwrap();

        assert_eq!(exit_code, 0);
    }

    #[tokio::test]
    async fn test_double_start_error() {
        let supervisor = ProcessSupervisor::new("sleep".to_string(), vec!["10".to_string()]);

        supervisor.start(3000).await.unwrap();

        // Second start should fail
        let result = supervisor.start(3000).await;
        assert!(result.is_err());

        // Clean up
        supervisor.stop().await.unwrap();
    }

    #[tokio::test]
    async fn test_stop_running_process() {
        let supervisor = ProcessSupervisor::new("sleep".to_string(), vec!["100".to_string()]);

        supervisor.start(3000).await.unwrap();

        tokio::time::sleep(Duration::from_millis(100)).await;

        supervisor.stop().await.unwrap();

        // Process should be stopped
        assert!(supervisor.child.lock().await.is_none());
    }

    #[tokio::test]
    async fn test_wait_for_start() {
        let supervisor = ProcessSupervisor::new("sleep".to_string(), vec!["1".to_string()]);

        let supervisor_clone = supervisor.clone();
        tokio::spawn(async move {
            tokio::time::sleep(Duration::from_millis(100)).await;
            supervisor_clone.start(3000).await.unwrap();
        });

        supervisor.wait_for_start().await;

        // Process should be running now
        assert!(supervisor.child.lock().await.is_some());

        // Clean up
        supervisor.wait().await.unwrap();
    }

    #[tokio::test]
    async fn test_invalid_command() {
        let supervisor = ProcessSupervisor::new("nonexistent_command_xyz".to_string(), vec![]);

        let result = supervisor.start(3000).await;
        assert!(result.is_err());
    }
}