thespian/src/remote/endpoint.zig

/// Parent-side endpoint actor.
///
/// Wraps a subprocess, accumulates framed CBOR from its stdout, and dispatches
/// decoded messages to local actors or proxies.
///
/// Outbound messages (local → endpoint):
///   {"send", from_id: u64, to_id: u64, payload}    — via proxy, by remote ID
///   {"send", from_id: u64, to_name: text, payload} — direct by name
///   {"proxy_exit", remote_id: u64, reason: text}   — proxy notifying of exit
const std = @import("std");
const tp = @import("thespian");
const cbor = @import("cbor");
const framing = @import("framing");
const protocol = @import("protocol");
const proxy = @import("proxy");

const subprocess = tp.subprocess;

const proc_tag = "endpoint_proc";

pub const Args = struct {
    allocator: std.mem.Allocator,
    /// CBOR-encoded argv array for the child binary.
    /// Must be heap-allocated; endpoint.init will free it.
    argv: []const u8,
};

const Endpoint = struct {
    allocator: std.mem.Allocator,
    proc: subprocess,
    accumulator: framing.Accumulator,
    /// Inbound proxy table: remote actor ID → local proxy pid.
    proxies: std.AutoHashMap(u64, tp.pid),
    receiver: tp.Receiver(*@This()),

    fn start(args: Args) tp.result {
        return init(args) catch |e| return tp.exit_error(e, @errorReturnTrace());
    }

    fn init(args: Args) !void {
        defer args.allocator.free(args.argv);
        const proc = try subprocess.init(args.allocator, tp.message{ .buf = args.argv }, proc_tag, .Pipe);
        const self = try args.allocator.create(@This());
        self.* = .{
            .allocator = args.allocator,
            .proc = proc,
            .accumulator = .{},
            .proxies = std.AutoHashMap(u64, tp.pid).init(args.allocator),
            .receiver = .init(receive_fn, deinit, self),
        };
        errdefer self.deinit();
        tp.receive(&self.receiver);
    }

    fn deinit(self: *@This()) void {
        // Exit all live proxies (best-effort; endpoint may already be shutting down).
        var it = self.proxies.valueIterator();
        while (it.next()) |p| {
            p.send(.{ "exit", "transport_closed" }) catch {};
            p.deinit();
        }
        self.proxies.deinit();
        self.proc.deinit();
        self.allocator.destroy(self);
    }

    fn receive_fn(self: *@This(), from: tp.pid_ref, m: tp.message) tp.result {
        return self.receive(from, m) catch |e| return tp.exit_error(e, @errorReturnTrace());
    }

    fn receive(self: *@This(), _: tp.pid_ref, m: tp.message) !void {
        var bytes: []const u8 = "";
        var from_id: u64 = 0;
        var to_id: u64 = 0;
        var to_name: []const u8 = "";
        var payload: []const u8 = "";
        var remote_id: u64 = 0;

        if (try m.match(.{ proc_tag, "stdout", tp.extract(&bytes) })) {
            if (self.accumulator.feed(bytes)) |frame| try self.dispatch_inbound(frame);
        } else if (try m.match(.{ proc_tag, "stderr", tp.any })) {
        } else if (try m.match(.{ "send", tp.extract(&from_id), tp.extract(&to_id), cbor.extract_cbor(&payload) })) {
            // Outbound send by remote ID — from a local proxy.
            try self.send_wire_by_id(from_id, to_id, payload);
        } else if (try m.match(.{ "send", tp.extract(&from_id), tp.extract(&to_name), cbor.extract_cbor(&payload) })) {
            // Outbound send_named — from a local actor addressing a well-known remote actor by name.
            try self.send_wire_named(from_id, to_name, payload);
        } else if (try m.match(.{ "proxy_exit", tp.extract(&remote_id), tp.any })) {
            // A local proxy has exited; remove it from the proxy table.
            if (self.proxies.fetchRemove(remote_id)) |entry| entry.value.deinit();
        } else if (try m.match(.{ proc_tag, "term", tp.any, tp.any })) {
            return tp.exit("transport_closed");
        } else {
            return tp.unexpected(m);
        }
    }

    /// Get the local proxy for remote_id, creating one if it doesn't exist yet.
    fn get_or_create_proxy(self: *@This(), remote_id: u64) !tp.pid_ref {
        if (self.proxies.getPtr(remote_id)) |p| return p.ref();
        const p = try tp.spawn(self.allocator, proxy.Args{
            .allocator = self.allocator,
            .endpoint = tp.self_pid().clone(),
            .remote_id = remote_id,
        }, proxy.start, "proxy");
        try self.proxies.put(remote_id, p);
        return self.proxies.getPtr(remote_id).?.ref();
    }

    fn dispatch_inbound(self: *@This(), frame: []const u8) !void {
        const msg = try protocol.decode(frame);
        switch (msg) {
            .send => |s| {
                // Ensure a local proxy exists for the sender.
                _ = try self.get_or_create_proxy(s.from_id);
                // Deliver payload to the local actor identified by to_id.
                // Requires the outbound ID table (not yet implemented); error for now.
                _ = s.to_id;
                _ = s.payload;
                return tp.exit_error(error.OutboundTableNotImplemented, null);
            },
            .send_named => |s| {
                // Ensure a local proxy exists for the sender so it can receive replies.
                _ = try self.get_or_create_proxy(s.from_id);
                const actor = tp.env.get().proc(s.to_name);
                try actor.send_raw(tp.message{ .buf = s.payload });
            },
            .exit => |e| {
                // Remote actor has exited; notify its local proxy.
                if (self.proxies.get(e.id)) |p|
                    p.send(.{ "exit", e.reason }) catch {};
            },
            .proxy_id, .link, .transport_error => return tp.exit_error(error.UnexpectedMessage, null),
        }
    }

    fn send_wire_by_id(self: *@This(), from_id: u64, to_id: u64, payload: []const u8) !void {
        var msg_buf: [framing.max_frame_size]u8 = undefined;
        var msg_stream: std.Io.Writer = .fixed(&msg_buf);
        try protocol.encode_send(&msg_stream, from_id, to_id, payload);

        var frame_buf: [framing.max_frame_size + 4]u8 = undefined;
        var frame_stream: std.Io.Writer = .fixed(&frame_buf);
        try framing.write_frame(&frame_stream, msg_stream.buffered());
        try self.proc.send(frame_stream.buffered());
    }

    fn send_wire_named(self: *@This(), from_id: u64, to_name: []const u8, payload: []const u8) !void {
        var msg_buf: [framing.max_frame_size]u8 = undefined;
        var msg_stream: std.Io.Writer = .fixed(&msg_buf);
        try protocol.encode_send_named(&msg_stream, from_id, to_name, payload);

        var frame_buf: [framing.max_frame_size + 4]u8 = undefined;
        var frame_stream: std.Io.Writer = .fixed(&frame_buf);
        try framing.write_frame(&frame_stream, msg_stream.buffered());
        try self.proc.send(frame_stream.buffered());
    }
};

pub const start = Endpoint.start;