authorgravatar for andrew@ziglang.orgAndrew Kelley <andrew@ziglang.org> 2023-03-01 18:08:43-05:00
committergravatar for noreply@github.comGitHub <noreply@github.com> 2023-03-01 18:08:43-05:00
log874d3a17ae0f270b3e3f0ece7839d483d749107d
tree42f724ff6e82b0c90bf3eb9f4600809ae9efeeb7
parent25b83188d06d4cc760c723db0f6ec65db96373f5
parentf2b15420ad595f77b6a3575dd7a6e85411dc69e9
signaturebadge-question-mark Signed by PGP key 4AEE18F83AFDEB23

Merge pull request #14744 from ziglang/std.io.poll

introduce std.io.poll

3 files changed, 280 insertions(+), 184 deletions(-)

lib/std/child_process.zig+29-184
......@@ -197,6 +197,19 @@ pub const ChildProcess = struct {
197197 stderr: []u8,
198198 };
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
200213 /// Collect the output from the process's stdout and stderr. Will return once all output
201214 /// has been collected. This does not mean that the process has ended. `wait` should still
202215 /// be called to wait for and clean up the process.
......@@ -210,195 +223,27 @@ pub const ChildProcess = struct {
210223 ) !void {
211224 debug.assert(child.stdout_behavior == .Pipe);
212225 debug.assert(child.stderr_behavior == .Pipe);
213 if (builtin.os.tag == .haiku) {
214 const stdout_in = child.stdout.?.reader();
215 const stderr_in = child.stderr.?.reader();
216
217 try stdout_in.readAllArrayList(stdout, max_output_bytes);
218 try stderr_in.readAllArrayList(stderr, max_output_bytes);
219 } else if (builtin.os.tag == .windows) {
220 try collectOutputWindows(child, stdout, stderr, max_output_bytes);
221 } else {
222 try collectOutputPosix(child, stdout, stderr, max_output_bytes);
223 }
224 }
225
226 fn collectOutputPosix(
227 child: ChildProcess,
228 stdout: *std.ArrayList(u8),
229 stderr: *std.ArrayList(u8),
230 max_output_bytes: usize,
231 ) !void {
232 var poll_fds = [_]os.pollfd{
233 .{ .fd = child.stdout.?.handle, .events = os.POLL.IN, .revents = undefined },
234 .{ .fd = child.stderr.?.handle, .events = os.POLL.IN, .revents = undefined },
235 };
236
237 var dead_fds: usize = 0;
238 // We ask for ensureTotalCapacity with this much extra space. This has more of an
239 // effect on small reads because once the reads start to get larger the amount
240 // of space an ArrayList will allocate grows exponentially.
241 const bump_amt = 512;
242
243 const err_mask = os.POLL.ERR | os.POLL.NVAL | os.POLL.HUP;
244
245 while (dead_fds < poll_fds.len) {
246 const events = try os.poll(&poll_fds, std.math.maxInt(i32));
247 if (events == 0) continue;
248
249 var remove_stdout = false;
250 var remove_stderr = false;
251 // Try reading whatever is available before checking the error
252 // conditions.
253 // It's still possible to read after a POLL.HUP is received, always
254 // check if there's some data waiting to be read first.
255 if (poll_fds[0].revents & os.POLL.IN != 0) {
256 // stdout is ready.
257 const new_capacity = std.math.min(stdout.items.len + bump_amt, max_output_bytes);
258 try stdout.ensureTotalCapacity(new_capacity);
259 const buf = stdout.unusedCapacitySlice();
260 if (buf.len == 0) return error.StdoutStreamTooLong;
261 const nread = try os.read(poll_fds[0].fd, buf);
262 stdout.items.len += nread;
263
264 // Remove the fd when the EOF condition is met.
265 remove_stdout = nread == 0;
266 } else {
267 remove_stdout = poll_fds[0].revents & err_mask != 0;
268 }
269
270 if (poll_fds[1].revents & os.POLL.IN != 0) {
271 // stderr is ready.
272 const new_capacity = std.math.min(stderr.items.len + bump_amt, max_output_bytes);
273 try stderr.ensureTotalCapacity(new_capacity);
274 const buf = stderr.unusedCapacitySlice();
275 if (buf.len == 0) return error.StderrStreamTooLong;
276 const nread = try os.read(poll_fds[1].fd, buf);
277 stderr.items.len += nread;
278
279 // Remove the fd when the EOF condition is met.
280 remove_stderr = nread == 0;
281 } else {
282 remove_stderr = poll_fds[1].revents & err_mask != 0;
283 }
284226
285 // Exclude the fds that signaled an error.
286 if (remove_stdout) {
287 poll_fds[0].fd = -1;
288 dead_fds += 1;
289 }
290 if (remove_stderr) {
291 poll_fds[1].fd = -1;
292 dead_fds += 1;
293 }
294 }
295 }
227 // we could make this work with multiple allocators but YAGNI
228 if (stdout.allocator.ptr != stderr.allocator.ptr or
229 stdout.allocator.vtable != stderr.allocator.vtable)
230 @panic("ChildProcess.collectOutput only supports 1 allocator");
296231
297 const WindowsAsyncReadResult = enum {
298 pending,
299 closed,
300 full,
301 };
302
303 fn windowsAsyncRead(
304 handle: windows.HANDLE,
305 overlapped: *windows.OVERLAPPED,
306 buf: *std.ArrayList(u8),
307 bump_amt: usize,
308 max_output_bytes: usize,
309 ) !WindowsAsyncReadResult {
310 while (true) {
311 const new_capacity = std.math.min(buf.items.len + bump_amt, max_output_bytes);
312 try buf.ensureTotalCapacity(new_capacity);
313 const next_buf = buf.unusedCapacitySlice();
314 if (next_buf.len == 0) return .full;
315 var read_bytes: u32 = undefined;
316 const read_result = windows.kernel32.ReadFile(handle, next_buf.ptr, math.cast(u32, next_buf.len) orelse maxInt(u32), &read_bytes, overlapped);
317 if (read_result == 0) return switch (windows.kernel32.GetLastError()) {
318 .IO_PENDING => .pending,
319 .BROKEN_PIPE => .closed,
320 else => |err| windows.unexpectedError(err),
321 };
322 buf.items.len += read_bytes;
323 }
324 }
325
326 fn collectOutputWindows(child: ChildProcess, stdout: *std.ArrayList(u8), stderr: *std.ArrayList(u8), max_output_bytes: usize) !void {
327 const bump_amt = 512;
328 const outs = [_]*std.ArrayList(u8){
329 stdout,
330 stderr,
331 };
332 const handles = [_]windows.HANDLE{
333 child.stdout.?.handle,
334 child.stderr.?.handle,
335 };
336
337 var overlapped = [_]windows.OVERLAPPED{
338 mem.zeroes(windows.OVERLAPPED),
339 mem.zeroes(windows.OVERLAPPED),
340 };
341
342 var wait_objects: [2]windows.HANDLE = undefined;
343 var wait_object_count: u2 = 0;
344
345 // we need to cancel all pending IO before returning so our OVERLAPPED values don't go out of scope
346 defer for (wait_objects[0..wait_object_count]) |o| {
347 _ = windows.kernel32.CancelIo(o);
348 };
232 var poller = std.io.poll(stdout.allocator, enum { stdout, stderr }, .{
233 .stdout = child.stdout.?,
234 .stderr = child.stderr.?,
235 });
236 defer poller.deinit();
349237
350 // Windows Async IO requires an initial call to ReadFile before waiting on the handle
351 for ([_]u1{ 0, 1 }) |i| {
352 switch (try windowsAsyncRead(handles[i], &overlapped[i], outs[i], bump_amt, max_output_bytes)) {
353 .pending => {
354 wait_objects[wait_object_count] = handles[i];
355 wait_object_count += 1;
356 },
357 .closed => {}, // don't add to the wait_objects list
358 .full => return if (i == 0) error.StdoutStreamTooLong else error.StderrStreamTooLong,
359 }
238 while (try poller.poll()) {
239 if (poller.fifo(.stdout).count > max_output_bytes)
240 return error.StdoutStreamTooLong;
241 if (poller.fifo(.stderr).count > max_output_bytes)
242 return error.StderrStreamTooLong;
360243 }
361244
362 while (wait_object_count > 0) {
363 const status = windows.kernel32.WaitForMultipleObjects(wait_object_count, &wait_objects, 0, windows.INFINITE);
364 if (status == windows.WAIT_FAILED) {
365 switch (windows.kernel32.GetLastError()) {
366 else => |err| return windows.unexpectedError(err),
367 }
368 }
369 if (status < windows.WAIT_OBJECT_0 or status > windows.WAIT_OBJECT_0 + wait_object_count - 1)
370 unreachable;
371
372 const wait_idx = status - windows.WAIT_OBJECT_0;
373
374 // this extra `i` index is needed to map the wait handle back to the stdout or stderr
375 // values since the wait_idx can change which handle it corresponds with
376 const i: u1 = if (wait_objects[wait_idx] == handles[0]) 0 else 1;
377
378 // remove completed event from the wait list
379 wait_object_count -= 1;
380 if (wait_idx == 0)
381 wait_objects[0] = wait_objects[1];
382
383 var read_bytes: u32 = undefined;
384 if (windows.kernel32.GetOverlappedResult(handles[i], &overlapped[i], &read_bytes, 0) == 0) {
385 switch (windows.kernel32.GetLastError()) {
386 .BROKEN_PIPE => continue,
387 else => |err| return windows.unexpectedError(err),
388 }
389 }
390
391 outs[i].items.len += read_bytes;
392
393 switch (try windowsAsyncRead(handles[i], &overlapped[i], outs[i], bump_amt, max_output_bytes)) {
394 .pending => {
395 wait_objects[wait_object_count] = handles[i];
396 wait_object_count += 1;
397 },
398 .closed => {}, // don't add to the wait_objects list
399 .full => return if (i == 0) error.StdoutStreamTooLong else error.StderrStreamTooLong,
400 }
401 }
245 stdout.* = fifoToOwnedArrayList(poller.fifo(.stdout));
246 stderr.* = fifoToOwnedArrayList(poller.fifo(.stderr));
402247 }
403248
404249 /// Spawns a child process, waits for it, collecting stdout and stderr, and then returns.
lib/std/heap/general_purpose_allocator.zig+1
......@@ -423,6 +423,7 @@ pub fn GeneralPurposeAllocator(comptime config: Config) type {
423423 }
424424 } else struct {};
425425
426 /// Returns true if there were leaks; false otherwise.
426427 pub fn deinit(self: *Self) bool {
427428 const leaks = if (config.safety) self.detectLeaks() else false;
428429 if (config.retain_metadata) {
lib/std/io.zig+250
......@@ -168,6 +168,256 @@ test "null_writer" {
168168 null_writer.writeAll("yay" ** 10) catch |err| switch (err) {};
169169}
170170
171pub fn poll(
172 allocator: std.mem.Allocator,
173 comptime StreamEnum: type,
174 files: PollFiles(StreamEnum),
175) Poller(StreamEnum) {
176 const enum_fields = @typeInfo(StreamEnum).Enum.fields;
177 var result: Poller(StreamEnum) = undefined;
178
179 if (builtin.os.tag == .windows) result.windows = .{
180 .first_read_done = false,
181 .overlapped = [1]os.windows.OVERLAPPED{
182 mem.zeroes(os.windows.OVERLAPPED),
183 } ** enum_fields.len,
184 .active = .{
185 .count = 0,
186 .handles_buf = undefined,
187 .stream_map = undefined,
188 },
189 };
190
191 inline for (0..enum_fields.len) |i| {
192 result.fifos[i] = .{
193 .allocator = allocator,
194 .buf = &.{},
195 .head = 0,
196 .count = 0,
197 };
198 if (builtin.os.tag == .windows) {
199 result.windows.active.handles_buf[i] = @field(files, enum_fields[i].name).handle;
200 } else {
201 result.poll_fds[i] = .{
202 .fd = @field(files, enum_fields[i].name).handle,
203 .events = os.POLL.IN,
204 .revents = undefined,
205 };
206 }
207 }
208 return result;
209}
210
211pub const PollFifo = std.fifo.LinearFifo(u8, .Dynamic);
212
213pub fn Poller(comptime StreamEnum: type) type {
214 return struct {
215 const enum_fields = @typeInfo(StreamEnum).Enum.fields;
216 const PollFd = if (builtin.os.tag == .windows) void else std.os.pollfd;
217
218 fifos: [enum_fields.len]PollFifo,
219 poll_fds: [enum_fields.len]PollFd,
220 windows: if (builtin.os.tag == .windows) struct {
221 first_read_done: bool,
222 overlapped: [enum_fields.len]os.windows.OVERLAPPED,
223 active: struct {
224 count: math.IntFittingRange(0, enum_fields.len),
225 handles_buf: [enum_fields.len]os.windows.HANDLE,
226 stream_map: [enum_fields.len]StreamEnum,
227
228 pub fn removeAt(self: *@This(), index: u32) void {
229 std.debug.assert(index < self.count);
230 for (index + 1..self.count) |i| {
231 self.handles_buf[i - 1] = self.handles_buf[i];
232 self.stream_map[i - 1] = self.stream_map[i];
233 }
234 self.count -= 1;
235 }
236 },
237 } else void,
238
239 const Self = @This();
240
241 pub fn deinit(self: *Self) void {
242 if (builtin.os.tag == .windows) {
243 // cancel any pending IO to prevent clobbering OVERLAPPED value
244 for (self.windows.active.handles_buf[0..self.windows.active.count]) |h| {
245 _ = os.windows.kernel32.CancelIo(h);
246 }
247 }
248 inline for (&self.fifos) |*q| q.deinit();
249 self.* = undefined;
250 }
251
252 pub fn poll(self: *Self) !bool {
253 if (builtin.os.tag == .windows) {
254 return pollWindows(self);
255 } else {
256 return pollPosix(self);
257 }
258 }
259
260 pub inline fn fifo(self: *Self, comptime which: StreamEnum) *PollFifo {
261 return &self.fifos[@enumToInt(which)];
262 }
263
264 fn pollWindows(self: *Self) !bool {
265 const bump_amt = 512;
266
267 if (!self.windows.first_read_done) {
268 // Windows Async IO requires an initial call to ReadFile before waiting on the handle
269 for (0..enum_fields.len) |i| {
270 const handle = self.windows.active.handles_buf[i];
271 switch (try windowsAsyncRead(
272 handle,
273 &self.windows.overlapped[i],
274 &self.fifos[i],
275 bump_amt,
276 )) {
277 .pending => {
278 self.windows.active.handles_buf[self.windows.active.count] = handle;
279 self.windows.active.stream_map[self.windows.active.count] = @intToEnum(StreamEnum, i);
280 self.windows.active.count += 1;
281 },
282 .closed => {}, // don't add to the wait_objects list
283 }
284 }
285 self.windows.first_read_done = true;
286 }
287
288 while (true) {
289 if (self.windows.active.count == 0) return false;
290
291 const status = os.windows.kernel32.WaitForMultipleObjects(
292 self.windows.active.count,
293 &self.windows.active.handles_buf,
294 0,
295 os.windows.INFINITE,
296 );
297 if (status == os.windows.WAIT_FAILED)
298 return os.windows.unexpectedError(os.windows.kernel32.GetLastError());
299
300 if (status < os.windows.WAIT_OBJECT_0 or status > os.windows.WAIT_OBJECT_0 + enum_fields.len - 1)
301 unreachable;
302
303 const active_idx = status - os.windows.WAIT_OBJECT_0;
304
305 const handle = self.windows.active.handles_buf[active_idx];
306 const stream_idx = @enumToInt(self.windows.active.stream_map[active_idx]);
307 var read_bytes: u32 = undefined;
308 if (0 == os.windows.kernel32.GetOverlappedResult(
309 handle,
310 &self.windows.overlapped[stream_idx],
311 &read_bytes,
312 0,
313 )) switch (os.windows.kernel32.GetLastError()) {
314 .BROKEN_PIPE => {
315 self.windows.active.removeAt(active_idx);
316 continue;
317 },
318 else => |err| return os.windows.unexpectedError(err),
319 };
320
321 self.fifos[stream_idx].update(read_bytes);
322
323 switch (try windowsAsyncRead(
324 handle,
325 &self.windows.overlapped[stream_idx],
326 &self.fifos[stream_idx],
327 bump_amt,
328 )) {
329 .pending => {},
330 .closed => self.windows.active.removeAt(active_idx),
331 }
332 return true;
333 }
334 }
335
336 fn pollPosix(self: *Self) !bool {
337 // We ask for ensureUnusedCapacity with this much extra space. This
338 // has more of an effect on small reads because once the reads
339 // start to get larger the amount of space an ArrayList will
340 // allocate grows exponentially.
341 const bump_amt = 512;
342
343 const err_mask = os.POLL.ERR | os.POLL.NVAL | os.POLL.HUP;
344
345 const events_len = try os.poll(&self.poll_fds, std.math.maxInt(i32));
346 if (events_len == 0) {
347 for (self.poll_fds) |poll_fd| {
348 if (poll_fd.fd != -1) return true;
349 } else return false;
350 }
351
352 var keep_polling = false;
353 inline for (&self.poll_fds, &self.fifos) |*poll_fd, *q| {
354 // Try reading whatever is available before checking the error
355 // conditions.
356 // It's still possible to read after a POLL.HUP is received,
357 // always check if there's some data waiting to be read first.
358 if (poll_fd.revents & os.POLL.IN != 0) {
359 const buf = try q.writableWithSize(bump_amt);
360 const amt = try os.read(poll_fd.fd, buf);
361 q.update(amt);
362 if (amt == 0) {
363 // Remove the fd when the EOF condition is met.
364 poll_fd.fd = -1;
365 } else {
366 keep_polling = true;
367 }
368 } else if (poll_fd.revents & err_mask != 0) {
369 // Exclude the fds that signaled an error.
370 poll_fd.fd = -1;
371 } else if (poll_fd.fd != -1) {
372 keep_polling = true;
373 }
374 }
375 return keep_polling;
376 }
377 };
378}
379
380fn windowsAsyncRead(
381 handle: os.windows.HANDLE,
382 overlapped: *os.windows.OVERLAPPED,
383 fifo: *PollFifo,
384 bump_amt: usize,
385) !enum { pending, closed } {
386 while (true) {
387 const buf = try fifo.writableWithSize(bump_amt);
388 var read_bytes: u32 = undefined;
389 const read_result = os.windows.kernel32.ReadFile(handle, buf.ptr, math.cast(u32, buf.len) orelse math.maxInt(u32), &read_bytes, overlapped);
390 if (read_result == 0) return switch (os.windows.kernel32.GetLastError()) {
391 .IO_PENDING => .pending,
392 .BROKEN_PIPE => .closed,
393 else => |err| os.windows.unexpectedError(err),
394 };
395 fifo.update(read_bytes);
396 }
397}
398
399/// Given an enum, returns a struct with fields of that enum, each field
400/// representing an I/O stream for polling.
401pub fn PollFiles(comptime StreamEnum: type) type {
402 const enum_fields = @typeInfo(StreamEnum).Enum.fields;
403 var struct_fields: [enum_fields.len]std.builtin.Type.StructField = undefined;
404 for (&struct_fields, enum_fields) |*struct_field, enum_field| {
405 struct_field.* = .{
406 .name = enum_field.name,
407 .type = fs.File,
408 .default_value = null,
409 .is_comptime = false,
410 .alignment = @alignOf(fs.File),
411 };
412 }
413 return @Type(.{ .Struct = .{
414 .layout = .Auto,
415 .fields = &struct_fields,
416 .decls = &.{},
417 .is_tuple = false,
418 } });
419}
420
171421test {
172422 _ = @import("io/bit_reader.zig");
173423 _ = @import("io/bit_writer.zig");