Skip to main content

dynamo_runtime/pipeline/network/
quic_response.rs

1// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2// SPDX-License-Identifier: Apache-2.0
3
4//! Fixed-lane QUIC transport for streaming worker responses.
5//!
6//! A worker keeps a fixed connection bundle to a frontend and distributes a
7//! fixed set of long-lived bidirectional streams (lanes) across it. Logical
8//! response frames are request-hashed to a lane once, queued in a bounded
9//! Tokio channel, and written in batches by the lane's sole writer task.
10
11use std::{
12    collections::{HashMap, VecDeque},
13    fmt,
14    net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr, UdpSocket},
15    sync::{
16        Arc, OnceLock,
17        atomic::{AtomicBool, AtomicUsize, Ordering},
18    },
19    time::Duration,
20};
21
22use anyhow::{Context as _, Result, anyhow, bail};
23use arc_swap::ArcSwapOption;
24use bytes::{BufMut, Bytes, BytesMut};
25use crossbeam_queue::SegQueue;
26use dashmap::DashMap;
27use parking_lot::Mutex;
28use prometheus::IntCounter;
29use quinn::crypto::rustls::{QuicClientConfig, QuicServerConfig};
30use rustls::{
31    DigitallySignedStruct, SignatureScheme,
32    client::danger::{HandshakeSignatureValid, ServerCertVerified, ServerCertVerifier},
33    pki_types::{CertificateDer, PrivatePkcs8KeyDer, ServerName, UnixTime},
34};
35use serde::{Deserialize, Serialize};
36use sha2::{Digest, Sha256};
37use socket2::{Domain, Protocol, Socket, Type};
38use tokio::{
39    io::{AsyncRead, AsyncReadExt, BufReader},
40    sync::{Notify, Semaphore, mpsc, oneshot},
41    time::{Instant, sleep_until},
42};
43use tokio_util::sync::CancellationToken;
44use uuid::Uuid;
45use xxhash_rust::xxh3::xxh3_64;
46
47use super::{
48    ConnectionInfo, RegisteredStream, ResponseStreamPrologue, StreamPrologueError, StreamReceiver,
49};
50use crate::{
51    config::environment_names::quic_response, discovery::EndpointInstanceId,
52    engine::AsyncEngineContext, pipeline::PipelineError,
53};
54
55pub const TRANSPORT_NAME: &str = "quic-response";
56const PROTOCOL_VERSION: u8 = 3;
57const ALPN: &[u8] = b"dynamo-response-v2";
58const BULK_CONNECTIONS: usize = 8;
59const BULK_LANES: usize = 8;
60const PRIORITY_CONNECTIONS: usize = 1;
61const LANE_QUEUE_CAPACITY: usize = 512;
62// Bulk writers can spend tens of milliseconds waiting for Quinn to accept a
63// large batch. Keep enough bounded slack to absorb those stalls without
64// parking every response task on the lane semaphore. Priority traffic stays
65// on the smaller queue because it never carries the long token stream.
66const BULK_LANE_QUEUE_CAPACITY: usize = 4_096;
67const MAX_BATCH_FRAMES: usize = 512;
68const DEFAULT_BATCH_INTERVAL_US: u64 = 5_000;
69const MAX_BATCH_INTERVAL_US: u64 = 100_000;
70const FRAME_HEADER_LEN: usize = 1 + 16 + 4;
71// A writer batch can contain hundreds of small response frames. Buffer reads
72// at the receiver so parsing those frames does not poll Quinn once per header
73// and payload and exhaust Tokio's cooperative task budget.
74const RECEIVE_BUFFER_CAPACITY: usize = 256 * 1024;
75const READ_BUDGET_FRAMES: usize = 16;
76const READ_BUDGET_BYTES: usize = 4 * 1024;
77// A single frontend UDP socket is limited by the host receive-buffer ceiling.
78// Reuse-port endpoints preserve one advertised address while spreading QUIC
79// connections and receive queues across several sockets.
80#[cfg(target_os = "linux")]
81const SERVER_ENDPOINTS: usize = 32;
82#[cfg(not(target_os = "linux"))]
83const SERVER_ENDPOINTS: usize = 1;
84const MAX_FRAME_PAYLOAD: usize = 32 * 1024 * 1024;
85const RESPONSE_BUFFER_CAPACITY: usize = 16_384;
86const MAX_RESPONSE_BUFFER_CAPACITY: usize = 65_536;
87const MAX_DEFERRED_RESPONSE_FRAMES: usize = 1_024;
88const MAX_DEFERRED_RESPONSE_BYTES: usize = 16 * 1024 * 1024;
89const TOMBSTONE_TTL: Duration = Duration::from_secs(5);
90const CLOSE_CODE_INVARIANT: quinn::VarInt = quinn::VarInt::from_u32(1);
91
92#[derive(Debug, Clone, Copy)]
93pub struct QuicResponseConfig {
94    pub batch_interval: Duration,
95    pub response_buffer_capacity: usize,
96}
97
98impl QuicResponseConfig {
99    pub fn from_env() -> Result<Self, PipelineError> {
100        let interval_us = parse_env_range(
101            quic_response::DYN_QUIC_RESPONSE_BATCH_INTERVAL_US,
102            DEFAULT_BATCH_INTERVAL_US,
103            0,
104            MAX_BATCH_INTERVAL_US,
105        )?;
106        let response_buffer_capacity = parse_env_range(
107            quic_response::DYN_QUIC_RESPONSE_BUFFER_CAPACITY,
108            RESPONSE_BUFFER_CAPACITY,
109            1,
110            MAX_RESPONSE_BUFFER_CAPACITY,
111        )?;
112        Ok(Self {
113            batch_interval: Duration::from_micros(interval_us),
114            response_buffer_capacity,
115        })
116    }
117}
118
119fn parse_env_range<T>(name: &str, default: T, min: T, max: T) -> Result<T, PipelineError>
120where
121    T: Copy + std::str::FromStr + PartialOrd + fmt::Display,
122{
123    let Some(raw) = std::env::var(name).ok().filter(|value| !value.is_empty()) else {
124        return Ok(default);
125    };
126    let value = raw
127        .parse::<T>()
128        .map_err(|_| PipelineError::Generic(format!("invalid {name}: '{raw}' is not a number")))?;
129    if value < min || value > max {
130        return Err(PipelineError::Generic(format!(
131            "invalid {name}: {value} is outside {min}..={max}"
132        )));
133    }
134    Ok(value)
135}
136
137#[derive(Debug, Clone, Serialize, Deserialize)]
138pub struct QuicResponseConnectionInfo {
139    version: u8,
140    address: String,
141    frontend_id: String,
142    registration_id: Uuid,
143    request_id: String,
144    certificate_sha256: String,
145}
146
147impl From<QuicResponseConnectionInfo> for ConnectionInfo {
148    fn from(value: QuicResponseConnectionInfo) -> Self {
149        Self {
150            transport: TRANSPORT_NAME.to_string(),
151            info: serde_json::to_string(&value).expect("QUIC connection info must serialize"),
152        }
153    }
154}
155
156impl TryFrom<ConnectionInfo> for QuicResponseConnectionInfo {
157    type Error = anyhow::Error;
158
159    fn try_from(value: ConnectionInfo) -> Result<Self> {
160        if value.transport != TRANSPORT_NAME {
161            bail!(
162                "expected {TRANSPORT_NAME} connection info, got {}",
163                value.transport
164            );
165        }
166        let info: Self = serde_json::from_str(&value.info)?;
167        if info.version != PROTOCOL_VERSION {
168            bail!(
169                "unsupported QUIC response protocol version {} (expected {})",
170                info.version,
171                PROTOCOL_VERSION
172            );
173        }
174        Ok(info)
175    }
176}
177
178#[derive(Debug, Clone, Copy, PartialEq, Eq)]
179#[repr(u8)]
180enum FrameKind {
181    Prologue = 1,
182    Data = 2,
183    /// Terminal pre-stream failure. Payload layout is [`QUIC_TYPED_ERROR_KIND`].
184    Error = 3,
185    End = 4,
186    Stop = 5,
187    Kill = 6,
188    Reset = 7,
189    FirstData = 8,
190    PriorityEnd = 9,
191    Bundle = 10,
192    Register = 11,
193    Registered = 12,
194}
195
196impl TryFrom<u8> for FrameKind {
197    type Error = anyhow::Error;
198
199    fn try_from(value: u8) -> Result<Self> {
200        match value {
201            1 => Ok(Self::Prologue),
202            2 => Ok(Self::Data),
203            3 => Ok(Self::Error),
204            4 => Ok(Self::End),
205            5 => Ok(Self::Stop),
206            6 => Ok(Self::Kill),
207            7 => Ok(Self::Reset),
208            8 => Ok(Self::FirstData),
209            9 => Ok(Self::PriorityEnd),
210            10 => Ok(Self::Bundle),
211            11 => Ok(Self::Register),
212            12 => Ok(Self::Registered),
213            _ => bail!("unknown QUIC response frame kind {value}"),
214        }
215    }
216}
217
218#[derive(Debug)]
219struct Frame {
220    kind: FrameKind,
221    registration_id: Uuid,
222    payload: Bytes,
223}
224
225impl Frame {
226    fn new(kind: FrameKind, registration_id: Uuid, payload: Bytes) -> Self {
227        Self {
228            kind,
229            registration_id,
230            payload,
231        }
232    }
233
234    fn header(&self) -> Bytes {
235        let mut header = BytesMut::with_capacity(FRAME_HEADER_LEN);
236        header.put_u8(self.kind as u8);
237        header.extend_from_slice(self.registration_id.as_bytes());
238        header.put_u32(self.payload.len() as u32);
239        header.freeze()
240    }
241}
242
243/// Typed [`FrameKind::Error`] payloads are JSON tagged with this `kind`.
244///
245/// Untyped errors stay raw UTF-8 so a mixed-version frontend still sees the
246/// worker's text. A current frontend accepts only this discriminant and
247/// recovers [`ResponseStreamPrologue`]. An older frontend reads the payload as
248/// UTF-8 and treats the object as the message. Untagged JSON, including a
249/// decoy `{"error","code"}` object, stays that message. The protocol version
250/// stays 3 because those bytes are already valid UTF-8 for an old decoder.
251const QUIC_TYPED_ERROR_KIND: &str = "dynamo.quic.typed_error";
252
253/// JSON object carrying [`QUIC_TYPED_ERROR_KIND`] and a [`ResponseStreamPrologue`].
254#[derive(Serialize, Deserialize)]
255struct QuicTypedErrorWire {
256    kind: String,
257    #[serde(flatten)]
258    prologue: ResponseStreamPrologue,
259}
260
261/// Encode a pre-stream error for [`FrameKind::Error`].
262fn encode_error_payload(error: StreamPrologueError) -> Bytes {
263    let Some(typed_error) = error.typed_error else {
264        return Bytes::from(error.message);
265    };
266    let wire = QuicTypedErrorWire {
267        kind: QUIC_TYPED_ERROR_KIND.to_string(),
268        prologue: ResponseStreamPrologue {
269            error: Some(error.message),
270            typed_error: Some(typed_error),
271        },
272    };
273    match serde_json::to_vec(&wire) {
274        Ok(bytes) => Bytes::from(bytes),
275        Err(err) => {
276            tracing::error!(
277                error = %err,
278                "QUIC error prologue did not serialize; sending untyped text"
279            );
280            Bytes::from(wire.prologue.error.unwrap_or_default())
281        }
282    }
283}
284
285/// Decode a [`FrameKind::Error`] payload.
286fn decode_error_payload(payload: &[u8]) -> Result<StreamPrologueError> {
287    if payload.first() == Some(&b'{')
288        && let Ok(wire) = serde_json::from_slice::<QuicTypedErrorWire>(payload)
289        && wire.kind == QUIC_TYPED_ERROR_KIND
290        && let Some(message) = wire.prologue.error
291    {
292        return Ok(StreamPrologueError {
293            message,
294            typed_error: wire.prologue.typed_error,
295        });
296    }
297    let message = String::from_utf8(payload.to_vec())
298        .context("QUIC response terminal error was not UTF-8")?;
299    Ok(StreamPrologueError::from_message(message))
300}
301
302async fn read_frame<R>(recv: &mut R) -> Result<Frame>
303where
304    R: AsyncRead + Unpin,
305{
306    let mut header = [0_u8; FRAME_HEADER_LEN];
307    recv.read_exact(&mut header).await?;
308    let kind = FrameKind::try_from(header[0])?;
309    let registration_id = Uuid::from_slice(&header[1..17])?;
310    let payload_len = u32::from_be_bytes(header[17..21].try_into().unwrap()) as usize;
311    if payload_len > MAX_FRAME_PAYLOAD {
312        bail!("QUIC response frame payload {payload_len} exceeds {MAX_FRAME_PAYLOAD}");
313    }
314    let mut payload = BytesMut::zeroed(payload_len);
315    recv.read_exact(&mut payload).await?;
316    Ok(Frame {
317        kind,
318        registration_id,
319        payload: payload.freeze(),
320    })
321}
322
323struct LaneSender(Arc<LaneQueue>);
324
325struct LaneReceiver(Arc<LaneQueue>);
326
327struct LaneQueue {
328    frames: SegQueue<Frame>,
329    slots: Semaphore,
330    ready: Notify,
331    scheduled: AtomicBool,
332    closed: AtomicBool,
333    senders: AtomicUsize,
334    #[cfg(test)]
335    notifications: AtomicUsize,
336}
337
338#[derive(Debug)]
339enum LaneTrySendError {
340    Closed(Frame),
341    Full(Frame),
342}
343
344impl LaneQueue {
345    fn new(capacity: usize) -> Arc<Self> {
346        Arc::new(Self {
347            frames: SegQueue::new(),
348            slots: Semaphore::new(capacity),
349            ready: Notify::new(),
350            scheduled: AtomicBool::new(false),
351            closed: AtomicBool::new(false),
352            senders: AtomicUsize::new(1),
353            #[cfg(test)]
354            notifications: AtomicUsize::new(0),
355        })
356    }
357
358    fn push_with_permit(
359        &self,
360        frame: Frame,
361        permit: tokio::sync::SemaphorePermit<'_>,
362    ) -> std::result::Result<(), Frame> {
363        if self.closed.load(Ordering::Acquire) {
364            return Err(frame);
365        }
366        self.frames.push(frame);
367        permit.forget();
368        if !self.scheduled.swap(true, Ordering::AcqRel) {
369            #[cfg(test)]
370            self.notifications.fetch_add(1, Ordering::Relaxed);
371            self.ready.notify_one();
372        }
373        Ok(())
374    }
375
376    fn try_send(&self, frame: Frame) -> std::result::Result<(), LaneTrySendError> {
377        let permit = match self.slots.try_acquire() {
378            Ok(permit) => permit,
379            Err(tokio::sync::TryAcquireError::Closed) => {
380                return Err(LaneTrySendError::Closed(frame));
381            }
382            Err(tokio::sync::TryAcquireError::NoPermits) => {
383                return Err(LaneTrySendError::Full(frame));
384            }
385        };
386        self.push_with_permit(frame, permit)
387            .map_err(LaneTrySendError::Closed)
388    }
389
390    async fn send(&self, frame: Frame) -> Result<()> {
391        let permit = self
392            .slots
393            .acquire()
394            .await
395            .map_err(|_| anyhow!("QUIC response lane closed"))?;
396        self.push_with_permit(frame, permit)
397            .map_err(|_| anyhow!("QUIC response lane closed"))
398    }
399
400    fn drain(&self, batch: &mut Vec<Frame>, limit: usize) -> usize {
401        let initial_len = batch.len();
402        while batch.len() - initial_len < limit {
403            let Some(frame) = self.frames.pop() else {
404                break;
405            };
406            batch.push(frame);
407        }
408        let count = batch.len() - initial_len;
409        if count != 0 {
410            self.slots.add_permits(count);
411        }
412        count
413    }
414
415    async fn recv_many(&self, batch: &mut Vec<Frame>, limit: usize) -> usize {
416        loop {
417            let mut notified = Box::pin(self.ready.notified());
418            notified.as_mut().enable();
419            let count = self.drain(batch, limit);
420            if count != 0 {
421                return count;
422            }
423            if self.closed.load(Ordering::Acquire) {
424                return 0;
425            }
426            // Producers suppress notifications while the writer is active. Mark
427            // it idle after observing an empty queue, then recheck to close the
428            // producer/consumer race without a lock.
429            self.scheduled.store(false, Ordering::Release);
430            if !self.frames.is_empty() {
431                self.scheduled.store(true, Ordering::Release);
432                continue;
433            }
434            notified.await;
435        }
436    }
437
438    fn close_sender(&self) {
439        self.closed.store(true, Ordering::Release);
440        self.slots.close();
441        self.ready.notify_waiters();
442    }
443
444    fn close_receiver(&self) {
445        self.closed.store(true, Ordering::Release);
446        let mut discarded = 0;
447        while self.frames.pop().is_some() {
448            discarded += 1;
449        }
450        if discarded != 0 {
451            self.slots.add_permits(discarded);
452        }
453        self.slots.close();
454        self.ready.notify_waiters();
455    }
456}
457
458impl Drop for LaneReceiver {
459    fn drop(&mut self) {
460        self.0.close_receiver();
461    }
462}
463
464impl Clone for LaneSender {
465    fn clone(&self) -> Self {
466        self.0.senders.fetch_add(1, Ordering::Relaxed);
467        Self(self.0.clone())
468    }
469}
470
471impl Drop for LaneSender {
472    fn drop(&mut self) {
473        if self.0.senders.fetch_sub(1, Ordering::AcqRel) == 1 {
474            self.0.close_sender();
475        }
476    }
477}
478
479impl LaneReceiver {
480    async fn recv_many(&mut self, batch: &mut Vec<Frame>, limit: usize) -> usize {
481        self.0.recv_many(batch, limit).await
482    }
483
484    #[cfg(test)]
485    async fn recv(&mut self) -> Option<Frame> {
486        let mut batch = Vec::with_capacity(1);
487        (self.recv_many(&mut batch, 1).await == 1).then(|| batch.pop().unwrap())
488    }
489}
490
491impl LaneSender {
492    fn try_send(&self, frame: Frame) -> std::result::Result<(), LaneTrySendError> {
493        self.0.try_send(frame)
494    }
495
496    async fn send(&self, frame: Frame) -> Result<()> {
497        self.0.send(frame).await
498    }
499}
500
501fn lane_queue(capacity: usize) -> (LaneSender, LaneReceiver) {
502    let queue = LaneQueue::new(capacity);
503    (LaneSender(queue.clone()), LaneReceiver(queue))
504}
505
506/// Fill `batch` with up to `max_batch_frames` frames. The first `recv_many`
507/// waits while idle;
508/// subsequent cancellation-safe `recv_many` calls race the single batch
509/// deadline. A zero interval therefore drains an already-ready burst but does
510/// not intentionally wait for more work.
511async fn receive_batch(
512    receiver: &mut LaneReceiver,
513    batch: &mut Vec<Frame>,
514    interval: Duration,
515    max_batch_frames: usize,
516) -> bool {
517    batch.clear();
518    if receiver.recv_many(batch, max_batch_frames).await == 0 {
519        return false;
520    }
521
522    if interval.is_zero() {
523        return true;
524    }
525    let deadline = Instant::now() + interval;
526    let sleep = sleep_until(deadline);
527    tokio::pin!(sleep);
528    while batch.len() < max_batch_frames {
529        tokio::select! {
530            biased;
531            _ = &mut sleep => break,
532            count = receiver.recv_many(batch, max_batch_frames - batch.len()) => {
533                if count == 0 {
534                    break;
535                }
536            }
537        }
538    }
539    true
540}
541
542const SERVER_STATE_SHARDS: usize = 64;
543
544#[derive(Default)]
545struct RegistrationShard {
546    pending: HashMap<Uuid, PendingResponse>,
547    active: HashMap<Uuid, ActiveResponse>,
548}
549
550#[derive(Default)]
551struct ServerIndexes {
552    registration_instance: HashMap<Uuid, EndpointInstanceId>,
553    instance_registrations: HashMap<EndpointInstanceId, Vec<Uuid>>,
554    removed_instances: HashMap<EndpointInstanceId, Instant>,
555    connection_bundles: HashMap<usize, Uuid>,
556    bundle_connections: HashMap<Uuid, HashMap<usize, quinn::Connection>>,
557}
558
559struct ServerState {
560    registrations: [Mutex<RegistrationShard>; SERVER_STATE_SHARDS],
561    indexes: Mutex<ServerIndexes>,
562}
563
564impl Default for ServerState {
565    fn default() -> Self {
566        Self {
567            registrations: std::array::from_fn(|_| Mutex::new(RegistrationShard::default())),
568            indexes: Mutex::new(ServerIndexes::default()),
569        }
570    }
571}
572
573impl ServerState {
574    fn registration(&self, registration_id: Uuid) -> &Mutex<RegistrationShard> {
575        &self.registrations[registration_id.as_u128() as usize % SERVER_STATE_SHARDS]
576    }
577}
578
579#[derive(Default)]
580struct DeferredResponse {
581    data: VecDeque<Bytes>,
582    bytes: usize,
583    end: bool,
584}
585
586impl DeferredResponse {
587    fn push(&mut self, payload: Bytes) -> Result<(), Bytes> {
588        let Some(bytes) = self.bytes.checked_add(payload.len()) else {
589            return Err(payload);
590        };
591        if self.data.len() >= MAX_DEFERRED_RESPONSE_FRAMES || bytes > MAX_DEFERRED_RESPONSE_BYTES {
592            return Err(payload);
593        }
594        self.bytes = bytes;
595        self.data.push_back(payload);
596        Ok(())
597    }
598
599    fn take_data(&mut self) -> VecDeque<Bytes> {
600        self.bytes = 0;
601        std::mem::take(&mut self.data)
602    }
603}
604
605struct PendingResponse {
606    context: Arc<dyn AsyncEngineContext>,
607    connection: oneshot::Sender<Result<StreamReceiver, StreamPrologueError>>,
608    bundle_id: Option<Uuid>,
609    deferred: DeferredResponse,
610    monitor_cancel: CancellationToken,
611    monitor_response: Option<oneshot::Sender<mpsc::Sender<Bytes>>>,
612}
613
614struct ActiveResponse {
615    sender: mpsc::Sender<Bytes>,
616    monitor_cancel: CancellationToken,
617    bundle_id: Uuid,
618    priority_ready: bool,
619    priority_draining: bool,
620    deferred: DeferredResponse,
621}
622
623fn prune_tombstones(indexes: &mut ServerIndexes, now: Instant) {
624    indexes
625        .removed_instances
626        .retain(|_, inserted| now.saturating_duration_since(*inserted) < TOMBSTONE_TTL);
627}
628
629pub struct QuicResponseServer {
630    endpoints: Arc<[quinn::Endpoint]>,
631    advertised_address: SocketAddr,
632    frontend_id: String,
633    certificate_sha256: String,
634    state: Arc<ServerState>,
635    response_buffer_capacity: usize,
636    shutdown: CancellationToken,
637}
638
639impl fmt::Debug for QuicResponseServer {
640    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
641        f.debug_struct("QuicResponseServer")
642            .field("address", &self.advertised_address)
643            .field("frontend_id", &self.frontend_id)
644            .finish_non_exhaustive()
645    }
646}
647
648impl QuicResponseServer {
649    pub fn new(
650        bind_address: SocketAddr,
651        advertised_address: SocketAddr,
652        shutdown: CancellationToken,
653    ) -> Result<Arc<Self>, PipelineError> {
654        Self::new_with_stream_window(bind_address, advertised_address, shutdown, None)
655    }
656
657    fn new_with_stream_window(
658        bind_address: SocketAddr,
659        advertised_address: SocketAddr,
660        shutdown: CancellationToken,
661        stream_receive_window: Option<u32>,
662    ) -> Result<Arc<Self>, PipelineError> {
663        let response_config = QuicResponseConfig::from_env()?;
664        Self::new_with_config(
665            bind_address,
666            advertised_address,
667            shutdown,
668            stream_receive_window,
669            response_config,
670        )
671    }
672
673    fn new_with_config(
674        bind_address: SocketAddr,
675        advertised_address: SocketAddr,
676        shutdown: CancellationToken,
677        stream_receive_window: Option<u32>,
678        response_config: QuicResponseConfig,
679    ) -> Result<Arc<Self>, PipelineError> {
680        let certified =
681            rcgen::generate_simple_self_signed(vec!["localhost".to_string()]).map_err(|error| {
682                PipelineError::Generic(format!("failed generating QUIC certificate: {error}"))
683            })?;
684        let cert_der = CertificateDer::from(certified.cert);
685        let fingerprint = Sha256::digest(cert_der.as_ref());
686        let certificate_sha256 = encode_hex(&fingerprint);
687        let key = PrivatePkcs8KeyDer::from(certified.key_pair.serialize_der());
688
689        let provider = Arc::new(rustls::crypto::ring::default_provider());
690        let mut crypto = rustls::ServerConfig::builder_with_provider(provider)
691            .with_safe_default_protocol_versions()
692            .map_err(|error| {
693                PipelineError::Generic(format!("failed configuring QUIC TLS: {error}"))
694            })?
695            .with_no_client_auth()
696            .with_single_cert(vec![cert_der], key.into())
697            .map_err(|error| {
698                PipelineError::Generic(format!("failed configuring QUIC certificate: {error}"))
699            })?;
700        crypto.alpn_protocols = vec![ALPN.to_vec()];
701
702        let mut server_config = quinn::ServerConfig::with_crypto(Arc::new(
703            QuicServerConfig::try_from(crypto).map_err(|error| {
704                PipelineError::Generic(format!("failed configuring QUIC server: {error}"))
705            })?,
706        ));
707        let transport = Arc::get_mut(&mut server_config.transport)
708            .expect("new QUIC server config has one transport owner");
709        transport.max_concurrent_bidi_streams((BULK_LANES as u32).into());
710        transport.max_concurrent_uni_streams(0_u8.into());
711        transport.keep_alive_interval(Some(Duration::from_secs(5)));
712        if let Some(window) = stream_receive_window {
713            transport.stream_receive_window(quinn::VarInt::from_u32(window));
714        }
715
716        let socket = bind_server_udp(bind_address, false).map_err(|error| {
717            PipelineError::Generic(format!(
718                "failed binding QUIC response server socket on {bind_address}: {error}"
719            ))
720        })?;
721        let bound_address = socket.local_addr().map_err(|error| {
722            PipelineError::Generic(format!(
723                "failed reading QUIC response server socket address: {error}"
724            ))
725        })?;
726        let mut first_socket = Some(socket);
727        let mut endpoints = Vec::with_capacity(SERVER_ENDPOINTS);
728        for index in 0..SERVER_ENDPOINTS {
729            let socket = if index == 0 {
730                first_socket
731                    .take()
732                    .expect("first QUIC response socket exists")
733            } else {
734                bind_server_udp(bound_address, true).map_err(|error| {
735                    PipelineError::Generic(format!(
736                        "failed binding QUIC response reuse-port socket {index} on {bound_address}: {error}"
737                    ))
738                })?
739            };
740            let endpoint = quinn::Endpoint::new(
741                quinn::EndpointConfig::default(),
742                Some(server_config.clone()),
743                socket,
744                Arc::new(quinn::TokioRuntime),
745            )
746            .map_err(|error| {
747                PipelineError::Generic(format!(
748                    "failed creating QUIC response endpoint {index} on {bound_address}: {error}"
749                ))
750            })?;
751            endpoints.push(endpoint);
752        }
753        let endpoints: Arc<[quinn::Endpoint]> = endpoints.into();
754        let advertised_address = if advertised_address.port() == 0 {
755            bound_address
756        } else {
757            advertised_address
758        };
759        let server = Arc::new(Self {
760            endpoints,
761            advertised_address,
762            frontend_id: Uuid::new_v4().to_string(),
763            certificate_sha256,
764            state: Arc::new(ServerState::default()),
765            response_buffer_capacity: response_config.response_buffer_capacity,
766            shutdown: shutdown.child_token(),
767        });
768        Self::spawn_accept_loop(&server);
769        Ok(server)
770    }
771
772    fn spawn_accept_loop(server: &Arc<Self>) {
773        for endpoint in server.endpoints.iter() {
774            let endpoint = endpoint.clone();
775            let state = server.state.clone();
776            let response_buffer_capacity = server.response_buffer_capacity;
777            let shutdown = server.shutdown.clone();
778            tokio::spawn(async move {
779                loop {
780                    tokio::select! {
781                        _ = shutdown.cancelled() => {
782                            endpoint.close(quinn::VarInt::from_u32(0), b"runtime shutdown");
783                            break;
784                        }
785                        incoming = endpoint.accept() => {
786                            let Some(incoming) = incoming else { break };
787                            let state = state.clone();
788                            tokio::spawn(async move {
789                                match incoming.await {
790                                    Ok(connection) => run_server_connection(
791                                        connection,
792                                        state,
793                                        response_buffer_capacity,
794                                    ).await,
795                                    Err(error) => tracing::warn!(%error, "QUIC response handshake failed"),
796                                }
797                            });
798                        }
799                    }
800                }
801            });
802        }
803    }
804
805    pub fn register_response(
806        &self,
807        context: Arc<dyn AsyncEngineContext>,
808    ) -> RegisteredStream<StreamReceiver> {
809        let registration_id = Uuid::new_v4();
810        let request_id = context.id().to_string();
811        let (pending_tx, pending_rx) = oneshot::channel();
812        self.state
813            .registration(registration_id)
814            .lock()
815            .pending
816            .insert(
817                registration_id,
818                PendingResponse {
819                    context,
820                    connection: pending_tx,
821                    bundle_id: None,
822                    deferred: DeferredResponse::default(),
823                    monitor_cancel: CancellationToken::new(),
824                    monitor_response: None,
825                },
826            );
827
828        let connection_info = QuicResponseConnectionInfo {
829            version: PROTOCOL_VERSION,
830            address: self.advertised_address.to_string(),
831            frontend_id: self.frontend_id.clone(),
832            registration_id,
833            request_id,
834            certificate_sha256: self.certificate_sha256.clone(),
835        }
836        .into();
837
838        let state = self.state.clone();
839        RegisteredStream::new(connection_info, pending_rx)
840            .with_registration_id(registration_id)
841            .with_cleanup(move || {
842                remove_registration(&state, registration_id);
843            })
844    }
845
846    pub async fn associate_instance(&self, registration_id: Uuid, id: &EndpointInstanceId) -> bool {
847        let mut indexes = self.state.indexes.lock();
848        let now = Instant::now();
849        prune_tombstones(&mut indexes, now);
850        if indexes.removed_instances.contains_key(id) {
851            drop(indexes);
852            remove_registration(&self.state, registration_id);
853            return false;
854        }
855        indexes
856            .registration_instance
857            .insert(registration_id, id.clone());
858        indexes
859            .instance_registrations
860            .entry(id.clone())
861            .or_default()
862            .push(registration_id);
863        true
864    }
865
866    pub async fn cancel_response(&self, registration_id: Uuid) {
867        remove_registration(&self.state, registration_id);
868    }
869
870    pub async fn cancel_instance_streams(&self, id: &EndpointInstanceId) -> usize {
871        let registrations = {
872            let mut indexes = self.state.indexes.lock();
873            let now = Instant::now();
874            prune_tombstones(&mut indexes, now);
875            indexes.removed_instances.insert(id.clone(), now);
876            let registrations = indexes
877                .instance_registrations
878                .remove(id)
879                .unwrap_or_default();
880            for registration_id in &registrations {
881                indexes.registration_instance.remove(registration_id);
882            }
883            registrations
884        };
885        let count = registrations.len();
886        for registration_id in registrations {
887            remove_registration_data(&self.state, registration_id);
888        }
889        count
890    }
891
892    pub async fn clear_instance_tombstone(&self, id: &EndpointInstanceId) {
893        self.state.indexes.lock().removed_instances.remove(id);
894    }
895}
896
897impl Drop for QuicResponseServer {
898    fn drop(&mut self) {
899        self.shutdown.cancel();
900        for endpoint in self.endpoints.iter() {
901            endpoint.close(quinn::VarInt::from_u32(0), b"response server dropped");
902        }
903    }
904}
905
906fn bind_server_udp(address: SocketAddr, join_reuseport: bool) -> std::io::Result<UdpSocket> {
907    let socket = Socket::new(
908        Domain::for_address(address),
909        Type::DGRAM,
910        Some(Protocol::UDP),
911    )?;
912    #[cfg(target_os = "linux")]
913    if join_reuseport {
914        socket.set_reuse_address(true)?;
915        socket.set_reuse_port(true)?;
916    }
917    socket.set_nonblocking(true)?;
918    socket.bind(&address.into())?;
919    // Bind the first endpoint exclusively so an ephemeral port cannot join an
920    // unrelated server's reuse-port group. Linux permits enabling reuse-port
921    // after that first bind, and the remaining endpoints can then join.
922    #[cfg(target_os = "linux")]
923    if !join_reuseport {
924        socket.set_reuse_address(true)?;
925        socket.set_reuse_port(true)?;
926    }
927    Ok(socket.into())
928}
929
930fn remove_registration_data(state: &ServerState, registration_id: Uuid) {
931    let mut registration = state.registration(registration_id).lock();
932    if let Some(pending) = registration.pending.remove(&registration_id) {
933        pending.monitor_cancel.cancel();
934    }
935    if let Some(active) = registration.active.remove(&registration_id) {
936        active.monitor_cancel.cancel();
937    }
938}
939
940fn remove_registration(state: &ServerState, registration_id: Uuid) {
941    let mut indexes = state.indexes.lock();
942    if let Some(instance) = indexes.registration_instance.remove(&registration_id)
943        && let Some(registrations) = indexes.instance_registrations.get_mut(&instance)
944    {
945        registrations.retain(|candidate| *candidate != registration_id);
946        if registrations.is_empty() {
947            indexes.instance_registrations.remove(&instance);
948        }
949    }
950    drop(indexes);
951    remove_registration_data(state, registration_id);
952}
953
954fn fail_registration(state: &ServerState, registration_id: Uuid, reason: &str) {
955    let mut indexes = state.indexes.lock();
956    if let Some(instance) = indexes.registration_instance.remove(&registration_id)
957        && let Some(registrations) = indexes.instance_registrations.get_mut(&instance)
958    {
959        registrations.retain(|candidate| *candidate != registration_id);
960        if registrations.is_empty() {
961            indexes.instance_registrations.remove(&instance);
962        }
963    }
964    drop(indexes);
965
966    let mut registration = state.registration(registration_id).lock();
967    if let Some(pending) = registration.pending.remove(&registration_id) {
968        pending.monitor_cancel.cancel();
969        let _ = pending.connection.send(Err(reason.into()));
970    }
971    if let Some(active) = registration.active.remove(&registration_id) {
972        active.monitor_cancel.cancel();
973    }
974}
975
976async fn run_server_connection(
977    connection: quinn::Connection,
978    state: Arc<ServerState>,
979    response_buffer_capacity: usize,
980) {
981    let connection_id = connection.stable_id();
982    crate::metrics::quic_response::track_connection(connection.clone());
983    tracing::debug!(connection_id, remote = %connection.remote_address(), "QUIC response connection established");
984    loop {
985        match connection.accept_bi().await {
986            Ok((send, recv)) => {
987                let lane_connection = connection.clone();
988                let lane_state = state.clone();
989                tokio::spawn(async move {
990                    let result = run_server_lane(
991                        send,
992                        recv,
993                        lane_connection.clone(),
994                        lane_state.clone(),
995                        response_buffer_capacity,
996                    )
997                    .await;
998                    if let Err(error) = result {
999                        tracing::warn!(connection_id, error = format!("{error:#}"), close_reason = ?lane_connection.close_reason(), "QUIC response lane failed; closing connection");
1000                        fail_server_connection(&lane_state, connection_id);
1001                        // Close explicitly if the lane failed before bundle registration.
1002                        lane_connection
1003                            .close(CLOSE_CODE_INVARIANT, b"response lane invariant failure");
1004                    }
1005                });
1006            }
1007            Err(error) => {
1008                if connection.close_reason().is_none() {
1009                    tracing::warn!(connection_id, %error, "QUIC response connection failed while accepting lane");
1010                }
1011                fail_server_connection(&state, connection_id);
1012                break;
1013            }
1014        }
1015    }
1016}
1017
1018fn register_server_connection_bundle(
1019    state: &ServerState,
1020    connection: quinn::Connection,
1021    bundle_id: Uuid,
1022) -> Result<()> {
1023    let connection_id = connection.stable_id();
1024    let mut indexes = state.indexes.lock();
1025    if let Some(existing) = indexes.connection_bundles.get(&connection_id)
1026        && *existing != bundle_id
1027    {
1028        bail!(
1029            "QUIC response connection {connection_id} changed bundle id from {existing} to {bundle_id}"
1030        );
1031    }
1032    indexes.connection_bundles.insert(connection_id, bundle_id);
1033    indexes
1034        .bundle_connections
1035        .entry(bundle_id)
1036        .or_default()
1037        .insert(connection_id, connection);
1038    Ok(())
1039}
1040
1041fn fail_server_connection(state: &ServerState, connection_id: usize) {
1042    let (bundle_id, connections) = {
1043        let mut indexes = state.indexes.lock();
1044        let Some(bundle_id) = indexes.connection_bundles.remove(&connection_id) else {
1045            return;
1046        };
1047        indexes
1048            .connection_bundles
1049            .retain(|_, candidate| *candidate != bundle_id);
1050        let connections = indexes
1051            .bundle_connections
1052            .remove(&bundle_id)
1053            .unwrap_or_default()
1054            .into_values()
1055            .collect::<Vec<_>>();
1056        (bundle_id, connections)
1057    };
1058    crate::metrics::quic_response::record_bundle_failure("frontend");
1059
1060    for connection in connections {
1061        connection.close(
1062            CLOSE_CODE_INVARIANT,
1063            b"response connection bundle invariant failure",
1064        );
1065    }
1066
1067    let doomed = state
1068        .registrations
1069        .iter()
1070        .flat_map(|shard| {
1071            let shard = shard.lock();
1072            shard
1073                .pending
1074                .iter()
1075                .filter_map(|(registration_id, pending)| {
1076                    (pending.bundle_id == Some(bundle_id)).then_some(*registration_id)
1077                })
1078                .chain(shard.active.iter().filter_map(|(registration_id, active)| {
1079                    (active.bundle_id == bundle_id).then_some(*registration_id)
1080                }))
1081                .collect::<Vec<_>>()
1082        })
1083        .collect::<Vec<_>>();
1084    for registration_id in doomed {
1085        fail_registration(
1086            state,
1087            registration_id,
1088            "QUIC response connection bundle failed before stream establishment",
1089        );
1090    }
1091}
1092
1093async fn run_server_lane(
1094    mut send: quinn::SendStream,
1095    recv: quinn::RecvStream,
1096    connection: quinn::Connection,
1097    state: Arc<ServerState>,
1098    response_buffer_capacity: usize,
1099) -> Result<()> {
1100    let mut recv = BufReader::with_capacity(RECEIVE_BUFFER_CAPACITY, recv);
1101    let bundle = read_frame(&mut recv).await?;
1102    if bundle.kind != FrameKind::Bundle || !bundle.payload.is_empty() {
1103        bail!("QUIC response lane did not start with a bundle frame");
1104    }
1105    let bundle_id = bundle.registration_id;
1106    register_server_connection_bundle(&state, connection, bundle_id)?;
1107
1108    let (control_tx, mut control_rx) = mpsc::channel::<Frame>(RESPONSE_BUFFER_CAPACITY);
1109    let writer = async move {
1110        while let Some(frame) = control_rx.recv().await {
1111            let mut chunks = [frame.header(), frame.payload];
1112            send.write_all_chunks(&mut chunks).await?;
1113        }
1114        Ok::<(), quinn::WriteError>(())
1115    };
1116
1117    let reader = async {
1118        let mut frames = READ_BUDGET_FRAMES;
1119        let mut bytes = READ_BUDGET_BYTES;
1120        loop {
1121            // Charge a bounded batch of small frames to the task budget.
1122            // Charging every token frame can leave already-received responses
1123            // queued behind repeated scheduler waits. Keep both bounds so
1124            // large frames still yield and the reverse-control writer runs.
1125            if frames >= READ_BUDGET_FRAMES || bytes >= READ_BUDGET_BYTES {
1126                tokio::task::consume_budget().await;
1127                frames = 0;
1128                bytes = 0;
1129            }
1130            let frame = read_frame(&mut recv).await?;
1131            frames += 1;
1132            bytes += FRAME_HEADER_LEN + frame.payload.len();
1133            process_server_frame(
1134                frame,
1135                bundle_id,
1136                &state,
1137                &control_tx,
1138                response_buffer_capacity,
1139            )
1140            .await?;
1141        }
1142        #[allow(unreachable_code)]
1143        Ok::<(), anyhow::Error>(())
1144    };
1145    // Keep both futures pinned across polls. This avoids a separate task wake
1146    // for registration acknowledgements without cancelling partial reads.
1147    tokio::pin!(reader, writer);
1148    tokio::select! {
1149        result = &mut reader => result,
1150        result = &mut writer => match result {
1151            Ok(()) => bail!("QUIC reverse-control writer exited unexpectedly"),
1152            Err(error) => Err(error.into()),
1153        },
1154    }
1155}
1156
1157async fn process_server_frame(
1158    frame: Frame,
1159    bundle_id: Uuid,
1160    state: &Arc<ServerState>,
1161    control_tx: &mpsc::Sender<Frame>,
1162    response_buffer_capacity: usize,
1163) -> Result<()> {
1164    match frame.kind {
1165        FrameKind::Register => {
1166            if !frame.payload.is_empty() {
1167                bail!("QUIC response Register frame carried a payload");
1168            }
1169            let monitor = {
1170                let mut registration = state.registration(frame.registration_id).lock();
1171                let Some(pending) = registration.pending.get_mut(&frame.registration_id) else {
1172                    send_control(control_tx, FrameKind::Reset, frame.registration_id)?;
1173                    return Ok(());
1174                };
1175                match pending.bundle_id {
1176                    Some(existing) if existing != bundle_id => {
1177                        bail!(
1178                            "QUIC response registration {} changed bundle id from {} to {}",
1179                            frame.registration_id,
1180                            existing,
1181                            bundle_id
1182                        );
1183                    }
1184                    Some(_) => {}
1185                    None => pending.bundle_id = Some(bundle_id),
1186                }
1187                if pending.monitor_response.is_none() {
1188                    let (response_tx, response_rx) = oneshot::channel();
1189                    pending.monitor_response = Some(response_tx);
1190                    Some((
1191                        pending.context.clone(),
1192                        response_rx,
1193                        pending.monitor_cancel.clone(),
1194                    ))
1195                } else {
1196                    None
1197                }
1198            };
1199            // Bundle ownership is visible before the worker is allowed to call
1200            // generate(), so a pre-prologue bundle failure can resolve the
1201            // pending provider with an error.
1202            send_control(control_tx, FrameKind::Registered, frame.registration_id)?;
1203            if let Some((context, response_rx, monitor_cancel)) = monitor {
1204                spawn_response_monitor(
1205                    context,
1206                    response_rx,
1207                    frame.registration_id,
1208                    control_tx.clone(),
1209                    monitor_cancel,
1210                );
1211            }
1212        }
1213        FrameKind::Prologue => {
1214            if !frame.payload.is_empty() {
1215                bail!("QUIC response Prologue frame carried a payload");
1216            }
1217            let (response_tx, response_rx) = mpsc::channel(response_buffer_capacity);
1218            let (pending_connection, monitor_response) = {
1219                let mut registration = state.registration(frame.registration_id).lock();
1220                let Some(pending) = registration.pending.remove(&frame.registration_id) else {
1221                    send_control(control_tx, FrameKind::Reset, frame.registration_id)?;
1222                    return Ok(());
1223                };
1224                let PendingResponse {
1225                    context: _,
1226                    connection,
1227                    bundle_id: pending_bundle,
1228                    deferred,
1229                    monitor_cancel,
1230                    monitor_response,
1231                } = pending;
1232                if pending_bundle != Some(bundle_id) {
1233                    bail!(
1234                        "QUIC response registration {} was not registered to bundle {}",
1235                        frame.registration_id,
1236                        bundle_id
1237                    );
1238                }
1239                registration.active.insert(
1240                    frame.registration_id,
1241                    ActiveResponse {
1242                        sender: response_tx.clone(),
1243                        monitor_cancel: monitor_cancel.clone(),
1244                        bundle_id,
1245                        priority_ready: false,
1246                        priority_draining: false,
1247                        deferred,
1248                    },
1249                );
1250                (connection, monitor_response)
1251            };
1252            let Some(monitor_response) = monitor_response else {
1253                bail!(
1254                    "QUIC response registration {} had no cancellation monitor",
1255                    frame.registration_id
1256                );
1257            };
1258            let _ = monitor_response.send(response_tx.clone());
1259            if pending_connection
1260                .send(Ok(StreamReceiver { rx: response_rx }))
1261                .is_err()
1262            {
1263                remove_registration(state, frame.registration_id);
1264                send_control(control_tx, FrameKind::Reset, frame.registration_id)?;
1265                return Ok(());
1266            }
1267        }
1268        FrameKind::Error => {
1269            let error = decode_error_payload(&frame.payload)?;
1270            let pending = state
1271                .registration(frame.registration_id)
1272                .lock()
1273                .pending
1274                .remove(&frame.registration_id);
1275            if let Some(pending) = pending {
1276                if let Some(pending_bundle) = pending.bundle_id
1277                    && pending_bundle != bundle_id
1278                {
1279                    bail!(
1280                        "QUIC response registration {} changed bundle id from {} to {}",
1281                        frame.registration_id,
1282                        pending_bundle,
1283                        bundle_id
1284                    );
1285                }
1286                let _ = pending.connection.send(Err(error));
1287                remove_registration(state, frame.registration_id);
1288            } else {
1289                send_control(control_tx, FrameKind::Reset, frame.registration_id)?;
1290            }
1291        }
1292        FrameKind::FirstData | FrameKind::PriorityEnd => {
1293            let first_payload = (frame.kind == FrameKind::FirstData).then_some(frame.payload);
1294            let active = {
1295                let mut registration = state.registration(frame.registration_id).lock();
1296                match registration.active.get_mut(&frame.registration_id) {
1297                    Some(active) if active.bundle_id != bundle_id => {
1298                        bail!(
1299                            "QUIC response registration {} changed bundle id from {} to {}",
1300                            frame.registration_id,
1301                            active.bundle_id,
1302                            bundle_id
1303                        );
1304                    }
1305                    Some(active) if active.priority_ready || active.priority_draining => None,
1306                    Some(active) => {
1307                        active.priority_draining = true;
1308                        Some(active.sender.clone())
1309                    }
1310                    None => None,
1311                }
1312            };
1313            let Some(sender) = active else {
1314                send_control(control_tx, FrameKind::Reset, frame.registration_id)?;
1315                return Ok(());
1316            };
1317
1318            if let Some(payload) = first_payload {
1319                let delivery = deliver_server_payload(&sender, payload);
1320                if let Delivery::Reset(reason) = delivery {
1321                    reset_registration(state, control_tx, frame.registration_id, reason)?;
1322                    return Ok(());
1323                }
1324            }
1325
1326            drain_deferred_response(state, frame.registration_id, &sender, control_tx).await?;
1327        }
1328        FrameKind::Data => {
1329            enum Disposition {
1330                Deferred,
1331                Deliver(mpsc::Sender<Bytes>, Bytes),
1332                Reset,
1333                Overflow,
1334            }
1335
1336            let disposition = {
1337                let mut registration = state.registration(frame.registration_id).lock();
1338                match registration.active.get_mut(&frame.registration_id) {
1339                    Some(active) if active.bundle_id != bundle_id => {
1340                        bail!(
1341                            "QUIC response registration {} changed bundle id from {} to {}",
1342                            frame.registration_id,
1343                            active.bundle_id,
1344                            bundle_id
1345                        );
1346                    }
1347                    Some(active) if !active.priority_ready => {
1348                        match active.deferred.push(frame.payload) {
1349                            Ok(()) => Disposition::Deferred,
1350                            Err(_) => Disposition::Overflow,
1351                        }
1352                    }
1353                    Some(active) => Disposition::Deliver(active.sender.clone(), frame.payload),
1354                    None => {
1355                        if let Some(pending) = registration.pending.get_mut(&frame.registration_id)
1356                        {
1357                            if let Some(pending_bundle) = pending.bundle_id
1358                                && pending_bundle != bundle_id
1359                            {
1360                                bail!(
1361                                    "QUIC response registration {} changed bundle id from {} to {}",
1362                                    frame.registration_id,
1363                                    pending_bundle,
1364                                    bundle_id
1365                                );
1366                            }
1367                            pending.bundle_id = Some(bundle_id);
1368                            match pending.deferred.push(frame.payload) {
1369                                Ok(()) => Disposition::Deferred,
1370                                Err(_) => Disposition::Overflow,
1371                            }
1372                        } else {
1373                            Disposition::Reset
1374                        }
1375                    }
1376                }
1377            };
1378            match disposition {
1379                Disposition::Deferred => {}
1380                Disposition::Deliver(sender, payload) => {
1381                    if let Delivery::Reset(reason) = deliver_server_payload(&sender, payload) {
1382                        reset_registration(state, control_tx, frame.registration_id, reason)?;
1383                    }
1384                }
1385                Disposition::Reset => {
1386                    send_control(control_tx, FrameKind::Reset, frame.registration_id)?;
1387                }
1388                Disposition::Overflow => {
1389                    tracing::warn!(
1390                        registration_id = %frame.registration_id,
1391                        "QUIC response deferred-data bound exceeded"
1392                    );
1393                    remove_registration(state, frame.registration_id);
1394                    send_control(control_tx, FrameKind::Reset, frame.registration_id)?;
1395                }
1396            }
1397        }
1398        FrameKind::End => {
1399            if !frame.payload.is_empty() {
1400                bail!("QUIC response End frame carried a payload");
1401            }
1402            let disposition = {
1403                let mut registration = state.registration(frame.registration_id).lock();
1404                match registration.active.get_mut(&frame.registration_id) {
1405                    Some(active) if active.bundle_id != bundle_id => {
1406                        bail!(
1407                            "QUIC response registration {} changed bundle id from {} to {}",
1408                            frame.registration_id,
1409                            active.bundle_id,
1410                            bundle_id
1411                        );
1412                    }
1413                    Some(active) if !active.priority_ready => {
1414                        active.deferred.end = true;
1415                        0
1416                    }
1417                    Some(_) => 1,
1418                    None => {
1419                        if let Some(pending) = registration.pending.get_mut(&frame.registration_id)
1420                        {
1421                            if let Some(pending_bundle) = pending.bundle_id
1422                                && pending_bundle != bundle_id
1423                            {
1424                                bail!(
1425                                    "QUIC response registration {} changed bundle id from {} to {}",
1426                                    frame.registration_id,
1427                                    pending_bundle,
1428                                    bundle_id
1429                                );
1430                            }
1431                            pending.bundle_id = Some(bundle_id);
1432                            pending.deferred.end = true;
1433                            0
1434                        } else {
1435                            2
1436                        }
1437                    }
1438                }
1439            };
1440            match disposition {
1441                0 => {}
1442                1 => remove_registration(state, frame.registration_id),
1443                2 => {
1444                    send_control(control_tx, FrameKind::Reset, frame.registration_id)?;
1445                }
1446                _ => unreachable!(),
1447            }
1448        }
1449        FrameKind::Reset => {
1450            if !frame.payload.is_empty() {
1451                bail!("QUIC response Reset frame carried a payload");
1452            }
1453            let registration = state.registration(frame.registration_id).lock();
1454            let registered_bundle = registration
1455                .pending
1456                .get(&frame.registration_id)
1457                .and_then(|pending| pending.bundle_id)
1458                .or_else(|| {
1459                    registration
1460                        .active
1461                        .get(&frame.registration_id)
1462                        .map(|active| active.bundle_id)
1463                });
1464            if registered_bundle != Some(bundle_id) {
1465                bail!(
1466                    "QUIC response registration {} was not registered to bundle {}",
1467                    frame.registration_id,
1468                    bundle_id
1469                );
1470            }
1471            drop(registration);
1472            remove_registration(state, frame.registration_id);
1473        }
1474        FrameKind::Stop | FrameKind::Kill | FrameKind::Registered | FrameKind::Bundle => {
1475            bail!("worker sent reverse-only control frame {:?}", frame.kind)
1476        }
1477    }
1478    Ok(())
1479}
1480
1481enum Delivery {
1482    Delivered,
1483    Reset(&'static str),
1484}
1485
1486fn deliver_server_payload(sender: &mpsc::Sender<Bytes>, payload: Bytes) -> Delivery {
1487    match sender.try_send(payload) {
1488        Ok(()) => Delivery::Delivered,
1489        Err(mpsc::error::TrySendError::Full(_)) => Delivery::Reset("full"),
1490        Err(mpsc::error::TrySendError::Closed(_)) => Delivery::Reset("closed"),
1491    }
1492}
1493
1494fn reset_registration(
1495    state: &ServerState,
1496    control_tx: &mpsc::Sender<Frame>,
1497    registration_id: Uuid,
1498    reason: &'static str,
1499) -> Result<()> {
1500    crate::metrics::quic_response::SERVER_MAILBOX_RESETS_TOTAL
1501        .with_label_values(&[reason])
1502        .inc();
1503    remove_registration(state, registration_id);
1504    send_control(control_tx, FrameKind::Reset, registration_id)
1505}
1506
1507async fn drain_deferred_response(
1508    state: &Arc<ServerState>,
1509    registration_id: Uuid,
1510    sender: &mpsc::Sender<Bytes>,
1511    control_tx: &mpsc::Sender<Frame>,
1512) -> Result<()> {
1513    loop {
1514        let (deferred, remove_after_drain) = {
1515            let mut registration = state.registration(registration_id).lock();
1516            let Some(active) = registration.active.get_mut(&registration_id) else {
1517                return Ok(());
1518            };
1519            if active.deferred.data.is_empty() {
1520                active.priority_ready = true;
1521                active.priority_draining = false;
1522                (VecDeque::new(), active.deferred.end)
1523            } else {
1524                (active.deferred.take_data(), false)
1525            }
1526        };
1527
1528        if deferred.is_empty() {
1529            if remove_after_drain {
1530                remove_registration(state, registration_id);
1531            }
1532            return Ok(());
1533        }
1534
1535        for payload in deferred {
1536            if let Delivery::Reset(reason) = deliver_server_payload(sender, payload) {
1537                reset_registration(state, control_tx, registration_id, reason)?;
1538                return Ok(());
1539            }
1540        }
1541    }
1542}
1543
1544fn spawn_response_monitor(
1545    context: Arc<dyn AsyncEngineContext>,
1546    response_rx: oneshot::Receiver<mpsc::Sender<Bytes>>,
1547    registration_id: Uuid,
1548    control_tx: mpsc::Sender<Frame>,
1549    cancel: CancellationToken,
1550) {
1551    tokio::spawn(async move {
1552        let killed = context.killed();
1553        let stopped = context.stopped();
1554        tokio::pin!(killed, stopped, response_rx);
1555        let mut stop_sent = false;
1556        let response_tx = loop {
1557            let kind = tokio::select! {
1558                _ = cancel.cancelled() => return,
1559                _ = &mut killed => Some(FrameKind::Kill),
1560                _ = &mut stopped, if !stop_sent => Some(FrameKind::Stop),
1561                response = &mut response_rx => match response {
1562                    Ok(response_tx) => break response_tx,
1563                    Err(_) => return,
1564                },
1565            };
1566            let Some(kind) = kind else { unreachable!() };
1567            if send_control(&control_tx, kind, registration_id).is_err() {
1568                return;
1569            }
1570            tracing::debug!(control_msg = ?kind, "issued control message");
1571            if kind == FrameKind::Stop {
1572                stop_sent = true;
1573            } else {
1574                return;
1575            }
1576        };
1577
1578        let closed = response_tx.closed();
1579        tokio::pin!(closed);
1580        loop {
1581            let kind = tokio::select! {
1582                _ = cancel.cancelled() => return,
1583                _ = &mut killed => FrameKind::Kill,
1584                _ = &mut stopped, if !stop_sent => FrameKind::Stop,
1585                _ = &mut closed => FrameKind::Reset,
1586            };
1587            if send_control(&control_tx, kind, registration_id).is_err() {
1588                return;
1589            }
1590            tracing::debug!(control_msg = ?kind, "issued control message");
1591            if kind == FrameKind::Stop {
1592                stop_sent = true;
1593            } else {
1594                return;
1595            }
1596        }
1597    });
1598}
1599
1600fn send_control(
1601    control_tx: &mpsc::Sender<Frame>,
1602    kind: FrameKind,
1603    registration_id: Uuid,
1604) -> Result<()> {
1605    control_tx
1606        .try_send(Frame::new(kind, registration_id, Bytes::new()))
1607        .map_err(|error| anyhow!("QUIC response reverse-control enqueue failed: {error}"))
1608}
1609
1610struct Lane {
1611    index: usize,
1612    sender: LaneSender,
1613}
1614
1615struct ResponseLanes {
1616    ordered: Arc<Lane>,
1617    priority: Arc<Lane>,
1618}
1619
1620struct ClientConnectionBundle {
1621    _endpoints: Arc<[quinn::Endpoint]>,
1622    connections: Arc<[quinn::Connection]>,
1623    lanes: Arc<[ResponseLanes]>,
1624    contexts: Arc<Mutex<HashMap<Uuid, Arc<ClientResponseContext>>>>,
1625    healthy: Arc<AtomicBool>,
1626}
1627
1628impl ClientConnectionBundle {
1629    fn is_healthy(&self) -> bool {
1630        self.healthy.load(Ordering::Acquire)
1631    }
1632}
1633
1634#[derive(Clone, Debug, Hash, PartialEq, Eq)]
1635struct ConnectionKey {
1636    address: String,
1637    frontend_id: String,
1638    certificate_sha256: String,
1639}
1640
1641struct ClientPoolEntry {
1642    current: ArcSwapOption<ClientConnectionBundle>,
1643    reconnect: tokio::sync::Mutex<()>,
1644}
1645
1646impl ClientPoolEntry {
1647    fn new() -> Self {
1648        Self {
1649            current: ArcSwapOption::empty(),
1650            reconnect: tokio::sync::Mutex::new(()),
1651        }
1652    }
1653}
1654
1655pub struct QuicResponseClientPool {
1656    config: QuicResponseConfig,
1657    connections: DashMap<ConnectionKey, Arc<ClientPoolEntry>>,
1658}
1659
1660static PROCESS_CLIENT_POOL: OnceLock<Arc<QuicResponseClientPool>> = OnceLock::new();
1661
1662pub fn process_client_pool_from_env() -> Result<Arc<QuicResponseClientPool>, PipelineError> {
1663    if let Some(pool) = PROCESS_CLIENT_POOL.get() {
1664        return Ok(pool.clone());
1665    }
1666    let pool = QuicResponseClientPool::from_env()?;
1667    Ok(PROCESS_CLIENT_POOL.get_or_init(|| pool).clone())
1668}
1669
1670impl fmt::Debug for QuicResponseClientPool {
1671    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
1672        f.debug_struct("QuicResponseClientPool")
1673            .field("config", &self.config)
1674            .finish_non_exhaustive()
1675    }
1676}
1677
1678impl QuicResponseClientPool {
1679    pub fn from_env() -> Result<Arc<Self>, PipelineError> {
1680        Ok(Arc::new(Self {
1681            config: QuicResponseConfig::from_env()?,
1682            connections: DashMap::new(),
1683        }))
1684    }
1685
1686    pub async fn sender(
1687        &self,
1688        context: Arc<dyn AsyncEngineContext>,
1689        connection_info: ConnectionInfo,
1690    ) -> Result<QuicResponseSender> {
1691        self.sender_with_cancellation_metric(context, connection_info, None)
1692            .await
1693    }
1694
1695    pub async fn sender_with_cancellation_metric(
1696        &self,
1697        context: Arc<dyn AsyncEngineContext>,
1698        connection_info: ConnectionInfo,
1699        cancellation_counter: Option<IntCounter>,
1700    ) -> Result<QuicResponseSender> {
1701        let info = QuicResponseConnectionInfo::try_from(connection_info)?;
1702        if info.request_id != context.id() {
1703            bail!(
1704                "QUIC response request id mismatch: connection has {}, context has {}",
1705                info.request_id,
1706                context.id()
1707            );
1708        }
1709
1710        let key = ConnectionKey {
1711            address: info.address.clone(),
1712            frontend_id: info.frontend_id.clone(),
1713            certificate_sha256: info.certificate_sha256.clone(),
1714        };
1715        let connection = self.connection(key).await?;
1716        let lane_index = xxh3_64(info.request_id.as_bytes()) as usize % connection.lanes.len();
1717        let lanes = &connection.lanes[lane_index];
1718        let (registered_tx, registered_rx) = oneshot::channel();
1719        let response_context = Arc::new(ClientResponseContext {
1720            context: context.clone(),
1721            cancellation_counter,
1722            cancellation_recorded: AtomicBool::new(false),
1723            registered: Mutex::new(Some(registered_tx)),
1724        });
1725        connection
1726            .contexts
1727            .lock()
1728            .insert(info.registration_id, response_context.clone());
1729        let mut sender = QuicResponseSender {
1730            lane: lanes.ordered.clone(),
1731            priority_lane: lanes.priority.clone(),
1732            contexts: connection.contexts.clone(),
1733            response_context,
1734            registration_id: info.registration_id,
1735            prologue_sent: false,
1736            terminated: false,
1737            first_data_sent: AtomicBool::new(false),
1738        };
1739
1740        if let Err(error) = sender
1741            .enqueue_on(
1742                &sender.priority_lane,
1743                Frame::new(FrameKind::Register, info.registration_id, Bytes::new()),
1744            )
1745            .await
1746        {
1747            connection.contexts.lock().remove(&info.registration_id);
1748            return Err(error);
1749        }
1750
1751        // Register ownership before generate(), even if cancellation raced the
1752        // handshake. The registered response monitor then propagates that
1753        // cancellation, and generate() keeps its established cancellation
1754        // semantics instead of being skipped by transport setup.
1755        let registered = registered_rx
1756            .await
1757            .map_err(|_| anyhow!("QUIC response registration acknowledgement was dropped"))?;
1758        if let Err(error) = registered {
1759            if context.is_stopped() && !context.is_killed() {
1760                let _ = sender.abort().await;
1761            } else {
1762                connection.contexts.lock().remove(&info.registration_id);
1763                let _ = sender.priority_lane.sender.try_send(Frame::new(
1764                    FrameKind::Error,
1765                    info.registration_id,
1766                    Bytes::from(error.clone()),
1767                ));
1768            }
1769            bail!(error);
1770        }
1771        Ok(sender)
1772    }
1773
1774    async fn connection(&self, key: ConnectionKey) -> Result<Arc<ClientConnectionBundle>> {
1775        let entry = self
1776            .connections
1777            .entry(key.clone())
1778            .or_insert_with(|| Arc::new(ClientPoolEntry::new()))
1779            .clone();
1780        if let Some(connection) = entry.current.load_full()
1781            && connection.is_healthy()
1782        {
1783            return Ok(connection);
1784        }
1785
1786        let _reconnect = entry.reconnect.lock().await;
1787        if let Some(connection) = entry.current.load_full()
1788            && connection.is_healthy()
1789        {
1790            return Ok(connection);
1791        }
1792        let connection = self.connect(&key).await?;
1793        entry.current.store(Some(connection.clone()));
1794        Ok(connection)
1795    }
1796
1797    async fn connect(&self, key: &ConnectionKey) -> Result<Arc<ClientConnectionBundle>> {
1798        let address: SocketAddr = key
1799            .address
1800            .parse()
1801            .with_context(|| format!("invalid QUIC response address {}", key.address))?;
1802        let expected_fingerprint = decode_fingerprint(&key.certificate_sha256)?;
1803        let client_config = pinned_client_config(expected_fingerprint)?;
1804        let priority_connection_index = BULK_CONNECTIONS;
1805        let total_connections = BULK_CONNECTIONS + PRIORITY_CONNECTIONS;
1806        let endpoints = self.client_endpoints(key, total_connections)?;
1807        let mut connections = Vec::with_capacity(total_connections);
1808        for endpoint in &endpoints {
1809            let connection = endpoint
1810                .connect_with(client_config.clone(), address, "localhost")?
1811                .await
1812                .with_context(|| {
1813                    format!("failed connecting QUIC response transport to {address}")
1814                })?;
1815            crate::metrics::quic_response::track_connection(connection.clone());
1816            connections.push(connection);
1817        }
1818        let connections: Arc<[quinn::Connection]> = connections.into();
1819
1820        let bundle_id = Uuid::new_v4();
1821        let contexts = Arc::new(Mutex::new(HashMap::new()));
1822        let healthy = Arc::new(AtomicBool::new(true));
1823        let mut lanes = Vec::with_capacity(BULK_LANES);
1824        for index in 0..BULK_LANES {
1825            let connection = connections[index % BULK_CONNECTIONS].clone();
1826            let (send, recv) = connection.open_bi().await?;
1827            let (lane_tx, lane_rx) = lane_queue(BULK_LANE_QUEUE_CAPACITY);
1828            let ordered = Arc::new(Lane {
1829                index,
1830                sender: lane_tx,
1831            });
1832            spawn_client_lane(
1833                send,
1834                recv,
1835                lane_rx,
1836                self.config.batch_interval,
1837                MAX_BATCH_FRAMES,
1838                bundle_id,
1839                connections.clone(),
1840                contexts.clone(),
1841                healthy.clone(),
1842            );
1843            let priority_connection = connections[priority_connection_index].clone();
1844            let (priority_send, priority_recv) = priority_connection.open_bi().await?;
1845            let (priority_tx, priority_rx) = lane_queue(LANE_QUEUE_CAPACITY);
1846            let priority = Arc::new(Lane {
1847                index,
1848                sender: priority_tx,
1849            });
1850            spawn_client_lane(
1851                priority_send,
1852                priority_recv,
1853                priority_rx,
1854                Duration::ZERO,
1855                1,
1856                bundle_id,
1857                connections.clone(),
1858                contexts.clone(),
1859                healthy.clone(),
1860            );
1861            lanes.push(ResponseLanes { ordered, priority });
1862        }
1863
1864        tracing::debug!(
1865            remote = %address,
1866            bulk_connections = BULK_CONNECTIONS,
1867            priority_connections = PRIORITY_CONNECTIONS,
1868            total_lanes = BULK_LANES,
1869            "QUIC response connection bundle and lanes ready"
1870        );
1871        Ok(Arc::new(ClientConnectionBundle {
1872            _endpoints: endpoints.into(),
1873            connections,
1874            lanes: lanes.into(),
1875            contexts,
1876            healthy,
1877        }))
1878    }
1879
1880    fn client_endpoints(&self, key: &ConnectionKey, count: usize) -> Result<Vec<quinn::Endpoint>> {
1881        let ipv6 = key.address.parse::<SocketAddr>()?.is_ipv6();
1882        let mut endpoints = Vec::with_capacity(count);
1883        for _ in 0..count {
1884            let bind = if ipv6 {
1885                SocketAddr::new(IpAddr::V6(Ipv6Addr::UNSPECIFIED), 0)
1886            } else {
1887                SocketAddr::new(IpAddr::V4(Ipv4Addr::UNSPECIFIED), 0)
1888            };
1889            endpoints.push(quinn::Endpoint::client(bind)?);
1890        }
1891        Ok(endpoints)
1892    }
1893}
1894
1895#[allow(clippy::too_many_arguments)]
1896fn spawn_client_lane(
1897    send: quinn::SendStream,
1898    recv: quinn::RecvStream,
1899    receiver: LaneReceiver,
1900    interval: Duration,
1901    max_batch_frames: usize,
1902    bundle_id: Uuid,
1903    connections: Arc<[quinn::Connection]>,
1904    contexts: Arc<Mutex<HashMap<Uuid, Arc<ClientResponseContext>>>>,
1905    healthy: Arc<AtomicBool>,
1906) {
1907    let writer_connections = connections.clone();
1908    let writer_contexts = contexts.clone();
1909    let writer_healthy = healthy.clone();
1910    tokio::spawn(async move {
1911        if let Err(error) = run_client_writer(
1912            send,
1913            receiver,
1914            interval,
1915            max_batch_frames,
1916            bundle_id,
1917            writer_contexts.clone(),
1918        )
1919        .await
1920        {
1921            fail_client_connection_bundle(
1922                &writer_connections,
1923                &writer_contexts,
1924                &writer_healthy,
1925                bundle_id,
1926                "writer",
1927                &format!("{error:#}"),
1928            );
1929        }
1930    });
1931
1932    tokio::spawn(async move {
1933        if let Err(error) = run_client_control_reader(recv, contexts.clone()).await {
1934            fail_client_connection_bundle(
1935                &connections,
1936                &contexts,
1937                &healthy,
1938                bundle_id,
1939                "control_reader",
1940                &format!("{error:#}"),
1941            );
1942        }
1943    });
1944}
1945
1946async fn run_client_writer(
1947    mut send: quinn::SendStream,
1948    mut receiver: LaneReceiver,
1949    interval: Duration,
1950    max_batch_frames: usize,
1951    bundle_id: Uuid,
1952    contexts: Arc<Mutex<HashMap<Uuid, Arc<ClientResponseContext>>>>,
1953) -> Result<()> {
1954    let bundle = Frame::new(FrameKind::Bundle, bundle_id, Bytes::new());
1955    let mut bundle_chunks = [bundle.header()];
1956    send.write_all_chunks(&mut bundle_chunks).await?;
1957
1958    let mut batch = Vec::with_capacity(max_batch_frames);
1959    let mut chunks = Vec::with_capacity(max_batch_frames * 2);
1960    while receive_batch(&mut receiver, &mut batch, interval, max_batch_frames).await {
1961        chunks.clear();
1962        for frame in &batch {
1963            chunks.push(frame.header());
1964            if !frame.payload.is_empty() {
1965                chunks.push(frame.payload.clone());
1966            }
1967        }
1968        send.write_all_chunks(&mut chunks).await?;
1969        if batch.iter().any(|frame| {
1970            matches!(
1971                frame.kind,
1972                FrameKind::Error | FrameKind::End | FrameKind::Reset
1973            )
1974        }) {
1975            let mut contexts = contexts.lock();
1976            for frame in &batch {
1977                if matches!(
1978                    frame.kind,
1979                    FrameKind::Error | FrameKind::End | FrameKind::Reset
1980                ) {
1981                    contexts.remove(&frame.registration_id);
1982                }
1983            }
1984        }
1985        batch.clear();
1986    }
1987    bail!("QUIC response lane queue closed unexpectedly")
1988}
1989
1990async fn run_client_control_reader(
1991    recv: quinn::RecvStream,
1992    contexts: Arc<Mutex<HashMap<Uuid, Arc<ClientResponseContext>>>>,
1993) -> Result<()> {
1994    let mut recv = BufReader::with_capacity(RECEIVE_BUFFER_CAPACITY, recv);
1995    loop {
1996        let frame = read_frame(&mut recv).await?;
1997        if !frame.payload.is_empty() {
1998            bail!("QUIC response reverse control carried a payload");
1999        }
2000        let entry = contexts.lock().get(&frame.registration_id).cloned();
2001        match frame.kind {
2002            FrameKind::Registered => {
2003                if let Some(entry) = entry
2004                    && let Some(registered) = entry.registered.lock().take()
2005                {
2006                    let _ = registered.send(Ok(()));
2007                }
2008            }
2009            FrameKind::Stop => {
2010                if let Some(entry) = entry {
2011                    entry.fail_registration("QUIC response registration was stopped");
2012                    entry.record_cancellation();
2013                    entry.context.stop();
2014                }
2015            }
2016            FrameKind::Kill | FrameKind::Reset => {
2017                if let Some(entry) = entry {
2018                    entry.fail_registration("QUIC response registration was reset");
2019                    entry.record_cancellation();
2020                    entry.context.kill();
2021                }
2022            }
2023            _ => bail!("frontend sent forward-only frame {:?}", frame.kind),
2024        }
2025    }
2026}
2027
2028fn fail_client_connection_bundle(
2029    connections: &[quinn::Connection],
2030    contexts: &Mutex<HashMap<Uuid, Arc<ClientResponseContext>>>,
2031    healthy: &AtomicBool,
2032    bundle_id: Uuid,
2033    failure_path: &'static str,
2034    reason: &str,
2035) {
2036    if !healthy.swap(false, Ordering::AcqRel) {
2037        return;
2038    }
2039    crate::metrics::quic_response::record_bundle_failure("worker");
2040    let remote = connections.first().map(quinn::Connection::remote_address);
2041    tracing::warn!(
2042        %bundle_id,
2043        failure_path,
2044        ?remote,
2045        reason,
2046        "QUIC response connection bundle invariant failed"
2047    );
2048    for (_, entry) in contexts.lock().drain() {
2049        entry.fail_registration(reason);
2050        entry.record_cancellation();
2051        entry.context.kill();
2052    }
2053    for connection in connections {
2054        connection.close(
2055            CLOSE_CODE_INVARIANT,
2056            b"response connection bundle invariant failure",
2057        );
2058    }
2059}
2060
2061pub struct QuicResponseSender {
2062    lane: Arc<Lane>,
2063    priority_lane: Arc<Lane>,
2064    contexts: Arc<Mutex<HashMap<Uuid, Arc<ClientResponseContext>>>>,
2065    response_context: Arc<ClientResponseContext>,
2066    registration_id: Uuid,
2067    prologue_sent: bool,
2068    terminated: bool,
2069    first_data_sent: AtomicBool,
2070}
2071
2072struct ClientResponseContext {
2073    context: Arc<dyn AsyncEngineContext>,
2074    cancellation_counter: Option<IntCounter>,
2075    cancellation_recorded: AtomicBool,
2076    registered: Mutex<Option<oneshot::Sender<std::result::Result<(), String>>>>,
2077}
2078
2079impl ClientResponseContext {
2080    fn fail_registration(&self, reason: &str) {
2081        if let Some(registered) = self.registered.lock().take() {
2082            let _ = registered.send(Err(reason.to_string()));
2083        }
2084    }
2085
2086    fn record_cancellation(&self) {
2087        if !self.cancellation_recorded.swap(true, Ordering::AcqRel)
2088            && let Some(counter) = &self.cancellation_counter
2089        {
2090            counter.inc();
2091        }
2092    }
2093}
2094
2095impl QuicResponseSender {
2096    #[cfg(test)]
2097    fn lane_index(&self) -> usize {
2098        self.lane.index
2099    }
2100
2101    pub async fn send_prologue(&mut self, error: Option<String>) -> Result<(), String> {
2102        self.send_prologue_typed(error.map(StreamPrologueError::from_message))
2103            .await
2104    }
2105
2106    pub async fn send_prologue_typed(
2107        &mut self,
2108        error: Option<StreamPrologueError>,
2109    ) -> Result<(), String> {
2110        if self.prologue_sent {
2111            return Err("QUIC response prologue already sent".to_string());
2112        }
2113        let (kind, payload, terminal) = match error {
2114            Some(error) => (FrameKind::Error, encode_error_payload(error), true),
2115            None => (FrameKind::Prologue, Bytes::new(), false),
2116        };
2117        self.enqueue_on(
2118            &self.priority_lane,
2119            Frame::new(kind, self.registration_id, payload),
2120        )
2121        .await
2122        .map_err(|error| error.to_string())?;
2123        self.prologue_sent = true;
2124        if terminal {
2125            self.terminated = true;
2126        }
2127        Ok(())
2128    }
2129
2130    pub async fn send(&self, payload: Bytes) -> Result<()> {
2131        if !self.prologue_sent || self.terminated {
2132            bail!("QUIC response sender is not open for data");
2133        }
2134        let first_data = !self.first_data_sent.swap(true, Ordering::AcqRel);
2135        if first_data
2136            && self.response_context.context.is_stopped()
2137            && !self.response_context.context.is_killed()
2138        {
2139            bail!("QUIC response context stopped before its first response");
2140        }
2141        let frame = if first_data {
2142            Frame::new(FrameKind::FirstData, self.registration_id, payload)
2143        } else {
2144            Frame::new(FrameKind::Data, self.registration_id, payload)
2145        };
2146        if first_data {
2147            self.enqueue_on(&self.priority_lane, frame).await
2148        } else {
2149            self.enqueue_on(&self.lane, frame).await
2150        }
2151    }
2152
2153    pub async fn finish(&mut self) -> Result<()> {
2154        if !self.prologue_sent || self.terminated {
2155            return Ok(());
2156        }
2157        if !self.first_data_sent.load(Ordering::Acquire) {
2158            self.enqueue_on(
2159                &self.priority_lane,
2160                Frame::new(FrameKind::PriorityEnd, self.registration_id, Bytes::new()),
2161            )
2162            .await?;
2163        }
2164        self.enqueue_on(
2165            &self.lane,
2166            Frame::new(FrameKind::End, self.registration_id, Bytes::new()),
2167        )
2168        .await?;
2169        self.terminated = true;
2170        Ok(())
2171    }
2172
2173    /// Close one cancelled logical response without ending its shared lane.
2174    pub async fn abort(&mut self) -> Result<()> {
2175        if self.terminated {
2176            return Ok(());
2177        }
2178        if !self.prologue_sent {
2179            self.enqueue_on(
2180                &self.priority_lane,
2181                Frame::new(FrameKind::Prologue, self.registration_id, Bytes::new()),
2182            )
2183            .await?;
2184            self.prologue_sent = true;
2185        }
2186        self.response_context.record_cancellation();
2187        self.enqueue_on(
2188            &self.priority_lane,
2189            Frame::new(FrameKind::Reset, self.registration_id, Bytes::new()),
2190        )
2191        .await?;
2192        self.terminated = true;
2193        Ok(())
2194    }
2195
2196    async fn enqueue_on(&self, lane: &Lane, frame: Frame) -> Result<()> {
2197        if frame.payload.len() > MAX_FRAME_PAYLOAD {
2198            bail!(
2199                "QUIC response frame payload {} exceeds {}",
2200                frame.payload.len(),
2201                MAX_FRAME_PAYLOAD
2202            );
2203        }
2204        match lane.sender.try_send(frame) {
2205            Ok(()) => Ok(()),
2206            Err(LaneTrySendError::Closed(_)) => {
2207                bail!("QUIC response lane closed")
2208            }
2209            Err(LaneTrySendError::Full(frame)) => lane.sender.send(frame).await,
2210        }
2211    }
2212}
2213
2214impl Drop for QuicResponseSender {
2215    fn drop(&mut self) {
2216        if !self.prologue_sent {
2217            self.contexts.lock().remove(&self.registration_id);
2218            return;
2219        }
2220        if self.terminated {
2221            return;
2222        }
2223        let sender = self.lane.sender.clone();
2224        let priority_sender = (!self.first_data_sent.load(Ordering::Acquire))
2225            .then(|| self.priority_lane.sender.clone());
2226        let registration_id = self.registration_id;
2227        if let Ok(runtime) = tokio::runtime::Handle::try_current() {
2228            runtime.spawn(async move {
2229                if let Some(priority_sender) = priority_sender {
2230                    let _ = priority_sender
2231                        .send(Frame::new(
2232                            FrameKind::PriorityEnd,
2233                            registration_id,
2234                            Bytes::new(),
2235                        ))
2236                        .await;
2237                }
2238                let _ = sender
2239                    .send(Frame::new(FrameKind::End, registration_id, Bytes::new()))
2240                    .await;
2241            });
2242        }
2243    }
2244}
2245
2246#[derive(Debug)]
2247struct PinnedCertificateVerifier {
2248    expected: [u8; 32],
2249    provider: Arc<rustls::crypto::CryptoProvider>,
2250}
2251
2252impl ServerCertVerifier for PinnedCertificateVerifier {
2253    fn verify_server_cert(
2254        &self,
2255        end_entity: &CertificateDer<'_>,
2256        _intermediates: &[CertificateDer<'_>],
2257        _server_name: &ServerName<'_>,
2258        _ocsp_response: &[u8],
2259        _now: UnixTime,
2260    ) -> std::result::Result<ServerCertVerified, rustls::Error> {
2261        let actual: [u8; 32] = Sha256::digest(end_entity.as_ref()).into();
2262        if actual != self.expected {
2263            return Err(rustls::Error::General(
2264                "QUIC response certificate fingerprint mismatch".to_string(),
2265            ));
2266        }
2267        Ok(ServerCertVerified::assertion())
2268    }
2269
2270    fn verify_tls12_signature(
2271        &self,
2272        message: &[u8],
2273        cert: &CertificateDer<'_>,
2274        dss: &DigitallySignedStruct,
2275    ) -> std::result::Result<HandshakeSignatureValid, rustls::Error> {
2276        rustls::crypto::verify_tls12_signature(
2277            message,
2278            cert,
2279            dss,
2280            &self.provider.signature_verification_algorithms,
2281        )
2282    }
2283
2284    fn verify_tls13_signature(
2285        &self,
2286        message: &[u8],
2287        cert: &CertificateDer<'_>,
2288        dss: &DigitallySignedStruct,
2289    ) -> std::result::Result<HandshakeSignatureValid, rustls::Error> {
2290        rustls::crypto::verify_tls13_signature(
2291            message,
2292            cert,
2293            dss,
2294            &self.provider.signature_verification_algorithms,
2295        )
2296    }
2297
2298    fn supported_verify_schemes(&self) -> Vec<SignatureScheme> {
2299        self.provider
2300            .signature_verification_algorithms
2301            .supported_schemes()
2302    }
2303}
2304
2305fn pinned_client_config(expected: [u8; 32]) -> Result<quinn::ClientConfig> {
2306    let provider = Arc::new(rustls::crypto::ring::default_provider());
2307    let verifier = Arc::new(PinnedCertificateVerifier {
2308        expected,
2309        provider: provider.clone(),
2310    });
2311    let mut crypto = rustls::ClientConfig::builder_with_provider(provider)
2312        .with_safe_default_protocol_versions()?
2313        .dangerous()
2314        .with_custom_certificate_verifier(verifier)
2315        .with_no_client_auth();
2316    crypto.alpn_protocols = vec![ALPN.to_vec()];
2317    let crypto = QuicClientConfig::try_from(crypto)?;
2318    Ok(quinn::ClientConfig::new(Arc::new(crypto)))
2319}
2320
2321fn encode_hex(bytes: &[u8]) -> String {
2322    use std::fmt::Write as _;
2323    let mut encoded = String::with_capacity(bytes.len() * 2);
2324    for byte in bytes {
2325        write!(&mut encoded, "{byte:02x}").expect("writing to String cannot fail");
2326    }
2327    encoded
2328}
2329
2330fn decode_fingerprint(encoded: &str) -> Result<[u8; 32]> {
2331    let bytes = encoded.as_bytes();
2332    if bytes.len() != 64 || !bytes.is_ascii() {
2333        bail!("invalid ASCII SHA-256 fingerprint length {}", bytes.len());
2334    }
2335    let mut decoded = [0_u8; 32];
2336    for (pair, output) in bytes.chunks_exact(2).zip(&mut decoded) {
2337        let high = decode_hex_digit(pair[0])?;
2338        let low = decode_hex_digit(pair[1])?;
2339        *output = (high << 4) | low;
2340    }
2341    Ok(decoded)
2342}
2343
2344fn decode_hex_digit(value: u8) -> Result<u8> {
2345    match value {
2346        b'0'..=b'9' => Ok(value - b'0'),
2347        b'a'..=b'f' => Ok(value - b'a' + 10),
2348        b'A'..=b'F' => Ok(value - b'A' + 10),
2349        _ => bail!("invalid SHA-256 fingerprint hex digit"),
2350    }
2351}
2352
2353#[cfg(test)]
2354mod tests {
2355    use super::*;
2356    use crate::{
2357        engine::AsyncEngineContextProvider,
2358        error::{BackendError, DynamoError, ErrorType},
2359        pipeline::Context as PipelineContext,
2360    };
2361
2362    fn frame(index: usize) -> Frame {
2363        Frame::new(FrameKind::Data, Uuid::nil(), Bytes::from(index.to_string()))
2364    }
2365
2366    fn test_pool() -> QuicResponseClientPool {
2367        QuicResponseClientPool {
2368            config: QuicResponseConfig {
2369                batch_interval: Duration::ZERO,
2370                response_buffer_capacity: RESPONSE_BUFFER_CAPACITY,
2371            },
2372            connections: DashMap::new(),
2373        }
2374    }
2375
2376    fn test_server(
2377        response_buffer_capacity: usize,
2378    ) -> (CancellationToken, Arc<QuicResponseServer>) {
2379        let shutdown = CancellationToken::new();
2380        let address = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 0);
2381        let server = QuicResponseServer::new_with_config(
2382            address,
2383            address,
2384            shutdown.clone(),
2385            None,
2386            QuicResponseConfig {
2387                batch_interval: Duration::ZERO,
2388                response_buffer_capacity,
2389            },
2390        )
2391        .unwrap();
2392        (shutdown, server)
2393    }
2394
2395    fn only_bundle(pool: &QuicResponseClientPool) -> Arc<ClientConnectionBundle> {
2396        assert_eq!(pool.connections.len(), 1);
2397        pool.connections
2398            .iter()
2399            .next()
2400            .unwrap()
2401            .current
2402            .load_full()
2403            .unwrap()
2404    }
2405
2406    #[tokio::test]
2407    async fn fixed_batch_cap_drains_512_frames() {
2408        let (tx, mut rx) = lane_queue(MAX_BATCH_FRAMES + 1);
2409        for index in 0..=MAX_BATCH_FRAMES {
2410            tx.send(frame(index)).await.unwrap();
2411        }
2412        let mut batch = Vec::with_capacity(MAX_BATCH_FRAMES);
2413        assert!(receive_batch(&mut rx, &mut batch, Duration::ZERO, MAX_BATCH_FRAMES).await);
2414        assert_eq!(batch.len(), MAX_BATCH_FRAMES);
2415        assert_eq!(batch[0].payload, Bytes::from_static(b"0"));
2416        assert_eq!(
2417            batch[MAX_BATCH_FRAMES - 1].payload,
2418            Bytes::from_static(b"511")
2419        );
2420    }
2421
2422    #[tokio::test]
2423    async fn lane_queue_preserves_capacity_order_and_close() {
2424        let (tx, mut rx) = lane_queue(LANE_QUEUE_CAPACITY);
2425        let queue = tx.0.clone();
2426        for index in 0..LANE_QUEUE_CAPACITY {
2427            tx.try_send(frame(index)).unwrap();
2428        }
2429        assert_eq!(queue.notifications.load(Ordering::Relaxed), 1);
2430        assert!(matches!(
2431            tx.try_send(frame(LANE_QUEUE_CAPACITY)),
2432            Err(LaneTrySendError::Full(_))
2433        ));
2434        let blocked = tokio::spawn({
2435            let tx = tx.clone();
2436            async move { tx.send(frame(LANE_QUEUE_CAPACITY)).await }
2437        });
2438        tokio::task::yield_now().await;
2439        assert!(!blocked.is_finished());
2440
2441        let mut batch = Vec::with_capacity(256);
2442        assert!(receive_batch(&mut rx, &mut batch, Duration::ZERO, 256).await);
2443        for (expected, actual) in (0..256).zip(&batch) {
2444            assert_eq!(actual.payload, Bytes::from(expected.to_string()));
2445        }
2446
2447        blocked.await.unwrap().unwrap();
2448        assert_eq!(queue.notifications.load(Ordering::Relaxed), 1);
2449        assert!(receive_batch(&mut rx, &mut batch, Duration::ZERO, LANE_QUEUE_CAPACITY).await);
2450        for (expected, actual) in (256..=LANE_QUEUE_CAPACITY).zip(&batch) {
2451            assert_eq!(actual.payload, Bytes::from(expected.to_string()));
2452        }
2453
2454        tx.try_send(frame(LANE_QUEUE_CAPACITY + 1)).unwrap();
2455        assert_eq!(queue.notifications.load(Ordering::Relaxed), 1);
2456        drop(tx);
2457        assert!(receive_batch(&mut rx, &mut batch, Duration::ZERO, LANE_QUEUE_CAPACITY).await);
2458        assert_eq!(batch.len(), 1);
2459        assert_eq!(
2460            batch[0].payload,
2461            Bytes::from((LANE_QUEUE_CAPACITY + 1).to_string())
2462        );
2463        assert!(!receive_batch(&mut rx, &mut batch, Duration::ZERO, LANE_QUEUE_CAPACITY).await);
2464    }
2465
2466    #[tokio::test]
2467    async fn lane_queue_close_has_no_lost_wakeup() {
2468        let (tx, mut rx) = lane_queue(1);
2469        let waiting = tokio::spawn(async move {
2470            let mut batch = Vec::new();
2471            rx.recv_many(&mut batch, 1).await
2472        });
2473        tokio::task::yield_now().await;
2474        drop(tx);
2475        assert_eq!(
2476            tokio::time::timeout(Duration::from_secs(1), waiting)
2477                .await
2478                .expect("closed queue must wake its receiver")
2479                .unwrap(),
2480            0
2481        );
2482    }
2483
2484    #[tokio::test(start_paused = true)]
2485    async fn frames_arriving_before_deadline_join_the_same_batch() {
2486        let (tx, mut rx) = lane_queue(128);
2487        tx.send(frame(0)).await.unwrap();
2488        let task = tokio::spawn(async move {
2489            let mut batch = Vec::with_capacity(MAX_BATCH_FRAMES);
2490            let wait = receive_batch(
2491                &mut rx,
2492                &mut batch,
2493                Duration::from_millis(1),
2494                MAX_BATCH_FRAMES,
2495            )
2496            .await;
2497            (wait, batch.len())
2498        });
2499        tokio::task::yield_now().await;
2500        tokio::time::advance(Duration::from_micros(500)).await;
2501        for index in 1..64 {
2502            tx.send(frame(index)).await.unwrap();
2503        }
2504        tokio::task::yield_now().await;
2505        let (wait, len) = task.await.unwrap();
2506        assert!(wait);
2507        assert_eq!(len, 64);
2508    }
2509
2510    #[tokio::test(start_paused = true)]
2511    async fn deadline_wins_when_a_frame_arrives_at_the_boundary() {
2512        let (tx, mut rx) = lane_queue(128);
2513        tx.send(frame(0)).await.unwrap();
2514        let task = tokio::spawn(async move {
2515            let mut batch = Vec::with_capacity(MAX_BATCH_FRAMES);
2516            let wait = receive_batch(
2517                &mut rx,
2518                &mut batch,
2519                Duration::from_millis(1),
2520                MAX_BATCH_FRAMES,
2521            )
2522            .await;
2523            (wait, batch, rx)
2524        });
2525        tokio::task::yield_now().await;
2526        tokio::time::advance(Duration::from_millis(1)).await;
2527        tx.send(frame(1)).await.unwrap();
2528        tokio::task::yield_now().await;
2529
2530        let (wait, batch, mut rx) = task.await.unwrap();
2531        assert!(wait);
2532        assert_eq!(batch.len(), 1);
2533        assert_eq!(rx.recv().await.unwrap().payload, Bytes::from_static(b"1"));
2534    }
2535
2536    #[tokio::test]
2537    async fn eight_connections_eight_lanes_carry_1000_ordered_responses() {
2538        let (shutdown, server) = test_server(RESPONSE_BUFFER_CAPACITY);
2539        let pool = test_pool();
2540
2541        let mut lane_by_request = HashMap::new();
2542        for index in 0..1_000 {
2543            let context = PipelineContext::new(());
2544            let request_id = context.id().to_string();
2545            let registered = server.register_response(context.context());
2546            let (connection_info, provider) = registered.into_parts();
2547            let mut sender = pool
2548                .sender(context.context(), connection_info)
2549                .await
2550                .unwrap();
2551            let expected_lane = xxh3_64(request_id.as_bytes()) as usize % 8;
2552            assert_eq!(sender.lane_index(), expected_lane);
2553            lane_by_request.insert(request_id, sender.lane_index());
2554            sender.send_prologue(None).await.unwrap();
2555            sender
2556                .send(Bytes::from(format!("{index}:first")))
2557                .await
2558                .unwrap();
2559            sender
2560                .send(Bytes::from(format!("{index}:second")))
2561                .await
2562                .unwrap();
2563            sender.finish().await.unwrap();
2564
2565            let mut receiver = provider.await.unwrap().unwrap();
2566            assert_eq!(
2567                receiver.rx.recv().await.unwrap(),
2568                Bytes::from(format!("{index}:first"))
2569            );
2570            assert_eq!(
2571                receiver.rx.recv().await.unwrap(),
2572                Bytes::from(format!("{index}:second"))
2573            );
2574            assert!(receiver.rx.recv().await.is_none());
2575        }
2576
2577        assert_eq!(lane_by_request.len(), 1_000);
2578        let bundle = only_bundle(&pool);
2579        assert_eq!(bundle.connections.len(), 9); // Eight bulk and one priority.
2580        assert_eq!(bundle.lanes.len(), 8);
2581        shutdown.cancel();
2582    }
2583
2584    #[tokio::test]
2585    async fn registration_ack_does_not_resolve_response_stream() {
2586        let (shutdown, server) = test_server(4);
2587        let context = PipelineContext::new(());
2588        let registered = server.register_response(context.context());
2589        let registration_id = registered.registration_id().unwrap();
2590        let (_, mut provider) = registered.into_parts();
2591        let bundle_id = Uuid::new_v4();
2592        let (control_tx, mut control_rx) = mpsc::channel(4);
2593
2594        process_server_frame(
2595            Frame::new(FrameKind::Register, registration_id, Bytes::new()),
2596            bundle_id,
2597            &server.state,
2598            &control_tx,
2599            4,
2600        )
2601        .await
2602        .unwrap();
2603        assert_eq!(control_rx.recv().await.unwrap().kind, FrameKind::Registered);
2604        assert!(matches!(
2605            provider.try_recv(),
2606            Err(oneshot::error::TryRecvError::Empty)
2607        ));
2608
2609        process_server_frame(
2610            Frame::new(FrameKind::Prologue, registration_id, Bytes::new()),
2611            bundle_id,
2612            &server.state,
2613            &control_tx,
2614            4,
2615        )
2616        .await
2617        .unwrap();
2618        assert!(provider.await.unwrap().is_ok());
2619        shutdown.cancel();
2620    }
2621
2622    #[tokio::test]
2623    async fn prologue_activation_is_atomic_with_concurrent_bulk_data() {
2624        let (shutdown, server) = test_server(4);
2625        let context = PipelineContext::new(());
2626        let registered = server.register_response(context.context());
2627        let registration_id = registered.registration_id().unwrap();
2628        let (_, provider) = registered.into_parts();
2629        let bundle_id = Uuid::new_v4();
2630        let (control_tx, mut control_rx) = mpsc::channel(8);
2631
2632        process_server_frame(
2633            Frame::new(FrameKind::Register, registration_id, Bytes::new()),
2634            bundle_id,
2635            &server.state,
2636            &control_tx,
2637            4,
2638        )
2639        .await
2640        .unwrap();
2641        assert_eq!(control_rx.recv().await.unwrap().kind, FrameKind::Registered);
2642
2643        let (prologue, bulk) = tokio::join!(
2644            process_server_frame(
2645                Frame::new(FrameKind::Prologue, registration_id, Bytes::new()),
2646                bundle_id,
2647                &server.state,
2648                &control_tx,
2649                4,
2650            ),
2651            process_server_frame(
2652                Frame::new(
2653                    FrameKind::Data,
2654                    registration_id,
2655                    Bytes::from_static(b"bulk"),
2656                ),
2657                bundle_id,
2658                &server.state,
2659                &control_tx,
2660                4,
2661            )
2662        );
2663        prologue.unwrap();
2664        bulk.unwrap();
2665
2666        process_server_frame(
2667            Frame::new(
2668                FrameKind::FirstData,
2669                registration_id,
2670                Bytes::from_static(b"first"),
2671            ),
2672            bundle_id,
2673            &server.state,
2674            &control_tx,
2675            4,
2676        )
2677        .await
2678        .unwrap();
2679
2680        let mut receiver = provider.await.unwrap().unwrap();
2681        assert_eq!(
2682            receiver.rx.recv().await.unwrap(),
2683            Bytes::from_static(b"first")
2684        );
2685        assert_eq!(
2686            receiver.rx.recv().await.unwrap(),
2687            Bytes::from_static(b"bulk")
2688        );
2689        shutdown.cancel();
2690    }
2691
2692    #[tokio::test]
2693    async fn priority_path_preserves_response_order() {
2694        let (shutdown, server) = test_server(RESPONSE_BUFFER_CAPACITY);
2695        let pool = test_pool();
2696
2697        let context = PipelineContext::new(());
2698        let registered = server.register_response(context.context());
2699        let (connection_info, provider) = registered.into_parts();
2700        let mut sender = pool
2701            .sender(context.context(), connection_info)
2702            .await
2703            .unwrap();
2704        sender.send_prologue(None).await.unwrap();
2705        sender.send(Bytes::from_static(b"first")).await.unwrap();
2706        sender.send(Bytes::from_static(b"second")).await.unwrap();
2707        sender.finish().await.unwrap();
2708
2709        let mut receiver = provider.await.unwrap().unwrap();
2710        assert_eq!(
2711            receiver.rx.recv().await.unwrap(),
2712            Bytes::from_static(b"first")
2713        );
2714        assert_eq!(
2715            receiver.rx.recv().await.unwrap(),
2716            Bytes::from_static(b"second")
2717        );
2718        assert!(receiver.rx.recv().await.is_none());
2719        shutdown.cancel();
2720    }
2721
2722    #[tokio::test]
2723    async fn receiver_drop_sends_logical_reset_without_resetting_lane() {
2724        let (shutdown, server) = test_server(RESPONSE_BUFFER_CAPACITY);
2725        let pool = test_pool();
2726        let context = PipelineContext::new(());
2727        let registered = server.register_response(context.context());
2728        let (info, provider) = registered.into_parts();
2729        let mut sender = pool.sender(context.context(), info).await.unwrap();
2730        sender.send_prologue(None).await.unwrap();
2731        let receiver = provider.await.unwrap().unwrap();
2732        drop(receiver);
2733
2734        let worker_context = context.context();
2735        tokio::time::timeout(Duration::from_secs(1), worker_context.killed())
2736            .await
2737            .expect("logical reset should kill the matching worker context");
2738        assert!(only_bundle(&pool).is_healthy());
2739
2740        // A sibling response still uses the same healthy connection.
2741        let sibling = PipelineContext::new(());
2742        let registered = server.register_response(sibling.context());
2743        let (info, provider) = registered.into_parts();
2744        let mut sibling_sender = pool.sender(sibling.context(), info).await.unwrap();
2745        sibling_sender.send_prologue(None).await.unwrap();
2746        sibling_sender.finish().await.unwrap();
2747        let mut receiver = provider.await.unwrap().unwrap();
2748        assert!(receiver.rx.recv().await.is_none());
2749        shutdown.cancel();
2750    }
2751
2752    #[tokio::test]
2753    async fn priority_barrier_reorders_early_bulk_frames() {
2754        let registration_id = Uuid::new_v4();
2755        let bundle_id = Uuid::new_v4();
2756        let (response_tx, mut response_rx) = mpsc::channel(4);
2757        let state = Arc::new(ServerState::default());
2758        state.registration(registration_id).lock().active.insert(
2759            registration_id,
2760            ActiveResponse {
2761                sender: response_tx,
2762                monitor_cancel: CancellationToken::new(),
2763                bundle_id,
2764                priority_ready: false,
2765                priority_draining: false,
2766                deferred: DeferredResponse::default(),
2767            },
2768        );
2769        let (control_tx, mut control_rx) = mpsc::channel(4);
2770
2771        process_server_frame(
2772            Frame::new(
2773                FrameKind::Data,
2774                registration_id,
2775                Bytes::from_static(b"second"),
2776            ),
2777            bundle_id,
2778            &state,
2779            &control_tx,
2780            4,
2781        )
2782        .await
2783        .unwrap();
2784        process_server_frame(
2785            Frame::new(FrameKind::End, registration_id, Bytes::new()),
2786            bundle_id,
2787            &state,
2788            &control_tx,
2789            4,
2790        )
2791        .await
2792        .unwrap();
2793        assert!(matches!(
2794            response_rx.try_recv(),
2795            Err(mpsc::error::TryRecvError::Empty)
2796        ));
2797
2798        process_server_frame(
2799            Frame::new(
2800                FrameKind::FirstData,
2801                registration_id,
2802                Bytes::from_static(b"first"),
2803            ),
2804            bundle_id,
2805            &state,
2806            &control_tx,
2807            4,
2808        )
2809        .await
2810        .unwrap();
2811
2812        assert_eq!(
2813            response_rx.recv().await.unwrap(),
2814            Bytes::from_static(b"first")
2815        );
2816        assert_eq!(
2817            response_rx.recv().await.unwrap(),
2818            Bytes::from_static(b"second")
2819        );
2820        assert!(response_rx.recv().await.is_none());
2821        assert!(matches!(
2822            control_rx.try_recv(),
2823            Err(mpsc::error::TryRecvError::Empty)
2824        ));
2825    }
2826
2827    #[test]
2828    fn deferred_response_enforces_frame_and_byte_bounds() {
2829        let mut frame_bounded = DeferredResponse::default();
2830        for _ in 0..MAX_DEFERRED_RESPONSE_FRAMES {
2831            frame_bounded.push(Bytes::new()).unwrap();
2832        }
2833        assert!(frame_bounded.push(Bytes::from_static(b"overflow")).is_err());
2834
2835        let mut byte_bounded = DeferredResponse::default();
2836        byte_bounded
2837            .push(Bytes::from(vec![0; MAX_DEFERRED_RESPONSE_BYTES]))
2838            .unwrap();
2839        assert!(byte_bounded.push(Bytes::from_static(b"overflow")).is_err());
2840    }
2841
2842    #[tokio::test]
2843    async fn terminal_generate_error_uses_ordered_lane_frame() {
2844        let (shutdown, server) = test_server(RESPONSE_BUFFER_CAPACITY);
2845        let pool = test_pool();
2846        let context = PipelineContext::new(());
2847        let registered = server.register_response(context.context());
2848        let (info, provider) = registered.into_parts();
2849        let mut sender = pool.sender(context.context(), info).await.unwrap();
2850        sender
2851            .send_prologue(Some("generate failed".to_string()))
2852            .await
2853            .unwrap();
2854        match provider.await.unwrap() {
2855            Err(error) => assert_eq!(&*error, "generate failed"),
2856            Ok(_) => panic!("terminal error unexpectedly opened a response stream"),
2857        }
2858        shutdown.cancel();
2859    }
2860
2861    #[test]
2862    fn error_payload_legacy_utf8_stays_untyped() {
2863        let decoded = decode_error_payload(b"generate failed").unwrap();
2864        assert_eq!(&*decoded, "generate failed");
2865        assert!(decoded.typed_error.is_none());
2866    }
2867
2868    #[test]
2869    fn error_payload_round_trips_typed_refusal() {
2870        let original = StreamPrologueError::new(
2871            "Generate Error: multimodal input is not supported by this backend",
2872            DynamoError::builder()
2873                .error_type(ErrorType::Backend(BackendError::InvalidArgument))
2874                .message("multimodal input is not supported by this backend")
2875                .build(),
2876        );
2877        let payload = encode_error_payload(original.clone());
2878        assert_eq!(payload.first(), Some(&b'{'));
2879        let encoded = std::str::from_utf8(&payload).expect("typed payload is JSON UTF-8");
2880        assert!(
2881            encoded.contains(QUIC_TYPED_ERROR_KIND),
2882            "typed payload must carry the wire discriminant; got {encoded}"
2883        );
2884
2885        let decoded = decode_error_payload(&payload).unwrap();
2886        assert_eq!(decoded, original);
2887        assert_eq!(
2888            decoded.typed_error.map(|e| e.error_type()),
2889            Some(ErrorType::Backend(BackendError::InvalidArgument))
2890        );
2891    }
2892
2893    #[test]
2894    fn error_payload_python_envelope_is_not_a_prologue() {
2895        let envelope = br#"{"message":"unsupported media type","code":415}"#;
2896        let decoded = decode_error_payload(envelope).unwrap();
2897        assert_eq!(decoded.message.as_bytes(), envelope);
2898        assert!(
2899            decoded.typed_error.is_none(),
2900            "a Python HTTP envelope must stay the untyped message, not a prologue"
2901        );
2902    }
2903
2904    #[test]
2905    fn error_payload_legacy_json_with_decoy_error_field_is_preserved() {
2906        let envelope = br#"{"error":"failed","code":500}"#;
2907        let decoded = decode_error_payload(envelope).unwrap();
2908        assert_eq!(
2909            decoded.message.as_bytes(),
2910            envelope,
2911            "untagged legacy JSON must be preserved byte-for-byte"
2912        );
2913        assert!(
2914            decoded.typed_error.is_none(),
2915            "a decoy error field must not be classified as a typed prologue"
2916        );
2917    }
2918
2919    #[tokio::test]
2920    async fn typed_generate_error_survives_the_quic_error_frame() {
2921        let (shutdown, server) = test_server(RESPONSE_BUFFER_CAPACITY);
2922        let pool = test_pool();
2923        let context = PipelineContext::new(());
2924        let registered = server.register_response(context.context());
2925        let (info, provider) = registered.into_parts();
2926        let mut sender = pool.sender(context.context(), info).await.unwrap();
2927        sender
2928            .send_prologue_typed(Some(StreamPrologueError::new(
2929                "Generate Error: multimodal input is not supported by this backend",
2930                DynamoError::builder()
2931                    .error_type(ErrorType::Backend(BackendError::InvalidArgument))
2932                    .message("multimodal input is not supported by this backend")
2933                    .build(),
2934            )))
2935            .await
2936            .unwrap();
2937        match provider.await.unwrap() {
2938            Err(error) => {
2939                assert!(error.contains("multimodal input is not supported"));
2940                assert_eq!(
2941                    error.typed_error.as_ref().map(|e| e.error_type()),
2942                    Some(ErrorType::Backend(BackendError::InvalidArgument)),
2943                    "the worker's error type must survive the QUIC error frame"
2944                );
2945            }
2946            Ok(_) => panic!("terminal error unexpectedly opened a response stream"),
2947        }
2948        shutdown.cancel();
2949    }
2950
2951    #[tokio::test]
2952    async fn cancelled_sender_closes_only_its_logical_response() {
2953        let (shutdown, server) = test_server(RESPONSE_BUFFER_CAPACITY);
2954        let pool = test_pool();
2955
2956        let cancelled = PipelineContext::new(());
2957        let registered = server.register_response(cancelled.context());
2958        let (info, provider) = registered.into_parts();
2959        let mut sender = pool.sender(cancelled.context(), info).await.unwrap();
2960        sender.abort().await.unwrap();
2961        let mut receiver = provider.await.unwrap().unwrap();
2962        assert!(receiver.rx.recv().await.is_none());
2963
2964        let sibling = PipelineContext::new(());
2965        let registered = server.register_response(sibling.context());
2966        let (info, provider) = registered.into_parts();
2967        let mut sibling_sender = pool.sender(sibling.context(), info).await.unwrap();
2968        sibling_sender.send_prologue(None).await.unwrap();
2969        sibling_sender.finish().await.unwrap();
2970        let mut receiver = provider.await.unwrap().unwrap();
2971        assert!(receiver.rx.recv().await.is_none());
2972        assert!(only_bundle(&pool).is_healthy());
2973        shutdown.cancel();
2974    }
2975
2976    #[tokio::test]
2977    async fn physical_connection_failure_kills_all_and_reconnects_once() {
2978        let (shutdown, server) = test_server(RESPONSE_BUFFER_CAPACITY);
2979        let pool = Arc::new(test_pool());
2980
2981        let first_context = PipelineContext::new(());
2982        let first = server.register_response(first_context.context());
2983        let first_info = first.connection_info.clone();
2984        let mut first_sender = pool
2985            .sender(first_context.context(), first_info.clone())
2986            .await
2987            .unwrap();
2988        first_sender.send_prologue(None).await.unwrap();
2989        let old_connections = only_bundle(&pool)
2990            .connections
2991            .iter()
2992            .cloned()
2993            .collect::<Vec<_>>();
2994        let old_ids = old_connections
2995            .iter()
2996            .map(quinn::Connection::stable_id)
2997            .collect::<Vec<_>>();
2998        old_connections[0].close(CLOSE_CODE_INVARIANT, b"test lane failure");
2999        let first_engine_context = first_context.context();
3000        tokio::time::timeout(Duration::from_secs(1), first_engine_context.killed())
3001            .await
3002            .expect("connection failure should kill every active context");
3003        tokio::time::timeout(Duration::from_secs(1), async {
3004            while old_connections
3005                .iter()
3006                .any(|connection| connection.close_reason().is_none())
3007            {
3008                tokio::task::yield_now().await;
3009            }
3010        })
3011        .await
3012        .expect("one physical failure must close the complete client bundle");
3013        let mut registrations = Vec::new();
3014        let mut providers = Vec::new();
3015        for _ in 0..16 {
3016            let context = PipelineContext::new(());
3017            let registered = server.register_response(context.context());
3018            let (connection_info, provider) = registered.into_parts();
3019            registrations.push((context, connection_info));
3020            providers.push(provider);
3021        }
3022        let senders =
3023            futures::future::join_all(registrations.into_iter().map(|(context, info)| {
3024                let pool = pool.clone();
3025                async move { pool.sender(context.context(), info).await }
3026            }))
3027            .await;
3028        let errors = senders
3029            .iter()
3030            .filter_map(|result| result.as_ref().err().map(ToString::to_string))
3031            .collect::<Vec<_>>();
3032        assert!(errors.is_empty(), "replacement sender errors: {errors:?}");
3033        let replacement = only_bundle(&pool);
3034        assert_eq!(
3035            replacement.connections.len(),
3036            BULK_CONNECTIONS + PRIORITY_CONNECTIONS
3037        );
3038        assert!(
3039            replacement
3040                .connections
3041                .iter()
3042                .all(|connection| !old_ids.contains(&connection.stable_id()))
3043        );
3044        assert_eq!(replacement.lanes.len(), 8);
3045        drop(providers);
3046        shutdown.cancel();
3047    }
3048
3049    #[tokio::test]
3050    async fn different_connection_keys_initialize_in_parallel() {
3051        let (shutdown, server) = test_server(RESPONSE_BUFFER_CAPACITY);
3052        let pool = test_pool();
3053
3054        let blocked_key = ConnectionKey {
3055            address: "127.0.0.1:1".to_string(),
3056            frontend_id: "blocked".to_string(),
3057            certificate_sha256: "00".repeat(32),
3058        };
3059        let blocked_entry = pool
3060            .connections
3061            .entry(blocked_key)
3062            .or_insert_with(|| Arc::new(ClientPoolEntry::new()))
3063            .clone();
3064        let blocked_reconnect = blocked_entry.reconnect.lock().await;
3065
3066        let context = PipelineContext::new(());
3067        let registered = server.register_response(context.context());
3068        let (info, provider) = registered.into_parts();
3069        let sender =
3070            tokio::time::timeout(Duration::from_secs(1), pool.sender(context.context(), info))
3071                .await
3072                .expect("one frontend key must not wait for another key's reconnect mutex")
3073                .unwrap();
3074
3075        drop(sender);
3076        drop(provider);
3077        drop(blocked_reconnect);
3078        shutdown.cancel();
3079    }
3080
3081    #[tokio::test]
3082    async fn bundle_failure_before_prologue_resolves_pending_provider() {
3083        let (shutdown, server) = test_server(RESPONSE_BUFFER_CAPACITY);
3084        let pool = test_pool();
3085        let context = PipelineContext::new(());
3086        let registered = server.register_response(context.context());
3087        let (info, provider) = registered.into_parts();
3088        let _sender = pool.sender(context.context(), info).await.unwrap();
3089        let bundle = only_bundle(&pool);
3090        bundle.connections[0].close(CLOSE_CODE_INVARIANT, b"test pre-prologue failure");
3091
3092        let result = tokio::time::timeout(Duration::from_secs(2), provider)
3093            .await
3094            .expect("pending provider must resolve after bundle failure")
3095            .unwrap();
3096        let error = match result {
3097            Err(error) => error,
3098            Ok(_) => panic!("pre-prologue bundle failure unexpectedly opened a stream"),
3099        };
3100        assert!(error.contains("bundle failed"));
3101        tokio::time::timeout(Duration::from_secs(2), context.context().killed())
3102            .await
3103            .expect("bundle failure must kill the worker request context");
3104        shutdown.cancel();
3105    }
3106
3107    #[tokio::test]
3108    async fn full_mailbox_resets_only_its_registration() {
3109        let state = Arc::new(ServerState::default());
3110        let bundle_id = Uuid::new_v4();
3111        let full_id = Uuid::new_v4();
3112        let sibling_id = Uuid::new_v4();
3113        let (full_tx, mut full_rx) = mpsc::channel(1);
3114        full_tx.try_send(Bytes::from_static(b"occupied")).unwrap();
3115        let (sibling_tx, mut sibling_rx) = mpsc::channel(1);
3116        for (registration_id, sender) in [(full_id, full_tx), (sibling_id, sibling_tx)] {
3117            state.registration(registration_id).lock().active.insert(
3118                registration_id,
3119                ActiveResponse {
3120                    sender,
3121                    monitor_cancel: CancellationToken::new(),
3122                    bundle_id,
3123                    priority_ready: true,
3124                    priority_draining: false,
3125                    deferred: DeferredResponse::default(),
3126                },
3127            );
3128        }
3129        let instance = EndpointInstanceId {
3130            namespace: "n".to_string(),
3131            component: "c".to_string(),
3132            endpoint: "e".to_string(),
3133            instance_id: 1,
3134        };
3135        {
3136            let mut indexes = state.indexes.lock();
3137            indexes
3138                .registration_instance
3139                .insert(full_id, instance.clone());
3140            indexes
3141                .instance_registrations
3142                .insert(instance.clone(), vec![full_id]);
3143        }
3144        let (control_tx, mut control_rx) = mpsc::channel(4);
3145
3146        process_server_frame(
3147            Frame::new(FrameKind::Data, full_id, Bytes::from_static(b"reset")),
3148            bundle_id,
3149            &state,
3150            &control_tx,
3151            1,
3152        )
3153        .await
3154        .unwrap();
3155        assert_eq!(control_rx.recv().await.unwrap().kind, FrameKind::Reset);
3156        assert!(
3157            !state
3158                .registration(full_id)
3159                .lock()
3160                .active
3161                .contains_key(&full_id)
3162        );
3163        assert!(
3164            !state
3165                .indexes
3166                .lock()
3167                .registration_instance
3168                .contains_key(&full_id)
3169        );
3170        assert_eq!(
3171            full_rx.recv().await.unwrap(),
3172            Bytes::from_static(b"occupied")
3173        );
3174
3175        process_server_frame(
3176            Frame::new(FrameKind::Data, sibling_id, Bytes::from_static(b"sibling")),
3177            bundle_id,
3178            &state,
3179            &control_tx,
3180            1,
3181        )
3182        .await
3183        .unwrap();
3184        assert_eq!(
3185            sibling_rx.recv().await.unwrap(),
3186            Bytes::from_static(b"sibling")
3187        );
3188    }
3189
3190    #[tokio::test]
3191    async fn priority_delivery_failures_remove_all_registration_indexes() {
3192        for deferred_failure in [false, true] {
3193            let state = Arc::new(ServerState::default());
3194            let bundle_id = Uuid::new_v4();
3195            let registration_id = Uuid::new_v4();
3196            let (response_tx, _response_rx) = mpsc::channel(1);
3197            let mut deferred = DeferredResponse::default();
3198            if deferred_failure {
3199                deferred.push(Bytes::from_static(b"bulk")).unwrap();
3200            } else {
3201                response_tx
3202                    .try_send(Bytes::from_static(b"occupied"))
3203                    .unwrap();
3204            }
3205            state.registration(registration_id).lock().active.insert(
3206                registration_id,
3207                ActiveResponse {
3208                    sender: response_tx,
3209                    monitor_cancel: CancellationToken::new(),
3210                    bundle_id,
3211                    priority_ready: false,
3212                    priority_draining: false,
3213                    deferred,
3214                },
3215            );
3216            let instance = EndpointInstanceId {
3217                namespace: "n".to_string(),
3218                component: "c".to_string(),
3219                endpoint: "e".to_string(),
3220                instance_id: registration_id.as_u128() as u64,
3221            };
3222            {
3223                let mut indexes = state.indexes.lock();
3224                indexes
3225                    .registration_instance
3226                    .insert(registration_id, instance.clone());
3227                indexes
3228                    .instance_registrations
3229                    .insert(instance.clone(), vec![registration_id]);
3230            }
3231            let (control_tx, mut control_rx) = mpsc::channel(4);
3232
3233            process_server_frame(
3234                Frame::new(
3235                    FrameKind::FirstData,
3236                    registration_id,
3237                    Bytes::from_static(b"first"),
3238                ),
3239                bundle_id,
3240                &state,
3241                &control_tx,
3242                1,
3243            )
3244            .await
3245            .unwrap();
3246
3247            assert_eq!(control_rx.recv().await.unwrap().kind, FrameKind::Reset);
3248            assert!(
3249                !state
3250                    .registration(registration_id)
3251                    .lock()
3252                    .active
3253                    .contains_key(&registration_id)
3254            );
3255            let indexes = state.indexes.lock();
3256            assert!(!indexes.registration_instance.contains_key(&registration_id));
3257            assert!(!indexes.instance_registrations.contains_key(&instance));
3258        }
3259    }
3260
3261    #[test]
3262    fn fingerprint_round_trip() {
3263        let fingerprint = [0x5a; 32];
3264        assert_eq!(
3265            decode_fingerprint(&encode_hex(&fingerprint)).unwrap(),
3266            fingerprint
3267        );
3268        assert!(decode_fingerprint(&"é".repeat(32)).is_err());
3269    }
3270
3271    #[tokio::test]
3272    async fn server_drop_cancels_accept_loops_and_closes_endpoints() {
3273        let (shutdown, server) = test_server(1);
3274        assert_eq!(server.endpoints.len(), SERVER_ENDPOINTS);
3275        let accept_shutdown = server.shutdown.clone();
3276        drop(server);
3277        assert!(accept_shutdown.is_cancelled());
3278        assert!(!shutdown.is_cancelled());
3279    }
3280
3281    #[cfg(target_os = "linux")]
3282    #[test]
3283    fn primary_udp_bind_is_exclusive_before_reuseport_siblings_join() {
3284        let ephemeral = "127.0.0.1:0".parse().unwrap();
3285        let primary = bind_server_udp(ephemeral, false).unwrap();
3286        let bound = primary.local_addr().unwrap();
3287        let sibling = bind_server_udp(bound, true).unwrap();
3288        assert_eq!(sibling.local_addr().unwrap(), bound);
3289
3290        let other_primary = bind_server_udp(ephemeral, false).unwrap();
3291        assert_ne!(other_primary.local_addr().unwrap(), bound);
3292    }
3293}