xand-utils 1.0.0

//! Unix specific code for process cleanup

use nix::{sys::signal::Signal, unistd::Pid};
use std::{
    io,
    os::unix::process::CommandExt,
    process::{Child, Command, ExitStatus},
};

pub struct KillChildOnDrop {
    pub command: Command,
    kill_mode: ChildKillMode,
    child: Option<Child>,
}

pub enum ChildKillMode {
    /// Sends the SIGTERM signal, and then waits for the process to terminate.
    /// Note that this may block indefinitely if the process does not respect
    /// SIGTERM.
    SIGTERM,
    /// Sends the SIGKILL signal, which forcefully kills the process. This is
    /// the default mode.
    SIGKILL,
}

impl Default for ChildKillMode {
    fn default() -> Self {
        ChildKillMode::SIGKILL
    }
}

impl KillChildOnDrop {
    /// Creates a new KillChildOnDrop with the default ChildKillMode: ChildKillMode::SIGKILL.
    pub fn new(command: Command) -> Self {
        KillChildOnDrop {
            command,
            kill_mode: ChildKillMode::default(),
            child: None,
        }
    }

    /// Creates a new KillChildOnDrop with the specified ChildKillMode.
    ///
    /// After sending the appropriate termination signal, the process will wait
    /// for termination. See the documentation for the relevant ChildKillMode
    /// for details.
    pub fn with_kill_mode(command: Command, kill_mode: ChildKillMode) -> Self {
        KillChildOnDrop {
            command,
            kill_mode,
            child: None,
        }
    }

    pub fn spawn(&mut self) -> io::Result<&mut Self> {
        self.child = Some(
            // Spoooooky -- pre_exec() is unsafe. Yarn in particular has forced our hand here by
            // being absolutely crap at forwarding signals to children, but it is useful for
            // any processes that might "go rogue"
            //
            // Rationale: `pre_exec` is `unsafe`. The closure passed to `pre_exec` only calls
            // `setsid`, which abides by the safety requirements of `pre_exec`.
            #[allow(unsafe_code)]
            unsafe {
                self.command
                    .pre_exec(|| {
                        // Force the processes to run in a new process group so they can be murdered
                        nix::unistd::setsid().unwrap();
                        Ok(())
                    })
                    .spawn()?
            },
        );
        Ok(self)
    }

    pub fn try_wait(&mut self) -> io::Result<Option<ExitStatus>> {
        match self.child.as_mut() {
            Some(ch) => ch.try_wait(),
            None => Err(io::Error::new(
                io::ErrorKind::NotFound,
                "child not yet initialized",
            )),
        }
    }

    pub fn kill(&mut self) -> io::Result<()> {
        match self.child.as_mut() {
            Some(c) => c.kill(),
            None => Ok(()),
        }
    }

    pub fn request_terminate(&mut self) -> io::Result<()> {
        match &mut self.child {
            Some(child) => {
                let pid = Pid::from_raw(child.id() as i32);
                match nix::unistd::getpgid(Some(pid)) {
                    // politely terminate the process group
                    Ok(pgroup) => nix::sys::signal::killpg(pgroup, Signal::SIGTERM)
                        .map_err(|e| io::Error::new(io::ErrorKind::Other, e)),
                    // no such process: nothing to do
                    Err(nix::Error::Sys(nix::errno::Errno::ESRCH)) => Ok(()),
                    Err(e) => Err(io::Error::new(io::ErrorKind::Other, e)),
                }
            }
            // child not yet initialized: nothing to do
            None => Ok(()),
        }
    }
}

impl Drop for KillChildOnDrop {
    fn drop(&mut self) {
        self.request_terminate()
            .expect("process group termination failed");

        if let Some(ref mut child) = &mut self.child {
            if let ChildKillMode::SIGKILL = self.kill_mode {
                let _ = child.kill();
            }

            // If wait isn't called we can leave zombie processes behind
            let _ = child.wait();
        }
    }
}

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

    use std::{
        process::{Command, ExitStatus},
        thread,
        time::{Duration, Instant},
    };

    #[test]
    fn terminate_long_command_before_finished() {
        let cmd = {
            let mut cmd = Command::new("sleep");
            cmd.arg("1h");
            cmd
        };
        let mut dropme = KillChildOnDrop::new(cmd);
        let process = dropme.spawn().expect("unable to spawn process");

        thread::sleep(Duration::new(1, 0));
        match process.try_wait() {
            Ok(Some(_)) => panic!("process unexpectedly terminated early"),
            Ok(None) => {
                // process is still running as expected
            }
            Err(e) => panic!("Error: {:?}", e),
        }

        process
            .request_terminate()
            .expect("process group termination failed");

        thread::sleep(Duration::new(1, 0));
        match process.try_wait() {
            Ok(Some(s)) => assert!(!s.success()),
            Ok(None) => panic!("process unexpectedly still running after SIGTERM"),
            Err(e) => panic!("Error: {:?}", e),
        }
    }

    #[test]
    fn capture_successful_exit_status_from_quick_command() {
        let cmd = Command::new("true");
        let status = run_command(cmd, default_timeout());
        assert!(status.success());
    }

    #[test]
    fn capture_unsuccessful_exit_status_from_quick_command() {
        let cmd = Command::new("false");
        let status = run_command(cmd, default_timeout());
        assert!(!status.success());
    }

    fn default_timeout() -> Duration {
        Duration::new(1, 0)
    }

    fn run_command(cmd: Command, timeout: Duration) -> ExitStatus {
        let start_time = Instant::now();

        let mut dropme = KillChildOnDrop::new(cmd);
        let process = dropme.spawn().expect("unable to spawn process");

        loop {
            match process.try_wait() {
                Ok(None) => {
                    if Instant::now().duration_since(start_time) > timeout {
                        panic!("process timed out");
                    } else {
                        thread::sleep(Duration::new(1, 0));
                    }
                }
                Ok(Some(s)) => return s,
                Err(e) => panic!("unexpected error while waiting on process: {}", e),
            }
        }
    }
}