1use 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;
62const 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;
71const RECEIVE_BUFFER_CAPACITY: usize = 256 * 1024;
75const READ_BUDGET_FRAMES: usize = 16;
76const READ_BUDGET_BYTES: usize = 4 * 1024;
77#[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 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
243const QUIC_TYPED_ERROR_KIND: &str = "dynamo.quic.typed_error";
252
253#[derive(Serialize, Deserialize)]
255struct QuicTypedErrorWire {
256 kind: String,
257 #[serde(flatten)]
258 prologue: ResponseStreamPrologue,
259}
260
261fn 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
285fn 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 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
506async 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 ®istrations {
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 #[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(®istration_id) {
933 pending.monitor_cancel.cancel();
934 }
935 if let Some(active) = registration.active.remove(®istration_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(®istration_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(®istration_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(®istration_id) {
968 pending.monitor_cancel.cancel();
969 let _ = pending.connection.send(Err(reason.into()));
970 }
971 if let Some(active) = registration.active.remove(®istration_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 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 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 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 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(®istration_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 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 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); 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 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(®istration_id)
3254 );
3255 let indexes = state.indexes.lock();
3256 assert!(!indexes.registration_instance.contains_key(®istration_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}