crates/stdx/src/process.rs

   1 //! Read both stdout and stderr of child without deadlocks.
   2 //!
   3 //! <https://github.com/rust-lang/cargo/blob/905af549966f23a9288e9993a85d1249a5436556/crates/cargo-util/src/read2.rs>
   4 //! <https://github.com/rust-lang/cargo/blob/58a961314437258065e23cb6316dfc121d96fb71/crates/cargo-util/src/process_builder.rs#L231>
   5
   6 use std::{
   7     io,
   8     process::{Command, Output, Stdio},
   9 };
  10
  11 pub fn streaming_output(
  12     mut cmd: Command,
  13     on_stdout_line: &mut dyn FnMut(&str),
  14     on_stderr_line: &mut dyn FnMut(&str),
  15 ) -> io::Result<Output> {
  16     let mut stdout = Vec::new();
  17     let mut stderr = Vec::new();
  18
  19     let cmd = cmd.stdout(Stdio::piped()).stderr(Stdio::piped()).stdin(Stdio::null());
  20
  21     let status = {
  22         let mut child = cmd.spawn()?;
  23         let out = child.stdout.take().unwrap();
  24         let err = child.stderr.take().unwrap();
  25         imp::read2(out, err, &mut |is_out, data, eof| {
  26             let idx = if eof {
  27                 data.len()
  28             } else {
  29                 match data.iter().rposition(|b| *b == b'\n') {
  30                     Some(i) => i + 1,
  31                     None => return,
  32                 }
  33             };
  34             {
  35                 // scope for new_lines
  36                 let new_lines = {
  37                     let dst = if is_out { &mut stdout } else { &mut stderr };
  38                     let start = dst.len();
  39                     let data = data.drain(..idx);
  40                     dst.extend(data);
  41                     &dst[start..]
  42                 };
  43                 for line in String::from_utf8_lossy(new_lines).lines() {
  44                     if is_out {
  45                         on_stdout_line(line);
  46                     } else {
  47                         on_stderr_line(line);
  48                     }
  49                 }
  50             }
  51         })?;
  52         child.wait()?
  53     };
  54
  55     Ok(Output { status, stdout, stderr })
  56 }
  57
  58 #[cfg(unix)]
  59 mod imp {
  60     use std::{
  61         io::{self, prelude::*},
  62         mem,
  63         os::unix::prelude::*,
  64         process::{ChildStderr, ChildStdout},
  65     };
  66
  67     pub(crate) fn read2(
  68         mut out_pipe: ChildStdout,
  69         mut err_pipe: ChildStderr,
  70         data: &mut dyn FnMut(bool, &mut Vec<u8>, bool),
  71     ) -> io::Result<()> {
  72         unsafe {
  73             libc::fcntl(out_pipe.as_raw_fd(), libc::F_SETFL, libc::O_NONBLOCK);
  74             libc::fcntl(err_pipe.as_raw_fd(), libc::F_SETFL, libc::O_NONBLOCK);
  75         }
  76
  77         let mut out_done = false;
  78         let mut err_done = false;
  79         let mut out = Vec::new();
  80         let mut err = Vec::new();
  81
  82         let mut fds: [libc::pollfd; 2] = unsafe { mem::zeroed() };
  83         fds[0].fd = out_pipe.as_raw_fd();
  84         fds[0].events = libc::POLLIN;
  85         fds[1].fd = err_pipe.as_raw_fd();
  86         fds[1].events = libc::POLLIN;
  87         let mut nfds = 2;
  88         let mut errfd = 1;
  89
  90         while nfds > 0 {
  91             // wait for either pipe to become readable using `select`
  92             let r = unsafe { libc::poll(fds.as_mut_ptr(), nfds, -1) };
  93             if r == -1 {
  94                 let err = io::Error::last_os_error();
  95                 if err.kind() == io::ErrorKind::Interrupted {
  96                     continue;
  97                 }
  98                 return Err(err);
  99             }
 100
 101             // Read as much as we can from each pipe, ignoring EWOULDBLOCK or
 102             // EAGAIN. If we hit EOF, then this will happen because the underlying
 103             // reader will return Ok(0), in which case we'll see `Ok` ourselves. In
 104             // this case we flip the other fd back into blocking mode and read
 105             // whatever's leftover on that file descriptor.
 106             let handle = |res: io::Result<_>| match res {
 107                 Ok(_) => Ok(true),
 108                 Err(e) => {
 109                     if e.kind() == io::ErrorKind::WouldBlock {
 110                         Ok(false)
 111                     } else {
 112                         Err(e)
 113                     }
 114                 }
 115             };
 116             if !err_done && fds[errfd].revents != 0 && handle(err_pipe.read_to_end(&mut err))? {
 117                 err_done = true;
 118                 nfds -= 1;
 119             }
 120             data(false, &mut err, err_done);
 121             if !out_done && fds[0].revents != 0 && handle(out_pipe.read_to_end(&mut out))? {
 122                 out_done = true;
 123                 fds[0].fd = err_pipe.as_raw_fd();
 124                 errfd = 0;
 125                 nfds -= 1;
 126             }
 127             data(true, &mut out, out_done);
 128         }
 129         Ok(())
 130     }
 131 }
 132
 133 #[cfg(windows)]
 134 mod imp {
 135     use std::{
 136         io,
 137         os::windows::prelude::*,
 138         process::{ChildStderr, ChildStdout},
 139         slice,
 140     };
 141
 142     use miow::{
 143         iocp::{CompletionPort, CompletionStatus},
 144         pipe::NamedPipe,
 145         Overlapped,
 146     };
 147     use winapi::shared::winerror::ERROR_BROKEN_PIPE;
 148
 149     struct Pipe<'a> {
 150         dst: &'a mut Vec<u8>,
 151         overlapped: Overlapped,
 152         pipe: NamedPipe,
 153         done: bool,
 154     }
 155
 156     pub(crate) fn read2(
 157         out_pipe: ChildStdout,
 158         err_pipe: ChildStderr,
 159         data: &mut dyn FnMut(bool, &mut Vec<u8>, bool),
 160     ) -> io::Result<()> {
 161         let mut out = Vec::new();
 162         let mut err = Vec::new();
 163
 164         let port = CompletionPort::new(1)?;
 165         port.add_handle(0, &out_pipe)?;
 166         port.add_handle(1, &err_pipe)?;
 167
 168         unsafe {
 169             let mut out_pipe = Pipe::new(out_pipe, &mut out);
 170             let mut err_pipe = Pipe::new(err_pipe, &mut err);
 171
 172             out_pipe.read()?;
 173             err_pipe.read()?;
 174
 175             let mut status = [CompletionStatus::zero(), CompletionStatus::zero()];
 176
 177             while !out_pipe.done || !err_pipe.done {
 178                 for status in port.get_many(&mut status, None)? {
 179                     if status.token() == 0 {
 180                         out_pipe.complete(status);
 181                         data(true, out_pipe.dst, out_pipe.done);
 182                         out_pipe.read()?;
 183                     } else {
 184                         err_pipe.complete(status);
 185                         data(false, err_pipe.dst, err_pipe.done);
 186                         err_pipe.read()?;
 187                     }
 188                 }
 189             }
 190
 191             Ok(())
 192         }
 193     }
 194
 195     impl<'a> Pipe<'a> {
 196         unsafe fn new<P: IntoRawHandle>(p: P, dst: &'a mut Vec<u8>) -> Pipe<'a> {
 197             Pipe {
 198                 dst,
 199                 pipe: NamedPipe::from_raw_handle(p.into_raw_handle()),
 200                 overlapped: Overlapped::zero(),
 201                 done: false,
 202             }
 203         }
 204
 205         unsafe fn read(&mut self) -> io::Result<()> {
 206             let dst = slice_to_end(self.dst);
 207             match self.pipe.read_overlapped(dst, self.overlapped.raw()) {
 208                 Ok(_) => Ok(()),
 209                 Err(e) => {
 210                     if e.raw_os_error() == Some(ERROR_BROKEN_PIPE as i32) {
 211                         self.done = true;
 212                         Ok(())
 213                     } else {
 214                         Err(e)
 215                     }
 216                 }
 217             }
 218         }
 219
 220         unsafe fn complete(&mut self, status: &CompletionStatus) {
 221             let prev = self.dst.len();
 222             self.dst.set_len(prev + status.bytes_transferred() as usize);
 223             if status.bytes_transferred() == 0 {
 224                 self.done = true;
 225             }
 226         }
 227     }
 228
 229     unsafe fn slice_to_end(v: &mut Vec<u8>) -> &mut [u8] {
 230         if v.capacity() == 0 {
 231             v.reserve(16);
 232         }
 233         if v.capacity() == v.len() {
 234             v.reserve(1);
 235         }
 236         slice::from_raw_parts_mut(v.as_mut_ptr().add(v.len()), v.capacity() - v.len())
 237     }
 238 }
 239
 240 #[cfg(target_arch = "wasm32")]
 241 mod imp {
 242     use std::{
 243         io,
 244         process::{ChildStderr, ChildStdout},
 245     };
 246
 247     pub(crate) fn read2(
 248         _out_pipe: ChildStdout,
 249         _err_pipe: ChildStderr,
 250         _data: &mut dyn FnMut(bool, &mut Vec<u8>, bool),
 251     ) -> io::Result<()> {
 252         panic!("no processes on wasm")
 253     }
 254 }