474 lines
17 KiB
Zig
474 lines
17 KiB
Zig
//! 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);
|
|
|
|
/// 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: [transport.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 != transport.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 = transport.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 = transport.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});
|
|
};
|
|
}
|
|
}
|
|
};
|
|
|
|
/// 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: *[transport.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 [transport.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;
|
|
|
|
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)));
|
|
}
|