/old_core/src/roon/connection.zig at 047caaf8b72945f5e9ca2f7756dd18cc31bc83dd

// Copyright 2025 Shota FUJI
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//
// SPDX-License-Identifier: Apache-2.0

const std = @import("std");

const moo = @import("moo");
const websocket = @import("websocket");

const Response = struct {
    wrote: *std.Thread.ResetEvent,

    /// `null` on OOM.
    data: ?[]const u8,

    /// If this field is `true`, `data` property does not have meaningful data.
    canceled: bool = false,
};

const ResponsesStore = std.AutoHashMap(i64, *Response);

const RequestHandler = *const fn (
    response_writer: anytype,
    meta: moo.Metadata,
    header_ctx: moo.HeaderParsingContext,
    message: []const u8,
) anyerror!bool;

pub const Connection = struct {
    allocator: std.mem.Allocator,
    thread_safe_allocator: *std.heap.ThreadSafeAllocator,
    // Making this non-pointer causes invalid read/write when thread terminates.
    ws: *websocket.Client,
    addr: []const u8,
    rng: std.Random.Xoshiro256,
    pool: *std.Thread.Pool,
    responses: ResponsesStore,
    responses_mutex: *std.Thread.Mutex,
    subscription_id_cnt: u64 = 0,
    state: State = .connecting,
    state_mutex: std.Thread.Mutex = .{},

    const State = enum {
        connecting,
        connected,
        closed,
    };

    pub const InitError = error{
        WebSocketClientCreationError,
        WebSocketHandshakeError,
        PRNGSeedGenerationFailure,
    } || std.mem.Allocator.Error;

    pub fn init(child_allocator: std.mem.Allocator, address: std.net.Address, pool: *std.Thread.Pool) InitError!@This() {
        var tsa = try child_allocator.create(std.heap.ThreadSafeAllocator);
        tsa.* = std.heap.ThreadSafeAllocator{ .child_allocator = child_allocator };
        const allocator = tsa.allocator();

        var addr = std.ArrayList(u8).init(allocator);
        defer addr.deinit();

        try addr.writer().print("{}", .{address});

        var addr_string = try addr.toOwnedSlice();
        errdefer allocator.free(addr_string);

        // Zig std always prints "<addr>:<port>" for IPv4 and IPv6
        const port_start = std.mem.lastIndexOfScalar(u8, addr_string, ':') orelse {
            unreachable;
        };

        var client = try allocator.create(websocket.Client);
        client.* = websocket.Client.init(allocator, .{
            .port = address.getPort(),
            .host = addr_string[0..port_start],
        }) catch return InitError.WebSocketClientCreationError;
        errdefer client.deinit();

        std.log.debug("Performing WebSocket handshake...", .{});
        client.handshake("/api", .{
            .timeout_ms = 1_000,
        }) catch return InitError.WebSocketHandshakeError;

        const responses = ResponsesStore.init(allocator);

        const responses_mutex = try allocator.create(std.Thread.Mutex);
        responses_mutex.* = std.Thread.Mutex{};

        return .{
            .allocator = allocator,
            .ws = client,
            .addr = addr_string,
            .rng = std.Random.DefaultPrng.init(seed: {
                var seed: u64 = undefined;
                std.posix.getrandom(std.mem.asBytes(&seed)) catch {
                    return InitError.PRNGSeedGenerationFailure;
                };
                break :seed seed;
            }),
            .pool = pool,
            .responses = responses,
            .responses_mutex = responses_mutex,
            .thread_safe_allocator = tsa,
        };
    }

    pub fn listen(self: *Connection, on_request: RequestHandler) std.Thread.SpawnError!void {
        std.log.debug("Spawning WebSocket request handler thread...", .{});
        self.pool.spawn(readLoop, .{ self, on_request }) catch |err| {
            // From user's perspective, normal spawn error and pool's error is no
            // different.
            std.log.warn(
                "Spawning a thread from pool failed, simplifying error ({s})",
                .{@errorName(err)},
            );
            return std.Thread.SpawnError.Unexpected;
        };
    }

    pub fn deinit(self: *Connection) void {
        const state = state: {
            self.state_mutex.lock();
            defer self.state_mutex.unlock();
            break :state self.state;
        };

        if (state != .closed) {
            self.close();
        }

        std.log.debug("Releasing WebSocket resources...", .{});
        self.ws.deinit();
        self.allocator.destroy(self.ws);

        std.log.debug("Releasing stale responses...", .{});
        self.responses_mutex.lock();
        var iter = self.responses.iterator();
        while (iter.next()) |entry| {
            std.log.debug("Releasing response handler for request ID={}...", .{entry.key_ptr.*});

            if (entry.value_ptr.*.data) |data| {
                self.allocator.free(data);
            }
        }
        self.responses.deinit();
        self.responses_mutex.unlock();

        std.log.debug("Releasing connection objects...", .{});
        self.allocator.destroy(self.responses_mutex);

        self.allocator.free(self.addr);
        self.thread_safe_allocator.child_allocator.destroy(self.thread_safe_allocator);
    }

    /// Closes connection and cancels active message listeners.
    fn close(self: *Connection) void {
        // Intentionally blocks `deinit` during closing procedure, so it'll run after closed.
        self.state_mutex.lock();
        defer self.state_mutex.unlock();

        std.log.debug("Closing WebSocket connection...", .{});

        // websocket.zig cannot terminate ongoing read via ".close()".
        // Manually shutting-down socket seems necessary.
        // Related: https://github.com/karlseguin/websocket.zig/issues/46
        std.posix.shutdown(self.ws.stream.stream.handle, .recv) catch |err| {
            std.log.warn("Failed to shutdown socket handle of WebSocket: {s}", .{@errorName(err)});
        };
        self.ws.close(.{ .code = 4000, .reason = "bye" }) catch |err| {
            std.log.warn("Unable to close WebSocket client, proceeding: {s}", .{@errorName(err)});
        };

        {
            self.responses_mutex.lock();
            defer self.responses_mutex.unlock();
            var iter = self.responses.iterator();
            while (iter.next()) |entry| {
                entry.value_ptr.*.canceled = true;
                entry.value_ptr.*.wrote.set();
            }
        }

        self.state = .closed;
    }

    pub fn newRequestId(self: *Connection) i64 {
        return self.rng.random().int(i64);
    }

    pub const RequestError = error{
        Canceled,
    };

    pub fn request(self: *Connection, request_id: i64, message: []u8) ![]const u8 {
        var wrote = std.Thread.ResetEvent{};
        var response = Response{
            .wrote = &wrote,
            .data = null,
        };

        {
            self.responses_mutex.lock();
            defer self.responses_mutex.unlock();
            try self.responses.put(request_id, &response);
        }

        defer {
            self.responses_mutex.lock();
            _ = self.responses.remove(request_id);
            self.responses_mutex.unlock();
        }

        try self.ws.writeBin(message);

        wrote.wait();

        if (response.canceled) {
            return RequestError.Canceled;
        }

        return response.data orelse return std.mem.Allocator.Error.OutOfMemory;
    }

    fn newSubscriptionId(self: *@This()) u64 {
        const id = self.subscription_id_cnt;
        self.subscription_id_cnt += 1;
        return id;
    }

    /// Returns an interface object for subscription management.
    ///
    /// `service_id` is an ID of the service.
    /// `subject` is an entity name to subscribe.
    ///
    /// The resulted service name would be `<service_id>/subscribe_<subject>` and
    /// `<service_id>/unsubscribe_<subject>`.
    ///
    /// Both `service_id` and `subject` must be pointer stable during lifetime of the
    /// returned struct.
    ///
    /// Caller have to call `.activate()` method on the returned struct to receive updates.
    /// To capture received updates, use `.next()` method on the returned struct.
    pub fn subscribe(self: *@This(), service_id: []const u8, subject: []const u8) Subscription {
        return Subscription{
            .request_id = self.newRequestId(),
            .subscription_id = self.newSubscriptionId(),
            .conn = self,
            .response = Response{
                .wrote = undefined,
                .data = null,
            },
            .response_wrote = std.Thread.ResetEvent{},
            .service_id = service_id,
            .subject = subject,
        };
    }

    pub const Subscription = struct {
        request_id: i64,
        subscription_id: u64,
        conn: *Connection,
        response: Response,
        response_wrote: std.Thread.ResetEvent,
        service_id: []const u8,
        subject: []const u8,
        request_message: []u8 = undefined,

        /// Sets up a response listener.
        pub fn activate(self: *@This(), allocator: std.mem.Allocator) !void {
            self.response.wrote = &self.response_wrote;

            {
                self.conn.responses_mutex.lock();
                defer self.conn.responses_mutex.unlock();

                try self.conn.responses.put(self.request_id, &self.response);
            }

            const service_tmpl = "{s}/subscribe_{s}";
            const service_len = std.fmt.count(service_tmpl, .{ self.service_id, self.subject });
            const service = try allocator.alloc(u8, service_len);
            defer allocator.free(service);
            var service_fbs = std.io.fixedBufferStream(service);
            try std.fmt.format(service_fbs.writer(), service_tmpl, .{ self.service_id, self.subject });

            const meta = moo.Metadata{
                .service = service_fbs.getWritten(),
                .verb = "REQUEST",
            };

            const body_tmpl = "{{\"subscription_key\":{}}}";
            const body_len = comptime std.fmt.count(body_tmpl, .{std.math.maxInt(u64)});
            var body_bytes: [body_len]u8 = undefined;
            var body_bytes_fbs = std.io.fixedBufferStream(&body_bytes);
            try std.fmt.format(body_bytes_fbs.writer(), body_tmpl, .{self.subscription_id});

            const body = moo.RawBody{
                .bytes = body_bytes_fbs.getWritten(),
            };

            var header = body.getHeader(self.request_id);
            header.content_type = "application/json";

            const req_buffer = try allocator.alloc(
                u8,
                meta.getEncodeSize() + header.getEncodeSize() + body.getEncodeSize(),
            );
            defer allocator.free(req_buffer);
            var req_fbs = std.io.fixedBufferStream(req_buffer);
            try moo.encode(req_fbs.writer(), meta, header, body);

            try self.conn.ws.writeBin(req_fbs.getWritten());
        }

        /// Release resources and unregister response handler.
        /// This does not sends unsubscribe command to the server.
        /// Call to this function after a call to `Connection.deinit()` is useless and invalid:
        /// `Connection.deinit()` releases this struct's resource as well.
        pub fn deinit(self: *@This()) void {
            self.conn.responses_mutex.lock();
            defer self.conn.responses_mutex.unlock();
            _ = self.conn.responses.remove(self.request_id);
        }

        /// Sends unsubscribe request to the server.
        pub fn unsubscribe(self: *@This(), allocator: std.mem.Allocator) !void {
            const request_id = self.conn.newRequestId();

            const service_tmpl = "{s}/unsubscribe_{s}";
            const service_len = std.fmt.count(service_tmpl, .{ self.service_id, self.subject });
            const service = try allocator.alloc(u8, service_len);
            defer allocator.free(service);
            var service_fbs = std.io.fixedBufferStream(service);
            try std.fmt.format(service_fbs.writer(), service_tmpl, .{ self.service_id, self.subject });

            const meta = moo.Metadata{
                .service = service_fbs.getWritten(),
                .verb = "REQUEST",
            };

            const body_tmpl = "{{\"subscription_key\":{}}}";
            const body_len = comptime std.fmt.count(body_tmpl, .{std.math.maxInt(u64)});
            var body_bytes: [body_len]u8 = undefined;
            var body_bytes_fbs = std.io.fixedBufferStream(&body_bytes);
            try std.fmt.format(body_bytes_fbs.writer(), body_tmpl, .{self.subscription_id});

            const body = moo.RawBody{
                .bytes = body_bytes_fbs.getWritten(),
            };

            var header = body.getHeader(request_id);
            header.content_type = "application/json";

            const req_buffer = try allocator.alloc(
                u8,
                meta.getEncodeSize() + header.getEncodeSize() + body.getEncodeSize(),
            );
            defer allocator.free(req_buffer);
            var req_fbs = std.io.fixedBufferStream(req_buffer);
            try moo.encode(req_fbs.writer(), meta, header, body);

            try self.conn.ws.writeBin(req_fbs.getWritten());
        }

        /// Returns incoming updates to the subject. If the subscription is unsubscribed,
        /// this function returns null.
        ///
        /// Calling this function before `.activate()` blocks the thread indefinetely.
        ///
        /// Caller have to release the returned memory using `Connection.allocator`.
        pub fn next(self: *@This(), allocator: std.mem.Allocator) std.mem.Allocator.Error!?[]const u8 {
            var entry = self.conn.responses.get(self.request_id) orelse {
                return null;
            };

            entry.wrote.wait();

            if (entry.canceled) {
                return null;
            }

            const data = entry.data orelse {
                return std.mem.Allocator.Error.OutOfMemory;
            };

            defer self.conn.allocator.free(data);

            {
                self.conn.responses_mutex.lock();
                defer self.conn.responses_mutex.unlock();
                entry.canceled = false;
                entry.data = null;
                entry.wrote.reset();
            }

            return try allocator.dupe(u8, data);
        }
    };
};

fn readLoop(conn: *Connection, on_request: RequestHandler) void {
    {
        conn.state_mutex.lock();
        defer conn.state_mutex.unlock();
        conn.state = .connected;
    }

    while (true) {
        const msg = (conn.ws.read() catch return) orelse unreachable;
        defer conn.ws.done(msg);

        switch (msg.type) {
            // NOTE: roon-node-api does not check whether message is binaryType.
            .text, .binary => {
                const meta, const header_ctx = moo.Metadata.parse(msg.data) catch |err| {
                    std.log.warn("Failed to parse MOO metadata: {s}", .{@errorName(err)});
                    continue;
                };

                // We just want to know Request-Id header.
                const header, _ = moo.NoBodyHeaders.parse(msg.data, header_ctx) catch |err| {
                    std.log.warn("Failed to parse MOO headers: {s}", .{@errorName(err)});
                    continue;
                };

                {
                    const response = response: {
                        conn.responses_mutex.lock();
                        defer conn.responses_mutex.unlock();
                        break :response conn.responses.get(header.request_id);
                    };

                    if (response) |store| {
                        if (store.wrote.isSet()) {
                            std.log.warn(
                                "Received more than one message having same Request-Id({d})",
                                .{header.request_id},
                            );
                            continue;
                        }

                        defer store.wrote.set();

                        const bytes = conn.allocator.dupe(u8, msg.data) catch |err| {
                            std.log.err("Unable to release incoming WS message: {s}", .{@errorName(err)});
                            return;
                        };

                        store.data = bytes;
                        continue;
                    }
                }

                var buffer = std.ArrayList(u8).init(conn.allocator);
                defer buffer.deinit();

                const wrote = on_request(buffer.writer(), meta, header_ctx, msg.data) catch |err| {
                    std.log.warn("Service server handler returned an error: {s}", .{@errorName(err)});
                    continue;
                };

                if (!wrote) {
                    std.log.info("Unhandled incoming request: service={s}\n", .{meta.service});
                    continue;
                }

                const bytes = buffer.toOwnedSlice() catch |err| {
                    std.log.warn("Unable to prepare response bytes: {s}\n", .{@errorName(err)});
                    continue;
                };
                defer conn.allocator.free(bytes);

                conn.ws.writeBin(bytes) catch |err| {
                    std.log.warn("Failed to write response message: {s}\n", .{@errorName(err)});
                    continue;
                };
            },
            .ping => conn.ws.writePong(msg.data) catch {},
            .pong => {},
            .close => {
                // TODO: Release Connection or notify to upstream
                conn.ws.close(.{}) catch return;
                return;
            },
        }
    }
}