mx_tester/
exec.rs

1use std::{ffi::OsStr, path::PathBuf, process::Stdio};
2
3use anyhow::{anyhow, Context, Error};
4use async_trait::async_trait;
5use ezexec::lookup::Shell;
6use log::{debug, info};
7use tokio::fs::OpenOptions;
8use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader, BufWriter};
9use tokio::process::Command;
10
11/// Utility class: run a script in a shell.
12///
13/// Based on ezexec, customized to improve the ability to log.
14pub struct Executor {
15    /// The shell used to execute the script.
16    shell: Shell,
17}
18impl Executor {
19    pub fn try_new() -> Result<Self, Error> {
20        let shell = ezexec::lookup::Shell::find()
21            .map_err(|e| anyhow!("Could not find a shell to execute command: {}", e))?;
22        Ok(Self { shell })
23    }
24
25    /// Prepare a `Command` from a script.
26    ///
27    /// The resulting `Command` will be ready to execute in the shell.
28    /// You may customize it with e.g. `env()`.
29    pub fn command<P>(&self, cmd: P) -> Result<Command, Error>
30    where
31        P: AsRef<str>,
32    {
33        // Lookup shell.
34        let shell: &OsStr = self.shell.as_ref();
35        let mut command = Command::new(shell);
36
37        // Prefix `command` with the strings we need to call the shell.
38        let cmd = cmd.as_ref();
39        let execstring_args = self
40            .shell
41            .execstring_args()
42            .map_err(|e| anyhow!("Could not find a shell string: {}", e))?;
43        let args = execstring_args.iter().chain(std::iter::once(&cmd));
44
45        command.args(args);
46        command.stdout(Stdio::piped());
47        command.stderr(Stdio::piped());
48
49        Ok(command)
50    }
51}
52
53/// Utility function: spawn an async task to asynchronously write the contents
54/// of a reader to both a file and a log.
55fn spawn_logger<T>(name: &'static str, reader: BufReader<T>, dest: PathBuf, command: &str)
56where
57    BufReader<T>: AsyncBufReadExt + Unpin,
58    T: 'static + Send,
59{
60    debug!("Storing {} logs in {:?}", name, dest);
61    let command = format!("\ncommand: {}\n", command.to_string());
62    tokio::task::spawn(async move {
63        let mut file = OpenOptions::new()
64            .create(true)
65            .append(true)
66            .open(dest)
67            .await
68            .with_context(|| format!("Could not create log file {}", name))?;
69        {
70            // Create a buffered writer, we don't want to hit the disk with
71            // every single byte.
72            let mut writer = BufWriter::new(&mut file);
73            writer
74                .write_all(command.as_bytes())
75                .await
76                .with_context(|| format!("Could not write log file {}", name))?;
77            writer
78                .flush()
79                .await
80                .with_context(|| format!("Could not write log file {}", name))?;
81
82            let mut lines = reader.lines();
83            while let Ok(Some(line)) = lines.next_line().await {
84                // Display logs.
85                info!("{}: {}", name, line);
86                // Write logs to `dest`.
87                writer
88                    .write_all(line.as_bytes())
89                    .await
90                    .with_context(|| format!("Could not write log file {}", name))?;
91                writer
92                    .write_all(b"\n")
93                    .await
94                    .with_context(|| format!("Could not write log file {}", name))?;
95                // Flush after each write, in case of crash.
96                writer
97                    .flush()
98                    .await
99                    .with_context(|| format!("Could not write log file {}", name))?;
100            }
101        }
102        let _ = file.sync_data().await;
103        Ok(()) as Result<(), anyhow::Error>
104    });
105}
106
107/// Extension trait for `Command`.
108#[async_trait]
109pub trait CommandExt {
110    /// Spawn a command, logging its stdout/stderr to files and to the env logger.
111    async fn spawn_logged(
112        &mut self,
113        log_dir: &PathBuf,
114        name: &'static str,
115        line: &str,
116    ) -> Result<(), Error>;
117}
118
119#[async_trait]
120impl CommandExt for Command {
121    async fn spawn_logged(
122        &mut self,
123        log_dir: &PathBuf,
124        name: &'static str,
125        line: &str,
126    ) -> Result<(), Error> {
127        let mut child = self
128            .spawn()
129            .with_context(|| format!("Could not spawn process for `{}`", name))?;
130        // Spawn background tasks to write down stdout.
131        if let Some(stdout) = child.stdout.take() {
132            let reader = BufReader::new(stdout);
133            let log_path = log_dir.join(format!("{name}.out", name = name));
134            spawn_logger(name, reader, log_path, line);
135        }
136        // Spawn background tasks to write down stderr.
137        if let Some(stderr) = child.stderr.take() {
138            let reader = BufReader::new(stderr);
139            let log_path = log_dir.join(format!("{name}.log", name = name));
140            spawn_logger(name, reader, log_path, line);
141        }
142        let status = child.wait().await.context("Child process not launched")?;
143        if status.success() {
144            return Ok(());
145        }
146        Err(anyhow!("Child `{}` failed: `{}`", name, status))
147    }
148}
mx_tester/exec.rs

mx_tester/
exec.rs
