Commit c01cfde688

Andrew Kelley <andrew@ziglang.org>
2024-05-24 05:22:58
std.process.Child: fix ZIG_PROGRESS env var handling
and properly dup2 the file descriptor to make it handle the case when other files are already open
1 parent 3a768bd
Changed files (2)
lib
lib/std/process/Child.zig
@@ -216,9 +216,9 @@ pub fn init(argv: []const []const u8, allocator: mem.Allocator) ChildProcess {
         .stdin = null,
         .stdout = null,
         .stderr = null,
-        .stdin_behavior = StdIo.Inherit,
-        .stdout_behavior = StdIo.Inherit,
-        .stderr_behavior = StdIo.Inherit,
+        .stdin_behavior = .Inherit,
+        .stdout_behavior = .Inherit,
+        .stderr_behavior = .Inherit,
         .expand_arg0 = .no_expand,
     };
 }
@@ -549,22 +549,22 @@ fn spawnPosix(self: *ChildProcess) SpawnError!void {
     // turns out, we `dup2` everything anyway, so there's no need!
     const pipe_flags: posix.O = .{ .CLOEXEC = true };
 
-    const stdin_pipe = if (self.stdin_behavior == StdIo.Pipe) try posix.pipe2(pipe_flags) else undefined;
-    errdefer if (self.stdin_behavior == StdIo.Pipe) {
+    const stdin_pipe = if (self.stdin_behavior == .Pipe) try posix.pipe2(pipe_flags) else undefined;
+    errdefer if (self.stdin_behavior == .Pipe) {
         destroyPipe(stdin_pipe);
     };
 
-    const stdout_pipe = if (self.stdout_behavior == StdIo.Pipe) try posix.pipe2(pipe_flags) else undefined;
-    errdefer if (self.stdout_behavior == StdIo.Pipe) {
+    const stdout_pipe = if (self.stdout_behavior == .Pipe) try posix.pipe2(pipe_flags) else undefined;
+    errdefer if (self.stdout_behavior == .Pipe) {
         destroyPipe(stdout_pipe);
     };
 
-    const stderr_pipe = if (self.stderr_behavior == StdIo.Pipe) try posix.pipe2(pipe_flags) else undefined;
-    errdefer if (self.stderr_behavior == StdIo.Pipe) {
+    const stderr_pipe = if (self.stderr_behavior == .Pipe) try posix.pipe2(pipe_flags) else undefined;
+    errdefer if (self.stderr_behavior == .Pipe) {
         destroyPipe(stderr_pipe);
     };
 
-    const any_ignore = (self.stdin_behavior == StdIo.Ignore or self.stdout_behavior == StdIo.Ignore or self.stderr_behavior == StdIo.Ignore);
+    const any_ignore = (self.stdin_behavior == .Ignore or self.stdout_behavior == .Ignore or self.stderr_behavior == .Ignore);
     const dev_null_fd = if (any_ignore)
         posix.openZ("/dev/null", .{ .ACCMODE = .RDWR }, 0) catch |err| switch (err) {
             error.PathAlreadyExists => unreachable,
@@ -609,35 +609,24 @@ 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 prog_fileno = 3;
+
     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]) },
-        };
+        const prog_fd: i32 = if (prog_pipe[1] == -1) -1 else prog_fileno;
         if (self.env_map) |env_map| {
-            break :m (try process.createEnviron(arena, .{
-                .env_map = env_map,
-                .extra_usizes = extra_usizes,
+            break :m (try process.createEnvironFromMap(arena, env_map, .{
+                .zig_progress_fd = prog_fd,
             })).ptr;
         } else if (builtin.link_libc) {
-            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;
-            }
+            break :m (try process.createEnvironFromExisting(arena, std.c.environ, .{
+                .zig_progress_fd = prog_fd,
+            })).ptr;
         } else if (builtin.output_mode == .Exe) {
             // Then we have Zig start code and this works.
-            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;
-            }
+            // TODO type-safety for null-termination of `os.environ`.
+            break :m (try process.createEnvironFromExisting(arena, @ptrCast(std.os.environ.ptr), .{
+                .zig_progress_fd = prog_fd,
+            })).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");
@@ -664,6 +653,12 @@ fn spawnPosix(self: *ChildProcess) SpawnError!void {
         setUpChildIo(self.stdin_behavior, stdin_pipe[0], posix.STDIN_FILENO, dev_null_fd) catch |err| forkChildErrReport(err_pipe[1], err);
         setUpChildIo(self.stdout_behavior, stdout_pipe[1], posix.STDOUT_FILENO, dev_null_fd) catch |err| forkChildErrReport(err_pipe[1], err);
         setUpChildIo(self.stderr_behavior, stderr_pipe[1], posix.STDERR_FILENO, dev_null_fd) catch |err| forkChildErrReport(err_pipe[1], err);
+        if (prog_pipe[1] != -1) posix.dup2(prog_pipe[1], prog_fileno) catch |err| forkChildErrReport(err_pipe[1], err);
+
+        if (prog_pipe[1] != -1) {
+            if (prog_pipe[0] != prog_fileno) posix.close(prog_pipe[0]);
+            if (prog_pipe[1] != prog_fileno) posix.close(prog_pipe[1]);
+        }
 
         if (self.cwd_dir) |cwd| {
             posix.fchdir(cwd.fd) catch |err| forkChildErrReport(err_pipe[1], err);
@@ -718,6 +713,9 @@ fn spawnPosix(self: *ChildProcess) SpawnError!void {
         posix.close(stderr_pipe[1]);
     }
 
+    if (prog_pipe[1] != -1) {
+        posix.close(prog_pipe[1]);
+    }
     self.progress_node.setIpcFd(prog_pipe[0]);
 }
 
lib/std/process.zig
@@ -1814,58 +1814,132 @@ test raiseFileDescriptorLimit {
 }
 
 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,
-    };
+    /// `null` means to leave the `ZIG_PROGRESS` environment variable unmodified.
+    /// If non-null, negative means to remove the environment variable, and >= 0
+    /// means to provide it with the given integer.
+    zig_progress_fd: ?i32 = null,
 };
 
 /// 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();
+/// expected by POSIX, from a hash map plus options.
+pub fn createEnvironFromMap(
+    arena: Allocator,
+    map: *const EnvMap,
+    options: CreateEnvironOptions,
+) Allocator.Error![:null]?[*:0]u8 {
+    const ZigProgressAction = enum { nothing, edit, delete, add };
+    const zig_progress_action: ZigProgressAction = a: {
+        const fd = options.zig_progress_fd orelse break :a .nothing;
+        const contains = map.get("ZIG_PROGRESS") != null;
+        if (fd >= 0) {
+            break :a if (contains) .edit else .add;
+        } else {
+            if (contains) break :a .delete;
         }
-        count += options.extra_usizes.len;
-        break :c count;
+        break :a .nothing;
     };
+
+    const envp_count: usize = @intCast(@as(isize, map.count()) + @as(isize, switch (zig_progress_action) {
+        .add => 1,
+        .delete => -1,
+        .nothing, .edit => 0,
+    }));
+
     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));
-        }
+    if (zig_progress_action == .add) {
+        envp_buf[i] = try std.fmt.allocPrintZ(arena, "ZIG_PROGRESS={d}", .{options.zig_progress_fd.?});
+        i += 1;
     }
 
-    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;
+    {
+        var it = map.iterator();
+        while (it.next()) |pair| {
+            if (mem.eql(u8, pair.key_ptr.*, "ZIG_PROGRESS")) switch (zig_progress_action) {
+                .add => unreachable,
+                .delete => continue,
+                .edit => {
+                    envp_buf[i] = try std.fmt.allocPrintZ(arena, "{s}={d}", .{
+                        pair.key_ptr.*, options.zig_progress_fd.?,
+                    });
+                    i += 1;
+                    continue;
+                },
+                .nothing => {},
+            };
 
-    if (options.env_map) |env_map| {
-        var it = env_map.iterator();
-        while (it.next()) |pair| : (i += 1) {
             envp_buf[i] = try std.fmt.allocPrintZ(arena, "{s}={s}", .{ pair.key_ptr.*, pair.value_ptr.* });
+            i += 1;
+        }
+    }
+
+    assert(i == envp_count);
+    return envp_buf;
+}
+
+/// Creates a null-deliminated environment variable block in the format
+/// expected by POSIX, from a hash map plus options.
+pub fn createEnvironFromExisting(
+    arena: Allocator,
+    existing: [*:null]const ?[*:0]const u8,
+    options: CreateEnvironOptions,
+) Allocator.Error![:null]?[*:0]u8 {
+    const existing_count, const contains_zig_progress = c: {
+        var count: usize = 0;
+        var contains = false;
+        while (existing[count]) |line| : (count += 1) {
+            contains = contains or mem.eql(u8, mem.sliceTo(line, '='), "ZIG_PROGRESS");
         }
+        break :c .{ count, contains };
+    };
+    const ZigProgressAction = enum { nothing, edit, delete, add };
+    const zig_progress_action: ZigProgressAction = a: {
+        const fd = options.zig_progress_fd orelse break :a .nothing;
+        if (fd >= 0) {
+            break :a if (contains_zig_progress) .edit else .add;
+        } else {
+            if (contains_zig_progress) break :a .delete;
+        }
+        break :a .nothing;
+    };
+
+    const envp_count: usize = @intCast(@as(isize, @intCast(existing_count)) + @as(isize, switch (zig_progress_action) {
+        .add => 1,
+        .delete => -1,
+        .nothing, .edit => 0,
+    }));
+
+    const envp_buf = try arena.allocSentinel(?[*:0]u8, envp_count, null);
+    var i: usize = 0;
+    var existing_index: usize = 0;
+
+    if (zig_progress_action == .add) {
+        envp_buf[i] = try std.fmt.allocPrintZ(arena, "ZIG_PROGRESS={d}", .{options.zig_progress_fd.?});
+        i += 1;
+    }
+
+    while (existing[existing_index]) |line| : (existing_index += 1) {
+        if (mem.eql(u8, mem.sliceTo(line, '='), "ZIG_PROGRESS")) switch (zig_progress_action) {
+            .add => unreachable,
+            .delete => continue,
+            .edit => {
+                envp_buf[i] = try std.fmt.allocPrintZ(arena, "ZIG_PROGRESS={d}", .{options.zig_progress_fd.?});
+                i += 1;
+                continue;
+            },
+            .nothing => {},
+        };
+        envp_buf[i] = try arena.dupeZ(u8, mem.span(line));
+        i += 1;
     }
 
     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 });
+pub fn createNullDelimitedEnvMap(arena: mem.Allocator, env_map: *const EnvMap) Allocator.Error![:null]?[*:0]u8 {
+    return createEnvironFromMap(arena, env_map, .{});
 }
 
 test createNullDelimitedEnvMap {