262 lines
10 KiB
Zig
262 lines
10 KiB
Zig
const std = @import("std");
|
|
const cbor = @import("cbor");
|
|
const tp = @import("thespian.zig");
|
|
|
|
pid: ?tp.pid,
|
|
stdin_behavior: std.process.Child.StdIo,
|
|
|
|
const Self = @This();
|
|
pub const max_chunk_size = 4096 - 32;
|
|
pub const Writer = std.io.Writer(*Self, error{Exit}, write);
|
|
pub const BufferedWriter = std.io.BufferedWriter(max_chunk_size, Writer);
|
|
|
|
pub fn init(a: std.mem.Allocator, argv: tp.message, tag: [:0]const u8, stdin_behavior: std.process.Child.StdIo) !Self {
|
|
return .{
|
|
.pid = try Proc.create(a, argv, tag, stdin_behavior),
|
|
.stdin_behavior = stdin_behavior,
|
|
};
|
|
}
|
|
|
|
pub fn deinit(self: *Self) void {
|
|
if (self.pid) |pid| {
|
|
pid.deinit();
|
|
self.pid = null;
|
|
}
|
|
}
|
|
|
|
pub fn write(self: *Self, bytes: []const u8) error{Exit}!usize {
|
|
try self.send(bytes);
|
|
return bytes.len;
|
|
}
|
|
|
|
pub fn send(self: *const Self, bytes_: []const u8) tp.result {
|
|
if (self.stdin_behavior != .Pipe) return tp.exit("cannot send to closed stdin");
|
|
const pid = if (self.pid) |pid| pid else return tp.exit_error(error.Closed, null);
|
|
var bytes = bytes_;
|
|
while (bytes.len > 0)
|
|
bytes = loop: {
|
|
if (bytes.len > max_chunk_size) {
|
|
try pid.send(.{ "stdin", bytes[0..max_chunk_size] });
|
|
break :loop bytes[max_chunk_size..];
|
|
} else {
|
|
try pid.send(.{ "stdin", bytes });
|
|
break :loop &[_]u8{};
|
|
}
|
|
};
|
|
}
|
|
|
|
pub fn close(self: *Self) tp.result {
|
|
defer self.deinit();
|
|
if (self.stdin_behavior == .Pipe)
|
|
if (self.pid) |pid| if (!pid.expired()) try pid.send(.{"stdin_close"});
|
|
}
|
|
|
|
pub fn term(self: *Self) tp.result {
|
|
defer self.deinit();
|
|
if (self.pid) |pid| if (!pid.expired()) try pid.send(.{"term"});
|
|
}
|
|
|
|
pub fn writer(self: *Self) Writer {
|
|
return .{ .context = self };
|
|
}
|
|
|
|
pub fn bufferedWriter(self: *Self) BufferedWriter {
|
|
return .{ .unbuffered_writer = self.writer() };
|
|
}
|
|
|
|
const Proc = struct {
|
|
a: std.mem.Allocator,
|
|
receiver: Receiver,
|
|
args: std.heap.ArenaAllocator,
|
|
parent: tp.pid,
|
|
child: std.process.Child,
|
|
tag: [:0]const u8,
|
|
stdin_buffer: std.ArrayList(u8),
|
|
fd_stdin: ?tp.file_descriptor = null,
|
|
fd_stdout: ?tp.file_descriptor = null,
|
|
fd_stderr: ?tp.file_descriptor = null,
|
|
write_pending: bool = false,
|
|
stdin_close_pending: bool = false,
|
|
|
|
const Receiver = tp.Receiver(*Proc);
|
|
|
|
fn create(a: std.mem.Allocator, argv: tp.message, tag: [:0]const u8, stdin_behavior: std.process.Child.StdIo) !tp.pid {
|
|
const self: *Proc = try a.create(Proc);
|
|
|
|
var args = std.heap.ArenaAllocator.init(a);
|
|
const args_a = args.allocator();
|
|
var iter = argv.buf;
|
|
var len = cbor.decodeArrayHeader(&iter) catch return error.InvalidArgument;
|
|
var argv_ = try args_a.alloc([]const u8, len);
|
|
var arg: []const u8 = undefined;
|
|
var i: usize = 0;
|
|
while (len > 0) : (len -= 1) {
|
|
if (!(cbor.matchString(&iter, &arg) catch return error.InvalidArgument))
|
|
return error.InvalidArgument;
|
|
argv_[i] = try args_a.dupe(u8, arg);
|
|
i += 1;
|
|
}
|
|
|
|
var child = std.process.Child.init(argv_, a);
|
|
child.stdin_behavior = stdin_behavior;
|
|
child.stdout_behavior = .Pipe;
|
|
child.stderr_behavior = .Pipe;
|
|
|
|
self.* = .{
|
|
.a = a,
|
|
.receiver = Receiver.init(receive, self),
|
|
.args = args,
|
|
.parent = tp.self_pid().clone(),
|
|
.child = child,
|
|
.tag = try a.dupeZ(u8, tag),
|
|
.stdin_buffer = std.ArrayList(u8).init(a),
|
|
};
|
|
return tp.spawn_link(a, self, Proc.start, tag);
|
|
}
|
|
|
|
fn deinit(self: *Proc) void {
|
|
self.args.deinit();
|
|
if (self.fd_stdin) |fd| fd.deinit();
|
|
if (self.fd_stdout) |fd| fd.deinit();
|
|
if (self.fd_stderr) |fd| fd.deinit();
|
|
self.stdin_buffer.deinit();
|
|
self.parent.deinit();
|
|
self.a.free(self.tag);
|
|
}
|
|
|
|
fn start(self: *Proc) tp.result {
|
|
errdefer self.deinit();
|
|
|
|
self.child.spawn() catch |e| {
|
|
try self.parent.send(.{ self.tag, "term", e, 1 });
|
|
return tp.exit_normal();
|
|
};
|
|
_ = self.args.reset(.free_all);
|
|
|
|
if (self.child.stdin_behavior == .Pipe)
|
|
self.fd_stdin = tp.file_descriptor.init("stdin", self.child.stdin.?.handle) catch |e| return tp.exit_error(e, @errorReturnTrace());
|
|
self.fd_stdout = tp.file_descriptor.init("stdout", self.child.stdout.?.handle) catch |e| return tp.exit_error(e, @errorReturnTrace());
|
|
self.fd_stderr = tp.file_descriptor.init("stderr", self.child.stderr.?.handle) catch |e| return tp.exit_error(e, @errorReturnTrace());
|
|
if (self.fd_stdout) |fd_stdout| fd_stdout.wait_read() catch |e| return tp.exit_error(e, @errorReturnTrace());
|
|
if (self.fd_stderr) |fd_stderr| fd_stderr.wait_read() catch |e| return tp.exit_error(e, @errorReturnTrace());
|
|
|
|
tp.receive(&self.receiver);
|
|
}
|
|
|
|
fn receive(self: *Proc, _: tp.pid_ref, m: tp.message) tp.result {
|
|
errdefer self.deinit();
|
|
var bytes: []u8 = "";
|
|
var err: i64 = 0;
|
|
var err_msg: []u8 = "";
|
|
if (try m.match(.{ "fd", "stdout", "read_ready" })) {
|
|
try self.dispatch_stdout();
|
|
if (self.fd_stdout) |fd_stdout| fd_stdout.wait_read() catch |e| return tp.exit_error(e, @errorReturnTrace());
|
|
} else if (try m.match(.{ "fd", "stderr", "read_ready" })) {
|
|
try self.dispatch_stderr();
|
|
if (self.fd_stderr) |fd_stderr| fd_stderr.wait_read() catch |e| return tp.exit_error(e, @errorReturnTrace());
|
|
} else if (try m.match(.{ "fd", "stdin", "write_ready" })) {
|
|
if (self.stdin_buffer.items.len > 0) {
|
|
if (self.child.stdin) |stdin| {
|
|
const written = stdin.write(self.stdin_buffer.items) catch |e| switch (e) {
|
|
error.WouldBlock => {
|
|
if (self.fd_stdin) |fd_stdin| {
|
|
fd_stdin.wait_write() catch |e_| return tp.exit_error(e_, @errorReturnTrace());
|
|
self.write_pending = true;
|
|
return;
|
|
} else return tp.exit_error(error.WouldBlock, @errorReturnTrace());
|
|
},
|
|
else => return tp.exit_error(e, @errorReturnTrace()),
|
|
};
|
|
self.write_pending = false;
|
|
defer {
|
|
if (self.stdin_close_pending and !self.write_pending)
|
|
self.stdin_close();
|
|
}
|
|
if (written == self.stdin_buffer.items.len) {
|
|
self.stdin_buffer.clearRetainingCapacity();
|
|
} else {
|
|
std.mem.copyForwards(u8, self.stdin_buffer.items, self.stdin_buffer.items[written..]);
|
|
self.stdin_buffer.items.len = self.stdin_buffer.items.len - written;
|
|
if (self.fd_stdin) |fd_stdin| {
|
|
fd_stdin.wait_write() catch |e| return tp.exit_error(e, @errorReturnTrace());
|
|
self.write_pending = true;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
} else if (try m.match(.{ "stdin", tp.extract(&bytes) })) {
|
|
if (self.fd_stdin) |fd_stdin| {
|
|
self.stdin_buffer.appendSlice(bytes) catch |e| return tp.exit_error(e, @errorReturnTrace());
|
|
fd_stdin.wait_write() catch |e| return tp.exit_error(e, @errorReturnTrace());
|
|
self.write_pending = true;
|
|
}
|
|
} else if (try m.match(.{"stdin_close"})) {
|
|
if (self.write_pending) {
|
|
self.stdin_close_pending = true;
|
|
} else {
|
|
self.stdin_close();
|
|
}
|
|
} else if (try m.match(.{"stdout_close"})) {
|
|
if (self.child.stdout) |*fd| {
|
|
fd.close();
|
|
self.child.stdout = null;
|
|
}
|
|
} else if (try m.match(.{"stderr_close"})) {
|
|
if (self.child.stderr) |*fd| {
|
|
fd.close();
|
|
self.child.stderr = null;
|
|
}
|
|
} else if (try m.match(.{"term"})) {
|
|
const term_ = self.child.kill() catch |e| return tp.exit_error(e, @errorReturnTrace());
|
|
return self.handle_term(term_);
|
|
} else if (try m.match(.{ "fd", tp.any, "read_error", tp.extract(&err), tp.extract(&err_msg) })) {
|
|
return tp.exit(err_msg);
|
|
}
|
|
}
|
|
|
|
fn stdin_close(self: *Proc) void {
|
|
if (self.child.stdin) |*fd| {
|
|
fd.close();
|
|
self.child.stdin = null;
|
|
tp.env.get().trace(tp.message.fmt(.{ self.tag, "stdin", "closed" }).to(tp.message.c_buffer_type));
|
|
}
|
|
}
|
|
|
|
fn dispatch_stdout(self: *Proc) tp.result {
|
|
var buffer: [max_chunk_size]u8 = undefined;
|
|
const stdout = self.child.stdout orelse return tp.exit("cannot read closed stdout");
|
|
const bytes = stdout.read(&buffer) catch |e| switch (e) {
|
|
error.WouldBlock => return,
|
|
else => return tp.exit_error(e, @errorReturnTrace()),
|
|
};
|
|
if (bytes == 0)
|
|
return self.handle_terminate();
|
|
try self.parent.send(.{ self.tag, "stdout", buffer[0..bytes] });
|
|
}
|
|
|
|
fn dispatch_stderr(self: *Proc) tp.result {
|
|
var buffer: [max_chunk_size]u8 = undefined;
|
|
const stderr = self.child.stderr orelse return tp.exit("cannot read closed stderr");
|
|
const bytes = stderr.read(&buffer) catch |e| switch (e) {
|
|
error.WouldBlock => return,
|
|
else => return tp.exit_error(e, @errorReturnTrace()),
|
|
};
|
|
if (bytes == 0)
|
|
return;
|
|
try self.parent.send(.{ self.tag, "stderr", buffer[0..bytes] });
|
|
}
|
|
|
|
fn handle_terminate(self: *Proc) tp.result {
|
|
return self.handle_term(self.child.wait() catch |e| return tp.exit_error(e, @errorReturnTrace()));
|
|
}
|
|
|
|
fn handle_term(self: *Proc, term_: std.process.Child.Term) tp.result {
|
|
(switch (term_) {
|
|
.Exited => |val| self.parent.send(.{ self.tag, "term", "exited", val }),
|
|
.Signal => |val| self.parent.send(.{ self.tag, "term", "signal", val }),
|
|
.Stopped => |val| self.parent.send(.{ self.tag, "term", "stop", val }),
|
|
.Unknown => |val| self.parent.send(.{ self.tag, "term", "unknown", val }),
|
|
}) catch {};
|
|
return tp.exit_normal();
|
|
}
|
|
};
|