moqtap_proxy/egress.rs
1//! The per-stream deferred-write deque — the engine behind `Delay` and
2//! `Hold`.
3//!
4//! A [`PendingQueue`] is a **local variable** in the pipe function that already
5//! owns both halves of the stream. No task, no channel, no `Semaphore`: the
6//! reader keeps `&mut send`, so a mirrored peer reset still leaves the
7//! destination *after* the bytes already written, which is what
8//! `proxy_reset.rs::upstream_reset_reaches_the_client_as_a_reset_with_the_same_code`
9//! (*data, **then** reset*) asserts. A writer-task design reverses that
10//! ordering and cannot be made to pass it.
11//!
12//! The queue is empty and non-allocating unless a timing action is used —
13//! `VecDeque::new` does not allocate — so a session that only passes,
14//! replaces or drops never touches this module's hot path and never links
15//! [`crate::release_timer`] into any executed code.
16//!
17//! # The shape the pipe loop takes
18//!
19//! ```ignore
20//! let can_read = pending.accepts_more(); // byte + unit budget, a bool
21//! let head_release = pending.head_release(); // Option<Release>, owned clone
22//!
23//! tokio::select! {
24//! // `can_read` chooses *inside* observe_source: true polls
25//! // `recv.read`, false polls `recv.received_reset()`, which consumes
26//! // no bytes. The source is never unobserved — see `observe_source`.
27//! source = observe_source(&mut recv, &mut buf, can_read, watch) => { /* enqueue-or-write-now */ }
28//! _ = egress::wait_release(head_release.clone(), &ctx.cancel), if head_release.is_some() => {
29//! let now = Instant::now();
30//! while let Some(unit) = pending.pop_next_due(now) {
31//! pending.record_release(&unit, now);
32//! egress::write_unit(unit, &mut send).await?;
33//! }
34//! }
35//! _ = ctx.cancel.cancelled() => { /* framer.finish() flush, then drain, then return */ }
36//! }
37//! ```
38//!
39//! Both branch *expressions* are free of any borrow of `pending`: one is a
40//! `bool`, the other an owned [`Release`] (one `Arc` clone plus one token
41//! clone). That is a **rule, not a compile error** — a branch future holding
42//! `&mut pending` whose winning body then calls `pending.pop_next_due()`
43//! does in fact compile, because `select!` drops the branch-future tuple
44//! before the handler runs. Keeping the expressions borrow-free is what lets
45//! [`wait_release`] be a free function shared by the release branch and by
46//! [`drain_honouring_release_times`], and lets a later edit add a fourth
47//! branch without re-arranging the other three.
48//!
49//! **Why [`tokio_util::sync::CancellationToken`] and not a `Notify` or a
50//! waker registration**: it is level-triggered. The `select!` may create and
51//! drop a fresh `cancelled()` future on every iteration with no lost-wakeup
52//! window and no re-registration, whereas a `Notify`-based wheel would need
53//! a global-mutex round trip per read chunk to re-register.
54//! `if head_release.is_some()` matters for the same reason from the other
55//! side: tokio does not evaluate a branch's async expression when its
56//! precondition is false, so an empty queue costs one `Option::is_some()`.
57//!
58//! # Cancellation
59//!
60//! Every wait in this module is a `select!` **branch**, never an arm body:
61//! an arm body is not preemptible, so an inline sleep there starves the
62//! cancel branch and a `Hold` on a gate nobody releases would pin session
63//! teardown for up to [`EgressConfig::max_hold`] (30 s by default). That
64//! applies to the two release-honouring drains as much as to the pipe loop,
65//! which is why [`drain_honouring_release_times`] is written as a `biased`
66//! race against cancellation with
67//! [`PendingQueue::drain_ignoring_release_times`] as its fallback.
68//!
69//! # Ordering is the queue's; readiness is the unit's
70//!
71//! The two are **separate properties**, and conflating them is what the
72//! first shape of this module got wrong.
73//!
74//! * *Ordering* belongs to the deque. [`PendingQueue::pop_next_due`] only
75//! ever looks at the **front**, so nothing can be written before
76//! everything ahead of it has been. That invariant needs no arithmetic
77//! at all, and it is absolute: anything else silently corrupts MoQT
78//! stream semantics.
79//! * *Readiness* belongs to the unit. A [`Pending`] is due when
80//! [`Pending::due_at`] has passed **or** its own `Hold` gate is open —
81//! see [`Pending::is_due`]. Nothing that happens to a unit ahead of it
82//! may rewrite that.
83//!
84//! So [`PendingQueue::push`] does **not** clamp `due_at` against the tail.
85//! An earlier shape did, and it turned a `Hold`'s
86//! [`EgressConfig::max_hold`] ceiling into a *deadline conferred on every
87//! successor*: releasing the gate woke only the unit carrying it, and
88//! objects queued behind — including the stream's FIN, which waits for the
89//! queue to drain — sat until `max_hold` (30 s by default). Worse, a
90//! `Delay`'s own deadline was overwritten by the ceiling of a `Hold` in
91//! front of it, so the deadline `Action::Delay` documents ("`arrived_at +
92//! by`, a deadline, not a spacing") was not the one the engine honoured.
93//!
94//! What the clamp *was* for is still wanted, and it survives as a separate
95//! field: [`Pending::expected_at`] is `max(due_at, tail.expected_at, now)`,
96//! the queue's estimate of when the unit will actually reach the wire given
97//! everything ahead of it. It is what [`Push::release_at`] reports as
98//! `Effect::Queued { release_at }` and what [`PendingQueue::record_release`]
99//! measures lateness against, so a unit queued behind a 300 ms delay is not
100//! logged as a 300 ms timer error. It never gates anything.
101//!
102//! # Bound
103//!
104//! [`PendingQueue::accepts_more`] is `false` once the queue holds
105//! [`EgressConfig::max_pending_bytes`] (1 MiB default) **or**
106//! [`MAX_PENDING_UNITS`] units. The pipe loop stops calling `recv.read` on
107//! `false`, so the queue cannot grow past that plus one in-flight read
108//! chunk. **It does not stop observing the source**: it swaps the read for
109//! `RecvStream::received_reset`, which consumes nothing, so the bound here
110//! is unchanged and a peer's `RESET_STREAM` is still seen at once. The
111//! unit cap is not redundant with the byte cap: an elided unit
112//! holds an ordering slot and carries zero bytes, so a hook that elides
113//! every object behind a never-released `Hold` would otherwise grow the
114//! deque without bound at zero bytes. Crossing either cap is reported once
115//! per stream — see [`Push::entered_backpressure`].
116
117use std::collections::VecDeque;
118use std::future::Future;
119use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
120use std::sync::{Arc, OnceLock};
121use std::time::{Duration, Instant};
122
123use bytes::Bytes;
124use tokio::sync::Notify;
125use tokio_util::sync::CancellationToken;
126
127use moqtap_client::transport::{SendStream, TransportError};
128
129use crate::action::{EgressConfig, Gate};
130use crate::error::ProxyError;
131use crate::instrument::Recorder;
132use crate::release_timer::{self, Deadline};
133use crate::shape::{Acquire, Class, Expiry, QueueDepth, Scheduler, ShapeRecorder};
134use crate::types::ProxySide;
135
136/// Hard cap on queued units, independent of the byte budget.
137///
138/// A [`Pending`] is on the order of 64 bytes, so this is roughly the same
139/// order of bookkeeping as the default 1 MiB byte budget, and it is far
140/// more objects than any realistic hold window. It exists because
141/// [`Item::Elided`] units carry zero bytes: without it, eliding every
142/// object behind a never-released `Hold` grows the deque without bound
143/// while `queued_bytes()` stays at zero.
144pub(crate) const MAX_PENDING_UNITS: usize = 8192;
145
146/// What writing to a destination can fail with.
147///
148/// Deliberately narrower than [`ProxyError`]: this module hands bytes to a
149/// transport and abandons streams, and those are the only two ways it can
150/// fail. Naming them separately is what lets a sink be implemented without
151/// naming a crate-wide error enum whose listener, TLS, certificate, qlog and
152/// draft-admission cases a write can never produce — and the recording sink
153/// below is the second implementor, so "without naming it" is not
154/// hypothetical.
155///
156/// Nothing a caller sees changes: the conversion below is applied by `?` at
157/// every boundary this module is called across, and it produces the same
158/// [`ProxyError`] the failure produced before.
159///
160/// One case, and the enum is still the right shape. Every failure reaching
161/// here today comes back from the transport, which is a fact about the two
162/// operations rather than about the one production implementor — so a second
163/// case would be a second *kind of failure*, not a second sink, and would want
164/// a name of its own. The recording sink below models a refused write as a
165/// refused write for the same reason: a test sink inventing a failure the
166/// transport cannot produce would be testing a path nothing takes.
167#[derive(Debug, thiserror::Error)]
168pub(crate) enum EgressError {
169 /// The transport refused a write or a reset.
170 #[error("transport error: {0}")]
171 Transport(#[from] TransportError),
172}
173
174impl From<EgressError> for ProxyError {
175 fn from(error: EgressError) -> Self {
176 let EgressError::Transport(source) = error;
177 ProxyError::Transport(source)
178 }
179}
180
181// ── The destination ─────────────────────────────────────────────────
182
183/// The write half this module drives.
184///
185/// A trait rather than a bare `&mut SendStream` for exactly one reason:
186/// the ordering, clamping, terminal and cancellation behaviour below is
187/// then testable in-process against a recording sink, with no QUIC
188/// handshake and no timing. The only production implementor is
189/// [`SendStream`].
190///
191/// `write_all` returns `impl Future + Send` rather than being an
192/// `async fn`: the forwarding tasks are spawned onto a multi-threaded
193/// runtime, so every future this module builds has to be `Send`, and an
194/// `async fn` in a trait carries no such bound.
195pub(crate) trait EgressSink {
196 /// Hand `buf` to the transport.
197 fn write_all(&mut self, buf: &[u8]) -> impl Future<Output = Result<(), EgressError>> + Send;
198 /// Abandon the stream with `code`, discarding anything still in the
199 /// local send buffer.
200 fn reset(&mut self, code: u64) -> Result<(), EgressError>;
201}
202
203impl EgressSink for SendStream {
204 // `async fn` here still satisfies the trait's `impl Future + Send`:
205 // the compiler checks the Send-ness of the generated future against
206 // the bound, which is the whole reason the bound is on the trait.
207 // Inherent methods win method resolution, so the call below is the
208 // transport's own `write_all`, not this trait's.
209 async fn write_all(&mut self, buf: &[u8]) -> Result<(), EgressError> {
210 SendStream::write_all(self, buf).await.map_err(EgressError::from)
211 }
212
213 fn reset(&mut self, code: u64) -> Result<(), EgressError> {
214 SendStream::reset(self, code).map_err(EgressError::from)
215 }
216}
217
218// ── Release handles ─────────────────────────────────────────────────
219
220/// What the head of a queue is waiting for. Owned, cheap to clone.
221///
222/// Cloning is one `Arc` clone plus one or two token clones. Nothing here
223/// borrows the queue, which is what keeps [`wait_release`] a free function
224/// usable from both the pipe loop's release branch and the drains.
225#[derive(Clone, Debug)]
226pub(crate) struct Release {
227 /// The wheel's handle for `release_at`. Registration happened once, in
228 /// [`PendingQueue::arm_head`]; awaiting this registers nothing.
229 deadline: Deadline,
230 /// A `Hold`'s gate, when the head carries one.
231 gate: Option<Gate>,
232}
233
234// Read only by this module's own tests today — `session.rs` awaits a
235// `Release` through `wait_release` and never inspects it, and `exec.rs`
236// never sees one. Kept because it is part of this type's deliberate read
237// surface and because the fast-path test asserts on the deadline's token
238// directly; the allow is on the one item rather than on the file, so
239// anything that becomes dead later is still a build failure.
240#[allow(dead_code)]
241impl Release {
242 /// The wheel handle, for callers that want to observe it directly.
243 pub(crate) fn deadline(&self) -> &Deadline {
244 &self.deadline
245 }
246}
247
248/// Resolve when the head unit may be written, or when the session is torn
249/// down — whichever comes first.
250///
251/// Races three level-triggered tokens: the wheel's deadline, the unit's
252/// `Hold` gate (a released gate makes a unit due early — see
253/// [`Pending::is_due`]) and session cancellation. It **registers nothing**:
254/// every one of the three is a `CancellationToken`, so creating and
255/// dropping the futures on each `select!` iteration is free of lost-wakeup
256/// races.
257///
258/// `None` parks until cancellation. Callers guard the branch with
259/// `if head_release.is_some()`, so that path is a safety net rather than a
260/// normal outcome: it must not resolve immediately, or a caller that forgot
261/// the guard would spin the loop hot on an empty queue.
262pub(crate) async fn wait_release(release: Option<Release>, cancel: &CancellationToken) {
263 let Some(release) = release else {
264 cancel.cancelled().await;
265 return;
266 };
267 match release.gate {
268 Some(gate) => {
269 tokio::select! {
270 () = release.deadline.token().cancelled() => {}
271 () = gate.wait() => {}
272 () = cancel.cancelled() => {}
273 }
274 }
275 None => {
276 tokio::select! {
277 () = release.deadline.token().cancelled() => {}
278 () = cancel.cancelled() => {}
279 }
280 }
281 }
282}
283
284// ── Queued units ────────────────────────────────────────────────────
285
286/// A positional stream ending.
287///
288/// Enqueued *behind* the current queue contents, so everything decided
289/// before it is written first. Nothing may be queued behind one.
290#[derive(Clone, Debug, PartialEq, Eq)]
291pub(crate) enum Terminal {
292 /// Write `prefix`, then `RESET_STREAM` with `code` —
293 /// [`crate::action::Action::Truncate`].
294 ///
295 /// The peer observes **at most** `prefix.len()` further bytes and may
296 /// observe none: quinn clears the receive assembler when the reset is
297 /// processed, and `reset()` discards the local send buffer too.
298 Truncate {
299 /// The prefix of the unit to write before resetting.
300 prefix: Bytes,
301 /// The application error code for the reset.
302 code: u64,
303 },
304 /// `RESET_STREAM` with `code` and nothing further —
305 /// [`crate::action::Action::ResetStream`].
306 Reset {
307 /// The application error code for the reset.
308 code: u64,
309 },
310}
311
312impl Terminal {
313 /// Bytes this terminal still owes the wire.
314 fn len(&self) -> usize {
315 match self {
316 Terminal::Truncate { prefix, .. } => prefix.len(),
317 Terminal::Reset { .. } => 0,
318 }
319 }
320}
321
322/// What a queued unit does when it is written.
323#[derive(Clone, Debug, PartialEq, Eq)]
324pub(crate) enum Item {
325 /// Write these bytes. Already whatever the action made of them —
326 /// verbatim for `Pass`, spliced for `ReplacePayload`, substituted for
327 /// `Replace`.
328 Write(Bytes),
329 /// Write nothing.
330 ///
331 /// An elided unit still takes a slot when the queue is non-empty: it
332 /// was decided at a point in the stream, and dropping it out of band
333 /// would let a later unit be written before an earlier one is.
334 Elided,
335 /// End the stream. See [`Terminal`].
336 Terminal(Terminal),
337}
338
339/// What the shaper attached to one queued unit, on a shaped stream.
340///
341/// `None` on every unit of every unshaped stream, which is what makes "a
342/// session with no `ShapeProfile` adds nothing to this path" a fact the type
343/// carries rather than a claim a reviewer checks.
344#[derive(Clone, Copy, Debug)]
345struct ShapeTag {
346 /// The row this unit's bytes are charged to, and the bucket it draws
347 /// from. Resolved per unit at admission — never per stream, because
348 /// `object_id` and `every_nth` are legal class selectors and a
349 /// stream-sticky class would evaluate them on object 0 alone.
350 class: Class,
351 /// The instant this unit goes out **anyway**: `arrived_at` plus the
352 /// profile's `max_hold`, or the engine's when the profile inherits it.
353 ///
354 /// Under the default `Expiry::Deliver` that is a clamp, not a drop — so
355 /// a starved class is late, never silently lossy, and "delivers zero
356 /// bytes" is always a statement about a sampling window.
357 expires_at: Instant,
358 /// Whether this unit has already been counted against
359 /// `starved_behind_other_class`. Once per unit, not once per wake.
360 starved_noted: bool,
361}
362
363/// One unit of traffic waiting for its release time.
364#[derive(Clone, Debug)]
365pub(crate) struct Pending {
366 /// The earliest instant **this** unit may be written, and the only instant
367 /// that gates it. Set once, by whoever built the unit, and never rewritten
368 /// by the queue — see the module doc's *Ordering is the queue's; readiness
369 /// is the unit's*.
370 due_at: Instant,
371 /// When the queue *expects* to write it, given everything ahead of it:
372 /// `max(due_at, tail.expected_at, now)` at push time.
373 ///
374 /// Reporting and lateness only. It gates nothing, and it is
375 /// deliberately an over-estimate for a unit queued behind a `Hold`,
376 /// whose gate may open long before its [`EgressConfig::max_hold`]
377 /// ceiling.
378 expected_at: Instant,
379 /// A `Hold`'s gate. Releasing it makes the unit due before `due_at`,
380 /// which for a held unit is the [`EgressConfig::max_hold`] ceiling.
381 gate: Option<Gate>,
382 /// What writing it does.
383 item: Item,
384 /// The shaper's tag, on a shaped stream only. Attached by
385 /// [`PendingQueue::push`] from the class the pipe loop most recently
386 /// resolved, so `exec`'s pushes carry it without `exec` naming a class.
387 shape: Option<ShapeTag>,
388}
389
390impl Pending {
391 /// Bytes to write no earlier than `due_at`.
392 pub(crate) fn bytes(raw: Bytes, due_at: Instant) -> Self {
393 Self { due_at, expected_at: due_at, gate: None, item: Item::Write(raw), shape: None }
394 }
395
396 /// An ordering placeholder that writes nothing.
397 pub(crate) fn elided(due_at: Instant) -> Self {
398 Self { due_at, expected_at: due_at, gate: None, item: Item::Elided, shape: None }
399 }
400
401 /// A positional terminal, due as soon as the queue ahead of it drains.
402 pub(crate) fn terminal(terminal: Terminal) -> Self {
403 let now = Instant::now();
404 Self {
405 due_at: now,
406 expected_at: now,
407 gate: None,
408 item: Item::Terminal(terminal),
409 shape: None,
410 }
411 }
412
413 /// Attach a `Hold`'s gate. `due_at` stays the `max_hold` ceiling.
414 #[must_use]
415 pub(crate) fn with_gate(mut self, gate: Gate) -> Self {
416 self.gate = Some(gate);
417 self
418 }
419
420 // The observers below are read by this module's tests and by nothing
421 // else: the queue answers `due_at` through `head_release()`, and a
422 // terminal announces itself through `Written::Terminated` rather than
423 // being asked. They stay as the deliberate read surface of a unit,
424 // with the allow scoped to them rather than to the file.
425 /// When this unit becomes due on its own account.
426 #[allow(dead_code)]
427 pub(crate) fn due_at(&self) -> Instant {
428 self.due_at
429 }
430
431 /// When the queue expected to write it. Never a gate; see the field.
432 #[allow(dead_code)]
433 pub(crate) fn expected_at(&self) -> Instant {
434 self.expected_at
435 }
436
437 /// What it does. Terminals end the stream.
438 #[allow(dead_code)]
439 pub(crate) fn item(&self) -> &Item {
440 &self.item
441 }
442
443 /// Whether this unit ends the stream.
444 #[allow(dead_code)]
445 pub(crate) fn is_terminal(&self) -> bool {
446 matches!(self.item, Item::Terminal(_))
447 }
448
449 /// Bytes this unit is holding on the queue's behalf.
450 pub(crate) fn len(&self) -> usize {
451 match &self.item {
452 Item::Write(raw) => raw.len(),
453 Item::Elided => 0,
454 Item::Terminal(t) => t.len(),
455 }
456 }
457
458 /// Whether this unit may be written at `now`, **on its own account**.
459 ///
460 /// Readiness only: whether anything is still queued ahead of it is the
461 /// deque's business, and [`PendingQueue::pop_next_due`] answers that by
462 /// looking at the front and nowhere else.
463 ///
464 /// A **released gate makes it due**, regardless of `due_at`. Without
465 /// that clause, releasing a `Hold` early wakes the release branch to
466 /// find nothing due and the loop spins hot until
467 /// [`EgressConfig::max_hold`] — one of the two ways this loop can spin
468 /// hot, and the one that is invisible from the call site.
469 ///
470 /// `expected_at` is deliberately **not** consulted: it is the queue's
471 /// estimate, and for a unit behind a `Hold` it is that hold's ceiling.
472 /// Gating on it is what stranded whole streams behind a released gate.
473 fn is_due(&self, now: Instant) -> bool {
474 self.due_at <= now || self.gate.as_ref().is_some_and(Gate::is_released)
475 }
476}
477
478/// What handing one unit to the transport did.
479#[derive(Clone, Copy, Debug, PartialEq, Eq)]
480pub(crate) enum Written {
481 /// This many bytes were handed to the transport.
482 Bytes(usize),
483 /// Nothing was written — an elided unit holding its ordering slot.
484 Nothing,
485 /// A terminal fired. The destination is reset; nothing further may be
486 /// written to it.
487 Terminated {
488 /// Bytes written ahead of the reset, on this unit.
489 forwarded: usize,
490 /// The code the stream was reset with.
491 code: u64,
492 },
493}
494
495/// Write one popped unit.
496///
497/// A failing prefix write on a [`Terminal::Truncate`] is propagated rather
498/// than swallowed: the caller mirrors it with `propagate_stop`, and the
499/// reset it would have sent is moot once the peer has stopped us. A failing
500/// `reset` is *not* propagated — it means the stream was already finished
501/// or reset, so there is nothing left to report and nothing left to do.
502pub(crate) async fn write_unit<S: EgressSink>(
503 unit: Pending,
504 send: &mut S,
505) -> Result<Written, EgressError> {
506 match unit.item {
507 Item::Write(raw) => {
508 send.write_all(&raw).await?;
509 Ok(Written::Bytes(raw.len()))
510 }
511 Item::Elided => Ok(Written::Nothing),
512 Item::Terminal(Terminal::Truncate { prefix, code }) => {
513 if !prefix.is_empty() {
514 send.write_all(&prefix).await?;
515 }
516 let _ = send.reset(code);
517 Ok(Written::Terminated { forwarded: prefix.len(), code })
518 }
519 Item::Terminal(Terminal::Reset { code }) => {
520 let _ = send.reset(code);
521 Ok(Written::Terminated { forwarded: 0, code })
522 }
523 }
524}
525
526// ── Delay and hold arithmetic ───────────────────────────────────────
527
528/// A resolved release deadline, and whether [`EgressConfig::max_hold`] cut
529/// it short.
530#[derive(Clone, Copy, Debug, PartialEq, Eq)]
531pub(crate) struct Deferral {
532 /// The deadline to put on the [`Pending`], before the queue's own
533 /// monotonic clamp.
534 pub(crate) release_at: Instant,
535 /// What was asked for.
536 pub(crate) requested: Duration,
537 /// What was applied.
538 pub(crate) applied: Duration,
539}
540
541impl Deferral {
542 /// Whether `max_hold` cut the request short — emit
543 /// `Impairment { HoldClamped { requested, applied } }`, once per
544 /// clamped unit.
545 pub(crate) fn was_clamped(&self) -> bool {
546 self.applied < self.requested
547 }
548}
549
550/// Resolve `Delay { by }` against [`EgressConfig::max_hold`].
551///
552/// A **deadline**, not a spacing: `arrived_at + by`. Two units arriving
553/// together with the same `by` are released together, not staggered.
554pub(crate) fn defer_by(arrived_at: Instant, by: Duration, config: &EgressConfig) -> Deferral {
555 let applied = by.min(config.max_hold);
556 Deferral { release_at: arrived_at + applied, requested: by, applied }
557}
558
559/// The ceiling on an [`crate::action::Action::Hold`]: a gate nobody
560/// releases is still written at `arrived_at + max_hold`, so a hook cannot
561/// pin a stream open forever.
562pub(crate) fn hold_ceiling(arrived_at: Instant, config: &EgressConfig) -> Instant {
563 arrived_at + config.max_hold
564}
565
566// ── The queue ───────────────────────────────────────────────────────
567
568/// What [`PendingQueue::push`] did.
569#[derive(Clone, Copy, Debug, PartialEq, Eq)]
570pub(crate) struct Push {
571 /// [`Pending::expected_at`] — the queue's estimate of when this unit
572 /// reaches the wire, which is the value to report as
573 /// `Effect::Queued { release_at }`. It accounts for everything already
574 /// queued, so it may be later than the action asked for, and for a unit
575 /// behind a `Hold` it is that hold's `max_hold` ceiling: an upper
576 /// bound, since a gate may open at any time before it.
577 ///
578 /// **Not** the unit's own deadline. That is [`Pending::due_at`], it is
579 /// the only thing that gates the unit, and `push` never rewrites it.
580 pub(crate) release_at: Instant,
581 /// `true` exactly once per stream: on the push after which
582 /// [`PendingQueue::accepts_more`] first turns `false`. Emit
583 /// `Impairment { EgressQueueFull { stream_id } }` on it, and only on
584 /// it — a queue that drains and fills again does not report twice.
585 pub(crate) entered_backpressure: bool,
586}
587
588/// How a drain ended.
589#[derive(Clone, Copy, Debug, PartialEq, Eq)]
590pub(crate) enum DrainOutcome {
591 /// Every queued unit was written.
592 Complete,
593 /// Cancellation won mid-drain. The remainder was flushed ignoring
594 /// release times; [`PendingQueue::unconfirmed_bytes`] afterwards is
595 /// what that fallback either could not write **or** wrote into a
596 /// transport that is being closed out from under it, and is what
597 /// `Impairment { QueuedBytesAtTeardown }` carries.
598 CancelledMidDrain,
599 /// A queued terminal fired. The destination is reset and nothing
600 /// further may be written to it.
601 Terminated {
602 /// Bytes written on the terminal unit itself, ahead of the reset.
603 forwarded: usize,
604 /// The code the stream was reset with.
605 code: u64,
606 },
607 /// A write failed. The failing unit and everything behind it are still
608 /// queued. Only the best-effort drain returns this; the
609 /// release-honouring drain propagates the error instead.
610 WriteFailed,
611 /// The teardown flush declined to write anything, because the session's
612 /// [`EgressGauge`] is discarding: a requested close ran out of drain
613 /// time, so what is left is reported rather than handed to a connection
614 /// that is about to be closed. The queue is untouched, so
615 /// [`PendingQueue::unconfirmed_bytes`] is exactly what was abandoned.
616 Discarded,
617}
618
619/// How many bytes one session's egress queues are holding, right now.
620///
621/// One gauge per session, shared by every [`PendingQueue`] the session
622/// builds. It exists for a single caller — a requested close, which has to
623/// know when there is nothing left to flush so it can stop waiting early
624/// — and it answers that question in the one unit that is worth waiting on.
625///
626/// # Why a byte gauge rather than a stream count
627///
628/// A session's live streams are already enumerated by
629/// [`StreamRegistry`](crate::control::StreamRegistry), and waiting for that
630/// to empty was the obvious alternative. It is the wrong signal twice over:
631/// a stream with an empty queue keeps its registration until its source
632/// FINs, so the wait would run to the deadline on a session that had
633/// nothing to flush at all; and a stream that ends while its queue is full
634/// drops its registration with the bytes still unwritten, so the wait would
635/// finish claiming a drain that lost data.
636///
637/// # Two-valued, and the second value is the whole bound
638///
639/// [`Self::wait_idle`] is the wait. [`Self::begin_discarding`] is what
640/// happens when the wait runs out: it puts every queue in the session into
641/// a mode where [`PendingQueue::drain_ignoring_release_times`] writes
642/// nothing at all.
643///
644/// That looks backwards next to the rest of this module, whose teardown rule is
645/// *delivered late beats lost silently*. The difference is that an ordinary
646/// teardown is not a deadline anybody set — the flush is the best remaining
647/// chance those bytes have. A close that has already been given its drain
648/// window and spent it is a deadline somebody set, and flushing past it makes
649/// the outcome *unreportable*: a `write_all` into quinn returns `Ok` as soon as
650/// the bytes are buffered and `Connection::close` discards that buffer, so the
651/// bytes are neither confirmably delivered nor confirmably lost, and no count
652/// of them adds up. Declining to write keeps the arithmetic exact — what the
653/// peer received plus what
654/// [`ImpairmentKind::QueuedBytesAtTeardown`](crate::event::ImpairmentKind::QueuedBytesAtTeardown)
655/// names equals what was queued when the close was requested — which is the
656/// property a gate can actually be written against.
657#[derive(Debug)]
658pub(crate) struct EgressGauge {
659 /// Bytes queued across every live [`PendingQueue`] of this session.
660 queued: AtomicUsize,
661 /// Pulsed whenever `queued` reaches zero.
662 idle: Notify,
663 /// Whether teardown flushes are declined — see the type's own doc.
664 discarding: AtomicBool,
665}
666
667impl EgressGauge {
668 /// A gauge reading zero, flushing normally.
669 pub(crate) fn new() -> Arc<Self> {
670 Arc::new(Self {
671 queued: AtomicUsize::new(0),
672 idle: Notify::new(),
673 discarding: AtomicBool::new(false),
674 })
675 }
676
677 /// Bytes queued across the session at this instant.
678 pub(crate) fn queued(&self) -> usize {
679 self.queued.load(Ordering::Acquire)
680 }
681
682 /// Whether teardown flushes are being declined.
683 pub(crate) fn is_discarding(&self) -> bool {
684 self.discarding.load(Ordering::Acquire)
685 }
686
687 /// Decline every later teardown flush on this session.
688 ///
689 /// One-way: nothing turns it off again, because the only caller is a
690 /// close whose drain window has expired and a session in that state is
691 /// on its way down.
692 pub(crate) fn begin_discarding(&self) {
693 self.discarding.store(true, Ordering::Release);
694 }
695
696 /// Charge `bytes` to the session total.
697 fn add(&self, bytes: usize) {
698 if bytes > 0 {
699 self.queued.fetch_add(bytes, Ordering::AcqRel);
700 }
701 }
702
703 /// Credit `bytes` back, waking a waiter when the session reaches zero.
704 fn sub(&self, bytes: usize) {
705 if bytes == 0 {
706 return;
707 }
708 // `saturating` in effect: a queue never credits back more than it
709 // charged, but an underflow here would wrap to a total that never
710 // reaches zero and would turn every later close into a full-length
711 // wait, so the arithmetic is written not to be able to.
712 let before = self
713 .queued
714 .fetch_update(Ordering::AcqRel, Ordering::Acquire, |n| Some(n.saturating_sub(bytes)))
715 .unwrap_or(0);
716 if before.saturating_sub(bytes) == 0 {
717 self.idle.notify_waiters();
718 }
719 }
720
721 /// Wait for the session's queues to empty, for at most `timeout`.
722 ///
723 /// Answers the bytes still queued when it returns: `0` when everything
724 /// drained, and the residue when the deadline won. That residue is the
725 /// figure a caller acts on — it is what will be abandoned — and it is
726 /// deliberately a byte count and not "did it time out", because the
727 /// two differ: a queue that empties on the last poll before the
728 /// deadline drained, and one that reports zero after the deadline
729 /// drained too.
730 ///
731 /// The registration is armed **before** the count is re-read, so a
732 /// queue that empties between the two is not a lost wakeup that costs
733 /// the caller its whole timeout.
734 pub(crate) async fn wait_idle(&self, timeout: Duration) -> usize {
735 let deadline = tokio::time::sleep(timeout);
736 tokio::pin!(deadline);
737 loop {
738 if self.queued() == 0 {
739 return 0;
740 }
741 let notified = self.idle.notified();
742 tokio::pin!(notified);
743 notified.as_mut().enable();
744 if self.queued() == 0 {
745 return 0;
746 }
747 tokio::select! {
748 () = &mut notified => {}
749 () = &mut deadline => return self.queued(),
750 }
751 }
752 }
753}
754
755/// A per-stream FIFO of units waiting for their release times.
756///
757/// Live as a local variable in `pipe_data_framed` / `pipe_control_mutating`
758/// for the lifetime of one stream direction.
759#[derive(Debug)]
760pub(crate) struct PendingQueue {
761 q: VecDeque<Pending>,
762 /// The wheel handle for `q.front()`. Invariant:
763 /// `head_deadline.is_some() == !q.is_empty()`.
764 head_deadline: Option<Deadline>,
765 queued_bytes: usize,
766 /// Bytes a teardown drain handed to the transport, which nothing can
767 /// confirm ever left it. See [`Self::unconfirmed_bytes`].
768 flushed_unconfirmed: usize,
769 /// The session-wide byte gauge this queue reports into, or `None` for a
770 /// queue built outside a session — every `#[cfg(test)]` queue in this
771 /// module, which has no session to report to.
772 ///
773 /// Mirrored rather than derived: nothing can walk a session's queues,
774 /// because each one is a local variable in the pipe function that owns
775 /// its stream. Every write to `queued_bytes` goes through
776 /// [`Self::charge`] or [`Self::credit`] so the mirror cannot drift from
777 /// the value it mirrors.
778 gauge: Option<Arc<EgressGauge>>,
779 config: EgressConfig,
780 counters: Arc<Recorder>,
781 /// Report-once latch for [`Push::entered_backpressure`].
782 backpressure_reported: bool,
783 /// The shaping profile's per-stream depth, when this stream's session
784 /// has one **and** its overflow policy is `Block`.
785 ///
786 /// `None` on every unshaped stream and under every non-blocking
787 /// overflow: `DropTail` needs the read branch to stay enabled or
788 /// nothing arrives to be dropped, and `ResetStream` is decided on the
789 /// arriving unit. See [`Self::accepts_more`].
790 shape_depth: Option<QueueDepth>,
791 /// The session's shaper, on a shaped **data** stream and nowhere else.
792 /// `None` on every unshaped stream and on both control pipes, which is what
793 /// makes *the control pipes are never shaped* structural: a control queue
794 /// has no scheduler to ask, so no control frame can reach a bucket even by
795 /// accident.
796 shaper: Option<Arc<Scheduler>>,
797 /// Where a shaped release records what it did. Carried beside the
798 /// scheduler rather than reached through it because the recorder is
799 /// always constructed and the scheduler is not.
800 shape_stats: Option<Arc<ShapeRecorder>>,
801 /// The side this stream's traffic **arrived** on, for the per-leg and
802 /// per-direction figures the recorder keeps.
803 ///
804 /// `Option`, and deliberately not a defaulted side: one recorder serves
805 /// both legs of a session, so a queue that guessed would charge real
806 /// bytes to the wrong side and the figure would still look plausible.
807 /// `Some` exactly when `shape_stats` is — the same builder sets both —
808 /// so every reader below takes them together.
809 ///
810 /// A side rather than a direction because the recorder needs the leg
811 /// too, and this queue is the only thing that knows which stream it is
812 /// serving. It stays the *arrival* side even for the release figures the
813 /// proxy charges to the leg a unit leaves by: the turn-around belongs to
814 /// the recorder, which is one place, rather than to every queue, which
815 /// is one per stream.
816 shape_side: Option<ProxySide>,
817 /// `max_hold` for this stream's shaped units: the profile's, or the
818 /// engine's when the profile inherits it.
819 shape_hold: Duration,
820 /// The class the pipe loop resolved for the unit it is currently
821 /// handling. Read — not taken — by [`Self::push`], so every push made
822 /// while one framed unit is being executed carries the same class, and
823 /// the next unit overwrites it.
824 unit_class: Class,
825 /// The gate the head is parked on because a [`Discipline`] is holding
826 /// its class back. Merged into [`Self::head_release`] so `wait_release`
827 /// races it exactly as it races a `Hold`'s gate.
828 ///
829 /// [`Discipline`]: crate::shape::Discipline
830 shape_gate: Option<Gate>,
831 /// Whether the shaper is currently holding the head back.
832 ///
833 /// It decides which gate [`Self::head_release`] reports, and that is a
834 /// **spin guard**, not bookkeeping. A `Hold` whose gate has been released
835 /// makes its unit due (`Pending::is_due`) and the gate stays released
836 /// forever; if the queue kept reporting it while the shaper was refusing
837 /// the same unit, `wait_release` would resolve on every `select!`
838 /// iteration and the pipe loop would spin hot until the bucket refilled.
839 /// Measured before this field existed: a barrier-held object on a 0-bps
840 /// class pinned a core for the whole five-second clamp and starved the
841 /// test's own task alongside it.
842 ///
843 /// Once a unit is due, the only things that may wake its stream are the
844 /// deadline the scheduler armed and the gate a discipline handed back.
845 shape_parked: bool,
846 /// The class this queue has declared demand for with the scheduler.
847 /// Exactly one declaration is outstanding at a time, and [`Drop`] is what
848 /// guarantees it is withdrawn.
849 shape_demand: Option<Class>,
850 /// Edge latch for `tokens_exhausted_episodes` — armed when a release is
851 /// refused for want of tokens, re-armed when one is granted.
852 tokens_dry: bool,
853 /// The index the next `starved_behind_other_class` sweep starts from, so
854 /// each queued unit is examined once over the queue's whole life rather
855 /// than once per wake.
856 starved_from: usize,
857 /// What the last [`Self::pop_next_due`] wants reported, for a caller that
858 /// has an `exec::Reporter` and an event sink. Taken, never accumulated:
859 /// at most one unit is released per call.
860 shape_report: Option<ShapeReport>,
861}
862
863/// What a shaped release wants the pipe loop to report.
864///
865/// The queue owns the decision and the caller owns the reporting: `egress`
866/// has no `Reporter`, no `ProxyObserver` and no session id, and giving it
867/// one would put the whole event surface behind a type whose job is a deque.
868#[derive(Clone, Copy, Debug, PartialEq, Eq)]
869pub(crate) enum ShapeReport {
870 /// A unit went out at its `max_hold` clamp rather than when its bucket
871 /// would have allowed it. Report `Impairment{HoldClamped}`, whose
872 /// cardinality is already *once per clamped unit*.
873 Clamped {
874 /// What the bucket asked for, or `None` when it named no instant at
875 /// all — a zero rate, or a unit larger than `burst_bytes`, where the
876 /// wait is genuinely unbounded.
877 ///
878 /// `None` rather than a sentinel duration: a finite figure here
879 /// would be a fabrication, and the largest representable one renders
880 /// as an eighteen-quintillion-second request that reads as an
881 /// encoding fault rather than as "unbounded".
882 requested: Option<Duration>,
883 /// What was applied: the clamp.
884 applied: Duration,
885 },
886 /// A unit outlived its clamp under `Expiry::ResetStream`, and the queue
887 /// has replaced everything with the reset. Report exactly one
888 /// `Shaped{Expired}`.
889 Expired,
890 /// The head could not be granted because its class's `burst_bytes` is
891 /// smaller than the head itself, so that class's configured rate can
892 /// never pace anything. Report `Impairment{ShapeBurstBelowUnit}`.
893 ///
894 /// The queue produces this on **every** refusal of such a unit and does
895 /// not deduplicate: the burst is a property of the session's profile,
896 /// not of this stream, so the report-once state belongs to the scheduler
897 /// and the caller claims it there through
898 /// `Scheduler::claim_burst_report`. A per-queue latch would report once
899 /// per stream for a fault that is the same fault on every stream.
900 ///
901 /// Carries the class rather than its name so this type stays `Copy`;
902 /// resolving a name here would allocate on a release path for an event
903 /// that is emitted once.
904 BurstBelowUnit {
905 /// The class whose rate is not being applied.
906 class: Class,
907 /// Its bucket cap, as configured.
908 burst_bytes: u64,
909 /// The wire size of the unit the cap could not cover.
910 unit_bytes: u64,
911 },
912}
913
914impl PendingQueue {
915 /// An empty queue. Allocates nothing until the first [`Self::push`].
916 pub(crate) fn new(config: EgressConfig, counters: Arc<Recorder>) -> Self {
917 Self {
918 q: VecDeque::new(),
919 head_deadline: None,
920 queued_bytes: 0,
921 flushed_unconfirmed: 0,
922 gauge: None,
923 shape_hold: config.max_hold,
924 config,
925 counters,
926 backpressure_reported: false,
927 shape_depth: None,
928 shaper: None,
929 shape_stats: None,
930 shape_side: None,
931 unit_class: Class::Unshapeable,
932 shape_gate: None,
933 shape_parked: false,
934 shape_demand: None,
935 tokens_dry: false,
936 starved_from: 1,
937 shape_report: None,
938 }
939 }
940
941 /// Report this queue's depth into its session's [`EgressGauge`].
942 ///
943 /// A builder for the same reason [`Self::with_shape_depth`] is: the
944 /// tests in this module build queues that belong to no session, and a
945 /// required parameter would make each of them invent one.
946 /// Installed by every pipe function, on every stream, shaped or not —
947 /// unlike a scheduler, which only the framed data pipe installs. A gauge
948 /// that only some queues reported into would answer *the session has
949 /// nothing left to flush* while a control stream still held a deferred
950 /// SUBSCRIBE_OK.
951 pub(crate) fn with_gauge(mut self, gauge: Arc<EgressGauge>) -> Self {
952 // A fresh queue is empty, so there is nothing to charge here. If
953 // that ever stops being true this has to charge `queued_bytes`.
954 debug_assert_eq!(self.queued_bytes, 0, "a queue takes its gauge before it takes bytes");
955 self.gauge = Some(gauge);
956 self
957 }
958
959 /// Add `bytes` to this queue's depth and to its session's gauge.
960 fn charge(&mut self, bytes: usize) {
961 self.queued_bytes = self.queued_bytes.saturating_add(bytes);
962 if let Some(gauge) = &self.gauge {
963 gauge.add(bytes);
964 }
965 }
966
967 /// Take `bytes` off this queue's depth and off its session's gauge.
968 fn credit(&mut self, bytes: usize) {
969 let bytes = bytes.min(self.queued_bytes);
970 self.queued_bytes -= bytes;
971 if let Some(gauge) = &self.gauge {
972 gauge.sub(bytes);
973 }
974 }
975
976 /// Whether this queue's session has stopped flushing at teardown.
977 ///
978 /// `false` for a queue with no gauge, which is every queue built
979 /// outside a session: nothing can have asked such a queue to stop.
980 fn discarding(&self) -> bool {
981 self.gauge.as_ref().is_some_and(|g| g.is_discarding())
982 }
983
984 /// Carry a shaping profile's per-stream depth into [`Self::accepts_more`].
985 ///
986 /// A builder rather than a fourth parameter on [`Self::new`], because
987 /// every caller but one passes `None` and the shape of the exception is
988 /// the point: exactly one construction site — the framed data pipe —
989 /// has a profile to install, and the ~25 unit tests below construct
990 /// queues that never had one.
991 ///
992 /// **Installed only under `Overflow::Block`**, which the caller
993 /// decides; this type deliberately does not name `Overflow`. What it
994 /// does is enforce a depth, and a depth that reaches it is one the
995 /// session has already decided should stop reads.
996 ///
997 /// Plumbed *into* `accepts_more` rather than checked beside it because
998 /// the once-per-stream backpressure latch is computed inside
999 /// [`Self::push`] from `accepts_more()`: a stream full by
1000 /// `QueueConfig::depth_objects` but under
1001 /// [`EgressConfig::max_pending_bytes`] would otherwise never trip it,
1002 /// and `Impairment{EgressQueueFull}` would have no producer on a shaped
1003 /// stream at all.
1004 pub(crate) fn with_shape_depth(mut self, depth: Option<QueueDepth>) -> Self {
1005 self.shape_depth = depth;
1006 self
1007 }
1008
1009 /// Install the session's shaper, making this a **paced** queue.
1010 /// A builder for the same reason [`Self::with_shape_depth`] is: exactly one
1011 /// construction site — the framed data pipe — has a shaper, and the ~25
1012 /// unit tests below build queues that never had one. Both control pipes and
1013 /// `pipe_data_passthrough` call [`Self::new`] and stop there, which is what
1014 /// makes *the control pipes are never shaped* a property of the type graph
1015 /// rather than of a review.
1016 ///
1017 /// What it changes, and only this: [`Self::pop_next_due`] asks
1018 /// `Scheduler::acquire` before it yields a shaped unit, and [`Self::push`]
1019 /// tags what it queues. Everything else — ordering, `due_at`, `Hold`
1020 /// gates, the two drains, `head_release`'s signature — is untouched.
1021 ///
1022 /// `side` travels with the recorder rather than beside it because the
1023 /// recorder is session-scoped and this queue is not: one `ShapeRecorder`
1024 /// is shared by the forwarding tasks of both legs, and the figures it
1025 /// keeps are per leg and per direction. Naming it here — at the one site
1026 /// that has both a profile and a stream — is what makes a downlink stall
1027 /// unattributable to the uplink.
1028 ///
1029 /// It is the side this stream's traffic **arrives** on, which is the
1030 /// only side a forwarding task is ever handed. What a release figure
1031 /// does with it — charge the leg the unit leaves by — is the recorder's
1032 /// business, not this queue's.
1033 pub(crate) fn with_shaper(
1034 mut self,
1035 shaper: Option<Arc<Scheduler>>,
1036 stats: Arc<ShapeRecorder>,
1037 side: ProxySide,
1038 ) -> Self {
1039 if let Some(shaper) = shaper {
1040 self.shape_hold = shaper.max_hold().unwrap_or(self.config.max_hold);
1041 self.shaper = Some(shaper);
1042 self.shape_stats = Some(stats);
1043 self.shape_side = Some(side);
1044 }
1045 self
1046 }
1047
1048 /// Whether this queue paces what it holds.
1049 ///
1050 /// Read by `exec::Engine::queue_is_busy`, which is what routes a shaped
1051 /// stream's units *through* the queue instead of writing them inline: a
1052 /// unit written straight to the transport never reaches
1053 /// [`Self::pop_next_due`], so on a shaped stream the fast path is exactly
1054 /// the path that cannot be paced.
1055 pub(crate) fn is_shaped(&self) -> bool {
1056 self.shaper.is_some()
1057 }
1058
1059 /// Name the class of the unit the pipe loop is about to queue.
1060 ///
1061 /// Set once per framed unit, immediately after classification, and read
1062 /// by every [`Self::push`] until the next unit overwrites it. That is why
1063 /// `exec` can queue a `Delay`'d object without ever naming a class: the
1064 /// tag is the loop's, the push is `exec`'s, and neither has to know about
1065 /// the other.
1066 ///
1067 /// A no-op on an unshaped queue, where nothing reads it.
1068 pub(crate) fn tag_unit(&mut self, class: Class) {
1069 self.unit_class = class;
1070 }
1071
1072 /// Whatever the last [`Self::pop_next_due`] wants the caller to report.
1073 pub(crate) fn take_shape_report(&mut self) -> Option<ShapeReport> {
1074 self.shape_report.take()
1075 }
1076
1077 /// Whether anything is waiting.
1078 pub(crate) fn is_empty(&self) -> bool {
1079 self.q.is_empty()
1080 }
1081
1082 /// How many units are waiting.
1083 ///
1084 /// Read by this module's tests and by `exec.rs`'s, which assert
1085 /// `DeferredEffects::len() == PendingQueue::len()`, and on the data
1086 /// path by shaping admission, which needs the object count the
1087 /// arriving unit would be the `n + 1`-th of. The pipe loops themselves
1088 /// ask [`Self::accepts_more`] and [`Self::is_empty`].
1089 pub(crate) fn len(&self) -> usize {
1090 self.q.len()
1091 }
1092
1093 /// Bytes the queue is holding, right now, still unwritten.
1094 ///
1095 /// Not the teardown report — see [`Self::unconfirmed_bytes`]. This is
1096 /// the live figure the byte budget and the tests are written against.
1097 pub(crate) fn queued_bytes(&self) -> usize {
1098 self.queued_bytes
1099 }
1100
1101 /// Bytes this queue cannot vouch for: everything a teardown drain
1102 /// handed to the transport, **plus** everything still queued.
1103 ///
1104 /// What `Impairment { QueuedBytesAtTeardown }` carries, and the reason
1105 /// it is not simply [`Self::queued_bytes`]. A cancelled session's pipe
1106 /// task races `run_with_transport`'s `client.close()` / `relay.close()`:
1107 /// [`Self::drain_ignoring_release_times`] can hand a 256 KiB
1108 /// object to quinn, have `write_all` return `Ok` because quinn buffered
1109 /// it, and have the connection torn down before a byte of it is on the
1110 /// wire. Reporting only the residue reports **zero** there, and the
1111 /// object is gone with no event at all — the one failure class this
1112 /// engine forbids outright: traffic is never silently gone.
1113 ///
1114 /// So a teardown flush counts as *unconfirmed*, not as delivered. The
1115 /// event says exactly this already: "they were flushed best-effort;
1116 /// `bytes` may not have reached the peer". A peer that did receive them
1117 /// gets a spurious impairment, which is the right side to err on: an
1118 /// over-report is a diagnosable nuisance, a silent loss is not.
1119 ///
1120 /// Zero on every path that is not a teardown, so a stream that drains
1121 /// at its release times reports nothing.
1122 pub(crate) fn unconfirmed_bytes(&self) -> usize {
1123 self.flushed_unconfirmed.saturating_add(self.queued_bytes)
1124 }
1125
1126 /// The engine knobs this queue was built with.
1127 pub(crate) fn config(&self) -> &EgressConfig {
1128 &self.config
1129 }
1130
1131 /// Whether the source may still be **read from**.
1132 ///
1133 /// `false` once the queue holds [`EgressConfig::max_pending_bytes`] or
1134 /// [`MAX_PENDING_UNITS`] units — the point at which `Delay` stops being
1135 /// a latency shift and becomes backpressure — **or** once it reaches a
1136 /// shaping profile's own depth, when one was installed by
1137 /// [`Self::with_shape_depth`]. A `bool`, deliberately: it is hoisted
1138 /// out of the `select!` and holds no borrow.
1139 ///
1140 /// **Read from, not observed.** This gates `recv.read`, which is the
1141 /// only call that consumes bytes and therefore the only one that grants
1142 /// the peer flow-control credit. It does *not* gate whether the pipe
1143 /// loop watches its source: `false` swaps the read for
1144 /// `RecvStream::received_reset`, so a peer's `RESET_STREAM` is mirrored
1145 /// at once however full this queue is. Gating the observation on this
1146 /// was the defect — a dry bucket under `Overflow::Block` held the read
1147 /// shut for `max_hold` (30 s by default) and the reset went unseen for
1148 /// the whole of it.
1149 ///
1150 /// Both terms are **per stream**, and both imply a non-empty deque, so
1151 /// `head_release()` stays `Some` and the release branch stays live. A
1152 /// shared cross-class budget must never reach this function: an empty
1153 /// own queue with a shared budget exhausted disables the read branch
1154 /// *and* the release branch, and the stream deadlocks until session
1155 /// teardown.
1156 pub(crate) fn accepts_more(&self) -> bool {
1157 let within_engine_budget =
1158 self.queued_bytes < self.config.max_pending_bytes && self.q.len() < MAX_PENDING_UNITS;
1159 let within_shape_depth = match self.shape_depth {
1160 Some(depth) => self.queued_bytes < depth.bytes && self.q.len() < depth.objects,
1161 None => true,
1162 };
1163 within_engine_budget && within_shape_depth
1164 }
1165
1166 /// What the head is waiting for, as an owned clone — see [`Release`].
1167 ///
1168 /// **`&self`, and it debits nothing.** It reports the deadline the
1169 /// scheduler last computed and the gate it last parked on; the decision
1170 /// itself belongs to [`Self::pop_next_due`], which is `&mut self` and
1171 /// runs exactly once per released unit. This function is called once
1172 /// per `select!` iteration from a hoisted local, so a debit here would
1173 /// be charged per read-wake rather than per byte written.
1174 ///
1175 /// **The FIN drain *is* on `pop_next_due`'s path**, and this sentence
1176 /// used to deny it. [`drain_honouring_release_times`] —
1177 /// the read arm's `None` branch — loops `pop_next_due`, debits the
1178 /// bucket, and applies the clamp and the expiry, which is why it takes
1179 /// an `on_shape` callback and reports them. The control pipes really
1180 /// are exempt, but by *construction* rather than by call graph: they
1181 /// install no [`Scheduler`] at all, so
1182 /// `shaping_grants_head` short-circuits on `shaper: None` and no
1183 /// control frame can reach a bucket even by accident. The only drain
1184 /// that is genuinely off the seam is
1185 /// [`Self::drain_ignoring_release_times`], which pops the front
1186 /// directly and consults nothing.
1187 pub(crate) fn head_release(&self) -> Option<Release> {
1188 let head = self.q.front()?;
1189 debug_assert!(
1190 self.head_deadline.is_some(),
1191 "a non-empty queue always has an armed head deadline",
1192 );
1193 Some(Release {
1194 deadline: self.head_deadline.clone()?,
1195 // Exactly one of the two, never both, and the choice is the spin
1196 // guard `shape_parked` documents: while the shaper is holding a
1197 // unit the shaper's gate is the only one that can usefully wake
1198 // it, because the unit's own `Hold` gate has already fired and
1199 // would resolve `wait_release` on every iteration forever.
1200 gate: if self.shape_parked { self.shape_gate.clone() } else { head.gate.clone() },
1201 })
1202 }
1203
1204 /// When the queue expects to have written the last unit in it, or
1205 /// `None` if it is empty.
1206 ///
1207 /// `None` rather than `Instant::now()`: the estimate exists only to say
1208 /// what a unit waits behind, and on an empty queue there is nothing
1209 /// ahead. Folding `now` in there would quietly shift every unit's
1210 /// reported release time forward to its push instant.
1211 fn tail_expected_at(&self) -> Option<Instant> {
1212 self.q.back().map(|u| u.expected_at)
1213 }
1214
1215 /// Register the head with the release wheel.
1216 ///
1217 /// Armed at the head's **`due_at`**, never its `expected_at`: the head
1218 /// has nothing ahead of it, so its own deadline is when it may go, and
1219 /// arming on the estimate would re-introduce the successor stall the
1220 /// module doc describes (a unit clamped behind a `Hold` would sleep to
1221 /// that hold's ceiling even after it became the head).
1222 ///
1223 /// The **only** caller of [`release_timer::arm_at`] besides
1224 /// [`Self::push`]'s use of it through here. `arm_at` starts nothing when
1225 /// the head is already due, so a queue of undelayed units never
1226 /// constructs the wheel.
1227 fn arm_head(&mut self) {
1228 self.head_deadline = self.q.front().map(|u| release_timer::arm_at(u.due_at));
1229 }
1230
1231 /// Re-arm the head at an instant the *scheduler* named, rather than at
1232 /// the head's own `due_at`.
1233 /// The head is due on its own account — that is how it reached the shaper —
1234 /// so re-arming at `due_at` would give an already-fired deadline and spin
1235 /// the release branch hot until the bucket refilled. This is what makes
1236 /// `head_release()` *the deadline the scheduler last computed*: it reads
1237 /// back exactly what was armed here.
1238 fn arm_head_at(&mut self, at: Instant) {
1239 self.head_deadline = Some(release_timer::arm_at(at));
1240 }
1241
1242 /// Queue a unit behind everything already waiting.
1243 ///
1244 /// **`due_at` is not touched.** The unit keeps the deadline it was
1245 /// built with, because that deadline is its own: a `Delay`'s
1246 /// `arrived_at + by`, a `Hold`'s `max_hold` ceiling, or `now` for a
1247 /// unit queued purely for ordering. Wire order is held by the deque —
1248 /// [`Self::pop_next_due`] only ever considers the *front* — so nothing
1249 /// has to be added to a successor's deadline to keep it behind its
1250 /// predecessor, and adding something is exactly what stranded a stream
1251 /// behind a released `Hold` (module doc).
1252 /// What *is* computed here is [`Pending::expected_at`], the queue's
1253 /// estimate of when this unit will reach the wire: `max(due_at,
1254 /// tail.expected_at, now)` on a non-empty queue, and `due_at` on an empty
1255 /// one. It is what [`Push::release_at`] returns, so `Effect::Queued {
1256 /// release_at }` still says *no earlier than everything ahead of you*, and
1257 /// it is what [`Self::record_release`] measures lateness against, so a
1258 /// `Pass` queued behind a 300 ms `Delay` is not logged as a 300 ms release
1259 /// error.
1260 ///
1261 /// A unit that is due now on an *empty* queue should not be pushed at
1262 /// all — the pipe loop writes it inline in the read arm, which is
1263 /// today's code and today's cost.
1264 pub(crate) fn push(&mut self, mut unit: Pending) -> Push {
1265 if let Some(tail) = self.tail_expected_at() {
1266 unit.expected_at = unit.due_at.max(tail).max(Instant::now());
1267 }
1268 // Tagged here, not at the call site, so that every producer of a
1269 // queued unit — this file's own terminals, `exec`'s `Delay`/`Hold`
1270 // pushes, and the pipe loop's own header and passthrough bytes — is
1271 // charged to the class the loop resolved, without any of them naming
1272 // one. The clock read is behind the shaping test, so an unshaped push
1273 // costs exactly what it did before.
1274 if self.shaper.is_some() {
1275 unit.shape = Some(ShapeTag {
1276 class: self.unit_class,
1277 expires_at: Instant::now() + self.shape_hold,
1278 starved_noted: false,
1279 });
1280 self.note_unshapeable_seen(&unit);
1281 }
1282 let release_at = unit.expected_at;
1283 let becomes_head = self.q.is_empty();
1284 self.charge(unit.len());
1285 self.q.push_back(unit);
1286 self.counters.note_egress_item_queued();
1287 if becomes_head {
1288 self.arm_head();
1289 }
1290 self.sync_shape_demand();
1291 let entered_backpressure = !self.accepts_more() && !self.backpressure_reported;
1292 if entered_backpressure {
1293 self.backpressure_reported = true;
1294 }
1295 Push { release_at, entered_backpressure }
1296 }
1297
1298 /// Pop the head if it is due at `now`, re-arming the new head.
1299 ///
1300 /// **The front and only the front.** That single fact is the whole
1301 /// ordering guarantee: a unit whose own deadline passed long ago is
1302 /// still not popped while anything is queued in front of it.
1303 ///
1304 /// Call it in a `while let` loop: `pop_next_due` yields one unit at a
1305 /// time and allocates nothing, where a `Vec`-returning `pop_due` would
1306 /// allocate on every release wake. Looping matters as much as it reads:
1307 /// one wake on a `Hold`'s gate has to drain the held unit *and* every
1308 /// unit behind it that is due on its own account, or releasing a gate
1309 /// resumes one object and strands the rest.
1310 ///
1311 /// # The release seam
1312 ///
1313 /// On a shaped queue — and only there — a unit that is due on its own
1314 /// account is then put to `Scheduler::acquire`, which debits its class's
1315 /// token bucket and applies the [`Discipline`]. This function is where
1316 /// that happens because it is `&mut self`, it is the sole ordering
1317 /// authority, and it is called **exactly once per released unit**;
1318 /// `head_release` is none of those things.
1319 ///
1320 /// A refusal is not an error and costs nothing: the head deadline is
1321 /// re-armed at whatever the scheduler named (clamped by the unit's own
1322 /// `max_hold`), the gate a discipline handed back is stored for
1323 /// `head_release`, and `None` is returned exactly as a not-yet-due head
1324 /// returns it. The caller's loop is unchanged.
1325 ///
1326 /// [`Discipline`]: crate::shape::Discipline
1327 pub(crate) fn pop_next_due(&mut self, now: Instant) -> Option<Pending> {
1328 if !self.q.front().is_some_and(|u| u.is_due(now)) {
1329 return None;
1330 }
1331 // A fresh decision every call: a gate stored by a previous refusal
1332 // says nothing about this one, and leaving it set would let
1333 // `head_release` hand out a gate nobody will release.
1334 self.shape_gate = None;
1335 self.shape_parked = false;
1336 if !self.shaping_grants_head(now) {
1337 return None;
1338 }
1339 let unit = self.q.pop_front()?;
1340 self.credit(unit.len());
1341 self.starved_from = self.starved_from.saturating_sub(1).max(1);
1342 self.note_delivered(&unit);
1343 self.arm_head();
1344 self.sync_shape_demand();
1345 Some(unit)
1346 }
1347
1348 /// Whether the shaper lets the head go at `now`, arming whatever it must
1349 /// wait for when it does not.
1350 ///
1351 /// `true` on every unshaped queue and for every unit the shaper does not
1352 /// own — a terminal, an elided ordering slot, a header — so the only
1353 /// thing a bucket can ever hold back is bytes.
1354 fn shaping_grants_head(&mut self, now: Instant) -> bool {
1355 let Some(shaper) = self.shaper.clone() else {
1356 return true;
1357 };
1358 let Some(head) = self.q.front() else {
1359 return true;
1360 };
1361 let Some(tag) = head.shape else {
1362 return true;
1363 };
1364 // Terminals bypass the pacer. A queued reset or truncate is the end
1365 // of the stream, not media: gating it would make teardown ordering
1366 // depend on a token bucket, which is exactly what turns the
1367 // data-then-reset tests into coin flips.
1368 if matches!(head.item, Item::Terminal(_)) {
1369 return true;
1370 }
1371 let bytes = head.len() as u64;
1372
1373 // The clamp comes first, and it is what makes starvation late rather
1374 // than lossy: a unit that has outlived `max_hold` is decided by
1375 // `Expiry`, never by a bucket that has already refused it.
1376 if now >= tag.expires_at {
1377 return self.expire_head(shaper.on_expiry());
1378 }
1379
1380 match shaper.acquire(tag.class, bytes, now) {
1381 Acquire::Now => {
1382 self.tokens_dry = false;
1383 true
1384 }
1385 Acquire::Later(at) => {
1386 self.note_tokens_dry(tag.class);
1387 self.park_head(at.min(tag.expires_at), tag, None);
1388 false
1389 }
1390 // No refill instant exists, so the clamp *is* the deadline. Arming
1391 // a fabricated far-future instant instead would arm a timer that
1392 // never fires, and the stream would look identical to one that
1393 // had simply hung.
1394 Acquire::Never => {
1395 self.note_tokens_dry(tag.class);
1396 self.park_head(tag.expires_at, tag, None);
1397 false
1398 }
1399 // The clamp is the deadline here too — nothing about *time* makes
1400 // a unit larger than the whole bucket fit — but this refusal is a
1401 // mis-sized burst rather than a rate doing its job, and the two
1402 // are otherwise identical from outside: same clamp, same counter,
1403 // same `HoldClamped`. The report is the only thing that separates
1404 // them, so it is set here and deduplicated by the scheduler.
1405 Acquire::LargerThanBurst { burst_bytes, unit_bytes } => {
1406 self.note_tokens_dry(tag.class);
1407 self.shape_report =
1408 Some(ShapeReport::BurstBelowUnit { class: tag.class, burst_bytes, unit_bytes });
1409 self.park_head(tag.expires_at, tag, None);
1410 false
1411 }
1412 Acquire::Starved(gate) => {
1413 self.park_head(tag.expires_at, tag, Some(gate));
1414 false
1415 }
1416 }
1417 }
1418
1419 /// Charge one **episode** of this class's own bucket being dry, on the
1420 /// edge only.
1421 ///
1422 /// Shared by the three refusals that are the bucket's rather than a
1423 /// discipline's, so the edge latch cannot be re-armed on one of them and
1424 /// forgotten on another. Re-armed by a grant, in `shaping_grants_head`.
1425 fn note_tokens_dry(&mut self, class: Class) {
1426 if self.tokens_dry {
1427 return;
1428 }
1429 self.tokens_dry = true;
1430 if let Some(stats) = &self.shape_stats {
1431 stats.note_tokens_exhausted(class);
1432 }
1433 }
1434
1435 /// Park the head until `at` — or until `gate` opens, when a discipline is
1436 /// what is holding it — and charge everything queued behind a unit of
1437 /// another class.
1438 fn park_head(&mut self, at: Instant, tag: ShapeTag, gate: Option<Gate>) {
1439 let by_discipline = gate.is_some();
1440 self.shape_gate = gate;
1441 self.shape_parked = true;
1442 self.arm_head_at(at);
1443 self.sync_shape_demand();
1444 if by_discipline {
1445 // The head itself waited on **another class**, not on its own
1446 // bucket. Same counter as a same-stream head of another class,
1447 // because to a caller they are one question — *was I held
1448 // up by my rate, or by somebody else?* — and the answer is the same
1449 // in both. What it is emphatically not is
1450 // `tokens_exhausted_episodes`: the bucket was never asked.
1451 self.note_head_starved();
1452 }
1453 self.note_starved_behind(tag.class);
1454 }
1455
1456 /// Charge the head once against `starved_behind_other_class`.
1457 fn note_head_starved(&mut self) {
1458 let Some(stats) = self.shape_stats.clone() else { return };
1459 let Some(tag) = self.q.front_mut().and_then(|u| u.shape.as_mut()) else { return };
1460 if tag.starved_noted {
1461 return;
1462 }
1463 tag.starved_noted = true;
1464 stats.note_starved(tag.class);
1465 }
1466
1467 /// Charge `starved_behind_other_class` for every queued unit that is
1468 /// waiting behind a unit of a *different* class.
1469 ///
1470 /// Once per unit over the queue's whole life, not once per wake: the
1471 /// sweep resumes from where the last one stopped and the cursor follows
1472 /// the front as units are popped, so the total work is linear in units
1473 /// queued rather than quadratic in wakes.
1474 ///
1475 /// Deliberately separate from `tokens_exhausted_episodes`. Head-gating
1476 /// makes configured shaping and head-of-line blocking look identical from
1477 /// outside, and this is the number that tells them apart.
1478 fn note_starved_behind(&mut self, head_class: Class) {
1479 let Some(stats) = self.shape_stats.clone() else { return };
1480 let len = self.q.len();
1481 for index in self.starved_from.max(1)..len {
1482 let Some(unit) = self.q.get_mut(index) else { continue };
1483 let Some(tag) = unit.shape.as_mut() else { continue };
1484 if tag.starved_noted || tag.class == head_class {
1485 continue;
1486 }
1487 tag.starved_noted = true;
1488 stats.note_starved(tag.class);
1489 }
1490 self.starved_from = len.max(1);
1491 }
1492
1493 /// The head outlived its clamp. Deliver it, or abandon the stream.
1494 ///
1495 /// Returns whether the caller may pop the head. Under
1496 /// [`Expiry::Deliver`] it may — that is the whole of "clamps and
1497 /// delivers", and it is why `objects_expired` is zero by default.
1498 fn expire_head(&mut self, expiry: Expiry) -> bool {
1499 match expiry {
1500 Expiry::Deliver => {
1501 // What the bucket *would* have asked for is gone by now — the
1502 // deadline it named is in the past, or it named none — so the
1503 // report states the clamp that was applied and leaves the
1504 // request absent rather than inventing a figure for it.
1505 self.shape_report =
1506 Some(ShapeReport::Clamped { requested: None, applied: self.shape_hold });
1507 true
1508 }
1509 Expiry::ResetStream { code } => {
1510 if let (Some(stats), Some(side)) = (&self.shape_stats, self.shape_side) {
1511 stats.note_expired(side);
1512 stats.note_stream_reset_by_shaping(side);
1513 }
1514 self.shape_report = Some(ShapeReport::Expired);
1515 // Everything queued is replaced by the reset: the stream is
1516 // being given up on, and writing part of it first would leave
1517 // the peer a prefix it cannot distinguish from a truncation.
1518 self.clear();
1519 self.q.push_back(Pending::terminal(Terminal::Reset { code }));
1520 self.arm_head();
1521 true
1522 }
1523 }
1524 }
1525
1526 /// Account for a queued unit **no rule could see**: the left-hand side of
1527 /// the conservation identity, for bytes the classifier never met.
1528 /// A stream header, an oversized object's passthrough chunk and a bypassed
1529 /// stream's tail carry no `ObjectMeta`, so `note_object_seen` is never
1530 /// reached for them and `bytes_shaped` would not count them. They are still
1531 /// bytes the shaper handled — it queued them, it ordered them, and it wrote
1532 /// them — so leaving them out would make `bytes_shaped` mean "bytes a rule
1533 /// saw" while
1534 /// [`ShapeStats::bytes_shaped`](crate::shape::ShapeStats::bytes_shaped)
1535 /// promises *every byte it accounted for*, and the `unshapeable` row would
1536 /// have nothing to be the other half of.
1537 ///
1538 /// **Here and not at the call site**, which is what makes the identity
1539 /// structural: `push` is the one place an unshapeable unit enters, this
1540 /// reads the same `unit.len()` [`Self::note_delivered`] will charge when
1541 /// it leaves, and both are gated on the same `ShapeTag`. `session.rs`
1542 /// cannot get it wrong because `session.rs` is not asked — the object
1543 /// arm and the unshown arm share one `exec` entry point, and only the
1544 /// tag tells them apart.
1545 ///
1546 /// Zero-length units — an elided ordering slot, a queued reset — add
1547 /// nothing on either side.
1548 fn note_unshapeable_seen(&self, unit: &Pending) {
1549 if self.unit_class != Class::Unshapeable {
1550 return;
1551 }
1552 let (Some(stats), Some(side)) = (&self.shape_stats, self.shape_side) else {
1553 return;
1554 };
1555 let bytes = unit.len();
1556 if bytes > 0 {
1557 stats.note_unshapeable_seen(side, bytes as u64);
1558 }
1559 }
1560
1561 /// Charge one released unit to its class.
1562 ///
1563 /// [`Class::Unshapeable`] is charged like any other row, and only
1564 /// because [`Self::note_unshapeable_seen`] now counts the same bytes
1565 /// into `bytes_shaped` on the way in: both terms move together or
1566 /// neither does, and this is the unit in which they moved.
1567 /// The consequence worth stating: unshapeable bytes are **not paced**.
1568 /// `shaping_grants_head` never asks a bucket about them — they carry no
1569 /// class a bucket is configured for — so an object too large for the framer
1570 /// to buffer bypasses every rate, and a class rate can be exceeded by
1571 /// exactly one oversized object. That is what "unshapeable" means; the row
1572 /// named for it is where the figure is reported, and it is deliberately a
1573 /// *separate* row from `default_class` so that *no rule claimed this unit*
1574 /// and "no rule could have" stay two answers.
1575 ///
1576 /// The row alone does not say *whose* rate went, which is the question an
1577 /// author with a 500 kbps video cap and a large initial segment is
1578 /// actually asking, so the pipe loop pairs it with
1579 /// `Impairment{ShapeUnpacedObject}` naming the class the stream's
1580 /// classified units are charged to.
1581 fn note_delivered(&self, unit: &Pending) {
1582 let (Some(stats), Some(side), Some(tag)) = (&self.shape_stats, self.shape_side, unit.shape)
1583 else {
1584 return;
1585 };
1586 let bytes = unit.len();
1587 if bytes > 0 {
1588 stats.note_delivered(side, tag.class, bytes as u64);
1589 }
1590 }
1591
1592 /// Tell the scheduler which class this queue is holding an unreleased
1593 /// head of, if any.
1594 ///
1595 /// **Demand, not depth.** One declaration per queue, replaced whenever
1596 /// the head's class changes and withdrawn when the queue empties — which
1597 /// includes [`Self::clear`], the teardown drain, and [`Drop`]. A
1598 /// declaration that outlived its stream would starve a lower-priority
1599 /// class for the rest of the session, and the failure would look like a
1600 /// hang rather than like a leak.
1601 fn sync_shape_demand(&mut self) {
1602 let Some(shaper) = &self.shaper else { return };
1603 let want = self.q.front().and_then(|u| u.shape).map(|t| t.class);
1604 if want == self.shape_demand {
1605 return;
1606 }
1607 if let Some(old) = self.shape_demand.take() {
1608 shaper.withdraw_demand(old);
1609 }
1610 if let Some(new) = want {
1611 shaper.declare_demand(new);
1612 self.shape_demand = Some(new);
1613 }
1614 }
1615
1616 /// Sample one deferred release's lateness onto the session counters.
1617 ///
1618 /// Call it **only** from the release branch, once per unit
1619 /// `pop_next_due` yields, before writing. Units written inline in the
1620 /// read arm are not releases and would dilute the distribution with
1621 /// zeros; drains ignore release times by design and recording their
1622 /// lateness would report teardown as a timing failure. Neither
1623 /// drain in this file calls it.
1624 /// Measured against [`Pending::expected_at`], not [`Pending::due_at`]: the
1625 /// question is *did the engine release when it said it would*, and what it
1626 /// said was `Effect::Queued { release_at }`. Sampling `due_at` instead
1627 /// would log the queueing delay of every unit behind a `Delay` as a timer
1628 /// error.
1629 ///
1630 /// `saturating_duration_since` because a unit released early — a
1631 /// `Hold`'s gate opening long before its ceiling, which is the normal
1632 /// case — is a zero, not a negative.
1633 pub(crate) fn record_release(&self, unit: &Pending, now: Instant) {
1634 self.counters.record_release(now.saturating_duration_since(unit.expected_at));
1635 }
1636
1637 /// Forget everything queued. Used after a terminal fires — the
1638 /// destination is reset, so nothing behind it can be written.
1639 ///
1640 /// Dropping the [`Release`]s abandons their wheel slots, which the
1641 /// wheel prunes at their own instants or at a sweep, without waking
1642 /// anything.
1643 pub(crate) fn clear(&mut self) {
1644 self.q.clear();
1645 self.credit(self.queued_bytes);
1646 self.head_deadline = None;
1647 self.starved_from = 1;
1648 // The shaping gate goes with the units it was holding, and the
1649 // demand goes with the head that declared it: a queue that is being
1650 // emptied has nothing to send, and leaving the declaration standing
1651 // would starve a lower class for the rest of the session.
1652 self.shape_gate = None;
1653 self.shape_parked = false;
1654 self.sync_shape_demand();
1655 }
1656
1657 /// Write everything queued **now**, in order, ignoring release times.
1658 ///
1659 /// The teardown path: delivered late beats lost silently. Best effort
1660 /// by construction — a write failure stops the drain and leaves the
1661 /// failing unit and everything behind it queued, so
1662 /// [`Self::queued_bytes`] reports exactly what did not reach the
1663 /// transport.
1664 ///
1665 /// Every byte it *does* hand over is added to
1666 /// [`Self::unconfirmed_bytes`], because a `write_all` that returns `Ok`
1667 /// into a connection another task is about to `close()` is not a
1668 /// delivery. That is the whole of the teardown-race loss: read
1669 /// [`Self::unconfirmed_bytes`] after this, never [`Self::queued_bytes`],
1670 /// on any path that runs because the session is going down.
1671 ///
1672 /// **Consults no bucket.** This is the teardown drain, and teardown
1673 /// bypasses the pacer entirely: it pops the front directly
1674 /// rather than through [`Self::pop_next_due`], so a token bucket can
1675 /// never gate a mirrored reset. That is what keeps the eleven
1676 /// data-then-teardown ordering tests independent of any configured rate.
1677 ///
1678 /// **And therefore reports no shaping, deliberately.** Every other
1679 /// release path takes a [`ShapeReport`] out of the queue and hands it to
1680 /// a reporter, because a clamp or an expiry that nothing says happened is
1681 /// indistinguishable from a profile that did nothing. This one is the
1682 /// exception, and it is an exception because it *applies* no shaping:
1683 /// [`Self::pop_next_due`] is the only producer of a `ShapeReport` and
1684 /// this function never calls it, so there is nothing to drain rather
1685 /// than something dropped. The slot is empty on entry as well — both
1686 /// callers that can reach it through a shaped queue
1687 /// ([`drain_honouring_release_times`]'s cancel arm and
1688 /// `release_due_units`) drain the report before the pop that produced it
1689 /// is written.
1690 ///
1691 /// **Known residue, stated rather than fixed:** popping the front
1692 /// directly also skips the `shape_parked` / `shape_gate` bookkeeping
1693 /// [`Self::pop_next_due`] maintains, so after this drain a
1694 /// [`Self::head_release`] would describe a park belonging to a unit that
1695 /// is gone. Nothing reads it — every caller returns immediately, the
1696 /// queue is empty or cleared, and `head_release` answers `None` on an
1697 /// empty deque — so there is no failure case here to fix. It is
1698 /// written down because "unreachable today" is a property of the
1699 /// callers, not of this function.
1700 pub(crate) async fn drain_ignoring_release_times<S: EgressSink>(
1701 &mut self,
1702 send: &mut S,
1703 ) -> DrainOutcome {
1704 // A close that was given a drain window and spent it has already
1705 // decided these bytes are not going out. Writing them anyway would move
1706 // them from *queued, and reported as abandoned* into *handed to a
1707 // connection that is closing*, which is the one state nothing
1708 // downstream can count: see [`EgressGauge`].
1709 if self.discarding() {
1710 return DrainOutcome::Discarded;
1711 }
1712 while let Some(unit) = self.q.front().cloned() {
1713 match write_unit(unit, send).await {
1714 Ok(Written::Terminated { forwarded, code }) => {
1715 // Everything behind a terminal is discarded by design,
1716 // not lost: the destination is reset. Only the prefix
1717 // this unit actually handed over is unconfirmed.
1718 self.flushed_unconfirmed = self.flushed_unconfirmed.saturating_add(forwarded);
1719 self.clear();
1720 return DrainOutcome::Terminated { forwarded, code };
1721 }
1722 Ok(written) => {
1723 if let Written::Bytes(n) = written {
1724 self.flushed_unconfirmed = self.flushed_unconfirmed.saturating_add(n);
1725 }
1726 let unit = self.q.pop_front().expect("front was just observed");
1727 self.credit(unit.len());
1728 self.starved_from = self.starved_from.saturating_sub(1).max(1);
1729 // Charged to its class like any other release: a flush that
1730 // `write_all` accepted is the only definition of *released
1731 // to the destination stream* this side of the transport,
1732 // and what it could not vouch for is separately reported as
1733 // `QueuedBytesAtTeardown`.
1734 self.note_delivered(&unit);
1735 self.arm_head();
1736 self.sync_shape_demand();
1737 }
1738 Err(_) => return DrainOutcome::WriteFailed,
1739 }
1740 }
1741 DrainOutcome::Complete
1742 }
1743}
1744
1745/// Withdraw whatever demand this queue still holds.
1746///
1747/// `Drop` rather than a call on each teardown path, and for the same reason
1748/// `StreamGuard` is: a forwarding task can end in at least five different
1749/// ways, any hand-written list of them is only as complete as its reader,
1750/// and `Drop` additionally covers a `?` return, a panicking task, and
1751/// `JoinSet::shutdown` dropping the future wholesale. A leaked declaration
1752/// is not a loud failure — it is a *different* stream stalling to
1753/// `max_hold`, on a class the reader was not looking at.
1754impl Drop for PendingQueue {
1755 fn drop(&mut self) {
1756 if let (Some(shaper), Some(class)) = (&self.shaper, self.shape_demand.take()) {
1757 shaper.withdraw_demand(class);
1758 }
1759 // The session gauge counts bytes that are still *somewhere*, and a
1760 // queue going away takes its remainder with it. Without this a
1761 // stream torn down with a full queue would leave the session's
1762 // total permanently above zero, and every later close would spend
1763 // its whole drain window waiting for a queue that no longer exists.
1764 self.credit(self.queued_bytes);
1765 }
1766}
1767
1768/// Write everything queued **at its release time**, racing cancellation.
1769///
1770/// The shape both release-honouring drains take — the FIN drain in the read
1771/// arm's `None` branch and the one a terminal runs before it resets. Both
1772/// sit inside `select!` arm bodies, which are not preemptible, so writing
1773/// either as a plain `while let` loop would let a `Hold` on a gate nobody
1774/// releases pin session teardown for up to [`EgressConfig::max_hold`] —
1775/// precisely the failure the deque was chosen over a writer task to avoid.
1776///
1777/// The inner `select!` is `biased` so cancellation wins deterministically
1778/// when both are ready. Unbiased, a cancelled session could keep picking
1779/// the release branch, find nothing due, and spin.
1780///
1781/// Release lateness is **not** sampled here: this drain honours release
1782/// times only until cancellation, and recording a teardown flush as a
1783/// timing sample would report teardown as a timing failure.
1784///
1785/// # Shaping **is** reported here
1786///
1787/// `on_shape` is called once per [`ShapeReport`] the drain's own
1788/// [`PendingQueue::pop_next_due`] produces, and it exists because this loop
1789/// is a full release seam and not a teardown flush: on a shaped queue every
1790/// pop below debits a token bucket and every clamp or expiry is decided
1791/// here. The FIN path — header, a few objects, FIN, the ordinary MoQT
1792/// subgroup shape — reaches the wire through *this* function and not
1793/// through `release_due_units`, so a drain that swallowed its reports would
1794/// apply the whole profile to the normal case and say nothing about it.
1795/// That is the silent no-op this signature exists to prevent, and the reason
1796/// the parameter is not optional: a caller cannot forget what it has to name.
1797///
1798/// Passing a no-op closure is correct only where no report can exist —
1799/// every `#[cfg(test)]` caller in this module builds an unshaped queue, and
1800/// an unshaped queue's `pop_next_due` never reaches a scheduler.
1801pub(crate) async fn drain_honouring_release_times<S, F>(
1802 pending: &mut PendingQueue,
1803 send: &mut S,
1804 cancel: &CancellationToken,
1805 mut on_shape: F,
1806) -> Result<DrainOutcome, EgressError>
1807where
1808 S: EgressSink,
1809 F: FnMut(ShapeReport),
1810{
1811 while let Some(release) = pending.head_release() {
1812 tokio::select! {
1813 biased;
1814 () = cancel.cancelled() => {
1815 return Ok(match pending.drain_ignoring_release_times(send).await {
1816 DrainOutcome::Terminated { forwarded, code } => {
1817 DrainOutcome::Terminated { forwarded, code }
1818 }
1819 _ => DrainOutcome::CancelledMidDrain,
1820 });
1821 }
1822 () = wait_release(Some(release), cancel) => {
1823 let now = Instant::now();
1824 while let Some(unit) = pending.pop_next_due(now) {
1825 // Before the write, exactly as `release_due_units` does:
1826 // a terminal write returns out of this loop, and an
1827 // expiry's report belongs to the very unit whose write
1828 // takes that return.
1829 if let Some(report) = pending.take_shape_report() {
1830 on_shape(report);
1831 }
1832 if let Written::Terminated { forwarded, code } = write_unit(unit, send).await? {
1833 pending.clear();
1834 return Ok(DrainOutcome::Terminated { forwarded, code });
1835 }
1836 }
1837 // A refusal reports too. The clamp is decided inside
1838 // `shaping_grants_head`, which runs whether or not the head
1839 // is yielded, so reading the report only after a successful
1840 // pop would lose the one case that matters.
1841 if let Some(report) = pending.take_shape_report() {
1842 on_shape(report);
1843 }
1844 }
1845 }
1846 }
1847 Ok(DrainOutcome::Complete)
1848}
1849
1850// ── Session close ───────────────────────────────────────────────────
1851
1852/// The default `(code, reason)` `run_with_transport` closes with when no
1853/// hook asked for anything else — today's hard-coded pair.
1854const DEFAULT_CLOSE: (u32, &[u8]) = (0, b"proxy session ended");
1855
1856/// Where [`crate::action::Action::CloseSession`] lands.
1857///
1858/// A `OnceLock<(u32, Bytes)>` plus the session's `CancellationToken`, and
1859/// deliberately **no `Transport` handles**: the transports own the tasks
1860/// that own the hooks that reach this, so holding them here would close a
1861/// reference cycle. The close itself happens where it already happens, in
1862/// `run_with_transport`, which reads [`Self::close_args`] instead of
1863/// hard-coding `(0, b"proxy session ended")`.
1864///
1865/// `CloseSession` is honoured at **every** site that returns an `Action`,
1866/// including `Site::StreamEnd` on both data and control streams: a close is
1867/// session-scoped, so no site can be the wrong one for it.
1868#[derive(Clone, Debug)]
1869pub(crate) struct SessionCloser {
1870 inner: Arc<CloserInner>,
1871}
1872
1873#[derive(Debug)]
1874struct CloserInner {
1875 request: OnceLock<(u32, Bytes, CloseOrigin)>,
1876 cancel: CancellationToken,
1877}
1878
1879/// Who asked for a session close.
1880///
1881/// Carried alongside the code and the reason because
1882/// [`ProxyEvent::SessionEnded`](crate::event::ProxyEvent::SessionEnded)
1883/// names the cause in prose, and the two callers are not interchangeable to
1884/// anyone reading that: a hook's `Action::CloseSession` is the run
1885/// under test deciding something, while
1886/// [`ProxyControl::close_session`](crate::control::ProxyControl::close_session)
1887/// is the operator outside it pulling the plug. The recorded pair alone
1888/// cannot tell them apart — both arrive as a `u32` and some bytes — and
1889/// before this every control-plane close was reported to observers as a
1890/// hook's, which is a sentence about the run that was simply untrue.
1891#[derive(Clone, Copy, Debug, PartialEq, Eq)]
1892pub(crate) enum CloseOrigin {
1893 /// A hook returned [`crate::action::Action::CloseSession`].
1894 Hook,
1895 /// [`ProxyControl::close_session`](crate::control::ProxyControl::close_session).
1896 ControlPlane,
1897}
1898
1899impl SessionCloser {
1900 /// A closer bound to the session's cancellation token.
1901 pub(crate) fn new(cancel: CancellationToken) -> Self {
1902 Self { inner: Arc::new(CloserInner { request: OnceLock::new(), cancel }) }
1903 }
1904
1905 /// Record a close request and cancel the session.
1906 ///
1907 /// The first request wins and returns `true`. A later one changes
1908 /// nothing and returns `false` — refuse it with
1909 /// [`Refusal::SessionAlreadyClosing`](crate::capability::Refusal::SessionAlreadyClosing)
1910 /// rather than letting the second reason overwrite the first.
1911 pub(crate) fn request(&self, code: u32, reason: Bytes) -> bool {
1912 let won = self.inner.request.set((code, reason, CloseOrigin::Hook)).is_ok();
1913 // Cancel unconditionally: a losing request still arrived after a
1914 // winning one, so the session is already on its way down and this
1915 // is idempotent.
1916 self.inner.cancel.cancel();
1917 won
1918 }
1919
1920 /// Record a close request **without** cancelling the session.
1921 ///
1922 /// The first request wins and returns `true`, exactly as
1923 /// [`Self::request`] does; the only difference is that the session
1924 /// keeps running afterwards.
1925 ///
1926 /// That difference is the whole reason this exists. A control-plane
1927 /// close has two steps — fix the code and reason, then give the egress
1928 /// queues a bounded window to flush — and [`Self::request`] cannot
1929 /// express the first without the second, because it cancels as it
1930 /// records. A caller that used it would end the session before the
1931 /// drain it just asked for had begun, and the drain window would be
1932 /// unobservable.
1933 ///
1934 /// Nothing else changes: `run_with_transport` still reads
1935 /// [`Self::close_args`] at teardown, so a session recorded this way
1936 /// closes with the recorded pair whenever it does close — whether that
1937 /// is the caller cancelling after its drain, or the peer going away
1938 /// first.
1939 pub(crate) fn record(&self, code: u32, reason: Bytes) -> bool {
1940 self.inner.request.set((code, reason, CloseOrigin::ControlPlane)).is_ok()
1941 }
1942
1943 /// Whether a close has been requested.
1944 ///
1945 /// `session.rs` reads [`Self::requested`] instead, because it wants the
1946 /// code and the reason as well; this stays as the cheap yes-or-no
1947 /// predicate for callers that need nothing more.
1948 #[allow(dead_code)]
1949 pub(crate) fn is_closing(&self) -> bool {
1950 self.inner.request.get().is_some()
1951 }
1952
1953 /// What was requested and by whom, if anything.
1954 ///
1955 /// The origin rides along rather than being a second accessor because
1956 /// the two are set together, under one `OnceLock`, and a caller that
1957 /// read them in two calls would be writing a race it does not have.
1958 pub(crate) fn requested(&self) -> Option<(u32, Bytes, CloseOrigin)> {
1959 self.inner.request.get().cloned()
1960 }
1961
1962 /// The `(code, reason)` to hand `Transport::close`, falling back to the
1963 /// pair the proxy has always used.
1964 /// The origin is deliberately dropped here: it is a fact about this
1965 /// process, and what goes on the wire is what the peer was told. A
1966 /// `CONNECTION_CLOSE` carrying *the control plane asked* would be this
1967 /// proxy leaking its own topology into somebody else's session.
1968 pub(crate) fn close_args(&self) -> (u32, Bytes) {
1969 self.inner
1970 .request
1971 .get()
1972 .map(|(code, reason, _)| (*code, reason.clone()))
1973 .unwrap_or_else(|| (DEFAULT_CLOSE.0, Bytes::from_static(DEFAULT_CLOSE.1)))
1974 }
1975}
1976
1977#[cfg(test)]
1978mod tests {
1979 use super::*;
1980
1981 /// An [`EgressSink`] that records instead of sending.
1982 #[derive(Debug, Default)]
1983 struct RecordingSink {
1984 writes: Vec<Bytes>,
1985 reset_code: Option<u64>,
1986 /// Fail every write once this many have succeeded.
1987 fail_after: Option<usize>,
1988 }
1989
1990 impl RecordingSink {
1991 fn written(&self) -> Bytes {
1992 let mut out = Vec::new();
1993 for w in &self.writes {
1994 out.extend_from_slice(w);
1995 }
1996 Bytes::from(out)
1997 }
1998 }
1999
2000 impl EgressSink for RecordingSink {
2001 fn write_all(
2002 &mut self,
2003 buf: &[u8],
2004 ) -> impl Future<Output = Result<(), EgressError>> + Send {
2005 let out = if self.fail_after.is_some_and(|n| self.writes.len() >= n) {
2006 Err(EgressError::Transport(TransportError::Write(
2007 "recording sink is full".to_owned(),
2008 )))
2009 } else {
2010 self.writes.push(Bytes::copy_from_slice(buf));
2011 Ok(())
2012 };
2013 async move { out }
2014 }
2015
2016 fn reset(&mut self, code: u64) -> Result<(), EgressError> {
2017 self.reset_code = Some(code);
2018 Ok(())
2019 }
2020 }
2021
2022 /// The `on_shape` every drain in this module passes.
2023 ///
2024 /// Sound only because every queue built here is **unshaped** — `queue()`
2025 /// and `queue_with()` call `PendingQueue::new` and never
2026 /// `with_shaper`, and `pop_next_due` reaches no scheduler without one —
2027 /// so no `ShapeReport` can exist to swallow. A `unreachable!` rather than
2028 /// an empty body, so a later edit that installs a shaper here finds out
2029 /// instead of inheriting that silent no-op a second time.
2030 fn no_shape_reports(report: ShapeReport) {
2031 unreachable!("an unshaped queue produces no shaping report, got {report:?}");
2032 }
2033
2034 fn queue() -> (PendingQueue, Arc<Recorder>) {
2035 let counters = Arc::new(Recorder::new());
2036 (PendingQueue::new(EgressConfig::default(), Arc::clone(&counters)), counters)
2037 }
2038
2039 fn queue_with(config: EgressConfig) -> (PendingQueue, Arc<Recorder>) {
2040 let counters = Arc::new(Recorder::new());
2041 (PendingQueue::new(config, Arc::clone(&counters)), counters)
2042 }
2043
2044 // ── shaped queues ───────────────────────────────────────────────
2045 //
2046 // Every queue above is unshaped, which is deliberate: the ~25 tests in
2047 // this module encode invariants that must hold whether or not a profile
2048 // is configured, and a shaper installed on all of them would stop being
2049 // separable from them. The handful below install one.
2050
2051 use crate::shape::{
2052 BucketConfig, ClassRule, DirectionStats, Discipline, Expiry, Matcher, Overflow,
2053 QueueConfig, ShapeProfile, ShapeRecorder,
2054 };
2055
2056 /// One named bucket. Struct literals rather than field assignment:
2057 /// `#[non_exhaustive]` does not apply inside the defining crate, and
2058 /// `..Default::default()` is what clippy's `field_reassign_with_default`
2059 /// asks for.
2060 fn test_bucket(rate: Option<u64>, burst: u64) -> BucketConfig {
2061 BucketConfig {
2062 name: "b".to_string(),
2063 rate_bps: rate,
2064 burst_bytes: burst,
2065 ..BucketConfig::default()
2066 }
2067 }
2068
2069 /// One class over bucket `b`, claiming every unit.
2070 fn test_class(name: &str, priority: u8) -> ClassRule {
2071 ClassRule {
2072 name: name.to_string(),
2073 bucket: "b".to_string(),
2074 matcher: Matcher::default(),
2075 priority,
2076 weight: 1,
2077 }
2078 }
2079
2080 /// A one-class profile over one bucket, with everything a shaped queue
2081 /// reads spelled out.
2082 fn profile(rate: Option<u64>, burst: u64, max_hold: Duration, expiry: Expiry) -> ShapeProfile {
2083 let queue = QueueConfig {
2084 max_hold: Some(max_hold),
2085 overflow: Overflow::Block,
2086 on_expiry: expiry,
2087 ..QueueConfig::default()
2088 };
2089 ShapeProfile::try_new(
2090 vec![test_bucket(rate, burst)],
2091 vec![test_class("only", 0)],
2092 queue,
2093 Discipline::Fifo,
2094 )
2095 .expect("the fixture names its own bucket")
2096 }
2097
2098 /// A queue paced by `profile`, plus the recorder it writes to.
2099 ///
2100 /// One leg, because one leg is all these tests are about; the leg the
2101 /// totals actually land on is
2102 /// [`two_legs_sharing_one_recorder_charge_their_own_side`].
2103 fn shaped(profile: ShapeProfile) -> (PendingQueue, Arc<ShapeRecorder>, Arc<Scheduler>) {
2104 let counters = Arc::new(Recorder::new());
2105 let stats = Arc::new(ShapeRecorder::for_profile(Some(&profile)));
2106 let shaper = Arc::new(Scheduler::new(profile));
2107 let mut q = PendingQueue::new(EgressConfig::default(), counters).with_shaper(
2108 Some(Arc::clone(&shaper)),
2109 Arc::clone(&stats),
2110 ProxySide::ClientToProxy,
2111 );
2112 q.tag_unit(Class::Rule(0));
2113 (q, stats, shaper)
2114 }
2115
2116 fn payload(n: usize) -> Bytes {
2117 Bytes::from(vec![0xAB; n])
2118 }
2119
2120 /// The seam. A shaped head that is due on its own account is still held
2121 /// back by its bucket, and the deadline `head_release` then reports is
2122 /// the one the **scheduler** computed — not the head's own `due_at`,
2123 /// which is already in the past.
2124 ///
2125 /// That last clause is the whole reason `arm_head_at` exists: re-arming
2126 /// at `due_at` would hand back an already-fired deadline and spin the
2127 /// release branch hot until the bucket refilled.
2128 ///
2129 /// *Ablation:* delete the `shaping_grants_head` call from
2130 /// `pop_next_due` — the first `is_none` assertion reddens with the unit
2131 /// popped, which is a configured rate that paces nothing.
2132 #[test]
2133 fn a_shaped_head_waits_for_its_bucket_and_reports_the_scheduler_s_deadline() {
2134 // 1000 bytes/s, 1000-byte burst: the first 1000-byte unit fits and
2135 // the second needs a full second.
2136 let (mut q, _stats, _s) =
2137 shaped(profile(Some(1_000), 1_000, Duration::from_secs(60), Expiry::Deliver));
2138 let now = Instant::now();
2139 q.push(Pending::bytes(payload(1_000), now));
2140 q.push(Pending::bytes(payload(1_000), now));
2141
2142 assert!(q.pop_next_due(now).is_some(), "the burst covers the first unit");
2143 assert!(
2144 q.pop_next_due(now).is_none(),
2145 "the burst is spent, so the second unit waits on a refill"
2146 );
2147 // The head is due on its own account — it was pushed at `now` — so a
2148 // deadline that had not been re-armed would already have fired.
2149 let release = q.head_release().expect("a non-empty queue always has one");
2150 assert!(
2151 !release.deadline().token().is_cancelled(),
2152 "the head must be armed at the refill instant, not at its own past due_at"
2153 );
2154 assert!(
2155 q.pop_next_due(now + Duration::from_secs(1)).is_some(),
2156 "a second's refill covers it"
2157 );
2158 assert!(q.is_empty());
2159 }
2160
2161 /// The spin guard, and it is a **measured** defect rather than a
2162 /// hypothetical one.
2163 ///
2164 /// A `Hold` whose gate has been released makes its unit due, and the gate
2165 /// stays released forever. If the queue kept reporting that gate while
2166 /// the shaper was refusing the same unit, `wait_release` would resolve on
2167 /// every `select!` iteration and the pipe loop would spin hot for the
2168 /// whole of `max_hold`. Measured before `shape_parked` existed: a
2169 /// barrier-held object on a 0-bps class pinned a core for five seconds
2170 /// and starved the test's own task alongside it, so
2171 /// `a_zero_rate_class_starves_and_the_other_flows` failed with the video
2172 /// stream *delivered at its clamp* rather than held.
2173 ///
2174 /// *Ablation:* restore `head.gate.clone().or_else(|| shape_gate)` in
2175 /// `head_release` — the `is_released` assertion reddens, and the
2176 /// integration gate above it goes from 30 ms to 5 s.
2177 #[test]
2178 fn a_shaper_holding_a_gated_unit_does_not_report_the_gate_that_freed_it() {
2179 let (mut q, _stats, _s) =
2180 shaped(profile(Some(0), 0, Duration::from_secs(60), Expiry::Deliver));
2181 let now = Instant::now();
2182 let gate = Gate::new();
2183 // A `Hold`'s shape: due at the ceiling, released early by its gate.
2184 q.push(Pending::bytes(payload(16), now + Duration::from_secs(30)).with_gate(gate.clone()));
2185
2186 assert!(q.pop_next_due(now).is_none(), "an unreleased gate leaves it not due");
2187 assert!(
2188 q.head_release().expect("queued").gate.is_some_and(|g| !g.is_released()),
2189 "before the shaper has an opinion, the unit's own gate is what it waits on"
2190 );
2191
2192 gate.release();
2193 assert!(q.pop_next_due(now).is_none(), "released and due, but the bucket refuses");
2194 let release = q.head_release().expect("still queued");
2195 assert!(
2196 !release.gate.is_some_and(|g| g.is_released()),
2197 "a released Hold gate must not be re-reported while the shaper holds the unit: \
2198 wait_release would resolve on every iteration and spin the pipe loop"
2199 );
2200 }
2201
2202 /// An **unshaped** queue is untouched by any of this: no scheduler, no
2203 /// tag, no clock read beyond what it already did.
2204 ///
2205 /// The control for every test above it. Without this row a bug that
2206 /// paced everything unconditionally would still leave the shaped tests
2207 /// green.
2208 #[test]
2209 fn an_unshaped_queue_pops_a_due_head_with_no_shaper_at_all() {
2210 let (mut q, _c) = queue();
2211 assert!(!q.is_shaped());
2212 let now = Instant::now();
2213 q.push(Pending::bytes(payload(1_000_000), now));
2214 assert!(q.pop_next_due(now).is_some(), "an unshaped queue has nothing to ask");
2215 }
2216
2217 /// A queued terminal bypasses the pacer, on a bucket that grants nothing
2218 /// at all.
2219 ///
2220 /// This is the teardown-bypasses-the-pacer rule at the one place it is
2221 /// not obvious: a `Truncate` or a hook's `ResetStream` is a *queued*
2222 /// unit, so it reaches `pop_next_due` like any other. Gating it would
2223 /// make the eleven data-then-teardown ordering tests depend on a token
2224 /// bucket.
2225 ///
2226 /// # Why the `Truncate` arm is the one that matters
2227 ///
2228 /// Found by the ablation, not by reading. A bare `Terminal::Reset`
2229 /// carries **zero bytes**, and a zero-byte charge is granted by any
2230 /// bucket — `tokens >= 0` holds even at rate zero — so a `Reset`-only
2231 /// fixture passes with the guard deleted and proves nothing. A
2232 /// `Truncate` owes its prefix to the wire, and that prefix is exactly
2233 /// what a dry bucket would sit on. Both are asserted, in that order, so
2234 /// the row says which of the two the guard is for.
2235 /// *Ablation, recorded:* remove the `Item::Terminal` early return from
2236 /// `shaping_grants_head`. The `Reset` still pops; the `Truncate` is held to
2237 /// `max_hold` and its `expect` reddens with *a truncate owes its prefix to
2238 /// the wire, not to a bucket*.
2239 #[test]
2240 fn a_terminal_is_never_gated_by_a_bucket() {
2241 let (mut q, _stats, _s) =
2242 shaped(profile(Some(0), 0, Duration::from_secs(60), Expiry::Deliver));
2243
2244 q.push(Pending::terminal(Terminal::Reset { code: 7 }));
2245 // `Pending::terminal` stamps its own `due_at` from the clock, so the
2246 // pop instant has to be taken after the push or the unit is simply
2247 // not due yet and this would pass without asking the shaper anything.
2248 let unit = q.pop_next_due(Instant::now()).expect("a reset consults no bucket");
2249 assert!(unit.is_terminal());
2250
2251 q.push(Pending::terminal(Terminal::Truncate { prefix: payload(64), code: 9 }));
2252 let unit = q
2253 .pop_next_due(Instant::now())
2254 .expect("a truncate owes its prefix to the wire, not to a bucket");
2255 assert_eq!(unit.len(), 64, "and the prefix goes with it");
2256 }
2257
2258 /// Under the default `Expiry::Deliver` a class that can never be
2259 /// granted from tokens **still delivers**, at `max_hold`, and says so.
2260 ///
2261 /// The two halves are one test because the report is what keeps the
2262 /// clamp from being a silent lateness, and the delivery is what keeps a
2263 /// starved class from being a silent loss.
2264 ///
2265 /// *Ablation:* make the `Expiry::Deliver` arm return `false` — the unit
2266 /// is never yielded, the `is_some` assertion reddens, and in a session
2267 /// that is a 0-bps class that hangs its stream instead of running late.
2268 ///
2269 /// The report is asserted as a **whole variant**, `requested` included,
2270 /// and not as `matches!(.., Clamped { applied, .. })`: a zero-rate
2271 /// bucket names no refill instant, so the request this clamp cut short
2272 /// is unbounded and the only honest value for it is absent. Written as
2273 /// a wildcard it passed against a `Duration::MAX` sentinel, which is
2274 /// what the field used to hold and what
2275 /// `an_unbounded_shaping_wait_reports_no_request_at_all` exists for.
2276 #[test]
2277 fn a_zero_rate_class_still_delivers_at_the_clamp_and_reports_it() {
2278 const HOLD: Duration = Duration::from_millis(80);
2279 let (mut q, stats, _s) = shaped(profile(Some(0), 0, HOLD, Expiry::Deliver));
2280 let now = Instant::now();
2281 q.push(Pending::bytes(payload(16), now));
2282
2283 assert!(q.pop_next_due(now).is_none(), "a zero-rate bucket grants nothing");
2284 assert_eq!(
2285 stats.snapshot().classes[0].tokens_exhausted_episodes,
2286 1,
2287 "and says its own bucket was dry, once per episode"
2288 );
2289 assert_eq!(q.take_shape_report(), None, "nothing is clamped before the clamp");
2290
2291 let unit = q
2292 .pop_next_due(now + HOLD + Duration::from_millis(1))
2293 .expect("Expiry::Deliver clamps and delivers");
2294 assert_eq!(unit.len(), 16);
2295 assert_eq!(
2296 q.take_shape_report(),
2297 Some(ShapeReport::Clamped { requested: None, applied: HOLD }),
2298 "a clamped release reports the clamp it applied, and reports the \
2299 request it cut short as absent because a zero-rate bucket named none"
2300 );
2301 assert_eq!(stats.snapshot().classes[0].bytes_delivered, 16);
2302 assert_eq!(stats.snapshot().objects_expired, 0, "Deliver never expires an object");
2303 }
2304
2305 /// A shaping clamp reports **no request at all**, on both of the two
2306 /// ways a bucket can decline to name a refill instant.
2307 /// The two legs are the two refusals that name no instant: a rate of zero,
2308 /// and a unit larger than the burst the bucket can ever hold. Both are
2309 /// released by the clamp and by nothing else, and neither has a duration to
2310 /// quote — the wait was unbounded. Leg (b) additionally has a **positive**
2311 /// rate, so it rules out *the report is absent because the class is
2312 /// switched off*.
2313 ///
2314 /// Asserted as a whole-variant equality, which is the point: the field
2315 /// used to carry `Duration::MAX` here, and every inequality an author
2316 /// would naturally write against it — `requested > applied` — is
2317 /// satisfied by that sentinel exactly as it is by a real request.
2318 /// The **pre-clamp** report is asserted per leg for the same reason, and it
2319 /// is where the two legs stop looking alike: a zero-rate class is silent
2320 /// until its clamp fires, and a class whose burst cannot cover one of its
2321 /// own units says so at the first refusal. That asymmetry is the only thing
2322 /// separating *this rate is being applied and it is low* from *this rate is
2323 /// not being applied at all*, since both produce this same `Clamped`
2324 /// afterwards.
2325 ///
2326 /// *Ablation, run:* restore the sentinel in `expire_head`'s
2327 /// `Expiry::Deliver` arm, `requested: Some(Duration::MAX)`. Both legs
2328 /// redden with
2329 /// `left: Some(Clamped { requested: Some(18446744073709551615.999999999s), applied: 80ms })`
2330 /// against `right: Some(Clamped { requested: None, applied: 80ms })`.
2331 #[test]
2332 fn an_unbounded_shaping_wait_reports_no_request_at_all() {
2333 const HOLD: Duration = Duration::from_millis(80);
2334 /// Bigger than the burst leg (b) configures, so the bucket can never
2335 /// hold enough for it however long it accrues.
2336 const OVERSIZED: usize = 4_096;
2337
2338 let legs = [
2339 ("a zero rate", profile(Some(0), 0, HOLD, Expiry::Deliver), 16, None),
2340 (
2341 "a unit above burst_bytes",
2342 profile(Some(64_000), 64, HOLD, Expiry::Deliver),
2343 OVERSIZED,
2344 Some(ShapeReport::BurstBelowUnit {
2345 class: Class::Rule(0),
2346 burst_bytes: 64,
2347 unit_bytes: OVERSIZED as u64,
2348 }),
2349 ),
2350 ];
2351
2352 for (label, profile, bytes, pre_clamp) in legs {
2353 let (mut q, _stats, _s) = shaped(profile);
2354 let now = Instant::now();
2355 q.push(Pending::bytes(payload(bytes), now));
2356
2357 assert!(q.pop_next_due(now).is_none(), "{label}: nothing may be granted up front");
2358 assert_eq!(
2359 q.take_shape_report(),
2360 pre_clamp,
2361 "{label}: nothing is *clamped* before the clamp, and only the \
2362 mis-sized burst says anything at all before it"
2363 );
2364
2365 let unit = q
2366 .pop_next_due(now + HOLD + Duration::from_millis(1))
2367 .unwrap_or_else(|| panic!("{label}: Expiry::Deliver clamps and delivers"));
2368 assert_eq!(unit.len(), bytes, "{label}: the whole unit goes out");
2369 assert_eq!(
2370 q.take_shape_report(),
2371 Some(ShapeReport::Clamped { requested: None, applied: HOLD }),
2372 "{label}: the bucket named no instant, so there is no request to \
2373 quote and the report must say so rather than name a sentinel"
2374 );
2375 }
2376 }
2377
2378 /// `Expiry::ResetStream` abandons the whole destination stream: the
2379 /// queue is replaced by the reset, the object is counted as expired and
2380 /// the stream as reset by shaping.
2381 ///
2382 /// Everything queued goes with it, and that is deliberate — writing part
2383 /// of a stream the shaper has given up on leaves the peer a prefix it
2384 /// cannot tell from a truncation.
2385 /// *Ablation:* drop the `note_expired` call — the `objects_expired`
2386 /// assertion reddens while the reset still happens, which is exactly the
2387 /// *it worked but nothing says so* failure the counter exists for.
2388 #[test]
2389 fn expiry_reset_stream_replaces_the_queue_with_its_reset() {
2390 const HOLD: Duration = Duration::from_millis(50);
2391 let (mut q, stats, _s) =
2392 shaped(profile(Some(0), 0, HOLD, Expiry::ResetStream { code: 0x2A }));
2393 let now = Instant::now();
2394 q.push(Pending::bytes(payload(16), now));
2395 q.push(Pending::bytes(payload(16), now));
2396
2397 let unit = q
2398 .pop_next_due(now + HOLD + Duration::from_millis(1))
2399 .expect("an expired head under ResetStream yields the reset");
2400 assert_eq!(unit.item(), &Item::Terminal(Terminal::Reset { code: 0x2A }));
2401 assert_eq!(q.take_shape_report(), Some(ShapeReport::Expired));
2402 assert!(q.is_empty(), "everything behind it went with the stream");
2403
2404 let snap = stats.snapshot();
2405 assert_eq!(snap.objects_expired, 1);
2406 assert_eq!(snap.streams_reset_by_shaping, 1);
2407 }
2408
2409 /// Two queues, one recorder, one leg each: what a queue charges lands on
2410 /// the leg it was built for.
2411 ///
2412 /// This is the seam between the split totals and the thing that knows
2413 /// which side a stream is on. The recorder is session-scoped and shared
2414 /// by every forwarding task, so it cannot work the direction out for
2415 /// itself; the queue is per stream direction and can. A recorder that
2416 /// could split but that nothing ever told which leg to charge would
2417 /// report one side doing all the work and the other doing none, which
2418 /// is indistinguishable from a session that really was one-way.
2419 ///
2420 /// The two legs are deliberately asymmetric in **opposite** senses —
2421 /// the uplink carries fewer bytes and takes the expiry, the downlink
2422 /// carries six times the bytes and expires nothing — so swapping the
2423 /// rows fails on every field rather than cancelling out.
2424 ///
2425 /// `objects_seen` is zero on both legs, and it is asserted rather than
2426 /// elided: it is charged where a unit is classified, which is the pipe
2427 /// loop's job and not this queue's, so a queue that started moving it
2428 /// would be counting the same object twice.
2429 ///
2430 /// *Ablation, recorded:* pass `ProxySide::ClientToProxy` for both queues — the
2431 /// uplink compare reddens with `left: DirectionStats { objects_seen: 0,
2432 /// bytes_shaped: 112, objects_expired: 1, streams_reset_by_shaping: 1,
2433 /// streams_with_mixed_classes: 0 }` against `right: DirectionStats { ..,
2434 /// bytes_shaped: 16, .. }`. Note what *does not* move: the aggregate is
2435 /// still 112 and the expiry is still counted once, so a one-legged
2436 /// recorder passes every figure a reader would have had before the
2437 /// split.
2438 #[test]
2439 fn two_legs_sharing_one_recorder_charge_their_own_side() {
2440 const HOLD: Duration = Duration::from_millis(50);
2441 // A bucket that grants nothing, so nothing leaves except through
2442 // the clamp, and the clamp is what the expiry arm decides.
2443 let p = profile(Some(0), 0, HOLD, Expiry::ResetStream { code: 0x2A });
2444 let counters = Arc::new(Recorder::new());
2445 let stats = Arc::new(ShapeRecorder::for_profile(Some(&p)));
2446 let shaper = Arc::new(Scheduler::new(p));
2447
2448 let mut up = PendingQueue::new(EgressConfig::default(), Arc::clone(&counters)).with_shaper(
2449 Some(Arc::clone(&shaper)),
2450 Arc::clone(&stats),
2451 ProxySide::ClientToProxy,
2452 );
2453 let mut down = PendingQueue::new(EgressConfig::default(), counters).with_shaper(
2454 Some(shaper),
2455 Arc::clone(&stats),
2456 ProxySide::RelayToProxy,
2457 );
2458
2459 let now = Instant::now();
2460 // Untagged, so these are the units no rule could see — the ones
2461 // charged to `bytes_shaped` as they are queued. Different sizes on
2462 // the two legs, so a leg that counted pushes rather than bytes is
2463 // separable from one that adds.
2464 up.push(Pending::bytes(payload(16), now));
2465 down.push(Pending::bytes(payload(48), now));
2466 down.push(Pending::bytes(payload(48), now));
2467
2468 // The uplink head outlives its clamp and the stream is abandoned.
2469 assert!(
2470 up.pop_next_due(now + HOLD + Duration::from_millis(1)).is_some(),
2471 "an expired head under ResetStream yields the reset"
2472 );
2473 // The downlink's head is inside its clamp and goes out normally —
2474 // unshapeable bytes charge no bucket — so nothing on that leg
2475 // expires and nothing on it is reset.
2476 assert!(down.pop_next_due(now).is_some(), "the downlink head is not past its clamp");
2477
2478 let snap = stats.snapshot();
2479 assert_eq!(
2480 snap.uplink,
2481 DirectionStats {
2482 objects_seen: 0,
2483 bytes_shaped: 16,
2484 objects_expired: 1,
2485 streams_reset_by_shaping: 1,
2486 streams_with_mixed_classes: 0,
2487 },
2488 "the uplink queue's bytes and its expiry are the uplink's"
2489 );
2490 assert_eq!(
2491 snap.downlink,
2492 DirectionStats {
2493 objects_seen: 0,
2494 bytes_shaped: 96,
2495 objects_expired: 0,
2496 streams_reset_by_shaping: 0,
2497 streams_with_mixed_classes: 0,
2498 },
2499 "the downlink queued more and gave nothing up: a stall on one leg \
2500 must not be reported on the other"
2501 );
2502 assert_eq!(snap.bytes_shaped, 112, "the aggregate is the two legs and nothing else");
2503 assert_eq!(snap.objects_expired, 1);
2504 assert_eq!(snap.streams_reset_by_shaping, 1);
2505 }
2506
2507 /// A unit queued behind a unit of a **different** class is counted once,
2508 /// against its own class, and never against `tokens_exhausted_episodes`
2509 /// — which belongs to the class whose bucket was actually dry.
2510 ///
2511 /// Two causes of waiting, two counters. Conflating them is what makes
2512 /// head-of-line blocking read as configured shaping.
2513 ///
2514 /// *Ablation:* charge `note_starved` to the *head's* class instead of
2515 /// the waiting unit's — the two assertions swap and both redden.
2516 #[test]
2517 fn a_unit_behind_another_class_is_counted_once_and_separately() {
2518 // Two classes, one bucket that grants nothing, so the head parks.
2519 let queue =
2520 QueueConfig { max_hold: Some(Duration::from_secs(60)), ..QueueConfig::default() };
2521 let p = ShapeProfile::try_new(
2522 vec![test_bucket(Some(0), 0)],
2523 vec![test_class("head", 0), test_class("behind", 0)],
2524 queue,
2525 Discipline::Fifo,
2526 )
2527 .expect("two uniquely named classes over one bucket");
2528
2529 let counters = Arc::new(Recorder::new());
2530 let stats = Arc::new(ShapeRecorder::for_profile(Some(&p)));
2531 let mut q = PendingQueue::new(EgressConfig::default(), counters).with_shaper(
2532 Some(Arc::new(Scheduler::new(p))),
2533 Arc::clone(&stats),
2534 ProxySide::ClientToProxy,
2535 );
2536
2537 let now = Instant::now();
2538 q.tag_unit(Class::Rule(0));
2539 q.push(Pending::bytes(payload(16), now));
2540 q.tag_unit(Class::Rule(1));
2541 q.push(Pending::bytes(payload(16), now));
2542 q.push(Pending::bytes(payload(16), now));
2543
2544 // Three refusals; the sweep must still charge each waiting unit once.
2545 for _ in 0..3 {
2546 assert!(q.pop_next_due(now).is_none());
2547 }
2548 let snap = stats.snapshot();
2549 assert_eq!(
2550 snap.classes[1].starved_behind_other_class, 2,
2551 "both units behind the other class's head are counted, once each"
2552 );
2553 assert_eq!(
2554 snap.classes[0].starved_behind_other_class, 0,
2555 "the head is not waiting behind anybody"
2556 );
2557 assert_eq!(
2558 snap.classes[0].tokens_exhausted_episodes, 1,
2559 "the dry bucket is the head's own, and it is one episode"
2560 );
2561 }
2562
2563 /// The demand a queue declares is withdrawn when it is dropped.
2564 ///
2565 /// Not a nicety: a declaration that outlives its stream starves every
2566 /// lower-priority class for the rest of the session, and the symptom is
2567 /// a *different* stream stalling to `max_hold`.
2568 ///
2569 /// *Ablation:* delete the `Drop` impl — the final `is_none` assertion
2570 /// reddens with `Starved(..)`, on a stream that no longer exists.
2571 #[test]
2572 fn dropping_a_queue_withdraws_the_demand_it_declared() {
2573 let queue =
2574 QueueConfig { max_hold: Some(Duration::from_secs(60)), ..QueueConfig::default() };
2575 let p = ShapeProfile::try_new(
2576 vec![test_bucket(None, 0)],
2577 vec![test_class("hi", 9), test_class("lo", 0)],
2578 queue,
2579 Discipline::StrictPriority,
2580 )
2581 .expect("two uniquely named classes over one bucket");
2582
2583 let counters = Arc::new(Recorder::new());
2584 let stats = Arc::new(ShapeRecorder::for_profile(Some(&p)));
2585 let shaper = Arc::new(Scheduler::new(p));
2586
2587 let now = Instant::now();
2588 {
2589 let mut hi = PendingQueue::new(EgressConfig::default(), Arc::clone(&counters))
2590 .with_shaper(
2591 Some(Arc::clone(&shaper)),
2592 Arc::clone(&stats),
2593 ProxySide::ClientToProxy,
2594 );
2595 hi.tag_unit(Class::Rule(0));
2596 // Due a minute out, so it stays queued and keeps declaring.
2597 hi.push(Pending::bytes(payload(16), now + Duration::from_secs(60)));
2598 assert!(
2599 matches!(shaper.acquire(Class::Rule(1), 16, now), Acquire::Starved(_)),
2600 "the high class is holding the bucket while its queue is alive"
2601 );
2602 }
2603 assert!(
2604 matches!(shaper.acquire(Class::Rule(1), 16, now), Acquire::Now),
2605 "the high class's stream is gone, so nothing is holding the low one back"
2606 );
2607 }
2608
2609 #[test]
2610 fn an_empty_queue_waits_for_nothing_and_accepts_more() {
2611 let (q, _c) = queue();
2612 assert!(q.is_empty());
2613 assert_eq!(q.len(), 0);
2614 assert_eq!(q.queued_bytes(), 0);
2615 assert!(q.accepts_more());
2616 assert!(q.head_release().is_none());
2617 }
2618
2619 #[test]
2620 fn a_later_unit_cannot_overtake_an_earlier_one() {
2621 let (mut q, _c) = queue();
2622 let now = Instant::now();
2623 let head =
2624 q.push(Pending::bytes(Bytes::from_static(b"0"), now + Duration::from_millis(300)));
2625 // A `Pass` — no delay at all — pushed behind it.
2626 let behind = q.push(Pending::bytes(Bytes::from_static(b"1"), now));
2627 assert_eq!(
2628 behind.release_at, head.release_at,
2629 "the reported release is the queue's estimate, and it accounts for the head",
2630 );
2631 // And a third with a smaller delay than the head's.
2632 let third =
2633 q.push(Pending::bytes(Bytes::from_static(b"2"), now + Duration::from_millis(10)));
2634 assert_eq!(third.release_at, head.release_at);
2635 assert_eq!(q.len(), 3);
2636 assert_eq!(q.queued_bytes(), 3);
2637 }
2638
2639 /// The estimate is reporting; `due_at` is readiness. Conflating the two
2640 /// is what stranded whole streams behind a released `Hold`.
2641 ///
2642 /// *Ablation (run, and it fails):* in [`PendingQueue::push`], clamp
2643 /// `unit.due_at` the way `unit.expected_at` is clamped. Units 1 and 2
2644 /// then report the head's 300 ms deadline as their own.
2645 #[test]
2646 fn the_queues_estimate_never_rewrites_a_units_own_deadline() {
2647 let (mut q, _c) = queue();
2648 let now = Instant::now();
2649 q.push(Pending::bytes(Bytes::from_static(b"0"), now + Duration::from_millis(300)));
2650 q.push(Pending::bytes(Bytes::from_static(b"1"), now));
2651 q.push(Pending::bytes(Bytes::from_static(b"2"), now + Duration::from_millis(10)));
2652
2653 // Popped at a time past everything, so the deque hands all three
2654 // back and each can be asked what it was actually built with.
2655 let far = now + Duration::from_secs(1);
2656 let units: Vec<Pending> = std::iter::from_fn(|| q.pop_next_due(far)).collect();
2657 assert_eq!(units.len(), 3);
2658 assert_eq!(units[0].due_at(), now + Duration::from_millis(300));
2659 assert_eq!(units[1].due_at(), now, "a `Pass` behind a delay keeps its own `now`");
2660 assert_eq!(units[2].due_at(), now + Duration::from_millis(10));
2661 // The estimate is the clamped one, on every unit.
2662 for unit in &units {
2663 assert_eq!(unit.expected_at(), units[0].due_at());
2664 }
2665 }
2666
2667 #[test]
2668 fn pushing_bumps_the_egress_counter_once_per_unit() {
2669 let (mut q, counters) = queue();
2670 let now = Instant::now();
2671 for _ in 0..4 {
2672 q.push(Pending::bytes(Bytes::from_static(b"x"), now));
2673 }
2674 assert_eq!(counters.snapshot().egress_items_queued, 4);
2675 }
2676
2677 #[test]
2678 fn a_unit_that_is_due_pops_and_one_that_is_not_does_not() {
2679 let (mut q, _c) = queue();
2680 let now = Instant::now();
2681 q.push(Pending::bytes(Bytes::from_static(b"soon"), now));
2682 q.push(Pending::bytes(Bytes::from_static(b"later"), now + Duration::from_secs(60)));
2683 let first = q.pop_next_due(now).expect("head is due");
2684 assert_eq!(first.len(), 4);
2685 assert!(q.pop_next_due(now).is_none(), "the tail is a minute out");
2686 assert_eq!(q.queued_bytes(), 5);
2687 assert!(q.head_release().is_some());
2688 }
2689
2690 #[test]
2691 fn a_released_gate_makes_a_unit_due_before_its_ceiling() {
2692 let (mut q, _c) = queue();
2693 let now = Instant::now();
2694 let gate = Gate::new();
2695 // The `max_hold` ceiling: 30 s out.
2696 q.push(
2697 Pending::bytes(Bytes::from_static(b"held"), hold_ceiling(now, q.config()))
2698 .with_gate(gate.clone()),
2699 );
2700 assert!(q.pop_next_due(Instant::now()).is_none());
2701 gate.release();
2702 let unit = q
2703 .pop_next_due(Instant::now())
2704 .expect("a released gate is due, or the release arm spins until max_hold");
2705 assert_eq!(unit.len(), 4);
2706 }
2707
2708 /// Releasing a gate frees the held unit **and the run behind it**.
2709 ///
2710 /// The head carries the gate; the three units behind it carry none, so
2711 /// the only thing that can make them due is their own deadline. They
2712 /// have one — `now` — and it is theirs, which is the whole point. When
2713 /// the queue stamped them with the hold's `max_hold` ceiling instead,
2714 /// one gate release wrote one object and stranded every object behind
2715 /// it, plus the stream's FIN, for 30 s.
2716 ///
2717 /// *Ablation (run, and it fails):* in [`PendingQueue::push`], write
2718 /// `unit.due_at = unit.due_at.max(tail).max(Instant::now())` alongside
2719 /// the `expected_at` line. Only the head pops, and `drained` is
2720 /// `[4]`.
2721 #[test]
2722 fn releasing_a_gate_frees_the_whole_run_queued_behind_it() {
2723 let (mut q, _c) = queue();
2724 let now = Instant::now();
2725 let gate = Gate::new();
2726 q.push(
2727 Pending::bytes(Bytes::from_static(b"held"), hold_ceiling(now, q.config()))
2728 .with_gate(gate.clone()),
2729 );
2730 // Three plain `Pass` units, decided after the hold, so they are
2731 // pushed for ordering and nothing else.
2732 for tail in [&b"1"[..], b"2", b"3"] {
2733 q.push(Pending::bytes(Bytes::copy_from_slice(tail), Instant::now()));
2734 }
2735 assert_eq!(q.len(), 4);
2736 assert!(q.pop_next_due(Instant::now()).is_none(), "the head holds the whole run");
2737
2738 gate.release();
2739 let at = Instant::now();
2740 let drained: Vec<usize> =
2741 std::iter::from_fn(|| q.pop_next_due(at)).map(|u| u.len()).collect();
2742 assert_eq!(
2743 drained,
2744 vec![4, 1, 1, 1],
2745 "one wake on the gate must free the head *and* everything behind it, in order",
2746 );
2747 assert!(q.is_empty(), "nothing may be left for the max_hold ceiling to release");
2748 assert_eq!(q.queued_bytes(), 0);
2749 }
2750
2751 /// A `Delay` queued behind a `Hold` keeps its **own** deadline.
2752 /// `Action::Delay` documents `arrived_at + by` — *a deadline, not a
2753 /// spacing* — and a `Hold` in front of it is not allowed to redefine that
2754 /// as the hold's ceiling. The unit is not due at the gate release, because
2755 /// its own 50 ms have not passed; it *is* due 50 ms later, rather than at
2756 /// the 20 s ceiling.
2757 ///
2758 /// *Ablation (run, and it fails):* in [`PendingQueue::push`], clamp
2759 /// `unit.due_at` as well as `unit.expected_at`. The second pop at
2760 /// `+60 ms` returns `None` and the unit waits out `max_hold`.
2761 #[test]
2762 fn a_delay_queued_behind_a_hold_keeps_its_own_deadline() {
2763 const CEILING: Duration = Duration::from_secs(20);
2764 const BY: Duration = Duration::from_millis(50);
2765
2766 let config = EgressConfig { max_hold: CEILING, ..EgressConfig::default() };
2767 let (mut q, _c) = queue_with(config);
2768 let now = Instant::now();
2769 let gate = Gate::new();
2770 q.push(
2771 Pending::bytes(Bytes::from_static(b"held"), hold_ceiling(now, q.config()))
2772 .with_gate(gate.clone()),
2773 );
2774 let queued = q.push(Pending::bytes(Bytes::from_static(b"late"), now + BY));
2775 assert!(
2776 queued.release_at >= now + CEILING,
2777 "the *estimate* stays conservative — the queue cannot know when a gate opens",
2778 );
2779
2780 gate.release();
2781 let head = q.pop_next_due(now + Duration::from_millis(10)).expect("a released gate is due");
2782 assert_eq!(head.len(), 4);
2783 assert!(
2784 q.pop_next_due(now + Duration::from_millis(10)).is_none(),
2785 "the delayed unit's own deadline still governs it: 10 ms is inside its {BY:?}",
2786 );
2787 let second = q
2788 .pop_next_due(now + BY + Duration::from_millis(10))
2789 .expect("arrived_at + by has passed, and that is the whole of the deadline");
2790 assert_eq!(second.len(), 4);
2791 assert_eq!(second.due_at(), now + BY, "its deadline was never rewritten");
2792 }
2793
2794 #[test]
2795 fn the_byte_budget_stops_the_read_branch_and_reports_once() {
2796 // A struct literal is legal here and not in `tests/`: this file is
2797 // inside the defining crate, where `#[non_exhaustive]` does not
2798 // apply.
2799 let config = EgressConfig { max_pending_bytes: 8, ..EgressConfig::default() };
2800 let (mut q, _c) = queue_with(config);
2801 let now = Instant::now();
2802 let first = q.push(Pending::bytes(Bytes::from_static(b"1234"), now));
2803 assert!(!first.entered_backpressure);
2804 assert!(q.accepts_more());
2805 let second = q.push(Pending::bytes(Bytes::from_static(b"5678"), now));
2806 assert!(second.entered_backpressure, "the transition into backpressure is reported");
2807 assert!(!q.accepts_more());
2808 let third = q.push(Pending::bytes(Bytes::from_static(b"9"), now));
2809 assert!(!third.entered_backpressure, "once per stream, not once per push");
2810 }
2811
2812 #[test]
2813 fn the_unit_cap_bounds_a_queue_of_zero_byte_units() {
2814 let (mut q, _c) = queue();
2815 let now = Instant::now();
2816 let mut reports = 0;
2817 for _ in 0..=MAX_PENDING_UNITS {
2818 if q.push(Pending::elided(now)).entered_backpressure {
2819 reports += 1;
2820 }
2821 }
2822 assert_eq!(q.queued_bytes(), 0, "elided units carry no bytes at all");
2823 assert!(!q.accepts_more(), "the byte budget alone would never stop this");
2824 assert_eq!(reports, 1);
2825 }
2826
2827 #[tokio::test]
2828 async fn an_elided_unit_writes_nothing_but_keeps_its_slot() {
2829 let (mut q, _c) = queue();
2830 let now = Instant::now();
2831 q.push(Pending::bytes(Bytes::from_static(b"a"), now));
2832 q.push(Pending::elided(now));
2833 q.push(Pending::bytes(Bytes::from_static(b"c"), now));
2834 let mut sink = RecordingSink::default();
2835 assert_eq!(q.drain_ignoring_release_times(&mut sink).await, DrainOutcome::Complete);
2836 assert_eq!(sink.written(), Bytes::from_static(b"ac"));
2837 assert_eq!(sink.writes.len(), 2);
2838 }
2839
2840 #[tokio::test]
2841 async fn a_truncate_terminal_writes_its_prefix_then_resets() {
2842 let (mut q, _c) = queue();
2843 let now = Instant::now();
2844 q.push(Pending::bytes(Bytes::from_static(b"ahead"), now));
2845 q.push(Pending::terminal(Terminal::Truncate {
2846 prefix: Bytes::from_static(b"pre"),
2847 code: 0x2,
2848 }));
2849 q.push(Pending::bytes(Bytes::from_static(b"never"), now));
2850 let mut sink = RecordingSink::default();
2851 let outcome = q.drain_ignoring_release_times(&mut sink).await;
2852 assert_eq!(outcome, DrainOutcome::Terminated { forwarded: 3, code: 0x2 });
2853 assert_eq!(sink.written(), Bytes::from_static(b"aheadpre"));
2854 assert_eq!(sink.reset_code, Some(0x2));
2855 assert!(q.is_empty(), "nothing behind a terminal is written");
2856 assert_eq!(q.queued_bytes(), 0);
2857 }
2858
2859 #[tokio::test]
2860 async fn a_reset_terminal_writes_nothing() {
2861 let mut sink = RecordingSink::default();
2862 let written = write_unit(Pending::terminal(Terminal::Reset { code: 7 }), &mut sink)
2863 .await
2864 .expect("reset never fails the write");
2865 assert_eq!(written, Written::Terminated { forwarded: 0, code: 7 });
2866 assert!(sink.writes.is_empty());
2867 assert_eq!(sink.reset_code, Some(7));
2868 }
2869
2870 #[tokio::test]
2871 async fn a_failed_drain_leaves_what_it_could_not_write_queued() {
2872 let (mut q, _c) = queue();
2873 let now = Instant::now();
2874 q.push(Pending::bytes(Bytes::from_static(b"aa"), now));
2875 q.push(Pending::bytes(Bytes::from_static(b"bbb"), now));
2876 q.push(Pending::bytes(Bytes::from_static(b"cccc"), now));
2877 let mut sink = RecordingSink { fail_after: Some(1), ..RecordingSink::default() };
2878 assert_eq!(q.drain_ignoring_release_times(&mut sink).await, DrainOutcome::WriteFailed);
2879 assert_eq!(sink.written(), Bytes::from_static(b"aa"));
2880 assert_eq!(q.len(), 2);
2881 // This is what `Impairment { QueuedBytesAtTeardown { bytes } }` carries.
2882 assert_eq!(q.queued_bytes(), 7);
2883 }
2884
2885 #[tokio::test]
2886 async fn the_honouring_drain_writes_in_order_at_release_time() {
2887 let (mut q, counters) = queue();
2888 let now = Instant::now();
2889 q.push(Pending::bytes(Bytes::from_static(b"0"), now + Duration::from_millis(30)));
2890 q.push(Pending::bytes(Bytes::from_static(b"1"), now));
2891 q.push(Pending::bytes(Bytes::from_static(b"2"), now));
2892 let cancel = CancellationToken::new();
2893 let mut sink = RecordingSink::default();
2894 let started = Instant::now();
2895 let outcome = drain_honouring_release_times(&mut q, &mut sink, &cancel, no_shape_reports)
2896 .await
2897 .expect("no write failed");
2898 assert_eq!(outcome, DrainOutcome::Complete);
2899 assert!(started.elapsed() >= Duration::from_millis(30), "release times were honoured");
2900 assert_eq!(sink.written(), Bytes::from_static(b"012"));
2901 assert_eq!(counters.snapshot().release_errors.count, 0, "drains are never release samples");
2902 assert_eq!(
2903 q.unconfirmed_bytes(),
2904 0,
2905 "a drain that ran to completion at its release times is not a teardown and owes \
2906 no `QueuedBytesAtTeardown`",
2907 );
2908 }
2909
2910 /// A teardown flush reports what it handed to a dying transport.
2911 ///
2912 /// The teardown-race loss, at queue level. `drain_ignoring_release_times`
2913 /// writes the held unit and `write_all` returns `Ok` — quinn buffered it —
2914 /// so **nothing is left queued**, and a report built on
2915 /// [`PendingQueue::queued_bytes`] says zero for exactly the case where
2916 /// `run_with_transport`'s `close()` then discards the buffer and the
2917 /// object never reaches the peer. `a_held_object_is_never_silently_lost_at_teardown`
2918 /// is the same claim over real QUIC; this is the mechanism, with no
2919 /// scheduler in it.
2920 ///
2921 /// *Ablation (run, and it fails):* make
2922 /// [`PendingQueue::unconfirmed_bytes`] return `self.queued_bytes`. It
2923 /// reports 0 and the last assertion goes red — which is precisely the
2924 /// silent loss, spelled out.
2925 #[tokio::test]
2926 async fn a_teardown_flush_reports_what_it_handed_to_a_dying_transport() {
2927 let (mut q, _c) = queue();
2928 // A hold nobody will release, so only the cancel fallback can move
2929 // it.
2930 q.push(
2931 Pending::bytes(Bytes::from_static(b"gone"), hold_ceiling(Instant::now(), q.config()))
2932 .with_gate(Gate::new()),
2933 );
2934 let cancel = CancellationToken::new();
2935 cancel.cancel();
2936 let mut sink = RecordingSink::default();
2937
2938 let outcome = drain_honouring_release_times(&mut q, &mut sink, &cancel, no_shape_reports)
2939 .await
2940 .expect("the fallback drain swallows write failures");
2941 assert_eq!(outcome, DrainOutcome::CancelledMidDrain);
2942 assert_eq!(sink.written(), Bytes::from_static(b"gone"), "delivered late beats lost");
2943 assert_eq!(q.queued_bytes(), 0, "the fallback handed everything to the transport");
2944 assert_eq!(
2945 q.unconfirmed_bytes(),
2946 4,
2947 "…and handing bytes to a transport the session is closing is not delivering \
2948 them: this is what `QueuedBytesAtTeardown` has to carry, or the object is \
2949 gone with no event at all",
2950 );
2951 }
2952
2953 #[tokio::test]
2954 async fn a_teardown_flush_that_could_not_write_reports_both_halves() {
2955 let (mut q, _c) = queue();
2956 let now = Instant::now();
2957 q.push(Pending::bytes(Bytes::from_static(b"aa"), now));
2958 q.push(Pending::bytes(Bytes::from_static(b"bbb"), now));
2959 q.push(Pending::bytes(Bytes::from_static(b"cccc"), now));
2960 let mut sink = RecordingSink { fail_after: Some(1), ..RecordingSink::default() };
2961 assert_eq!(q.drain_ignoring_release_times(&mut sink).await, DrainOutcome::WriteFailed);
2962 assert_eq!(q.queued_bytes(), 7, "what never reached the transport");
2963 assert_eq!(q.unconfirmed_bytes(), 9, "…plus the two bytes that did, and may be lost");
2964 }
2965
2966 #[tokio::test]
2967 async fn cancelling_a_drain_that_is_waiting_on_a_gate_falls_back_at_once() {
2968 let (mut q, _c) = queue();
2969 let now = Instant::now();
2970 // A hold nobody will ever release, at the 30 s ceiling.
2971 q.push(
2972 Pending::bytes(Bytes::from_static(b"held"), hold_ceiling(now, q.config()))
2973 .with_gate(Gate::new()),
2974 );
2975 let cancel = CancellationToken::new();
2976 let waker = cancel.clone();
2977 tokio::spawn(async move {
2978 tokio::time::sleep(Duration::from_millis(20)).await;
2979 waker.cancel();
2980 });
2981 let mut sink = RecordingSink::default();
2982 let started = Instant::now();
2983 let outcome = drain_honouring_release_times(&mut q, &mut sink, &cancel, no_shape_reports)
2984 .await
2985 .expect("no write failed");
2986 assert_eq!(outcome, DrainOutcome::CancelledMidDrain);
2987 assert!(
2988 started.elapsed() < Duration::from_secs(2),
2989 "teardown must not wait out max_hold; took {:?}",
2990 started.elapsed(),
2991 );
2992 assert_eq!(
2993 sink.written(),
2994 Bytes::from_static(b"held"),
2995 "delivered late beats lost silently"
2996 );
2997 assert_eq!(q.queued_bytes(), 0);
2998 }
2999
3000 #[tokio::test]
3001 async fn an_already_cancelled_drain_does_not_spin() {
3002 let (mut q, _c) = queue();
3003 q.push(Pending::bytes(Bytes::from_static(b"x"), Instant::now() + Duration::from_secs(60)));
3004 let cancel = CancellationToken::new();
3005 cancel.cancel();
3006 let mut sink = RecordingSink::default();
3007 let started = Instant::now();
3008 let outcome = drain_honouring_release_times(&mut q, &mut sink, &cancel, no_shape_reports)
3009 .await
3010 .unwrap();
3011 assert_eq!(outcome, DrainOutcome::CancelledMidDrain);
3012 assert!(started.elapsed() < Duration::from_secs(1), "the biased arm must win");
3013 assert_eq!(sink.written(), Bytes::from_static(b"x"));
3014 }
3015
3016 #[tokio::test]
3017 async fn wait_release_resolves_on_the_gate_the_deadline_or_the_cancel() {
3018 let (mut q, _c) = queue();
3019 let far = Instant::now() + Duration::from_secs(60);
3020
3021 // 1. The gate.
3022 let gate = Gate::new();
3023 q.push(Pending::bytes(Bytes::from_static(b"g"), far).with_gate(gate.clone()));
3024 gate.release();
3025 let cancel = CancellationToken::new();
3026 wait_release(q.head_release(), &cancel).await;
3027
3028 // 2. The deadline.
3029 let (mut q, _c) = queue();
3030 q.push(Pending::bytes(
3031 Bytes::from_static(b"d"),
3032 Instant::now() + Duration::from_millis(20),
3033 ));
3034 let started = Instant::now();
3035 wait_release(q.head_release(), &cancel).await;
3036 assert!(started.elapsed() >= Duration::from_millis(20));
3037
3038 // 3. The session cancel.
3039 let (mut q, _c) = queue();
3040 q.push(Pending::bytes(Bytes::from_static(b"c"), far));
3041 cancel.cancel();
3042 wait_release(q.head_release(), &cancel).await;
3043
3044 // 4. `None` parks until cancellation rather than resolving at once.
3045 wait_release(None, &cancel).await;
3046 }
3047
3048 #[tokio::test]
3049 async fn the_pipe_loop_shape_from_the_contract_compiles_and_keeps_order() {
3050 // The pipe loop's `select!`, with the read half standing in for
3051 // `recv.read` and a recording sink for `send`. What it proves is
3052 // the shape: both branch expressions are borrow-free, the read
3053 // guard and the release branch coexist, and a delayed head holds
3054 // everything behind it.
3055 let (mut pending, counters) = queue();
3056 let cancel = CancellationToken::new();
3057 let mut sink = RecordingSink::default();
3058 let (tx, mut rx) = tokio::sync::mpsc::channel::<Bytes>(8);
3059
3060 tokio::spawn(async move {
3061 for chunk in [&b"0"[..], b"1", b"2", b"3"] {
3062 tx.send(Bytes::from_static(chunk)).await.expect("receiver lives");
3063 }
3064 // Stay open past the head's release time, or the source FINs
3065 // first and the FIN drain — not the release branch — is what
3066 // writes everything.
3067 tokio::time::sleep(Duration::from_millis(200)).await;
3068 });
3069
3070 let mut first = true;
3071 loop {
3072 let can_read = pending.accepts_more();
3073 let head_release = pending.head_release();
3074
3075 tokio::select! {
3076 chunk = rx.recv(), if can_read => {
3077 match chunk {
3078 Some(raw) => {
3079 // Object 0 is delayed; 1.. are not, and must
3080 // still not overtake it.
3081 let release_at = if std::mem::take(&mut first) {
3082 Instant::now() + Duration::from_millis(40)
3083 } else {
3084 Instant::now()
3085 };
3086 if pending.is_empty() && release_at <= Instant::now() {
3087 // Today's fast path: written inline in the read arm.
3088 write_unit(Pending::bytes(raw, release_at), &mut sink)
3089 .await
3090 .expect("recording sink");
3091 } else {
3092 pending.push(Pending::bytes(raw, release_at));
3093 }
3094 }
3095 None => {
3096 let outcome =
3097 drain_honouring_release_times(&mut pending, &mut sink, &cancel, no_shape_reports)
3098 .await
3099 .expect("recording sink");
3100 assert_eq!(outcome, DrainOutcome::Complete);
3101 break;
3102 }
3103 }
3104 }
3105 () = wait_release(head_release.clone(), &cancel), if head_release.is_some() => {
3106 let now = Instant::now();
3107 while let Some(unit) = pending.pop_next_due(now) {
3108 pending.record_release(&unit, now);
3109 write_unit(unit, &mut sink).await.expect("recording sink");
3110 }
3111 }
3112 () = cancel.cancelled() => unreachable!("nothing cancels this test"),
3113 }
3114 }
3115
3116 assert_eq!(
3117 sink.written(),
3118 Bytes::from_static(b"0123"),
3119 "byte equality is the ordering assertion"
3120 );
3121 let snap = counters.snapshot();
3122 assert_eq!(snap.egress_items_queued, 4, "every unit went through the deque");
3123 assert_eq!(
3124 snap.release_errors.count, 4,
3125 "all four share the head's clamped release time, so one wake pops all four \
3126 — and only the release branch samples",
3127 );
3128 }
3129
3130 #[test]
3131 fn defer_by_is_a_deadline_and_clamps_to_max_hold() {
3132 let config =
3133 EgressConfig { max_hold: Duration::from_millis(100), ..EgressConfig::default() };
3134 let now = Instant::now();
3135
3136 let short = defer_by(now, Duration::from_millis(10), &config);
3137 assert_eq!(short.release_at, now + Duration::from_millis(10));
3138 assert!(!short.was_clamped());
3139
3140 let long = defer_by(now, Duration::from_secs(5), &config);
3141 assert_eq!(long.release_at, now + Duration::from_millis(100));
3142 assert!(long.was_clamped());
3143 assert_eq!(long.requested, Duration::from_secs(5));
3144 assert_eq!(long.applied, Duration::from_millis(100));
3145
3146 // Two units arriving together with the same `by` release together,
3147 // not at +by and +2by.
3148 let a = defer_by(now, Duration::from_millis(10), &config);
3149 let b = defer_by(now, Duration::from_millis(10), &config);
3150 assert_eq!(a.release_at, b.release_at);
3151
3152 assert_eq!(hold_ceiling(now, &config), now + Duration::from_millis(100));
3153 }
3154
3155 #[test]
3156 fn the_first_close_request_wins() {
3157 let cancel = CancellationToken::new();
3158 let closer = SessionCloser::new(cancel.clone());
3159 assert!(!closer.is_closing());
3160 assert_eq!(
3161 closer.close_args(),
3162 (0, Bytes::from_static(b"proxy session ended")),
3163 "the default is what run_with_transport has always sent",
3164 );
3165
3166 assert!(closer.request(3, Bytes::from_static(b"protocol violation")));
3167 assert!(cancel.is_cancelled());
3168 assert!(closer.is_closing());
3169 assert!(!closer.request(1, Bytes::from_static(b"too late")));
3170 assert_eq!(closer.close_args(), (3, Bytes::from_static(b"protocol violation")));
3171 assert_eq!(
3172 closer.requested(),
3173 Some((3, Bytes::from_static(b"protocol violation"), CloseOrigin::Hook)),
3174 "a close that arrived through `request` is a hook's, and the session's own \
3175 `SessionEnded` reason says so in those words",
3176 );
3177 }
3178
3179 #[test]
3180 fn an_undelayed_head_arms_a_deadline_that_never_touches_the_wheel() {
3181 // `arm_at` on an already-passed instant returns an
3182 // already-cancelled `Deadline` without constructing the wheel, so
3183 // a session that never delays never links it into executed code.
3184 //
3185 // Asserted on the deadline's own token rather than on
3186 // `release_timer::started()`: that is process-global state, and
3187 // another test in this binary may legitimately have started the
3188 // wheel on another thread between the two reads.
3189 let (mut q, _c) = queue();
3190 q.push(Pending::bytes(Bytes::from_static(b"now"), Instant::now()));
3191 let release = q.head_release().expect("one unit queued");
3192 assert!(
3193 release.deadline().token().is_cancelled(),
3194 "an already-due head must not be registered with the wheel",
3195 );
3196
3197 let (mut q, _c) = queue();
3198 q.push(Pending::bytes(
3199 Bytes::from_static(b"later"),
3200 Instant::now() + Duration::from_secs(60),
3201 ));
3202 let release = q.head_release().expect("one unit queued");
3203 assert!(!release.deadline().token().is_cancelled(), "a future head is registered");
3204 assert!(release_timer::started(), "…and registering is what starts the wheel");
3205 }
3206
3207 // ── The session-wide byte gauge ─────────────────────────────────
3208
3209 /// A queue reporting into a gauge, plus the gauge.
3210 fn gauged() -> (PendingQueue, Arc<EgressGauge>) {
3211 let gauge = EgressGauge::new();
3212 let q = PendingQueue::new(EgressConfig::default(), Arc::new(Recorder::new()))
3213 .with_gauge(Arc::clone(&gauge));
3214 (q, gauge)
3215 }
3216
3217 /// The gauge is the sum of what the queues hold, on every route bytes
3218 /// take out of one.
3219 ///
3220 /// It has to be exact in both directions or the one caller that reads
3221 /// it — a requested close deciding whether it may stop waiting —
3222 /// either stops early on bytes that are still queued, or waits out its
3223 /// whole window on bytes that left long ago.
3224 ///
3225 /// *Ablation, recorded:* drop the `self.credit(self.queued_bytes)` from
3226 /// `PendingQueue`'s `Drop`. The final assertion goes red with
3227 ///
3228 /// ```text
3229 /// assertion `left == right` failed: a queue that goes away takes its
3230 /// remainder with it, or every later close waits out its whole window
3231 /// left: 4
3232 /// right: 0
3233 /// ```
3234 ///
3235 /// — a session total permanently above zero, on a stream that no longer
3236 /// exists.
3237 #[tokio::test]
3238 async fn the_gauge_follows_every_route_bytes_leave_a_queue_by() {
3239 let (mut q, gauge) = gauged();
3240 let now = Instant::now();
3241 assert_eq!(gauge.queued(), 0);
3242
3243 q.push(Pending::bytes(Bytes::from_static(b"aaa"), now));
3244 q.push(Pending::bytes(Bytes::from_static(b"bb"), now));
3245 assert_eq!(gauge.queued(), 5, "two pushes, five bytes");
3246
3247 // Route one: released at its release time.
3248 q.pop_next_due(now).expect("both are due");
3249 assert_eq!(gauge.queued(), 2);
3250
3251 // Route two: flushed by the teardown drain.
3252 let mut sink = RecordingSink::default();
3253 assert_eq!(q.drain_ignoring_release_times(&mut sink).await, DrainOutcome::Complete);
3254 assert_eq!(gauge.queued(), 0);
3255
3256 // Route three: cleared wholesale, as a terminal or a shaping reset
3257 // does.
3258 q.push(Pending::bytes(Bytes::from_static(b"cccc"), now));
3259 assert_eq!(gauge.queued(), 4);
3260 q.clear();
3261 assert_eq!(gauge.queued(), 0);
3262
3263 // Route four: the stream ends holding bytes.
3264 q.push(Pending::bytes(Bytes::from_static(b"dddd"), now));
3265 assert_eq!(gauge.queued(), 4);
3266 drop(q);
3267 assert_eq!(
3268 gauge.queued(),
3269 0,
3270 "a queue that goes away takes its remainder with it, or every later close waits out \
3271 its whole window"
3272 );
3273 }
3274
3275 /// The wait ends when the queues empty, and answers zero.
3276 #[tokio::test]
3277 async fn the_drain_wait_ends_the_moment_the_queues_empty() {
3278 let (mut q, gauge) = gauged();
3279 let now = Instant::now();
3280 q.push(Pending::bytes(Bytes::from_static(b"payload"), now));
3281
3282 // A generous window, so finishing quickly can only be the queue
3283 // emptying and not the deadline.
3284 let waiting = tokio::spawn({
3285 let gauge = Arc::clone(&gauge);
3286 async move { gauge.wait_idle(Duration::from_secs(30)).await }
3287 });
3288
3289 tokio::task::yield_now().await;
3290 let mut sink = RecordingSink::default();
3291 assert_eq!(q.drain_ignoring_release_times(&mut sink).await, DrainOutcome::Complete);
3292
3293 let stranded = tokio::time::timeout(Duration::from_secs(5), waiting)
3294 .await
3295 .expect("the wait must end on the queue emptying, not on its own deadline")
3296 .expect("the waiting task must not panic");
3297 assert_eq!(stranded, 0, "nothing was left, so nothing is abandoned");
3298 }
3299
3300 /// The wait ends at its deadline holding the exact residue, and an
3301 /// expired window makes the teardown flush write nothing.
3302 ///
3303 /// The second half is what keeps the arithmetic closed. A flush at that
3304 /// point hands bytes to a connection that is about to discard its
3305 /// buffer, so they are neither confirmably delivered nor confirmably
3306 /// lost; declining to write leaves `unconfirmed_bytes` equal to what
3307 /// was abandoned, exactly.
3308 #[tokio::test]
3309 async fn an_expired_window_reports_the_residue_and_writes_nothing() {
3310 let (mut q, gauge) = gauged();
3311 let now = Instant::now();
3312 q.push(Pending::bytes(Bytes::from_static(b"held"), now + Duration::from_secs(60)));
3313
3314 let stranded = gauge.wait_idle(Duration::from_millis(20)).await;
3315 assert_eq!(stranded, 4, "the window closed on four bytes that had not been released");
3316
3317 gauge.begin_discarding();
3318 let mut sink = RecordingSink::default();
3319 assert_eq!(
3320 q.drain_ignoring_release_times(&mut sink).await,
3321 DrainOutcome::Discarded,
3322 "past the deadline the flush declines rather than best-efforts"
3323 );
3324 assert!(sink.writes.is_empty(), "nothing reached the transport");
3325 assert_eq!(
3326 q.unconfirmed_bytes(),
3327 4,
3328 "so the reported figure is what was abandoned, with nothing handed anywhere"
3329 );
3330 assert_eq!(
3331 q.queued_bytes(),
3332 q.unconfirmed_bytes(),
3333 "and the two agree, which is the point"
3334 );
3335 }
3336
3337 /// A close recorded on the closer does **not** end the session.
3338 ///
3339 /// `request` cancels as it records, which is right for a hook: the hook
3340 /// asked for the session to end. A control-plane close has to fix the
3341 /// pair first and let the session keep running through its drain
3342 /// window, and it would have no window at all if recording cancelled.
3343 ///
3344 /// It also checks who the close is *reported as*.
3345 /// `run_with_transport` turns the recorded triple into the prose on
3346 /// [`ProxyEvent::SessionEnded`](crate::event::ProxyEvent::SessionEnded),
3347 /// and both callers reach the same `OnceLock`, so without an origin
3348 /// beside the pair every operator-initiated close was reported to
3349 /// observers as a hook's decision.
3350 ///
3351 /// *Ablation, recorded:* have `record` store `CloseOrigin::Hook` — the
3352 /// state the type was in before this field existed, where the two
3353 /// callers are indistinguishable once they have written. This test goes
3354 /// red with the real message
3355 ///
3356 /// ```text
3357 /// assertion `left == right` failed
3358 /// left: Some((7, b"asked", Hook))
3359 /// right: Some((7, b"asked", ControlPlane))
3360 /// ```
3361 #[test]
3362 fn recording_a_close_leaves_the_session_running() {
3363 let cancel = CancellationToken::new();
3364 let closer = SessionCloser::new(cancel.clone());
3365
3366 assert!(closer.record(7, Bytes::from_static(b"asked")));
3367 assert!(!cancel.is_cancelled(), "the drain window has not even started yet");
3368 assert_eq!(closer.close_args(), (7, Bytes::from_static(b"asked")));
3369
3370 // First writer wins here exactly as it does through `request`, so a
3371 // hook and a control plane racing cannot end up with a code from
3372 // one and a reason from the other.
3373 assert!(!closer.record(9, Bytes::from_static(b"second")));
3374 assert!(!closer.request(9, Bytes::from_static(b"second")));
3375 assert_eq!(closer.close_args(), (7, Bytes::from_static(b"asked")));
3376 assert!(cancel.is_cancelled(), "…and `request` still cancels, losing or not");
3377
3378 // The losing `request` above is a hook's, and it lost. The origin
3379 // has to lose with it: a session whose close was recorded by the
3380 // control plane and then re-asked for by a hook must not report the
3381 // hook's name against the control plane's code and reason, which is
3382 // the mislabel this field exists to stop.
3383 assert_eq!(
3384 closer.requested(),
3385 Some((7, Bytes::from_static(b"asked"), CloseOrigin::ControlPlane)),
3386 );
3387 }
3388}