milestone 31: concurrent upstream exchanges, dot session reuse, queue metrics

src/upstream/dot_client.zig

@@ -10,6 +10,16 @@
 //! local resource error or a cancellation must never reach the pool as a peer
 //! fault, so the concrete error is unwrapped from `error.ReadFailed` /
 //! `error.WriteFailed` before it is classified.
+//!
+//! A client keeps its connection open between exchanges (RFC 7858 §3.4). An
+//! upstream is free to drop an idle one at any time and says nothing first, so
+//! staleness is detected at use: an exchange that fails on a *reused* session,
+//! before any byte of its response arrived, with a connection-lifecycle cause,
+//! is retried exactly once on a fresh dial and counted through
+//! `reuse_recoveries` rather than against the upstream's health. Every other
+//! failure is final, and every final failure closes the session — after a
+//! failed send, receive or validation, the framing and TLS state are
+//! untrustworthy and the next exchange must start from a fresh dial.
 
 const std = @import("std");
 const net = std.Io.net;
@@ -41,6 +51,9 @@ const Diagnostic = struct {
         connect_failed: anyerror,
         handshake_failed: Handshake,
         bundle_load_failed: anyerror,
+        /// A reused session turned out to be dead. Normal operation, so this
+        /// one is only ever logged at debug; the counter is the real surface.
+        stale_session: anyerror,
 
         const Handshake = struct {
             verify_name: []const u8,
@@ -65,6 +78,10 @@ const Diagnostic = struct {
                 tls_client.classify(hs.cause),
             }),
             .bundle_load_failed => |err| try w.print("CA bundle load failed: {s}", .{@errorName(err)}),
+            .stale_session => |err| try w.print(
+                "reused session was stale ({s}), redialing once",
+                .{@errorName(err)},
+            ),
         }
     }
 };
@@ -93,6 +110,25 @@ pub const DotClient = struct {
     bundle_lock: *std.Io.RwLock,
     /// Caller-owned. One `DotClient` is used by one task at a time.
     buffers: Buffers,
+    /// Points at the owning pool entry's counter, which lives in the
+    /// composition root's stable storage; null in `nxdns check` probes, which
+    /// have no pool to report through. Incremented when a stale reused session
+    /// is recovered by a redial.
+    reuse_recoveries: ?*std.atomic.Value(u64),
+    /// The connection this client keeps open between exchanges, or null when it
+    /// holds none. Owned here: nothing else may close either half.
+    ///
+    /// Emplaced, never assigned from a local: `TlsStream` hands `tls.Client`
+    /// pointers into its own reader and writer fields, so a `Session` built on
+    /// the stack and copied in would leave the TLS client pointing at the dead
+    /// copy. `dial` sets this to `.{ .stream = …, .tls = undefined }` first and
+    /// runs the handshake through the stored payload.
+    session: ?Session = null,
+
+    const Session = struct {
+        stream: net.Stream,
+        tls: tls_client.TlsStream,
+    };
 
     pub const Buffers = struct {
         /// Plaintext read buffer.
@@ -116,6 +152,7 @@ pub const DotClient = struct {
         gpa: std.mem.Allocator,
         bundle: *Certificate.Bundle,
         bundle_lock: *std.Io.RwLock,
+        reuse_recoveries: ?*std.atomic.Value(u64),
         buffers: Buffers,
     ) DotClient {
         std.debug.assert(endpoint.scheme == .dot);
@@ -130,9 +167,25 @@ pub const DotClient = struct {
             .bundle = bundle,
             .bundle_lock = bundle_lock,
             .buffers = buffers,
+            .reuse_recoveries = reuse_recoveries,
         };
     }
 
+    /// Releases whatever this client holds open between exchanges: TLS first,
+    /// so the peer gets a close_notify, then the socket underneath it.
+    ///
+    /// Idempotent and safe with no session open, because that is how every
+    /// caller uses it — `Upstreams.deinit` closes clients that may never have
+    /// dialed, the `nxdns check` probe loop closes on its error paths, and
+    /// `exchange` closes after a failure it has already classified.
+    pub fn close(self: *DotClient, io: std.Io) void {
+        if (self.session) |*session| {
+            transport.closeBlocked(io, &session.tls);
+            transport.closeBlocked(io, &session.stream);
+            self.session = null;
+        }
+    }
+
     pub fn client(self: *DotClient) transport.Client {
         return .{ .ptr = self, .exchangeFn = exchangeFn };
     }
@@ -155,14 +208,15 @@ pub const DotClient = struct {
         return self.exchange(io, query, response_buf);
     }
 
-    /// One TCP connection and one TLS handshake per exchange, both closed
-    /// before returning.
+    /// One framed query and one framed reply over a session this client keeps
+    /// open, dialing one only when it holds none.
     ///
-    /// Connection reuse is deliberately not built. At household query rates the
-    /// saved round trips are worth less than what a per-exchange connection
-    /// buys: the pool's per-attempt budget stays a plain race against one task,
-    /// the failover path never has to reason about a half-dead pooled socket,
-    /// and every failure is attributable to exactly one exchange.
+    /// A reused session may already be dead: the upstream is entitled to close
+    /// an idle connection and RFC 7858 keepalive is advisory, so the first sign
+    /// is this exchange failing. That case — and only that case — is retried
+    /// once on a fresh dial; see `retryDecision` for the three conditions. The
+    /// caller's attempt budget covers the whole call including that redial,
+    /// because the pool races this function as a whole.
     pub fn exchange(
         self: *DotClient,
         io: std.Io,
@@ -174,6 +228,49 @@ pub const DotClient = struct {
         // local error rather than a silently truncated frame.
         if (query.len > transport.max_message_len) return error.BufferTooSmall;
 
+        const reused = self.session != null;
+        if (!reused) try self.dial(io);
+
+        const failure = switch (self.transact(query, response_buf)) {
+            .ok => |reply| return reply,
+            .failed => |failure| failure,
+        };
+
+        switch (retryDecision(reused, failure.received_any, failure.cause)) {
+            .final => {
+                self.close(io);
+                return transport.mapPhase(failure.cause, failure.phase);
+            },
+            .retry => {},
+        }
+
+        // Debug, not warn: an upstream dropping an idle connection is routine,
+        // and one line per idle timeout on a household resolver is log spam.
+        // `reuse_recoveries` is what makes the churn visible.
+        log.debug("{f}", .{self.diagnose(.{ .stale_session = failure.cause })});
+        self.close(io);
+        try self.dial(io);
+
+        switch (self.transact(query, response_buf)) {
+            .ok => |reply| {
+                if (self.reuse_recoveries) |counter| _ = counter.fetchAdd(1, .monotonic);
+                return reply;
+            },
+            // The retry's outcome is the exchange's outcome: one redial, never
+            // two.
+            .failed => |retried| {
+                self.close(io);
+                return transport.mapPhase(retried.cause, retried.phase);
+            },
+        }
+    }
+
+    /// Opens a session and leaves it in `self.session`, or leaves `self.session`
+    /// null and returns the classified failure. No partially initialized
+    /// session ever survives this call.
+    fn dial(self: *DotClient, io: std.Io) transport.ExchangeError!void {
+        std.debug.assert(self.session == null);
+
         const address = resolveAddress(self.endpoint) catch |err| {
             log.warn("{f}", .{self.diagnose(.not_an_ip_literal)});
             return err;
@@ -181,22 +278,24 @@ pub const DotClient = struct {
 
         try self.ensureBundle(io);
 
-        var stream = address.connect(io, .{ .mode = .stream }) catch |err| {
+        const stream = address.connect(io, .{ .mode = .stream }) catch |err| {
             log.debug("{f}", .{self.diagnose(.{ .connect_failed = err })});
             return transport.mapPhase(err, error.ConnectFailed);
         };
-        defer transport.closeBlocked(io, &stream);
 
-        // `TlsStream` is pinned: it holds its reader and writer by value and the
-        // TLS client points at them, so it must not move after `init`.
-        var tls_stream: tls_client.TlsStream = undefined;
+        // Emplaced before the handshake, never built beside it and copied in:
+        // `TlsStream` is pinned by the pointers `tls.Client` holds into its own
+        // reader and writer fields, so `init` has to run at the address the
+        // session will live at.
+        self.session = .{ .stream = stream, .tls = undefined };
+        const session = &self.session.?;
         // `concreteHandshake` reads these two fields when the handshake reports
         // `error.ReadFailed` / `error.WriteFailed`. `TlsStream.init` sets them
         // before it can produce either error, but clearing them here keeps that
         // out of this file's correctness argument.
-        tls_stream.stream_reader.err = null;
-        tls_stream.stream_writer.err = null;
-        tls_stream.init(io, &stream, self.bundle, self.bundle_lock, self.gpa, .{
+        session.tls.stream_reader.err = null;
+        session.tls.stream_writer.err = null;
+        session.tls.init(io, &session.stream, self.bundle, self.bundle_lock, self.gpa, .{
             .host = self.verify_name,
             .ca = .system,
             .read_buffer = self.buffers.tls_read,
@@ -204,36 +303,72 @@ pub const DotClient = struct {
             .stream_read_buffer = self.buffers.stream_read,
             .stream_write_buffer = self.buffers.stream_write,
         }) catch |err| {
-            const cause = concreteHandshake(&tls_stream, err);
+            // Order matters: the concrete cause lives in the in-place
+            // `TlsStream`'s two error fields, so clearing the optional first
+            // would destroy the only record that a cancellation or a local
+            // resource failure — not the peer — ended the handshake.
+            const cause = concreteHandshake(&session.tls, err);
+            transport.closeBlocked(io, &session.stream);
+            self.session = null;
             log.warn("{f}", .{self.diagnose(.{ .handshake_failed = .{
                 .verify_name = self.verify_name,
                 .cause = cause,
             } })});
             return transport.mapPhase(cause, error.TlsFailed);
         };
-        defer transport.closeBlocked(io, &tls_stream);
+    }
+
+    /// One query and one reply on the open session, with enough detail on
+    /// failure for `retryDecision` to rule on it. Leaves the session open
+    /// either way; closing a failed one is `exchange`'s job, because only it
+    /// knows whether the failure is final.
+    fn transact(self: *DotClient, query: []const u8, response_buf: []u8) Transact {
+        const session = &self.session.?;
+        const tls_stream = &session.tls;
 
         const writer = tls_stream.writer();
         const prefix = transport.framePrefix(@intCast(query.len));
-        writer.writeAll(&prefix) catch |err| return sendFailure(&tls_stream, err);
-        writer.writeAll(query) catch |err| return sendFailure(&tls_stream, err);
+        writer.writeAll(&prefix) catch |err| return sendFailed(tls_stream, err);
+        writer.writeAll(query) catch |err| return sendFailed(tls_stream, err);
         // `TlsStream.flush`, not `writer.flush`: the latter leaves the encrypted
         // record in the socket writer's buffer and the query never leaves this
         // process.
-        tls_stream.flush() catch |err| return sendFailure(&tls_stream, err);
+        tls_stream.flush() catch |err| return sendFailed(tls_stream, err);
 
         const reader = tls_stream.reader();
+        // Byte at a time, not `readSliceAll`: that call is all-or-nothing, so a
+        // failure after the first prefix byte would be indistinguishable from
+        // one before it, and "no response byte received yet" is one of the three
+        // conditions a retry needs.
         var prefix_bytes: [transport.prefix_len]u8 = undefined;
-        reader.readSliceAll(&prefix_bytes) catch |err| return receiveFailure(&tls_stream, err);
+        for (&prefix_bytes, 0..) |*byte, received| {
+            byte.* = reader.takeByte() catch |err|
+                return receiveFailed(tls_stream, err, received != 0);
+        }
 
         const len = transport.parsePrefix(prefix_bytes);
-        if (len == 0) return error.BadResponse;
-        if (len > response_buf.len) return error.ResponseTooLarge;
+        if (len == 0) return .{ .failed = .{
+            .cause = error.BadResponse,
+            .phase = error.BadResponse,
+            .received_any = true,
+        } };
+        if (len > response_buf.len) return .{ .failed = .{
+            .cause = error.ResponseTooLarge,
+            .phase = error.ResponseTooLarge,
+            .received_any = true,
+        } };
         reader.readSliceAll(response_buf[0..len]) catch |err|
-            return receiveFailure(&tls_stream, err);
+            return receiveFailed(tls_stream, err, true);
 
-        try transport.validateResponse(query, response_buf[0..len]);
-        return response_buf[0..len];
+        transport.validateResponse(query, response_buf[0..len]) catch |err| return .{ .failed = .{
+            .cause = err,
+            .phase = switch (err) {
+                error.BadResponse => error.BadResponse,
+                error.ResponseMismatch => error.ResponseMismatch,
+            },
+            .received_any = true,
+        } };
+        return .{ .ok = response_buf[0..len] };
     }
 
     /// Loads the system CA bundle before the handshake, so that a failure to
@@ -301,14 +436,72 @@ fn concreteWrite(stream: *tls_client.TlsStream, err: anyerror) anyerror {
     return err;
 }
 
-fn sendFailure(stream: *tls_client.TlsStream, err: anyerror) transport.ExchangeError {
-    return transport.mapPhase(concreteWrite(stream, err), error.SendFailed);
+/// The outcome of one `transact`, as a value: a reply, or everything
+/// `retryDecision` needs to rule on the failure.
+const Transact = union(enum) {
+    ok: []u8,
+    failed: Failure,
+
+    const Failure = struct {
+        /// Already unwrapped out of `error.ReadFailed` / `error.WriteFailed`.
+        cause: anyerror,
+        /// The peer fault this phase means, unless `cause` belongs to this
+        /// process.
+        phase: transport.PeerFault,
+        /// Whether any byte of this exchange's response had arrived.
+        received_any: bool,
+    };
+};
+
+fn sendFailed(stream: *tls_client.TlsStream, err: anyerror) Transact {
+    return .{
+        .failed = .{
+            .cause = concreteWrite(stream, err),
+            .phase = error.SendFailed,
+            // Nothing is read before the query is out, so a send failure is always
+            // pre-first-byte.
+            .received_any = false,
+        },
+    };
 }
 
-fn receiveFailure(stream: *tls_client.TlsStream, err: anyerror) transport.ExchangeError {
-    return transport.mapPhase(concreteRead(stream, err), error.ReceiveFailed);
+fn receiveFailed(stream: *tls_client.TlsStream, err: anyerror, received_any: bool) Transact {
+    return .{ .failed = .{
+        .cause = concreteRead(stream, err),
+        .phase = error.ReceiveFailed,
+        .received_any = received_any,
+    } };
 }
 
+/// Whether a failed exchange may be retried once on a fresh dial.
+///
+/// Pure, and separate from the wire for that reason: the three conditions are
+/// the whole of the reuse contract, and a table test is the only way to see all
+/// of them at once.
+///
+/// `reused` — a session this call dialed itself was never idle, so its failure
+/// is the upstream's answer, not a stale connection. `bytes_received` — once a
+/// reply has started, re-sending the query would be a second question, and the
+/// failure is the upstream's. `cause` — only the ways a connection ends
+/// (`error.EndOfStream` is a clean close_notify, `error.TlsConnectionTruncated`
+/// a close without one, `error.ConnectionResetByPeer` and `error.BrokenPipe`
+/// the socket-level pair). A TLS alert, a certificate fault, a local resource
+/// failure and a cancellation all keep their meaning and are final.
+fn retryDecision(reused: bool, bytes_received: bool, cause: anyerror) RetryDecision {
+    if (!reused) return .final;
+    if (bytes_received) return .final;
+    return switch (cause) {
+        error.EndOfStream,
+        error.TlsConnectionTruncated,
+        error.ConnectionResetByPeer,
+        error.BrokenPipe,
+        => .retry,
+        else => .final,
+    };
+}
+
+const RetryDecision = enum { retry, final };
+
 const testing = std.testing;
 
 fn expectDiagnostic(expected: []const u8, url: []const u8, detail: Diagnostic.Detail) !void {
@@ -343,6 +536,11 @@ test "every diagnostic line redacts the url it names the upstream by" {
         "tls://dns.example/abcd12?apikey=s3cr3t",
         .{ .bundle_load_failed = error.FileNotFound },
     );
+    try expectDiagnostic(
+        "dot upstream 'tls://dns.example:853': reused session was stale (EndOfStream), redialing once",
+        "tls://dns.example:853/",
+        .{ .stale_session = error.EndOfStream },
+    );
     try expectDiagnostic(
         "dot upstream 'tls://dns.example': TLS handshake as 'one.one.one.one' failed: " ++
             "CertificateHostMismatch (certificate)",
@@ -497,27 +695,126 @@ test "the handshake unwrap passes other errors through untouched" {
     );
 }
 
+/// What `exchange` would return for a failure `transact` reported.
+fn mappedFailure(outcome: Transact) transport.ExchangeError {
+    return transport.mapPhase(outcome.failed.cause, outcome.failed.phase);
+}
+
 test "the send and receive unwraps prefer the stored cause" {
     var send = stubStream(null, error.Canceled, null);
     try testing.expectEqual(
         transport.ExchangeError.Canceled,
-        sendFailure(&send, error.WriteFailed),
+        mappedFailure(sendFailed(&send, error.WriteFailed)),
     );
 
     // The TLS client's own error wins over the socket reader's.
     var receive = stubStream(error.ConnectionResetByPeer, null, error.TlsAlert);
     try testing.expectEqual(
         transport.ExchangeError.ReceiveFailed,
-        receiveFailure(&receive, error.ReadFailed),
+        mappedFailure(receiveFailed(&receive, error.ReadFailed, false)),
     );
 
     var socket = stubStream(error.SystemResources, null, null);
     try testing.expectEqual(
         transport.ExchangeError.SystemResources,
-        receiveFailure(&socket, error.ReadFailed),
+        mappedFailure(receiveFailed(&socket, error.ReadFailed, true)),
     );
 }
 
+test "a send failure is always pre-first-byte, and a receive failure reports what it read" {
+    // The retry rule reads `received_any`, so where it comes from is part of the
+    // contract rather than an incidental field: nothing is read before the query
+    // is on the wire, and the receive side is told by its caller.
+    var stream = stubStream(error.ConnectionResetByPeer, error.ConnectionResetByPeer, null);
+    try testing.expect(!sendFailed(&stream, error.WriteFailed).failed.received_any);
+    try testing.expect(!receiveFailed(&stream, error.ReadFailed, false).failed.received_any);
+    try testing.expect(receiveFailed(&stream, error.ReadFailed, true).failed.received_any);
+}
+
+test "a fresh session is never retried, whatever failed" {
+    // The redial has nothing to fix: this call dialed the connection itself, so
+    // the failure is the upstream's answer rather than a stale socket.
+    for ([_]anyerror{
+        error.EndOfStream,
+        error.TlsConnectionTruncated,
+        error.ConnectionResetByPeer,
+        error.BrokenPipe,
+        error.TlsAlert,
+        error.Canceled,
+    }) |cause| {
+        try testing.expectEqual(RetryDecision.final, retryDecision(false, false, cause));
+        try testing.expectEqual(RetryDecision.final, retryDecision(false, true, cause));
+    }
+}
+
+test "a reused session is retried once for the ways a connection ends" {
+    for ([_]anyerror{
+        error.EndOfStream,
+        error.TlsConnectionTruncated,
+        error.ConnectionResetByPeer,
+        error.BrokenPipe,
+    }) |cause| {
+        try testing.expectEqual(RetryDecision.retry, retryDecision(true, false, cause));
+    }
+}
+
+test "a reused session that already answered is never retried" {
+    // Re-sending the query after a byte of the reply arrived would ask the
+    // upstream a second question, so a mid-reply failure is final however the
+    // connection died.
+    for ([_]anyerror{
+        error.EndOfStream,
+        error.TlsConnectionTruncated,
+        error.ConnectionResetByPeer,
+        error.BrokenPipe,
+    }) |cause| {
+        try testing.expectEqual(RetryDecision.final, retryDecision(true, true, cause));
+    }
+}
+
+test "a reused session is not retried for a fault that says something" {
+    // A TLS alert, a certificate fault, a local resource failure, a
+    // cancellation and a bad frame all keep their meaning: none of them is an
+    // idle connection going away, so none is worth a second dial.
+    for ([_]anyerror{
+        error.TlsAlert,
+        error.TlsBadRecordMac,
+        error.CertificateExpired,
+        error.SystemResources,
+        error.OutOfMemory,
+        error.Canceled,
+        error.BadResponse,
+        error.ResponseMismatch,
+        error.ResponseTooLarge,
+        error.ConnectionRefused,
+        error.Timeout,
+    }) |cause| {
+        try testing.expectEqual(RetryDecision.final, retryDecision(true, false, cause));
+    }
+}
+
+test "a validation failure is final and its phase survives the mapping" {
+    // `transact` reports these with `received_any` set, so the decision is final
+    // by two of the three conditions at once, and the peer fault the pool
+    // records is the validation error itself rather than a receive failure.
+    const outcomes = [_]struct { cause: anyerror, phase: transport.PeerFault }{
+        .{ .cause = error.BadResponse, .phase = error.BadResponse },
+        .{ .cause = error.ResponseMismatch, .phase = error.ResponseMismatch },
+        .{ .cause = error.ResponseTooLarge, .phase = error.ResponseTooLarge },
+    };
+    for (outcomes) |outcome| {
+        try testing.expectEqual(RetryDecision.final, retryDecision(true, true, outcome.cause));
+        try testing.expectEqual(
+            @as(transport.ExchangeError, outcome.phase),
+            mappedFailure(.{ .failed = .{
+                .cause = outcome.cause,
+                .phase = outcome.phase,
+                .received_any = true,
+            } }),
+        );
+    }
+}
+
 test "a CA bundle scan failure keeps local resource errors out of the peer fault group" {
     // `Certificate.Bundle.rescan` reaches these through `Allocator.Error`,
     // `Io.File.OpenError` and `Io.UnexpectedError`.
@@ -565,7 +862,7 @@ test "DotClient satisfies the Client interface" {
 
     // `init` asserts `endpoint.scheme == .dot`; a `.doh` endpoint trips
     // `std.debug.assert`, which a test cannot catch in-process.
-    var dot: DotClient = .init(try .parse("tls://9.9.9.9:853"), "", gpa, &bundle, &bundle_lock, .{
+    var dot: DotClient = .init(try .parse("tls://9.9.9.9:853"), "", gpa, &bundle, &bundle_lock, null, .{
         .tls_read = buffer[0..chunk],
         .tls_write = buffer[chunk .. 2 * chunk],
         .stream_read = buffer[2 * chunk .. 3 * chunk],
@@ -574,6 +871,12 @@ test "DotClient satisfies the Client interface" {
 
     try testing.expectEqual(transport.Scheme.dot, dot.endpoint.scheme);
     try testing.expectEqualStrings("9.9.9.9", dot.endpoint.host);
+    // A client is wired with nothing open, so `init` can return by value and
+    // the composition root can copy the result into its array: the pinned TLS
+    // state only exists once `exchange` has dialed.
+    try testing.expect(dot.session == null);
+    dot.close(undefined);
+    try testing.expect(dot.session == null);
 
     const iface: transport.Client = dot.client();
     try testing.expectEqual(@as(*anyopaque, @ptrCast(&dot)), iface.ptr);
@@ -598,7 +901,7 @@ test "a tls_name replaces the verification name and leaves the dial target alone
     };
 
     const endpoint: transport.Endpoint = try .parse("tls://1.1.1.1:853");
-    const named: DotClient = .init(endpoint, "one.one.one.one", gpa, &bundle, &bundle_lock, buffers);
+    const named: DotClient = .init(endpoint, "one.one.one.one", gpa, &bundle, &bundle_lock, null, buffers);
     try testing.expectEqualStrings("one.one.one.one", named.verify_name);
     try testing.expectEqualStrings("1.1.1.1", named.endpoint.host);
 
@@ -606,6 +909,6 @@ test "a tls_name replaces the verification name and leaves the dial target alone
     try testing.expectEqualSlices(u8, &.{ 1, 1, 1, 1 }, &address.ip4.bytes);
     try testing.expectEqual(@as(u16, 853), address.ip4.port);
 
-    const plain: DotClient = .init(endpoint, "", gpa, &bundle, &bundle_lock, buffers);
+    const plain: DotClient = .init(endpoint, "", gpa, &bundle, &bundle_lock, null, buffers);
     try testing.expectEqualStrings("1.1.1.1", plain.verify_name);
 }