aboutsummaryrefslogtreecommitdiff
path: root/crates/stdx
diff options
context:
space:
mode:
Diffstat (limited to 'crates/stdx')
-rw-r--r--crates/stdx/Cargo.toml5
-rw-r--r--crates/stdx/src/lib.rs27
-rw-r--r--crates/stdx/src/process.rs238
3 files changed, 263 insertions, 7 deletions
diff --git a/crates/stdx/Cargo.toml b/crates/stdx/Cargo.toml
index d28b5e658..f78c5da7c 100644
--- a/crates/stdx/Cargo.toml
+++ b/crates/stdx/Cargo.toml
@@ -10,10 +10,15 @@ edition = "2018"
10doctest = false 10doctest = false
11 11
12[dependencies] 12[dependencies]
13libc = "0.2.93"
13backtrace = { version = "0.3.44", optional = true } 14backtrace = { version = "0.3.44", optional = true }
14always-assert = { version = "0.1.2", features = ["log"] } 15always-assert = { version = "0.1.2", features = ["log"] }
15# Think twice before adding anything here 16# Think twice before adding anything here
16 17
18[target.'cfg(windows)'.dependencies]
19miow = "0.3.6"
20winapi = "0.3.9"
21
17[features] 22[features]
18# Uncomment to enable for the whole crate graph 23# Uncomment to enable for the whole crate graph
19# default = [ "backtrace" ] 24# default = [ "backtrace" ]
diff --git a/crates/stdx/src/lib.rs b/crates/stdx/src/lib.rs
index b0a18d58d..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.
2use std::{cmp::Ordering, ops, process, time::Instant}; 2use std::{cmp::Ordering, ops, time::Instant};
3 3
4mod macros; 4mod macros;
5pub mod process;
5pub mod panic_context; 6pub mod panic_context;
6 7
7pub use always_assert::{always, never}; 8pub use always_assert::{always, never};
@@ -178,18 +179,30 @@ where
178 start..start + len 179 start..start + len
179} 180}
180 181
182pub 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
181#[repr(transparent)] 194#[repr(transparent)]
182pub struct JodChild(pub process::Child); 195pub struct JodChild(pub std::process::Child);
183 196
184impl ops::Deref for JodChild { 197impl ops::Deref for JodChild {
185 type Target = process::Child; 198 type Target = std::process::Child;
186 fn deref(&self) -> &process::Child { 199 fn deref(&self) -> &std::process::Child {
187 &self.0 200 &self.0
188 } 201 }
189} 202}
190 203
191impl ops::DerefMut for JodChild { 204impl ops::DerefMut for JodChild {
192 fn deref_mut(&mut self) -> &mut process::Child { 205 fn deref_mut(&mut self) -> &mut std::process::Child {
193 &mut self.0 206 &mut self.0
194 } 207 }
195} 208}
@@ -202,9 +215,9 @@ impl Drop for JodChild {
202} 215}
203 216
204impl JodChild { 217impl JodChild {
205 pub fn into_inner(self) -> process::Child { 218 pub fn into_inner(self) -> std::process::Child {
206 // SAFETY: repr transparent 219 // SAFETY: repr transparent
207 unsafe { std::mem::transmute::<JodChild, process::Child>(self) } 220 unsafe { std::mem::transmute::<JodChild, std::process::Child>(self) }
208 } 221 }
209} 222}
210 223
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
6use std::{
7 io,
8 process::{Command, Output, Stdio},
9};
10
11pub 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)]
59mod 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)]
134mod 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}