1use 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
68pub const INITIAL_BACKOFF: Duration = Duration::from_millis(50);
70
71pub const MAX_BACKOFF: Duration = Duration::from_secs(5);
73
74pub const RESERVED_SHARE: (usize, usize) = (3, 4);
79
80pub const SHED_LOG_INTERVAL: Duration = Duration::from_secs(1);
85
86#[derive(Debug, Clone, Copy, PartialEq, Eq)]
89pub struct OutboxLimits {
90 pub items: usize,
92 pub bytes: usize,
95}
96
97#[derive(Debug, Clone, Copy, PartialEq, Eq)]
100pub enum LossReason {
101 Shed,
103 Rejected,
106 Unreachable,
108 Shutdown,
110}
111
112impl LossReason {
113 #[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#[derive(Debug, Clone, Copy, PartialEq, Eq, thiserror::Error)]
129#[error("outbox {outbox} is full: {admitted} of {limit} reservable {unit} admitted")]
130pub struct OutboxFull {
131 pub outbox: &'static str,
133 pub admitted: usize,
136 pub limit: usize,
138 pub unit: &'static str,
140}
141
142pub trait OutboxSink: Send + Sync + 'static {
149 type Item: Send + 'static;
151 type Error: Display + Send + 'static;
153 const NAME: &'static str;
155 const LANES: usize = 1;
159
160 #[must_use]
166 fn lane(item: &Self::Item) -> u64 {
167 let _ = item;
168 0
169 }
170
171 #[must_use]
178 fn weight(item: &Self::Item) -> usize {
179 let _ = item;
180 0
181 }
182
183 fn write(
194 &self,
195 tenant: DataTenantId,
196 items: &[Self::Item],
197 ) -> impl Future<Output = Result<(), Self::Error>> + Send;
198}
199
200type Queued<S> = (DataTenantId, <S as OutboxSink>::Item, Load);
202
203type QueueSender<S> = mpsc::UnboundedSender<Queued<S>>;
205
206pub struct Outbox<S: OutboxSink> {
208 queue: RwLock<Option<QueueSender<S>>>,
213 pending: Arc<AtomicUsize>,
215 admission: Arc<Admission>,
217 limits: OutboxLimits,
219 shed: Mutex<HashMap<DataTenantId, ShedLog>>,
224 idle: Arc<Notify>,
227 abandon: CancellationToken,
229 writer: TaskTracker,
231 sink: Arc<S>,
233}
234
235impl<S: OutboxSink> std::fmt::Debug for Outbox<S> {
238 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 #[must_use]
265 pub fn new(sink: S, concurrency: usize, limits: OutboxLimits) -> Arc<Self> {
266 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 #[must_use]
307 pub fn sink(&self) -> &S {
308 &self.sink
309 }
310
311 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 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 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 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 const fn reservable(&self) -> usize {
424 self.limits.items * RESERVED_SHARE.0 / RESERVED_SHARE.1
425 }
426
427 const fn reservable_bytes(&self) -> usize {
430 self.limits.bytes / RESERVED_SHARE.1 * RESERVED_SHARE.0
431 }
432
433 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 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 #[must_use]
481 pub fn pending(&self) -> usize {
482 self.pending.load(Ordering::Acquire)
483 }
484
485 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 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 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
576pub struct Reservation<S: OutboxSink> {
583 outbox: Arc<Outbox<S>>,
585 slots: AtomicUsize,
587}
588
589impl<S: OutboxSink> Reservation<S> {
590 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
608impl<S: OutboxSink> Drop for Reservation<S> {
610 fn drop(&mut self) {
612 self.outbox
613 .admission
614 .reserved
615 .fetch_sub(*self.slots.get_mut(), Ordering::AcqRel);
616 }
617}
618
619impl<S: OutboxSink> std::fmt::Debug for Reservation<S> {
621 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#[derive(Debug, Default)]
637struct Admission {
638 unreserved: AtomicUsize,
640 reserved: AtomicUsize,
642 staged_reserved: AtomicUsize,
645 unreserved_bytes: AtomicUsize,
647 reserved_bytes: AtomicUsize,
650}
651
652impl Admission {
653 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#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
667struct Load {
668 items: usize,
670 reserved: usize,
672 bytes: usize,
674 reserved_bytes: usize,
676}
677
678impl Load {
679 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 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#[derive(Debug, Default)]
701struct ShedLog {
702 logged_at: Option<Instant>,
704 unlogged: usize,
706}
707
708fn 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
722type Backlog<I> = (Vec<I>, Load);
724
725type Lane = (DataTenantId, usize);
728
729type Written<S> = (Lane, Vec<<S as OutboxSink>::Item>, Result<(), String>);
734
735struct OutboxWriter<S: OutboxSink> {
740 sink: Arc<S>,
742 requests: mpsc::UnboundedReceiver<Queued<S>>,
744 waiting: HashMap<Lane, Backlog<S::Item>>,
746 failed: HashMap<Lane, Backlog<S::Item>>,
749 writing: JoinSet<Written<S>>,
751 in_flight: HashMap<Id, (Lane, Load, Instant)>,
754 retry_at: HashMap<Lane, (Instant, Duration)>,
757 concurrency: usize,
759 pending: Arc<AtomicUsize>,
761 admission: Arc<Admission>,
763 idle: Arc<Notify>,
765 abandon: CancellationToken,
767}
768
769impl<S: OutboxSink> OutboxWriter<S> {
770 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 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 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 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 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 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 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 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 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 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
998async 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
1019fn 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
1028pub 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
1040fn 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 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 #[derive(Default)]
1072 struct MemorySink {
1073 state: Arc<Mutex<MemoryState>>,
1075 hang: Arc<Notify>,
1077 started: Arc<Notify>,
1080 }
1081
1082 #[derive(Default)]
1084 struct MemoryState {
1085 written: HashMap<DataTenantId, Vec<u32>>,
1087 failing: HashSet<DataTenantId>,
1089 hanging: HashSet<DataTenantId>,
1091 panicking: HashSet<DataTenantId>,
1093 display_panicking: HashSet<DataTenantId>,
1096 failures: usize,
1098 attempts: Vec<(DataTenantId, Vec<u32>)>,
1100 }
1101
1102 #[derive(Debug)]
1104 struct MemoryError {
1105 panics_on_display: bool,
1108 }
1109
1110 impl Display for MemoryError {
1111 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 fn weight(item: &u32) -> usize {
1131 *item as usize
1132 }
1133
1134 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 fn lock(&self) -> MutexGuard<'_, MemoryState> {
1193 self.state
1194 .lock()
1195 .expect("memory sink state is never poisoned")
1196 }
1197 }
1198
1199 struct SlowDrop(Mutex<std::sync::mpsc::Receiver<()>>);
1202
1203 impl Drop for SlowDrop {
1204 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 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 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 #[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 #[derive(Default)]
1260 struct LaneSink {
1261 gate: Notify,
1263 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 fn lane(item: &u32) -> u64 {
1275 u64::from(item % 2)
1276 }
1277
1278 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 #[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 const ROOMY: OutboxLimits = items(1_024);
1331
1332 const fn items(items: usize) -> OutboxLimits {
1334 OutboxLimits {
1335 items,
1336 bytes: usize::MAX,
1337 }
1338 }
1339
1340 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 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 #[derive(Default)]
1368 struct TestMetrics {
1369 values: Mutex<HashMap<String, Arc<AtomicU64>>>,
1372 }
1373
1374 impl TestMetrics {
1375 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 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 fn counter(&self, name: &str) -> u64 {
1405 self.raw(name)
1406 }
1407
1408 fn gauge(&self, name: &str) -> f64 {
1410 f64::from_bits(self.raw(name))
1411 }
1412 }
1413
1414 impl Recorder for TestMetrics {
1415 fn describe_counter(&self, _: KeyName, _: Option<Unit>, _: SharedString) {}
1418
1419 fn describe_gauge(&self, _: KeyName, _: Option<Unit>, _: SharedString) {}
1421
1422 fn describe_histogram(&self, _: KeyName, _: Option<Unit>, _: SharedString) {}
1425
1426 fn register_counter(&self, key: &Key, _: &Metadata<'_>) -> Counter {
1429 Counter::from_arc(self.cell(key))
1430 }
1431
1432 fn register_gauge(&self, key: &Key, _: &Metadata<'_>) -> Gauge {
1435 Gauge::from_arc(self.cell(key))
1436 }
1437
1438 fn register_histogram(&self, _: &Key, _: &Metadata<'_>) -> Histogram {
1441 Histogram::noop()
1442 }
1443 }
1444
1445 fn within(seconds: u64) -> Instant {
1447 Instant::now() + Duration::from_secs(seconds)
1448 }
1449
1450 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 #[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 #[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 #[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 #[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 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 #[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 #[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 #[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 #[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 #[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 #[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 #[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 #[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 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 #[derive(Default)]
1927 struct ShedRecords {
1928 counts: Mutex<Vec<u64>>,
1930 }
1931
1932 struct ShedField<'a>(&'a mut Option<u64>);
1934
1935 impl tracing::field::Visit for ShedField<'_> {
1936 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 fn record_debug(&mut self, _: &tracing::field::Field, _: &dyn fmt::Debug) {}
1945 }
1946
1947 impl tracing::Subscriber for ShedRecords {
1948 fn enabled(&self, _: &tracing::Metadata<'_>) -> bool {
1950 true
1951 }
1952
1953 fn new_span(&self, _: &tracing::span::Attributes<'_>) -> tracing::span::Id {
1955 tracing::span::Id::from_u64(1)
1956 }
1957
1958 fn record(&self, _: &tracing::span::Id, _: &tracing::span::Record<'_>) {}
1960
1961 fn record_follows_from(&self, _: &tracing::span::Id, _: &tracing::span::Id) {}
1963
1964 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 fn enter(&self, _: &tracing::span::Id) {}
1975
1976 fn exit(&self, _: &tracing::span::Id) {}
1978 }
1979
1980 #[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}