Skip to main content

wyrd_runtime/
outbox.rs

1//! Generic non-blocking, per-tenant batched outbox.
2//!
3//! A server surface that must record derived work without delaying the request
4//! that produced it hands each item to [`Outbox::stage`], which only enqueues
5//! it. One background writer per outbox groups queued items by tenant and by
6//! the sink's [`OutboxSink::lane`], and writes each lane's whole backlog
7//! through its [`OutboxSink`] in one all-or-nothing call. Each lane has at
8//! most one write in flight, so items of one lane are written in arrival
9//! order; a tenant's lanes write concurrently. At most `concurrency` writes
10//! run at once, and the tenant with the fewest writes in flight is
11//! dispatched first, so a busy tenant cannot hold every slot while another
12//! waits, and a slow tenant delays only its own items.
13//!
14//! A failed write is never dropped. A write that returns an error or panics
15//! keeps its exact items as that lane's next write, apart from anything that
16//! arrived meanwhile, and the lane is retried after a backoff that starts at
17//! [`INITIAL_BACKOFF`] and doubles up to [`MAX_BACKOFF`]; other lanes keep
18//! writing. Later arrivals wait for the write after the retry succeeds, so a
19//! sink deriving identity from its slice presents the same identity on every
20//! attempt. Because a write may be retried after an unknown commit outcome,
21//! every sink write must be safe to repeat.
22//!
23//! Every outbox has finite [`OutboxLimits`]: a most admitted items and a
24//! most admitted bytes, as each item's [`OutboxSink::weight`]. Admitted
25//! items are those queued, waiting, retrying, or being written, plus slots
26//! reserved and not yet staged. An item stays admitted until it is written
27//! or lost, never merely dequeued, and a retried slice is admitted once
28//! however often it is retried. Both limits are split in two shares. A
29//! producer that must refuse its caller rather than lose an item takes a
30//! [`Reservation`] first; [`Outbox::reserve`] refuses with [`OutboxFull`]
31//! once reserved slots and the unwritten items staged through them would
32//! exceed [`RESERVED_SHARE`] of the item limit, or the bytes of those items
33//! already reach that share of the byte limit. Unreserved items from
34//! [`Outbox::stage`] hold the rest of both, so neither share can crowd out
35//! the other. `stage` never fails or waits: past either unreserved limit an
36//! item is shed, counted lost, and logged at most once per tenant every
37//! [`SHED_LOG_INTERVAL`] with the count shed since the last record, so a
38//! full outbox never blocks or fails the operation staging it.
39//!
40//! Nothing is preallocated. Every lost item is counted once in
41//! `outbox_events_lost_total{outbox, reason}` with its [`LossReason`]: shed at
42//! an unreserved limit, staged after shutdown began or unwritten at its
43//! deadline, or consumed unwritten by the sink. Items lost when the process
44//! stops abruptly are never counted. Written items and their bytes are
45//! counted in `outbox_items_written_total{outbox}` and
46//! `outbox_bytes_written_total{outbox}`, every failed write attempt in
47//! `outbox_write_failures_total{outbox}`, and the items not yet written are
48//! exported as the `outbox_pending{outbox}` gauge.
49
50use std::any::Any;
51use std::collections::HashMap;
52use std::fmt::Display;
53use std::future::{Future, pending, poll_fn};
54use std::panic::{AssertUnwindSafe, catch_unwind};
55use std::pin::pin;
56use std::sync::atomic::{AtomicUsize, Ordering};
57use std::sync::{Arc, Mutex, PoisonError, RwLock};
58use std::task::Poll;
59use std::time::{Duration, Instant};
60
61use tokio::sync::{Notify, mpsc};
62use tokio::task::{Id, JoinError, JoinSet};
63use tokio::time::{sleep_until, timeout_at};
64use tokio_util::sync::CancellationToken;
65use tokio_util::task::TaskTracker;
66use wyrd_spec::DataTenantId;
67
68/// Delay before the first retry of a lane whose write failed.
69pub const INITIAL_BACKOFF: Duration = Duration::from_millis(50);
70
71/// Longest delay between retries of a tenant whose writes keep failing.
72pub const MAX_BACKOFF: Duration = Duration::from_secs(5);
73
74/// Share of an outbox's capacity, as a fraction `(numerator, denominator)`,
75/// that reservations may hold; the rest is unreserved staging's, so a
76/// refusing producer's burst cannot crowd out audit decisions and their
77/// traffic cannot refuse the producer.
78pub const RESERVED_SHARE: (usize, usize) = (3, 4);
79
80/// Shortest time between two shed records of one outbox and tenant.
81///
82/// A saturated producer sheds every item it stages; one record per interval
83/// carrying the count keeps the log readable and its cost flat.
84pub const SHED_LOG_INTERVAL: Duration = Duration::from_secs(1);
85
86/// The most an [`Outbox`] admits at once, by item count and by accounted
87/// bytes; reaching either limit of a share sheds or refuses.
88#[derive(Debug, Clone, Copy, PartialEq, Eq)]
89pub struct OutboxLimits {
90    /// Most items admitted, including unspent reserved slots.
91    pub items: usize,
92    /// Most bytes admitted, as the sum of [`OutboxSink::weight`] over the
93    /// admitted items.
94    pub bytes: usize,
95}
96
97/// Why an outbox item was lost, the `reason` label of
98/// `outbox_events_lost_total`.
99#[derive(Debug, Clone, Copy, PartialEq, Eq)]
100pub enum LossReason {
101    /// Refused at an unreserved staging limit.
102    Shed,
103    /// Consumed by the sink without being written: a terminal refusal or an
104    /// item that could not be encoded, which no retry could cure.
105    Rejected,
106    /// The process reaches no destination to write to.
107    Unreachable,
108    /// Staged after shutdown began, or still unwritten at its deadline.
109    Shutdown,
110}
111
112impl LossReason {
113    /// The fixed metric label value.
114    #[must_use]
115    pub const fn as_str(self) -> &'static str {
116        match self {
117            Self::Shed => "shed",
118            Self::Rejected => "rejected",
119            Self::Unreachable => "unreachable",
120            Self::Shutdown => "shutdown",
121        }
122    }
123}
124
125/// Refusal of a [`Reservation`] because the outbox holds its reserved share.
126///
127/// Nothing was admitted, so the caller may retry the whole operation later.
128#[derive(Debug, Clone, Copy, PartialEq, Eq, thiserror::Error)]
129#[error("outbox {outbox} is full: {admitted} of {limit} reservable {unit} admitted")]
130pub struct OutboxFull {
131    /// The refusing outbox's [`OutboxSink::NAME`].
132    pub outbox: &'static str,
133    /// Reserved slots and unwritten reserved items, or the bytes of those
134    /// items, when the reservation was refused.
135    pub admitted: usize,
136    /// Items or bytes reservations may hold in this outbox.
137    pub limit: usize,
138    /// Which limit refused: `"items"` or `"bytes"`.
139    pub unit: &'static str,
140}
141
142/// One destination an outbox writes to.
143///
144/// Implementations write everything one lane of a tenant has queued in one
145/// call: in one transaction, or as idempotent writes whose repeat is
146/// absorbed. The outbox owns queueing, grouping, concurrency, retry, and
147/// shutdown.
148pub trait OutboxSink: Send + Sync + 'static {
149    /// The item callers stage.
150    type Item: Send + 'static;
151    /// The failure a write reports; logged, counted, and retried.
152    type Error: Display + Send + 'static;
153    /// Label for this outbox's metrics and logs.
154    const NAME: &'static str;
155    /// Lanes each tenant's items are spread over. Lanes of one tenant write
156    /// concurrently; the default single lane writes a tenant one slice at a
157    /// time.
158    const LANES: usize = 1;
159
160    /// The lane `item` is written in, taken modulo [`Self::LANES`].
161    ///
162    /// Items of one tenant and lane are written in arrival order, one write
163    /// at a time, so items whose writes must observe each other in arrival
164    /// order, such as requests sharing an idempotency key, must share a lane.
165    #[must_use]
166    fn lane(item: &Self::Item) -> u64 {
167        let _ = item;
168        0
169    }
170
171    /// The bytes `item` is accounted as against [`OutboxLimits::bytes`].
172    ///
173    /// At least the length of every payload the item holds, such as captured
174    /// request and response bodies; encoding overhead need not be counted.
175    /// The default of zero suits items bounded by their producer's input
176    /// limits, which the item limit then bounds.
177    #[must_use]
178    fn weight(item: &Self::Item) -> usize {
179        let _ = item;
180        0
181    }
182
183    /// Writes one lane's items of one tenant, in order, in one transaction.
184    ///
185    /// A write may be retried after its commit outcome is unknown, so writing
186    /// the same items twice must leave the same durable state as writing them
187    /// once.
188    ///
189    /// # Errors
190    ///
191    /// Returns the sink's error when any part of the write fails; nothing of
192    /// the batch may then be durable unless a repeat would skip it.
193    fn write(
194        &self,
195        tenant: DataTenantId,
196        items: &[Self::Item],
197    ) -> impl Future<Output = Result<(), Self::Error>> + Send;
198}
199
200/// One queued item: its tenant, the item, and its admitted [`Load`].
201type Queued<S> = (DataTenantId, <S as OutboxSink>::Item, Load);
202
203/// Sending half of an outbox queue.
204type QueueSender<S> = mpsc::UnboundedSender<Queued<S>>;
205
206/// The handle callers stage through; owns the queue and its writer task.
207pub struct Outbox<S: OutboxSink> {
208    /// Sending half of the unbounded queue the writer drains, taken by
209    /// shutdown. `stage` sends under the read lock and shutdown takes it under
210    /// the write lock, so every item is either queued before the fence and
211    /// drained, or refused after it.
212    queue: RwLock<Option<QueueSender<S>>>,
213    /// Items queued, waiting, or being written, not yet written or lost.
214    pending: Arc<AtomicUsize>,
215    /// Admission counts the writer releases written and lost items from.
216    admission: Arc<Admission>,
217    /// Most items and bytes this outbox admits at once.
218    limits: OutboxLimits,
219    /// Per tenant, items shed since the last shed record and when that record
220    /// was written; see [`SHED_LOG_INTERVAL`].
221    // ponytail: one small entry per tenant that ever shed, never pruned;
222    // prune idle entries if tenant counts reach the millions.
223    shed: Mutex<HashMap<DataTenantId, ShedLog>>,
224    /// Signalled whenever a settled write or abandonment leaves nothing
225    /// pending.
226    idle: Arc<Notify>,
227    /// Asks the writer to give up at the shutdown deadline.
228    abandon: CancellationToken,
229    /// Tracks the writer so shutdown can wait for it.
230    writer: TaskTracker,
231    /// The sink the writer submits through, shared for inspection.
232    sink: Arc<S>,
233}
234
235/// Names the outbox and its pending count, so values carrying a handle — a
236/// staged write that hands off to a second outbox — stay debuggable.
237impl<S: OutboxSink> std::fmt::Debug for Outbox<S> {
238    /// Write the sink's name and a relaxed snapshot of the pending count;
239    /// queued items and the sink itself are never formatted.
240    ///
241    /// # Errors
242    /// Returns the formatter's error when it rejects a write.
243    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
244        formatter
245            .debug_struct("Outbox")
246            .field("name", &S::NAME)
247            .field("pending", &self.pending.load(Ordering::Relaxed))
248            .finish_non_exhaustive()
249    }
250}
251
252impl<S: OutboxSink> Outbox<S> {
253    /// Creates the outbox around `sink` and starts its writer.
254    ///
255    /// `concurrency` is the most writes run at once across every tenant and
256    /// lane, which bounds the connections the sink may hold; `limits` are the
257    /// most items and bytes admitted at once, which bound the memory queued,
258    /// waiting, retrying, and in-flight work may hold.
259    ///
260    /// # Panics
261    ///
262    /// Panics when called outside a Tokio runtime, because the writer task is
263    /// spawned immediately.
264    #[must_use]
265    pub fn new(sink: S, concurrency: usize, limits: OutboxLimits) -> Arc<Self> {
266        // The channel is unbounded because `admission` already bounds it.
267        let (queue, requests) = mpsc::unbounded_channel();
268        let pending = Arc::new(AtomicUsize::new(0));
269        let admission = Arc::new(Admission::default());
270        let idle = Arc::new(Notify::new());
271        let abandon = CancellationToken::new();
272        let writer = TaskTracker::new();
273        let sink = Arc::new(sink);
274        writer.spawn(
275            OutboxWriter {
276                sink: Arc::clone(&sink),
277                requests,
278                waiting: HashMap::new(),
279                failed: HashMap::new(),
280                writing: JoinSet::new(),
281                in_flight: HashMap::new(),
282                retry_at: HashMap::new(),
283                concurrency: concurrency.max(1),
284                pending: Arc::clone(&pending),
285                admission: Arc::clone(&admission),
286                idle: Arc::clone(&idle),
287                abandon: abandon.clone(),
288            }
289            .run(),
290        );
291        writer.close();
292        Arc::new(Self {
293            queue: RwLock::new(Some(queue)),
294            pending,
295            admission,
296            limits,
297            shed: Mutex::default(),
298            idle,
299            abandon,
300            writer,
301            sink,
302        })
303    }
304
305    /// Borrows the sink this outbox writes through.
306    #[must_use]
307    pub fn sink(&self) -> &S {
308        &self.sink
309    }
310
311    /// Queues one item for `tenant` without waiting and never fails.
312    ///
313    /// The item is admitted and counted pending before it is sent, so the
314    /// writer can never release it first. Admission happens under the queue's
315    /// read lock, so nothing is admitted once shutdown has taken the queue.
316    /// When either unreserved limit would be passed the item is shed: counted
317    /// lost with [`LossReason::Shed`] and recorded by [`Self::shed`]. Once
318    /// shutdown has begun, or if the writer is gone, the item is counted lost
319    /// with [`LossReason::Shutdown`] and logged.
320    pub fn stage(&self, tenant: DataTenantId, item: impl Into<S::Item>) {
321        let item = item.into();
322        let bytes = S::weight(&item);
323        let queue = self.queue.read().unwrap_or_else(PoisonError::into_inner);
324        let Some(sender) = queue.as_ref() else {
325            drop(queue);
326            lose_late::<S>(tenant);
327            return;
328        };
329        let admission = &self.admission;
330        if admit(
331            &admission.unreserved,
332            1,
333            self.limits.items - self.reservable(),
334        )
335        .is_err()
336        {
337            drop(queue);
338            self.shed(tenant);
339            return;
340        }
341        let byte_limit = self.limits.bytes - self.reservable_bytes();
342        if admit(&admission.unreserved_bytes, bytes, byte_limit).is_err() {
343            admission.unreserved.fetch_sub(1, Ordering::AcqRel);
344            drop(queue);
345            self.shed(tenant);
346            return;
347        }
348        self.enqueue(sender, tenant, item, Load::one(false, bytes));
349    }
350
351    /// Counts one shed item of `tenant` and writes a shed record when the
352    /// tenant's last one is at least [`SHED_LOG_INTERVAL`] old; the record
353    /// carries every item shed since the previous one.
354    fn shed(&self, tenant: DataTenantId) {
355        count_lost::<S>(1, LossReason::Shed);
356        let now = Instant::now();
357        let shed = {
358            let mut logs = self.shed.lock().unwrap_or_else(PoisonError::into_inner);
359            let log = logs.entry(tenant).or_default();
360            log.unlogged += 1;
361            if log
362                .logged_at
363                .is_some_and(|at| now.duration_since(at) < SHED_LOG_INTERVAL)
364            {
365                return;
366            }
367            log.logged_at = Some(now);
368            std::mem::take(&mut log.unlogged)
369        };
370        self.log_shed(tenant, shed);
371    }
372
373    /// Writes one shed record for `shed` items of `tenant`.
374    fn log_shed(&self, tenant: DataTenantId, shed: usize) {
375        tracing::error!(
376            outbox = S::NAME,
377            %tenant,
378            shed,
379            max_items = self.limits.items,
380            max_bytes = self.limits.bytes,
381            "outbox is full; staged items were shed"
382        );
383    }
384
385    /// Reserves `slots` items of capacity for a producer that must refuse its
386    /// caller instead of losing what it stages.
387    ///
388    /// The reservation is admitted at once and holds its slots until each is
389    /// staged through it or it is dropped, which releases the unused rest. An
390    /// item staged through it keeps its slot until it is written or lost.
391    /// Unreserved items never count against reservations.
392    ///
393    /// # Errors
394    ///
395    /// Returns [`OutboxFull`], admitting nothing, when the reservation would
396    /// take reservations past [`RESERVED_SHARE`] of the capacity.
397    pub fn reserve(self: &Arc<Self>, slots: usize) -> Result<Reservation<S>, OutboxFull> {
398        let byte_limit = self.reservable_bytes();
399        let bytes = self.admission.reserved_bytes.load(Ordering::Acquire);
400        if bytes >= byte_limit {
401            return Err(OutboxFull {
402                outbox: S::NAME,
403                admitted: bytes,
404                limit: byte_limit,
405                unit: "bytes",
406            });
407        }
408        let limit = self.reservable();
409        admit(&self.admission.reserved, slots, limit).map_err(|admitted| OutboxFull {
410            outbox: S::NAME,
411            admitted,
412            limit,
413            unit: "items",
414        })?;
415        Ok(Reservation {
416            outbox: Arc::clone(self),
417            slots: AtomicUsize::new(slots),
418        })
419    }
420
421    /// The most slots reservations may hold: [`RESERVED_SHARE`] of the
422    /// item limit, rounded down.
423    const fn reservable(&self) -> usize {
424        self.limits.items * RESERVED_SHARE.0 / RESERVED_SHARE.1
425    }
426
427    /// The bytes at which reservations are refused: [`RESERVED_SHARE`] of
428    /// the byte limit, rounded down, computed without overflowing.
429    const fn reservable_bytes(&self) -> usize {
430        self.limits.bytes / RESERVED_SHARE.1 * RESERVED_SHARE.0
431    }
432
433    /// Sends one item on a spent reserved slot to the writer.
434    ///
435    /// Its bytes join the reserved share unconditionally, because the slot
436    /// already promised the producer its item is kept; [`Self::reserve`]
437    /// refuses new reservations once that share's bytes reach their limit.
438    /// Once shutdown has begun the item is released from its slot, counted
439    /// lost with [`LossReason::Shutdown`], and logged.
440    fn send_reserved(&self, tenant: DataTenantId, item: S::Item) {
441        let bytes = S::weight(&item);
442        let queue = self.queue.read().unwrap_or_else(PoisonError::into_inner);
443        let Some(sender) = queue.as_ref() else {
444            drop(queue);
445            self.admission.release(Load::one(true, 0));
446            lose_late::<S>(tenant);
447            return;
448        };
449        self.admission
450            .reserved_bytes
451            .fetch_add(bytes, Ordering::AcqRel);
452        self.enqueue(sender, tenant, item, Load::one(true, bytes));
453    }
454
455    /// Sends one admitted item to the writer, counting it pending first.
456    ///
457    /// Called under the queue's read lock. If the writer is gone, the item is
458    /// released from its shares, counted lost with [`LossReason::Shutdown`],
459    /// and logged.
460    fn enqueue(&self, sender: &QueueSender<S>, tenant: DataTenantId, item: S::Item, load: Load) {
461        let gauge = metrics::gauge!("outbox_pending", "outbox" => S::NAME);
462        self.pending.fetch_add(1, Ordering::AcqRel);
463        self.admission
464            .staged_reserved
465            .fetch_add(load.reserved, Ordering::AcqRel);
466        gauge.increment(1.0);
467        if sender.send((tenant, item, load)).is_ok() {
468            return;
469        }
470        self.pending.fetch_sub(1, Ordering::AcqRel);
471        self.admission
472            .staged_reserved
473            .fetch_sub(load.reserved, Ordering::AcqRel);
474        gauge.decrement(1.0);
475        self.admission.release(load);
476        lose_late::<S>(tenant);
477    }
478
479    /// Returns the number of items not yet written.
480    #[must_use]
481    pub fn pending(&self) -> usize {
482        self.pending.load(Ordering::Acquire)
483    }
484
485    /// Waits until `deadline` for everything staged so far to be written,
486    /// without closing the outbox, and returns how many items remain.
487    ///
488    /// The outbox keeps accepting items throughout; callers that read the
489    /// sink's destination right after staging use this to observe their items.
490    pub async fn settle(&self, deadline: Instant) -> usize {
491        loop {
492            let idle = self.idle.notified();
493            tokio::pin!(idle);
494            idle.as_mut().enable();
495            if self.pending() == 0 {
496                return 0;
497            }
498            if timeout_at(deadline.into(), idle).await.is_err() {
499                return self.pending();
500            }
501        }
502    }
503
504    /// Stops taking items, keeps writing and retrying until `deadline`, and
505    /// returns how many items were lost at the deadline.
506    ///
507    /// The staging sender is dropped before the first await, which is a
508    /// one-way fence: everything staged before it drains, and every later
509    /// [`Outbox::stage`] is refused and counted lost without entering the
510    /// queue. Shed counts not yet recorded are recorded at the fence.
511    /// Items still unwritten at the deadline are abandoned with the
512    /// writer, which cancels any in-flight write; once the writer is gone
513    /// they are counted in `outbox_events_lost_total` with
514    /// [`LossReason::Shutdown`], logged, and released
515    /// from `pending()` and `outbox_pending`, which both end at zero, and
516    /// released from their admission shares. When shutdown is called
517    /// concurrently, exactly one caller reports the loss and settles both
518    /// shares, and the others report zero. Unspent [`Reservation`] slots stay
519    /// held until their reservation drops.
520    pub async fn shutdown(&self, deadline: Instant) -> usize {
521        drop(
522            self.queue
523                .write()
524                .unwrap_or_else(PoisonError::into_inner)
525                .take(),
526        );
527        let unlogged: Vec<_> = self
528            .shed
529            .lock()
530            .unwrap_or_else(PoisonError::into_inner)
531            .iter_mut()
532            .filter(|(_, log)| log.unlogged > 0)
533            .map(|(tenant, log)| (*tenant, std::mem::take(&mut log.unlogged)))
534            .collect();
535        for (tenant, shed) in unlogged {
536            self.log_shed(tenant, shed);
537        }
538        if timeout_at(deadline.into(), self.writer.wait())
539            .await
540            .is_err()
541        {
542            self.abandon.cancel();
543            self.writer.wait().await;
544        }
545        // With the queue closed and the writer gone nothing else moves these
546        // counts, and admission only happens under the queue lock, so the
547        // admitted bytes are exactly the abandoned items'. The caller that
548        // takes a nonzero `pending` owns its reserved part and settles every
549        // share alone.
550        let lost = self.pending.swap(0, Ordering::AcqRel);
551        if lost == 0 {
552            return 0;
553        }
554        let admission = &self.admission;
555        let reserved = admission.staged_reserved.swap(0, Ordering::AcqRel);
556        let reserved_bytes = admission.reserved_bytes.swap(0, Ordering::AcqRel);
557        let bytes = admission.unreserved_bytes.swap(0, Ordering::AcqRel) + reserved_bytes;
558        admission.release(Load {
559            items: lost,
560            reserved,
561            bytes,
562            reserved_bytes,
563        });
564        count_lost::<S>(lost, LossReason::Shutdown);
565        metrics::gauge!("outbox_pending", "outbox" => S::NAME).decrement(lost as f64);
566        tracing::error!(
567            outbox = S::NAME,
568            lost,
569            "outbox shutdown deadline passed with items unwritten"
570        );
571        self.idle.notify_waiters();
572        lost
573    }
574}
575
576/// Capacity an [`Outbox`] admitted ahead of staging, from [`Outbox::reserve`].
577///
578/// Each [`Reservation::stage`] spends one slot; dropping the reservation
579/// releases the slots it never spent. A value carrying a reservation, such as
580/// a write that hands off to a second outbox, holds that capacity until the
581/// value is staged or dropped.
582pub struct Reservation<S: OutboxSink> {
583    /// The outbox the slots were admitted in.
584    outbox: Arc<Outbox<S>>,
585    /// Slots admitted and not yet staged.
586    slots: AtomicUsize,
587}
588
589impl<S: OutboxSink> Reservation<S> {
590    /// Stages one item for `tenant` on a reserved slot without waiting.
591    ///
592    /// With no slot left, the item is staged like [`Outbox::stage`]. Once
593    /// shutdown has begun the item is counted lost, as any late stage is.
594    pub fn stage(&self, tenant: DataTenantId, item: impl Into<S::Item>) {
595        let spent = self
596            .slots
597            .try_update(Ordering::AcqRel, Ordering::Acquire, |slots| {
598                slots.checked_sub(1)
599            });
600        if spent.is_ok() {
601            self.outbox.send_reserved(tenant, item.into());
602        } else {
603            self.outbox.stage(tenant, item);
604        }
605    }
606}
607
608/// Releases the slots this reservation never spent.
609impl<S: OutboxSink> Drop for Reservation<S> {
610    /// Returns the unspent slots to the reserved share.
611    fn drop(&mut self) {
612        self.outbox
613            .admission
614            .reserved
615            .fetch_sub(*self.slots.get_mut(), Ordering::AcqRel);
616    }
617}
618
619/// Names the outbox and the unspent slots, never the outbox's contents.
620impl<S: OutboxSink> std::fmt::Debug for Reservation<S> {
621    /// Write the sink's name and a relaxed snapshot of the unspent slots.
622    ///
623    /// # Errors
624    /// Returns the formatter's error when it rejects a write.
625    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
626        formatter
627            .debug_struct("Reservation")
628            .field("outbox", &S::NAME)
629            .field("slots", &self.slots.load(Ordering::Relaxed))
630            .finish()
631    }
632}
633
634/// The two shares an [`Outbox`] admits items and bytes into, shared with its
635/// writer.
636#[derive(Debug, Default)]
637struct Admission {
638    /// Unreserved items not yet written or lost.
639    unreserved: AtomicUsize,
640    /// Unspent reserved slots plus reserved items not yet written or lost.
641    reserved: AtomicUsize,
642    /// Reserved items pending in the writer, which shutdown releases from
643    /// `reserved` when it abandons them.
644    staged_reserved: AtomicUsize,
645    /// Bytes of unreserved items not yet written or lost.
646    unreserved_bytes: AtomicUsize,
647    /// Bytes of reserved items not yet written or lost; unspent slots hold
648    /// none.
649    reserved_bytes: AtomicUsize,
650}
651
652impl Admission {
653    /// Releases written or lost items from their shares.
654    fn release(&self, load: Load) {
655        self.reserved.fetch_sub(load.reserved, Ordering::AcqRel);
656        self.unreserved
657            .fetch_sub(load.items - load.reserved, Ordering::AcqRel);
658        self.reserved_bytes
659            .fetch_sub(load.reserved_bytes, Ordering::AcqRel);
660        self.unreserved_bytes
661            .fetch_sub(load.bytes - load.reserved_bytes, Ordering::AcqRel);
662    }
663}
664
665/// What a set of admitted items holds of the [`Admission`] shares.
666#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
667struct Load {
668    /// Items.
669    items: usize,
670    /// Items admitted on reserved slots.
671    reserved: usize,
672    /// Accounted bytes of every item.
673    bytes: usize,
674    /// Accounted bytes of the reserved items.
675    reserved_bytes: usize,
676}
677
678impl Load {
679    /// The load of one item of `bytes`, `reserved` or not.
680    const fn one(reserved: bool, bytes: usize) -> Self {
681        Self {
682            items: 1,
683            reserved: reserved as usize,
684            bytes,
685            reserved_bytes: if reserved { bytes } else { 0 },
686        }
687    }
688
689    /// Adds `other` into this load.
690    const fn add(&mut self, other: Self) {
691        self.items += other.items;
692        self.reserved += other.reserved;
693        self.bytes += other.bytes;
694        self.reserved_bytes += other.reserved_bytes;
695    }
696}
697
698/// When a tenant's last shed record was written and how many items it has
699/// shed since.
700#[derive(Debug, Default)]
701struct ShedLog {
702    /// When the last record was written; `None` before the first.
703    logged_at: Option<Instant>,
704    /// Items shed since that record.
705    unlogged: usize,
706}
707
708/// Admits `count` items into `share` unless that takes it past `limit`.
709///
710/// # Errors
711///
712/// Returns the share's count, admitting nothing, when `count` more would
713/// exceed `limit`.
714fn admit(share: &AtomicUsize, count: usize, limit: usize) -> Result<(), usize> {
715    share
716        .try_update(Ordering::AcqRel, Ordering::Acquire, |admitted| {
717            admitted.checked_add(count).filter(|total| *total <= limit)
718        })
719        .map(|_| ())
720}
721
722/// One lane's items in arrival order and their admitted [`Load`].
723type Backlog<I> = (Vec<I>, Load);
724
725/// A tenant and one of its [`OutboxSink::LANES`]: the unit written in arrival
726/// order, one write at a time.
727type Lane = (DataTenantId, usize);
728
729/// Outcome of one spawned lane write: the lane, its items, and the result,
730/// whose error is the sink's error or the message of a contained panic. The
731/// items come back only from a failed write, for its retry; a written slice's
732/// items are freed off the writer loop and come back empty.
733type Written<S> = (Lane, Vec<<S as OutboxSink>::Item>, Result<(), String>);
734
735/// The writer task behind one [`Outbox`].
736///
737/// Moved into its task by [`Outbox::new`]; it owns the receiving half of the
738/// queue, the per-lane backlogs, the in-flight writes, and the retry state.
739struct OutboxWriter<S: OutboxSink> {
740    /// Destination every write goes to.
741    sink: Arc<S>,
742    /// Receiving half of the queue.
743    requests: mpsc::UnboundedReceiver<Queued<S>>,
744    /// Received items per lane, in arrival order, not being written.
745    waiting: HashMap<Lane, Backlog<S::Item>>,
746    /// The exact slice of each lane whose last write failed; it is that
747    /// lane's next write, unmerged with `waiting`.
748    failed: HashMap<Lane, Backlog<S::Item>>,
749    /// In-flight lane writes.
750    writing: JoinSet<Written<S>>,
751    /// Lane, admitted load, and dispatch time of each in-flight write, keyed
752    /// by its task.
753    in_flight: HashMap<Id, (Lane, Load, Instant)>,
754    /// Earliest retry time and current backoff of each lane whose last
755    /// write failed.
756    retry_at: HashMap<Lane, (Instant, Duration)>,
757    /// Most writes run at once.
758    concurrency: usize,
759    /// Count shared with the handle; decremented once items are written.
760    pending: Arc<AtomicUsize>,
761    /// Admission shares shared with the handle; released with `pending`.
762    admission: Arc<Admission>,
763    /// Woken whenever a settled write leaves nothing pending.
764    idle: Arc<Notify>,
765    /// Cancelled when shutdown reaches its deadline.
766    abandon: CancellationToken,
767}
768
769impl<S: OutboxSink> OutboxWriter<S> {
770    /// The writer loop: receives items, dispatches lane writes, settles
771    /// finished ones, and wakes for retries, until shutdown has dropped the
772    /// sender and nothing is left, or shutdown abandons it at the deadline.
773    ///
774    /// Abandoning returns at once; dropping the writer cancels its in-flight
775    /// writes and discards their items, which shutdown then counts lost.
776    async fn run(mut self) {
777        let mut received = Vec::new();
778        let mut open = true;
779        loop {
780            self.dispatch();
781            if !open && self.waiting.is_empty() && self.failed.is_empty() && self.writing.is_empty()
782            {
783                return;
784            }
785            let retry = self.next_retry();
786            tokio::select! {
787                biased;
788                () = self.abandon.cancelled() => return,
789                Some(done) = self.writing.join_next_with_id(), if !self.writing.is_empty() => {
790                    self.finish(done);
791                }
792                count = self.requests.recv_many(&mut received, usize::MAX), if open => {
793                    open = count != 0;
794                    let lanes = S::LANES.max(1) as u64;
795                    for (tenant, item, load) in received.drain(..) {
796                        // The remainder is below `LANES`, a `usize`.
797                        let lane = usize::try_from(S::lane(&item) % lanes).unwrap_or_default();
798                        let (items, backlog) = self.waiting.entry((tenant, lane)).or_default();
799                        items.push(item);
800                        backlog.add(load);
801                    }
802                }
803                () = async {
804                    match retry {
805                        Some(at) => sleep_until(at.into()).await,
806                        None => pending().await,
807                    }
808                } => {}
809            }
810        }
811    }
812
813    /// Starts one write for each lane that has items waiting, has no write
814    /// in flight, and is not backing off, while fewer than `concurrency` writes
815    /// run. Ready lanes are taken in order of their tenant's writes in flight
816    /// plus the lane's rank among that tenant's ready lanes, so every tenant's
817    /// first lane goes before any tenant's second. A lane with a failed slice
818    /// writes exactly that slice again; otherwise a write takes the lane's
819    /// whole backlog.
820    ///
821    /// The write task keeps its items through a panicking sink, including a
822    /// panic while rendering the sink's error, and returns them with the panic
823    /// as the error, so they are retried like any failed write.
824    fn dispatch(&mut self) {
825        if self.writing.len() >= self.concurrency {
826            return;
827        }
828        let now = Instant::now();
829        let mut busy: HashMap<DataTenantId, usize> = HashMap::new();
830        for ((tenant, _), ..) in self.in_flight.values() {
831            *busy.entry(*tenant).or_default() += 1;
832        }
833        let mut ready: Vec<(usize, Lane)> = self
834            .waiting
835            .keys()
836            .chain(
837                self.failed
838                    .keys()
839                    .filter(|lane| !self.waiting.contains_key(*lane)),
840            )
841            .filter(|lane| self.is_ready(**lane, now))
842            .map(|lane| {
843                let order = busy.entry(lane.0).or_default();
844                *order += 1;
845                (*order, *lane)
846            })
847            .collect();
848        ready.sort_unstable_by_key(|(order, _)| *order);
849        for (_, lane) in ready {
850            if self.writing.len() >= self.concurrency {
851                return;
852            }
853            let Some((items, load)) = self
854                .failed
855                .remove(&lane)
856                .or_else(|| self.waiting.remove(&lane))
857            else {
858                continue;
859            };
860            let (tenant, _) = lane;
861            let sink = Arc::clone(&self.sink);
862            let task = self.writing.spawn(async move {
863                // The constructing closure borrows `items`, so it is dropped
864                // before awaiting, and the write before `items` is returned.
865                let outcome = {
866                    let write = catch_unwind(AssertUnwindSafe(|| sink.write(tenant, &items)));
867                    match write {
868                        Ok(write) => contain_panic(write).await,
869                        Err(panic) => Err(panic),
870                    }
871                };
872                // Rendering a sink error runs sink code too, so it is
873                // contained as well and `items` is still returned.
874                let result = catch_unwind(AssertUnwindSafe(|| match outcome {
875                    Ok(written) => written.map_err(|error| error.to_string()),
876                    Err(panic) => Err(format!("sink write panicked: {}", panic_message(&*panic))),
877                }))
878                .unwrap_or_else(|panic| {
879                    Err(format!(
880                        "sink error rendering panicked: {}",
881                        panic_message(&*panic)
882                    ))
883                });
884                if result.is_ok() {
885                    // A large slice's items can take as long to free as to
886                    // write; freeing them on the blocking pool keeps the
887                    // writer loop free to receive and dispatch the next write.
888                    drop(tokio::task::spawn_blocking(move || drop(items)));
889                    return (lane, Vec::new(), result);
890                }
891                (lane, items, result)
892            });
893            self.in_flight
894                .insert(task.id(), (lane, load, Instant::now()));
895        }
896    }
897
898    /// Whether `lane` has no write in flight and its backoff has elapsed.
899    fn is_ready(&self, lane: Lane, now: Instant) -> bool {
900        !self.in_flight.values().any(|(busy, ..)| *busy == lane)
901            && self.retry_at.get(&lane).is_none_or(|(at, _)| *at <= now)
902    }
903
904    /// The earliest future retry time of a failed lane, or `None` when no
905    /// timer is needed because every write slot is busy or nobody backs off.
906    fn next_retry(&self) -> Option<Instant> {
907        if self.writing.len() >= self.concurrency {
908            return None;
909        }
910        let now = Instant::now();
911        self.retry_at
912            .iter()
913            .filter(|(lane, (at, _))| *at > now && self.failed.contains_key(*lane))
914            .map(|(_, (at, _))| *at)
915            .min()
916    }
917
918    /// Settles one finished write.
919    ///
920    /// Every joined write, succeeded or failed, records its dispatch-to-settle
921    /// time in `outbox_write_seconds`. Success clears the lane's backoff,
922    /// counts the items and bytes written, and releases the slots of the
923    /// items it wrote, which the write task already handed to the blocking
924    /// pool to free. Failure,
925    /// including a contained sink panic, keeps the items as the lane's next
926    /// write, unmerged with later arrivals, doubles its backoff, and counts
927    /// and logs the attempt. A task that failed to join lost its items, which
928    /// are counted lost with [`LossReason::Rejected`]; with sink panics
929    /// contained that needs a panic that escapes containment, such as one
930    /// raised while dropping a panic payload.
931    fn finish(&mut self, done: Result<(Id, Written<S>), JoinError>) {
932        let ((lane, items, result), load) = match done {
933            Ok((id, written)) => {
934                let load =
935                    self.in_flight
936                        .remove(&id)
937                        .map_or_else(Load::default, |(_, load, started)| {
938                            metrics::histogram!("outbox_write_seconds", "outbox" => S::NAME)
939                                .record(started.elapsed().as_secs_f64());
940                            load
941                        });
942                (written, load)
943            }
944            Err(error) => {
945                if let Some(((tenant, lane), load, _)) = self.in_flight.remove(&error.id()) {
946                    count_lost::<S>(load.items, LossReason::Rejected);
947                    tracing::error!(outbox = S::NAME, %tenant, lane, %error, lost = load.items, "outbox write task failed");
948                    self.release(load);
949                }
950                return;
951            }
952        };
953        let error = match result {
954            Ok(()) => {
955                self.retry_at.remove(&lane);
956                metrics::counter!("outbox_items_written_total", "outbox" => S::NAME)
957                    .increment(load.items as u64);
958                metrics::counter!("outbox_bytes_written_total", "outbox" => S::NAME)
959                    .increment(load.bytes as u64);
960                self.release(load);
961                return;
962            }
963            Err(error) => error,
964        };
965        let backoff = self
966            .retry_at
967            .get(&lane)
968            .map_or(INITIAL_BACKOFF, |(_, last)| (*last * 2).min(MAX_BACKOFF));
969        self.retry_at
970            .insert(lane, (Instant::now() + backoff, backoff));
971        metrics::counter!("outbox_write_failures_total", "outbox" => S::NAME).increment(1);
972        tracing::error!(
973            outbox = S::NAME,
974            tenant = %lane.0,
975            lane = lane.1,
976            %error,
977            items = items.len(),
978            retry_in_ms = backoff.as_millis(),
979            "outbox write failed; retrying"
980        );
981        self.failed.insert(lane, (items, load));
982    }
983
984    /// Releases written or lost items from pending and their shares, and
985    /// wakes settlers when nothing remains pending.
986    fn release(&self, load: Load) {
987        metrics::gauge!("outbox_pending", "outbox" => S::NAME).decrement(load.items as f64);
988        self.admission
989            .staged_reserved
990            .fetch_sub(load.reserved, Ordering::AcqRel);
991        self.admission.release(load);
992        if self.pending.fetch_sub(load.items, Ordering::AcqRel) == load.items {
993            self.idle.notify_waiters();
994        }
995    }
996}
997
998/// Polls `future` to completion and returns its output, or the payload of a
999/// panic raised while polling it.
1000///
1001/// Containing the panic inside the polled task, rather than letting it unwind
1002/// the task, keeps the caller's owned state, such as a write's items, intact.
1003///
1004/// # Errors
1005///
1006/// Returns the panic payload when polling `future` panics; the future is then
1007/// dropped without being polled again.
1008async fn contain_panic<F: Future>(future: F) -> Result<F::Output, Box<dyn Any + Send>> {
1009    let mut future = pin!(future);
1010    poll_fn(
1011        |context| match catch_unwind(AssertUnwindSafe(|| future.as_mut().poll(context))) {
1012            Ok(poll) => poll.map(Ok),
1013            Err(panic) => Poll::Ready(Err(panic)),
1014        },
1015    )
1016    .await
1017}
1018
1019/// The message of a panic payload, when it is a string.
1020fn panic_message(panic: &(dyn Any + Send)) -> &str {
1021    panic
1022        .downcast_ref::<&str>()
1023        .copied()
1024        .or_else(|| panic.downcast_ref::<String>().map(String::as_str))
1025        .unwrap_or("non-string panic payload")
1026}
1027
1028/// Counts `count` items of outbox `S` as lost in
1029/// `outbox_events_lost_total{outbox, reason}`.
1030///
1031/// The outbox counts shed, shutdown, and late-staging loss itself; a sink
1032/// calls this for items it consumes without writing, such as a terminal
1033/// rejection that a retry could never cure, and a producer for items it
1034/// sheds at its own limit before staging.
1035pub fn count_lost<S: OutboxSink>(count: usize, reason: LossReason) {
1036    metrics::counter!("outbox_events_lost_total", "outbox" => S::NAME, "reason" => reason.as_str())
1037        .increment(count as u64);
1038}
1039
1040/// Counts and logs one item of `tenant` lost because it was staged after
1041/// shutdown began or the writer is gone.
1042fn lose_late<S: OutboxSink>(tenant: DataTenantId) {
1043    count_lost::<S>(1, LossReason::Shutdown);
1044    tracing::error!(outbox = S::NAME, %tenant, "outbox item staged after shutdown was lost");
1045}
1046
1047#[cfg(test)]
1048mod tests {
1049    //! The generic outbox against an in-memory sink whose tenants can be made
1050    //! to fail on demand.
1051
1052    use std::collections::{HashMap, HashSet};
1053    use std::fmt::{self, Display, Formatter};
1054    use std::future::{Future, poll_fn};
1055    use std::pin::pin;
1056    use std::sync::atomic::Ordering;
1057    use std::sync::{Arc, Mutex, MutexGuard};
1058    use std::task::Poll;
1059    use std::time::{Duration, Instant};
1060
1061    use metrics::atomics::AtomicU64;
1062    use metrics::{
1063        Counter, Gauge, Histogram, Key, KeyName, Metadata, Recorder, SharedString, Unit,
1064    };
1065    use tokio::sync::Notify;
1066    use wyrd_spec::DataTenantId;
1067
1068    use super::{Outbox, OutboxLimits, OutboxSink};
1069
1070    /// In-memory sink recording each tenant's written items in order.
1071    #[derive(Default)]
1072    struct MemorySink {
1073        /// Shared record of written items and failing tenants.
1074        state: Arc<Mutex<MemoryState>>,
1075        /// Writes of hanging tenants block until notified, simulating a hang.
1076        hang: Arc<Notify>,
1077        /// Notified as each write starts, so a test can wait for a dispatch
1078        /// without polling.
1079        started: Arc<Notify>,
1080    }
1081
1082    /// What the memory sink has written and which tenants fail or hang.
1083    #[derive(Default)]
1084    struct MemoryState {
1085        /// Written items per tenant, in write order.
1086        written: HashMap<DataTenantId, Vec<u32>>,
1087        /// Tenants whose writes fail.
1088        failing: HashSet<DataTenantId>,
1089        /// Tenants whose writes hang.
1090        hanging: HashSet<DataTenantId>,
1091        /// Tenants whose next write panics once.
1092        panicking: HashSet<DataTenantId>,
1093        /// Tenants whose next write fails once with an error whose `Display`
1094        /// panics.
1095        display_panicking: HashSet<DataTenantId>,
1096        /// Failed write attempts.
1097        failures: usize,
1098        /// Every write attempt's tenant and slice, in attempt order.
1099        attempts: Vec<(DataTenantId, Vec<u32>)>,
1100    }
1101
1102    /// Failure the memory sink reports for an injected write failure.
1103    #[derive(Debug)]
1104    struct MemoryError {
1105        /// Whether rendering this error panics, simulating a sink error type
1106        /// whose `Display` implementation is faulty.
1107        panics_on_display: bool,
1108    }
1109
1110    impl Display for MemoryError {
1111        /// Renders the injected failure message.
1112        ///
1113        /// # Panics
1114        ///
1115        /// Panics when `panics_on_display` is set, so a test can prove the
1116        /// writer contains a panic raised while rendering a sink error.
1117        fn fmt(&self, formatter: &mut Formatter<'_>) -> fmt::Result {
1118            assert!(!self.panics_on_display, "injected error display panic");
1119            formatter.write_str("injected failure")
1120        }
1121    }
1122
1123    impl OutboxSink for MemorySink {
1124        type Item = u32;
1125        type Error = MemoryError;
1126        const NAME: &'static str = "test";
1127
1128        /// Accounts each item as its own value in bytes, so a test sizes
1129        /// items by choosing their values.
1130        fn weight(item: &u32) -> usize {
1131            *item as usize
1132        }
1133
1134        /// Writes `items` for `tenant`, applying the injections configured in
1135        /// [`MemoryState`] in a fixed order the tests depend on.
1136        ///
1137        /// The order is: notify `started` that a write was dispatched; if the
1138        /// tenant is hanging, wait for `hang`; if the tenant has a one-shot
1139        /// panic injected, consume it and panic; if the tenant has a one-shot
1140        /// panicking-display failure injected, consume it, count a failure,
1141        /// and fail; if the tenant is failing, count a failure and fail;
1142        /// otherwise append `items` to the tenant's written record.
1143        ///
1144        /// # Errors
1145        ///
1146        /// Returns a [`MemoryError`] when the tenant is failing or has a
1147        /// panicking-display failure injected; that error's `Display` panics
1148        /// only in the latter case.
1149        ///
1150        /// # Panics
1151        ///
1152        /// Panics once, after the hang stage, when a panic is injected for the
1153        /// tenant.
1154        ///
1155        /// # Cancellation
1156        ///
1157        /// Dropping the future while it waits for `hang` leaves the state
1158        /// unchanged, so the batch is neither written nor counted as failed.
1159        async fn write(&self, tenant: DataTenantId, items: &[u32]) -> Result<(), MemoryError> {
1160            self.started.notify_one();
1161            let hangs = self.lock().hanging.contains(&tenant);
1162            if hangs {
1163                self.hang.notified().await;
1164            }
1165            let panics = self.lock().panicking.remove(&tenant);
1166            assert!(!panics, "injected sink panic");
1167            let mut state = self.lock();
1168            state.attempts.push((tenant, items.to_vec()));
1169            if state.display_panicking.remove(&tenant) {
1170                state.failures += 1;
1171                return Err(MemoryError {
1172                    panics_on_display: true,
1173                });
1174            }
1175            if state.failing.contains(&tenant) {
1176                state.failures += 1;
1177                return Err(MemoryError {
1178                    panics_on_display: false,
1179                });
1180            }
1181            state
1182                .written
1183                .entry(tenant)
1184                .or_default()
1185                .extend_from_slice(items);
1186            Ok(())
1187        }
1188    }
1189
1190    impl MemorySink {
1191        /// Locks the shared state, which no test holds across an await.
1192        fn lock(&self) -> MutexGuard<'_, MemoryState> {
1193            self.state
1194                .lock()
1195                .expect("memory sink state is never poisoned")
1196        }
1197    }
1198
1199    /// An item whose drop blocks until its gate opens or ten seconds pass,
1200    /// standing in for a slice that is slow to free.
1201    struct SlowDrop(Mutex<std::sync::mpsc::Receiver<()>>);
1202
1203    impl Drop for SlowDrop {
1204        /// Blocks until the gate opens, its sender is dropped, or ten
1205        /// seconds pass.
1206        fn drop(&mut self) {
1207            if let Ok(gate) = self.0.lock() {
1208                let _ = gate.recv_timeout(Duration::from_secs(10));
1209            }
1210        }
1211    }
1212
1213    /// A sink that accepts every write of [`SlowDrop`] items.
1214    struct SlowDropSink;
1215
1216    impl OutboxSink for SlowDropSink {
1217        type Item = SlowDrop;
1218        type Error = MemoryError;
1219        const NAME: &'static str = "slow_drop";
1220
1221        /// Accepts `items` without inspecting them.
1222        ///
1223        /// # Errors
1224        ///
1225        /// Never fails.
1226        fn write(
1227            &self,
1228            _tenant: DataTenantId,
1229            _items: &[SlowDrop],
1230        ) -> impl Future<Output = Result<(), MemoryError>> + Send {
1231            std::future::ready(Ok(()))
1232        }
1233    }
1234
1235    /// Proves freeing a written slice does not hold up the writer: while the
1236    /// first slice's item is still blocked in its drop, a slice staged after
1237    /// it is written and settles.
1238    ///
1239    /// # Panics
1240    ///
1241    /// Panics when the second slice waits for the first slice's drop. Two
1242    /// workers keep a blocked writer loop from also blocking the test.
1243    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
1244    async fn freeing_a_written_slice_does_not_hold_up_the_next_write() {
1245        let outbox = Outbox::new(SlowDropSink, 1, ROOMY);
1246        let tenant = DataTenantId::new_v7();
1247        let (gate, held) = std::sync::mpsc::channel();
1248        outbox.stage(tenant, SlowDrop(Mutex::new(held)));
1249        let first = outbox.settle(Instant::now() + Duration::from_secs(2)).await;
1250        let (_, open) = std::sync::mpsc::channel();
1251        outbox.stage(tenant, SlowDrop(Mutex::new(open)));
1252        let second = outbox.settle(Instant::now() + Duration::from_secs(2)).await;
1253        drop(gate);
1254        assert_eq!((first, second), (0, 0), "the next write waited for a drop");
1255    }
1256
1257    /// A two-lane sink, odd items in lane 1 and even in lane 0, that records
1258    /// each written slice and holds the slice `[0]` until `gate` is notified.
1259    #[derive(Default)]
1260    struct LaneSink {
1261        /// Releases the held `[0]` write.
1262        gate: Notify,
1263        /// Every written slice, in completion order.
1264        written: Mutex<Vec<Vec<u32>>>,
1265    }
1266
1267    impl OutboxSink for LaneSink {
1268        type Item = u32;
1269        type Error = MemoryError;
1270        const NAME: &'static str = "lanes";
1271        const LANES: usize = 2;
1272
1273        /// An item's parity is its lane.
1274        fn lane(item: &u32) -> u64 {
1275            u64::from(item % 2)
1276        }
1277
1278        /// Records `items` once written, first waiting for `gate` when the
1279        /// slice is exactly `[0]`.
1280        ///
1281        /// # Errors
1282        ///
1283        /// Never fails.
1284        ///
1285        /// # Panics
1286        ///
1287        /// Panics if the record's lock is poisoned.
1288        async fn write(&self, _tenant: DataTenantId, items: &[u32]) -> Result<(), MemoryError> {
1289            if items == [0] {
1290                self.gate.notified().await;
1291            }
1292            self.written
1293                .lock()
1294                .expect("unpoisoned")
1295                .push(items.to_vec());
1296            Ok(())
1297        }
1298    }
1299
1300    /// One tenant's lanes write concurrently while each lane keeps arrival
1301    /// order: lane 1 commits while lane 0's first write hangs, and lane 0's
1302    /// later items wait for that write instead of passing it.
1303    ///
1304    /// # Panics
1305    ///
1306    /// Panics when lane 1 waits for lane 0, or lane 0 writes out of order.
1307    #[tokio::test]
1308    async fn a_tenants_lanes_write_concurrently_and_each_keeps_order() {
1309        let outbox = Outbox::new(LaneSink::default(), 2, ROOMY);
1310        let tenant = DataTenantId::new_v7();
1311        outbox.stage(tenant, 0_u32);
1312        outbox.stage(tenant, 1_u32);
1313        let deadline = within(10);
1314        while outbox.sink().written.lock().expect("unpoisoned").is_empty() {
1315            assert!(Instant::now() < deadline, "lane 1 waited for lane 0");
1316            tokio::time::sleep(Duration::from_millis(5)).await;
1317        }
1318        outbox.stage(tenant, 2_u32);
1319        outbox.stage(tenant, 4_u32);
1320        outbox.sink().gate.notify_one();
1321
1322        assert_eq!(outbox.settle(within(10)).await, 0);
1323        assert_eq!(
1324            *outbox.sink().written.lock().expect("unpoisoned"),
1325            vec![vec![1], vec![0], vec![2, 4]]
1326        );
1327    }
1328
1329    /// Capacity of test outboxes that do not exercise the bound.
1330    const ROOMY: OutboxLimits = items(1_024);
1331
1332    /// Limits of `items` items and unlimited bytes.
1333    const fn items(items: usize) -> OutboxLimits {
1334        OutboxLimits {
1335            items,
1336            bytes: usize::MAX,
1337        }
1338    }
1339
1340    /// Starts an outbox over a fresh memory sink and returns both.
1341    fn outbox(concurrency: usize) -> (Arc<Outbox<MemorySink>>, Arc<Mutex<MemoryState>>) {
1342        let sink = MemorySink::default();
1343        let state = Arc::clone(&sink.state);
1344        (Outbox::new(sink, concurrency, ROOMY), state)
1345    }
1346
1347    /// Starts an outbox over `sink` and returns it with the sink's state and
1348    /// write-start signal.
1349    fn outbox_over(
1350        sink: MemorySink,
1351    ) -> (
1352        Arc<Outbox<MemorySink>>,
1353        Arc<Mutex<MemoryState>>,
1354        Arc<Notify>,
1355    ) {
1356        let state = Arc::clone(&sink.state);
1357        let started = Arc::clone(&sink.started);
1358        (Outbox::new(sink, 4, ROOMY), state, started)
1359    }
1360
1361    /// Thread-local metrics recorder keeping every counter and gauge by name,
1362    /// and by name and `reason` label when it has one.
1363    ///
1364    /// Installed with `metrics::set_default_local_recorder` in a
1365    /// current-thread test, it also sees the outbox writer's updates because
1366    /// every task runs on the test thread.
1367    #[derive(Default)]
1368    struct TestMetrics {
1369        /// Raw value of each metric, keyed `name` or `name{reason=...}`;
1370        /// gauges store `f64` bits.
1371        values: Mutex<HashMap<String, Arc<AtomicU64>>>,
1372    }
1373
1374    impl TestMetrics {
1375        /// The shared value cell of the metric `key` names, one per `reason`
1376        /// label value.
1377        fn cell(&self, key: &Key) -> Arc<AtomicU64> {
1378            let name = key
1379                .labels()
1380                .find(|label| label.key() == "reason")
1381                .map_or_else(
1382                    || key.name().to_owned(),
1383                    |reason| format!("{}{{reason={}}}", key.name(), reason.value()),
1384                );
1385            let mut values = self.values.lock().expect("test metrics are never poisoned");
1386            Arc::clone(values.entry(name).or_default())
1387        }
1388
1389        /// The raw value of metric `name`, summed over its `reason` labels
1390        /// unless `name` names one as `name{reason=...}`; zero when it was
1391        /// never touched.
1392        fn raw(&self, name: &str) -> u64 {
1393            let labelled = format!("{name}{{");
1394            self.values
1395                .lock()
1396                .expect("test metrics are never poisoned")
1397                .iter()
1398                .filter(|(key, _)| *key == name || key.starts_with(&labelled))
1399                .map(|(_, value)| value.load(Ordering::Acquire))
1400                .sum()
1401        }
1402
1403        /// The value of counter `name`.
1404        fn counter(&self, name: &str) -> u64 {
1405            self.raw(name)
1406        }
1407
1408        /// The value of gauge `name`.
1409        fn gauge(&self, name: &str) -> f64 {
1410            f64::from_bits(self.raw(name))
1411        }
1412    }
1413
1414    impl Recorder for TestMetrics {
1415        /// Intentionally ignores counter descriptions; tests assert values
1416        /// only.
1417        fn describe_counter(&self, _: KeyName, _: Option<Unit>, _: SharedString) {}
1418
1419        /// Intentionally ignores gauge descriptions; tests assert values only.
1420        fn describe_gauge(&self, _: KeyName, _: Option<Unit>, _: SharedString) {}
1421
1422        /// Intentionally ignores histogram descriptions; histograms are not
1423        /// recorded.
1424        fn describe_histogram(&self, _: KeyName, _: Option<Unit>, _: SharedString) {}
1425
1426        /// Records the counter by metric name, ignoring labels, so every
1427        /// registration of one name shares a single value cell.
1428        fn register_counter(&self, key: &Key, _: &Metadata<'_>) -> Counter {
1429            Counter::from_arc(self.cell(key))
1430        }
1431
1432        /// Records the gauge by metric name, ignoring labels, storing its
1433        /// `f64` value as bits in a shared cell.
1434        fn register_gauge(&self, key: &Key, _: &Metadata<'_>) -> Gauge {
1435            Gauge::from_arc(self.cell(key))
1436        }
1437
1438        /// Intentionally records nothing: the outbox emits no histogram the
1439        /// tests assert on.
1440        fn register_histogram(&self, _: &Key, _: &Metadata<'_>) -> Histogram {
1441            Histogram::noop()
1442        }
1443    }
1444
1445    /// A deadline `seconds` from now.
1446    fn within(seconds: u64) -> Instant {
1447        Instant::now() + Duration::from_secs(seconds)
1448    }
1449
1450    /// Waits until the sink has recorded at least `failures` failed writes.
1451    async fn await_failures(state: &Mutex<MemoryState>, failures: usize) {
1452        let deadline = within(10);
1453        while state.lock().expect("unpoisoned").failures < failures {
1454            assert!(Instant::now() < deadline, "the write never failed");
1455            tokio::time::sleep(Duration::from_millis(5)).await;
1456        }
1457    }
1458
1459    /// A failed write is retried and committed exactly once, in order, ahead of
1460    /// the tenant's later items, and `pending` returns to zero.
1461    #[tokio::test]
1462    async fn a_failed_write_is_retried_once_in_order_ahead_of_later_items() {
1463        let (outbox, state) = outbox(4);
1464        let tenant = DataTenantId::new_v7();
1465        state.lock().expect("unpoisoned").failing.insert(tenant);
1466        outbox.stage(tenant, 1_u32);
1467        outbox.stage(tenant, 2_u32);
1468        await_failures(&state, 2).await;
1469        outbox.stage(tenant, 3_u32);
1470        assert_eq!(outbox.pending(), 3, "failed items stay pending");
1471        state.lock().expect("unpoisoned").failing.clear();
1472
1473        assert_eq!(outbox.settle(within(10)).await, 0, "the retry drains");
1474        assert_eq!(
1475            state.lock().expect("unpoisoned").written.get(&tenant),
1476            Some(&vec![1, 2, 3]),
1477            "retried items commit once, ahead of later items"
1478        );
1479        assert_eq!(outbox.pending(), 0);
1480    }
1481
1482    /// A retry writes exactly the failed slice: every attempt sees the same
1483    /// items, and an item staged after the failure is written separately once
1484    /// the retry succeeds, so a sink deriving identity from its slice presents
1485    /// one identity across attempts.
1486    #[tokio::test]
1487    async fn a_retry_writes_the_identical_failed_slice_without_later_items() {
1488        let (outbox, state) = outbox(4);
1489        let tenant = DataTenantId::new_v7();
1490        state.lock().expect("unpoisoned").failing.insert(tenant);
1491        outbox.stage(tenant, 1_u32);
1492        outbox.stage(tenant, 2_u32);
1493        await_failures(&state, 1).await;
1494        outbox.stage(tenant, 3_u32);
1495        await_failures(&state, 2).await;
1496        state.lock().expect("unpoisoned").failing.clear();
1497
1498        assert_eq!(outbox.settle(within(10)).await, 0, "the retry drains");
1499        let state = state.lock().expect("unpoisoned");
1500        let (last, retries) = state
1501            .attempts
1502            .split_last()
1503            .expect("the later item was written");
1504        assert_eq!(last, &(tenant, vec![3]), "the later item is its own write");
1505        assert!(retries.len() >= 3, "two failures and one successful retry");
1506        assert!(
1507            retries
1508                .iter()
1509                .all(|attempt| attempt == &(tenant, vec![1, 2])),
1510            "every retry is the identical failed slice: {retries:?}"
1511        );
1512        assert_eq!(state.written.get(&tenant), Some(&vec![1, 2, 3]));
1513    }
1514
1515    /// A tenant whose writes keep failing does not delay another tenant.
1516    #[tokio::test]
1517    async fn a_failing_tenant_does_not_delay_another_tenant() {
1518        let (outbox, state) = outbox(1);
1519        let failing = DataTenantId::new_v7();
1520        let healthy = DataTenantId::new_v7();
1521        state.lock().expect("unpoisoned").failing.insert(failing);
1522        outbox.stage(failing, 1_u32);
1523        await_failures(&state, 1).await;
1524        outbox.stage(healthy, 2_u32);
1525
1526        let deadline = within(10);
1527        while !state
1528            .lock()
1529            .expect("unpoisoned")
1530            .written
1531            .contains_key(&healthy)
1532        {
1533            assert!(Instant::now() < deadline, "the healthy tenant was delayed");
1534            tokio::time::sleep(Duration::from_millis(5)).await;
1535        }
1536        assert_eq!(outbox.pending(), 1, "only the failing tenant's item waits");
1537        state.lock().expect("unpoisoned").failing.clear();
1538        assert_eq!(outbox.settle(within(10)).await, 0);
1539    }
1540
1541    /// The bound covers in-flight and retrying work, not just the queue: with
1542    /// one tenant's write hanging and another's retrying, a reservation past
1543    /// the reserved share is refused, retries do not charge their slice
1544    /// again, and capacity returns once the work is written.
1545    ///
1546    /// # Panics
1547    /// Panics when a reservation within the share is refused, one past it is
1548    /// admitted, the refusal reports other counts, or the work is not written
1549    /// once in order.
1550    #[tokio::test]
1551    async fn bounded_staging_covers_retrying_and_inflight_work() {
1552        let sink = MemorySink::default();
1553        let (hang, started) = (Arc::clone(&sink.hang), Arc::clone(&sink.started));
1554        let state = Arc::clone(&sink.state);
1555        // Reservations may hold three of these four items.
1556        let outbox = Outbox::new(sink, 4, items(4));
1557        let (hanging, failing) = (DataTenantId::new_v7(), DataTenantId::new_v7());
1558        {
1559            let mut state = state.lock().expect("unpoisoned");
1560            state.hanging.insert(hanging);
1561            state.failing.insert(failing);
1562        }
1563        let reserved = outbox.reserve(1).expect("room for one reservation");
1564        reserved.stage(hanging, 1_u32);
1565        started.notified().await;
1566        let retried = outbox.reserve(2).expect("room for a two-item reservation");
1567        retried.stage(failing, 2_u32);
1568        retried.stage(failing, 3_u32);
1569        await_failures(&state, 3).await;
1570        assert_eq!(outbox.pending(), 3, "a retried slice is admitted once");
1571
1572        let refused = outbox
1573            .reserve(1)
1574            .expect_err("in-flight and retrying work fill the share");
1575        assert_eq!(
1576            (refused.admitted, refused.limit, refused.unit),
1577            (3, 3, "items")
1578        );
1579        outbox.stage(failing, 4_u32);
1580        outbox.stage(failing, 5_u32);
1581        assert_eq!(
1582            outbox.pending(),
1583            4,
1584            "unreserved staging past capacity is lost"
1585        );
1586
1587        {
1588            let mut state = state.lock().expect("unpoisoned");
1589            state.hanging.clear();
1590            state.failing.clear();
1591        }
1592        hang.notify_one();
1593        assert_eq!(
1594            outbox.settle(within(10)).await,
1595            0,
1596            "the bounded work drains"
1597        );
1598        drop((reserved, retried));
1599        let recovered = outbox
1600            .reserve(3)
1601            .expect("written work releases its capacity");
1602        recovered.stage(hanging, 6_u32);
1603        assert_eq!(outbox.settle(within(10)).await, 0);
1604        let state = state.lock().expect("unpoisoned");
1605        assert_eq!(state.written[&hanging], vec![1, 6]);
1606        assert_eq!(state.written[&failing], vec![2, 3, 4]);
1607    }
1608
1609    /// Unreserved staging on a full outbox, as audit decisions use it, returns
1610    /// at once and counts the item lost; reservations cannot take the share
1611    /// kept for it, unreserved items never count against reservations, and a
1612    /// dropped reservation returns its unspent slots.
1613    ///
1614    /// # Panics
1615    /// Panics when a full outbox blocks or keeps an unreserved item, the lost
1616    /// count differs, or reservations are admitted or refused against the
1617    /// wrong share.
1618    #[tokio::test(flavor = "current_thread")]
1619    async fn bounded_audit_staging_remains_nonblocking() {
1620        let metrics = TestMetrics::default();
1621        let _metrics = metrics::set_default_local_recorder(&metrics);
1622        let sink = MemorySink::default();
1623        let state = Arc::clone(&sink.state);
1624        let outbox = Outbox::new(sink, 4, items(4));
1625        let tenant = DataTenantId::new_v7();
1626        state.lock().expect("unpoisoned").hanging.insert(tenant);
1627
1628        let unspent = outbox.reserve(3).expect("the reserved share");
1629        outbox
1630            .reserve(1)
1631            .expect_err("reservations hold their whole share");
1632        outbox.stage(tenant, 1_u32);
1633        assert_eq!(
1634            outbox.pending(),
1635            1,
1636            "the kept share admits unreserved staging"
1637        );
1638        outbox.stage(tenant, 2_u32);
1639        assert_eq!(
1640            outbox.pending(),
1641            1,
1642            "a full unreserved share refuses without waiting"
1643        );
1644        assert_eq!(
1645            metrics.counter("outbox_events_lost_total{reason=shed}"),
1646            1,
1647            "the refusal is counted as shed"
1648        );
1649
1650        drop(unspent);
1651        outbox
1652            .reserve(3)
1653            .expect("a dropped reservation returns its slots despite the hanging item");
1654    }
1655
1656    /// Shutdown keeps retrying until the deadline and writes what recovers.
1657    #[tokio::test]
1658    async fn shutdown_flushes_items_that_recover_before_the_deadline() {
1659        let (outbox, state) = outbox(4);
1660        let tenant = DataTenantId::new_v7();
1661        state.lock().expect("unpoisoned").failing.insert(tenant);
1662        outbox.stage(tenant, 1_u32);
1663        await_failures(&state, 1).await;
1664        let shutdown = {
1665            let outbox = Arc::clone(&outbox);
1666            tokio::spawn(async move { outbox.shutdown(within(10)).await })
1667        };
1668        await_failures(&state, 2).await;
1669        state.lock().expect("unpoisoned").failing.clear();
1670
1671        assert_eq!(shutdown.await.expect("shutdown joins"), 0, "nothing lost");
1672        assert_eq!(
1673            state.lock().expect("unpoisoned").written.get(&tenant),
1674            Some(&vec![1])
1675        );
1676        outbox.stage(tenant, 2_u32);
1677        assert_eq!(outbox.pending(), 0, "a shut-down outbox accepts nothing");
1678    }
1679
1680    /// Shutdown gives up at its deadline, reports failing and hanging items as
1681    /// lost exactly once, and leaves nothing pending in the count or gauge.
1682    #[tokio::test(flavor = "current_thread")]
1683    async fn shutdown_counts_items_unwritten_at_the_deadline_as_lost() {
1684        let metrics = TestMetrics::default();
1685        let _metrics = metrics::set_default_local_recorder(&metrics);
1686        let (outbox, state) = outbox(4);
1687        let failing = DataTenantId::new_v7();
1688        let hanging = DataTenantId::new_v7();
1689        {
1690            let mut state = state.lock().expect("unpoisoned");
1691            state.failing.insert(failing);
1692            state.hanging.insert(hanging);
1693        }
1694        outbox.stage(failing, 1_u32);
1695        outbox.stage(failing, 2_u32);
1696        outbox.stage(hanging, 3_u32);
1697        await_failures(&state, 1).await;
1698
1699        let lost = outbox
1700            .shutdown(Instant::now() + Duration::from_millis(200))
1701            .await;
1702        assert_eq!(lost, 3, "every unwritten item is reported lost");
1703        assert!(state.lock().expect("unpoisoned").written.is_empty());
1704        assert_eq!(
1705            metrics.counter("outbox_events_lost_total{reason=shutdown}"),
1706            3,
1707            "lost once, at shutdown"
1708        );
1709        assert_eq!(outbox.pending(), 0, "abandoned items are not pending");
1710        assert!(
1711            metrics.gauge("outbox_pending").abs() < f64::EPSILON,
1712            "the pending gauge returns to zero"
1713        );
1714    }
1715
1716    /// Concurrent shutdowns over reserved and unreserved items unwritten at
1717    /// the deadline report the loss once, settle both shares exactly, and
1718    /// leave an unspent reservation holding its slot until it drops.
1719    ///
1720    /// # Panics
1721    /// Panics when a reservation is refused, the loss is reported other than
1722    /// once, or a share or count does not settle.
1723    #[tokio::test(flavor = "current_thread")]
1724    async fn concurrent_shutdown_releases_reserved_losses_once() {
1725        let metrics = TestMetrics::default();
1726        let _metrics = metrics::set_default_local_recorder(&metrics);
1727        let (outbox, state) = outbox(4);
1728        let hanging = DataTenantId::new_v7();
1729        state.lock().expect("unpoisoned").hanging.insert(hanging);
1730        let held = outbox.reserve(1).expect("an unspent slot is admitted");
1731        let spent = outbox.reserve(2).expect("two reserved slots are admitted");
1732        spent.stage(hanging, 1_u32);
1733        spent.stage(hanging, 2_u32);
1734        drop(spent);
1735        outbox.stage(hanging, 3_u32);
1736
1737        let deadline = Instant::now() + Duration::from_millis(200);
1738        let (first, second) = tokio::join!(outbox.shutdown(deadline), outbox.shutdown(deadline));
1739        assert_eq!(first + second, 3, "the loss is reported by one caller");
1740        assert_eq!(first.min(second), 0, "the other caller reports zero");
1741        assert_eq!(
1742            metrics.counter("outbox_events_lost_total{reason=shutdown}"),
1743            3,
1744            "lost once, at shutdown"
1745        );
1746        assert_eq!(outbox.pending(), 0, "abandoned items are not pending");
1747        assert_eq!(outbox.admission.staged_reserved.load(Ordering::Acquire), 0);
1748        assert_eq!(outbox.admission.unreserved.load(Ordering::Acquire), 0);
1749        assert_eq!(
1750            outbox.admission.reserved.load(Ordering::Acquire),
1751            1,
1752            "only the unspent reservation still holds a slot"
1753        );
1754        drop(held);
1755        assert_eq!(outbox.admission.reserved.load(Ordering::Acquire), 0);
1756    }
1757
1758    /// A sink panic is retried like a failed write: the batch stays pending,
1759    /// commits once ahead of a later item, the writer keeps running, nothing is
1760    /// counted lost, and `pending` returns to zero.
1761    #[tokio::test(flavor = "current_thread")]
1762    async fn a_panicking_write_is_retried_once_in_order_without_loss() {
1763        let metrics = TestMetrics::default();
1764        let _metrics = metrics::set_default_local_recorder(&metrics);
1765        let (outbox, state, started) = outbox_over(MemorySink::default());
1766        let tenant = DataTenantId::new_v7();
1767        state.lock().expect("unpoisoned").panicking.insert(tenant);
1768        outbox.stage(tenant, 1_u32);
1769        started.notified().await;
1770        outbox.stage(tenant, 2_u32);
1771        assert_eq!(outbox.pending(), 2, "the panicked batch stays pending");
1772
1773        assert_eq!(outbox.settle(within(10)).await, 0, "the retry drains");
1774        assert_eq!(
1775            state.lock().expect("unpoisoned").written.get(&tenant),
1776            Some(&vec![1, 2]),
1777            "the panicked batch commits once, ahead of the later item"
1778        );
1779        assert_eq!(metrics.counter("outbox_write_failures_total"), 1);
1780        assert_eq!(metrics.counter("outbox_events_lost_total"), 0);
1781        assert!(metrics.gauge("outbox_pending").abs() < f64::EPSILON);
1782    }
1783
1784    /// A sink error whose `Display` panics is retried like a failed write: the
1785    /// batch stays pending, commits once ahead of a later item, the writer
1786    /// keeps running, the failure is counted once, nothing is counted lost,
1787    /// and `pending` and its gauge return to zero.
1788    #[tokio::test(flavor = "current_thread")]
1789    async fn a_panicking_error_display_is_retried_once_in_order_without_loss() {
1790        let metrics = TestMetrics::default();
1791        let _metrics = metrics::set_default_local_recorder(&metrics);
1792        let (outbox, state, started) = outbox_over(MemorySink::default());
1793        let tenant = DataTenantId::new_v7();
1794        state
1795            .lock()
1796            .expect("unpoisoned")
1797            .display_panicking
1798            .insert(tenant);
1799        outbox.stage(tenant, 1_u32);
1800        started.notified().await;
1801        outbox.stage(tenant, 2_u32);
1802        assert_eq!(outbox.pending(), 2, "the failed batch stays pending");
1803
1804        assert_eq!(outbox.settle(within(10)).await, 0, "the retry drains");
1805        assert_eq!(
1806            state.lock().expect("unpoisoned").written.get(&tenant),
1807            Some(&vec![1, 2]),
1808            "the failed batch commits once, ahead of the later item"
1809        );
1810        assert_eq!(metrics.counter("outbox_write_failures_total"), 1);
1811        assert_eq!(metrics.counter("outbox_events_lost_total"), 0);
1812        assert!(metrics.gauge("outbox_pending").abs() < f64::EPSILON);
1813    }
1814
1815    /// Shutdown fences admission when it begins: items staged before it drain
1816    /// even behind a blocked write, and a later stage is refused and counted
1817    /// lost without reaching the sink.
1818    #[tokio::test(flavor = "current_thread")]
1819    async fn shutdown_refuses_items_staged_after_it_begins() {
1820        let metrics = TestMetrics::default();
1821        let _metrics = metrics::set_default_local_recorder(&metrics);
1822        let sink = MemorySink::default();
1823        let hang = Arc::clone(&sink.hang);
1824        let (outbox, state, started) = outbox_over(sink);
1825        let tenant = DataTenantId::new_v7();
1826        state.lock().expect("unpoisoned").hanging.insert(tenant);
1827        outbox.stage(tenant, 1_u32);
1828        started.notified().await;
1829        outbox.stage(tenant, 2_u32);
1830
1831        let mut shutdown = pin!(outbox.shutdown(within(10)));
1832        let begun = poll_fn(|context| Poll::Ready(shutdown.as_mut().poll(context).is_pending()));
1833        assert!(begun.await, "shutdown waits for the blocked write");
1834        outbox.stage(tenant, 3_u32);
1835        assert_eq!(outbox.pending(), 2, "the late item never entered the queue");
1836        assert_eq!(
1837            metrics.counter("outbox_events_lost_total{reason=shutdown}"),
1838            1
1839        );
1840
1841        state.lock().expect("unpoisoned").hanging.clear();
1842        hang.notify_one();
1843        assert_eq!(shutdown.await, 0, "pre-fence items drain");
1844        assert_eq!(
1845            state.lock().expect("unpoisoned").written.get(&tenant),
1846            Some(&vec![1, 2]),
1847            "only pre-fence items reach the sink"
1848        );
1849        assert_eq!(outbox.pending(), 0);
1850        assert!(metrics.gauge("outbox_pending").abs() < f64::EPSILON);
1851    }
1852
1853    /// Large items shed on bytes before the item limit is reached, small ones
1854    /// are still admitted in the bytes left, written items release their
1855    /// bytes, and reservations are refused on bytes once reserved items hold
1856    /// their byte share, while a spent slot always keeps its item.
1857    ///
1858    /// # Panics
1859    /// Panics when an item is admitted past an unreserved byte limit, shed
1860    /// within it, counted under another reason, or a reservation is admitted
1861    /// or refused against the wrong byte share.
1862    #[tokio::test(flavor = "current_thread")]
1863    async fn large_items_shed_on_bytes_before_items() {
1864        let metrics = TestMetrics::default();
1865        let _metrics = metrics::set_default_local_recorder(&metrics);
1866        let sink = MemorySink::default();
1867        let (hang, state) = (Arc::clone(&sink.hang), Arc::clone(&sink.state));
1868        // Unreserved staging holds 100 bytes; reservations 300.
1869        let outbox = Outbox::new(
1870            sink,
1871            4,
1872            OutboxLimits {
1873                items: 1_024,
1874                bytes: 400,
1875            },
1876        );
1877        let tenant = DataTenantId::new_v7();
1878        state.lock().expect("unpoisoned").hanging.insert(tenant);
1879
1880        outbox.stage(tenant, 60_u32);
1881        outbox.stage(tenant, 60_u32);
1882        assert_eq!(
1883            outbox.pending(),
1884            1,
1885            "the second large item is shed on bytes"
1886        );
1887        outbox.stage(tenant, 40_u32);
1888        assert_eq!(outbox.pending(), 2, "a small item fits the bytes left");
1889        assert_eq!(metrics.counter("outbox_events_lost_total{reason=shed}"), 1);
1890
1891        let reservation = outbox.reserve(2).expect("reserved bytes are free");
1892        reservation.stage(tenant, 300_u32);
1893        let full = outbox
1894            .reserve(1)
1895            .expect_err("reserved items hold the reserved byte share");
1896        assert_eq!((full.admitted, full.limit, full.unit), (300, 300, "bytes"));
1897        reservation.stage(tenant, 50_u32);
1898        assert_eq!(
1899            outbox.pending(),
1900            4,
1901            "a spent slot keeps its item past the share"
1902        );
1903        drop(reservation);
1904
1905        state.lock().expect("unpoisoned").hanging.clear();
1906        hang.notify_one();
1907        assert_eq!(outbox.settle(within(10)).await, 0);
1908        assert_eq!(metrics.counter("outbox_items_written_total"), 4);
1909        assert_eq!(metrics.counter("outbox_bytes_written_total"), 450);
1910        assert_eq!(outbox.admission.unreserved_bytes.load(Ordering::Acquire), 0);
1911        assert_eq!(outbox.admission.reserved_bytes.load(Ordering::Acquire), 0);
1912        outbox.stage(tenant, 100_u32);
1913        outbox
1914            .reserve(1)
1915            .expect("written items release their bytes");
1916        assert_eq!(outbox.settle(within(10)).await, 0);
1917        assert_eq!(
1918            metrics.counter("outbox_events_lost_total"),
1919            1,
1920            "nothing else was lost"
1921        );
1922    }
1923
1924    /// Shed records of one tenant collected by [`ShedRecords`]: the `shed`
1925    /// count of every record, in order.
1926    #[derive(Default)]
1927    struct ShedRecords {
1928        /// The `shed` field of each record.
1929        counts: Mutex<Vec<u64>>,
1930    }
1931
1932    /// Reads the `shed` field of an event.
1933    struct ShedField<'a>(&'a mut Option<u64>);
1934
1935    impl tracing::field::Visit for ShedField<'_> {
1936        /// Keeps the value of the `shed` field.
1937        fn record_u64(&mut self, field: &tracing::field::Field, value: u64) {
1938            if field.name() == "shed" {
1939                *self.0 = Some(value);
1940            }
1941        }
1942
1943        /// Ignores every field not recorded as a `u64`.
1944        fn record_debug(&mut self, _: &tracing::field::Field, _: &dyn fmt::Debug) {}
1945    }
1946
1947    impl tracing::Subscriber for ShedRecords {
1948        /// Sees every event.
1949        fn enabled(&self, _: &tracing::Metadata<'_>) -> bool {
1950            true
1951        }
1952
1953        /// Assigns every span the same id; spans are not inspected.
1954        fn new_span(&self, _: &tracing::span::Attributes<'_>) -> tracing::span::Id {
1955            tracing::span::Id::from_u64(1)
1956        }
1957
1958        /// Ignores span records.
1959        fn record(&self, _: &tracing::span::Id, _: &tracing::span::Record<'_>) {}
1960
1961        /// Ignores span links.
1962        fn record_follows_from(&self, _: &tracing::span::Id, _: &tracing::span::Id) {}
1963
1964        /// Keeps the `shed` count of every event carrying one.
1965        fn event(&self, event: &tracing::Event<'_>) {
1966            let mut shed = None;
1967            event.record(&mut ShedField(&mut shed));
1968            if let Some(shed) = shed {
1969                self.counts.lock().expect("unpoisoned").push(shed);
1970            }
1971        }
1972
1973        /// Ignores span entry.
1974        fn enter(&self, _: &tracing::span::Id) {}
1975
1976        /// Ignores span exit.
1977        fn exit(&self, _: &tracing::span::Id) {}
1978    }
1979
1980    /// A tenant shedding continuously writes one shed record per interval,
1981    /// and shutdown records the rest, so the records' counts sum to the
1982    /// `shed` loss counter.
1983    ///
1984    /// # Panics
1985    /// Panics when more than one record is written within the interval, a
1986    /// shed item is left unrecorded after shutdown, or the counts disagree.
1987    #[tokio::test(flavor = "current_thread")]
1988    async fn shed_records_are_aggregated_per_tenant_and_interval() {
1989        let metrics = TestMetrics::default();
1990        let _metrics = metrics::set_default_local_recorder(&metrics);
1991        let records = Arc::new(ShedRecords::default());
1992        let _records = tracing::subscriber::set_default(Arc::clone(&records));
1993        let sink = MemorySink::default();
1994        let state = Arc::clone(&sink.state);
1995        let outbox = Outbox::new(sink, 4, items(4));
1996        let (tenant, other) = (DataTenantId::new_v7(), DataTenantId::new_v7());
1997        state.lock().expect("unpoisoned").hanging.insert(tenant);
1998        outbox.stage(tenant, 1_u32);
1999        for _ in 0..5 {
2000            outbox.stage(tenant, 2_u32);
2001        }
2002        outbox.stage(other, 3_u32);
2003        assert_eq!(
2004            *records.counts.lock().expect("unpoisoned"),
2005            vec![1, 1],
2006            "one record per tenant within the interval"
2007        );
2008
2009        outbox.shutdown(Instant::now()).await;
2010        let counts = records.counts.lock().expect("unpoisoned").clone();
2011        assert_eq!(counts, vec![1, 1, 4], "shutdown records the rest");
2012        assert_eq!(
2013            counts.iter().sum::<u64>(),
2014            metrics.counter("outbox_events_lost_total{reason=shed}")
2015        );
2016    }
2017}