use std::collections::VecDeque;
use std::future::Future;
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
use std::sync::{Arc, OnceLock};
use std::time::{Duration, Instant};
use bytes::Bytes;
use tokio::sync::Notify;
use tokio_util::sync::CancellationToken;
use moqtap_client::transport::{SendStream, TransportError};
use crate::action::{EgressConfig, Gate};
use crate::error::ProxyError;
use crate::instrument::Recorder;
use crate::release_timer::{self, Deadline};
use crate::shape::{Acquire, Class, Expiry, QueueDepth, Scheduler, ShapeRecorder};
use crate::types::ProxySide;
pub(crate) const MAX_PENDING_UNITS: usize = 8192;
#[derive(Debug, thiserror::Error)]
pub(crate) enum EgressError {
#[error("transport error: {0}")]
Transport(#[from] TransportError),
}
impl From<EgressError> for ProxyError {
fn from(error: EgressError) -> Self {
let EgressError::Transport(source) = error;
ProxyError::Transport(source)
}
}
pub(crate) trait EgressSink {
fn write_all(&mut self, buf: &[u8]) -> impl Future<Output = Result<(), EgressError>> + Send;
fn reset(&mut self, code: u64) -> Result<(), EgressError>;
}
impl EgressSink for SendStream {
async fn write_all(&mut self, buf: &[u8]) -> Result<(), EgressError> {
SendStream::write_all(self, buf).await.map_err(EgressError::from)
}
fn reset(&mut self, code: u64) -> Result<(), EgressError> {
SendStream::reset(self, code).map_err(EgressError::from)
}
}
#[derive(Clone, Debug)]
pub(crate) struct Release {
deadline: Deadline,
gate: Option<Gate>,
}
#[allow(dead_code)]
impl Release {
pub(crate) fn deadline(&self) -> &Deadline {
&self.deadline
}
}
pub(crate) async fn wait_release(release: Option<Release>, cancel: &CancellationToken) {
let Some(release) = release else {
cancel.cancelled().await;
return;
};
match release.gate {
Some(gate) => {
tokio::select! {
() = release.deadline.token().cancelled() => {}
() = gate.wait() => {}
() = cancel.cancelled() => {}
}
}
None => {
tokio::select! {
() = release.deadline.token().cancelled() => {}
() = cancel.cancelled() => {}
}
}
}
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub(crate) enum Terminal {
Truncate {
prefix: Bytes,
code: u64,
},
Reset {
code: u64,
},
}
impl Terminal {
fn len(&self) -> usize {
match self {
Terminal::Truncate { prefix, .. } => prefix.len(),
Terminal::Reset { .. } => 0,
}
}
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub(crate) enum Item {
Write(Bytes),
Elided,
Terminal(Terminal),
}
#[derive(Clone, Copy, Debug)]
struct ShapeTag {
class: Class,
expires_at: Instant,
starved_noted: bool,
}
#[derive(Clone, Debug)]
pub(crate) struct Pending {
due_at: Instant,
expected_at: Instant,
gate: Option<Gate>,
item: Item,
shape: Option<ShapeTag>,
}
impl Pending {
pub(crate) fn bytes(raw: Bytes, due_at: Instant) -> Self {
Self { due_at, expected_at: due_at, gate: None, item: Item::Write(raw), shape: None }
}
pub(crate) fn elided(due_at: Instant) -> Self {
Self { due_at, expected_at: due_at, gate: None, item: Item::Elided, shape: None }
}
pub(crate) fn terminal(terminal: Terminal) -> Self {
let now = Instant::now();
Self {
due_at: now,
expected_at: now,
gate: None,
item: Item::Terminal(terminal),
shape: None,
}
}
#[must_use]
pub(crate) fn with_gate(mut self, gate: Gate) -> Self {
self.gate = Some(gate);
self
}
#[allow(dead_code)]
pub(crate) fn due_at(&self) -> Instant {
self.due_at
}
#[allow(dead_code)]
pub(crate) fn expected_at(&self) -> Instant {
self.expected_at
}
#[allow(dead_code)]
pub(crate) fn item(&self) -> &Item {
&self.item
}
#[allow(dead_code)]
pub(crate) fn is_terminal(&self) -> bool {
matches!(self.item, Item::Terminal(_))
}
pub(crate) fn len(&self) -> usize {
match &self.item {
Item::Write(raw) => raw.len(),
Item::Elided => 0,
Item::Terminal(t) => t.len(),
}
}
fn is_due(&self, now: Instant) -> bool {
self.due_at <= now || self.gate.as_ref().is_some_and(Gate::is_released)
}
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub(crate) enum Written {
Bytes(usize),
Nothing,
Terminated {
forwarded: usize,
code: u64,
},
}
pub(crate) async fn write_unit<S: EgressSink>(
unit: Pending,
send: &mut S,
) -> Result<Written, EgressError> {
match unit.item {
Item::Write(raw) => {
send.write_all(&raw).await?;
Ok(Written::Bytes(raw.len()))
}
Item::Elided => Ok(Written::Nothing),
Item::Terminal(Terminal::Truncate { prefix, code }) => {
if !prefix.is_empty() {
send.write_all(&prefix).await?;
}
let _ = send.reset(code);
Ok(Written::Terminated { forwarded: prefix.len(), code })
}
Item::Terminal(Terminal::Reset { code }) => {
let _ = send.reset(code);
Ok(Written::Terminated { forwarded: 0, code })
}
}
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub(crate) struct Deferral {
pub(crate) release_at: Instant,
pub(crate) requested: Duration,
pub(crate) applied: Duration,
}
impl Deferral {
pub(crate) fn was_clamped(&self) -> bool {
self.applied < self.requested
}
}
pub(crate) fn defer_by(arrived_at: Instant, by: Duration, config: &EgressConfig) -> Deferral {
let applied = by.min(config.max_hold);
Deferral { release_at: arrived_at + applied, requested: by, applied }
}
pub(crate) fn hold_ceiling(arrived_at: Instant, config: &EgressConfig) -> Instant {
arrived_at + config.max_hold
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub(crate) struct Push {
pub(crate) release_at: Instant,
pub(crate) entered_backpressure: bool,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub(crate) enum DrainOutcome {
Complete,
CancelledMidDrain,
Terminated {
forwarded: usize,
code: u64,
},
WriteFailed,
Discarded,
}
#[derive(Debug)]
pub(crate) struct EgressGauge {
queued: AtomicUsize,
idle: Notify,
discarding: AtomicBool,
}
impl EgressGauge {
pub(crate) fn new() -> Arc<Self> {
Arc::new(Self {
queued: AtomicUsize::new(0),
idle: Notify::new(),
discarding: AtomicBool::new(false),
})
}
pub(crate) fn queued(&self) -> usize {
self.queued.load(Ordering::Acquire)
}
pub(crate) fn is_discarding(&self) -> bool {
self.discarding.load(Ordering::Acquire)
}
pub(crate) fn begin_discarding(&self) {
self.discarding.store(true, Ordering::Release);
}
fn add(&self, bytes: usize) {
if bytes > 0 {
self.queued.fetch_add(bytes, Ordering::AcqRel);
}
}
fn sub(&self, bytes: usize) {
if bytes == 0 {
return;
}
let before = self
.queued
.fetch_update(Ordering::AcqRel, Ordering::Acquire, |n| Some(n.saturating_sub(bytes)))
.unwrap_or(0);
if before.saturating_sub(bytes) == 0 {
self.idle.notify_waiters();
}
}
pub(crate) async fn wait_idle(&self, timeout: Duration) -> usize {
let deadline = tokio::time::sleep(timeout);
tokio::pin!(deadline);
loop {
if self.queued() == 0 {
return 0;
}
let notified = self.idle.notified();
tokio::pin!(notified);
notified.as_mut().enable();
if self.queued() == 0 {
return 0;
}
tokio::select! {
() = &mut notified => {}
() = &mut deadline => return self.queued(),
}
}
}
}
#[derive(Debug)]
pub(crate) struct PendingQueue {
q: VecDeque<Pending>,
head_deadline: Option<Deadline>,
queued_bytes: usize,
flushed_unconfirmed: usize,
gauge: Option<Arc<EgressGauge>>,
config: EgressConfig,
counters: Arc<Recorder>,
backpressure_reported: bool,
shape_depth: Option<QueueDepth>,
shaper: Option<Arc<Scheduler>>,
shape_stats: Option<Arc<ShapeRecorder>>,
shape_side: Option<ProxySide>,
shape_hold: Duration,
unit_class: Class,
shape_gate: Option<Gate>,
shape_parked: bool,
shape_demand: Option<Class>,
tokens_dry: bool,
starved_from: usize,
shape_report: Option<ShapeReport>,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub(crate) enum ShapeReport {
Clamped {
requested: Option<Duration>,
applied: Duration,
},
Expired,
BurstBelowUnit {
class: Class,
burst_bytes: u64,
unit_bytes: u64,
},
}
impl PendingQueue {
pub(crate) fn new(config: EgressConfig, counters: Arc<Recorder>) -> Self {
Self {
q: VecDeque::new(),
head_deadline: None,
queued_bytes: 0,
flushed_unconfirmed: 0,
gauge: None,
shape_hold: config.max_hold,
config,
counters,
backpressure_reported: false,
shape_depth: None,
shaper: None,
shape_stats: None,
shape_side: None,
unit_class: Class::Unshapeable,
shape_gate: None,
shape_parked: false,
shape_demand: None,
tokens_dry: false,
starved_from: 1,
shape_report: None,
}
}
pub(crate) fn with_gauge(mut self, gauge: Arc<EgressGauge>) -> Self {
debug_assert_eq!(self.queued_bytes, 0, "a queue takes its gauge before it takes bytes");
self.gauge = Some(gauge);
self
}
fn charge(&mut self, bytes: usize) {
self.queued_bytes = self.queued_bytes.saturating_add(bytes);
if let Some(gauge) = &self.gauge {
gauge.add(bytes);
}
}
fn credit(&mut self, bytes: usize) {
let bytes = bytes.min(self.queued_bytes);
self.queued_bytes -= bytes;
if let Some(gauge) = &self.gauge {
gauge.sub(bytes);
}
}
fn discarding(&self) -> bool {
self.gauge.as_ref().is_some_and(|g| g.is_discarding())
}
pub(crate) fn with_shape_depth(mut self, depth: Option<QueueDepth>) -> Self {
self.shape_depth = depth;
self
}
pub(crate) fn with_shaper(
mut self,
shaper: Option<Arc<Scheduler>>,
stats: Arc<ShapeRecorder>,
side: ProxySide,
) -> Self {
if let Some(shaper) = shaper {
self.shape_hold = shaper.max_hold().unwrap_or(self.config.max_hold);
self.shaper = Some(shaper);
self.shape_stats = Some(stats);
self.shape_side = Some(side);
}
self
}
pub(crate) fn is_shaped(&self) -> bool {
self.shaper.is_some()
}
pub(crate) fn tag_unit(&mut self, class: Class) {
self.unit_class = class;
}
pub(crate) fn take_shape_report(&mut self) -> Option<ShapeReport> {
self.shape_report.take()
}
pub(crate) fn is_empty(&self) -> bool {
self.q.is_empty()
}
pub(crate) fn len(&self) -> usize {
self.q.len()
}
pub(crate) fn queued_bytes(&self) -> usize {
self.queued_bytes
}
pub(crate) fn unconfirmed_bytes(&self) -> usize {
self.flushed_unconfirmed.saturating_add(self.queued_bytes)
}
pub(crate) fn config(&self) -> &EgressConfig {
&self.config
}
pub(crate) fn accepts_more(&self) -> bool {
let within_engine_budget =
self.queued_bytes < self.config.max_pending_bytes && self.q.len() < MAX_PENDING_UNITS;
let within_shape_depth = match self.shape_depth {
Some(depth) => self.queued_bytes < depth.bytes && self.q.len() < depth.objects,
None => true,
};
within_engine_budget && within_shape_depth
}
pub(crate) fn head_release(&self) -> Option<Release> {
let head = self.q.front()?;
debug_assert!(
self.head_deadline.is_some(),
"a non-empty queue always has an armed head deadline",
);
Some(Release {
deadline: self.head_deadline.clone()?,
gate: if self.shape_parked { self.shape_gate.clone() } else { head.gate.clone() },
})
}
fn tail_expected_at(&self) -> Option<Instant> {
self.q.back().map(|u| u.expected_at)
}
fn arm_head(&mut self) {
self.head_deadline = self.q.front().map(|u| release_timer::arm_at(u.due_at));
}
fn arm_head_at(&mut self, at: Instant) {
self.head_deadline = Some(release_timer::arm_at(at));
}
pub(crate) fn push(&mut self, mut unit: Pending) -> Push {
if let Some(tail) = self.tail_expected_at() {
unit.expected_at = unit.due_at.max(tail).max(Instant::now());
}
if self.shaper.is_some() {
unit.shape = Some(ShapeTag {
class: self.unit_class,
expires_at: Instant::now() + self.shape_hold,
starved_noted: false,
});
self.note_unshapeable_seen(&unit);
}
let release_at = unit.expected_at;
let becomes_head = self.q.is_empty();
self.charge(unit.len());
self.q.push_back(unit);
self.counters.note_egress_item_queued();
if becomes_head {
self.arm_head();
}
self.sync_shape_demand();
let entered_backpressure = !self.accepts_more() && !self.backpressure_reported;
if entered_backpressure {
self.backpressure_reported = true;
}
Push { release_at, entered_backpressure }
}
pub(crate) fn pop_next_due(&mut self, now: Instant) -> Option<Pending> {
if !self.q.front().is_some_and(|u| u.is_due(now)) {
return None;
}
self.shape_gate = None;
self.shape_parked = false;
if !self.shaping_grants_head(now) {
return None;
}
let unit = self.q.pop_front()?;
self.credit(unit.len());
self.starved_from = self.starved_from.saturating_sub(1).max(1);
self.note_delivered(&unit);
self.arm_head();
self.sync_shape_demand();
Some(unit)
}
fn shaping_grants_head(&mut self, now: Instant) -> bool {
let Some(shaper) = self.shaper.clone() else {
return true;
};
let Some(head) = self.q.front() else {
return true;
};
let Some(tag) = head.shape else {
return true;
};
if matches!(head.item, Item::Terminal(_)) {
return true;
}
let bytes = head.len() as u64;
if now >= tag.expires_at {
return self.expire_head(shaper.on_expiry());
}
match shaper.acquire(tag.class, bytes, now) {
Acquire::Now => {
self.tokens_dry = false;
true
}
Acquire::Later(at) => {
self.note_tokens_dry(tag.class);
self.park_head(at.min(tag.expires_at), tag, None);
false
}
Acquire::Never => {
self.note_tokens_dry(tag.class);
self.park_head(tag.expires_at, tag, None);
false
}
Acquire::LargerThanBurst { burst_bytes, unit_bytes } => {
self.note_tokens_dry(tag.class);
self.shape_report =
Some(ShapeReport::BurstBelowUnit { class: tag.class, burst_bytes, unit_bytes });
self.park_head(tag.expires_at, tag, None);
false
}
Acquire::Starved(gate) => {
self.park_head(tag.expires_at, tag, Some(gate));
false
}
}
}
fn note_tokens_dry(&mut self, class: Class) {
if self.tokens_dry {
return;
}
self.tokens_dry = true;
if let Some(stats) = &self.shape_stats {
stats.note_tokens_exhausted(class);
}
}
fn park_head(&mut self, at: Instant, tag: ShapeTag, gate: Option<Gate>) {
let by_discipline = gate.is_some();
self.shape_gate = gate;
self.shape_parked = true;
self.arm_head_at(at);
self.sync_shape_demand();
if by_discipline {
self.note_head_starved();
}
self.note_starved_behind(tag.class);
}
fn note_head_starved(&mut self) {
let Some(stats) = self.shape_stats.clone() else { return };
let Some(tag) = self.q.front_mut().and_then(|u| u.shape.as_mut()) else { return };
if tag.starved_noted {
return;
}
tag.starved_noted = true;
stats.note_starved(tag.class);
}
fn note_starved_behind(&mut self, head_class: Class) {
let Some(stats) = self.shape_stats.clone() else { return };
let len = self.q.len();
for index in self.starved_from.max(1)..len {
let Some(unit) = self.q.get_mut(index) else { continue };
let Some(tag) = unit.shape.as_mut() else { continue };
if tag.starved_noted || tag.class == head_class {
continue;
}
tag.starved_noted = true;
stats.note_starved(tag.class);
}
self.starved_from = len.max(1);
}
fn expire_head(&mut self, expiry: Expiry) -> bool {
match expiry {
Expiry::Deliver => {
self.shape_report =
Some(ShapeReport::Clamped { requested: None, applied: self.shape_hold });
true
}
Expiry::ResetStream { code } => {
if let (Some(stats), Some(side)) = (&self.shape_stats, self.shape_side) {
stats.note_expired(side);
stats.note_stream_reset_by_shaping(side);
}
self.shape_report = Some(ShapeReport::Expired);
self.clear();
self.q.push_back(Pending::terminal(Terminal::Reset { code }));
self.arm_head();
true
}
}
}
fn note_unshapeable_seen(&self, unit: &Pending) {
if self.unit_class != Class::Unshapeable {
return;
}
let (Some(stats), Some(side)) = (&self.shape_stats, self.shape_side) else {
return;
};
let bytes = unit.len();
if bytes > 0 {
stats.note_unshapeable_seen(side, bytes as u64);
}
}
fn note_delivered(&self, unit: &Pending) {
let (Some(stats), Some(side), Some(tag)) = (&self.shape_stats, self.shape_side, unit.shape)
else {
return;
};
let bytes = unit.len();
if bytes > 0 {
stats.note_delivered(side, tag.class, bytes as u64);
}
}
fn sync_shape_demand(&mut self) {
let Some(shaper) = &self.shaper else { return };
let want = self.q.front().and_then(|u| u.shape).map(|t| t.class);
if want == self.shape_demand {
return;
}
if let Some(old) = self.shape_demand.take() {
shaper.withdraw_demand(old);
}
if let Some(new) = want {
shaper.declare_demand(new);
self.shape_demand = Some(new);
}
}
pub(crate) fn record_release(&self, unit: &Pending, now: Instant) {
self.counters.record_release(now.saturating_duration_since(unit.expected_at));
}
pub(crate) fn clear(&mut self) {
self.q.clear();
self.credit(self.queued_bytes);
self.head_deadline = None;
self.starved_from = 1;
self.shape_gate = None;
self.shape_parked = false;
self.sync_shape_demand();
}
pub(crate) async fn drain_ignoring_release_times<S: EgressSink>(
&mut self,
send: &mut S,
) -> DrainOutcome {
if self.discarding() {
return DrainOutcome::Discarded;
}
while let Some(unit) = self.q.front().cloned() {
match write_unit(unit, send).await {
Ok(Written::Terminated { forwarded, code }) => {
self.flushed_unconfirmed = self.flushed_unconfirmed.saturating_add(forwarded);
self.clear();
return DrainOutcome::Terminated { forwarded, code };
}
Ok(written) => {
if let Written::Bytes(n) = written {
self.flushed_unconfirmed = self.flushed_unconfirmed.saturating_add(n);
}
let unit = self.q.pop_front().expect("front was just observed");
self.credit(unit.len());
self.starved_from = self.starved_from.saturating_sub(1).max(1);
self.note_delivered(&unit);
self.arm_head();
self.sync_shape_demand();
}
Err(_) => return DrainOutcome::WriteFailed,
}
}
DrainOutcome::Complete
}
}
impl Drop for PendingQueue {
fn drop(&mut self) {
if let (Some(shaper), Some(class)) = (&self.shaper, self.shape_demand.take()) {
shaper.withdraw_demand(class);
}
self.credit(self.queued_bytes);
}
}
pub(crate) async fn drain_honouring_release_times<S, F>(
pending: &mut PendingQueue,
send: &mut S,
cancel: &CancellationToken,
mut on_shape: F,
) -> Result<DrainOutcome, EgressError>
where
S: EgressSink,
F: FnMut(ShapeReport),
{
while let Some(release) = pending.head_release() {
tokio::select! {
biased;
() = cancel.cancelled() => {
return Ok(match pending.drain_ignoring_release_times(send).await {
DrainOutcome::Terminated { forwarded, code } => {
DrainOutcome::Terminated { forwarded, code }
}
_ => DrainOutcome::CancelledMidDrain,
});
}
() = wait_release(Some(release), cancel) => {
let now = Instant::now();
while let Some(unit) = pending.pop_next_due(now) {
if let Some(report) = pending.take_shape_report() {
on_shape(report);
}
if let Written::Terminated { forwarded, code } = write_unit(unit, send).await? {
pending.clear();
return Ok(DrainOutcome::Terminated { forwarded, code });
}
}
if let Some(report) = pending.take_shape_report() {
on_shape(report);
}
}
}
}
Ok(DrainOutcome::Complete)
}
const DEFAULT_CLOSE: (u32, &[u8]) = (0, b"proxy session ended");
#[derive(Clone, Debug)]
pub(crate) struct SessionCloser {
inner: Arc<CloserInner>,
}
#[derive(Debug)]
struct CloserInner {
request: OnceLock<(u32, Bytes, CloseOrigin)>,
cancel: CancellationToken,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub(crate) enum CloseOrigin {
Hook,
ControlPlane,
}
impl SessionCloser {
pub(crate) fn new(cancel: CancellationToken) -> Self {
Self { inner: Arc::new(CloserInner { request: OnceLock::new(), cancel }) }
}
pub(crate) fn request(&self, code: u32, reason: Bytes) -> bool {
let won = self.inner.request.set((code, reason, CloseOrigin::Hook)).is_ok();
self.inner.cancel.cancel();
won
}
pub(crate) fn record(&self, code: u32, reason: Bytes) -> bool {
self.inner.request.set((code, reason, CloseOrigin::ControlPlane)).is_ok()
}
#[allow(dead_code)]
pub(crate) fn is_closing(&self) -> bool {
self.inner.request.get().is_some()
}
pub(crate) fn requested(&self) -> Option<(u32, Bytes, CloseOrigin)> {
self.inner.request.get().cloned()
}
pub(crate) fn close_args(&self) -> (u32, Bytes) {
self.inner
.request
.get()
.map(|(code, reason, _)| (*code, reason.clone()))
.unwrap_or_else(|| (DEFAULT_CLOSE.0, Bytes::from_static(DEFAULT_CLOSE.1)))
}
}
#[cfg(test)]
mod tests {
use super::*;
#[derive(Debug, Default)]
struct RecordingSink {
writes: Vec<Bytes>,
reset_code: Option<u64>,
fail_after: Option<usize>,
}
impl RecordingSink {
fn written(&self) -> Bytes {
let mut out = Vec::new();
for w in &self.writes {
out.extend_from_slice(w);
}
Bytes::from(out)
}
}
impl EgressSink for RecordingSink {
fn write_all(
&mut self,
buf: &[u8],
) -> impl Future<Output = Result<(), EgressError>> + Send {
let out = if self.fail_after.is_some_and(|n| self.writes.len() >= n) {
Err(EgressError::Transport(TransportError::Write(
"recording sink is full".to_owned(),
)))
} else {
self.writes.push(Bytes::copy_from_slice(buf));
Ok(())
};
async move { out }
}
fn reset(&mut self, code: u64) -> Result<(), EgressError> {
self.reset_code = Some(code);
Ok(())
}
}
fn no_shape_reports(report: ShapeReport) {
unreachable!("an unshaped queue produces no shaping report, got {report:?}");
}
fn queue() -> (PendingQueue, Arc<Recorder>) {
let counters = Arc::new(Recorder::new());
(PendingQueue::new(EgressConfig::default(), Arc::clone(&counters)), counters)
}
fn queue_with(config: EgressConfig) -> (PendingQueue, Arc<Recorder>) {
let counters = Arc::new(Recorder::new());
(PendingQueue::new(config, Arc::clone(&counters)), counters)
}
use crate::shape::{
BucketConfig, ClassRule, DirectionStats, Discipline, Expiry, Matcher, Overflow,
QueueConfig, ShapeProfile, ShapeRecorder,
};
fn test_bucket(rate: Option<u64>, burst: u64) -> BucketConfig {
BucketConfig {
name: "b".to_string(),
rate_bps: rate,
burst_bytes: burst,
..BucketConfig::default()
}
}
fn test_class(name: &str, priority: u8) -> ClassRule {
ClassRule {
name: name.to_string(),
bucket: "b".to_string(),
matcher: Matcher::default(),
priority,
weight: 1,
}
}
fn profile(rate: Option<u64>, burst: u64, max_hold: Duration, expiry: Expiry) -> ShapeProfile {
let queue = QueueConfig {
max_hold: Some(max_hold),
overflow: Overflow::Block,
on_expiry: expiry,
..QueueConfig::default()
};
ShapeProfile::try_new(
vec![test_bucket(rate, burst)],
vec![test_class("only", 0)],
queue,
Discipline::Fifo,
)
.expect("the fixture names its own bucket")
}
fn shaped(profile: ShapeProfile) -> (PendingQueue, Arc<ShapeRecorder>, Arc<Scheduler>) {
let counters = Arc::new(Recorder::new());
let stats = Arc::new(ShapeRecorder::for_profile(Some(&profile)));
let shaper = Arc::new(Scheduler::new(profile));
let mut q = PendingQueue::new(EgressConfig::default(), counters).with_shaper(
Some(Arc::clone(&shaper)),
Arc::clone(&stats),
ProxySide::ClientToProxy,
);
q.tag_unit(Class::Rule(0));
(q, stats, shaper)
}
fn payload(n: usize) -> Bytes {
Bytes::from(vec![0xAB; n])
}
#[test]
fn a_shaped_head_waits_for_its_bucket_and_reports_the_scheduler_s_deadline() {
let (mut q, _stats, _s) =
shaped(profile(Some(1_000), 1_000, Duration::from_secs(60), Expiry::Deliver));
let now = Instant::now();
q.push(Pending::bytes(payload(1_000), now));
q.push(Pending::bytes(payload(1_000), now));
assert!(q.pop_next_due(now).is_some(), "the burst covers the first unit");
assert!(
q.pop_next_due(now).is_none(),
"the burst is spent, so the second unit waits on a refill"
);
let release = q.head_release().expect("a non-empty queue always has one");
assert!(
!release.deadline().token().is_cancelled(),
"the head must be armed at the refill instant, not at its own past due_at"
);
assert!(
q.pop_next_due(now + Duration::from_secs(1)).is_some(),
"a second's refill covers it"
);
assert!(q.is_empty());
}
#[test]
fn a_shaper_holding_a_gated_unit_does_not_report_the_gate_that_freed_it() {
let (mut q, _stats, _s) =
shaped(profile(Some(0), 0, Duration::from_secs(60), Expiry::Deliver));
let now = Instant::now();
let gate = Gate::new();
q.push(Pending::bytes(payload(16), now + Duration::from_secs(30)).with_gate(gate.clone()));
assert!(q.pop_next_due(now).is_none(), "an unreleased gate leaves it not due");
assert!(
q.head_release().expect("queued").gate.is_some_and(|g| !g.is_released()),
"before the shaper has an opinion, the unit's own gate is what it waits on"
);
gate.release();
assert!(q.pop_next_due(now).is_none(), "released and due, but the bucket refuses");
let release = q.head_release().expect("still queued");
assert!(
!release.gate.is_some_and(|g| g.is_released()),
"a released Hold gate must not be re-reported while the shaper holds the unit: \
wait_release would resolve on every iteration and spin the pipe loop"
);
}
#[test]
fn an_unshaped_queue_pops_a_due_head_with_no_shaper_at_all() {
let (mut q, _c) = queue();
assert!(!q.is_shaped());
let now = Instant::now();
q.push(Pending::bytes(payload(1_000_000), now));
assert!(q.pop_next_due(now).is_some(), "an unshaped queue has nothing to ask");
}
#[test]
fn a_terminal_is_never_gated_by_a_bucket() {
let (mut q, _stats, _s) =
shaped(profile(Some(0), 0, Duration::from_secs(60), Expiry::Deliver));
q.push(Pending::terminal(Terminal::Reset { code: 7 }));
let unit = q.pop_next_due(Instant::now()).expect("a reset consults no bucket");
assert!(unit.is_terminal());
q.push(Pending::terminal(Terminal::Truncate { prefix: payload(64), code: 9 }));
let unit = q
.pop_next_due(Instant::now())
.expect("a truncate owes its prefix to the wire, not to a bucket");
assert_eq!(unit.len(), 64, "and the prefix goes with it");
}
#[test]
fn a_zero_rate_class_still_delivers_at_the_clamp_and_reports_it() {
const HOLD: Duration = Duration::from_millis(80);
let (mut q, stats, _s) = shaped(profile(Some(0), 0, HOLD, Expiry::Deliver));
let now = Instant::now();
q.push(Pending::bytes(payload(16), now));
assert!(q.pop_next_due(now).is_none(), "a zero-rate bucket grants nothing");
assert_eq!(
stats.snapshot().classes[0].tokens_exhausted_episodes,
1,
"and says its own bucket was dry, once per episode"
);
assert_eq!(q.take_shape_report(), None, "nothing is clamped before the clamp");
let unit = q
.pop_next_due(now + HOLD + Duration::from_millis(1))
.expect("Expiry::Deliver clamps and delivers");
assert_eq!(unit.len(), 16);
assert_eq!(
q.take_shape_report(),
Some(ShapeReport::Clamped { requested: None, applied: HOLD }),
"a clamped release reports the clamp it applied, and reports the \
request it cut short as absent because a zero-rate bucket named none"
);
assert_eq!(stats.snapshot().classes[0].bytes_delivered, 16);
assert_eq!(stats.snapshot().objects_expired, 0, "Deliver never expires an object");
}
#[test]
fn an_unbounded_shaping_wait_reports_no_request_at_all() {
const HOLD: Duration = Duration::from_millis(80);
const OVERSIZED: usize = 4_096;
let legs = [
("a zero rate", profile(Some(0), 0, HOLD, Expiry::Deliver), 16, None),
(
"a unit above burst_bytes",
profile(Some(64_000), 64, HOLD, Expiry::Deliver),
OVERSIZED,
Some(ShapeReport::BurstBelowUnit {
class: Class::Rule(0),
burst_bytes: 64,
unit_bytes: OVERSIZED as u64,
}),
),
];
for (label, profile, bytes, pre_clamp) in legs {
let (mut q, _stats, _s) = shaped(profile);
let now = Instant::now();
q.push(Pending::bytes(payload(bytes), now));
assert!(q.pop_next_due(now).is_none(), "{label}: nothing may be granted up front");
assert_eq!(
q.take_shape_report(),
pre_clamp,
"{label}: nothing is *clamped* before the clamp, and only the \
mis-sized burst says anything at all before it"
);
let unit = q
.pop_next_due(now + HOLD + Duration::from_millis(1))
.unwrap_or_else(|| panic!("{label}: Expiry::Deliver clamps and delivers"));
assert_eq!(unit.len(), bytes, "{label}: the whole unit goes out");
assert_eq!(
q.take_shape_report(),
Some(ShapeReport::Clamped { requested: None, applied: HOLD }),
"{label}: the bucket named no instant, so there is no request to \
quote and the report must say so rather than name a sentinel"
);
}
}
#[test]
fn expiry_reset_stream_replaces_the_queue_with_its_reset() {
const HOLD: Duration = Duration::from_millis(50);
let (mut q, stats, _s) =
shaped(profile(Some(0), 0, HOLD, Expiry::ResetStream { code: 0x2A }));
let now = Instant::now();
q.push(Pending::bytes(payload(16), now));
q.push(Pending::bytes(payload(16), now));
let unit = q
.pop_next_due(now + HOLD + Duration::from_millis(1))
.expect("an expired head under ResetStream yields the reset");
assert_eq!(unit.item(), &Item::Terminal(Terminal::Reset { code: 0x2A }));
assert_eq!(q.take_shape_report(), Some(ShapeReport::Expired));
assert!(q.is_empty(), "everything behind it went with the stream");
let snap = stats.snapshot();
assert_eq!(snap.objects_expired, 1);
assert_eq!(snap.streams_reset_by_shaping, 1);
}
#[test]
fn two_legs_sharing_one_recorder_charge_their_own_side() {
const HOLD: Duration = Duration::from_millis(50);
let p = profile(Some(0), 0, HOLD, Expiry::ResetStream { code: 0x2A });
let counters = Arc::new(Recorder::new());
let stats = Arc::new(ShapeRecorder::for_profile(Some(&p)));
let shaper = Arc::new(Scheduler::new(p));
let mut up = PendingQueue::new(EgressConfig::default(), Arc::clone(&counters)).with_shaper(
Some(Arc::clone(&shaper)),
Arc::clone(&stats),
ProxySide::ClientToProxy,
);
let mut down = PendingQueue::new(EgressConfig::default(), counters).with_shaper(
Some(shaper),
Arc::clone(&stats),
ProxySide::RelayToProxy,
);
let now = Instant::now();
up.push(Pending::bytes(payload(16), now));
down.push(Pending::bytes(payload(48), now));
down.push(Pending::bytes(payload(48), now));
assert!(
up.pop_next_due(now + HOLD + Duration::from_millis(1)).is_some(),
"an expired head under ResetStream yields the reset"
);
assert!(down.pop_next_due(now).is_some(), "the downlink head is not past its clamp");
let snap = stats.snapshot();
assert_eq!(
snap.uplink,
DirectionStats {
objects_seen: 0,
bytes_shaped: 16,
objects_expired: 1,
streams_reset_by_shaping: 1,
streams_with_mixed_classes: 0,
},
"the uplink queue's bytes and its expiry are the uplink's"
);
assert_eq!(
snap.downlink,
DirectionStats {
objects_seen: 0,
bytes_shaped: 96,
objects_expired: 0,
streams_reset_by_shaping: 0,
streams_with_mixed_classes: 0,
},
"the downlink queued more and gave nothing up: a stall on one leg \
must not be reported on the other"
);
assert_eq!(snap.bytes_shaped, 112, "the aggregate is the two legs and nothing else");
assert_eq!(snap.objects_expired, 1);
assert_eq!(snap.streams_reset_by_shaping, 1);
}
#[test]
fn a_unit_behind_another_class_is_counted_once_and_separately() {
let queue =
QueueConfig { max_hold: Some(Duration::from_secs(60)), ..QueueConfig::default() };
let p = ShapeProfile::try_new(
vec![test_bucket(Some(0), 0)],
vec![test_class("head", 0), test_class("behind", 0)],
queue,
Discipline::Fifo,
)
.expect("two uniquely named classes over one bucket");
let counters = Arc::new(Recorder::new());
let stats = Arc::new(ShapeRecorder::for_profile(Some(&p)));
let mut q = PendingQueue::new(EgressConfig::default(), counters).with_shaper(
Some(Arc::new(Scheduler::new(p))),
Arc::clone(&stats),
ProxySide::ClientToProxy,
);
let now = Instant::now();
q.tag_unit(Class::Rule(0));
q.push(Pending::bytes(payload(16), now));
q.tag_unit(Class::Rule(1));
q.push(Pending::bytes(payload(16), now));
q.push(Pending::bytes(payload(16), now));
for _ in 0..3 {
assert!(q.pop_next_due(now).is_none());
}
let snap = stats.snapshot();
assert_eq!(
snap.classes[1].starved_behind_other_class, 2,
"both units behind the other class's head are counted, once each"
);
assert_eq!(
snap.classes[0].starved_behind_other_class, 0,
"the head is not waiting behind anybody"
);
assert_eq!(
snap.classes[0].tokens_exhausted_episodes, 1,
"the dry bucket is the head's own, and it is one episode"
);
}
#[test]
fn dropping_a_queue_withdraws_the_demand_it_declared() {
let queue =
QueueConfig { max_hold: Some(Duration::from_secs(60)), ..QueueConfig::default() };
let p = ShapeProfile::try_new(
vec![test_bucket(None, 0)],
vec![test_class("hi", 9), test_class("lo", 0)],
queue,
Discipline::StrictPriority,
)
.expect("two uniquely named classes over one bucket");
let counters = Arc::new(Recorder::new());
let stats = Arc::new(ShapeRecorder::for_profile(Some(&p)));
let shaper = Arc::new(Scheduler::new(p));
let now = Instant::now();
{
let mut hi = PendingQueue::new(EgressConfig::default(), Arc::clone(&counters))
.with_shaper(
Some(Arc::clone(&shaper)),
Arc::clone(&stats),
ProxySide::ClientToProxy,
);
hi.tag_unit(Class::Rule(0));
hi.push(Pending::bytes(payload(16), now + Duration::from_secs(60)));
assert!(
matches!(shaper.acquire(Class::Rule(1), 16, now), Acquire::Starved(_)),
"the high class is holding the bucket while its queue is alive"
);
}
assert!(
matches!(shaper.acquire(Class::Rule(1), 16, now), Acquire::Now),
"the high class's stream is gone, so nothing is holding the low one back"
);
}
#[test]
fn an_empty_queue_waits_for_nothing_and_accepts_more() {
let (q, _c) = queue();
assert!(q.is_empty());
assert_eq!(q.len(), 0);
assert_eq!(q.queued_bytes(), 0);
assert!(q.accepts_more());
assert!(q.head_release().is_none());
}
#[test]
fn a_later_unit_cannot_overtake_an_earlier_one() {
let (mut q, _c) = queue();
let now = Instant::now();
let head =
q.push(Pending::bytes(Bytes::from_static(b"0"), now + Duration::from_millis(300)));
let behind = q.push(Pending::bytes(Bytes::from_static(b"1"), now));
assert_eq!(
behind.release_at, head.release_at,
"the reported release is the queue's estimate, and it accounts for the head",
);
let third =
q.push(Pending::bytes(Bytes::from_static(b"2"), now + Duration::from_millis(10)));
assert_eq!(third.release_at, head.release_at);
assert_eq!(q.len(), 3);
assert_eq!(q.queued_bytes(), 3);
}
#[test]
fn the_queues_estimate_never_rewrites_a_units_own_deadline() {
let (mut q, _c) = queue();
let now = Instant::now();
q.push(Pending::bytes(Bytes::from_static(b"0"), now + Duration::from_millis(300)));
q.push(Pending::bytes(Bytes::from_static(b"1"), now));
q.push(Pending::bytes(Bytes::from_static(b"2"), now + Duration::from_millis(10)));
let far = now + Duration::from_secs(1);
let units: Vec<Pending> = std::iter::from_fn(|| q.pop_next_due(far)).collect();
assert_eq!(units.len(), 3);
assert_eq!(units[0].due_at(), now + Duration::from_millis(300));
assert_eq!(units[1].due_at(), now, "a `Pass` behind a delay keeps its own `now`");
assert_eq!(units[2].due_at(), now + Duration::from_millis(10));
for unit in &units {
assert_eq!(unit.expected_at(), units[0].due_at());
}
}
#[test]
fn pushing_bumps_the_egress_counter_once_per_unit() {
let (mut q, counters) = queue();
let now = Instant::now();
for _ in 0..4 {
q.push(Pending::bytes(Bytes::from_static(b"x"), now));
}
assert_eq!(counters.snapshot().egress_items_queued, 4);
}
#[test]
fn a_unit_that_is_due_pops_and_one_that_is_not_does_not() {
let (mut q, _c) = queue();
let now = Instant::now();
q.push(Pending::bytes(Bytes::from_static(b"soon"), now));
q.push(Pending::bytes(Bytes::from_static(b"later"), now + Duration::from_secs(60)));
let first = q.pop_next_due(now).expect("head is due");
assert_eq!(first.len(), 4);
assert!(q.pop_next_due(now).is_none(), "the tail is a minute out");
assert_eq!(q.queued_bytes(), 5);
assert!(q.head_release().is_some());
}
#[test]
fn a_released_gate_makes_a_unit_due_before_its_ceiling() {
let (mut q, _c) = queue();
let now = Instant::now();
let gate = Gate::new();
q.push(
Pending::bytes(Bytes::from_static(b"held"), hold_ceiling(now, q.config()))
.with_gate(gate.clone()),
);
assert!(q.pop_next_due(Instant::now()).is_none());
gate.release();
let unit = q
.pop_next_due(Instant::now())
.expect("a released gate is due, or the release arm spins until max_hold");
assert_eq!(unit.len(), 4);
}
#[test]
fn releasing_a_gate_frees_the_whole_run_queued_behind_it() {
let (mut q, _c) = queue();
let now = Instant::now();
let gate = Gate::new();
q.push(
Pending::bytes(Bytes::from_static(b"held"), hold_ceiling(now, q.config()))
.with_gate(gate.clone()),
);
for tail in [&b"1"[..], b"2", b"3"] {
q.push(Pending::bytes(Bytes::copy_from_slice(tail), Instant::now()));
}
assert_eq!(q.len(), 4);
assert!(q.pop_next_due(Instant::now()).is_none(), "the head holds the whole run");
gate.release();
let at = Instant::now();
let drained: Vec<usize> =
std::iter::from_fn(|| q.pop_next_due(at)).map(|u| u.len()).collect();
assert_eq!(
drained,
vec![4, 1, 1, 1],
"one wake on the gate must free the head *and* everything behind it, in order",
);
assert!(q.is_empty(), "nothing may be left for the max_hold ceiling to release");
assert_eq!(q.queued_bytes(), 0);
}
#[test]
fn a_delay_queued_behind_a_hold_keeps_its_own_deadline() {
const CEILING: Duration = Duration::from_secs(20);
const BY: Duration = Duration::from_millis(50);
let config = EgressConfig { max_hold: CEILING, ..EgressConfig::default() };
let (mut q, _c) = queue_with(config);
let now = Instant::now();
let gate = Gate::new();
q.push(
Pending::bytes(Bytes::from_static(b"held"), hold_ceiling(now, q.config()))
.with_gate(gate.clone()),
);
let queued = q.push(Pending::bytes(Bytes::from_static(b"late"), now + BY));
assert!(
queued.release_at >= now + CEILING,
"the *estimate* stays conservative — the queue cannot know when a gate opens",
);
gate.release();
let head = q.pop_next_due(now + Duration::from_millis(10)).expect("a released gate is due");
assert_eq!(head.len(), 4);
assert!(
q.pop_next_due(now + Duration::from_millis(10)).is_none(),
"the delayed unit's own deadline still governs it: 10 ms is inside its {BY:?}",
);
let second = q
.pop_next_due(now + BY + Duration::from_millis(10))
.expect("arrived_at + by has passed, and that is the whole of the deadline");
assert_eq!(second.len(), 4);
assert_eq!(second.due_at(), now + BY, "its deadline was never rewritten");
}
#[test]
fn the_byte_budget_stops_the_read_branch_and_reports_once() {
let config = EgressConfig { max_pending_bytes: 8, ..EgressConfig::default() };
let (mut q, _c) = queue_with(config);
let now = Instant::now();
let first = q.push(Pending::bytes(Bytes::from_static(b"1234"), now));
assert!(!first.entered_backpressure);
assert!(q.accepts_more());
let second = q.push(Pending::bytes(Bytes::from_static(b"5678"), now));
assert!(second.entered_backpressure, "the transition into backpressure is reported");
assert!(!q.accepts_more());
let third = q.push(Pending::bytes(Bytes::from_static(b"9"), now));
assert!(!third.entered_backpressure, "once per stream, not once per push");
}
#[test]
fn the_unit_cap_bounds_a_queue_of_zero_byte_units() {
let (mut q, _c) = queue();
let now = Instant::now();
let mut reports = 0;
for _ in 0..=MAX_PENDING_UNITS {
if q.push(Pending::elided(now)).entered_backpressure {
reports += 1;
}
}
assert_eq!(q.queued_bytes(), 0, "elided units carry no bytes at all");
assert!(!q.accepts_more(), "the byte budget alone would never stop this");
assert_eq!(reports, 1);
}
#[tokio::test]
async fn an_elided_unit_writes_nothing_but_keeps_its_slot() {
let (mut q, _c) = queue();
let now = Instant::now();
q.push(Pending::bytes(Bytes::from_static(b"a"), now));
q.push(Pending::elided(now));
q.push(Pending::bytes(Bytes::from_static(b"c"), now));
let mut sink = RecordingSink::default();
assert_eq!(q.drain_ignoring_release_times(&mut sink).await, DrainOutcome::Complete);
assert_eq!(sink.written(), Bytes::from_static(b"ac"));
assert_eq!(sink.writes.len(), 2);
}
#[tokio::test]
async fn a_truncate_terminal_writes_its_prefix_then_resets() {
let (mut q, _c) = queue();
let now = Instant::now();
q.push(Pending::bytes(Bytes::from_static(b"ahead"), now));
q.push(Pending::terminal(Terminal::Truncate {
prefix: Bytes::from_static(b"pre"),
code: 0x2,
}));
q.push(Pending::bytes(Bytes::from_static(b"never"), now));
let mut sink = RecordingSink::default();
let outcome = q.drain_ignoring_release_times(&mut sink).await;
assert_eq!(outcome, DrainOutcome::Terminated { forwarded: 3, code: 0x2 });
assert_eq!(sink.written(), Bytes::from_static(b"aheadpre"));
assert_eq!(sink.reset_code, Some(0x2));
assert!(q.is_empty(), "nothing behind a terminal is written");
assert_eq!(q.queued_bytes(), 0);
}
#[tokio::test]
async fn a_reset_terminal_writes_nothing() {
let mut sink = RecordingSink::default();
let written = write_unit(Pending::terminal(Terminal::Reset { code: 7 }), &mut sink)
.await
.expect("reset never fails the write");
assert_eq!(written, Written::Terminated { forwarded: 0, code: 7 });
assert!(sink.writes.is_empty());
assert_eq!(sink.reset_code, Some(7));
}
#[tokio::test]
async fn a_failed_drain_leaves_what_it_could_not_write_queued() {
let (mut q, _c) = queue();
let now = Instant::now();
q.push(Pending::bytes(Bytes::from_static(b"aa"), now));
q.push(Pending::bytes(Bytes::from_static(b"bbb"), now));
q.push(Pending::bytes(Bytes::from_static(b"cccc"), now));
let mut sink = RecordingSink { fail_after: Some(1), ..RecordingSink::default() };
assert_eq!(q.drain_ignoring_release_times(&mut sink).await, DrainOutcome::WriteFailed);
assert_eq!(sink.written(), Bytes::from_static(b"aa"));
assert_eq!(q.len(), 2);
assert_eq!(q.queued_bytes(), 7);
}
#[tokio::test]
async fn the_honouring_drain_writes_in_order_at_release_time() {
let (mut q, counters) = queue();
let now = Instant::now();
q.push(Pending::bytes(Bytes::from_static(b"0"), now + Duration::from_millis(30)));
q.push(Pending::bytes(Bytes::from_static(b"1"), now));
q.push(Pending::bytes(Bytes::from_static(b"2"), now));
let cancel = CancellationToken::new();
let mut sink = RecordingSink::default();
let started = Instant::now();
let outcome = drain_honouring_release_times(&mut q, &mut sink, &cancel, no_shape_reports)
.await
.expect("no write failed");
assert_eq!(outcome, DrainOutcome::Complete);
assert!(started.elapsed() >= Duration::from_millis(30), "release times were honoured");
assert_eq!(sink.written(), Bytes::from_static(b"012"));
assert_eq!(counters.snapshot().release_errors.count, 0, "drains are never release samples");
assert_eq!(
q.unconfirmed_bytes(),
0,
"a drain that ran to completion at its release times is not a teardown and owes \
no `QueuedBytesAtTeardown`",
);
}
#[tokio::test]
async fn a_teardown_flush_reports_what_it_handed_to_a_dying_transport() {
let (mut q, _c) = queue();
q.push(
Pending::bytes(Bytes::from_static(b"gone"), hold_ceiling(Instant::now(), q.config()))
.with_gate(Gate::new()),
);
let cancel = CancellationToken::new();
cancel.cancel();
let mut sink = RecordingSink::default();
let outcome = drain_honouring_release_times(&mut q, &mut sink, &cancel, no_shape_reports)
.await
.expect("the fallback drain swallows write failures");
assert_eq!(outcome, DrainOutcome::CancelledMidDrain);
assert_eq!(sink.written(), Bytes::from_static(b"gone"), "delivered late beats lost");
assert_eq!(q.queued_bytes(), 0, "the fallback handed everything to the transport");
assert_eq!(
q.unconfirmed_bytes(),
4,
"…and handing bytes to a transport the session is closing is not delivering \
them: this is what `QueuedBytesAtTeardown` has to carry, or the object is \
gone with no event at all",
);
}
#[tokio::test]
async fn a_teardown_flush_that_could_not_write_reports_both_halves() {
let (mut q, _c) = queue();
let now = Instant::now();
q.push(Pending::bytes(Bytes::from_static(b"aa"), now));
q.push(Pending::bytes(Bytes::from_static(b"bbb"), now));
q.push(Pending::bytes(Bytes::from_static(b"cccc"), now));
let mut sink = RecordingSink { fail_after: Some(1), ..RecordingSink::default() };
assert_eq!(q.drain_ignoring_release_times(&mut sink).await, DrainOutcome::WriteFailed);
assert_eq!(q.queued_bytes(), 7, "what never reached the transport");
assert_eq!(q.unconfirmed_bytes(), 9, "…plus the two bytes that did, and may be lost");
}
#[tokio::test]
async fn cancelling_a_drain_that_is_waiting_on_a_gate_falls_back_at_once() {
let (mut q, _c) = queue();
let now = Instant::now();
q.push(
Pending::bytes(Bytes::from_static(b"held"), hold_ceiling(now, q.config()))
.with_gate(Gate::new()),
);
let cancel = CancellationToken::new();
let waker = cancel.clone();
tokio::spawn(async move {
tokio::time::sleep(Duration::from_millis(20)).await;
waker.cancel();
});
let mut sink = RecordingSink::default();
let started = Instant::now();
let outcome = drain_honouring_release_times(&mut q, &mut sink, &cancel, no_shape_reports)
.await
.expect("no write failed");
assert_eq!(outcome, DrainOutcome::CancelledMidDrain);
assert!(
started.elapsed() < Duration::from_secs(2),
"teardown must not wait out max_hold; took {:?}",
started.elapsed(),
);
assert_eq!(
sink.written(),
Bytes::from_static(b"held"),
"delivered late beats lost silently"
);
assert_eq!(q.queued_bytes(), 0);
}
#[tokio::test]
async fn an_already_cancelled_drain_does_not_spin() {
let (mut q, _c) = queue();
q.push(Pending::bytes(Bytes::from_static(b"x"), Instant::now() + Duration::from_secs(60)));
let cancel = CancellationToken::new();
cancel.cancel();
let mut sink = RecordingSink::default();
let started = Instant::now();
let outcome = drain_honouring_release_times(&mut q, &mut sink, &cancel, no_shape_reports)
.await
.unwrap();
assert_eq!(outcome, DrainOutcome::CancelledMidDrain);
assert!(started.elapsed() < Duration::from_secs(1), "the biased arm must win");
assert_eq!(sink.written(), Bytes::from_static(b"x"));
}
#[tokio::test]
async fn wait_release_resolves_on_the_gate_the_deadline_or_the_cancel() {
let (mut q, _c) = queue();
let far = Instant::now() + Duration::from_secs(60);
let gate = Gate::new();
q.push(Pending::bytes(Bytes::from_static(b"g"), far).with_gate(gate.clone()));
gate.release();
let cancel = CancellationToken::new();
wait_release(q.head_release(), &cancel).await;
let (mut q, _c) = queue();
q.push(Pending::bytes(
Bytes::from_static(b"d"),
Instant::now() + Duration::from_millis(20),
));
let started = Instant::now();
wait_release(q.head_release(), &cancel).await;
assert!(started.elapsed() >= Duration::from_millis(20));
let (mut q, _c) = queue();
q.push(Pending::bytes(Bytes::from_static(b"c"), far));
cancel.cancel();
wait_release(q.head_release(), &cancel).await;
wait_release(None, &cancel).await;
}
#[tokio::test]
async fn the_pipe_loop_shape_from_the_contract_compiles_and_keeps_order() {
let (mut pending, counters) = queue();
let cancel = CancellationToken::new();
let mut sink = RecordingSink::default();
let (tx, mut rx) = tokio::sync::mpsc::channel::<Bytes>(8);
tokio::spawn(async move {
for chunk in [&b"0"[..], b"1", b"2", b"3"] {
tx.send(Bytes::from_static(chunk)).await.expect("receiver lives");
}
tokio::time::sleep(Duration::from_millis(200)).await;
});
let mut first = true;
loop {
let can_read = pending.accepts_more();
let head_release = pending.head_release();
tokio::select! {
chunk = rx.recv(), if can_read => {
match chunk {
Some(raw) => {
let release_at = if std::mem::take(&mut first) {
Instant::now() + Duration::from_millis(40)
} else {
Instant::now()
};
if pending.is_empty() && release_at <= Instant::now() {
write_unit(Pending::bytes(raw, release_at), &mut sink)
.await
.expect("recording sink");
} else {
pending.push(Pending::bytes(raw, release_at));
}
}
None => {
let outcome =
drain_honouring_release_times(&mut pending, &mut sink, &cancel, no_shape_reports)
.await
.expect("recording sink");
assert_eq!(outcome, DrainOutcome::Complete);
break;
}
}
}
() = wait_release(head_release.clone(), &cancel), if head_release.is_some() => {
let now = Instant::now();
while let Some(unit) = pending.pop_next_due(now) {
pending.record_release(&unit, now);
write_unit(unit, &mut sink).await.expect("recording sink");
}
}
() = cancel.cancelled() => unreachable!("nothing cancels this test"),
}
}
assert_eq!(
sink.written(),
Bytes::from_static(b"0123"),
"byte equality is the ordering assertion"
);
let snap = counters.snapshot();
assert_eq!(snap.egress_items_queued, 4, "every unit went through the deque");
assert_eq!(
snap.release_errors.count, 4,
"all four share the head's clamped release time, so one wake pops all four \
— and only the release branch samples",
);
}
#[test]
fn defer_by_is_a_deadline_and_clamps_to_max_hold() {
let config =
EgressConfig { max_hold: Duration::from_millis(100), ..EgressConfig::default() };
let now = Instant::now();
let short = defer_by(now, Duration::from_millis(10), &config);
assert_eq!(short.release_at, now + Duration::from_millis(10));
assert!(!short.was_clamped());
let long = defer_by(now, Duration::from_secs(5), &config);
assert_eq!(long.release_at, now + Duration::from_millis(100));
assert!(long.was_clamped());
assert_eq!(long.requested, Duration::from_secs(5));
assert_eq!(long.applied, Duration::from_millis(100));
let a = defer_by(now, Duration::from_millis(10), &config);
let b = defer_by(now, Duration::from_millis(10), &config);
assert_eq!(a.release_at, b.release_at);
assert_eq!(hold_ceiling(now, &config), now + Duration::from_millis(100));
}
#[test]
fn the_first_close_request_wins() {
let cancel = CancellationToken::new();
let closer = SessionCloser::new(cancel.clone());
assert!(!closer.is_closing());
assert_eq!(
closer.close_args(),
(0, Bytes::from_static(b"proxy session ended")),
"the default is what run_with_transport has always sent",
);
assert!(closer.request(3, Bytes::from_static(b"protocol violation")));
assert!(cancel.is_cancelled());
assert!(closer.is_closing());
assert!(!closer.request(1, Bytes::from_static(b"too late")));
assert_eq!(closer.close_args(), (3, Bytes::from_static(b"protocol violation")));
assert_eq!(
closer.requested(),
Some((3, Bytes::from_static(b"protocol violation"), CloseOrigin::Hook)),
"a close that arrived through `request` is a hook's, and the session's own \
`SessionEnded` reason says so in those words",
);
}
#[test]
fn an_undelayed_head_arms_a_deadline_that_never_touches_the_wheel() {
let (mut q, _c) = queue();
q.push(Pending::bytes(Bytes::from_static(b"now"), Instant::now()));
let release = q.head_release().expect("one unit queued");
assert!(
release.deadline().token().is_cancelled(),
"an already-due head must not be registered with the wheel",
);
let (mut q, _c) = queue();
q.push(Pending::bytes(
Bytes::from_static(b"later"),
Instant::now() + Duration::from_secs(60),
));
let release = q.head_release().expect("one unit queued");
assert!(!release.deadline().token().is_cancelled(), "a future head is registered");
assert!(release_timer::started(), "…and registering is what starts the wheel");
}
fn gauged() -> (PendingQueue, Arc<EgressGauge>) {
let gauge = EgressGauge::new();
let q = PendingQueue::new(EgressConfig::default(), Arc::new(Recorder::new()))
.with_gauge(Arc::clone(&gauge));
(q, gauge)
}
#[tokio::test]
async fn the_gauge_follows_every_route_bytes_leave_a_queue_by() {
let (mut q, gauge) = gauged();
let now = Instant::now();
assert_eq!(gauge.queued(), 0);
q.push(Pending::bytes(Bytes::from_static(b"aaa"), now));
q.push(Pending::bytes(Bytes::from_static(b"bb"), now));
assert_eq!(gauge.queued(), 5, "two pushes, five bytes");
q.pop_next_due(now).expect("both are due");
assert_eq!(gauge.queued(), 2);
let mut sink = RecordingSink::default();
assert_eq!(q.drain_ignoring_release_times(&mut sink).await, DrainOutcome::Complete);
assert_eq!(gauge.queued(), 0);
q.push(Pending::bytes(Bytes::from_static(b"cccc"), now));
assert_eq!(gauge.queued(), 4);
q.clear();
assert_eq!(gauge.queued(), 0);
q.push(Pending::bytes(Bytes::from_static(b"dddd"), now));
assert_eq!(gauge.queued(), 4);
drop(q);
assert_eq!(
gauge.queued(),
0,
"a queue that goes away takes its remainder with it, or every later close waits out \
its whole window"
);
}
#[tokio::test]
async fn the_drain_wait_ends_the_moment_the_queues_empty() {
let (mut q, gauge) = gauged();
let now = Instant::now();
q.push(Pending::bytes(Bytes::from_static(b"payload"), now));
let waiting = tokio::spawn({
let gauge = Arc::clone(&gauge);
async move { gauge.wait_idle(Duration::from_secs(30)).await }
});
tokio::task::yield_now().await;
let mut sink = RecordingSink::default();
assert_eq!(q.drain_ignoring_release_times(&mut sink).await, DrainOutcome::Complete);
let stranded = tokio::time::timeout(Duration::from_secs(5), waiting)
.await
.expect("the wait must end on the queue emptying, not on its own deadline")
.expect("the waiting task must not panic");
assert_eq!(stranded, 0, "nothing was left, so nothing is abandoned");
}
#[tokio::test]
async fn an_expired_window_reports_the_residue_and_writes_nothing() {
let (mut q, gauge) = gauged();
let now = Instant::now();
q.push(Pending::bytes(Bytes::from_static(b"held"), now + Duration::from_secs(60)));
let stranded = gauge.wait_idle(Duration::from_millis(20)).await;
assert_eq!(stranded, 4, "the window closed on four bytes that had not been released");
gauge.begin_discarding();
let mut sink = RecordingSink::default();
assert_eq!(
q.drain_ignoring_release_times(&mut sink).await,
DrainOutcome::Discarded,
"past the deadline the flush declines rather than best-efforts"
);
assert!(sink.writes.is_empty(), "nothing reached the transport");
assert_eq!(
q.unconfirmed_bytes(),
4,
"so the reported figure is what was abandoned, with nothing handed anywhere"
);
assert_eq!(
q.queued_bytes(),
q.unconfirmed_bytes(),
"and the two agree, which is the point"
);
}
#[test]
fn recording_a_close_leaves_the_session_running() {
let cancel = CancellationToken::new();
let closer = SessionCloser::new(cancel.clone());
assert!(closer.record(7, Bytes::from_static(b"asked")));
assert!(!cancel.is_cancelled(), "the drain window has not even started yet");
assert_eq!(closer.close_args(), (7, Bytes::from_static(b"asked")));
assert!(!closer.record(9, Bytes::from_static(b"second")));
assert!(!closer.request(9, Bytes::from_static(b"second")));
assert_eq!(closer.close_args(), (7, Bytes::from_static(b"asked")));
assert!(cancel.is_cancelled(), "…and `request` still cancels, losing or not");
assert_eq!(
closer.requested(),
Some((7, Bytes::from_static(b"asked"), CloseOrigin::ControlPlane)),
);
}
}