Commit f07116404a

Andrew Kelley <andrew@ziglang.org>
2024-05-23 23:10:03
std.Progress: child process sends updates via IPC
1 parent ed36470
Changed files (3)
lib/std/process/Child.zig
@@ -12,6 +12,7 @@ const EnvMap = std.process.EnvMap;
 const maxInt = std.math.maxInt;
 const assert = std.debug.assert;
 const native_os = builtin.os.tag;
+const Allocator = std.mem.Allocator;
 const ChildProcess = @This();
 
 pub const Id = switch (native_os) {
@@ -92,6 +93,13 @@ request_resource_usage_statistics: bool = false,
 /// `spawn`.
 resource_usage_statistics: ResourceUsageStatistics = .{},
 
+/// When populated, a pipe will be created for the child process to
+/// communicate progress back to the parent. The file descriptor of the
+/// write end of the pipe will be specified in the `ZIG_PROGRESS`
+/// environment variable inside the child process. The progress reported by
+/// the child will be attached to this progress node in the parent process.
+parent_progress_node: std.Progress.Node = .{ .index = .none },
+
 pub const ResourceUsageStatistics = struct {
     rusage: @TypeOf(rusage_init) = rusage_init,
 
@@ -572,6 +580,16 @@ fn spawnPosix(self: *ChildProcess) SpawnError!void {
         if (any_ignore) posix.close(dev_null_fd);
     }
 
+    const prog_pipe: [2]posix.fd_t = p: {
+        if (self.parent_progress_node.index == .none) {
+            break :p .{ -1, -1 };
+        } else {
+            // No CLOEXEC because the child needs access to this file descriptor.
+            break :p try posix.pipe2(.{});
+        }
+    };
+    errdefer destroyPipe(prog_pipe);
+
     var arena_allocator = std.heap.ArenaAllocator.init(self.allocator);
     defer arena_allocator.deinit();
     const arena = arena_allocator.allocator();
@@ -588,16 +606,35 @@ fn spawnPosix(self: *ChildProcess) SpawnError!void {
     const argv_buf = try arena.allocSentinel(?[*:0]const u8, self.argv.len, null);
     for (self.argv, 0..) |arg, i| argv_buf[i] = (try arena.dupeZ(u8, arg)).ptr;
 
-    const envp = m: {
+    const envp: [*:null]const ?[*:0]const u8 = m: {
+        const extra_usizes: []const process.CreateEnvironOptions.ExtraUsize = if (prog_pipe[1] == -1) &.{} else &.{
+            .{ .name = "ZIG_PROGRESS", .value = @intCast(prog_pipe[1]) },
+        };
         if (self.env_map) |env_map| {
-            const envp_buf = try process.createNullDelimitedEnvMap(arena, env_map);
-            break :m envp_buf.ptr;
+            break :m (try process.createEnviron(arena, .{
+                .env_map = env_map,
+                .extra_usizes = extra_usizes,
+            })).ptr;
         } else if (builtin.link_libc) {
-            break :m std.c.environ;
+            if (extra_usizes.len == 0) {
+                break :m std.c.environ;
+            } else {
+                break :m (try process.createEnviron(arena, .{
+                    .existing = std.c.environ,
+                    .extra_usizes = extra_usizes,
+                })).ptr;
+            }
         } else if (builtin.output_mode == .Exe) {
             // Then we have Zig start code and this works.
-            // TODO type-safety for null-termination of `os.environ`.
-            break :m @as([*:null]const ?[*:0]const u8, @ptrCast(std.os.environ.ptr));
+            if (extra_usizes.len == 0) {
+                break :m @ptrCast(std.os.environ.ptr);
+            } else {
+                break :m (try process.createEnviron(arena, .{
+                    // TODO type-safety for null-termination of `os.environ`.
+                    .existing = @ptrCast(std.os.environ.ptr),
+                    .extra_usizes = extra_usizes,
+                })).ptr;
+            }
         } else {
             // TODO come up with a solution for this.
             @compileError("missing std lib enhancement: ChildProcess implementation has no way to collect the environment variables to forward to the child process");
@@ -962,7 +999,7 @@ fn setUpChildIo(stdio: StdIo, pipe_fd: i32, std_fileno: i32, dev_null_fd: i32) !
 }
 
 fn destroyPipe(pipe: [2]posix.fd_t) void {
-    posix.close(pipe[0]);
+    if (pipe[0] != -1) posix.close(pipe[0]);
     if (pipe[0] != pipe[1]) posix.close(pipe[1]);
 }
 
lib/std/process.zig
@@ -431,6 +431,29 @@ pub fn hasEnvVarConstant(comptime key: []const u8) bool {
     }
 }
 
+pub const ParseEnvVarIntError = std.fmt.ParseIntError || error{EnvironmentVariableNotFound};
+
+/// Parses an environment variable as an integer.
+///
+/// Since the key is comptime-known, no allocation is needed.
+///
+/// On Windows, `key` must be valid UTF-8.
+pub fn parseEnvVarInt(comptime key: []const u8, comptime I: type, base: u8) ParseEnvVarIntError!I {
+    if (native_os == .windows) {
+        const key_w = comptime std.unicode.utf8ToUtf16LeStringLiteral(key);
+        const text = getenvW(key_w) orelse return error.EnvironmentVariableNotFound;
+        // For this implementation perhaps std.fmt.parseInt can be expanded to be generic across
+        // []u8 and []u16 like how many std.mem functions work.
+        _ = text;
+        @compileError("TODO implement this");
+    } else if (native_os == .wasi and !builtin.link_libc) {
+        @compileError("parseEnvVarInt is not supported for WASI without libc");
+    } else {
+        const text = posix.getenv(key) orelse return error.EnvironmentVariableNotFound;
+        return std.fmt.parseInt(I, text, base);
+    }
+}
+
 pub const HasEnvVarError = error{
     OutOfMemory,
 
@@ -1790,24 +1813,61 @@ test raiseFileDescriptorLimit {
     raiseFileDescriptorLimit();
 }
 
-pub fn createNullDelimitedEnvMap(arena: mem.Allocator, env_map: *const EnvMap) ![:null]?[*:0]u8 {
-    const envp_count = env_map.count();
+pub const CreateEnvironOptions = struct {
+    env_map: ?*const EnvMap = null,
+    existing: ?[*:null]const ?[*:0]const u8 = null,
+    extra_usizes: []const ExtraUsize = &.{},
+
+    pub const ExtraUsize = struct {
+        name: []const u8,
+        value: usize,
+    };
+};
+
+/// Creates a null-deliminated environment variable block in the format
+/// expected by POSIX, by combining all the sources of key-value pairs together
+/// from `options`.
+pub fn createEnviron(arena: Allocator, options: CreateEnvironOptions) Allocator.Error![:null]?[*:0]u8 {
+    const envp_count = c: {
+        var count: usize = 0;
+        if (options.existing) |env| {
+            while (env[count]) |_| : (count += 1) {}
+        }
+        if (options.env_map) |env_map| {
+            count += env_map.count();
+        }
+        count += options.extra_usizes.len;
+        break :c count;
+    };
     const envp_buf = try arena.allocSentinel(?[*:0]u8, envp_count, null);
-    {
+    var i: usize = 0;
+
+    if (options.existing) |env| {
+        while (env[i]) |line| : (i += 1) {
+            envp_buf[i] = try arena.dupeZ(u8, mem.span(line));
+        }
+    }
+
+    for (options.extra_usizes, envp_buf[i..][0..options.extra_usizes.len]) |extra_usize, *out| {
+        out.* = try std.fmt.allocPrintZ(arena, "{s}={d}", .{ extra_usize.name, extra_usize.value });
+    }
+    i += options.extra_usizes.len;
+
+    if (options.env_map) |env_map| {
         var it = env_map.iterator();
-        var i: usize = 0;
         while (it.next()) |pair| : (i += 1) {
-            const env_buf = try arena.allocSentinel(u8, pair.key_ptr.len + pair.value_ptr.len + 1, 0);
-            @memcpy(env_buf[0..pair.key_ptr.len], pair.key_ptr.*);
-            env_buf[pair.key_ptr.len] = '=';
-            @memcpy(env_buf[pair.key_ptr.len + 1 ..][0..pair.value_ptr.len], pair.value_ptr.*);
-            envp_buf[i] = env_buf.ptr;
+            envp_buf[i] = try std.fmt.allocPrintZ(arena, "{s}={s}", .{ pair.key_ptr.*, pair.value_ptr.* });
         }
-        assert(i == envp_count);
     }
+
+    assert(i == envp_count);
     return envp_buf;
 }
 
+pub fn createNullDelimitedEnvMap(arena: mem.Allocator, env_map: *const EnvMap) ![:null]?[*:0]u8 {
+    return createEnviron(arena, .{ .env_map = env_map });
+}
+
 test createNullDelimitedEnvMap {
     const allocator = testing.allocator;
     var envmap = EnvMap.init(allocator);
lib/std/Progress.zig
@@ -74,7 +74,7 @@ pub const Options = struct {
 pub const Node = struct {
     index: OptionalIndex,
 
-    pub const max_name_len = 38;
+    pub const max_name_len = 40;
 
     const Storage = extern struct {
         /// Little endian.
@@ -268,17 +268,7 @@ var node_freelist_buffer: [default_node_storage_buffer_len]Node.OptionalIndex =
 pub fn start(options: Options) Node {
     // Ensure there is only 1 global Progress object.
     assert(global_progress.node_end_index == 0);
-    const stderr = std.io.getStdErr();
-    if (stderr.supportsAnsiEscapeCodes()) {
-        global_progress.terminal = stderr;
-        global_progress.supports_ansi_escape_codes = true;
-    } else if (builtin.os.tag == .windows and stderr.isTty()) {
-        global_progress.is_windows_terminal = true;
-        global_progress.terminal = stderr;
-    } else if (builtin.os.tag != .windows) {
-        // we are in a "dumb" terminal like in acme or writing to a file
-        global_progress.terminal = stderr;
-    }
+
     @memset(global_progress.node_parents, .unused);
     const root_node = Node.init(@enumFromInt(0), .none, options.root_name, options.estimated_total_items);
     global_progress.done = false;
@@ -289,22 +279,51 @@ pub fn start(options: Options) Node {
     global_progress.refresh_rate_ns = options.refresh_rate_ns;
     global_progress.initial_delay_ns = options.initial_delay_ns;
 
-    var act: posix.Sigaction = .{
-        .handler = .{ .sigaction = handleSigWinch },
-        .mask = posix.empty_sigset,
-        .flags = (posix.SA.SIGINFO | posix.SA.RESTART),
-    };
-    posix.sigaction(posix.SIG.WINCH, &act, null) catch {
-        global_progress.terminal = null;
-        return root_node;
-    };
-
-    if (global_progress.terminal != null) {
-        if (std.Thread.spawn(.{}, updateThreadRun, .{})) |thread| {
+    if (std.process.parseEnvVarInt("ZIG_PROGRESS", u31, 10)) |ipc_fd| {
+        if (std.Thread.spawn(.{}, ipcThreadRun, .{ipc_fd})) |thread| {
             global_progress.update_thread = thread;
-        } else |_| {
-            global_progress.terminal = null;
+        } else |err| {
+            std.log.warn("failed to spawn IPC thread for communicating progress to parent: {s}", .{@errorName(err)});
+            return .{ .index = .none };
         }
+    } else |env_err| switch (env_err) {
+        error.EnvironmentVariableNotFound => {
+            const stderr = std.io.getStdErr();
+            if (stderr.supportsAnsiEscapeCodes()) {
+                global_progress.terminal = stderr;
+                global_progress.supports_ansi_escape_codes = true;
+            } else if (builtin.os.tag == .windows and stderr.isTty()) {
+                global_progress.is_windows_terminal = true;
+                global_progress.terminal = stderr;
+            } else if (builtin.os.tag != .windows) {
+                // we are in a "dumb" terminal like in acme or writing to a file
+                global_progress.terminal = stderr;
+            }
+
+            if (global_progress.terminal == null) {
+                return .{ .index = .none };
+            }
+
+            var act: posix.Sigaction = .{
+                .handler = .{ .sigaction = handleSigWinch },
+                .mask = posix.empty_sigset,
+                .flags = (posix.SA.SIGINFO | posix.SA.RESTART),
+            };
+            posix.sigaction(posix.SIG.WINCH, &act, null) catch |err| {
+                std.log.warn("failed to install SIGWINCH signal handler for noticing terminal resizes: {s}", .{@errorName(err)});
+            };
+
+            if (std.Thread.spawn(.{}, updateThreadRun, .{})) |thread| {
+                global_progress.update_thread = thread;
+            } else |err| {
+                std.log.warn("unable to spawn thread for printing progress to terminal: {s}", .{@errorName(err)});
+                return .{ .index = .none };
+            }
+        },
+        else => |e| {
+            std.log.warn("invalid ZIG_PROGRESS file descriptor integer: {s}", .{@errorName(e)});
+            return .{ .index = .none };
+        },
     }
 
     return root_node;
@@ -326,12 +345,10 @@ fn updateThreadRun() void {
         const resize_flag = wait(global_progress.initial_delay_ns);
         maybeUpdateSize(resize_flag);
 
-        const buffer = b: {
-            if (@atomicLoad(bool, &global_progress.done, .seq_cst))
-                return clearTerminal();
+        if (@atomicLoad(bool, &global_progress.done, .seq_cst))
+            return clearTerminal();
 
-            break :b computeRedraw();
-        };
+        const buffer = computeRedraw();
         write(buffer);
     }
 
@@ -339,16 +356,36 @@ fn updateThreadRun() void {
         const resize_flag = wait(global_progress.refresh_rate_ns);
         maybeUpdateSize(resize_flag);
 
-        const buffer = b: {
-            if (@atomicLoad(bool, &global_progress.done, .seq_cst))
-                return clearTerminal();
+        if (@atomicLoad(bool, &global_progress.done, .seq_cst))
+            return clearTerminal();
 
-            break :b computeRedraw();
-        };
+        const buffer = computeRedraw();
         write(buffer);
     }
 }
 
+fn ipcThreadRun(fd: posix.fd_t) void {
+    {
+        _ = wait(global_progress.initial_delay_ns);
+
+        if (@atomicLoad(bool, &global_progress.done, .seq_cst))
+            return;
+
+        const serialized = serialize();
+        writeIpc(fd, serialized);
+    }
+
+    while (true) {
+        _ = wait(global_progress.refresh_rate_ns);
+
+        if (@atomicLoad(bool, &global_progress.done, .seq_cst))
+            return clearTerminal();
+
+        const serialized = serialize();
+        writeIpc(fd, serialized);
+    }
+}
+
 const start_sync = "\x1b[?2026h";
 const up_one_line = "\x1bM";
 const clear = "\x1b[J";
@@ -400,11 +437,17 @@ const Children = struct {
     sibling: Node.OptionalIndex,
 };
 
-fn computeRedraw() []u8 {
-    // TODO make this configurable
-    var serialized_node_parents_buffer: [default_node_storage_buffer_len]Node.Parent = undefined;
-    var serialized_node_storage_buffer: [default_node_storage_buffer_len]Node.Storage = undefined;
-    var serialized_node_map_buffer: [default_node_storage_buffer_len]Node.Index = undefined;
+// TODO make this configurable
+var serialized_node_parents_buffer: [default_node_storage_buffer_len]Node.Parent = undefined;
+var serialized_node_storage_buffer: [default_node_storage_buffer_len]Node.Storage = undefined;
+var serialized_node_map_buffer: [default_node_storage_buffer_len]Node.Index = undefined;
+
+const Serialized = struct {
+    parents: []Node.Parent,
+    storage: []Node.Storage,
+};
+
+fn serialize() Serialized {
     var serialized_len: usize = 0;
 
     // Iterate all of the nodes and construct a serializable copy of the state that can be examined
@@ -447,12 +490,21 @@ fn computeRedraw() []u8 {
         };
     }
 
+    return .{
+        .parents = serialized_node_parents,
+        .storage = serialized_node_storage,
+    };
+}
+
+fn computeRedraw() []u8 {
+    const serialized = serialize();
+
     var children_buffer: [default_node_storage_buffer_len]Children = undefined;
-    const children = children_buffer[0..serialized_len];
+    const children = children_buffer[0..serialized.parents.len];
 
     @memset(children, .{ .child = .none, .sibling = .none });
 
-    for (serialized_node_parents, 0..) |parent, child_index_usize| {
+    for (serialized.parents, 0..) |parent, child_index_usize| {
         const child_index: Node.Index = @enumFromInt(child_index_usize);
         assert(parent != .unused);
         const parent_index = parent.unwrap() orelse continue;
@@ -478,7 +530,7 @@ fn computeRedraw() []u8 {
     i = computeClear(buf, i);
 
     const root_node_index: Node.Index = @enumFromInt(0);
-    i = computeNode(buf, i, serialized_node_storage, serialized_node_parents, children, root_node_index);
+    i = computeNode(buf, i, serialized, children, root_node_index);
 
     // Truncate trailing newline.
     if (buf[i - 1] == '\n') i -= 1;
@@ -492,15 +544,14 @@ fn computeRedraw() []u8 {
 fn computePrefix(
     buf: []u8,
     start_i: usize,
-    serialized_node_storage: []const Node.Storage,
-    serialized_node_parents: []const Node.Parent,
+    serialized: Serialized,
     children: []const Children,
     node_index: Node.Index,
 ) usize {
     var i = start_i;
-    const parent_index = serialized_node_parents[@intFromEnum(node_index)].unwrap() orelse return i;
-    if (serialized_node_parents[@intFromEnum(parent_index)] == .none) return i;
-    i = computePrefix(buf, i, serialized_node_storage, serialized_node_parents, children, parent_index);
+    const parent_index = serialized.parents[@intFromEnum(node_index)].unwrap() orelse return i;
+    if (serialized.parents[@intFromEnum(parent_index)] == .none) return i;
+    i = computePrefix(buf, i, serialized, children, parent_index);
     if (children[@intFromEnum(parent_index)].sibling == .none) {
         buf[i..][0..3].* = "   ".*;
         i += 3;
@@ -514,19 +565,18 @@ fn computePrefix(
 fn computeNode(
     buf: []u8,
     start_i: usize,
-    serialized_node_storage: []const Node.Storage,
-    serialized_node_parents: []const Node.Parent,
+    serialized: Serialized,
     children: []const Children,
     node_index: Node.Index,
 ) usize {
     var i = start_i;
-    i = computePrefix(buf, i, serialized_node_storage, serialized_node_parents, children, node_index);
+    i = computePrefix(buf, i, serialized, children, node_index);
 
-    const storage = &serialized_node_storage[@intFromEnum(node_index)];
+    const storage = &serialized.storage[@intFromEnum(node_index)];
     const estimated_total = storage.estimated_total_count;
     const completed_items = storage.completed_count;
     const name = if (std.mem.indexOfScalar(u8, &storage.name, 0)) |end| storage.name[0..end] else &storage.name;
-    const parent = serialized_node_parents[@intFromEnum(node_index)];
+    const parent = serialized.parents[@intFromEnum(node_index)];
 
     if (parent != .none) {
         if (children[@intFromEnum(node_index)].sibling == .none) {
@@ -555,11 +605,11 @@ fn computeNode(
     global_progress.newline_count += 1;
 
     if (children[@intFromEnum(node_index)].child.unwrap()) |child| {
-        i = computeNode(buf, i, serialized_node_storage, serialized_node_parents, children, child);
+        i = computeNode(buf, i, serialized, children, child);
     }
 
     if (children[@intFromEnum(node_index)].sibling.unwrap()) |sibling| {
-        i = computeNode(buf, i, serialized_node_storage, serialized_node_parents, children, sibling);
+        i = computeNode(buf, i, serialized, children, sibling);
     }
 
     return i;
@@ -572,6 +622,27 @@ fn write(buf: []const u8) void {
     };
 }
 
+fn writeIpc(fd: posix.fd_t, serialized: Serialized) void {
+    assert(serialized.parents.len == serialized.storage.len);
+    const header = std.mem.asBytes(&serialized.parents.len);
+    const storage = std.mem.sliceAsBytes(serialized.storage);
+    const parents = std.mem.sliceAsBytes(serialized.parents);
+
+    var vecs: [3]std.posix.iovec_const = .{
+        .{ .base = header.ptr, .len = header.len },
+        .{ .base = storage.ptr, .len = storage.len },
+        .{ .base = parents.ptr, .len = parents.len },
+    };
+
+    // TODO: if big endian, byteswap
+    // this is needed because the parent or child process might be running in qemu
+
+    const file: std.fs.File = .{ .handle = fd };
+    file.writevAll(&vecs) catch |err| {
+        std.log.warn("failed to send progress to parent process: {s}", .{@errorName(err)});
+    };
+}
+
 fn maybeUpdateSize(resize_flag: bool) void {
     if (!resize_flag) return;