revive-differential-tests/crates/node/src/helpers/process.rs

use std::{
    fs::{File, OpenOptions},
    io::{BufRead, BufReader, Write},
    path::Path,
    process::{Child, Command},
    time::{Duration, Instant},
};

use anyhow::{Context, Result, bail};

/// A wrapper around processes which allows for their stdout and stderr to be logged and flushed
/// when the process is dropped.
#[derive(Debug)]
pub struct Process {
    /// The handle of the child process.
    child: Child,

    /// The file that stdout is being logged to.
    stdout_logs_file: File,

    /// The file that stderr is being logged to.
    stderr_logs_file: File,
}

impl Process {
    pub fn new(
        log_file_prefix: impl Into<Option<&'static str>>,
        logs_directory: impl AsRef<Path>,
        binary_path: impl AsRef<Path>,
        command_building_callback: impl FnOnce(&mut Command, File, File),
        process_readiness_wait_behavior: ProcessReadinessWaitBehavior,
    ) -> Result<Self> {
        let log_file_prefix = log_file_prefix.into();

        let (stdout_file_name, stderr_file_name) = match log_file_prefix {
            Some(prefix) => (
                format!("{prefix}_stdout.log"),
                format!("{prefix}_stderr.log"),
            ),
            None => ("stdout.log".to_string(), "stderr.log".to_string()),
        };

        let stdout_logs_file_path = logs_directory.as_ref().join(stdout_file_name);
        let stderr_logs_file_path = logs_directory.as_ref().join(stderr_file_name);

        let stdout_logs_file = OpenOptions::new()
            .write(true)
            .truncate(true)
            .create(true)
            .open(stdout_logs_file_path.as_path())
            .context("Failed to open the stdout logs file")?;
        let stderr_logs_file = OpenOptions::new()
            .write(true)
            .truncate(true)
            .create(true)
            .open(stderr_logs_file_path.as_path())
            .context("Failed to open the stderr logs file")?;

        let mut command = {
            let stdout_logs_file = stdout_logs_file
                .try_clone()
                .context("Failed to clone the stdout logs file")?;
            let stderr_logs_file = stderr_logs_file
                .try_clone()
                .context("Failed to clone the stderr logs file")?;

            let mut command = Command::new(binary_path.as_ref());
            command_building_callback(&mut command, stdout_logs_file, stderr_logs_file);
            command
        };
        let mut child = command
            .spawn()
            .context("Failed to spawn the built command")?;

        match process_readiness_wait_behavior {
            ProcessReadinessWaitBehavior::NoStartupWait => {}
            ProcessReadinessWaitBehavior::WaitDuration(duration) => std::thread::sleep(duration),
            ProcessReadinessWaitBehavior::TimeBoundedWaitFunction {
                max_wait_duration,
                mut check_function,
            } => {
                let spawn_time = Instant::now();

                let stdout_logs_file = OpenOptions::new()
                    .read(true)
                    .open(stdout_logs_file_path)
                    .context("Failed to open the stdout logs file")?;
                let stderr_logs_file = OpenOptions::new()
                    .read(true)
                    .open(stderr_logs_file_path)
                    .context("Failed to open the stderr logs file")?;

                let mut stdout_lines = BufReader::new(stdout_logs_file).lines();
                let mut stderr_lines = BufReader::new(stderr_logs_file).lines();

                let mut stdout = String::new();
                let mut stderr = String::new();

                loop {
                    let stdout_line = stdout_lines.next().and_then(Result::ok);
                    let stderr_line = stderr_lines.next().and_then(Result::ok);

                    if let Some(stdout_line) = stdout_line.as_ref() {
                        stdout.push_str(stdout_line);
                        stdout.push('\n');
                    }
                    if let Some(stderr_line) = stderr_line.as_ref() {
                        stderr.push_str(stderr_line);
                        stderr.push('\n');
                    }

                    let check_result =
                        check_function(stdout_line.as_deref(), stderr_line.as_deref()).context(
                            format!(
                                "Failed to wait for the process to be ready - {stdout} - {stderr}"
                            ),
                        )?;

                    if check_result {
                        break;
                    }

                    if Instant::now().duration_since(spawn_time) > max_wait_duration {
                        bail!(
                            "Waited for the process to start but it failed to start in time. stderr {stderr} - stdout {stdout}"
                        )
                    }
                }
            }
            ProcessReadinessWaitBehavior::WaitForCommandToExit => {
                if !child
                    .wait()
                    .context("Failed waiting for process to finish")?
                    .success()
                {
                    anyhow::bail!("Failed to spawn command");
                }
            }
        }

        Ok(Self {
            child,
            stdout_logs_file,
            stderr_logs_file,
        })
    }
}

impl Drop for Process {
    fn drop(&mut self) {
        self.child.kill().expect("Failed to kill the process");
        self.stdout_logs_file
            .flush()
            .expect("Failed to flush the stdout logs file");
        self.stderr_logs_file
            .flush()
            .expect("Failed to flush the stderr logs file");
    }
}

pub enum ProcessReadinessWaitBehavior {
    /// The process does not require any kind of wait after it's been spawned and can be used
    /// straight away.
    NoStartupWait,

    /// Waits for the command to exit.
    WaitForCommandToExit,

    /// The process does require some amount of wait duration after it's been started.
    WaitDuration(Duration),

    /// The process requires a time bounded wait function which is a function of the lines that
    /// appear in the log files.
    TimeBoundedWaitFunction {
        /// The maximum amount of time to wait for the check function to return true.
        max_wait_duration: Duration,

        /// The function to use to check if the process spawned is ready to use or not. This
        /// function should return the following in the following cases:
        ///
        /// - `Ok(true)`: Returned when the condition the process is waiting for has been fulfilled
        ///   and the wait is completed.
        /// - `Ok(false)`: The process is not ready yet but it might be ready in the future.
        /// - `Err`: The process is not ready yet and will not be ready in the future as it appears
        ///   that it has encountered an error when it was being spawned.
        ///
        /// The first argument is a line from stdout and the second argument is a line from stderr.
        #[allow(clippy::type_complexity)]
        check_function: Box<dyn FnMut(Option<&str>, Option<&str>) -> anyhow::Result<bool>>,
    },
}