//! The TCP/53 listener.
//!
//! RFC 1035 §4.2.2 frames every message with a 2-byte big-endian length, and
//! RFC 7766 §6.2.1.1 lets one connection carry several queries. Both are
//! implemented here: a connection is answered serially until the client closes
//! it or the idle budget runs out.
//!
//! Connection slots are fixed and pre-allocated. Over capacity the listener
//! closes the new stream immediately and counts it; it never queues, and it
//! never allocates per connection.
//!
//! No stream read or write in 0.16.0 accepts a timeout, so every per-connection
//! operation is raced against `Options.idle_timeout` through `std.Io.Select` and
//! the loser is canceled.

const std = @import("std");
const handler = @import("handler.zig");
const transport = @import("../upstream/transport.zig");

const log = std.log.scoped(.tcp_server);

/// RFC 1035 §4.2.2: the message length prefix is two bytes, big-endian.
pub const prefix_len = 2;

/// The stream buffers only stage the framing bytes. A message longer than this
/// is read straight into `Conn.query` and written straight from `Conn.reply`,
/// so making them larger would buy nothing.
const stream_buffer_len = 1024;

/// How long the accept loop waits after an unexpected accept failure, so a
/// persistent one cannot turn the loop into a spin.
const retry_delay: std.Io.Clock.Duration = .{ .raw = .fromMilliseconds(100), .clock = .awake };

pub const Options = struct {
    max_connections: u16 = 64,
    /// RFC 7766 §6.2.3 recommends a few seconds of idle tolerance.
    idle_timeout: std.Io.Clock.Duration = .{ .raw = .fromSeconds(10), .clock = .awake },
};

pub const Stats = struct {
    accepted: std.atomic.Value(u64) = .init(0),
    rejected_at_capacity: std.atomic.Value(u64) = .init(0),
    rejected_at_shutdown: std.atomic.Value(u64) = .init(0),
    accept_errors: std.atomic.Value(u64) = .init(0),
    connection_errors: std.atomic.Value(u64) = .init(0),
    idle_timeouts: std.atomic.Value(u64) = .init(0),
};

/// Lifecycle of the accept loop. `serve` claims `.serving`, `deinit` publishes
/// `.closing`, and the two meet at `stopped` so no task touches a connection
/// slot after it is freed.
const State = enum(u32) { idle, serving, closing };

/// `.closing` exists so `deinit` never shuts down a descriptor that its own
/// task is about to close: the transition to `.closing` happens under the mutex
/// before the close, and `deinit` only touches `.active` slots.
const ConnState = enum { free, active, closing };

/// What the accept loop does with a stream it has just accepted.
const Claim = union(enum) {
    /// The stream owns `conns[index]`.
    slot: usize,
    /// Every slot is taken. The stream is closed and the loop continues.
    at_capacity,
    /// `deinit` has started. The stream is closed and the loop returns.
    shutting_down,
};

pub const TcpServer = struct {
    server: std.Io.net.Server,
    handler: *handler.Handler,
    conns: []Conn,
    mutex: std.Io.Mutex,
    /// Guarded by `mutex`. `deinit` sets it in the same critical section that
    /// shuts the active connections down, so a stream that arrives after that
    /// scan can never claim a slot the scan will not visit again.
    shutdown_begun: bool,
    options: Options,
    stats: Stats,
    state: std.atomic.Value(State),
    stopped: std.Io.Event,

    /// One slot is ~131 KiB, so the default 64 connections cost ~8.4 MiB, which
    /// is inside the PLAN §18 budget. The two message buffers cannot be shared
    /// or shrunk: the handler holds the query while the reply is built, and
    /// both ceilings are the 65535 bytes the length prefix can express.
    pub const Conn = struct {
        query: [transport.max_message_len]u8,
        reply: [transport.max_message_len]u8,
        read_buf: [stream_buffer_len]u8,
        write_buf: [stream_buffer_len]u8,
        stream: std.Io.net.Stream,
        /// Guarded by `TcpServer.mutex`.
        state: ConnState,
    };

    pub const ListenError = std.Io.net.IpAddress.ListenError || error{OutOfMemory};

    pub fn listen(
        gpa: std.mem.Allocator,
        io: std.Io,
        address: std.Io.net.IpAddress,
        h: *handler.Handler,
        options: Options,
    ) ListenError!TcpServer {
        std.debug.assert(options.max_connections > 0);

        const conns = try gpa.alloc(Conn, options.max_connections);
        errdefer gpa.free(conns);
        for (conns) |*conn| conn.state = .free;

        const local = address;
        const server = try local.listen(io, .{ .reuse_address = true });

        return .{
            .server = server,
            .handler = h,
            .conns = conns,
            .mutex = .init,
            .shutdown_begun = false,
            .options = options,
            .stats = .{},
            .state = .init(.idle),
            .stopped = .unset,
        };
    }

    /// The kernel-assigned address. A port of 0 in `listen` resolves here.
    pub fn boundAddress(self: *const TcpServer) std.Io.net.IpAddress {
        return self.server.socket.address;
    }

    /// Accept loop. Returns when the task is canceled or `deinit` stops it.
    pub fn serve(self: *TcpServer, io: std.Io) void {
        if (self.state.cmpxchgStrong(.idle, .serving, .acq_rel, .acquire) != null) return;

        var group: std.Io.Group = .init;
        self.acceptLoop(io, &group);

        // A reply that is half written is worse than no reply, so the live
        // connections are awaited even when this task is being canceled.
        const prev = io.swapCancelProtection(.blocked);
        group.await(io) catch |err| switch (err) {
            error.Canceled => unreachable,
        };
        _ = io.swapCancelProtection(prev);

        self.stopped.set(io);
    }

    pub fn deinit(self: *TcpServer, gpa: std.mem.Allocator, io: std.Io) void {
        const was_serving = self.state.swap(.closing, .acq_rel) == .serving;

        // Shutting the listening socket down is the documented way to unblock a
        // pending `accept`: it fails with `error.SocketNotListening`.
        const listener: std.Io.net.Stream = .{ .socket = self.server.socket };
        listener.shutdown(io, .both) catch |err| {
            log.debug("tcp listener shutdown failed: {t}", .{err});
        };

        // A live connection is blocked in a read that only the idle budget
        // would end, which is seconds away. Shutting each one down bounds this,
        // and the same critical section closes the door on new connections.
        self.beginShutdown(io);

        if (was_serving) self.stopped.waitUncancelable(io);

        self.server.deinit(io);
        gpa.free(self.conns);
        self.* = undefined;
    }

    fn acceptLoop(self: *TcpServer, io: std.Io, group: *std.Io.Group) void {
        while (self.state.load(.acquire) == .serving) {
            const stream = self.server.accept(io) catch |err| switch (err) {
                error.Canceled, error.SocketNotListening => return,
                else => {
                    bump(&self.stats.accept_errors);
                    log.debug("tcp accept failed: {t}", .{err});
                    retry_delay.sleep(io) catch return;
                    continue;
                },
            };

            const index = switch (self.claim(io, stream)) {
                .slot => |index| index,
                // Refusing now is honest; a queue would only hide the overload.
                .at_capacity => {
                    bump(&self.stats.rejected_at_capacity);
                    stream.close(io);
                    continue;
                },
                // `deinit` will not see this stream in any slot, so serving it
                // would hold `deinit` for the whole idle budget.
                .shutting_down => {
                    bump(&self.stats.rejected_at_shutdown);
                    stream.close(io);
                    return;
                },
            };

            group.concurrent(io, serveConn, .{ self, io, index }) catch |err| switch (err) {
                error.ConcurrencyUnavailable => {
                    bump(&self.stats.rejected_at_capacity);
                    self.finish(io, index);
                    continue;
                },
            };

            bump(&self.stats.accepted);
        }
    }

    fn serveConn(self: *TcpServer, io: std.Io, index: usize) void {
        defer self.finish(io, index);

        const conn = &self.conns[index];
        var reader = conn.stream.reader(io, &conn.read_buf);
        var writer = conn.stream.writer(io, &conn.write_buf);
        const budget = self.options.idle_timeout;

        while (true) {
            var prefix: [prefix_len]u8 = undefined;
            var got: usize = 0;
            switch (race(io, budget, readPrefix, .{ &reader.interface, &prefix, &got })) {
                .ok => {},
                .timed_out => {
                    bump(&self.stats.idle_timeouts);
                    return;
                },
                .canceled => return,
                .failed => {
                    bump(&self.stats.connection_errors);
                    return;
                },
            }

            // A client that closes between messages has finished asking, which
            // is the normal end of a connection, not a failure.
            if (got == 0) return;
            if (got != prefix_len) {
                bump(&self.stats.connection_errors);
                return;
            }

            // RFC 1035 §4.2.2 gives no meaning to a zero-length message, and
            // the prefix is a u16 so it can never exceed `max_message_len`.
            const len = parsePrefix(prefix);
            if (len == 0) {
                bump(&self.stats.connection_errors);
                return;
            }

            switch (race(io, budget, readBody, .{ &reader.interface, conn.query[0..len] })) {
                .ok => {},
                .canceled => return,
                // A half-sent message is a broken peer, not an idle one.
                .timed_out, .failed => {
                    bump(&self.stats.connection_errors);
                    return;
                },
            }

            const bytes = switch (self.handler.handle(io, .tcp, conn.query[0..len], &conn.reply)) {
                // There is no framing for "no answer", so the connection ends.
                .drop => return,
                .reply => |b| b,
            };

            const out = framePrefix(@intCast(bytes.len));
            switch (race(io, budget, writeReply, .{ &writer.interface, &out, bytes })) {
                .ok => {},
                .canceled => return,
                .timed_out, .failed => {
                    bump(&self.stats.connection_errors);
                    return;
                },
            }
        }
    }

    fn claim(self: *TcpServer, io: std.Io, stream: std.Io.net.Stream) Claim {
        // Uncancelable: this section takes no Io and never blocks on a peer, so
        // it cannot deadlock, and losing the lock mid-update would leak a slot.
        self.mutex.lockUncancelable(io);
        defer self.mutex.unlock(io);

        const outcome = decideClaim(self.conns, self.shutdown_begun);
        switch (outcome) {
            .slot => |index| {
                self.conns[index].stream = stream;
                self.conns[index].state = .active;
            },
            .at_capacity, .shutting_down => {},
        }
        return outcome;
    }

    fn finish(self: *TcpServer, io: std.Io, index: usize) void {
        const conn = &self.conns[index];

        self.mutex.lockUncancelable(io);
        conn.state = .closing;
        self.mutex.unlock(io);

        // The socket is released even when this task is being torn down: the
        // next cancelable call would otherwise skip the close.
        const prev = io.swapCancelProtection(.blocked);
        conn.stream.close(io);
        _ = io.swapCancelProtection(prev);

        self.mutex.lockUncancelable(io);
        conn.state = .free;
        self.mutex.unlock(io);
    }

    /// Closes the door on new connections and unblocks the live ones. Both
    /// happen under one hold of the mutex: a `claim` that runs before this
    /// leaves an `.active` slot the loop below shuts down, and a `claim` that
    /// runs after it reads `shutdown_begun` and takes no slot at all.
    fn beginShutdown(self: *TcpServer, io: std.Io) void {
        self.mutex.lockUncancelable(io);
        defer self.mutex.unlock(io);

        self.shutdown_begun = true;

        for (self.conns) |*conn| {
            if (conn.state != .active) continue;
            conn.stream.shutdown(io, .both) catch |err| {
                log.debug("tcp connection shutdown failed: {t}", .{err});
            };
        }
    }
};

/// RFC 1035 §4.2.2: the message length as a 2-byte big-endian integer.
pub fn framePrefix(len: u16) [prefix_len]u8 {
    var out: [prefix_len]u8 = undefined;
    std.mem.writeInt(u16, &out, len, .big);
    return out;
}

pub fn parsePrefix(bytes: [prefix_len]u8) u16 {
    return std.mem.readInt(u16, &bytes, .big);
}

/// The capacity rule, without the mutex, so it is testable without a backend.
fn firstFree(conns: []const TcpServer.Conn) ?usize {
    for (conns, 0..) |*conn, index| {
        if (conn.state == .free) return index;
    }
    return null;
}

/// The whole claim rule, without the mutex. Shutdown outranks capacity: a free
/// slot is still refused once `deinit` has passed the connections.
fn decideClaim(conns: []const TcpServer.Conn, shutdown_begun: bool) Claim {
    if (shutdown_begun) return .shutting_down;
    const index = firstFree(conns) orelse return .at_capacity;
    return .{ .slot = index };
}

const Outcome = union(enum) {
    op: anyerror!void,
    expiry: std.Io.Cancelable!void,
};

const Result = enum { ok, timed_out, failed, canceled };

/// Runs one connection operation against the idle budget and cancels the loser.
fn race(
    io: std.Io,
    budget: std.Io.Clock.Duration,
    comptime f: anytype,
    args: std.meta.ArgsTuple(@TypeOf(f)),
) Result {
    var outcomes: [2]Outcome = undefined;
    var select: std.Io.Select(Outcome) = .init(io, &outcomes);
    defer select.cancelDiscard();

    select.concurrent(.op, f, args) catch |err| switch (err) {
        error.ConcurrencyUnavailable => return .failed,
    };
    select.concurrent(.expiry, expire, .{ io, budget }) catch |err| switch (err) {
        error.ConcurrencyUnavailable => return .failed,
    };

    return switch (select.await() catch return .canceled) {
        .op => |result| if (result) |_| .ok else |err| switch (err) {
            error.Canceled => .canceled,
            else => .failed,
        },
        // A canceled sleep means this task is being torn down, not that the
        // client went idle.
        .expiry => |result| if (result) |_| .timed_out else |_| .canceled,
    };
}

fn expire(io: std.Io, budget: std.Io.Clock.Duration) std.Io.Cancelable!void {
    return budget.sleep(io);
}

/// `readSliceShort` rather than `readSliceAll`: a zero-length read is a client
/// that closed cleanly between messages, and only a partial prefix is an error.
fn readPrefix(reader: *std.Io.Reader, buf: *[prefix_len]u8, out_len: *usize) anyerror!void {
    out_len.* = try reader.readSliceShort(buf);
}

fn readBody(reader: *std.Io.Reader, buf: []u8) anyerror!void {
    return reader.readSliceAll(buf);
}

fn writeReply(writer: *std.Io.Writer, prefix: *const [prefix_len]u8, bytes: []const u8) anyerror!void {
    try writer.writeAll(prefix);
    try writer.writeAll(bytes);
    try writer.flush();
}

fn bump(counter: *std.atomic.Value(u64)) void {
    _ = counter.fetchAdd(1, .monotonic);
}

const testing = std.testing;

test "the length prefix is big-endian and round-trips" {
    try testing.expectEqualSlices(u8, &.{ 0x00, 0x00 }, &framePrefix(0));
    try testing.expectEqualSlices(u8, &.{ 0x01, 0x00 }, &framePrefix(256));
    try testing.expectEqualSlices(u8, &.{ 0xff, 0xff }, &framePrefix(65535));

    for ([_]u16{ 0, 1, 12, 512, 4096, 65534, 65535 }) |len| {
        try testing.expectEqual(len, parsePrefix(framePrefix(len)));
    }
}

test "the prefix ceiling is the message ceiling" {
    try testing.expectEqual(@as(u16, transport.max_message_len), parsePrefix(.{ 0xff, 0xff }));
}

fn testConns(count: usize) ![]TcpServer.Conn {
    const conns = try testing.allocator.alloc(TcpServer.Conn, count);
    for (conns) |*conn| conn.state = .free;
    return conns;
}

test "the connection pool hands out every slot once" {
    const conns = try testConns(3);
    defer testing.allocator.free(conns);

    for (0..conns.len) |expected| {
        const index = firstFree(conns) orelse return error.TestUnexpectedResult;
        try testing.expectEqual(expected, index);
        conns[index].state = .active;
    }
}

test "a full connection pool refuses instead of growing" {
    const conns = try testConns(2);
    defer testing.allocator.free(conns);

    for (conns) |*conn| conn.state = .active;
    try testing.expectEqual(@as(?usize, null), firstFree(conns));
}

test "a closing slot is not reused until it is free" {
    const conns = try testConns(2);
    defer testing.allocator.free(conns);

    conns[0].state = .active;
    conns[1].state = .closing;
    try testing.expectEqual(@as(?usize, null), firstFree(conns));

    conns[1].state = .free;
    try testing.expectEqual(@as(?usize, 1), firstFree(conns));
}

test "a claim takes the first free slot before shutdown" {
    const conns = try testConns(2);
    defer testing.allocator.free(conns);

    conns[0].state = .active;
    try testing.expectEqual(@as(usize, 1), decideClaim(conns, false).slot);
}

test "a claim after shutdown is refused even with a free slot" {
    const conns = try testConns(2);
    defer testing.allocator.free(conns);

    try testing.expectEqual(.shutting_down, std.meta.activeTag(decideClaim(conns, true)));

    // The refusal must not consume the slot: `deinit` frees it, nothing else.
    try testing.expectEqual(@as(?usize, 0), firstFree(conns));
}

test "shutdown outranks capacity" {
    const conns = try testConns(1);
    defer testing.allocator.free(conns);

    conns[0].state = .active;
    try testing.expectEqual(.at_capacity, std.meta.activeTag(decideClaim(conns, false)));
    try testing.expectEqual(.shutting_down, std.meta.activeTag(decideClaim(conns, true)));
}