tmc_langs_framework/
command.rs

1//! Custom wrapper for Command that supports timeouts and contains custom error handling.
2
3use crate::{TmcError, error::CommandError};
4use std::{
5    ffi::OsStr,
6    fs::File,
7    io::{Read, Write},
8    thread::JoinHandle,
9    time::Duration,
10};
11pub use subprocess::ExitStatus;
12use subprocess::{Exec, PopenError, Redirection};
13
14/// Wrapper around subprocess::Exec
15#[must_use]
16pub struct TmcCommand {
17    exec: Exec,
18    stdin: Option<String>,
19}
20
21impl TmcCommand {
22    /// Creates a new command
23    pub fn new(cmd: impl AsRef<OsStr>) -> Self {
24        Self {
25            exec: Exec::cmd(cmd).env("LANG", "en_US.UTF-8"),
26            stdin: None,
27        }
28    }
29
30    /// Creates a new command with piped stdout/stderr.
31    pub fn piped(cmd: impl AsRef<OsStr>) -> Self {
32        Self {
33            exec: Exec::cmd(cmd)
34                .stdout(Redirection::Pipe)
35                .stderr(Redirection::Pipe)
36                .env("LANG", "en_US.UTF-8"),
37            stdin: None,
38        }
39    }
40
41    /// Allows modification of the internal command without providing access to it.
42    pub fn with(self, f: impl FnOnce(Exec) -> Exec) -> Self {
43        Self {
44            exec: f(self.exec),
45            ..self
46        }
47    }
48
49    /// Gives the command data to write into stdin.
50    pub fn set_stdin_data(self, data: String) -> Self {
51        Self {
52            exec: self.exec.stdin(Redirection::Pipe),
53            stdin: Some(data),
54        }
55    }
56
57    // executes the given command and collects its output
58    fn execute(self, timeout: Option<Duration>, checked: bool) -> Result<Output, TmcError> {
59        let cmd = self.exec.to_cmdline_lossy();
60        log::info!("executing {cmd}");
61
62        let Self { exec, stdin } = self;
63
64        // starts executing the command
65        let mut popen = exec.popen().map_err(|e| popen_to_tmc_err(cmd.clone(), e))?;
66        let stdin_handle = spawn_writer(popen.stdin.take(), stdin);
67        let stdout_handle = spawn_reader(popen.stdout.take());
68        let stderr_handle = spawn_reader(popen.stderr.take());
69
70        let exit_status = if let Some(timeout) = timeout {
71            // timeout set
72            let exit_status = popen
73                .wait_timeout(timeout)
74                .map_err(|e| popen_to_tmc_err(cmd.clone(), e))?;
75
76            match exit_status {
77                Some(exit_status) => exit_status,
78                None => {
79                    // None means that we timed out
80                    popen
81                        .terminate()
82                        .map_err(|e| CommandError::Terminate(cmd.clone(), e))?;
83                    let stdout = stdout_handle
84                        .join()
85                        .expect("the thread should not be able to panic");
86                    let stderr = stderr_handle
87                        .join()
88                        .expect("the thread should not be able to panic");
89                    return Err(TmcError::Command(CommandError::TimeOut {
90                        command: cmd,
91                        timeout,
92                        stdout: String::from_utf8_lossy(&stdout).into_owned(),
93                        stderr: String::from_utf8_lossy(&stderr).into_owned(),
94                    }));
95                }
96            }
97        } else {
98            // no timeout, block until done
99            popen.wait().map_err(|e| popen_to_tmc_err(cmd.clone(), e))?
100        };
101
102        log::info!("finished executing {cmd}");
103        stdin_handle
104            .join()
105            .expect("the thread should not be able to panic");
106        let stdout = stdout_handle
107            .join()
108            .expect("the thread should not be able to panic");
109        let stderr = stderr_handle
110            .join()
111            .expect("the thread should not be able to panic");
112
113        // on success, log stdout trace and stderr debug
114        // on failure if checked, log warn
115        // on failure if not checked, log debug
116        if !exit_status.success() {
117            // if checked is set, error with failed exit status
118            if checked {
119                log::warn!("stdout: {}", String::from_utf8_lossy(&stdout).into_owned());
120                log::warn!("stderr: {}", String::from_utf8_lossy(&stderr).into_owned());
121                return Err(CommandError::Failed {
122                    command: cmd,
123                    status: exit_status,
124                    stdout: String::from_utf8_lossy(&stdout).into_owned(),
125                    stderr: String::from_utf8_lossy(&stderr).into_owned(),
126                }
127                .into());
128            } else {
129                log::debug!("stdout: {}", String::from_utf8_lossy(&stdout).into_owned());
130                log::debug!("stderr: {}", String::from_utf8_lossy(&stderr).into_owned());
131            }
132        } else {
133            log::trace!("stdout: {}", String::from_utf8_lossy(&stdout).into_owned());
134            log::debug!("stderr: {}", String::from_utf8_lossy(&stderr).into_owned());
135        }
136        Ok(Output {
137            status: exit_status,
138            stdout,
139            stderr,
140        })
141    }
142
143    /// Executes the command and waits for its output.
144    pub fn status(self) -> Result<ExitStatus, TmcError> {
145        self.execute(None, false).map(|o| o.status)
146    }
147
148    /// Executes the command and waits for its output.
149    pub fn output(self) -> Result<Output, TmcError> {
150        self.execute(None, false)
151    }
152
153    /// Executes the command and waits for its output and errors if the status is not successful.
154    pub fn output_checked(self) -> Result<Output, TmcError> {
155        self.execute(None, true)
156    }
157
158    /// Executes the command and waits for its output with the given timeout.
159    pub fn output_with_timeout(self, timeout: Duration) -> Result<Output, TmcError> {
160        self.execute(Some(timeout), false)
161    }
162
163    /// Executes the command and waits for its output with the given timeout and errors if the status is not successful.
164    pub fn output_with_timeout_checked(self, timeout: Duration) -> Result<Output, TmcError> {
165        self.execute(Some(timeout), true)
166    }
167}
168
169// it's assumed the thread will never panic
170fn spawn_writer(file: Option<File>, data: Option<String>) -> JoinHandle<()> {
171    std::thread::spawn(move || {
172        if let Some(mut file) = file {
173            if let Some(data) = data {
174                log::debug!("writing data");
175                if let Err(err) = file.write_all(data.as_bytes()) {
176                    log::error!("failed to write data in writer thread: {err}");
177                }
178            }
179        }
180    })
181}
182
183// it's assumed the thread will never panic
184fn spawn_reader(file: Option<File>) -> JoinHandle<Vec<u8>> {
185    std::thread::spawn(move || {
186        if let Some(mut file) = file {
187            let mut buf = vec![];
188            if let Err(err) = file.read_to_end(&mut buf) {
189                log::error!("failed to read data in reader thread: {err}");
190            }
191            buf
192        } else {
193            vec![]
194        }
195    })
196}
197
198// convenience function to convert an error while checking for command not found error
199fn popen_to_tmc_err(cmd: String, err: PopenError) -> TmcError {
200    if let PopenError::IoError(io) = &err {
201        if let std::io::ErrorKind::NotFound = io.kind() {
202            TmcError::Command(CommandError::NotFound { cmd, source: err })
203        } else {
204            TmcError::Command(CommandError::FailedToRun(cmd, err))
205        }
206    } else {
207        TmcError::Command(CommandError::Popen(cmd, err))
208    }
209}
210
211#[derive(Debug)]
212pub struct Output {
213    pub status: ExitStatus,
214    pub stdout: Vec<u8>,
215    pub stderr: Vec<u8>,
216}
217
218#[cfg(test)]
219mod test {
220    use super::*;
221
222    #[test]
223    fn timeout() {
224        let cmd = TmcCommand::piped("sleep").with(|e| e.arg("2"));
225        assert!(matches!(
226            cmd.output_with_timeout(Duration::from_nanos(1)),
227            Err(TmcError::Command(CommandError::TimeOut { .. }))
228        ));
229    }
230
231    #[test]
232    fn not_found() {
233        let cmd = TmcCommand::piped("nonexistent command");
234        assert!(matches!(
235            cmd.output(),
236            Err(TmcError::Command(CommandError::NotFound { .. }))
237        ));
238    }
239}
tmc_langs_framework/command.rs

tmc_langs_framework/
command.rs
