src/tools/compiletest/src/read2.rs

   1 // FIXME: This is a complete copy of `cargo/src/cargo/util/read2.rs`
   2 // Consider unify the read2() in libstd, cargo and this to prevent further code duplication.
   3
   4 pub use self::imp::read2;
   5 use std::io;
   6 use std::process::{Child, Output};
   7
   8 pub fn read2_abbreviated(mut child: Child) -> io::Result<Output> {
   9     use io::Write;
  10     use std::mem::replace;
  11
  12     const HEAD_LEN: usize = 160 * 1024;
  13     const TAIL_LEN: usize = 256 * 1024;
  14
  15     enum ProcOutput {
  16         Full(Vec<u8>),
  17         Abbreviated { head: Vec<u8>, skipped: usize, tail: Box<[u8]> },
  18     }
  19
  20     impl ProcOutput {
  21         fn extend(&mut self, data: &[u8]) {
  22             let new_self = match *self {
  23                 ProcOutput::Full(ref mut bytes) => {
  24                     bytes.extend_from_slice(data);
  25                     let new_len = bytes.len();
  26                     if new_len <= HEAD_LEN + TAIL_LEN {
  27                         return;
  28                     }
  29                     let tail = bytes.split_off(new_len - TAIL_LEN).into_boxed_slice();
  30                     let head = replace(bytes, Vec::new());
  31                     let skipped = new_len - HEAD_LEN - TAIL_LEN;
  32                     ProcOutput::Abbreviated { head, skipped, tail }
  33                 }
  34                 ProcOutput::Abbreviated { ref mut skipped, ref mut tail, .. } => {
  35                     *skipped += data.len();
  36                     if data.len() <= TAIL_LEN {
  37                         tail[..data.len()].copy_from_slice(data);
  38                         tail.rotate_left(data.len());
  39                     } else {
  40                         tail.copy_from_slice(&data[(data.len() - TAIL_LEN)..]);
  41                     }
  42                     return;
  43                 }
  44             };
  45             *self = new_self;
  46         }
  47
  48         fn into_bytes(self) -> Vec<u8> {
  49             match self {
  50                 ProcOutput::Full(bytes) => bytes,
  51                 ProcOutput::Abbreviated { mut head, skipped, tail } => {
  52                     write!(&mut head, "\n\n<<<<<< SKIPPED {} BYTES >>>>>>\n\n", skipped).unwrap();
  53                     head.extend_from_slice(&tail);
  54                     head
  55                 }
  56             }
  57         }
  58     }
  59
  60     let mut stdout = ProcOutput::Full(Vec::new());
  61     let mut stderr = ProcOutput::Full(Vec::new());
  62
  63     drop(child.stdin.take());
  64     read2(
  65         child.stdout.take().unwrap(),
  66         child.stderr.take().unwrap(),
  67         &mut |is_stdout, data, _| {
  68             if is_stdout { &mut stdout } else { &mut stderr }.extend(data);
  69             data.clear();
  70         },
  71     )?;
  72     let status = child.wait()?;
  73
  74     Ok(Output { status, stdout: stdout.into_bytes(), stderr: stderr.into_bytes() })
  75 }
  76
  77 #[cfg(not(any(unix, windows)))]
  78 mod imp {
  79     use std::io::{self, Read};
  80     use std::process::{ChildStderr, ChildStdout};
  81
  82     pub fn read2(
  83         out_pipe: ChildStdout,
  84         err_pipe: ChildStderr,
  85         data: &mut dyn FnMut(bool, &mut Vec<u8>, bool),
  86     ) -> io::Result<()> {
  87         let mut buffer = Vec::new();
  88         out_pipe.read_to_end(&mut buffer)?;
  89         data(true, &mut buffer, true);
  90         buffer.clear();
  91         err_pipe.read_to_end(&mut buffer)?;
  92         data(false, &mut buffer, true);
  93         Ok(())
  94     }
  95 }
  96
  97 #[cfg(unix)]
  98 mod imp {
  99     use std::io;
 100     use std::io::prelude::*;
 101     use std::mem;
 102     use std::os::unix::prelude::*;
 103     use std::process::{ChildStderr, ChildStdout};
 104
 105     pub fn read2(
 106         mut out_pipe: ChildStdout,
 107         mut err_pipe: ChildStderr,
 108         data: &mut dyn FnMut(bool, &mut Vec<u8>, bool),
 109     ) -> io::Result<()> {
 110         unsafe {
 111             libc::fcntl(out_pipe.as_raw_fd(), libc::F_SETFL, libc::O_NONBLOCK);
 112             libc::fcntl(err_pipe.as_raw_fd(), libc::F_SETFL, libc::O_NONBLOCK);
 113         }
 114
 115         let mut out_done = false;
 116         let mut err_done = false;
 117         let mut out = Vec::new();
 118         let mut err = Vec::new();
 119
 120         let mut fds: [libc::pollfd; 2] = unsafe { mem::zeroed() };
 121         fds[0].fd = out_pipe.as_raw_fd();
 122         fds[0].events = libc::POLLIN;
 123         fds[1].fd = err_pipe.as_raw_fd();
 124         fds[1].events = libc::POLLIN;
 125         let mut nfds = 2;
 126         let mut errfd = 1;
 127
 128         while nfds > 0 {
 129             // wait for either pipe to become readable using `select`
 130             let r = unsafe { libc::poll(fds.as_mut_ptr(), nfds, -1) };
 131             if r == -1 {
 132                 let err = io::Error::last_os_error();
 133                 if err.kind() == io::ErrorKind::Interrupted {
 134                     continue;
 135                 }
 136                 return Err(err);
 137             }
 138
 139             // Read as much as we can from each pipe, ignoring EWOULDBLOCK or
 140             // EAGAIN. If we hit EOF, then this will happen because the underlying
 141             // reader will return Ok(0), in which case we'll see `Ok` ourselves. In
 142             // this case we flip the other fd back into blocking mode and read
 143             // whatever's leftover on that file descriptor.
 144             let handle = |res: io::Result<_>| match res {
 145                 Ok(_) => Ok(true),
 146                 Err(e) => {
 147                     if e.kind() == io::ErrorKind::WouldBlock {
 148                         Ok(false)
 149                     } else {
 150                         Err(e)
 151                     }
 152                 }
 153             };
 154             if !err_done && fds[errfd].revents != 0 && handle(err_pipe.read_to_end(&mut err))? {
 155                 err_done = true;
 156                 nfds -= 1;
 157             }
 158             data(false, &mut err, err_done);
 159             if !out_done && fds[0].revents != 0 && handle(out_pipe.read_to_end(&mut out))? {
 160                 out_done = true;
 161                 fds[0].fd = err_pipe.as_raw_fd();
 162                 errfd = 0;
 163                 nfds -= 1;
 164             }
 165             data(true, &mut out, out_done);
 166         }
 167         Ok(())
 168     }
 169 }
 170
 171 #[cfg(windows)]
 172 mod imp {
 173     use std::io;
 174     use std::os::windows::prelude::*;
 175     use std::process::{ChildStderr, ChildStdout};
 176     use std::slice;
 177
 178     use miow::iocp::{CompletionPort, CompletionStatus};
 179     use miow::pipe::NamedPipe;
 180     use miow::Overlapped;
 181     use winapi::shared::winerror::ERROR_BROKEN_PIPE;
 182
 183     struct Pipe<'a> {
 184         dst: &'a mut Vec<u8>,
 185         overlapped: Overlapped,
 186         pipe: NamedPipe,
 187         done: bool,
 188     }
 189
 190     pub fn read2(
 191         out_pipe: ChildStdout,
 192         err_pipe: ChildStderr,
 193         data: &mut dyn FnMut(bool, &mut Vec<u8>, bool),
 194     ) -> io::Result<()> {
 195         let mut out = Vec::new();
 196         let mut err = Vec::new();
 197
 198         let port = CompletionPort::new(1)?;
 199         port.add_handle(0, &out_pipe)?;
 200         port.add_handle(1, &err_pipe)?;
 201
 202         unsafe {
 203             let mut out_pipe = Pipe::new(out_pipe, &mut out);
 204             let mut err_pipe = Pipe::new(err_pipe, &mut err);
 205
 206             out_pipe.read()?;
 207             err_pipe.read()?;
 208
 209             let mut status = [CompletionStatus::zero(), CompletionStatus::zero()];
 210
 211             while !out_pipe.done || !err_pipe.done {
 212                 for status in port.get_many(&mut status, None)? {
 213                     if status.token() == 0 {
 214                         out_pipe.complete(status);
 215                         data(true, out_pipe.dst, out_pipe.done);
 216                         out_pipe.read()?;
 217                     } else {
 218                         err_pipe.complete(status);
 219                         data(false, err_pipe.dst, err_pipe.done);
 220                         err_pipe.read()?;
 221                     }
 222                 }
 223             }
 224
 225             Ok(())
 226         }
 227     }
 228
 229     impl<'a> Pipe<'a> {
 230         unsafe fn new<P: IntoRawHandle>(p: P, dst: &'a mut Vec<u8>) -> Pipe<'a> {
 231             Pipe {
 232                 dst: dst,
 233                 pipe: NamedPipe::from_raw_handle(p.into_raw_handle()),
 234                 overlapped: Overlapped::zero(),
 235                 done: false,
 236             }
 237         }
 238
 239         unsafe fn read(&mut self) -> io::Result<()> {
 240             let dst = slice_to_end(self.dst);
 241             match self.pipe.read_overlapped(dst, self.overlapped.raw()) {
 242                 Ok(_) => Ok(()),
 243                 Err(e) => {
 244                     if e.raw_os_error() == Some(ERROR_BROKEN_PIPE as i32) {
 245                         self.done = true;
 246                         Ok(())
 247                     } else {
 248                         Err(e)
 249                     }
 250                 }
 251             }
 252         }
 253
 254         unsafe fn complete(&mut self, status: &CompletionStatus) {
 255             let prev = self.dst.len();
 256             self.dst.set_len(prev + status.bytes_transferred() as usize);
 257             if status.bytes_transferred() == 0 {
 258                 self.done = true;
 259             }
 260         }
 261     }
 262
 263     unsafe fn slice_to_end(v: &mut Vec<u8>) -> &mut [u8] {
 264         if v.capacity() == 0 {
 265             v.reserve(16);
 266         }
 267         if v.capacity() == v.len() {
 268             v.reserve(1);
 269         }
 270         slice::from_raw_parts_mut(v.as_mut_ptr().offset(v.len() as isize), v.capacity() - v.len())
 271     }
 272 }