1use std::any::Any;
64use std::collections::{HashMap, VecDeque};
65use std::fmt::Display;
66use std::future::{Future, pending, poll_fn};
67use std::panic::{AssertUnwindSafe, catch_unwind};
68use std::pin::pin;
69use std::sync::atomic::{AtomicUsize, Ordering, fence};
70use std::sync::{Arc, Mutex, PoisonError, RwLock};
71use std::task::Poll;
72use std::time::{Duration, Instant};
73
74use tokio::sync::{Notify, mpsc};
75use tokio::task::{Id, JoinError, JoinSet};
76use tokio::time::{sleep_until, timeout_at};
77use tokio_util::sync::CancellationToken;
78use tokio_util::task::TaskTracker;
79use uuid::Uuid;
80use wyrd_spec::DataTenantId;
81
82pub const INITIAL_BACKOFF: Duration = Duration::from_millis(50);
84
85pub const MAX_BACKOFF: Duration = Duration::from_secs(5);
87
88pub const RESERVED_SHARE: (usize, usize) = (3, 4);
93
94pub const SHED_LOG_INTERVAL: Duration = Duration::from_secs(1);
99
100#[derive(Debug, Clone, Copy, PartialEq, Eq)]
103pub struct OutboxLimits {
104 pub items: usize,
106 pub bytes: usize,
109}
110
111#[derive(Debug, Clone, Copy, PartialEq, Eq)]
114pub enum LossReason {
115 Shed,
117 Rejected,
120 Unreachable,
122 Shutdown,
124}
125
126impl LossReason {
127 #[must_use]
129 pub const fn as_str(self) -> &'static str {
130 match self {
131 Self::Shed => "shed",
132 Self::Rejected => "rejected",
133 Self::Unreachable => "unreachable",
134 Self::Shutdown => "shutdown",
135 }
136 }
137}
138
139#[derive(Debug, Clone, Copy, PartialEq, Eq, thiserror::Error)]
143#[error("outbox {outbox} is full: {admitted} of {limit} reservable {unit} admitted")]
144pub struct OutboxFull {
145 pub outbox: &'static str,
147 pub admitted: usize,
150 pub limit: usize,
152 pub unit: &'static str,
154}
155
156pub trait OutboxSink: Send + Sync + 'static {
163 type Item: Send + 'static;
165 type Error: Display + Send + 'static;
167 const NAME: &'static str;
169 const LANES: usize = 1;
173 const LINGER: Duration = Duration::ZERO;
177 const BATCH_ITEMS: usize = usize::MAX;
180 const BATCH_BYTES: usize = usize::MAX;
187
188 #[must_use]
194 fn lane(item: &Self::Item) -> u64 {
195 let _ = item;
196 0
197 }
198
199 #[must_use]
206 fn weight(item: &Self::Item) -> usize {
207 let _ = item;
208 0
209 }
210
211 fn write(
225 &self,
226 tenant: DataTenantId,
227 slice: Uuid,
228 items: &[Self::Item],
229 ) -> impl Future<Output = Result<(), Self::Error>> + Send;
230}
231
232type Queued<S> = (DataTenantId, <S as OutboxSink>::Item, Load);
234
235type QueueSender<S> = mpsc::UnboundedSender<Queued<S>>;
237
238pub struct Outbox<S: OutboxSink> {
240 queue: RwLock<Option<QueueSender<S>>>,
245 pending: Arc<AtomicUsize>,
247 admission: Arc<Admission>,
249 limits: OutboxLimits,
251 shed: Mutex<HashMap<DataTenantId, ShedLog>>,
256 idle: Arc<Notify>,
259 waiters: Mutex<Waiters>,
262 abandon: CancellationToken,
264 writer: TaskTracker,
266 sink: Arc<S>,
268}
269
270impl<S: OutboxSink> std::fmt::Debug for Outbox<S> {
273 fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
279 formatter
280 .debug_struct("Outbox")
281 .field("name", &S::NAME)
282 .field("pending", &self.pending.load(Ordering::Relaxed))
283 .finish_non_exhaustive()
284 }
285}
286
287impl<S: OutboxSink> Outbox<S> {
288 #[must_use]
300 pub fn new(sink: S, concurrency: usize, limits: OutboxLimits) -> Arc<Self> {
301 let (queue, requests) = mpsc::unbounded_channel();
303 let pending = Arc::new(AtomicUsize::new(0));
304 let admission = Arc::new(Admission::default());
305 let idle = Arc::new(Notify::new());
306 let abandon = CancellationToken::new();
307 let writer = TaskTracker::new();
308 let sink = Arc::new(sink);
309 writer.spawn(
310 OutboxWriter {
311 sink: Arc::clone(&sink),
312 requests,
313 waiting: HashMap::new(),
314 since: HashMap::new(),
315 closed: false,
316 failed: HashMap::new(),
317 writing: JoinSet::new(),
318 in_flight: HashMap::new(),
319 retry_at: HashMap::new(),
320 concurrency: concurrency.max(1),
321 pending: Arc::clone(&pending),
322 admission: Arc::clone(&admission),
323 idle: Arc::clone(&idle),
324 abandon: abandon.clone(),
325 }
326 .run(),
327 );
328 writer.close();
329 Arc::new(Self {
330 queue: RwLock::new(Some(queue)),
331 pending,
332 admission,
333 limits,
334 shed: Mutex::default(),
335 idle,
336 waiters: Mutex::default(),
337 abandon,
338 writer,
339 sink,
340 })
341 }
342
343 #[must_use]
345 pub fn sink(&self) -> &S {
346 &self.sink
347 }
348
349 pub fn stage(&self, tenant: DataTenantId, item: impl Into<S::Item>) {
359 let item = item.into();
360 let bytes = S::weight(&item);
361 let queue = self.queue.read().unwrap_or_else(PoisonError::into_inner);
362 let Some(sender) = queue.as_ref() else {
363 drop(queue);
364 lose_late::<S>(tenant);
365 return;
366 };
367 let admission = &self.admission;
368 if admit(
369 &admission.unreserved,
370 1,
371 self.limits.items - self.reservable(),
372 )
373 .is_err()
374 {
375 drop(queue);
376 self.shed(tenant);
377 return;
378 }
379 let byte_limit = self.limits.bytes - self.reservable_bytes();
380 if admit(&admission.unreserved_bytes, bytes, byte_limit).is_err() {
381 admission.unreserved.fetch_sub(1, Ordering::AcqRel);
382 drop(queue);
383 self.shed(tenant);
384 return;
385 }
386 self.enqueue(sender, tenant, item, Load::one(false, bytes));
387 }
388
389 fn shed(&self, tenant: DataTenantId) {
393 count_lost::<S>(1, LossReason::Shed);
394 let now = Instant::now();
395 let shed = {
396 let mut logs = self.shed.lock().unwrap_or_else(PoisonError::into_inner);
397 let log = logs.entry(tenant).or_default();
398 log.unlogged += 1;
399 if log
400 .logged_at
401 .is_some_and(|at| now.duration_since(at) < SHED_LOG_INTERVAL)
402 {
403 return;
404 }
405 log.logged_at = Some(now);
406 std::mem::take(&mut log.unlogged)
407 };
408 self.log_shed(tenant, shed);
409 }
410
411 fn log_shed(&self, tenant: DataTenantId, shed: usize) {
413 tracing::error!(
414 outbox = S::NAME,
415 %tenant,
416 shed,
417 max_items = self.limits.items,
418 max_bytes = self.limits.bytes,
419 "outbox is full; staged items were shed"
420 );
421 }
422
423 pub fn reserve(self: &Arc<Self>, slots: usize) -> Result<Reservation<S>, OutboxFull> {
436 let byte_limit = self.reservable_bytes();
437 let bytes = self.admission.reserved_bytes.load(Ordering::Acquire);
438 if bytes >= byte_limit {
439 return Err(OutboxFull {
440 outbox: S::NAME,
441 admitted: bytes,
442 limit: byte_limit,
443 unit: "bytes",
444 });
445 }
446 let limit = self.reservable();
447 admit(&self.admission.reserved, slots, limit).map_err(|admitted| OutboxFull {
448 outbox: S::NAME,
449 admitted,
450 limit,
451 unit: "items",
452 })?;
453 Ok(Reservation {
454 outbox: Arc::clone(self),
455 slots: AtomicUsize::new(slots),
456 })
457 }
458
459 pub async fn reserve_until(
476 self: &Arc<Self>,
477 slots: usize,
478 deadline: Instant,
479 ) -> Result<Reservation<S>, OutboxFull> {
480 let ticket = {
481 let mut waiters = self.waiters();
482 if waiters.queue.is_empty() {
483 match self.reserve(slots) {
484 Ok(reservation) => return Ok(reservation),
485 Err(full) if Instant::now() >= deadline => return Err(full),
486 Err(full) => waiters.refused = Some(full),
487 }
488 }
489 waiters.join(&self.admission)
490 };
491 let _place = Place {
492 outbox: self,
493 ticket,
494 };
495 loop {
496 let freed = self.admission.freed.notified();
497 tokio::pin!(freed);
498 freed.as_mut().enable();
499 fence(Ordering::SeqCst);
500 let refused = {
501 let mut waiters = self.waiters();
502 if waiters.queue.front() == Some(&ticket) {
503 match self.reserve(slots) {
504 Ok(reservation) => return Ok(reservation),
505 Err(full) => waiters.refused = Some(full),
506 }
507 }
508 waiters.refused
509 };
510 if timeout_at(deadline.into(), freed).await.is_err() {
511 return Err(refused.unwrap_or(OutboxFull {
512 outbox: S::NAME,
513 admitted: self.admission.reserved.load(Ordering::Acquire),
514 limit: self.reservable(),
515 unit: "items",
516 }));
517 }
518 }
519 }
520
521 fn waiters(&self) -> std::sync::MutexGuard<'_, Waiters> {
523 self.waiters.lock().unwrap_or_else(PoisonError::into_inner)
524 }
525
526 const fn reservable(&self) -> usize {
529 self.limits.items * RESERVED_SHARE.0 / RESERVED_SHARE.1
530 }
531
532 const fn reservable_bytes(&self) -> usize {
535 self.limits.bytes / RESERVED_SHARE.1 * RESERVED_SHARE.0
536 }
537
538 fn send_reserved(&self, tenant: DataTenantId, item: S::Item) {
546 let bytes = S::weight(&item);
547 let queue = self.queue.read().unwrap_or_else(PoisonError::into_inner);
548 let Some(sender) = queue.as_ref() else {
549 drop(queue);
550 self.admission.release(Load::one(true, 0));
551 lose_late::<S>(tenant);
552 return;
553 };
554 self.admission
555 .reserved_bytes
556 .fetch_add(bytes, Ordering::AcqRel);
557 self.enqueue(sender, tenant, item, Load::one(true, bytes));
558 }
559
560 fn enqueue(&self, sender: &QueueSender<S>, tenant: DataTenantId, item: S::Item, load: Load) {
566 let gauge = metrics::gauge!("outbox_pending", "outbox" => S::NAME);
567 self.pending.fetch_add(1, Ordering::AcqRel);
568 self.admission
569 .staged_reserved
570 .fetch_add(load.reserved, Ordering::AcqRel);
571 gauge.increment(1.0);
572 if sender.send((tenant, item, load)).is_ok() {
573 return;
574 }
575 self.pending.fetch_sub(1, Ordering::AcqRel);
576 self.admission
577 .staged_reserved
578 .fetch_sub(load.reserved, Ordering::AcqRel);
579 gauge.decrement(1.0);
580 self.admission.release(load);
581 lose_late::<S>(tenant);
582 }
583
584 #[must_use]
586 pub fn pending(&self) -> usize {
587 self.pending.load(Ordering::Acquire)
588 }
589
590 pub async fn settle(&self, deadline: Instant) -> usize {
596 loop {
597 let idle = self.idle.notified();
598 tokio::pin!(idle);
599 idle.as_mut().enable();
600 if self.pending() == 0 {
601 return 0;
602 }
603 if timeout_at(deadline.into(), idle).await.is_err() {
604 return self.pending();
605 }
606 }
607 }
608
609 pub async fn shutdown(&self, deadline: Instant) -> usize {
626 drop(
627 self.queue
628 .write()
629 .unwrap_or_else(PoisonError::into_inner)
630 .take(),
631 );
632 let unlogged: Vec<_> = self
633 .shed
634 .lock()
635 .unwrap_or_else(PoisonError::into_inner)
636 .iter_mut()
637 .filter(|(_, log)| log.unlogged > 0)
638 .map(|(tenant, log)| (*tenant, std::mem::take(&mut log.unlogged)))
639 .collect();
640 for (tenant, shed) in unlogged {
641 self.log_shed(tenant, shed);
642 }
643 if timeout_at(deadline.into(), self.writer.wait())
644 .await
645 .is_err()
646 {
647 self.abandon.cancel();
648 self.writer.wait().await;
649 }
650 let lost = self.pending.swap(0, Ordering::AcqRel);
656 if lost == 0 {
657 return 0;
658 }
659 let admission = &self.admission;
660 let reserved = admission.staged_reserved.swap(0, Ordering::AcqRel);
661 let reserved_bytes = admission.reserved_bytes.swap(0, Ordering::AcqRel);
662 let bytes = admission.unreserved_bytes.swap(0, Ordering::AcqRel) + reserved_bytes;
663 admission.release(Load {
664 items: lost,
665 reserved,
666 bytes,
667 reserved_bytes,
668 });
669 count_lost::<S>(lost, LossReason::Shutdown);
670 metrics::gauge!("outbox_pending", "outbox" => S::NAME).decrement(lost as f64);
671 tracing::error!(
672 outbox = S::NAME,
673 lost,
674 "outbox shutdown deadline passed with items unwritten"
675 );
676 self.idle.notify_waiters();
677 lost
678 }
679}
680
681pub struct Reservation<S: OutboxSink> {
688 outbox: Arc<Outbox<S>>,
690 slots: AtomicUsize,
692}
693
694impl<S: OutboxSink> Reservation<S> {
695 pub fn stage(&self, tenant: DataTenantId, item: impl Into<S::Item>) {
700 let spent = self
701 .slots
702 .try_update(Ordering::AcqRel, Ordering::Acquire, |slots| {
703 slots.checked_sub(1)
704 });
705 if spent.is_ok() {
706 self.outbox.send_reserved(tenant, item.into());
707 } else {
708 self.outbox.stage(tenant, item);
709 }
710 }
711}
712
713impl<S: OutboxSink> Drop for Reservation<S> {
715 fn drop(&mut self) {
717 self.outbox
718 .admission
719 .reserved
720 .fetch_sub(*self.slots.get_mut(), Ordering::AcqRel);
721 if *self.slots.get_mut() > 0 {
722 self.outbox.admission.wake();
723 }
724 }
725}
726
727impl<S: OutboxSink> std::fmt::Debug for Reservation<S> {
729 fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
734 formatter
735 .debug_struct("Reservation")
736 .field("outbox", &S::NAME)
737 .field("slots", &self.slots.load(Ordering::Relaxed))
738 .finish()
739 }
740}
741
742#[derive(Debug, Default)]
745struct Admission {
746 unreserved: AtomicUsize,
748 reserved: AtomicUsize,
750 staged_reserved: AtomicUsize,
753 unreserved_bytes: AtomicUsize,
755 reserved_bytes: AtomicUsize,
758 freed: Notify,
761 waiting: AtomicUsize,
764}
765
766impl Admission {
767 fn release(&self, load: Load) {
769 self.reserved.fetch_sub(load.reserved, Ordering::AcqRel);
770 self.unreserved
771 .fetch_sub(load.items - load.reserved, Ordering::AcqRel);
772 self.reserved_bytes
773 .fetch_sub(load.reserved_bytes, Ordering::AcqRel);
774 self.unreserved_bytes
775 .fetch_sub(load.bytes - load.reserved_bytes, Ordering::AcqRel);
776 self.wake();
777 }
778
779 fn wake(&self) {
786 fence(Ordering::SeqCst);
787 if self.waiting.load(Ordering::SeqCst) > 0 {
788 self.freed.notify_waiters();
789 }
790 }
791}
792
793#[derive(Debug, Default)]
795struct Waiters {
796 next: u64,
798 queue: VecDeque<u64>,
800 refused: Option<OutboxFull>,
802}
803
804impl Waiters {
805 fn join(&mut self, admission: &Admission) -> u64 {
807 let ticket = self.next;
808 self.next += 1;
809 self.queue.push_back(ticket);
810 admission.waiting.store(self.queue.len(), Ordering::SeqCst);
811 ticket
812 }
813}
814
815struct Place<'a, S: OutboxSink> {
818 outbox: &'a Outbox<S>,
820 ticket: u64,
822}
823
824impl<S: OutboxSink> Drop for Place<'_, S> {
825 fn drop(&mut self) {
828 let mut waiters = self.outbox.waiters();
829 let head = waiters.queue.front() == Some(&self.ticket);
830 waiters.queue.retain(|ticket| *ticket != self.ticket);
831 let admission = &self.outbox.admission;
832 admission
833 .waiting
834 .store(waiters.queue.len(), Ordering::SeqCst);
835 if waiters.queue.is_empty() {
836 waiters.refused = None;
837 } else if head {
838 admission.freed.notify_waiters();
839 }
840 }
841}
842
843#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
845struct Load {
846 items: usize,
848 reserved: usize,
850 bytes: usize,
852 reserved_bytes: usize,
854}
855
856impl Load {
857 const fn one(reserved: bool, bytes: usize) -> Self {
859 Self {
860 items: 1,
861 reserved: reserved as usize,
862 bytes,
863 reserved_bytes: if reserved { bytes } else { 0 },
864 }
865 }
866
867 const fn add(&mut self, other: Self) {
869 self.items += other.items;
870 self.reserved += other.reserved;
871 self.bytes += other.bytes;
872 self.reserved_bytes += other.reserved_bytes;
873 }
874
875 const fn sub(&mut self, part: Self) {
877 self.items -= part.items;
878 self.reserved -= part.reserved;
879 self.bytes -= part.bytes;
880 self.reserved_bytes -= part.reserved_bytes;
881 }
882}
883
884#[derive(Debug, Default)]
887struct ShedLog {
888 logged_at: Option<Instant>,
890 unlogged: usize,
892}
893
894fn admit(share: &AtomicUsize, count: usize, limit: usize) -> Result<(), usize> {
901 share
902 .try_update(Ordering::AcqRel, Ordering::Acquire, |admitted| {
903 admitted.checked_add(count).filter(|total| *total <= limit)
904 })
905 .map(|_| ())
906}
907
908type Backlog<I> = (Vec<I>, Load);
910
911type Waiting<I> = (Vec<(I, Load)>, Load);
914
915type Lane = (DataTenantId, usize);
918
919type Written<S> = (Lane, Uuid, Vec<<S as OutboxSink>::Item>, Result<(), String>);
925
926struct OutboxWriter<S: OutboxSink> {
931 sink: Arc<S>,
933 requests: mpsc::UnboundedReceiver<Queued<S>>,
935 waiting: HashMap<Lane, Waiting<S::Item>>,
937 since: HashMap<Lane, Instant>,
940 closed: bool,
943 failed: HashMap<Lane, (Uuid, Backlog<S::Item>)>,
946 writing: JoinSet<Written<S>>,
948 in_flight: HashMap<Id, (Lane, Load, Instant)>,
951 retry_at: HashMap<Lane, (Instant, Duration)>,
954 concurrency: usize,
956 pending: Arc<AtomicUsize>,
958 admission: Arc<Admission>,
960 idle: Arc<Notify>,
962 abandon: CancellationToken,
964}
965
966impl<S: OutboxSink> OutboxWriter<S> {
967 async fn run(mut self) {
974 let mut received = Vec::new();
975 let mut open = true;
976 loop {
977 let now = Instant::now();
980 self.dispatch(now);
981 if !open && self.waiting.is_empty() && self.failed.is_empty() && self.writing.is_empty()
982 {
983 return;
984 }
985 let retry = self.next_wake(now);
986 tokio::select! {
987 biased;
988 () = self.abandon.cancelled() => return,
989 Some(done) = self.writing.join_next_with_id(), if !self.writing.is_empty() => {
990 self.finish(done);
991 }
992 count = self.requests.recv_many(&mut received, usize::MAX), if open => {
993 open = count != 0;
994 self.closed = !open;
995 let lanes = S::LANES.max(1) as u64;
996 let now = Instant::now();
997 for (tenant, item, load) in received.drain(..) {
998 let lane = usize::try_from(S::lane(&item) % lanes).unwrap_or_default();
1000 self.since.entry((tenant, lane)).or_insert(now);
1001 let (items, backlog) = self.waiting.entry((tenant, lane)).or_default();
1002 items.push((item, load));
1003 backlog.add(load);
1004 }
1005 }
1006 () = async {
1007 match retry {
1008 Some(at) => sleep_until(at.into()).await,
1009 None => pending().await,
1010 }
1011 } => {}
1012 }
1013 }
1014 }
1015
1016 fn dispatch(&mut self, now: Instant) {
1029 if self.writing.len() >= self.concurrency {
1030 return;
1031 }
1032 let mut busy: HashMap<DataTenantId, usize> = HashMap::new();
1033 for ((tenant, _), ..) in self.in_flight.values() {
1034 *busy.entry(*tenant).or_default() += 1;
1035 }
1036 let mut ready: Vec<(usize, Lane)> = self
1037 .waiting
1038 .keys()
1039 .chain(
1040 self.failed
1041 .keys()
1042 .filter(|lane| !self.waiting.contains_key(*lane)),
1043 )
1044 .filter(|lane| self.is_ready(**lane, now))
1045 .map(|lane| {
1046 let order = busy.entry(lane.0).or_default();
1047 *order += 1;
1048 (*order, *lane)
1049 })
1050 .collect();
1051 ready.sort_unstable_by_key(|(order, _)| *order);
1052 for (_, lane) in ready {
1053 if self.writing.len() >= self.concurrency {
1054 return;
1055 }
1056 let Some((slice, (items, load))) = self
1057 .failed
1058 .remove(&lane)
1059 .or_else(|| self.take(lane).map(|backlog| (Uuid::now_v7(), backlog)))
1060 else {
1061 continue;
1062 };
1063 let (tenant, _) = lane;
1064 let sink = Arc::clone(&self.sink);
1065 let task = self.writing.spawn(async move {
1066 let outcome = {
1069 let write =
1070 catch_unwind(AssertUnwindSafe(|| sink.write(tenant, slice, &items)));
1071 match write {
1072 Ok(write) => contain_panic(write).await,
1073 Err(panic) => Err(panic),
1074 }
1075 };
1076 let result = catch_unwind(AssertUnwindSafe(|| match outcome {
1079 Ok(written) => written.map_err(|error| error.to_string()),
1080 Err(panic) => Err(format!("sink write panicked: {}", panic_message(&*panic))),
1081 }))
1082 .unwrap_or_else(|panic| {
1083 Err(format!(
1084 "sink error rendering panicked: {}",
1085 panic_message(&*panic)
1086 ))
1087 });
1088 if result.is_ok() {
1089 drop(tokio::task::spawn_blocking(move || drop(items)));
1093 return (lane, slice, Vec::new(), result);
1094 }
1095 (lane, slice, items, result)
1096 });
1097 self.in_flight
1098 .insert(task.id(), (lane, load, Instant::now()));
1099 }
1100 }
1101
1102 fn take(&mut self, lane: Lane) -> Option<Backlog<S::Item>> {
1107 let (queued, total) = self.waiting.get_mut(&lane)?;
1108 let mut load = Load::default();
1109 let count = queued
1110 .iter()
1111 .position(|(_, next)| {
1112 let over = load.items > 0 && load.bytes.saturating_add(next.bytes) > S::BATCH_BYTES;
1113 if !over {
1114 load.add(*next);
1115 }
1116 over
1117 })
1118 .unwrap_or(queued.len());
1119 let items = if count == queued.len() {
1120 self.since.remove(&lane);
1121 self.waiting.remove(&lane)?.0
1122 } else {
1123 total.sub(load);
1124 queued.drain(..count).collect()
1125 };
1126 Some((items.into_iter().map(|(item, _)| item).collect(), load))
1127 }
1128
1129 fn is_ready(&self, lane: Lane, now: Instant) -> bool {
1135 !self.in_flight.values().any(|(busy, ..)| *busy == lane)
1136 && self.retry_at.get(&lane).is_none_or(|(at, _)| *at <= now)
1137 && (self.closed || self.failed.contains_key(&lane) || self.lingered(lane, now))
1138 }
1139
1140 fn lingered(&self, lane: Lane, now: Instant) -> bool {
1143 self.linger_until(lane).is_none_or(|at| at <= now)
1144 || self.waiting.get(&lane).is_some_and(|(items, load)| {
1145 items.len() >= S::BATCH_ITEMS || load.bytes >= S::BATCH_BYTES
1146 })
1147 }
1148
1149 fn linger_until(&self, lane: Lane) -> Option<Instant> {
1152 self.since.get(&lane).map(|since| *since + S::LINGER)
1153 }
1154
1155 fn next_wake(&self, now: Instant) -> Option<Instant> {
1163 if self.writing.len() >= self.concurrency {
1164 return None;
1165 }
1166 let retries = self
1167 .retry_at
1168 .iter()
1169 .filter(|(lane, (at, _))| *at > now && self.failed.contains_key(*lane))
1170 .map(|(_, (at, _))| *at);
1171 let lingering = self
1172 .since
1173 .keys()
1174 .filter_map(|lane| self.linger_until(*lane))
1175 .filter(|at| *at > now);
1176 retries.chain(lingering).min()
1177 }
1178
1179 fn finish(&mut self, done: Result<(Id, Written<S>), JoinError>) {
1193 let ((lane, slice, items, result), load) = match done {
1194 Ok((id, written)) => {
1195 let load =
1196 self.in_flight
1197 .remove(&id)
1198 .map_or_else(Load::default, |(_, load, started)| {
1199 metrics::histogram!("outbox_write_seconds", "outbox" => S::NAME)
1200 .record(started.elapsed().as_secs_f64());
1201 load
1202 });
1203 (written, load)
1204 }
1205 Err(error) => {
1206 if let Some(((tenant, lane), load, _)) = self.in_flight.remove(&error.id()) {
1207 count_lost::<S>(load.items, LossReason::Rejected);
1208 tracing::error!(outbox = S::NAME, %tenant, lane, %error, lost = load.items, "outbox write task failed");
1209 self.release(load);
1210 }
1211 return;
1212 }
1213 };
1214 let error = match result {
1215 Ok(()) => {
1216 self.retry_at.remove(&lane);
1217 metrics::counter!("outbox_items_written_total", "outbox" => S::NAME)
1218 .increment(load.items as u64);
1219 metrics::counter!("outbox_bytes_written_total", "outbox" => S::NAME)
1220 .increment(load.bytes as u64);
1221 self.release(load);
1222 return;
1223 }
1224 Err(error) => error,
1225 };
1226 let backoff = self
1227 .retry_at
1228 .get(&lane)
1229 .map_or(INITIAL_BACKOFF, |(_, last)| (*last * 2).min(MAX_BACKOFF));
1230 self.retry_at
1231 .insert(lane, (Instant::now() + backoff, backoff));
1232 metrics::counter!("outbox_write_failures_total", "outbox" => S::NAME).increment(1);
1233 tracing::error!(
1234 outbox = S::NAME,
1235 tenant = %lane.0,
1236 lane = lane.1,
1237 %error,
1238 items = items.len(),
1239 retry_in_ms = backoff.as_millis(),
1240 "outbox write failed; retrying"
1241 );
1242 self.failed.insert(lane, (slice, (items, load)));
1243 }
1244
1245 fn release(&self, load: Load) {
1248 metrics::gauge!("outbox_pending", "outbox" => S::NAME).decrement(load.items as f64);
1249 self.admission
1250 .staged_reserved
1251 .fetch_sub(load.reserved, Ordering::AcqRel);
1252 self.admission.release(load);
1253 if self.pending.fetch_sub(load.items, Ordering::AcqRel) == load.items {
1254 self.idle.notify_waiters();
1255 }
1256 }
1257}
1258
1259async fn contain_panic<F: Future>(future: F) -> Result<F::Output, Box<dyn Any + Send>> {
1270 let mut future = pin!(future);
1271 poll_fn(
1272 |context| match catch_unwind(AssertUnwindSafe(|| future.as_mut().poll(context))) {
1273 Ok(poll) => poll.map(Ok),
1274 Err(panic) => Poll::Ready(Err(panic)),
1275 },
1276 )
1277 .await
1278}
1279
1280fn panic_message(panic: &(dyn Any + Send)) -> &str {
1282 panic
1283 .downcast_ref::<&str>()
1284 .copied()
1285 .or_else(|| panic.downcast_ref::<String>().map(String::as_str))
1286 .unwrap_or("non-string panic payload")
1287}
1288
1289pub fn count_lost<S: OutboxSink>(count: usize, reason: LossReason) {
1297 metrics::counter!("outbox_events_lost_total", "outbox" => S::NAME, "reason" => reason.as_str())
1298 .increment(count as u64);
1299}
1300
1301fn lose_late<S: OutboxSink>(tenant: DataTenantId) {
1304 count_lost::<S>(1, LossReason::Shutdown);
1305 tracing::error!(outbox = S::NAME, %tenant, "outbox item staged after shutdown was lost");
1306}
1307
1308#[cfg(test)]
1309mod tests {
1310 use std::collections::{HashMap, HashSet};
1314 use std::fmt::{self, Display, Formatter};
1315 use std::future::{Future, poll_fn};
1316 use std::pin::pin;
1317 use std::sync::atomic::Ordering;
1318 use std::sync::{Arc, Mutex, MutexGuard};
1319 use std::task::Poll;
1320 use std::time::{Duration, Instant};
1321
1322 use metrics::atomics::AtomicU64;
1323 use metrics::{
1324 Counter, Gauge, Histogram, Key, KeyName, Metadata, Recorder, SharedString, Unit,
1325 };
1326 use tokio::sync::Notify;
1327 use uuid::Uuid;
1328 use wyrd_spec::DataTenantId;
1329
1330 use super::{Outbox, OutboxLimits, OutboxSink};
1331
1332 #[derive(Default)]
1334 struct MemorySink {
1335 state: Arc<Mutex<MemoryState>>,
1337 hang: Arc<Notify>,
1339 started: Arc<Notify>,
1342 }
1343
1344 #[derive(Default)]
1346 struct MemoryState {
1347 written: HashMap<DataTenantId, Vec<u32>>,
1349 failing: HashSet<DataTenantId>,
1351 hanging: HashSet<DataTenantId>,
1353 panicking: HashSet<DataTenantId>,
1355 display_panicking: HashSet<DataTenantId>,
1358 failures: usize,
1360 attempts: Vec<(DataTenantId, Uuid, Vec<u32>)>,
1363 }
1364
1365 #[derive(Debug)]
1367 struct MemoryError {
1368 panics_on_display: bool,
1371 }
1372
1373 impl Display for MemoryError {
1374 fn fmt(&self, formatter: &mut Formatter<'_>) -> fmt::Result {
1381 assert!(!self.panics_on_display, "injected error display panic");
1382 formatter.write_str("injected failure")
1383 }
1384 }
1385
1386 impl OutboxSink for MemorySink {
1387 type Item = u32;
1388 type Error = MemoryError;
1389 const NAME: &'static str = "test";
1390
1391 fn weight(item: &u32) -> usize {
1394 *item as usize
1395 }
1396
1397 async fn write(
1423 &self,
1424 tenant: DataTenantId,
1425 slice: Uuid,
1426 items: &[u32],
1427 ) -> Result<(), MemoryError> {
1428 self.started.notify_one();
1429 let hangs = self.lock().hanging.contains(&tenant);
1430 if hangs {
1431 self.hang.notified().await;
1432 }
1433 let panics = self.lock().panicking.remove(&tenant);
1434 assert!(!panics, "injected sink panic");
1435 let mut state = self.lock();
1436 state.attempts.push((tenant, slice, items.to_vec()));
1437 if state.display_panicking.remove(&tenant) {
1438 state.failures += 1;
1439 return Err(MemoryError {
1440 panics_on_display: true,
1441 });
1442 }
1443 if state.failing.contains(&tenant) {
1444 state.failures += 1;
1445 return Err(MemoryError {
1446 panics_on_display: false,
1447 });
1448 }
1449 state
1450 .written
1451 .entry(tenant)
1452 .or_default()
1453 .extend_from_slice(items);
1454 Ok(())
1455 }
1456 }
1457
1458 impl MemorySink {
1459 fn lock(&self) -> MutexGuard<'_, MemoryState> {
1461 self.state
1462 .lock()
1463 .expect("memory sink state is never poisoned")
1464 }
1465 }
1466
1467 struct SlowDrop(Mutex<std::sync::mpsc::Receiver<()>>);
1470
1471 impl Drop for SlowDrop {
1472 fn drop(&mut self) {
1475 if let Ok(gate) = self.0.lock() {
1476 let _ = gate.recv_timeout(Duration::from_secs(10));
1477 }
1478 }
1479 }
1480
1481 struct SlowDropSink;
1483
1484 impl OutboxSink for SlowDropSink {
1485 type Item = SlowDrop;
1486 type Error = MemoryError;
1487 const NAME: &'static str = "slow_drop";
1488
1489 fn write(
1495 &self,
1496 _tenant: DataTenantId,
1497 _slice: Uuid,
1498 _items: &[SlowDrop],
1499 ) -> impl Future<Output = Result<(), MemoryError>> + Send {
1500 std::future::ready(Ok(()))
1501 }
1502 }
1503
1504 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
1513 async fn freeing_a_written_slice_does_not_hold_up_the_next_write() {
1514 let outbox = Outbox::new(SlowDropSink, 1, ROOMY);
1515 let tenant = DataTenantId::new_v7();
1516 let (gate, held) = std::sync::mpsc::channel();
1517 outbox.stage(tenant, SlowDrop(Mutex::new(held)));
1518 let first = outbox.settle(Instant::now() + Duration::from_secs(2)).await;
1519 let (_, open) = std::sync::mpsc::channel();
1520 outbox.stage(tenant, SlowDrop(Mutex::new(open)));
1521 let second = outbox.settle(Instant::now() + Duration::from_secs(2)).await;
1522 drop(gate);
1523 assert_eq!((first, second), (0, 0), "the next write waited for a drop");
1524 }
1525
1526 #[derive(Default)]
1529 struct LingerSink {
1530 written: Mutex<Vec<Vec<u32>>>,
1532 }
1533
1534 impl OutboxSink for LingerSink {
1535 type Item = u32;
1536 type Error = MemoryError;
1537 const NAME: &'static str = "linger";
1538 const LINGER: Duration = Duration::from_millis(300);
1539 const BATCH_ITEMS: usize = 3;
1540
1541 fn write(
1551 &self,
1552 _tenant: DataTenantId,
1553 _slice: Uuid,
1554 items: &[u32],
1555 ) -> impl Future<Output = Result<(), MemoryError>> + Send {
1556 self.written
1557 .lock()
1558 .expect("unpoisoned")
1559 .push(items.to_vec());
1560 std::future::ready(Ok(()))
1561 }
1562 }
1563
1564 #[tokio::test]
1573 async fn a_lingering_lane_batches_trickled_items_and_writes_a_full_batch_at_once() {
1574 let outbox = Outbox::new(LingerSink::default(), 1, ROOMY);
1575 let tenant = DataTenantId::new_v7();
1576 outbox.stage(tenant, 1_u32);
1577 outbox.stage(tenant, 2_u32);
1578 assert_eq!(
1579 outbox.settle(within(5)).await,
1580 0,
1581 "the lingering lane settles"
1582 );
1583 let started = Instant::now();
1584 for item in 3_u32..=5 {
1585 outbox.stage(tenant, item);
1586 }
1587 assert_eq!(outbox.settle(within(5)).await, 0, "the full batch settles");
1588 assert!(
1589 started.elapsed() < LingerSink::LINGER,
1590 "a full batch waited out the linger"
1591 );
1592 let written = outbox.sink().written.lock().expect("unpoisoned").clone();
1593 assert_eq!(written, [vec![1, 2], vec![3, 4, 5]]);
1594 }
1595
1596 #[tokio::test]
1606 async fn a_lane_short_of_its_linger_keeps_a_wake_at_its_deadline() {
1607 let (_queue, requests) = tokio::sync::mpsc::unbounded_channel();
1608 let mut writer = super::OutboxWriter {
1609 sink: Arc::new(LingerSink::default()),
1610 requests,
1611 waiting: HashMap::new(),
1612 since: HashMap::new(),
1613 closed: false,
1614 failed: HashMap::new(),
1615 writing: tokio::task::JoinSet::new(),
1616 in_flight: HashMap::new(),
1617 retry_at: HashMap::new(),
1618 concurrency: 1,
1619 pending: Arc::default(),
1620 admission: Arc::default(),
1621 idle: Arc::default(),
1622 abandon: tokio_util::sync::CancellationToken::new(),
1623 };
1624 let lane = (DataTenantId::new_v7(), 0);
1625 let since = Instant::now()
1628 .checked_sub(LingerSink::LINGER + Duration::from_secs(1))
1629 .expect("the monotonic clock is past one linger");
1630 writer.since.insert(lane, since);
1631 writer
1632 .waiting
1633 .entry(lane)
1634 .or_default()
1635 .0
1636 .push((1_u32, super::Load::default()));
1637 let deadline = since + LingerSink::LINGER;
1638 let short = deadline
1639 .checked_sub(Duration::from_nanos(1))
1640 .expect("the deadline is past one tick");
1641 writer.dispatch(short);
1642 assert!(writer.writing.is_empty(), "a lingering lane was written");
1643 assert_eq!(
1644 writer.next_wake(short),
1645 Some(deadline),
1646 "a lingering lane has no wake at its deadline"
1647 );
1648 writer.dispatch(deadline);
1649 assert_eq!(
1650 writer.writing.len(),
1651 1,
1652 "the lane is written at its deadline"
1653 );
1654 }
1655
1656 #[derive(Default)]
1659 struct LaneSink {
1660 gate: Notify,
1662 written: Mutex<Vec<Vec<u32>>>,
1664 }
1665
1666 impl OutboxSink for LaneSink {
1667 type Item = u32;
1668 type Error = MemoryError;
1669 const NAME: &'static str = "lanes";
1670 const LANES: usize = 2;
1671
1672 fn lane(item: &u32) -> u64 {
1674 u64::from(item % 2)
1675 }
1676
1677 async fn write(
1688 &self,
1689 _tenant: DataTenantId,
1690 _slice: Uuid,
1691 items: &[u32],
1692 ) -> Result<(), MemoryError> {
1693 if items == [0] {
1694 self.gate.notified().await;
1695 }
1696 self.written
1697 .lock()
1698 .expect("unpoisoned")
1699 .push(items.to_vec());
1700 Ok(())
1701 }
1702 }
1703
1704 #[tokio::test]
1712 async fn a_tenants_lanes_write_concurrently_and_each_keeps_order() {
1713 let outbox = Outbox::new(LaneSink::default(), 2, ROOMY);
1714 let tenant = DataTenantId::new_v7();
1715 outbox.stage(tenant, 0_u32);
1716 outbox.stage(tenant, 1_u32);
1717 let deadline = within(10);
1718 while outbox.sink().written.lock().expect("unpoisoned").is_empty() {
1719 assert!(Instant::now() < deadline, "lane 1 waited for lane 0");
1720 tokio::time::sleep(Duration::from_millis(5)).await;
1721 }
1722 outbox.stage(tenant, 2_u32);
1723 outbox.stage(tenant, 4_u32);
1724 outbox.sink().gate.notify_one();
1725
1726 assert_eq!(outbox.settle(within(10)).await, 0);
1727 assert_eq!(
1728 *outbox.sink().written.lock().expect("unpoisoned"),
1729 vec![vec![1], vec![0], vec![2, 4]]
1730 );
1731 }
1732
1733 #[derive(Default)]
1737 struct ByteSink {
1738 started: Notify,
1740 gate: Notify,
1742 written: Mutex<Vec<Vec<u32>>>,
1744 }
1745
1746 impl OutboxSink for ByteSink {
1747 type Item = u32;
1748 type Error = MemoryError;
1749 const NAME: &'static str = "bytes";
1750 const BATCH_BYTES: usize = 10;
1751
1752 fn weight(item: &u32) -> usize {
1754 *item as usize
1755 }
1756
1757 async fn write(
1768 &self,
1769 _tenant: DataTenantId,
1770 _slice: Uuid,
1771 items: &[u32],
1772 ) -> Result<(), MemoryError> {
1773 if items == [0] {
1774 self.started.notify_one();
1775 self.gate.notified().await;
1776 }
1777 self.written
1778 .lock()
1779 .expect("unpoisoned")
1780 .push(items.to_vec());
1781 Ok(())
1782 }
1783 }
1784
1785 #[tokio::test]
1794 async fn a_heavy_backlog_is_written_in_bounded_slices_in_order() {
1795 let outbox = Outbox::new(ByteSink::default(), 1, ROOMY);
1796 let tenant = DataTenantId::new_v7();
1797 outbox.stage(tenant, 0_u32);
1798 outbox.sink().started.notified().await;
1799 for item in [4_u32, 5, 1, 20, 3] {
1800 outbox.stage(tenant, item);
1801 }
1802 outbox.sink().gate.notify_one();
1803
1804 assert_eq!(outbox.settle(within(10)).await, 0);
1805 assert_eq!(
1806 *outbox.sink().written.lock().expect("unpoisoned"),
1807 vec![vec![0], vec![4, 5, 1], vec![20], vec![3]]
1808 );
1809 }
1810
1811 const ROOMY: OutboxLimits = items(1_024);
1813
1814 const fn items(items: usize) -> OutboxLimits {
1816 OutboxLimits {
1817 items,
1818 bytes: usize::MAX,
1819 }
1820 }
1821
1822 fn outbox(concurrency: usize) -> (Arc<Outbox<MemorySink>>, Arc<Mutex<MemoryState>>) {
1824 let sink = MemorySink::default();
1825 let state = Arc::clone(&sink.state);
1826 (Outbox::new(sink, concurrency, ROOMY), state)
1827 }
1828
1829 fn outbox_over(
1832 sink: MemorySink,
1833 ) -> (
1834 Arc<Outbox<MemorySink>>,
1835 Arc<Mutex<MemoryState>>,
1836 Arc<Notify>,
1837 ) {
1838 let state = Arc::clone(&sink.state);
1839 let started = Arc::clone(&sink.started);
1840 (Outbox::new(sink, 4, ROOMY), state, started)
1841 }
1842
1843 #[derive(Default)]
1850 struct TestMetrics {
1851 values: Mutex<HashMap<String, Arc<AtomicU64>>>,
1854 }
1855
1856 impl TestMetrics {
1857 fn cell(&self, key: &Key) -> Arc<AtomicU64> {
1860 let name = key
1861 .labels()
1862 .find(|label| label.key() == "reason")
1863 .map_or_else(
1864 || key.name().to_owned(),
1865 |reason| format!("{}{{reason={}}}", key.name(), reason.value()),
1866 );
1867 let mut values = self.values.lock().expect("test metrics are never poisoned");
1868 Arc::clone(values.entry(name).or_default())
1869 }
1870
1871 fn raw(&self, name: &str) -> u64 {
1875 let labelled = format!("{name}{{");
1876 self.values
1877 .lock()
1878 .expect("test metrics are never poisoned")
1879 .iter()
1880 .filter(|(key, _)| *key == name || key.starts_with(&labelled))
1881 .map(|(_, value)| value.load(Ordering::Acquire))
1882 .sum()
1883 }
1884
1885 fn counter(&self, name: &str) -> u64 {
1887 self.raw(name)
1888 }
1889
1890 fn gauge(&self, name: &str) -> f64 {
1892 f64::from_bits(self.raw(name))
1893 }
1894 }
1895
1896 impl Recorder for TestMetrics {
1897 fn describe_counter(&self, _: KeyName, _: Option<Unit>, _: SharedString) {}
1900
1901 fn describe_gauge(&self, _: KeyName, _: Option<Unit>, _: SharedString) {}
1903
1904 fn describe_histogram(&self, _: KeyName, _: Option<Unit>, _: SharedString) {}
1907
1908 fn register_counter(&self, key: &Key, _: &Metadata<'_>) -> Counter {
1911 Counter::from_arc(self.cell(key))
1912 }
1913
1914 fn register_gauge(&self, key: &Key, _: &Metadata<'_>) -> Gauge {
1917 Gauge::from_arc(self.cell(key))
1918 }
1919
1920 fn register_histogram(&self, _: &Key, _: &Metadata<'_>) -> Histogram {
1923 Histogram::noop()
1924 }
1925 }
1926
1927 fn within(seconds: u64) -> Instant {
1929 Instant::now() + Duration::from_secs(seconds)
1930 }
1931
1932 async fn await_failures(state: &Mutex<MemoryState>, failures: usize) {
1934 let deadline = within(10);
1935 while state.lock().expect("unpoisoned").failures < failures {
1936 assert!(Instant::now() < deadline, "the write never failed");
1937 tokio::time::sleep(Duration::from_millis(5)).await;
1938 }
1939 }
1940
1941 #[tokio::test]
1944 async fn a_failed_write_is_retried_once_in_order_ahead_of_later_items() {
1945 let (outbox, state) = outbox(4);
1946 let tenant = DataTenantId::new_v7();
1947 state.lock().expect("unpoisoned").failing.insert(tenant);
1948 outbox.stage(tenant, 1_u32);
1949 outbox.stage(tenant, 2_u32);
1950 await_failures(&state, 2).await;
1951 outbox.stage(tenant, 3_u32);
1952 assert_eq!(outbox.pending(), 3, "failed items stay pending");
1953 state.lock().expect("unpoisoned").failing.clear();
1954
1955 assert_eq!(outbox.settle(within(10)).await, 0, "the retry drains");
1956 assert_eq!(
1957 state.lock().expect("unpoisoned").written.get(&tenant),
1958 Some(&vec![1, 2, 3]),
1959 "retried items commit once, ahead of later items"
1960 );
1961 assert_eq!(outbox.pending(), 0);
1962 }
1963
1964 #[tokio::test]
1969 async fn a_retry_writes_the_identical_failed_slice_without_later_items() {
1970 let (outbox, state) = outbox(4);
1971 let tenant = DataTenantId::new_v7();
1972 state.lock().expect("unpoisoned").failing.insert(tenant);
1973 outbox.stage(tenant, 1_u32);
1974 outbox.stage(tenant, 2_u32);
1975 await_failures(&state, 1).await;
1976 outbox.stage(tenant, 3_u32);
1977 await_failures(&state, 2).await;
1978 state.lock().expect("unpoisoned").failing.clear();
1979
1980 assert_eq!(outbox.settle(within(10)).await, 0, "the retry drains");
1981 let state = state.lock().expect("unpoisoned");
1982 let (last, retries) = state
1983 .attempts
1984 .split_last()
1985 .expect("the later item was written");
1986 assert_eq!(
1987 (last.0, &last.2),
1988 (tenant, &vec![3]),
1989 "the later item is its own write"
1990 );
1991 assert!(retries.len() >= 3, "two failures and one successful retry");
1992 let slice = retries[0].1;
1993 assert!(
1994 retries
1995 .iter()
1996 .all(|attempt| attempt == &(tenant, slice, vec![1, 2])),
1997 "every retry is the identical failed slice under one key: {retries:?}"
1998 );
1999 assert_ne!(last.1, slice, "the later write has its own key");
2000 assert_eq!(state.written.get(&tenant), Some(&vec![1, 2, 3]));
2001 }
2002
2003 #[tokio::test]
2005 async fn a_failing_tenant_does_not_delay_another_tenant() {
2006 let (outbox, state) = outbox(1);
2007 let failing = DataTenantId::new_v7();
2008 let healthy = DataTenantId::new_v7();
2009 state.lock().expect("unpoisoned").failing.insert(failing);
2010 outbox.stage(failing, 1_u32);
2011 await_failures(&state, 1).await;
2012 outbox.stage(healthy, 2_u32);
2013
2014 let deadline = within(10);
2015 while !state
2016 .lock()
2017 .expect("unpoisoned")
2018 .written
2019 .contains_key(&healthy)
2020 {
2021 assert!(Instant::now() < deadline, "the healthy tenant was delayed");
2022 tokio::time::sleep(Duration::from_millis(5)).await;
2023 }
2024 assert_eq!(outbox.pending(), 1, "only the failing tenant's item waits");
2025 state.lock().expect("unpoisoned").failing.clear();
2026 assert_eq!(outbox.settle(within(10)).await, 0);
2027 }
2028
2029 #[tokio::test]
2039 async fn bounded_staging_covers_retrying_and_inflight_work() {
2040 let sink = MemorySink::default();
2041 let (hang, started) = (Arc::clone(&sink.hang), Arc::clone(&sink.started));
2042 let state = Arc::clone(&sink.state);
2043 let outbox = Outbox::new(sink, 4, items(4));
2045 let (hanging, failing) = (DataTenantId::new_v7(), DataTenantId::new_v7());
2046 {
2047 let mut state = state.lock().expect("unpoisoned");
2048 state.hanging.insert(hanging);
2049 state.failing.insert(failing);
2050 }
2051 let reserved = outbox.reserve(1).expect("room for one reservation");
2052 reserved.stage(hanging, 1_u32);
2053 started.notified().await;
2054 let retried = outbox.reserve(2).expect("room for a two-item reservation");
2055 retried.stage(failing, 2_u32);
2056 retried.stage(failing, 3_u32);
2057 await_failures(&state, 3).await;
2058 assert_eq!(outbox.pending(), 3, "a retried slice is admitted once");
2059
2060 let refused = outbox
2061 .reserve(1)
2062 .expect_err("in-flight and retrying work fill the share");
2063 assert_eq!(
2064 (refused.admitted, refused.limit, refused.unit),
2065 (3, 3, "items")
2066 );
2067 outbox.stage(failing, 4_u32);
2068 outbox.stage(failing, 5_u32);
2069 assert_eq!(
2070 outbox.pending(),
2071 4,
2072 "unreserved staging past capacity is lost"
2073 );
2074
2075 {
2076 let mut state = state.lock().expect("unpoisoned");
2077 state.hanging.clear();
2078 state.failing.clear();
2079 }
2080 hang.notify_one();
2081 assert_eq!(
2082 outbox.settle(within(10)).await,
2083 0,
2084 "the bounded work drains"
2085 );
2086 drop((reserved, retried));
2087 let recovered = outbox
2088 .reserve(3)
2089 .expect("written work releases its capacity");
2090 recovered.stage(hanging, 6_u32);
2091 assert_eq!(outbox.settle(within(10)).await, 0);
2092 let state = state.lock().expect("unpoisoned");
2093 assert_eq!(state.written[&hanging], vec![1, 6]);
2094 assert_eq!(state.written[&failing], vec![2, 3, 4]);
2095 }
2096
2097 #[tokio::test(flavor = "current_thread")]
2107 async fn bounded_audit_staging_remains_nonblocking() {
2108 let metrics = TestMetrics::default();
2109 let _metrics = metrics::set_default_local_recorder(&metrics);
2110 let sink = MemorySink::default();
2111 let state = Arc::clone(&sink.state);
2112 let outbox = Outbox::new(sink, 4, items(4));
2113 let tenant = DataTenantId::new_v7();
2114 state.lock().expect("unpoisoned").hanging.insert(tenant);
2115
2116 let unspent = outbox.reserve(3).expect("the reserved share");
2117 outbox
2118 .reserve(1)
2119 .expect_err("reservations hold their whole share");
2120 outbox.stage(tenant, 1_u32);
2121 assert_eq!(
2122 outbox.pending(),
2123 1,
2124 "the kept share admits unreserved staging"
2125 );
2126 outbox.stage(tenant, 2_u32);
2127 assert_eq!(
2128 outbox.pending(),
2129 1,
2130 "a full unreserved share refuses without waiting"
2131 );
2132 assert_eq!(
2133 metrics.counter("outbox_events_lost_total{reason=shed}"),
2134 1,
2135 "the refusal is counted as shed"
2136 );
2137
2138 drop(unspent);
2139 outbox
2140 .reserve(3)
2141 .expect("a dropped reservation returns its slots despite the hanging item");
2142 }
2143
2144 #[tokio::test(flavor = "current_thread")]
2147 async fn waiting_reservations_are_admitted_in_arrival_order() {
2148 let outbox = Outbox::new(MemorySink::default(), 4, items(4));
2149 let mut held = (0..3)
2150 .map(|_| outbox.reserve(1).expect("the reserved share"))
2151 .collect::<Vec<_>>();
2152 outbox
2153 .reserve_until(1, Instant::now())
2154 .await
2155 .expect_err("a passed deadline refuses at once");
2156
2157 let waiter = |outbox: &Arc<Outbox<MemorySink>>| {
2158 let outbox = Arc::clone(outbox);
2159 tokio::spawn(async move { outbox.reserve_until(1, within(10)).await })
2160 };
2161 let first = waiter(&outbox);
2162 tokio::task::yield_now().await;
2163 let second = waiter(&outbox);
2164 tokio::task::yield_now().await;
2165
2166 drop(held.pop());
2167 let first = first
2168 .await
2169 .expect("joined")
2170 .expect("the first waiter takes the freed slot");
2171 tokio::task::yield_now().await;
2172 assert!(!second.is_finished(), "the second waiter queues behind it");
2173
2174 drop(held.pop());
2175 let _second = second
2176 .await
2177 .expect("joined")
2178 .expect("the next freed slot admits the second waiter");
2179 drop(first);
2180 }
2181
2182 #[tokio::test(flavor = "current_thread")]
2186 async fn a_later_caller_never_takes_capacity_freed_for_the_head() {
2187 let outbox = Outbox::new(MemorySink::default(), 4, items(4));
2188 let mut held = (0..3)
2189 .map(|_| outbox.reserve(1).expect("the reserved share"))
2190 .collect::<Vec<_>>();
2191 let head = {
2192 let outbox = Arc::clone(&outbox);
2193 tokio::spawn(async move { outbox.reserve_until(2, within(10)).await })
2194 };
2195 tokio::task::yield_now().await;
2196
2197 drop(held.pop());
2198 tokio::task::yield_now().await;
2199 outbox
2200 .reserve_until(1, Instant::now() + Duration::from_millis(50))
2201 .await
2202 .expect_err("a later caller does not take the head's freed slot");
2203 assert!(!head.is_finished(), "the head still needs a second slot");
2204
2205 let later = {
2206 let outbox = Arc::clone(&outbox);
2207 tokio::spawn(async move { outbox.reserve_until(1, within(10)).await })
2208 };
2209 tokio::task::yield_now().await;
2210 drop(held.pop());
2211 let head = head
2212 .await
2213 .expect("joined")
2214 .expect("the head is admitted once both slots are free");
2215 tokio::task::yield_now().await;
2216 assert!(!later.is_finished(), "the later caller still queues");
2217 drop(head);
2218 later
2219 .await
2220 .expect("joined")
2221 .expect("the head's departure admits the next caller");
2222 }
2223
2224 #[tokio::test]
2226 async fn shutdown_flushes_items_that_recover_before_the_deadline() {
2227 let (outbox, state) = outbox(4);
2228 let tenant = DataTenantId::new_v7();
2229 state.lock().expect("unpoisoned").failing.insert(tenant);
2230 outbox.stage(tenant, 1_u32);
2231 await_failures(&state, 1).await;
2232 let shutdown = {
2233 let outbox = Arc::clone(&outbox);
2234 tokio::spawn(async move { outbox.shutdown(within(10)).await })
2235 };
2236 await_failures(&state, 2).await;
2237 state.lock().expect("unpoisoned").failing.clear();
2238
2239 assert_eq!(shutdown.await.expect("shutdown joins"), 0, "nothing lost");
2240 assert_eq!(
2241 state.lock().expect("unpoisoned").written.get(&tenant),
2242 Some(&vec![1])
2243 );
2244 outbox.stage(tenant, 2_u32);
2245 assert_eq!(outbox.pending(), 0, "a shut-down outbox accepts nothing");
2246 }
2247
2248 #[tokio::test(flavor = "current_thread")]
2251 async fn shutdown_counts_items_unwritten_at_the_deadline_as_lost() {
2252 let metrics = TestMetrics::default();
2253 let _metrics = metrics::set_default_local_recorder(&metrics);
2254 let (outbox, state) = outbox(4);
2255 let failing = DataTenantId::new_v7();
2256 let hanging = DataTenantId::new_v7();
2257 {
2258 let mut state = state.lock().expect("unpoisoned");
2259 state.failing.insert(failing);
2260 state.hanging.insert(hanging);
2261 }
2262 outbox.stage(failing, 1_u32);
2263 outbox.stage(failing, 2_u32);
2264 outbox.stage(hanging, 3_u32);
2265 await_failures(&state, 1).await;
2266
2267 let lost = outbox
2268 .shutdown(Instant::now() + Duration::from_millis(200))
2269 .await;
2270 assert_eq!(lost, 3, "every unwritten item is reported lost");
2271 assert!(state.lock().expect("unpoisoned").written.is_empty());
2272 assert_eq!(
2273 metrics.counter("outbox_events_lost_total{reason=shutdown}"),
2274 3,
2275 "lost once, at shutdown"
2276 );
2277 assert_eq!(outbox.pending(), 0, "abandoned items are not pending");
2278 assert!(
2279 metrics.gauge("outbox_pending").abs() < f64::EPSILON,
2280 "the pending gauge returns to zero"
2281 );
2282 }
2283
2284 #[tokio::test(flavor = "current_thread")]
2292 async fn concurrent_shutdown_releases_reserved_losses_once() {
2293 let metrics = TestMetrics::default();
2294 let _metrics = metrics::set_default_local_recorder(&metrics);
2295 let (outbox, state) = outbox(4);
2296 let hanging = DataTenantId::new_v7();
2297 state.lock().expect("unpoisoned").hanging.insert(hanging);
2298 let held = outbox.reserve(1).expect("an unspent slot is admitted");
2299 let spent = outbox.reserve(2).expect("two reserved slots are admitted");
2300 spent.stage(hanging, 1_u32);
2301 spent.stage(hanging, 2_u32);
2302 drop(spent);
2303 outbox.stage(hanging, 3_u32);
2304
2305 let deadline = Instant::now() + Duration::from_millis(200);
2306 let (first, second) = tokio::join!(outbox.shutdown(deadline), outbox.shutdown(deadline));
2307 assert_eq!(first + second, 3, "the loss is reported by one caller");
2308 assert_eq!(first.min(second), 0, "the other caller reports zero");
2309 assert_eq!(
2310 metrics.counter("outbox_events_lost_total{reason=shutdown}"),
2311 3,
2312 "lost once, at shutdown"
2313 );
2314 assert_eq!(outbox.pending(), 0, "abandoned items are not pending");
2315 assert_eq!(outbox.admission.staged_reserved.load(Ordering::Acquire), 0);
2316 assert_eq!(outbox.admission.unreserved.load(Ordering::Acquire), 0);
2317 assert_eq!(
2318 outbox.admission.reserved.load(Ordering::Acquire),
2319 1,
2320 "only the unspent reservation still holds a slot"
2321 );
2322 drop(held);
2323 assert_eq!(outbox.admission.reserved.load(Ordering::Acquire), 0);
2324 }
2325
2326 #[tokio::test(flavor = "current_thread")]
2330 async fn a_panicking_write_is_retried_once_in_order_without_loss() {
2331 let metrics = TestMetrics::default();
2332 let _metrics = metrics::set_default_local_recorder(&metrics);
2333 let (outbox, state, started) = outbox_over(MemorySink::default());
2334 let tenant = DataTenantId::new_v7();
2335 state.lock().expect("unpoisoned").panicking.insert(tenant);
2336 outbox.stage(tenant, 1_u32);
2337 started.notified().await;
2338 outbox.stage(tenant, 2_u32);
2339 assert_eq!(outbox.pending(), 2, "the panicked batch stays pending");
2340
2341 assert_eq!(outbox.settle(within(10)).await, 0, "the retry drains");
2342 assert_eq!(
2343 state.lock().expect("unpoisoned").written.get(&tenant),
2344 Some(&vec![1, 2]),
2345 "the panicked batch commits once, ahead of the later item"
2346 );
2347 assert_eq!(metrics.counter("outbox_write_failures_total"), 1);
2348 assert_eq!(metrics.counter("outbox_events_lost_total"), 0);
2349 assert!(metrics.gauge("outbox_pending").abs() < f64::EPSILON);
2350 }
2351
2352 #[tokio::test(flavor = "current_thread")]
2357 async fn a_panicking_error_display_is_retried_once_in_order_without_loss() {
2358 let metrics = TestMetrics::default();
2359 let _metrics = metrics::set_default_local_recorder(&metrics);
2360 let (outbox, state, started) = outbox_over(MemorySink::default());
2361 let tenant = DataTenantId::new_v7();
2362 state
2363 .lock()
2364 .expect("unpoisoned")
2365 .display_panicking
2366 .insert(tenant);
2367 outbox.stage(tenant, 1_u32);
2368 started.notified().await;
2369 outbox.stage(tenant, 2_u32);
2370 assert_eq!(outbox.pending(), 2, "the failed batch stays pending");
2371
2372 assert_eq!(outbox.settle(within(10)).await, 0, "the retry drains");
2373 assert_eq!(
2374 state.lock().expect("unpoisoned").written.get(&tenant),
2375 Some(&vec![1, 2]),
2376 "the failed batch commits once, ahead of the later item"
2377 );
2378 assert_eq!(metrics.counter("outbox_write_failures_total"), 1);
2379 assert_eq!(metrics.counter("outbox_events_lost_total"), 0);
2380 assert!(metrics.gauge("outbox_pending").abs() < f64::EPSILON);
2381 }
2382
2383 #[tokio::test(flavor = "current_thread")]
2387 async fn shutdown_refuses_items_staged_after_it_begins() {
2388 let metrics = TestMetrics::default();
2389 let _metrics = metrics::set_default_local_recorder(&metrics);
2390 let sink = MemorySink::default();
2391 let hang = Arc::clone(&sink.hang);
2392 let (outbox, state, started) = outbox_over(sink);
2393 let tenant = DataTenantId::new_v7();
2394 state.lock().expect("unpoisoned").hanging.insert(tenant);
2395 outbox.stage(tenant, 1_u32);
2396 started.notified().await;
2397 outbox.stage(tenant, 2_u32);
2398
2399 let mut shutdown = pin!(outbox.shutdown(within(10)));
2400 let begun = poll_fn(|context| Poll::Ready(shutdown.as_mut().poll(context).is_pending()));
2401 assert!(begun.await, "shutdown waits for the blocked write");
2402 outbox.stage(tenant, 3_u32);
2403 assert_eq!(outbox.pending(), 2, "the late item never entered the queue");
2404 assert_eq!(
2405 metrics.counter("outbox_events_lost_total{reason=shutdown}"),
2406 1
2407 );
2408
2409 state.lock().expect("unpoisoned").hanging.clear();
2410 hang.notify_one();
2411 assert_eq!(shutdown.await, 0, "pre-fence items drain");
2412 assert_eq!(
2413 state.lock().expect("unpoisoned").written.get(&tenant),
2414 Some(&vec![1, 2]),
2415 "only pre-fence items reach the sink"
2416 );
2417 assert_eq!(outbox.pending(), 0);
2418 assert!(metrics.gauge("outbox_pending").abs() < f64::EPSILON);
2419 }
2420
2421 #[tokio::test(flavor = "current_thread")]
2431 async fn large_items_shed_on_bytes_before_items() {
2432 let metrics = TestMetrics::default();
2433 let _metrics = metrics::set_default_local_recorder(&metrics);
2434 let sink = MemorySink::default();
2435 let (hang, state) = (Arc::clone(&sink.hang), Arc::clone(&sink.state));
2436 let outbox = Outbox::new(
2438 sink,
2439 4,
2440 OutboxLimits {
2441 items: 1_024,
2442 bytes: 400,
2443 },
2444 );
2445 let tenant = DataTenantId::new_v7();
2446 state.lock().expect("unpoisoned").hanging.insert(tenant);
2447
2448 outbox.stage(tenant, 60_u32);
2449 outbox.stage(tenant, 60_u32);
2450 assert_eq!(
2451 outbox.pending(),
2452 1,
2453 "the second large item is shed on bytes"
2454 );
2455 outbox.stage(tenant, 40_u32);
2456 assert_eq!(outbox.pending(), 2, "a small item fits the bytes left");
2457 assert_eq!(metrics.counter("outbox_events_lost_total{reason=shed}"), 1);
2458
2459 let reservation = outbox.reserve(2).expect("reserved bytes are free");
2460 reservation.stage(tenant, 300_u32);
2461 let full = outbox
2462 .reserve(1)
2463 .expect_err("reserved items hold the reserved byte share");
2464 assert_eq!((full.admitted, full.limit, full.unit), (300, 300, "bytes"));
2465 reservation.stage(tenant, 50_u32);
2466 assert_eq!(
2467 outbox.pending(),
2468 4,
2469 "a spent slot keeps its item past the share"
2470 );
2471 drop(reservation);
2472
2473 state.lock().expect("unpoisoned").hanging.clear();
2474 hang.notify_one();
2475 assert_eq!(outbox.settle(within(10)).await, 0);
2476 assert_eq!(metrics.counter("outbox_items_written_total"), 4);
2477 assert_eq!(metrics.counter("outbox_bytes_written_total"), 450);
2478 assert_eq!(outbox.admission.unreserved_bytes.load(Ordering::Acquire), 0);
2479 assert_eq!(outbox.admission.reserved_bytes.load(Ordering::Acquire), 0);
2480 outbox.stage(tenant, 100_u32);
2481 outbox
2482 .reserve(1)
2483 .expect("written items release their bytes");
2484 assert_eq!(outbox.settle(within(10)).await, 0);
2485 assert_eq!(
2486 metrics.counter("outbox_events_lost_total"),
2487 1,
2488 "nothing else was lost"
2489 );
2490 }
2491
2492 #[derive(Default)]
2495 struct ShedRecords {
2496 counts: Mutex<Vec<u64>>,
2498 }
2499
2500 struct ShedField<'a>(&'a mut Option<u64>);
2502
2503 impl tracing::field::Visit for ShedField<'_> {
2504 fn record_u64(&mut self, field: &tracing::field::Field, value: u64) {
2506 if field.name() == "shed" {
2507 *self.0 = Some(value);
2508 }
2509 }
2510
2511 fn record_debug(&mut self, _: &tracing::field::Field, _: &dyn fmt::Debug) {}
2513 }
2514
2515 impl tracing::Subscriber for ShedRecords {
2516 fn enabled(&self, _: &tracing::Metadata<'_>) -> bool {
2518 true
2519 }
2520
2521 fn new_span(&self, _: &tracing::span::Attributes<'_>) -> tracing::span::Id {
2523 tracing::span::Id::from_u64(1)
2524 }
2525
2526 fn record(&self, _: &tracing::span::Id, _: &tracing::span::Record<'_>) {}
2528
2529 fn record_follows_from(&self, _: &tracing::span::Id, _: &tracing::span::Id) {}
2531
2532 fn event(&self, event: &tracing::Event<'_>) {
2534 let mut shed = None;
2535 event.record(&mut ShedField(&mut shed));
2536 if let Some(shed) = shed {
2537 self.counts.lock().expect("unpoisoned").push(shed);
2538 }
2539 }
2540
2541 fn enter(&self, _: &tracing::span::Id) {}
2543
2544 fn exit(&self, _: &tracing::span::Id) {}
2546 }
2547
2548 #[tokio::test(flavor = "current_thread")]
2556 async fn shed_records_are_aggregated_per_tenant_and_interval() {
2557 let metrics = TestMetrics::default();
2558 let _metrics = metrics::set_default_local_recorder(&metrics);
2559 let records = Arc::new(ShedRecords::default());
2560 let _records = tracing::subscriber::set_default(Arc::clone(&records));
2561 let sink = MemorySink::default();
2562 let state = Arc::clone(&sink.state);
2563 let outbox = Outbox::new(sink, 4, items(4));
2564 let (tenant, other) = (DataTenantId::new_v7(), DataTenantId::new_v7());
2565 state.lock().expect("unpoisoned").hanging.insert(tenant);
2566 outbox.stage(tenant, 1_u32);
2567 for _ in 0..5 {
2568 outbox.stage(tenant, 2_u32);
2569 }
2570 outbox.stage(other, 3_u32);
2571 assert_eq!(
2572 *records.counts.lock().expect("unpoisoned"),
2573 vec![1, 1],
2574 "one record per tenant within the interval"
2575 );
2576
2577 outbox.shutdown(Instant::now()).await;
2578 let counts = records.counts.lock().expect("unpoisoned").clone();
2579 assert_eq!(counts, vec![1, 1, 4], "shutdown records the rest");
2580 assert_eq!(
2581 counts.iter().sum::<u64>(),
2582 metrics.counter("outbox_events_lost_total{reason=shed}")
2583 );
2584 }
2585}