mokhtar/nxdns
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases 1 Wiki Activity

resolver transport: udp/tcp servers, doh/dot clients, pool failover with health

Browse Source
This commit is contained in:
mokhtar
2026-08-01 12:36:50 +02:00
parent 346f2dc502
commit 17d0401f8a
15 changed files with 6062 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/server/tcp_server.zig
+501
View File
@@ -0,0 +1,501 @@
//! 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)));
}