801 lines · 26.1 KB · Zig
RawDownload
const std = @import("std");const builtin = @import("builtin");const io_mod = @import("../shared/io.zig"); const Allocator = std.mem.Allocator; const Identity = union(enum) {    linux_start_ticks: u64,    macos_unique_id: u64,     fn eql(self: Identity, other: Identity) bool {        return switch (self) {            .linux_start_ticks => |ticks| switch (other) {                .linux_start_ticks => |other_ticks| ticks == other_ticks,                else => false,            },            .macos_unique_id => |unique_id| switch (other) {                .macos_unique_id => |other_unique_id| unique_id == other_unique_id,                else => false,            },        };    }}; const TrackedProcess = struct {    pid: std.posix.pid_t,    identity: Identity,}; const ProcessSnapshot = struct {    identity: Identity,    parent_pid: std.posix.pid_t,    parent_unique_id: ?u64 = null,}; pub const DeliverySummary = struct {    delivered: usize = 0,    incomplete: bool = false,}; const ProcessGroupState = union(enum) {    found: std.posix.pid_t,    vanished,    unavailable,}; const SystemSignalEffects = struct {    fn capture(alloc: Allocator, pid: std.posix.pid_t) !ProcessSnapshot {        return captureSnapshot(alloc, pid);    }     fn processGroup(pid: std.posix.pid_t) ProcessGroupState {        return inspectProcessGroup(pid);    }     fn send(pid: std.posix.pid_t, signal: std.posix.SIG) std.posix.KillError!void {        return std.posix.kill(pid, signal);    }}; /// Tracks a command root and descendants across new sessions or process/// groups. Identity checks guard both traversal and signaling against PID/// reuse.pub const Tracker = struct {    alloc: Allocator,    root: ?TrackedProcess = null,    processes: std.ArrayList(TrackedProcess) = .empty,    macos_child_buffer: []std.posix.pid_t = &.{},    macos_pid_buffer: []std.posix.pid_t = &.{},     pub fn init(alloc: Allocator) !Tracker {        var tracker = Tracker{ .alloc = alloc };        errdefer tracker.deinit();        if (comptime builtin.os.tag == .macos) {            const reported = Darwin.proc_listchildpids(0, null, 0);            const capacity: usize = if (reported > 0)                @max(@as(usize, @intCast(reported)) + 256, 1024)            else                1024;            tracker.macos_child_buffer = try alloc.alloc(                std.posix.pid_t,                capacity,            );            const process_count = Darwin.proc_listallpids(null, 0);            const process_capacity: usize = if (process_count > 0)                @max(@as(usize, @intCast(process_count)) + 256, 1024)            else                1024;            tracker.macos_pid_buffer = try alloc.alloc(                std.posix.pid_t,                process_capacity,            );        }        return tracker;    }     pub fn deinit(self: *Tracker) void {        self.processes.deinit(self.alloc);        if (self.macos_child_buffer.len > 0) {            self.alloc.free(self.macos_child_buffer);        }        if (self.macos_pid_buffer.len > 0) {            self.alloc.free(self.macos_pid_buffer);        }        self.* = undefined;    }     pub fn refresh(self: *Tracker, root_pid: std.posix.pid_t) !void {        const root_snapshot: ?ProcessSnapshot = captureSnapshot(self.alloc, root_pid) catch |err| switch (err) {            error.ProcessNotFound => null,            else => return err,        };        var traverse_root = false;        if (root_snapshot) |snapshot| {            if (self.root) |root| {                traverse_root = root.pid == root_pid and                    root.identity.eql(snapshot.identity);            } else {                self.root = .{                    .pid = root_pid,                    .identity = snapshot.identity,                };                traverse_root = true;            }        }        if (traverse_root) try self.appendDirectChildren(self.root.?);        if (comptime builtin.os.tag == .macos) {            if (root_snapshot == null) try self.refreshLineageProcesses();        }         var parent_index: usize = 0;        while (parent_index < self.processes.items.len) : (parent_index += 1) {            const parent = self.processes.items[parent_index];            const actual = captureSnapshot(self.alloc, parent.pid) catch continue;            if (!shouldTraverseParent(parent.identity, actual.identity)) continue;            try self.appendDirectChildren(parent);        }    }     pub fn refreshAdditionalRoot(        self: *Tracker,        root_pid: std.posix.pid_t,    ) !void {        const snapshot = captureSnapshot(self.alloc, root_pid) catch |err| switch (err) {            error.ProcessNotFound => return,            else => return err,        };        try self.appendDirectChildren(.{            .pid = root_pid,            .identity = snapshot.identity,        });    }     pub fn refreshLineageProcesses(self: *Tracker) !void {        if (comptime builtin.os.tag != .macos) return;        const reported = Darwin.proc_listallpids(null, 0);        if (reported <= 0) return;        const required = @as(usize, @intCast(reported)) + 256;        if (required > self.macos_pid_buffer.len) {            self.macos_pid_buffer = try self.alloc.realloc(                self.macos_pid_buffer,                required,            );        }        const count = Darwin.proc_listallpids(            self.macos_pid_buffer.ptr,            @intCast(self.macos_pid_buffer.len * @sizeOf(std.posix.pid_t)),        );        if (count <= 0) return;        const process_count = @min(            @as(usize, @intCast(count)),            self.macos_pid_buffer.len,        );        var changed = true;        while (changed) {            changed = false;            for (self.macos_pid_buffer[0..process_count]) |pid| {                if (pid <= 0 or pid == std.c.getpid()) continue;                if (try self.trackLineageProcess(pid)) changed = true;            }        }    }     pub fn signalAll(self: *Tracker, signal: std.posix.SIG) usize {        return self.signalProcessesChecked(signal, null).delivered;    }     pub fn signalOutsideProcessGroup(        self: *Tracker,        signal: std.posix.SIG,        preserved_group: std.posix.pid_t,    ) usize {        return self.signalProcessesChecked(signal, preserved_group).delivered;    }     pub fn signalOutsideProcessGroupChecked(        self: *Tracker,        signal: std.posix.SIG,        preserved_group: std.posix.pid_t,    ) DeliverySummary {        return self.signalProcessesChecked(signal, preserved_group);    }     fn signalProcessesChecked(        self: *Tracker,        signal: std.posix.SIG,        preserved_group: ?std.posix.pid_t,    ) DeliverySummary {        return self.signalProcessesWith(            signal,            preserved_group,            SystemSignalEffects,        );    }     fn signalProcessesWith(        self: *Tracker,        signal: std.posix.SIG,        preserved_group: ?std.posix.pid_t,        comptime Effects: type,    ) DeliverySummary {        var summary: DeliverySummary = .{};        var index = self.processes.items.len;        while (index > 0) {            index -= 1;            self.signalTrackedProcessWith(                self.processes.items[index],                signal,                preserved_group,                &summary,                Effects,            );        }        if (self.root) |root| {            self.signalTrackedProcessWith(                root,                signal,                preserved_group,                &summary,                Effects,            );        }        return summary;    }     fn signalTrackedProcessWith(        self: *Tracker,        process: TrackedProcess,        signal: std.posix.SIG,        preserved_group: ?std.posix.pid_t,        summary: *DeliverySummary,        comptime Effects: type,    ) void {        const actual = Effects.capture(self.alloc, process.pid) catch |err| {            if (err != error.ProcessNotFound) summary.incomplete = true;            return;        };        if (!process.identity.eql(actual.identity)) return;        const process_group = switch (Effects.processGroup(process.pid)) {            .found => |value| value,            .vanished => return,            .unavailable => {                summary.incomplete = true;                return;            },        };        if (!shouldSignalProcess(process_group, preserved_group)) return;        Effects.send(process.pid, signal) catch |err| switch (err) {            error.ProcessNotFound => return,            else => {                summary.incomplete = true;                return;            },        };        summary.delivered += 1;    }     pub fn anyAlive(self: *Tracker) bool {        if (self.root) |root| {            const actual: ?ProcessSnapshot = captureSnapshot(self.alloc, root.pid) catch null;            if (actual) |snapshot| {                if (root.identity.eql(snapshot.identity)) return true;            }        }        for (self.processes.items) |process| {            const actual = captureSnapshot(self.alloc, process.pid) catch continue;            if (process.identity.eql(actual.identity)) return true;        }        return false;    }     fn appendDirectChildren(        self: *Tracker,        parent: TrackedProcess,    ) !void {        switch (builtin.os.tag) {            .linux => try self.appendLinuxChildren(parent),            .macos => try self.appendMacOSChildren(parent),            else => return error.ProcessTreeUnsupported,        }    }     fn appendLinuxChildren(        self: *Tracker,        parent: TrackedProcess,    ) !void {        if (comptime builtin.os.tag != .linux) return error.ProcessTreeUnsupported;        if (!try self.parentIdentityMatches(parent)) return;        const task_path = try std.fmt.allocPrint(            self.alloc,            "/proc/{d}/task",            .{parent.pid},        );        defer self.alloc.free(task_path);        var task_dir = std.Io.Dir.openDirAbsolute(            io_mod.getIo(),            task_path,            .{ .iterate = true },        ) catch |err| switch (err) {            error.FileNotFound => return,            else => return err,        };        defer task_dir.close(io_mod.getIo());        var tasks = task_dir.iterate();        while (try tasks.next(io_mod.getIo())) |entry| {            const tid = std.fmt.parseInt(                std.posix.pid_t,                entry.name,                10,            ) catch continue;            if (tid <= 0) continue;            try self.appendLinuxTaskChildren(parent, tid);        }    }     fn appendLinuxTaskChildren(        self: *Tracker,        parent: TrackedProcess,        tid: std.posix.pid_t,    ) !void {        const path = try std.fmt.allocPrint(            self.alloc,            "/proc/{d}/task/{d}/children",            .{ parent.pid, tid },        );        defer self.alloc.free(path);        var file = std.Io.Dir.openFileAbsolute(io_mod.getIo(), path, .{}) catch |err| switch (err) {            error.FileNotFound => return,            else => return err,        };        defer file.close(io_mod.getIo());        var buffer: [64 * 1024]u8 = undefined;        const read_len = readLinuxChildrenFile(file, &buffer) catch |err| switch (err) {            error.ProcessNotFound => return,            else => return err,        };        if (!try self.parentIdentityMatches(parent)) return;        var children = std.mem.tokenizeAny(u8, buffer[0..read_len], " \t\r\n");        while (children.next()) |pid_text| {            const pid = std.fmt.parseInt(std.posix.pid_t, pid_text, 10) catch continue;            try self.trackChild(pid, parent.pid);        }    }     fn appendMacOSChildren(        self: *Tracker,        parent: TrackedProcess,    ) !void {        if (comptime builtin.os.tag != .macos) return error.ProcessTreeUnsupported;        if (!try self.parentIdentityMatches(parent)) return;        const count = Darwin.proc_listchildpids(            parent.pid,            self.macos_child_buffer.ptr,            @intCast(self.macos_child_buffer.len * @sizeOf(std.posix.pid_t)),        );        if (count <= 0) return;        const child_count = @min(            @as(usize, @intCast(count)),            self.macos_child_buffer.len,        );        if (!try self.parentIdentityMatches(parent)) return;        for (self.macos_child_buffer[0..child_count]) |pid| {            if (pid > 0) try self.trackChild(pid, parent.pid);        }    }     fn parentIdentityMatches(self: *Tracker, parent: TrackedProcess) !bool {        const snapshot = captureSnapshot(self.alloc, parent.pid) catch |err| switch (err) {            error.ProcessNotFound => return false,            else => return err,        };        return parent.identity.eql(snapshot.identity);    }     fn trackChild(        self: *Tracker,        pid: std.posix.pid_t,        expected_parent_pid: std.posix.pid_t,    ) !void {        const snapshot = captureSnapshot(self.alloc, pid) catch |err| switch (err) {            error.ProcessNotFound => return,            else => return err,        };        if (!snapshotBelongsToParent(snapshot, expected_parent_pid)) return;        if (self.root) |root| {            if (root.pid == pid and root.identity.eql(snapshot.identity)) return;        }        for (self.processes.items) |*process| {            if (process.pid != pid) continue;            process.identity = snapshot.identity;            return;        }        try self.processes.append(self.alloc, .{            .pid = pid,            .identity = snapshot.identity,        });    }     fn trackLineageProcess(self: *Tracker, pid: std.posix.pid_t) !bool {        const snapshot = captureSnapshot(self.alloc, pid) catch return false;        const parent_unique_id = snapshot.parent_unique_id orelse return false;        if (!self.containsMacOSUniqueId(parent_unique_id)) return false;        if (self.root) |root| {            if (root.pid == pid and root.identity.eql(snapshot.identity)) return false;        }        for (self.processes.items) |*process| {            if (process.pid != pid) continue;            if (process.identity.eql(snapshot.identity)) return false;            process.identity = snapshot.identity;            return true;        }        try self.processes.append(self.alloc, .{            .pid = pid,            .identity = snapshot.identity,        });        return true;    }     fn containsMacOSUniqueId(self: *Tracker, unique_id: u64) bool {        if (self.root) |root| {            if (identityHasMacOSUniqueId(root.identity, unique_id)) return true;        }        for (self.processes.items) |process| {            if (identityHasMacOSUniqueId(process.identity, unique_id)) return true;        }        return false;    }}; fn identityHasMacOSUniqueId(identity: Identity, unique_id: u64) bool {    return switch (identity) {        .macos_unique_id => |actual| actual == unique_id,        else => false,    };} fn shouldTraverseParent(expected: Identity, actual: Identity) bool {    return expected.eql(actual);} fn snapshotBelongsToParent(    snapshot: ProcessSnapshot,    expected_parent_pid: std.posix.pid_t,) bool {    return snapshot.parent_pid == expected_parent_pid;} fn shouldSignalProcess(    process_group: ?std.posix.pid_t,    preserved_group: ?std.posix.pid_t,) bool {    const preserved = preserved_group orelse return true;    const actual = process_group orelse return false;    return actual != preserved;} fn inspectProcessGroup(pid: std.posix.pid_t) ProcessGroupState {    if (comptime builtin.os.tag == .windows or builtin.os.tag == .wasi) {        return .unavailable;    }    const process_group = getpgid(pid);    if (process_group >= 0) return .{ .found = process_group };    return switch (std.c.errno(process_group)) {        .SRCH => .vanished,        else => .unavailable,    };} extern "c" fn getpgid(pid: std.posix.pid_t) std.posix.pid_t; fn readLinuxChildrenFile(file: std.Io.File, buffer: []u8) !usize {    if (comptime builtin.os.tag != .linux) return error.ProcessTreeUnsupported;    while (true) {        const result = std.posix.system.read(file.handle, buffer.ptr, buffer.len);        switch (std.posix.errno(result)) {            .SUCCESS => return @intCast(result),            .INTR => continue,            .SRCH, .NOENT => return error.ProcessNotFound,            else => return error.ProcessTreeInspectionFailed,        }

Lines 1–500 of 801.