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

src/server/doh_server.zig

@@ -11,7 +11,7 @@
 //!
 //! The handshake runs under the same race budget tcp_server applies to its
 //! reads (ruling 3's rationale): a client that connects and never handshakes
-//! must not pin one of the 64 slots forever.
+//! must not pin a slot forever.
 //!
 //! Over capacity the raw TCP stream is closed and counted, with no response.
 //! The web listener's 503 is not possible here: HTTP exists only on the far
@@ -19,6 +19,11 @@
 //! certificate acquisition — exactly what an over-capacity listener does not
 //! have to spend. Closing is the cheapest honest refusal.
 //!
+//! A query holds a permit from the shared query budget from its decoded body
+//! to its sent response. With none free the connection is closed unanswered
+//! and counted, the TCP/53 rule: a 503 would be the one DNS refusal that
+//! differs by transport.
+//!
 //! Only `/dns-query` exists. This is not the admin API, so errors are plain
 //! text, not JSON, and there are no cache-control headers (LAN, no
 //! intermediaries).
@@ -34,6 +39,7 @@ const doh_client = @import("../upstream/doh_client.zig");
 const handler = @import("handler.zig");
 const listener = @import("listener.zig");
 const model = @import("../config/model.zig");
+const query_budget = @import("query_budget.zig");
 const tls_server = @import("../platform/tls_server.zig");
 const transport = @import("../upstream/transport.zig");
 const upstream_owner = @import("../upstream/owner.zig");
@@ -52,12 +58,10 @@ pub const alpn_protocols: [*:null]const ?[*:0]const u8 = &.{"http/1.1"};
 const recv_buffer_len = 8 * 1024;
 const send_buffer_len = 4 * 1024;
 
-pub const default_max_connections: u16 = 64;
-
 const allow_header: http.Header = .{ .name = "allow", .value = "GET, POST" };
 
 pub const Options = struct {
-    max_connections: u16 = default_max_connections,
+    max_connections: u16 = listener.dns_max_connections,
     /// The budget for the TLS handshake and for the wait on the next request
     /// head of a keep-alive connection (milestone-16 ruling 10). A request
     /// already being served has no timeout, like the web listener: the port is
@@ -87,27 +91,23 @@ pub const Snapshot = struct {
     idle_timeouts: u64,
     connection_errors: u64,
     bad_requests: u64,
+    dropped_no_permit: u64,
 };
 
 pub const DohServer = struct {
     core: listener.Core(Config),
     handler: *handler.Handler,
+    budget: *query_budget.QueryBudget,
     certs: *cert_store.CertStore,
     options: Options,
     stats: Stats,
 
-    /// One slot is ~150 KiB, so the default 64 connections cost ~9.4 MiB. The
-    /// two message buffers cannot shrink: a POST body and the reply both go up
-    /// to the 65535 bytes a DNS message can be. The `ServerStream` plaintext
-    /// buffers belong to the core; its `read_buf` doubles as the HTTP head cap
-    /// (see `recv_buffer_len`).
+    /// The `ServerStream` plaintext buffers belong to the core; its `read_buf`
+    /// doubles as the HTTP head cap (see `recv_buffer_len`).
     pub const Payload = struct {
-        /// The decoded query: a POST body or a GET `dns` parameter.
+        /// The decoded query: a POST body or a GET `dns` parameter, up to the
+        /// 65535 bytes a DNS message can be. Read before a permit is taken.
         query: [transport.max_message_len]u8,
-        reply: [transport.max_message_len]u8,
-        /// The handler's per-query working memory. A connection is answered
-        /// serially, so one query uses it at a time.
-        scratch: handler.Scratch,
         /// Valid between a successful `ServerStream.accept` and the
         /// `close(gpa)` in `serveOne`'s defer.
         tls: tls_server.ServerStream,
@@ -131,12 +131,14 @@ pub const DohServer = struct {
         io: std.Io,
         listen_address: net.IpAddress,
         h: *handler.Handler,
+        budget: *query_budget.QueryBudget,
         certs: *cert_store.CertStore,
         options: Options,
     ) ListenError!DohServer {
         return .{
             .core = try listener.Core(Config).listen(gpa, io, listen_address, options.max_connections),
             .handler = h,
+            .budget = budget,
             .certs = certs,
             .options = options,
             .stats = .{},
@@ -163,6 +165,7 @@ pub const DohServer = struct {
             .idle_timeouts = core.idle_timeouts.load(.monotonic),
             .connection_errors = core.connection_errors.load(.monotonic),
             .bad_requests = self.stats.bad_requests.load(.monotonic),
+            .dropped_no_permit = core.dropped_no_permit.load(.monotonic),
         };
     }
 
@@ -192,7 +195,7 @@ pub const DohServer = struct {
         const stage: Handshake = .{ .conn = conn, .gpa = gpa, .ctx = &entry.ctx, .io = io };
         switch (listener.handshakeStage(io, self.options.idle_timeout, stage)) {
             .ok => {},
-            .canceled => return,
+            .canceled, .unavailable => return,
             // Milestone-16 ruling 9: a stalled handshake is refused like a broken
             // one, the DoT arrangement. `idle_timeouts` belongs to the keep-alive
             // wait below, so the two listeners export the same names for the
@@ -220,9 +223,10 @@ pub const DohServer = struct {
                     listener.bump(&stats.idle_timeouts);
                     return;
                 },
-                // Cancellation is shutdown; `.failed` here is only the wrapper
-                // failing to start, which costs nothing and counts as nothing.
-                .canceled, .failed => return,
+                // Cancellation is shutdown, and `.unavailable` is counted by
+                // the race. `.failed` cannot happen: `receiveHeadInto` stores
+                // the head's error instead of returning it.
+                .canceled, .unavailable, .failed => return,
             }
 
             var request = head catch |err| switch (err) {
@@ -385,13 +389,19 @@ pub const DohServer = struct {
         request: *http.Server.Request,
         query: []const u8,
     ) error{ WriteFailed, HttpExpectationFailed }!Next {
+        const work = self.budget.tryAcquire(io) orelse {
+            listener.bump(&self.core.stats.dropped_no_permit);
+            return .close;
+        };
+        defer self.budget.release(io, work);
+
         const outcome = self.handler.handle(
             io,
             .tcp,
             address.NetAddress.fromIp(conn.peer),
             query,
-            &conn.payload.reply,
-            &conn.payload.scratch,
+            &work.reply,
+            &work.scratch,
         );
         switch (outcome) {
             .drop => return self.refuse(request, .bad_request, "bad request\n", &.{}, false),
@@ -626,6 +636,7 @@ const Harness = struct {
     upstream: FailingUpstream,
     upstream_owner: upstream_owner.Borrowed,
     h: handler.Handler,
+    permits: query_budget.QueryBudget,
     server: DohServer,
     group: std.Io.Group,
 
@@ -659,8 +670,11 @@ const Harness = struct {
             .local_tables = &hx.tables,
         };
 
+        hx.permits = try .init(testing.allocator, 4);
+        errdefer hx.permits.deinit();
+
         const listen_address: net.IpAddress = try .parse("127.0.0.1", 0);
-        hx.server = try DohServer.listen(testing.allocator, hio, listen_address, &hx.h, &hx.store, options);
+        hx.server = try DohServer.listen(testing.allocator, hio, listen_address, &hx.h, &hx.permits, &hx.store, options);
         errdefer hx.server.deinit(hio);
 
         hx.group = .init;
@@ -673,6 +687,7 @@ const Harness = struct {
         hx.group.await(hio) catch |err| switch (err) {
             error.Canceled => unreachable,
         };
+        hx.permits.deinit();
         hx.store.deinit(hio);
         hx.tables.deinit(testing.allocator);
         hx.tmp.cleanup();