std.Progress: child process sends updates via IPC

This commit is contained in:
Andrew Kelley 2024-05-23 14:10:03 -07:00
parent ed36470af1
commit f07116404a
3 changed files with 244 additions and 76 deletions

View File

@ -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;

View File

@ -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();
const envp_buf = try arena.allocSentinel(?[*:0]u8, envp_count, null);
{
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;
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));
}
assert(i == envp_count);
}
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();
while (it.next()) |pair| : (i += 1) {
envp_buf[i] = try std.fmt.allocPrintZ(arena, "{s}={s}", .{ pair.key_ptr.*, pair.value_ptr.* });
}
}
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);

View File

@ -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]);
}