sopht 0.3.0

//! contains all structures and functions related to the [`State`] structure.
//! this includes the implementations of all the server-side process management
//! functions like `start` and `status` as well as things like [`RestartPolicy`]
//! and [`ManagedProcess`]
//!
use crate::{Error, Process, Result, PID};
use serde::{Deserialize, Serialize};
use std::fmt::{self, Display, Formatter};
use std::path::PathBuf;
use std::time::Duration;

/// arguments to [`State::start`]
#[derive(PartialEq, Serialize, Deserialize, Debug)]
pub struct StartArgs {
    /// the program to run
    pub prog: String,
    /// arguments to the program
    pub args: Vec<String>,
    /// the restart policy to start the command with
    pub restart_policy: RestartPolicy,
    // working directory for the process
    pub working_dir: PathBuf,
}

/// arguments to [`State::change_restart_policy`]
#[derive(PartialEq, Serialize, Deserialize, Debug)]
pub struct ChangeRestartPolicyArgs {
    /// the pid of the process
    pub pid: PID,
    /// the restart policy to set
    pub restart_policy: RestartPolicy,
}

/// arguments to [`State::send`]
#[derive(PartialEq, Serialize, Deserialize, Debug)]
pub struct SendArgs {
    /// the pid of the process
    pub pid: PID,
    /// the input to send to the process. does not have end with a newline
    pub input: String,
}

/// wrapper around [`Process`] with some additional information required for
/// the process manager to operate
struct ManagedProcess {
    /// the process itself
    process: Process,
    /// contains the process' new PID if it has been restarted. if this field is
    /// `None`, then the process has never been restarted
    restarted: Option<PID>,
    /// whether or not this process was ever considered started
    started: bool,
    /// describes when to restart the process automatically. see
    /// [`RestartPolicy`] for more details
    restart_policy: RestartPolicy,
}

/// the state of the sopht process manager. contains all processes started and
/// all other configuration/state
pub struct State {
    /// the list of processes managed by sopht
    processes: Vec<ManagedProcess>,
    /// how long a process must stay alive for it to be considered started
    stay_alive_duration: Duration,
}

/// how sopht internally represents the exit status of a process
#[derive(Serialize, Deserialize, PartialEq, Clone, Debug)]
pub enum ExitStatus {
    /// the process is running or its exit status is unconfirmed
    NotExited,
    /// the process exited with a successful exit code
    Success,
    /// the process exited with an exit code indicating failure
    Failure,
}

/// detailed status of a `Process`. this structure is generated by
/// [`State::status_of`]
#[derive(Serialize, Deserialize, Clone, Debug)]
pub struct ProcessStatus {
    /// the PID of the process
    pub pid: PID,
    /// the path to the program the process is running at
    pub prog: String,
    /// the arguments to the program
    pub args: Vec<String>,
    /// whether or not the program is alive. a value of `None` means its status
    /// is unknown
    pub exit_status: ExitStatus,
    /// whether or not the process was ever started
    pub started: bool,
    /// if the process has been restarted, what its new PID is
    pub restarted: Option<PID>,
    /// the restart policy of the process. see [`RestartPolicy`] for more
    /// details on restart policies
    pub restart_policy: RestartPolicy,
}

/// this structure is generated by [`State::upkeep`] on success and contains
/// information about what upkeep was performed (i.e which processes were
/// restarted)
pub struct UpkeepInfo {
    /// which processes were restarted in an upkeep cycle. note that all of
    /// these PIDs are dead
    pub restarted: Vec<PID>,
}

/// describes the conditions under which a process will be restarted in an
/// upkeep cycle. this has no effect on [`State::restart`]
#[derive(PartialEq, Serialize, Deserialize, Copy, Clone, Debug)]
pub enum RestartPolicy {
    /// do not restart the process
    Never,
    /// restart the process if it stops for any reason, including successful
    /// exit codes
    Always,
    /// restart the process if it exits with a nonzero exit code. the single
    /// parameter is the *retry count*, which will be 1 less in each
    /// new instance of the process. a process with a retry count of 0 will not
    /// be restarted
    OnFailure(u32),
}

impl State {
    /// creates a new [`State`] with the specified configuration
    pub fn new(stay_alive_secs: u64) -> Self {
        let processes = Vec::new();
        let stay_alive_duration = Duration::from_secs(stay_alive_secs);
        Self {
            processes,
            stay_alive_duration,
        }
    }

    /// starts a process given a command string and adds it to the processes
    /// list. the process *must* stay alive for the `stay_alive_duration` to be
    /// considered started. if the process does not stay alive long enough, a
    /// [`StayAliveError`](Error::StayAliveError) will be returned. the
    /// process's output will still be collected and it will still be assigned a
    /// PID and stored in the process list. if the process fail to start,
    /// a [`ProcessError`](Error::ProcessError) will be returned
    pub fn start(&mut self, args: &StartArgs) -> Result<ProcessStatus> {
        let mut managed = ManagedProcess {
            process: Process::start(&args.prog, &args.args, &args.working_dir)?,
            restarted: None,
            started: false,
            restart_policy: args.restart_policy,
        };
        std::thread::sleep(self.stay_alive_duration);

        if let Some(true) = managed.process.is_alive() {
            managed.started = true;
        } else {
            return Err(Error::StayAliveError);
        }

        self.processes.push(managed);
        let pid = self.processes.len() - 1;
        Ok(self.status_of(pid))
    }

    /// retrieves the output of the specified process as a slice of its lines of
    /// output. this function can fail if the PID supplied is out of range with
    /// an [`InvalidPID`](Error::InvalidPID) error
    pub fn output(&mut self, args: PID) -> Result<&Vec<u8>> {
        let pid = self.check_pid(args)?;

        self.processes[pid].process.gather_output();

        Ok(&self.processes[pid].process.output)
    }

    /// retrieves the status of all processes unless `count` is specified, in
    /// which case the last `count` processes' status
    pub fn status(&mut self, count: &Option<usize>) -> Vec<ProcessStatus> {
        let mut v = Vec::new();
        let len = self.processes.len();
        let mut count = count.unwrap_or(len);
        if len < count {
            count = len;
        }
        for pid in (len - count)..self.processes.len() {
            let s = self.status_of(pid);
            v.push(s);
        }

        v
    }

    /// attempts to stop a process given its PID. this function can fail if the
    /// PID supplied is out of range ([`InvalidPID`](Error::InvalidPID)) or if
    /// killing the process fails ([`ProcessError`](Error::ProcessError)).
    /// additionally, this function sets the restart policy of the process to
    /// [`Never`](RestartPolicy::Never) so it won't be restarted
    pub fn stop(&mut self, args: PID) -> Result<ProcessStatus> {
        let pid = self.check_pid(args)?;
        //self.processes[pid].restart_policy = RestartPolicy::Never;
        match self.processes[pid].process.kill() {
            Ok(_) => Ok(self.status_of(pid)),
            Err(e) => Err(e),
        }
    }

    /// stops a particular process and then start it again with a new PID. this
    /// function just calls [`State::start`] and [`State::stop`] internally,
    /// so it can fail in exactly the same ways those functions can. the `retry`
    /// parameter signifies whether or not this restart is a manual restart or
    /// an automatic retry by [`State::upkeep`]. if `retry` is true, then the
    /// restart policy will be copied to the restartd process. in the case of
    /// the [`RestartPolicy::OnFailure`] policy, then this restart will count
    /// against its retries and subtract 1 from it
    pub fn restart(&mut self, args: PID, retry: bool) -> Result<ProcessStatus> {
        let pid = self.check_pid(args)?;
        let stopped = self.stop(args)?;
        let restart_policy = if retry {
            if let RestartPolicy::OnFailure(n) = stopped.restart_policy {
                RestartPolicy::OnFailure(n - 1)
            } else {
                stopped.restart_policy
            }
        } else {
            stopped.restart_policy
        };
        let status = self.start(&StartArgs {
            prog: stopped.prog.to_owned(),
            args: stopped.args.clone(),
            restart_policy,
            working_dir: self.processes[pid].process.working_dir.clone(),
        })?;
        self.processes[stopped.pid].restarted = Some(status.pid);

        Ok(status)
    }

    /// send input to a process given its PID and input. this function can fail
    /// if the PID is invalid ([`InvalidPID`](Error::InvalidPID)) or if the
    /// input fails to send ([`ProcessError`](Error::ProcessError))
    pub fn send(&mut self, args: &SendArgs) -> Result<()> {
        let pid = self.check_pid(args.pid)?;
        self.processes[pid].process.send_input(&args.input)
    }

    /// generates a [`ProcessStatus`] for a PID. <strong>this function will
    /// panic if supplied an invalid PID</strong>, so checking the PID is up to
    /// the caller!
    pub fn status_of(&mut self, pid: PID) -> ProcessStatus {
        let proc = self
            .processes
            .get_mut(pid)
            .expect("attempted to get status of invalid PID");

        let prog = proc.process.prog.clone();
        let args = proc.process.args.clone();
        let exit_status = match proc.process.exit_status() {
            Some(e) => match e.success() {
                true => ExitStatus::Success,
                false => ExitStatus::Failure,
            },
            None => ExitStatus::NotExited,
        };
        let started = proc.started;
        let restarted = proc.restarted;
        let restart_policy = proc.restart_policy;

        ProcessStatus {
            pid,
            prog,
            args,
            exit_status,
            started,
            restarted,
            restart_policy,
        }
    }

    /// changes the restart policy of the specified process to the specified
    /// policy
    pub fn change_restart_policy(
        &mut self,
        args: &ChangeRestartPolicyArgs,
    ) -> Result<ProcessStatus> {
        let pid = self.check_pid(args.pid)?;

        self.processes[pid].restart_policy = args.restart_policy;

        Ok(self.status_of(pid))
    }

    /// utility method for checking if a PID is valid. fails with
    /// [`InvalidPID`](Error::InvalidPID) if the PID is invalid and returns the
    /// PID itself otherwise
    pub fn check_pid(&self, pid: PID) -> Result<PID> {
        if pid >= self.processes.len() {
            Err(Error::InvalidPID)
        } else {
            Ok(pid)
        }
    }

    /// handles any periodic upkeep the state needs. currently, this is only
    /// restarting any crashed processes that have their restart policy set.
    /// returns an `UpkeepInfo` structure containing which processes were
    /// restarted
    pub fn upkeep(&mut self) -> Result<UpkeepInfo> {
        let mut failed_processes = Vec::new();
        // build the list of failed processes to consider restarting.
        // if a process has exited and has not been restarted yet, then consider
        // it failed
        for (i, proc) in self.processes.iter_mut().enumerate() {
            if proc.process.exit_status().is_some() && proc.restarted.is_none()
            {
                failed_processes.push(i);
            }
        }

        let mut restarts = Vec::new();

        for pid in &failed_processes {
            let status = self.status_of(*pid);
            match status.restart_policy {
                RestartPolicy::Never => continue,
                RestartPolicy::Always => {
                    self.restart(*pid, false)?;
                    restarts.push(*pid);
                }
                RestartPolicy::OnFailure(n) => {
                    // if it's out of retries, don't restart it
                    if n == 0 {
                        continue;
                    }

                    // and if it exited successfully, don't restart it
                    if status.exit_status == ExitStatus::Success {
                        continue;
                    }

                    if self.restart(*pid, true).is_ok() {
                        restarts.push(*pid);
                        break;
                    }
                }
            }
        }

        Ok(UpkeepInfo {
            restarted: restarts,
        })
    }
}

impl Display for RestartPolicy {
    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
        match self {
            RestartPolicy::Never => write!(f, "never"),
            RestartPolicy::Always => write!(f, "always"),
            RestartPolicy::OnFailure(retries) => {
                write!(f, "on failure ({})", retries)
            }
        }
    }
}

impl Drop for State {
    fn drop(&mut self) {
        for proc in &mut self.processes {
            proc.process
                .kill()
                .expect("failed to kill process during drop");
            proc.process
                .child
                .wait()
                .expect("failed to wait on process during drop");
        }
    }
}