//! The UDP/53 listener.
//!
//! One receive loop hands each datagram to a task that holds a permit from the
//! shared query budget. A datagram that finds no permit is dropped and counted,
//! never queued: an unbounded queue only moves the failure from "visible drop"
//! to "invisible latency", which is the failure mode PLAN §1 names.
//!
//! The permit carries the query, the reply and the handler's scratch, so
//! answering a datagram allocates nothing.

const std = @import("std");
const address = @import("../platform/address.zig");
const concurrency = @import("../platform/concurrency.zig");
const handler = @import("handler.zig");
const query_budget = @import("query_budget.zig");

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

pub const max_datagram = query_budget.max_datagram;

/// How long a receive waits before the loop re-reads the stop flag. `deinit`
/// therefore returns within this much time, without depending on a blocked
/// `receive` being unblocked by a close — a guarantee POSIX does not make.
const poll_interval: std.Io.Clock.Duration = .{ .raw = .fromMilliseconds(100), .clock = .awake };

pub const Stats = struct {
    received: std.atomic.Value(u64) = .init(0),
    /// The datagram arrived with `flags.trunc`: the kernel discarded its tail.
    dropped_oversize: std.atomic.Value(u64) = .init(0),
    /// Every permit of the shared query budget was taken.
    dropped_no_slot: std.atomic.Value(u64) = .init(0),
    /// The handler answered `.drop` — an upstream exchange canceled by
    /// shutdown. The handler counts its own malformed drops; this counter is
    /// what keeps the listener's "every dropped datagram is counted" rule true
    /// for the drops it cannot see. A datagram that got a permit but no task
    /// is counted under `nxdns_concurrency_unavailable_total{site="udp_query"}`.
    dropped_handler: std.atomic.Value(u64) = .init(0),
    receive_errors: std.atomic.Value(u64) = .init(0),
    send_errors: std.atomic.Value(u64) = .init(0),
};

/// A plain copy of `Stats`, the shape `metrics.counterGroup` walks for the
/// `nxdns_udp_server_*` families. Every counter is exported: the module doc
/// promises that a dropped datagram is counted, and a count nothing can read is
/// not a count.
pub const Snapshot = struct {
    received: u64,
    dropped_oversize: u64,
    dropped_no_slot: u64,
    dropped_handler: u64,
    receive_errors: u64,
    send_errors: u64,
};

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

pub const UdpServer = struct {
    socket: std.Io.net.Socket,
    handler: *handler.Handler,
    budget: *query_budget.QueryBudget,
    stats: Stats,
    state: std.atomic.Value(State),
    stopped: std.Io.Event,

    pub const BindError = std.Io.net.IpAddress.BindError;

    pub fn bind(
        io: std.Io,
        bind_address: std.Io.net.IpAddress,
        h: *handler.Handler,
        budget: *query_budget.QueryBudget,
    ) BindError!UdpServer {
        const socket = try bind_address.bind(io, .{ .mode = .dgram });
        return .{
            .socket = socket,
            .handler = h,
            .budget = budget,
            .stats = .{},
            .state = .init(.idle),
            .stopped = .unset,
        };
    }

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

    /// The counters, read one at a time. A scrape that lands mid-datagram can
    /// see a receive counted before its drop is; a lock would buy a consistency
    /// no consumer needs.
    pub fn snapshotStats(self: *const UdpServer) Snapshot {
        return .{
            .received = self.stats.received.load(.monotonic),
            .dropped_oversize = self.stats.dropped_oversize.load(.monotonic),
            .dropped_no_slot = self.stats.dropped_no_slot.load(.monotonic),
            .dropped_handler = self.stats.dropped_handler.load(.monotonic),
            .receive_errors = self.stats.receive_errors.load(.monotonic),
            .send_errors = self.stats.send_errors.load(.monotonic),
        };
    }

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

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

        // A reply that is half sent is worse than no reply, so the in-flight
        // tasks 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: *UdpServer, io: std.Io) void {
        // `serve` reads this between receives, so it stops on its own and the
        // socket is closed only once nothing can be reading from it.
        if (self.state.swap(.closing, .acq_rel) == .serving) self.stopped.waitUncancelable(io);
        self.socket.close(io);
        self.* = undefined;
    }

    fn receiveLoop(self: *UdpServer, io: std.Io, group: *std.Io.Group) void {
        // The datagram is received here, without a permit, so an idle listener
        // holds none; a query that gets one is copied into it.
        var datagram: [max_datagram]u8 = undefined;

        while (self.state.load(.acquire) == .serving) {
            const msg = self.socket.receiveTimeout(io, &datagram, .{ .duration = poll_interval }) catch |err| {
                switch (err) {
                    error.Canceled => return,
                    // The stop flag is re-read at the top of the loop.
                    error.Timeout => continue,
                    error.ConcurrencyUnavailable => concurrency.unavailable(.udp_receive),
                    else => {
                        bump(&self.stats.receive_errors);
                        log.debug("udp receive failed: {t}", .{err});
                    },
                }
                // One bad datagram must not kill the listener, and a persistent
                // error must not turn it into a spin.
                poll_interval.sleep(io) catch return;
                continue;
            };

            bump(&self.stats.received);

            // The kernel already threw the tail away, so the message cannot be
            // parsed and any answer would be a guess.
            if (msg.flags.trunc) {
                bump(&self.stats.dropped_oversize);
                continue;
            }

            const work = self.budget.tryAcquire(io) orelse {
                bump(&self.stats.dropped_no_slot);
                continue;
            };
            @memcpy(work.datagram[0..msg.data.len], msg.data);

            group.concurrent(io, respondOne, .{ self, io, work, msg.from, msg.data.len }) catch |err| switch (err) {
                error.ConcurrencyUnavailable => {
                    concurrency.unavailable(.udp_query);
                    self.budget.release(io, work);
                },
            };
        }
    }

    fn respondOne(
        self: *UdpServer,
        io: std.Io,
        work: *query_budget.Work,
        from: std.Io.net.IpAddress,
        len: usize,
    ) void {
        defer self.budget.release(io, work);

        const outcome = self.handler.handle(
            io,
            .udp,
            address.NetAddress.fromIp(from),
            work.datagram[0..len],
            &work.reply,
            &work.scratch,
        );
        switch (outcome) {
            .drop => bump(&self.stats.dropped_handler),
            .reply => |bytes| self.socket.send(io, &from, bytes) catch |err| {
                bump(&self.stats.send_errors);
                log.debug("udp send to a client failed: {t}", .{err});
            },
        }
    }
};

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