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
+187 -7
View File
@@ -13,7 +13,9 @@ const std = @import("std");
const build_options = @import("build_options");
const net = std.Io.net;
const concurrency = @import("../platform/concurrency.zig");
const handler = @import("handler.zig");
const query_budget = @import("query_budget.zig");
const tcp_server = @import("tcp_server.zig");
const model = @import("../config/model.zig");
const response = @import("../filter/response.zig");
@@ -179,7 +181,9 @@ test "two length-prefixed queries share one connection" {
var h = bareHandler(h_owner.client(fake.client()));
const listen_address: net.IpAddress = try .parse("127.0.0.1", 0);
var server = try tcp_server.TcpServer.listen(gpa, io, listen_address, &h, .{ .max_connections = 2 });
var permits: query_budget.QueryBudget = try .init(gpa, 4);
defer permits.deinit();
var server = try tcp_server.TcpServer.listen(gpa, io, listen_address, &h, &permits, .{ .max_connections = 2 });
const server_address = server.boundAddress();
var group: std.Io.Group = .init;
@@ -213,7 +217,9 @@ test "the claimed slot records the connecting client" {
var h = bareHandler(h_owner.client(fake.client()));
const listen_address: net.IpAddress = try .parse("127.0.0.1", 0);
var server = try tcp_server.TcpServer.listen(gpa, io, listen_address, &h, .{ .max_connections = 2 });
var permits: query_budget.QueryBudget = try .init(gpa, 4);
defer permits.deinit();
var server = try tcp_server.TcpServer.listen(gpa, io, listen_address, &h, &permits, .{ .max_connections = 2 });
const server_address = server.boundAddress();
var group: std.Io.Group = .init;
@@ -295,7 +301,9 @@ test "a canceled serve does not wait for a live connection" {
// The idle budget is the whole time a drain would have to wait out, so it
// is set far beyond any patience this test run has: if `serve` waits for
// the connection instead of canceling it, the wait never ends.
var server = try tcp_server.TcpServer.listen(gpa, io, listen_address, &h, .{
var permits: query_budget.QueryBudget = try .init(gpa, 4);
defer permits.deinit();
var server = try tcp_server.TcpServer.listen(gpa, io, listen_address, &h, &permits, .{
.max_connections = 2,
.idle_timeout = .{ .raw = .fromSeconds(600), .clock = .awake },
});
@@ -346,7 +354,9 @@ test "the connections peak is the most slots held at once" {
var h = bareHandler(h_owner.client(fake.client()));
const listen_address: net.IpAddress = try .parse("127.0.0.1", 0);
var server = try tcp_server.TcpServer.listen(gpa, io, listen_address, &h, .{
var permits: query_budget.QueryBudget = try .init(gpa, 4);
defer permits.deinit();
var server = try tcp_server.TcpServer.listen(gpa, io, listen_address, &h, &permits, .{
.max_connections = 3,
.idle_timeout = .{ .raw = .fromSeconds(600), .clock = .awake },
});
@@ -401,7 +411,9 @@ test "an idle connection is closed and counted" {
var h = bareHandler(h_owner.client(fake.client()));
const listen_address: net.IpAddress = try .parse("127.0.0.1", 0);
var server = try tcp_server.TcpServer.listen(gpa, io, listen_address, &h, .{
var permits: query_budget.QueryBudget = try .init(gpa, 4);
defer permits.deinit();
var server = try tcp_server.TcpServer.listen(gpa, io, listen_address, &h, &permits, .{
.max_connections = 2,
.idle_timeout = short_idle,
});
@@ -436,7 +448,9 @@ test "a zero-length message is a connection error" {
var h = bareHandler(h_owner.client(fake.client()));
const listen_address: net.IpAddress = try .parse("127.0.0.1", 0);
var server = try tcp_server.TcpServer.listen(gpa, io, listen_address, &h, .{
var permits: query_budget.QueryBudget = try .init(gpa, 4);
defer permits.deinit();
var server = try tcp_server.TcpServer.listen(gpa, io, listen_address, &h, &permits, .{
.max_connections = 2,
.idle_timeout = short_idle,
});
@@ -490,7 +504,9 @@ test "deinit ends a serve loop that is blocked on accept" {
var h = bareHandler(h_owner.client(fake.client()));
const listen_address: net.IpAddress = try .parse("127.0.0.1", 0);
var server = try tcp_server.TcpServer.listen(gpa, io, listen_address, &h, .{ .max_connections = 2 });
var permits: query_budget.QueryBudget = try .init(gpa, 4);
defer permits.deinit();
var server = try tcp_server.TcpServer.listen(gpa, io, listen_address, &h, &permits, .{ .max_connections = 2 });
var group: std.Io.Group = .init;
try group.concurrent(io, tcp_server.TcpServer.serve, .{ &server, io });
@@ -501,3 +517,167 @@ test "deinit ends a serve loop that is blocked on accept" {
error.Canceled => unreachable,
};
}
/// Writes one framed query and returns the frame the server answered with.
fn askOn(reader: *std.Io.Reader, writer: *std.Io.Writer) ![]const u8 {
try writer.writeAll(&transport.framePrefix(@intCast(query_bytes.len)));
try writer.writeAll(query_bytes);
try writer.flush();
const len = transport.parsePrefix((try reader.takeArray(transport.prefix_len)).*);
return reader.take(len);
}
/// One answered query, then the connection idles with every permit back in the
/// budget; a second query that finds none taken ends the connection unanswered.
fn keepAliveThenStarved(io: std.Io, address: net.IpAddress, permits: *query_budget.QueryBudget) anyerror!void {
var stream = try address.connect(io, .{ .mode = .stream });
defer stream.close(io);
var read_buf: [1024]u8 = undefined;
var write_buf: [1024]u8 = undefined;
var reader = stream.reader(io, &read_buf);
var writer = stream.writer(io, &write_buf);
try expectAnswersQuery(try askOn(&reader.interface, &writer.interface));
// The server returns the permit just after its reply's last byte, which can
// reach this side first. `bounded` caps the wait.
const tick: std.Io.Clock.Duration = .{ .raw = .fromMilliseconds(1), .clock = .awake };
while (permits.taken(io) != 0) try tick.sleep(io);
const held = permits.tryAcquire(io).?;
defer permits.release(io, held);
try writer.interface.writeAll(&transport.framePrefix(@intCast(query_bytes.len)));
try writer.interface.writeAll(query_bytes);
try writer.interface.flush();
var sink: [64]u8 = undefined;
if (try reader.interface.readSliceShort(&sink) != 0) return error.TestUnexpectedBytes;
}
test "a keep-alive connection holds no permit between queries, and a query without one ends it" {
if (!build_options.integration) return error.SkipZigTest;
const gpa = testing.allocator;
var threaded: std.Io.Threaded = .init(gpa, .{});
defer threaded.deinit();
const io = threaded.io();
var fake: FakeUpstream = .{ .reply = response_bytes };
var h_owner: upstream_owner.Borrowed = .{};
var h = bareHandler(h_owner.client(fake.client()));
const listen_address: net.IpAddress = try .parse("127.0.0.1", 0);
var permits: query_budget.QueryBudget = try .init(gpa, 1);
defer permits.deinit();
var server = try tcp_server.TcpServer.listen(gpa, io, listen_address, &h, &permits, .{ .max_connections = 2 });
const server_address = server.boundAddress();
var group: std.Io.Group = .init;
try group.concurrent(io, tcp_server.TcpServer.serve, .{ &server, io });
defer {
server.deinit(io);
group.await(io) catch |err| switch (err) {
error.Canceled => unreachable,
};
}
try bounded(io, keepAliveThenStarved, .{ io, server_address, &permits });
try testing.expectEqual(@as(u64, 1), server.core.stats.dropped_no_permit.load(.monotonic));
try testing.expectEqual(@as(u64, 0), server.core.stats.connection_errors.load(.monotonic));
try testing.expectEqual(@as(u64, 1), h.stats.queries.load(.monotonic));
}
/// Holds the only slot with an answered connection, then opens a second one,
/// which the server closes without reading.
fn oneHeldOneRefused(io: std.Io, address: net.IpAddress) anyerror!void {
var held = try address.connect(io, .{ .mode = .stream });
defer held.close(io);
var read_buf: [1024]u8 = undefined;
var write_buf: [1024]u8 = undefined;
var reader = held.reader(io, &read_buf);
var writer = held.writer(io, &write_buf);
try expectAnswersQuery(try askOn(&reader.interface, &writer.interface));
var refused = try address.connect(io, .{ .mode = .stream });
defer refused.close(io);
var refused_buf: [64]u8 = undefined;
var refused_reader = refused.reader(io, &refused_buf);
var sink: [64]u8 = undefined;
if (try refused_reader.interface.readSliceShort(&sink) != 0) return error.TestUnexpectedBytes;
}
test "a listener at its connection cap refuses and counts" {
if (!build_options.integration) return error.SkipZigTest;
const gpa = testing.allocator;
var threaded: std.Io.Threaded = .init(gpa, .{});
defer threaded.deinit();
const io = threaded.io();
var fake: FakeUpstream = .{ .reply = response_bytes };
var h_owner: upstream_owner.Borrowed = .{};
var h = bareHandler(h_owner.client(fake.client()));
const listen_address: net.IpAddress = try .parse("127.0.0.1", 0);
var permits: query_budget.QueryBudget = try .init(gpa, 4);
defer permits.deinit();
var server = try tcp_server.TcpServer.listen(gpa, io, listen_address, &h, &permits, .{ .max_connections = 1 });
const server_address = server.boundAddress();
var group: std.Io.Group = .init;
try group.concurrent(io, tcp_server.TcpServer.serve, .{ &server, io });
defer {
server.deinit(io);
group.await(io) catch |err| switch (err) {
error.Canceled => unreachable,
};
}
try bounded(io, oneHeldOneRefused, .{ io, server_address });
try testing.expectEqual(@as(u64, 1), server.core.stats.rejected_at_capacity.load(.monotonic));
try testing.expectEqual(@as(u64, 1), server.core.stats.connections.load(.monotonic));
}
test "a connection that gets no task is closed and counted, not refused at capacity" {
if (!build_options.integration) return error.SkipZigTest;
const gpa = testing.allocator;
// The accept loop is the one task the server's pool allows. The client
// runs on a pool of its own, so `bounded` has the tasks it needs.
var server_threaded: std.Io.Threaded = .init(gpa, .{ .concurrent_limit = .limited(1) });
defer server_threaded.deinit();
const server_io = server_threaded.io();
var client_threaded: std.Io.Threaded = .init(gpa, .{});
defer client_threaded.deinit();
const io = client_threaded.io();
var fake: FakeUpstream = .{ .reply = response_bytes };
var h_owner: upstream_owner.Borrowed = .{};
var h = bareHandler(h_owner.client(fake.client()));
const listen_address: net.IpAddress = try .parse("127.0.0.1", 0);
var permits: query_budget.QueryBudget = try .init(gpa, 4);
defer permits.deinit();
var server = try tcp_server.TcpServer.listen(gpa, server_io, listen_address, &h, &permits, .{ .max_connections = 2 });
const server_address = server.boundAddress();
var group: std.Io.Group = .init;
try group.concurrent(server_io, tcp_server.TcpServer.serve, .{ &server, server_io });
defer {
server.deinit(server_io);
group.await(server_io) catch |err| switch (err) {
error.Canceled => unreachable,
};
}
const before = concurrency.snapshot().get(.connection);
try bounded(io, waitForServerClose, .{ io, server_address });
try testing.expectEqual(@as(u64, 1), concurrency.snapshot().get(.connection) - before);
try testing.expectEqual(@as(u64, 0), server.core.stats.rejected_at_capacity.load(.monotonic));
try testing.expectEqual(@as(u64, 0), server.core.stats.connections.load(.monotonic));
}