Commit 41a5ad28c9

Andrew Kelley <andrew@ziglang.org>
2023-03-03 23:27:53
std: child process API supports rusage data
1 parent 8b054e1
Changed files (4)
lib/std/os/linux.zig
@@ -944,6 +944,16 @@ pub fn waitpid(pid: pid_t, status: *u32, flags: u32) usize {
     return syscall4(.wait4, @bitCast(usize, @as(isize, pid)), @ptrToInt(status), flags, 0);
 }
 
+pub fn wait4(pid: pid_t, status: *u32, flags: u32, usage: ?*rusage) usize {
+    return syscall4(
+        .wait4,
+        @bitCast(usize, @as(isize, pid)),
+        @ptrToInt(status),
+        flags,
+        @ptrToInt(usage),
+    );
+}
+
 pub fn waitid(id_type: P, id: i32, infop: *siginfo_t, flags: u32) usize {
     return syscall5(.waitid, @enumToInt(id_type), @bitCast(usize, @as(isize, id)), @ptrToInt(infop), flags, 0);
 }
lib/std/c.zig
@@ -153,7 +153,8 @@ pub extern "c" fn linkat(oldfd: c.fd_t, oldpath: [*:0]const u8, newfd: c.fd_t, n
 pub extern "c" fn unlink(path: [*:0]const u8) c_int;
 pub extern "c" fn unlinkat(dirfd: c.fd_t, path: [*:0]const u8, flags: c_uint) c_int;
 pub extern "c" fn getcwd(buf: [*]u8, size: usize) ?[*]u8;
-pub extern "c" fn waitpid(pid: c.pid_t, stat_loc: ?*c_int, options: c_int) c.pid_t;
+pub extern "c" fn waitpid(pid: c.pid_t, status: ?*c_int, options: c_int) c.pid_t;
+pub extern "c" fn wait4(pid: c.pid_t, status: ?*c_int, options: c_int, ru: ?*c.rusage) c.pid_t;
 pub extern "c" fn fork() c_int;
 pub extern "c" fn access(path: [*:0]const u8, mode: c_uint) c_int;
 pub extern "c" fn faccessat(dirfd: c.fd_t, path: [*:0]const u8, mode: c_uint, flags: c_uint) c_int;
lib/std/child_process.zig
@@ -17,6 +17,7 @@ const Os = std.builtin.Os;
 const TailQueue = std.TailQueue;
 const maxInt = std.math.maxInt;
 const assert = std.debug.assert;
+const is_darwin = builtin.target.isDarwin();
 
 pub const ChildProcess = struct {
     pub const Id = switch (builtin.os.tag) {
@@ -70,6 +71,43 @@ pub const ChildProcess = struct {
     /// Darwin-only. Start child process in suspended state as if SIGSTOP was sent.
     start_suspended: bool = false,
 
+    /// Set to true to obtain rusage information for the child process.
+    /// Depending on the target platform and implementation status, the
+    /// requested statistics may or may not be available. If they are
+    /// available, then the `resource_usage_statistics` field will be populated
+    /// after calling `wait`.
+    /// On Linux, this obtains rusage statistics from wait4().
+    request_resource_usage_statistics: bool = false,
+
+    /// This is available after calling wait if
+    /// `request_resource_usage_statistics` was set to `true` before calling
+    /// `spawn`.
+    resource_usage_statistics: ResourceUsageStatistics = .{},
+
+    pub const ResourceUsageStatistics = struct {
+        rusage: @TypeOf(rusage_init) = rusage_init,
+
+        /// Returns the peak resident set size of the child process, in bytes,
+        /// if available.
+        pub inline fn getMaxRss(rus: ResourceUsageStatistics) ?usize {
+            switch (builtin.os.tag) {
+                .linux => {
+                    if (rus.rusage) |ru| {
+                        return @intCast(usize, ru.maxrss) * 1024;
+                    } else {
+                        return null;
+                    }
+                },
+                else => return null,
+            }
+        }
+
+        const rusage_init = switch (builtin.os.tag) {
+            .linux => @as(?std.os.rusage, null),
+            else => {},
+        };
+    };
+
     pub const Arg0Expand = os.Arg0Expand;
 
     pub const SpawnError = error{
@@ -332,7 +370,16 @@ pub const ChildProcess = struct {
     }
 
     fn waitUnwrapped(self: *ChildProcess) !void {
-        const res: os.WaitPidResult = os.waitpid(self.id, 0);
+        const res: os.WaitPidResult = res: {
+            if (builtin.os.tag == .linux and self.request_resource_usage_statistics) {
+                var ru: std.os.rusage = undefined;
+                const res = os.wait4(self.id, 0, &ru);
+                self.resource_usage_statistics.rusage = ru;
+                break :res res;
+            }
+
+            break :res os.waitpid(self.id, 0);
+        };
         const status = res.status;
         self.cleanupStreams();
         self.handleWaitResult(status);
lib/std/os.zig
@@ -4000,8 +4000,28 @@ pub const WaitPidResult = struct {
 pub fn waitpid(pid: pid_t, flags: u32) WaitPidResult {
     const Status = if (builtin.link_libc) c_int else u32;
     var status: Status = undefined;
+    const coerced_flags = if (builtin.link_libc) @intCast(c_int, flags) else flags;
     while (true) {
-        const rc = system.waitpid(pid, &status, if (builtin.link_libc) @intCast(c_int, flags) else flags);
+        const rc = system.waitpid(pid, &status, coerced_flags);
+        switch (errno(rc)) {
+            .SUCCESS => return .{
+                .pid = @intCast(pid_t, rc),
+                .status = @bitCast(u32, status),
+            },
+            .INTR => continue,
+            .CHILD => unreachable, // The process specified does not exist. It would be a race condition to handle this error.
+            .INVAL => unreachable, // Invalid flags.
+            else => unreachable,
+        }
+    }
+}
+
+pub fn wait4(pid: pid_t, flags: u32, ru: ?*rusage) WaitPidResult {
+    const Status = if (builtin.link_libc) c_int else u32;
+    var status: Status = undefined;
+    const coerced_flags = if (builtin.link_libc) @intCast(c_int, flags) else flags;
+    while (true) {
+        const rc = system.wait4(pid, &status, coerced_flags, ru);
         switch (errno(rc)) {
             .SUCCESS => return .{
                 .pid = @intCast(pid_t, rc),