server: every dns listener shares one budget of 32 queries, stream listeners take 16 connections, the task pool is capped from those bounds, and work it cannot start is counted by site

This commit is contained in:
2026-10-22 15:24:10 +02:00
parent e4e405f8e9
commit 63cc5f8dc3
23 changed files with 983 additions and 324 deletions
+51 -183
View File
@@ -1,49 +1,39 @@
//! The UDP/53 listener.
//!
//! One receive loop feeds a bounded pool of in-flight tasks. A datagram that
//! finds no free slot 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.
//! 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.
//!
//! Every slot carries its own query and reply buffer, so answering a datagram
//! allocates nothing. The pool is sized once in `bind` and never grows.
//! 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 transport = @import("../upstream/transport.zig");
const query_budget = @import("query_budget.zig");
const log = std.log.scoped(.udp_server);
/// The EDNS ceiling this server advertises. A larger datagram arrives
/// truncated and is dropped rather than guessed at.
pub const max_datagram = 4096;
comptime {
// `handler.handle` asserts this range on the buffers it is given.
std.debug.assert(max_datagram >= handler.udp_limit_min);
std.debug.assert(max_datagram >= handler.udp_limit_max);
std.debug.assert(transport.max_message_len >= handler.udp_limit_min);
}
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 Options = struct {
max_in_flight: u16 = 64,
};
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.
/// 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),
@@ -63,69 +53,31 @@ pub const Snapshot = struct {
};
/// Lifecycle of the receive loop. `serve` claims `.serving`, `deinit` publishes
/// `.closing`, and the two meet at `stopped` so no task touches the slots after
/// they are freed.
/// `.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,
slots: []Slot,
mutex: std.Io.Mutex,
budget: *query_budget.QueryBudget,
stats: Stats,
state: std.atomic.Value(State),
stopped: std.Io.Event,
pub const Slot = struct {
query: [max_datagram]u8,
/// The buffer the handler receives the upstream answer into, so it
/// holds a whole DNS message, not a whole datagram. A reply between
/// `max_datagram` and `transport.max_message_len` must reach the
/// handler intact: only the handler can turn it into the TC=1 reply
/// that makes the client retry over TCP (S5.3). A 4096-byte buffer
/// would instead fail the exchange with `error.ResponseTooLarge` and
/// answer SERVFAIL, which is wrong.
///
/// The datagram actually sent stays bounded by `udpLimit` inside the
/// handler, so nothing larger than 4096 bytes leaves this socket.
///
/// Cost: 4096 + 65535 + the scratch below ≈ 74 KiB per slot, so the
/// default 64 slots hold ≈ 4.6 MiB. The PLAN §18 budget is 100 MiB with
/// ~1M blocked domains, so this pool takes about 5% of it.
reply: [transport.max_message_len]u8,
/// The handler's per-query working memory. It belongs to the slot so
/// that answering a datagram still allocates nothing, and one slot
/// serves one query at a time.
scratch: handler.Scratch,
from: std.Io.net.IpAddress,
len: usize,
/// Guarded by `UdpServer.mutex`.
in_use: bool,
};
pub const BindError = std.Io.net.IpAddress.BindError || error{OutOfMemory};
pub const BindError = std.Io.net.IpAddress.BindError;
pub fn bind(
gpa: std.mem.Allocator,
io: std.Io,
bind_address: std.Io.net.IpAddress,
h: *handler.Handler,
options: Options,
budget: *query_budget.QueryBudget,
) BindError!UdpServer {
std.debug.assert(options.max_in_flight > 0);
const slots = try gpa.alloc(Slot, options.max_in_flight);
errdefer gpa.free(slots);
for (slots) |*slot| slot.in_use = false;
const local = bind_address;
const socket = try local.bind(io, .{ .mode = .dgram });
const socket = try bind_address.bind(io, .{ .mode = .dgram });
return .{
.socket = socket,
.handler = h,
.slots = slots,
.mutex = .init,
.budget = budget,
.stats = .{},
.state = .init(.idle),
.stopped = .unset,
@@ -139,7 +91,7 @@ pub const UdpServer = struct {
/// 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, and the receive loop takes that lock per datagram.
/// no consumer needs.
pub fn snapshotStats(self: *const UdpServer) Snapshot {
return .{
.received = self.stats.received.load(.monotonic),
@@ -169,172 +121,88 @@ pub const UdpServer = struct {
self.stopped.set(io);
}
pub fn deinit(self: *UdpServer, gpa: std.mem.Allocator, io: std.Io) void {
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);
gpa.free(self.slots);
self.* = undefined;
}
fn receiveLoop(self: *UdpServer, io: std.Io, group: *std.Io.Group) void {
// A datagram that finds no free slot still has to leave the socket
// buffer, or the loop would spin on the same undeliverable packet.
var overflow: [max_datagram]u8 = undefined;
// 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 claimed = self.claim(io);
const buffer = if (claimed) |index| &self.slots[index].query else &overflow;
const msg = self.socket.receiveTimeout(io, buffer, .{ .duration = poll_interval }) catch |err| {
if (claimed) |index| self.release(io, index);
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;
},
}
// 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);
const index = claimed orelse {
bump(&self.stats.dropped_no_slot);
continue;
};
// 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);
self.release(io, index);
continue;
}
const slot = &self.slots[index];
slot.from = msg.from;
slot.len = msg.data.len;
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, index }) catch |err| switch (err) {
group.concurrent(io, respondOne, .{ self, io, work, msg.from, msg.data.len }) catch |err| switch (err) {
error.ConcurrencyUnavailable => {
bump(&self.stats.dropped_no_slot);
self.release(io, index);
concurrency.unavailable(.udp_query);
self.budget.release(io, work);
},
};
}
}
fn respondOne(self: *UdpServer, io: std.Io, index: usize) void {
const slot = &self.slots[index];
defer self.release(io, index);
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(slot.from),
slot.query[0..slot.len],
&slot.reply,
&slot.scratch,
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, &slot.from, bytes) catch |err| {
.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 claim(self: *UdpServer, io: std.Io) ?usize {
// 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 index = firstFree(self.slots) orelse return null;
self.slots[index].in_use = true;
return index;
}
fn release(self: *UdpServer, io: std.Io, index: usize) void {
self.mutex.lockUncancelable(io);
defer self.mutex.unlock(io);
self.slots[index].in_use = false;
}
};
/// The capacity rule, without the mutex, so it is testable without a backend.
fn firstFree(slots: []const UdpServer.Slot) ?usize {
for (slots, 0..) |*slot, index| {
if (!slot.in_use) return index;
}
return null;
}
fn bump(counter: *std.atomic.Value(u64)) void {
_ = counter.fetchAdd(1, .monotonic);
}
const testing = std.testing;
fn testSlots(count: usize) ![]UdpServer.Slot {
const slots = try testing.allocator.alloc(UdpServer.Slot, count);
for (slots) |*slot| slot.in_use = false;
return slots;
}
test "the slot pool hands out every index once" {
const slots = try testSlots(3);
defer testing.allocator.free(slots);
for (0..slots.len) |expected| {
const index = firstFree(slots) orelse return error.TestUnexpectedResult;
try testing.expectEqual(expected, index);
slots[index].in_use = true;
}
}
test "a full slot pool refuses instead of growing" {
const slots = try testSlots(2);
defer testing.allocator.free(slots);
for (slots) |*slot| slot.in_use = true;
try testing.expectEqual(@as(?usize, null), firstFree(slots));
}
test "a slot's reply buffer holds a whole DNS message" {
const slots = try testSlots(1);
defer testing.allocator.free(slots);
// The handler asserts both ends of this range on the buffer it is given,
// and an upstream reply of up to 65535 bytes has to fit so the handler can
// truncate it instead of the exchange failing.
try testing.expectEqual(@as(usize, transport.max_message_len), slots[0].reply.len);
try testing.expect(slots[0].reply.len >= handler.udp_limit_min);
try testing.expect(slots[0].reply.len > max_datagram);
}
test "the default slot pool stays inside the memory budget" {
// 4096 + 65535 + scratch ≈ 74 KiB per slot; 64 slots ≈ 4.6 MiB, against the
// 100 MiB of PLAN §18.
const options: Options = .{};
const pool_bytes = @sizeOf(UdpServer.Slot) * @as(usize, options.max_in_flight);
try testing.expect(pool_bytes < 8 * 1024 * 1024);
}
test "a released slot is handed out again" {
const slots = try testSlots(2);
defer testing.allocator.free(slots);
for (slots) |*slot| slot.in_use = true;
slots[1].in_use = false;
try testing.expectEqual(@as(?usize, 1), firstFree(slots));
}