diff options
Diffstat (limited to 'crates/stdx/src')
-rw-r--r-- | crates/stdx/src/lib.rs | 31 | ||||
-rw-r--r-- | crates/stdx/src/process.rs | 238 |
2 files changed, 264 insertions, 5 deletions
diff --git a/crates/stdx/src/lib.rs b/crates/stdx/src/lib.rs index d26be4853..857567a85 100644 --- a/crates/stdx/src/lib.rs +++ b/crates/stdx/src/lib.rs | |||
@@ -1,7 +1,8 @@ | |||
1 | //! Missing batteries for standard libraries. | 1 | //! Missing batteries for standard libraries. |
2 | use std::{cmp::Ordering, ops, process, time::Instant}; | 2 | use std::{cmp::Ordering, ops, time::Instant}; |
3 | 3 | ||
4 | mod macros; | 4 | mod macros; |
5 | pub mod process; | ||
5 | pub mod panic_context; | 6 | pub mod panic_context; |
6 | 7 | ||
7 | pub use always_assert::{always, never}; | 8 | pub use always_assert::{always, never}; |
@@ -178,17 +179,30 @@ where | |||
178 | start..start + len | 179 | start..start + len |
179 | } | 180 | } |
180 | 181 | ||
181 | pub struct JodChild(pub process::Child); | 182 | pub fn defer<F: FnOnce()>(f: F) -> impl Drop { |
183 | struct D<F: FnOnce()>(Option<F>); | ||
184 | impl<F: FnOnce()> Drop for D<F> { | ||
185 | fn drop(&mut self) { | ||
186 | if let Some(f) = self.0.take() { | ||
187 | f() | ||
188 | } | ||
189 | } | ||
190 | } | ||
191 | D(Some(f)) | ||
192 | } | ||
193 | |||
194 | #[repr(transparent)] | ||
195 | pub struct JodChild(pub std::process::Child); | ||
182 | 196 | ||
183 | impl ops::Deref for JodChild { | 197 | impl ops::Deref for JodChild { |
184 | type Target = process::Child; | 198 | type Target = std::process::Child; |
185 | fn deref(&self) -> &process::Child { | 199 | fn deref(&self) -> &std::process::Child { |
186 | &self.0 | 200 | &self.0 |
187 | } | 201 | } |
188 | } | 202 | } |
189 | 203 | ||
190 | impl ops::DerefMut for JodChild { | 204 | impl ops::DerefMut for JodChild { |
191 | fn deref_mut(&mut self) -> &mut process::Child { | 205 | fn deref_mut(&mut self) -> &mut std::process::Child { |
192 | &mut self.0 | 206 | &mut self.0 |
193 | } | 207 | } |
194 | } | 208 | } |
@@ -200,6 +214,13 @@ impl Drop for JodChild { | |||
200 | } | 214 | } |
201 | } | 215 | } |
202 | 216 | ||
217 | impl JodChild { | ||
218 | pub fn into_inner(self) -> std::process::Child { | ||
219 | // SAFETY: repr transparent | ||
220 | unsafe { std::mem::transmute::<JodChild, std::process::Child>(self) } | ||
221 | } | ||
222 | } | ||
223 | |||
203 | #[cfg(test)] | 224 | #[cfg(test)] |
204 | mod tests { | 225 | mod tests { |
205 | use super::*; | 226 | use super::*; |
diff --git a/crates/stdx/src/process.rs b/crates/stdx/src/process.rs new file mode 100644 index 000000000..b0fa12f76 --- /dev/null +++ b/crates/stdx/src/process.rs | |||
@@ -0,0 +1,238 @@ | |||
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 | } | ||