authorgravatar for johnnymarler@gmail.comJonathan Marler <johnnymarler@gmail.com> 2023-02-28 14:10:44-07:00
committergravatar for andrew@ziglang.orgAndrew Kelley <andrew@ziglang.org> 2023-03-01 12:21:53-05:00
log138e8b162aeecc69cc62f0822e73b1862b7d07f6
tree3aed1af1481222e9a4406dad9da039960e345084
parent4f58a80735b47e6b98ed6f73cc9a0a772cdc6fcd

std.child_process: use std.io.poll for collectOutput

1 files changed, 31 insertions(+), 179 deletions(-)

lib/std/child_process.zig+31-179
...@@ -197,6 +197,19 @@ pub const ChildProcess = struct {...@@ -197,6 +197,19 @@ pub const ChildProcess = struct {
197 stderr: []u8,197 stderr: []u8,
198 };198 };
199199
200 fn fifoToOwnedArrayList(fifo: *std.io.PollFifo) std.ArrayList(u8) {
201 if (fifo.head > 0) {
202 std.mem.copy(u8, fifo.buf[0..fifo.count], fifo.buf[fifo.head .. fifo.head + fifo.count]);
203 }
204 const result = std.ArrayList(u8){
205 .items = fifo.buf[0..fifo.count],
206 .capacity = fifo.buf.len,
207 .allocator = fifo.allocator,
208 };
209 fifo.* = std.io.PollFifo.init(fifo.allocator);
210 return result;
211 }
212
200 /// Collect the output from the process's stdout and stderr. Will return once all output213 /// Collect the output from the process's stdout and stderr. Will return once all output
201 /// has been collected. This does not mean that the process has ended. `wait` should still214 /// has been collected. This does not mean that the process has ended. `wait` should still
202 /// be called to wait for and clean up the process.215 /// be called to wait for and clean up the process.
...@@ -210,189 +223,28 @@ pub const ChildProcess = struct {...@@ -210,189 +223,28 @@ pub const ChildProcess = struct {
210 ) !void {223 ) !void {
211 debug.assert(child.stdout_behavior == .Pipe);224 debug.assert(child.stdout_behavior == .Pipe);
212 debug.assert(child.stderr_behavior == .Pipe);225 debug.assert(child.stderr_behavior == .Pipe);
213 if (builtin.os.tag == .windows) {
214 try collectOutputWindows(child, stdout, stderr, max_output_bytes);
215 } else {
216 try collectOutputPosix(child, stdout, stderr, max_output_bytes);
217 }
218 }
219226
220 fn collectOutputPosix(227 // we could make this work with multiple allocators but YAGNI
221 child: ChildProcess,228 if (stdout.allocator.ptr != stderr.allocator.ptr or
222 stdout: *std.ArrayList(u8),229 stdout.allocator.vtable != stderr.allocator.vtable)
223 stderr: *std.ArrayList(u8),230 @panic("ChildProcess.collectOutput only supports 1 allocator");
224 max_output_bytes: usize,
225 ) !void {
226 var poll_fds = [_]os.pollfd{
227 .{ .fd = child.stdout.?.handle, .events = os.POLL.IN, .revents = undefined },
228 .{ .fd = child.stderr.?.handle, .events = os.POLL.IN, .revents = undefined },
229 };
230231
231 var dead_fds: usize = 0;232 var poller = std.io.poll(stdout.allocator, enum { stdout, stderr }, .{
232 // We ask for ensureTotalCapacity with this much extra space. This has more of an233 .stdout = child.stdout.?,
233 // effect on small reads because once the reads start to get larger the amount234 .stderr = child.stderr.?,
234 // of space an ArrayList will allocate grows exponentially.235 });
235 const bump_amt = 512;236 defer poller.deinit();
236237
237 const err_mask = os.POLL.ERR | os.POLL.NVAL | os.POLL.HUP;238 while (!poller.done()) {
238239 try poller.poll();
239 while (dead_fds < poll_fds.len) {240 if (poller.fifo(.stdout).count > max_output_bytes)
240 const events = try os.poll(&poll_fds, std.math.maxInt(i32));241 return error.StdoutStreamTooLong;
241 if (events == 0) continue;242 if (poller.fifo(.stderr).count > max_output_bytes)
242243 return error.StderrStreamTooLong;
243 var remove_stdout = false;
244 var remove_stderr = false;
245 // Try reading whatever is available before checking the error
246 // conditions.
247 // It's still possible to read after a POLL.HUP is received, always
248 // check if there's some data waiting to be read first.
249 if (poll_fds[0].revents & os.POLL.IN != 0) {
250 // stdout is ready.
251 const new_capacity = std.math.min(stdout.items.len + bump_amt, max_output_bytes);
252 try stdout.ensureTotalCapacity(new_capacity);
253 const buf = stdout.unusedCapacitySlice();
254 if (buf.len == 0) return error.StdoutStreamTooLong;
255 const nread = try os.read(poll_fds[0].fd, buf);
256 stdout.items.len += nread;
257
258 // Remove the fd when the EOF condition is met.
259 remove_stdout = nread == 0;
260 } else {
261 remove_stdout = poll_fds[0].revents & err_mask != 0;
262 }
263
264 if (poll_fds[1].revents & os.POLL.IN != 0) {
265 // stderr is ready.
266 const new_capacity = std.math.min(stderr.items.len + bump_amt, max_output_bytes);
267 try stderr.ensureTotalCapacity(new_capacity);
268 const buf = stderr.unusedCapacitySlice();
269 if (buf.len == 0) return error.StderrStreamTooLong;
270 const nread = try os.read(poll_fds[1].fd, buf);
271 stderr.items.len += nread;
272
273 // Remove the fd when the EOF condition is met.
274 remove_stderr = nread == 0;
275 } else {
276 remove_stderr = poll_fds[1].revents & err_mask != 0;
277 }
278
279 // Exclude the fds that signaled an error.
280 if (remove_stdout) {
281 poll_fds[0].fd = -1;
282 dead_fds += 1;
283 }
284 if (remove_stderr) {
285 poll_fds[1].fd = -1;
286 dead_fds += 1;
287 }
288 }
289 }
290
291 const WindowsAsyncReadResult = enum {
292 pending,
293 closed,
294 full,
295 };
296
297 fn windowsAsyncRead(
298 handle: windows.HANDLE,
299 overlapped: *windows.OVERLAPPED,
300 buf: *std.ArrayList(u8),
301 bump_amt: usize,
302 max_output_bytes: usize,
303 ) !WindowsAsyncReadResult {
304 while (true) {
305 const new_capacity = std.math.min(buf.items.len + bump_amt, max_output_bytes);
306 try buf.ensureTotalCapacity(new_capacity);
307 const next_buf = buf.unusedCapacitySlice();
308 if (next_buf.len == 0) return .full;
309 var read_bytes: u32 = undefined;
310 const read_result = windows.kernel32.ReadFile(handle, next_buf.ptr, math.cast(u32, next_buf.len) orelse maxInt(u32), &read_bytes, overlapped);
311 if (read_result == 0) return switch (windows.kernel32.GetLastError()) {
312 .IO_PENDING => .pending,
313 .BROKEN_PIPE => .closed,
314 else => |err| windows.unexpectedError(err),
315 };
316 buf.items.len += read_bytes;
317 }
318 }
319
320 fn collectOutputWindows(child: ChildProcess, stdout: *std.ArrayList(u8), stderr: *std.ArrayList(u8), max_output_bytes: usize) !void {
321 const bump_amt = 512;
322 const outs = [_]*std.ArrayList(u8){
323 stdout,
324 stderr,
325 };
326 const handles = [_]windows.HANDLE{
327 child.stdout.?.handle,
328 child.stderr.?.handle,
329 };
330
331 var overlapped = [_]windows.OVERLAPPED{
332 mem.zeroes(windows.OVERLAPPED),
333 mem.zeroes(windows.OVERLAPPED),
334 };
335
336 var wait_objects: [2]windows.HANDLE = undefined;
337 var wait_object_count: u2 = 0;
338
339 // we need to cancel all pending IO before returning so our OVERLAPPED values don't go out of scope
340 defer for (wait_objects[0..wait_object_count]) |o| {
341 _ = windows.kernel32.CancelIo(o);
342 };
343
344 // Windows Async IO requires an initial call to ReadFile before waiting on the handle
345 for ([_]u1{ 0, 1 }) |i| {
346 switch (try windowsAsyncRead(handles[i], &overlapped[i], outs[i], bump_amt, max_output_bytes)) {
347 .pending => {
348 wait_objects[wait_object_count] = handles[i];
349 wait_object_count += 1;
350 },
351 .closed => {}, // don't add to the wait_objects list
352 .full => return if (i == 0) error.StdoutStreamTooLong else error.StderrStreamTooLong,
353 }
354 }244 }
355245
356 while (wait_object_count > 0) {246 stdout.* = fifoToOwnedArrayList(poller.fifo(.stdout));
357 const status = windows.kernel32.WaitForMultipleObjects(wait_object_count, &wait_objects, 0, windows.INFINITE);247 stderr.* = fifoToOwnedArrayList(poller.fifo(.stderr));
358 if (status == windows.WAIT_FAILED) {
359 switch (windows.kernel32.GetLastError()) {
360 else => |err| return windows.unexpectedError(err),
361 }
362 }
363 if (status < windows.WAIT_OBJECT_0 or status > windows.WAIT_OBJECT_0 + wait_object_count - 1)
364 unreachable;
365
366 const wait_idx = status - windows.WAIT_OBJECT_0;
367
368 // this extra `i` index is needed to map the wait handle back to the stdout or stderr
369 // values since the wait_idx can change which handle it corresponds with
370 const i: u1 = if (wait_objects[wait_idx] == handles[0]) 0 else 1;
371
372 // remove completed event from the wait list
373 wait_object_count -= 1;
374 if (wait_idx == 0)
375 wait_objects[0] = wait_objects[1];
376
377 var read_bytes: u32 = undefined;
378 if (windows.kernel32.GetOverlappedResult(handles[i], &overlapped[i], &read_bytes, 0) == 0) {
379 switch (windows.kernel32.GetLastError()) {
380 .BROKEN_PIPE => continue,
381 else => |err| return windows.unexpectedError(err),
382 }
383 }
384
385 outs[i].items.len += read_bytes;
386
387 switch (try windowsAsyncRead(handles[i], &overlapped[i], outs[i], bump_amt, max_output_bytes)) {
388 .pending => {
389 wait_objects[wait_object_count] = handles[i];
390 wait_object_count += 1;
391 },
392 .closed => {}, // don't add to the wait_objects list
393 .full => return if (i == 0) error.StdoutStreamTooLong else error.StderrStreamTooLong,
394 }
395 }
396 }248 }
397249
398 /// Spawns a child process, waits for it, collecting stdout and stderr, and then returns.250 /// Spawns a child process, waits for it, collecting stdout and stderr, and then returns.