1use std::io::{self, Read, Write};
2use std::process::{Command, ExitStatus, Stdio};
3use std::time::Duration;
4
5use async_trait::async_trait;
6use blocking::unblock;
7use wait_timeout::ChildExt as _;
8
9#[derive(Debug)]
10pub struct SpawnAsyncOutput {
11 pub stdout: Vec<u8>,
12 pub stderr: Vec<u8>,
13 pub exit_status: ExitStatus,
14}
15
16#[async_trait]
17pub trait ChildExt {
18 async fn spawn_async(
19 &mut self,
20 stdin: Option<Vec<u8>>,
21 timeout: Option<Duration>,
22 max_size: Option<usize>,
23 ) -> io::Result<SpawnAsyncOutput>;
24}
25
26#[async_trait]
27impl ChildExt for Command {
28 async fn spawn_async(
29 &mut self,
30 stdin: Option<Vec<u8>>,
31 timeout: Option<Duration>,
32 max_size: Option<usize>,
33 ) -> io::Result<SpawnAsyncOutput> {
34 let command = self.stdout(Stdio::piped()).stderr(Stdio::piped());
35 if stdin.is_some() {
36 command.stdin(Stdio::piped());
37 }
38
39 let mut child = command.spawn()?;
40
41 unblock(move || {
42 if let Some(stdin_bytes) = stdin {
43 let stdin = match child.stdin.as_mut() {
44 Some(stdin) => stdin,
45 None => unreachable!("never"),
46 };
47 stdin.write_all(&stdin_bytes[..])?;
48 }
49
50 let exit_status =
51 match child.wait_timeout(timeout.unwrap_or_else(|| Duration::from_millis(1000)))? {
52 Some(exit_status) => exit_status,
53 None => {
54 child.kill()?;
56 child.wait()?;
57
58 return Err(io::Error::new(io::ErrorKind::TimedOut, "run timeout"));
59 }
60 };
61
62 let max_size = max_size.unwrap_or(2048);
63 let mut buf = vec![0; max_size];
64 buf.resize(max_size, 0);
65
66 let stdout = child
67 .stdout
68 .as_mut()
69 .ok_or_else(|| io::Error::new(io::ErrorKind::Other, "never"))?;
70 let n = stdout.read(&mut buf)?;
71 let stdout_bytes = buf[..n].to_vec();
72
73 let stderr = child
74 .stderr
75 .as_mut()
76 .ok_or_else(|| io::Error::new(io::ErrorKind::Other, "never"))?;
77 let n = stderr.read(&mut buf)?;
78 let stderr_bytes = buf[..n].to_vec();
79
80 Ok(SpawnAsyncOutput {
81 stdout: stdout_bytes,
82 stderr: stderr_bytes,
83 exit_status,
84 })
85 })
86 .await
87 }
88}