use std::any::Any;
use std::collections::{HashMap, VecDeque};
use std::fmt::Display;
use std::future::{Future, pending, poll_fn};
use std::panic::{AssertUnwindSafe, catch_unwind};
use std::pin::pin;
use std::sync::atomic::{AtomicUsize, Ordering, fence};
use std::sync::{Arc, Mutex, PoisonError, RwLock};
use std::task::Poll;
use std::time::{Duration, Instant};
use tokio::sync::{Notify, mpsc};
use tokio::task::{Id, JoinError, JoinSet};
use tokio::time::{sleep_until, timeout_at};
use tokio_util::sync::CancellationToken;
use tokio_util::task::TaskTracker;
use uuid::Uuid;
use wyrd_spec::DataTenantId;
pub const INITIAL_BACKOFF: Duration = Duration::from_millis(50);
pub const MAX_BACKOFF: Duration = Duration::from_secs(5);
pub const RESERVED_SHARE: (usize, usize) = (3, 4);
pub const SHED_LOG_INTERVAL: Duration = Duration::from_secs(1);
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct OutboxLimits {
pub items: usize,
pub bytes: usize,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum LossReason {
Shed,
Rejected,
Unreachable,
Shutdown,
}
impl LossReason {
#[must_use]
pub const fn as_str(self) -> &'static str {
match self {
Self::Shed => "shed",
Self::Rejected => "rejected",
Self::Unreachable => "unreachable",
Self::Shutdown => "shutdown",
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, thiserror::Error)]
#[error("outbox {outbox} is full: {admitted} of {limit} reservable {unit} admitted")]
pub struct OutboxFull {
pub outbox: &'static str,
pub admitted: usize,
pub limit: usize,
pub unit: &'static str,
}
pub trait OutboxSink: Send + Sync + 'static {
type Item: Send + 'static;
type Error: Display + Send + 'static;
const NAME: &'static str;
const LANES: usize = 1;
const LINGER: Duration = Duration::ZERO;
const BATCH_ITEMS: usize = usize::MAX;
const BATCH_BYTES: usize = usize::MAX;
#[must_use]
fn lane(item: &Self::Item) -> u64 {
let _ = item;
0
}
#[must_use]
fn weight(item: &Self::Item) -> usize {
let _ = item;
0
}
fn write(
&self,
tenant: DataTenantId,
slice: Uuid,
items: &[Self::Item],
) -> impl Future<Output = Result<(), Self::Error>> + Send;
}
type Queued<S> = (DataTenantId, <S as OutboxSink>::Item, Load);
type QueueSender<S> = mpsc::UnboundedSender<Queued<S>>;
pub struct Outbox<S: OutboxSink> {
queue: RwLock<Option<QueueSender<S>>>,
pending: Arc<AtomicUsize>,
admission: Arc<Admission>,
limits: OutboxLimits,
shed: Mutex<HashMap<DataTenantId, ShedLog>>,
idle: Arc<Notify>,
waiters: Mutex<Waiters>,
abandon: CancellationToken,
writer: TaskTracker,
sink: Arc<S>,
}
impl<S: OutboxSink> std::fmt::Debug for Outbox<S> {
fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
formatter
.debug_struct("Outbox")
.field("name", &S::NAME)
.field("pending", &self.pending.load(Ordering::Relaxed))
.finish_non_exhaustive()
}
}
impl<S: OutboxSink> Outbox<S> {
#[must_use]
pub fn new(sink: S, concurrency: usize, limits: OutboxLimits) -> Arc<Self> {
let (queue, requests) = mpsc::unbounded_channel();
let pending = Arc::new(AtomicUsize::new(0));
let admission = Arc::new(Admission::default());
let idle = Arc::new(Notify::new());
let abandon = CancellationToken::new();
let writer = TaskTracker::new();
let sink = Arc::new(sink);
writer.spawn(
OutboxWriter {
sink: Arc::clone(&sink),
requests,
waiting: HashMap::new(),
since: HashMap::new(),
closed: false,
failed: HashMap::new(),
writing: JoinSet::new(),
in_flight: HashMap::new(),
retry_at: HashMap::new(),
concurrency: concurrency.max(1),
pending: Arc::clone(&pending),
admission: Arc::clone(&admission),
idle: Arc::clone(&idle),
abandon: abandon.clone(),
}
.run(),
);
writer.close();
Arc::new(Self {
queue: RwLock::new(Some(queue)),
pending,
admission,
limits,
shed: Mutex::default(),
idle,
waiters: Mutex::default(),
abandon,
writer,
sink,
})
}
#[must_use]
pub fn sink(&self) -> &S {
&self.sink
}
pub fn stage(&self, tenant: DataTenantId, item: impl Into<S::Item>) {
let item = item.into();
let bytes = S::weight(&item);
let queue = self.queue.read().unwrap_or_else(PoisonError::into_inner);
let Some(sender) = queue.as_ref() else {
drop(queue);
lose_late::<S>(tenant);
return;
};
let admission = &self.admission;
if admit(
&admission.unreserved,
1,
self.limits.items - self.reservable(),
)
.is_err()
{
drop(queue);
self.shed(tenant);
return;
}
let byte_limit = self.limits.bytes - self.reservable_bytes();
if admit(&admission.unreserved_bytes, bytes, byte_limit).is_err() {
admission.unreserved.fetch_sub(1, Ordering::AcqRel);
drop(queue);
self.shed(tenant);
return;
}
self.enqueue(sender, tenant, item, Load::one(false, bytes));
}
fn shed(&self, tenant: DataTenantId) {
count_lost::<S>(1, LossReason::Shed);
let now = Instant::now();
let shed = {
let mut logs = self.shed.lock().unwrap_or_else(PoisonError::into_inner);
let log = logs.entry(tenant).or_default();
log.unlogged += 1;
if log
.logged_at
.is_some_and(|at| now.duration_since(at) < SHED_LOG_INTERVAL)
{
return;
}
log.logged_at = Some(now);
std::mem::take(&mut log.unlogged)
};
self.log_shed(tenant, shed);
}
fn log_shed(&self, tenant: DataTenantId, shed: usize) {
tracing::error!(
outbox = S::NAME,
%tenant,
shed,
max_items = self.limits.items,
max_bytes = self.limits.bytes,
"outbox is full; staged items were shed"
);
}
pub fn reserve(self: &Arc<Self>, slots: usize) -> Result<Reservation<S>, OutboxFull> {
let byte_limit = self.reservable_bytes();
let bytes = self.admission.reserved_bytes.load(Ordering::Acquire);
if bytes >= byte_limit {
return Err(OutboxFull {
outbox: S::NAME,
admitted: bytes,
limit: byte_limit,
unit: "bytes",
});
}
let limit = self.reservable();
admit(&self.admission.reserved, slots, limit).map_err(|admitted| OutboxFull {
outbox: S::NAME,
admitted,
limit,
unit: "items",
})?;
Ok(Reservation {
outbox: Arc::clone(self),
slots: AtomicUsize::new(slots),
})
}
pub async fn reserve_until(
self: &Arc<Self>,
slots: usize,
deadline: Instant,
) -> Result<Reservation<S>, OutboxFull> {
let ticket = {
let mut waiters = self.waiters();
if waiters.queue.is_empty() {
match self.reserve(slots) {
Ok(reservation) => return Ok(reservation),
Err(full) if Instant::now() >= deadline => return Err(full),
Err(full) => waiters.refused = Some(full),
}
}
waiters.join(&self.admission)
};
let _place = Place {
outbox: self,
ticket,
};
loop {
let freed = self.admission.freed.notified();
tokio::pin!(freed);
freed.as_mut().enable();
fence(Ordering::SeqCst);
let refused = {
let mut waiters = self.waiters();
if waiters.queue.front() == Some(&ticket) {
match self.reserve(slots) {
Ok(reservation) => return Ok(reservation),
Err(full) => waiters.refused = Some(full),
}
}
waiters.refused
};
if timeout_at(deadline.into(), freed).await.is_err() {
return Err(refused.unwrap_or(OutboxFull {
outbox: S::NAME,
admitted: self.admission.reserved.load(Ordering::Acquire),
limit: self.reservable(),
unit: "items",
}));
}
}
}
fn waiters(&self) -> std::sync::MutexGuard<'_, Waiters> {
self.waiters.lock().unwrap_or_else(PoisonError::into_inner)
}
const fn reservable(&self) -> usize {
self.limits.items * RESERVED_SHARE.0 / RESERVED_SHARE.1
}
const fn reservable_bytes(&self) -> usize {
self.limits.bytes / RESERVED_SHARE.1 * RESERVED_SHARE.0
}
fn send_reserved(&self, tenant: DataTenantId, item: S::Item) {
let bytes = S::weight(&item);
let queue = self.queue.read().unwrap_or_else(PoisonError::into_inner);
let Some(sender) = queue.as_ref() else {
drop(queue);
self.admission.release(Load::one(true, 0));
lose_late::<S>(tenant);
return;
};
self.admission
.reserved_bytes
.fetch_add(bytes, Ordering::AcqRel);
self.enqueue(sender, tenant, item, Load::one(true, bytes));
}
fn enqueue(&self, sender: &QueueSender<S>, tenant: DataTenantId, item: S::Item, load: Load) {
let gauge = metrics::gauge!("outbox_pending", "outbox" => S::NAME);
self.pending.fetch_add(1, Ordering::AcqRel);
self.admission
.staged_reserved
.fetch_add(load.reserved, Ordering::AcqRel);
gauge.increment(1.0);
if sender.send((tenant, item, load)).is_ok() {
return;
}
self.pending.fetch_sub(1, Ordering::AcqRel);
self.admission
.staged_reserved
.fetch_sub(load.reserved, Ordering::AcqRel);
gauge.decrement(1.0);
self.admission.release(load);
lose_late::<S>(tenant);
}
#[must_use]
pub fn pending(&self) -> usize {
self.pending.load(Ordering::Acquire)
}
pub async fn settle(&self, deadline: Instant) -> usize {
loop {
let idle = self.idle.notified();
tokio::pin!(idle);
idle.as_mut().enable();
if self.pending() == 0 {
return 0;
}
if timeout_at(deadline.into(), idle).await.is_err() {
return self.pending();
}
}
}
pub async fn shutdown(&self, deadline: Instant) -> usize {
drop(
self.queue
.write()
.unwrap_or_else(PoisonError::into_inner)
.take(),
);
let unlogged: Vec<_> = self
.shed
.lock()
.unwrap_or_else(PoisonError::into_inner)
.iter_mut()
.filter(|(_, log)| log.unlogged > 0)
.map(|(tenant, log)| (*tenant, std::mem::take(&mut log.unlogged)))
.collect();
for (tenant, shed) in unlogged {
self.log_shed(tenant, shed);
}
if timeout_at(deadline.into(), self.writer.wait())
.await
.is_err()
{
self.abandon.cancel();
self.writer.wait().await;
}
let lost = self.pending.swap(0, Ordering::AcqRel);
if lost == 0 {
return 0;
}
let admission = &self.admission;
let reserved = admission.staged_reserved.swap(0, Ordering::AcqRel);
let reserved_bytes = admission.reserved_bytes.swap(0, Ordering::AcqRel);
let bytes = admission.unreserved_bytes.swap(0, Ordering::AcqRel) + reserved_bytes;
admission.release(Load {
items: lost,
reserved,
bytes,
reserved_bytes,
});
count_lost::<S>(lost, LossReason::Shutdown);
metrics::gauge!("outbox_pending", "outbox" => S::NAME).decrement(lost as f64);
tracing::error!(
outbox = S::NAME,
lost,
"outbox shutdown deadline passed with items unwritten"
);
self.idle.notify_waiters();
lost
}
}
pub struct Reservation<S: OutboxSink> {
outbox: Arc<Outbox<S>>,
slots: AtomicUsize,
}
impl<S: OutboxSink> Reservation<S> {
pub fn stage(&self, tenant: DataTenantId, item: impl Into<S::Item>) {
let spent = self
.slots
.try_update(Ordering::AcqRel, Ordering::Acquire, |slots| {
slots.checked_sub(1)
});
if spent.is_ok() {
self.outbox.send_reserved(tenant, item.into());
} else {
self.outbox.stage(tenant, item);
}
}
}
impl<S: OutboxSink> Drop for Reservation<S> {
fn drop(&mut self) {
self.outbox
.admission
.reserved
.fetch_sub(*self.slots.get_mut(), Ordering::AcqRel);
if *self.slots.get_mut() > 0 {
self.outbox.admission.wake();
}
}
}
impl<S: OutboxSink> std::fmt::Debug for Reservation<S> {
fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
formatter
.debug_struct("Reservation")
.field("outbox", &S::NAME)
.field("slots", &self.slots.load(Ordering::Relaxed))
.finish()
}
}
#[derive(Debug, Default)]
struct Admission {
unreserved: AtomicUsize,
reserved: AtomicUsize,
staged_reserved: AtomicUsize,
unreserved_bytes: AtomicUsize,
reserved_bytes: AtomicUsize,
freed: Notify,
waiting: AtomicUsize,
}
impl Admission {
fn release(&self, load: Load) {
self.reserved.fetch_sub(load.reserved, Ordering::AcqRel);
self.unreserved
.fetch_sub(load.items - load.reserved, Ordering::AcqRel);
self.reserved_bytes
.fetch_sub(load.reserved_bytes, Ordering::AcqRel);
self.unreserved_bytes
.fetch_sub(load.bytes - load.reserved_bytes, Ordering::AcqRel);
self.wake();
}
fn wake(&self) {
fence(Ordering::SeqCst);
if self.waiting.load(Ordering::SeqCst) > 0 {
self.freed.notify_waiters();
}
}
}
#[derive(Debug, Default)]
struct Waiters {
next: u64,
queue: VecDeque<u64>,
refused: Option<OutboxFull>,
}
impl Waiters {
fn join(&mut self, admission: &Admission) -> u64 {
let ticket = self.next;
self.next += 1;
self.queue.push_back(ticket);
admission.waiting.store(self.queue.len(), Ordering::SeqCst);
ticket
}
}
struct Place<'a, S: OutboxSink> {
outbox: &'a Outbox<S>,
ticket: u64,
}
impl<S: OutboxSink> Drop for Place<'_, S> {
fn drop(&mut self) {
let mut waiters = self.outbox.waiters();
let head = waiters.queue.front() == Some(&self.ticket);
waiters.queue.retain(|ticket| *ticket != self.ticket);
let admission = &self.outbox.admission;
admission
.waiting
.store(waiters.queue.len(), Ordering::SeqCst);
if waiters.queue.is_empty() {
waiters.refused = None;
} else if head {
admission.freed.notify_waiters();
}
}
}
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
struct Load {
items: usize,
reserved: usize,
bytes: usize,
reserved_bytes: usize,
}
impl Load {
const fn one(reserved: bool, bytes: usize) -> Self {
Self {
items: 1,
reserved: reserved as usize,
bytes,
reserved_bytes: if reserved { bytes } else { 0 },
}
}
const fn add(&mut self, other: Self) {
self.items += other.items;
self.reserved += other.reserved;
self.bytes += other.bytes;
self.reserved_bytes += other.reserved_bytes;
}
const fn sub(&mut self, part: Self) {
self.items -= part.items;
self.reserved -= part.reserved;
self.bytes -= part.bytes;
self.reserved_bytes -= part.reserved_bytes;
}
}
#[derive(Debug, Default)]
struct ShedLog {
logged_at: Option<Instant>,
unlogged: usize,
}
fn admit(share: &AtomicUsize, count: usize, limit: usize) -> Result<(), usize> {
share
.try_update(Ordering::AcqRel, Ordering::Acquire, |admitted| {
admitted.checked_add(count).filter(|total| *total <= limit)
})
.map(|_| ())
}
type Backlog<I> = (Vec<I>, Load);
type Waiting<I> = (Vec<(I, Load)>, Load);
type Lane = (DataTenantId, usize);
type Written<S> = (Lane, Uuid, Vec<<S as OutboxSink>::Item>, Result<(), String>);
struct OutboxWriter<S: OutboxSink> {
sink: Arc<S>,
requests: mpsc::UnboundedReceiver<Queued<S>>,
waiting: HashMap<Lane, Waiting<S::Item>>,
since: HashMap<Lane, Instant>,
closed: bool,
failed: HashMap<Lane, (Uuid, Backlog<S::Item>)>,
writing: JoinSet<Written<S>>,
in_flight: HashMap<Id, (Lane, Load, Instant)>,
retry_at: HashMap<Lane, (Instant, Duration)>,
concurrency: usize,
pending: Arc<AtomicUsize>,
admission: Arc<Admission>,
idle: Arc<Notify>,
abandon: CancellationToken,
}
impl<S: OutboxSink> OutboxWriter<S> {
async fn run(mut self) {
let mut received = Vec::new();
let mut open = true;
loop {
let now = Instant::now();
self.dispatch(now);
if !open && self.waiting.is_empty() && self.failed.is_empty() && self.writing.is_empty()
{
return;
}
let retry = self.next_wake(now);
tokio::select! {
biased;
() = self.abandon.cancelled() => return,
Some(done) = self.writing.join_next_with_id(), if !self.writing.is_empty() => {
self.finish(done);
}
count = self.requests.recv_many(&mut received, usize::MAX), if open => {
open = count != 0;
self.closed = !open;
let lanes = S::LANES.max(1) as u64;
let now = Instant::now();
for (tenant, item, load) in received.drain(..) {
let lane = usize::try_from(S::lane(&item) % lanes).unwrap_or_default();
self.since.entry((tenant, lane)).or_insert(now);
let (items, backlog) = self.waiting.entry((tenant, lane)).or_default();
items.push((item, load));
backlog.add(load);
}
}
() = async {
match retry {
Some(at) => sleep_until(at.into()).await,
None => pending().await,
}
} => {}
}
}
}
fn dispatch(&mut self, now: Instant) {
if self.writing.len() >= self.concurrency {
return;
}
let mut busy: HashMap<DataTenantId, usize> = HashMap::new();
for ((tenant, _), ..) in self.in_flight.values() {
*busy.entry(*tenant).or_default() += 1;
}
let mut ready: Vec<(usize, Lane)> = self
.waiting
.keys()
.chain(
self.failed
.keys()
.filter(|lane| !self.waiting.contains_key(*lane)),
)
.filter(|lane| self.is_ready(**lane, now))
.map(|lane| {
let order = busy.entry(lane.0).or_default();
*order += 1;
(*order, *lane)
})
.collect();
ready.sort_unstable_by_key(|(order, _)| *order);
for (_, lane) in ready {
if self.writing.len() >= self.concurrency {
return;
}
let Some((slice, (items, load))) = self
.failed
.remove(&lane)
.or_else(|| self.take(lane).map(|backlog| (Uuid::now_v7(), backlog)))
else {
continue;
};
let (tenant, _) = lane;
let sink = Arc::clone(&self.sink);
let task = self.writing.spawn(async move {
let outcome = {
let write =
catch_unwind(AssertUnwindSafe(|| sink.write(tenant, slice, &items)));
match write {
Ok(write) => contain_panic(write).await,
Err(panic) => Err(panic),
}
};
let result = catch_unwind(AssertUnwindSafe(|| match outcome {
Ok(written) => written.map_err(|error| error.to_string()),
Err(panic) => Err(format!("sink write panicked: {}", panic_message(&*panic))),
}))
.unwrap_or_else(|panic| {
Err(format!(
"sink error rendering panicked: {}",
panic_message(&*panic)
))
});
if result.is_ok() {
drop(tokio::task::spawn_blocking(move || drop(items)));
return (lane, slice, Vec::new(), result);
}
(lane, slice, items, result)
});
self.in_flight
.insert(task.id(), (lane, load, Instant::now()));
}
}
fn take(&mut self, lane: Lane) -> Option<Backlog<S::Item>> {
let (queued, total) = self.waiting.get_mut(&lane)?;
let mut load = Load::default();
let count = queued
.iter()
.position(|(_, next)| {
let over = load.items > 0 && load.bytes.saturating_add(next.bytes) > S::BATCH_BYTES;
if !over {
load.add(*next);
}
over
})
.unwrap_or(queued.len());
let items = if count == queued.len() {
self.since.remove(&lane);
self.waiting.remove(&lane)?.0
} else {
total.sub(load);
queued.drain(..count).collect()
};
Some((items.into_iter().map(|(item, _)| item).collect(), load))
}
fn is_ready(&self, lane: Lane, now: Instant) -> bool {
!self.in_flight.values().any(|(busy, ..)| *busy == lane)
&& self.retry_at.get(&lane).is_none_or(|(at, _)| *at <= now)
&& (self.closed || self.failed.contains_key(&lane) || self.lingered(lane, now))
}
fn lingered(&self, lane: Lane, now: Instant) -> bool {
self.linger_until(lane).is_none_or(|at| at <= now)
|| self.waiting.get(&lane).is_some_and(|(items, load)| {
items.len() >= S::BATCH_ITEMS || load.bytes >= S::BATCH_BYTES
})
}
fn linger_until(&self, lane: Lane) -> Option<Instant> {
self.since.get(&lane).map(|since| *since + S::LINGER)
}
fn next_wake(&self, now: Instant) -> Option<Instant> {
if self.writing.len() >= self.concurrency {
return None;
}
let retries = self
.retry_at
.iter()
.filter(|(lane, (at, _))| *at > now && self.failed.contains_key(*lane))
.map(|(_, (at, _))| *at);
let lingering = self
.since
.keys()
.filter_map(|lane| self.linger_until(*lane))
.filter(|at| *at > now);
retries.chain(lingering).min()
}
fn finish(&mut self, done: Result<(Id, Written<S>), JoinError>) {
let ((lane, slice, items, result), load) = match done {
Ok((id, written)) => {
let load =
self.in_flight
.remove(&id)
.map_or_else(Load::default, |(_, load, started)| {
metrics::histogram!("outbox_write_seconds", "outbox" => S::NAME)
.record(started.elapsed().as_secs_f64());
load
});
(written, load)
}
Err(error) => {
if let Some(((tenant, lane), load, _)) = self.in_flight.remove(&error.id()) {
count_lost::<S>(load.items, LossReason::Rejected);
tracing::error!(outbox = S::NAME, %tenant, lane, %error, lost = load.items, "outbox write task failed");
self.release(load);
}
return;
}
};
let error = match result {
Ok(()) => {
self.retry_at.remove(&lane);
metrics::counter!("outbox_items_written_total", "outbox" => S::NAME)
.increment(load.items as u64);
metrics::counter!("outbox_bytes_written_total", "outbox" => S::NAME)
.increment(load.bytes as u64);
self.release(load);
return;
}
Err(error) => error,
};
let backoff = self
.retry_at
.get(&lane)
.map_or(INITIAL_BACKOFF, |(_, last)| (*last * 2).min(MAX_BACKOFF));
self.retry_at
.insert(lane, (Instant::now() + backoff, backoff));
metrics::counter!("outbox_write_failures_total", "outbox" => S::NAME).increment(1);
tracing::error!(
outbox = S::NAME,
tenant = %lane.0,
lane = lane.1,
%error,
items = items.len(),
retry_in_ms = backoff.as_millis(),
"outbox write failed; retrying"
);
self.failed.insert(lane, (slice, (items, load)));
}
fn release(&self, load: Load) {
metrics::gauge!("outbox_pending", "outbox" => S::NAME).decrement(load.items as f64);
self.admission
.staged_reserved
.fetch_sub(load.reserved, Ordering::AcqRel);
self.admission.release(load);
if self.pending.fetch_sub(load.items, Ordering::AcqRel) == load.items {
self.idle.notify_waiters();
}
}
}
async fn contain_panic<F: Future>(future: F) -> Result<F::Output, Box<dyn Any + Send>> {
let mut future = pin!(future);
poll_fn(
|context| match catch_unwind(AssertUnwindSafe(|| future.as_mut().poll(context))) {
Ok(poll) => poll.map(Ok),
Err(panic) => Poll::Ready(Err(panic)),
},
)
.await
}
fn panic_message(panic: &(dyn Any + Send)) -> &str {
panic
.downcast_ref::<&str>()
.copied()
.or_else(|| panic.downcast_ref::<String>().map(String::as_str))
.unwrap_or("non-string panic payload")
}
pub fn count_lost<S: OutboxSink>(count: usize, reason: LossReason) {
metrics::counter!("outbox_events_lost_total", "outbox" => S::NAME, "reason" => reason.as_str())
.increment(count as u64);
}
fn lose_late<S: OutboxSink>(tenant: DataTenantId) {
count_lost::<S>(1, LossReason::Shutdown);
tracing::error!(outbox = S::NAME, %tenant, "outbox item staged after shutdown was lost");
}
#[cfg(test)]
mod tests {
use std::collections::{HashMap, HashSet};
use std::fmt::{self, Display, Formatter};
use std::future::{Future, poll_fn};
use std::pin::pin;
use std::sync::atomic::Ordering;
use std::sync::{Arc, Mutex, MutexGuard};
use std::task::Poll;
use std::time::{Duration, Instant};
use metrics::atomics::AtomicU64;
use metrics::{
Counter, Gauge, Histogram, Key, KeyName, Metadata, Recorder, SharedString, Unit,
};
use tokio::sync::Notify;
use uuid::Uuid;
use wyrd_spec::DataTenantId;
use super::{Outbox, OutboxLimits, OutboxSink};
#[derive(Default)]
struct MemorySink {
state: Arc<Mutex<MemoryState>>,
hang: Arc<Notify>,
started: Arc<Notify>,
}
#[derive(Default)]
struct MemoryState {
written: HashMap<DataTenantId, Vec<u32>>,
failing: HashSet<DataTenantId>,
hanging: HashSet<DataTenantId>,
panicking: HashSet<DataTenantId>,
display_panicking: HashSet<DataTenantId>,
failures: usize,
attempts: Vec<(DataTenantId, Uuid, Vec<u32>)>,
}
#[derive(Debug)]
struct MemoryError {
panics_on_display: bool,
}
impl Display for MemoryError {
fn fmt(&self, formatter: &mut Formatter<'_>) -> fmt::Result {
assert!(!self.panics_on_display, "injected error display panic");
formatter.write_str("injected failure")
}
}
impl OutboxSink for MemorySink {
type Item = u32;
type Error = MemoryError;
const NAME: &'static str = "test";
fn weight(item: &u32) -> usize {
*item as usize
}
async fn write(
&self,
tenant: DataTenantId,
slice: Uuid,
items: &[u32],
) -> Result<(), MemoryError> {
self.started.notify_one();
let hangs = self.lock().hanging.contains(&tenant);
if hangs {
self.hang.notified().await;
}
let panics = self.lock().panicking.remove(&tenant);
assert!(!panics, "injected sink panic");
let mut state = self.lock();
state.attempts.push((tenant, slice, items.to_vec()));
if state.display_panicking.remove(&tenant) {
state.failures += 1;
return Err(MemoryError {
panics_on_display: true,
});
}
if state.failing.contains(&tenant) {
state.failures += 1;
return Err(MemoryError {
panics_on_display: false,
});
}
state
.written
.entry(tenant)
.or_default()
.extend_from_slice(items);
Ok(())
}
}
impl MemorySink {
fn lock(&self) -> MutexGuard<'_, MemoryState> {
self.state
.lock()
.expect("memory sink state is never poisoned")
}
}
struct SlowDrop(Mutex<std::sync::mpsc::Receiver<()>>);
impl Drop for SlowDrop {
fn drop(&mut self) {
if let Ok(gate) = self.0.lock() {
let _ = gate.recv_timeout(Duration::from_secs(10));
}
}
}
struct SlowDropSink;
impl OutboxSink for SlowDropSink {
type Item = SlowDrop;
type Error = MemoryError;
const NAME: &'static str = "slow_drop";
fn write(
&self,
_tenant: DataTenantId,
_slice: Uuid,
_items: &[SlowDrop],
) -> impl Future<Output = Result<(), MemoryError>> + Send {
std::future::ready(Ok(()))
}
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn freeing_a_written_slice_does_not_hold_up_the_next_write() {
let outbox = Outbox::new(SlowDropSink, 1, ROOMY);
let tenant = DataTenantId::new_v7();
let (gate, held) = std::sync::mpsc::channel();
outbox.stage(tenant, SlowDrop(Mutex::new(held)));
let first = outbox.settle(Instant::now() + Duration::from_secs(2)).await;
let (_, open) = std::sync::mpsc::channel();
outbox.stage(tenant, SlowDrop(Mutex::new(open)));
let second = outbox.settle(Instant::now() + Duration::from_secs(2)).await;
drop(gate);
assert_eq!((first, second), (0, 0), "the next write waited for a drop");
}
#[derive(Default)]
struct LingerSink {
written: Mutex<Vec<Vec<u32>>>,
}
impl OutboxSink for LingerSink {
type Item = u32;
type Error = MemoryError;
const NAME: &'static str = "linger";
const LINGER: Duration = Duration::from_millis(300);
const BATCH_ITEMS: usize = 3;
fn write(
&self,
_tenant: DataTenantId,
_slice: Uuid,
items: &[u32],
) -> impl Future<Output = Result<(), MemoryError>> + Send {
self.written
.lock()
.expect("unpoisoned")
.push(items.to_vec());
std::future::ready(Ok(()))
}
}
#[tokio::test]
async fn a_lingering_lane_batches_trickled_items_and_writes_a_full_batch_at_once() {
let outbox = Outbox::new(LingerSink::default(), 1, ROOMY);
let tenant = DataTenantId::new_v7();
outbox.stage(tenant, 1_u32);
outbox.stage(tenant, 2_u32);
assert_eq!(
outbox.settle(within(5)).await,
0,
"the lingering lane settles"
);
let started = Instant::now();
for item in 3_u32..=5 {
outbox.stage(tenant, item);
}
assert_eq!(outbox.settle(within(5)).await, 0, "the full batch settles");
assert!(
started.elapsed() < LingerSink::LINGER,
"a full batch waited out the linger"
);
let written = outbox.sink().written.lock().expect("unpoisoned").clone();
assert_eq!(written, [vec![1, 2], vec![3, 4, 5]]);
}
#[tokio::test]
async fn a_lane_short_of_its_linger_keeps_a_wake_at_its_deadline() {
let (_queue, requests) = tokio::sync::mpsc::unbounded_channel();
let mut writer = super::OutboxWriter {
sink: Arc::new(LingerSink::default()),
requests,
waiting: HashMap::new(),
since: HashMap::new(),
closed: false,
failed: HashMap::new(),
writing: tokio::task::JoinSet::new(),
in_flight: HashMap::new(),
retry_at: HashMap::new(),
concurrency: 1,
pending: Arc::default(),
admission: Arc::default(),
idle: Arc::default(),
abandon: tokio_util::sync::CancellationToken::new(),
};
let lane = (DataTenantId::new_v7(), 0);
let since = Instant::now()
.checked_sub(LingerSink::LINGER + Duration::from_secs(1))
.expect("the monotonic clock is past one linger");
writer.since.insert(lane, since);
writer
.waiting
.entry(lane)
.or_default()
.0
.push((1_u32, super::Load::default()));
let deadline = since + LingerSink::LINGER;
let short = deadline
.checked_sub(Duration::from_nanos(1))
.expect("the deadline is past one tick");
writer.dispatch(short);
assert!(writer.writing.is_empty(), "a lingering lane was written");
assert_eq!(
writer.next_wake(short),
Some(deadline),
"a lingering lane has no wake at its deadline"
);
writer.dispatch(deadline);
assert_eq!(
writer.writing.len(),
1,
"the lane is written at its deadline"
);
}
#[derive(Default)]
struct LaneSink {
gate: Notify,
written: Mutex<Vec<Vec<u32>>>,
}
impl OutboxSink for LaneSink {
type Item = u32;
type Error = MemoryError;
const NAME: &'static str = "lanes";
const LANES: usize = 2;
fn lane(item: &u32) -> u64 {
u64::from(item % 2)
}
async fn write(
&self,
_tenant: DataTenantId,
_slice: Uuid,
items: &[u32],
) -> Result<(), MemoryError> {
if items == [0] {
self.gate.notified().await;
}
self.written
.lock()
.expect("unpoisoned")
.push(items.to_vec());
Ok(())
}
}
#[tokio::test]
async fn a_tenants_lanes_write_concurrently_and_each_keeps_order() {
let outbox = Outbox::new(LaneSink::default(), 2, ROOMY);
let tenant = DataTenantId::new_v7();
outbox.stage(tenant, 0_u32);
outbox.stage(tenant, 1_u32);
let deadline = within(10);
while outbox.sink().written.lock().expect("unpoisoned").is_empty() {
assert!(Instant::now() < deadline, "lane 1 waited for lane 0");
tokio::time::sleep(Duration::from_millis(5)).await;
}
outbox.stage(tenant, 2_u32);
outbox.stage(tenant, 4_u32);
outbox.sink().gate.notify_one();
assert_eq!(outbox.settle(within(10)).await, 0);
assert_eq!(
*outbox.sink().written.lock().expect("unpoisoned"),
vec![vec![1], vec![0], vec![2, 4]]
);
}
#[derive(Default)]
struct ByteSink {
started: Notify,
gate: Notify,
written: Mutex<Vec<Vec<u32>>>,
}
impl OutboxSink for ByteSink {
type Item = u32;
type Error = MemoryError;
const NAME: &'static str = "bytes";
const BATCH_BYTES: usize = 10;
fn weight(item: &u32) -> usize {
*item as usize
}
async fn write(
&self,
_tenant: DataTenantId,
_slice: Uuid,
items: &[u32],
) -> Result<(), MemoryError> {
if items == [0] {
self.started.notify_one();
self.gate.notified().await;
}
self.written
.lock()
.expect("unpoisoned")
.push(items.to_vec());
Ok(())
}
}
#[tokio::test]
async fn a_heavy_backlog_is_written_in_bounded_slices_in_order() {
let outbox = Outbox::new(ByteSink::default(), 1, ROOMY);
let tenant = DataTenantId::new_v7();
outbox.stage(tenant, 0_u32);
outbox.sink().started.notified().await;
for item in [4_u32, 5, 1, 20, 3] {
outbox.stage(tenant, item);
}
outbox.sink().gate.notify_one();
assert_eq!(outbox.settle(within(10)).await, 0);
assert_eq!(
*outbox.sink().written.lock().expect("unpoisoned"),
vec![vec![0], vec![4, 5, 1], vec![20], vec![3]]
);
}
const ROOMY: OutboxLimits = items(1_024);
const fn items(items: usize) -> OutboxLimits {
OutboxLimits {
items,
bytes: usize::MAX,
}
}
fn outbox(concurrency: usize) -> (Arc<Outbox<MemorySink>>, Arc<Mutex<MemoryState>>) {
let sink = MemorySink::default();
let state = Arc::clone(&sink.state);
(Outbox::new(sink, concurrency, ROOMY), state)
}
fn outbox_over(
sink: MemorySink,
) -> (
Arc<Outbox<MemorySink>>,
Arc<Mutex<MemoryState>>,
Arc<Notify>,
) {
let state = Arc::clone(&sink.state);
let started = Arc::clone(&sink.started);
(Outbox::new(sink, 4, ROOMY), state, started)
}
#[derive(Default)]
struct TestMetrics {
values: Mutex<HashMap<String, Arc<AtomicU64>>>,
}
impl TestMetrics {
fn cell(&self, key: &Key) -> Arc<AtomicU64> {
let name = key
.labels()
.find(|label| label.key() == "reason")
.map_or_else(
|| key.name().to_owned(),
|reason| format!("{}{{reason={}}}", key.name(), reason.value()),
);
let mut values = self.values.lock().expect("test metrics are never poisoned");
Arc::clone(values.entry(name).or_default())
}
fn raw(&self, name: &str) -> u64 {
let labelled = format!("{name}{{");
self.values
.lock()
.expect("test metrics are never poisoned")
.iter()
.filter(|(key, _)| *key == name || key.starts_with(&labelled))
.map(|(_, value)| value.load(Ordering::Acquire))
.sum()
}
fn counter(&self, name: &str) -> u64 {
self.raw(name)
}
fn gauge(&self, name: &str) -> f64 {
f64::from_bits(self.raw(name))
}
}
impl Recorder for TestMetrics {
fn describe_counter(&self, _: KeyName, _: Option<Unit>, _: SharedString) {}
fn describe_gauge(&self, _: KeyName, _: Option<Unit>, _: SharedString) {}
fn describe_histogram(&self, _: KeyName, _: Option<Unit>, _: SharedString) {}
fn register_counter(&self, key: &Key, _: &Metadata<'_>) -> Counter {
Counter::from_arc(self.cell(key))
}
fn register_gauge(&self, key: &Key, _: &Metadata<'_>) -> Gauge {
Gauge::from_arc(self.cell(key))
}
fn register_histogram(&self, _: &Key, _: &Metadata<'_>) -> Histogram {
Histogram::noop()
}
}
fn within(seconds: u64) -> Instant {
Instant::now() + Duration::from_secs(seconds)
}
async fn await_failures(state: &Mutex<MemoryState>, failures: usize) {
let deadline = within(10);
while state.lock().expect("unpoisoned").failures < failures {
assert!(Instant::now() < deadline, "the write never failed");
tokio::time::sleep(Duration::from_millis(5)).await;
}
}
#[tokio::test]
async fn a_failed_write_is_retried_once_in_order_ahead_of_later_items() {
let (outbox, state) = outbox(4);
let tenant = DataTenantId::new_v7();
state.lock().expect("unpoisoned").failing.insert(tenant);
outbox.stage(tenant, 1_u32);
outbox.stage(tenant, 2_u32);
await_failures(&state, 2).await;
outbox.stage(tenant, 3_u32);
assert_eq!(outbox.pending(), 3, "failed items stay pending");
state.lock().expect("unpoisoned").failing.clear();
assert_eq!(outbox.settle(within(10)).await, 0, "the retry drains");
assert_eq!(
state.lock().expect("unpoisoned").written.get(&tenant),
Some(&vec![1, 2, 3]),
"retried items commit once, ahead of later items"
);
assert_eq!(outbox.pending(), 0);
}
#[tokio::test]
async fn a_retry_writes_the_identical_failed_slice_without_later_items() {
let (outbox, state) = outbox(4);
let tenant = DataTenantId::new_v7();
state.lock().expect("unpoisoned").failing.insert(tenant);
outbox.stage(tenant, 1_u32);
outbox.stage(tenant, 2_u32);
await_failures(&state, 1).await;
outbox.stage(tenant, 3_u32);
await_failures(&state, 2).await;
state.lock().expect("unpoisoned").failing.clear();
assert_eq!(outbox.settle(within(10)).await, 0, "the retry drains");
let state = state.lock().expect("unpoisoned");
let (last, retries) = state
.attempts
.split_last()
.expect("the later item was written");
assert_eq!(
(last.0, &last.2),
(tenant, &vec![3]),
"the later item is its own write"
);
assert!(retries.len() >= 3, "two failures and one successful retry");
let slice = retries[0].1;
assert!(
retries
.iter()
.all(|attempt| attempt == &(tenant, slice, vec![1, 2])),
"every retry is the identical failed slice under one key: {retries:?}"
);
assert_ne!(last.1, slice, "the later write has its own key");
assert_eq!(state.written.get(&tenant), Some(&vec![1, 2, 3]));
}
#[tokio::test]
async fn a_failing_tenant_does_not_delay_another_tenant() {
let (outbox, state) = outbox(1);
let failing = DataTenantId::new_v7();
let healthy = DataTenantId::new_v7();
state.lock().expect("unpoisoned").failing.insert(failing);
outbox.stage(failing, 1_u32);
await_failures(&state, 1).await;
outbox.stage(healthy, 2_u32);
let deadline = within(10);
while !state
.lock()
.expect("unpoisoned")
.written
.contains_key(&healthy)
{
assert!(Instant::now() < deadline, "the healthy tenant was delayed");
tokio::time::sleep(Duration::from_millis(5)).await;
}
assert_eq!(outbox.pending(), 1, "only the failing tenant's item waits");
state.lock().expect("unpoisoned").failing.clear();
assert_eq!(outbox.settle(within(10)).await, 0);
}
#[tokio::test]
async fn bounded_staging_covers_retrying_and_inflight_work() {
let sink = MemorySink::default();
let (hang, started) = (Arc::clone(&sink.hang), Arc::clone(&sink.started));
let state = Arc::clone(&sink.state);
let outbox = Outbox::new(sink, 4, items(4));
let (hanging, failing) = (DataTenantId::new_v7(), DataTenantId::new_v7());
{
let mut state = state.lock().expect("unpoisoned");
state.hanging.insert(hanging);
state.failing.insert(failing);
}
let reserved = outbox.reserve(1).expect("room for one reservation");
reserved.stage(hanging, 1_u32);
started.notified().await;
let retried = outbox.reserve(2).expect("room for a two-item reservation");
retried.stage(failing, 2_u32);
retried.stage(failing, 3_u32);
await_failures(&state, 3).await;
assert_eq!(outbox.pending(), 3, "a retried slice is admitted once");
let refused = outbox
.reserve(1)
.expect_err("in-flight and retrying work fill the share");
assert_eq!(
(refused.admitted, refused.limit, refused.unit),
(3, 3, "items")
);
outbox.stage(failing, 4_u32);
outbox.stage(failing, 5_u32);
assert_eq!(
outbox.pending(),
4,
"unreserved staging past capacity is lost"
);
{
let mut state = state.lock().expect("unpoisoned");
state.hanging.clear();
state.failing.clear();
}
hang.notify_one();
assert_eq!(
outbox.settle(within(10)).await,
0,
"the bounded work drains"
);
drop((reserved, retried));
let recovered = outbox
.reserve(3)
.expect("written work releases its capacity");
recovered.stage(hanging, 6_u32);
assert_eq!(outbox.settle(within(10)).await, 0);
let state = state.lock().expect("unpoisoned");
assert_eq!(state.written[&hanging], vec![1, 6]);
assert_eq!(state.written[&failing], vec![2, 3, 4]);
}
#[tokio::test(flavor = "current_thread")]
async fn bounded_audit_staging_remains_nonblocking() {
let metrics = TestMetrics::default();
let _metrics = metrics::set_default_local_recorder(&metrics);
let sink = MemorySink::default();
let state = Arc::clone(&sink.state);
let outbox = Outbox::new(sink, 4, items(4));
let tenant = DataTenantId::new_v7();
state.lock().expect("unpoisoned").hanging.insert(tenant);
let unspent = outbox.reserve(3).expect("the reserved share");
outbox
.reserve(1)
.expect_err("reservations hold their whole share");
outbox.stage(tenant, 1_u32);
assert_eq!(
outbox.pending(),
1,
"the kept share admits unreserved staging"
);
outbox.stage(tenant, 2_u32);
assert_eq!(
outbox.pending(),
1,
"a full unreserved share refuses without waiting"
);
assert_eq!(
metrics.counter("outbox_events_lost_total{reason=shed}"),
1,
"the refusal is counted as shed"
);
drop(unspent);
outbox
.reserve(3)
.expect("a dropped reservation returns its slots despite the hanging item");
}
#[tokio::test(flavor = "current_thread")]
async fn waiting_reservations_are_admitted_in_arrival_order() {
let outbox = Outbox::new(MemorySink::default(), 4, items(4));
let mut held = (0..3)
.map(|_| outbox.reserve(1).expect("the reserved share"))
.collect::<Vec<_>>();
outbox
.reserve_until(1, Instant::now())
.await
.expect_err("a passed deadline refuses at once");
let waiter = |outbox: &Arc<Outbox<MemorySink>>| {
let outbox = Arc::clone(outbox);
tokio::spawn(async move { outbox.reserve_until(1, within(10)).await })
};
let first = waiter(&outbox);
tokio::task::yield_now().await;
let second = waiter(&outbox);
tokio::task::yield_now().await;
drop(held.pop());
let first = first
.await
.expect("joined")
.expect("the first waiter takes the freed slot");
tokio::task::yield_now().await;
assert!(!second.is_finished(), "the second waiter queues behind it");
drop(held.pop());
let _second = second
.await
.expect("joined")
.expect("the next freed slot admits the second waiter");
drop(first);
}
#[tokio::test(flavor = "current_thread")]
async fn a_later_caller_never_takes_capacity_freed_for_the_head() {
let outbox = Outbox::new(MemorySink::default(), 4, items(4));
let mut held = (0..3)
.map(|_| outbox.reserve(1).expect("the reserved share"))
.collect::<Vec<_>>();
let head = {
let outbox = Arc::clone(&outbox);
tokio::spawn(async move { outbox.reserve_until(2, within(10)).await })
};
tokio::task::yield_now().await;
drop(held.pop());
tokio::task::yield_now().await;
outbox
.reserve_until(1, Instant::now() + Duration::from_millis(50))
.await
.expect_err("a later caller does not take the head's freed slot");
assert!(!head.is_finished(), "the head still needs a second slot");
let later = {
let outbox = Arc::clone(&outbox);
tokio::spawn(async move { outbox.reserve_until(1, within(10)).await })
};
tokio::task::yield_now().await;
drop(held.pop());
let head = head
.await
.expect("joined")
.expect("the head is admitted once both slots are free");
tokio::task::yield_now().await;
assert!(!later.is_finished(), "the later caller still queues");
drop(head);
later
.await
.expect("joined")
.expect("the head's departure admits the next caller");
}
#[tokio::test]
async fn shutdown_flushes_items_that_recover_before_the_deadline() {
let (outbox, state) = outbox(4);
let tenant = DataTenantId::new_v7();
state.lock().expect("unpoisoned").failing.insert(tenant);
outbox.stage(tenant, 1_u32);
await_failures(&state, 1).await;
let shutdown = {
let outbox = Arc::clone(&outbox);
tokio::spawn(async move { outbox.shutdown(within(10)).await })
};
await_failures(&state, 2).await;
state.lock().expect("unpoisoned").failing.clear();
assert_eq!(shutdown.await.expect("shutdown joins"), 0, "nothing lost");
assert_eq!(
state.lock().expect("unpoisoned").written.get(&tenant),
Some(&vec![1])
);
outbox.stage(tenant, 2_u32);
assert_eq!(outbox.pending(), 0, "a shut-down outbox accepts nothing");
}
#[tokio::test(flavor = "current_thread")]
async fn shutdown_counts_items_unwritten_at_the_deadline_as_lost() {
let metrics = TestMetrics::default();
let _metrics = metrics::set_default_local_recorder(&metrics);
let (outbox, state) = outbox(4);
let failing = DataTenantId::new_v7();
let hanging = DataTenantId::new_v7();
{
let mut state = state.lock().expect("unpoisoned");
state.failing.insert(failing);
state.hanging.insert(hanging);
}
outbox.stage(failing, 1_u32);
outbox.stage(failing, 2_u32);
outbox.stage(hanging, 3_u32);
await_failures(&state, 1).await;
let lost = outbox
.shutdown(Instant::now() + Duration::from_millis(200))
.await;
assert_eq!(lost, 3, "every unwritten item is reported lost");
assert!(state.lock().expect("unpoisoned").written.is_empty());
assert_eq!(
metrics.counter("outbox_events_lost_total{reason=shutdown}"),
3,
"lost once, at shutdown"
);
assert_eq!(outbox.pending(), 0, "abandoned items are not pending");
assert!(
metrics.gauge("outbox_pending").abs() < f64::EPSILON,
"the pending gauge returns to zero"
);
}
#[tokio::test(flavor = "current_thread")]
async fn concurrent_shutdown_releases_reserved_losses_once() {
let metrics = TestMetrics::default();
let _metrics = metrics::set_default_local_recorder(&metrics);
let (outbox, state) = outbox(4);
let hanging = DataTenantId::new_v7();
state.lock().expect("unpoisoned").hanging.insert(hanging);
let held = outbox.reserve(1).expect("an unspent slot is admitted");
let spent = outbox.reserve(2).expect("two reserved slots are admitted");
spent.stage(hanging, 1_u32);
spent.stage(hanging, 2_u32);
drop(spent);
outbox.stage(hanging, 3_u32);
let deadline = Instant::now() + Duration::from_millis(200);
let (first, second) = tokio::join!(outbox.shutdown(deadline), outbox.shutdown(deadline));
assert_eq!(first + second, 3, "the loss is reported by one caller");
assert_eq!(first.min(second), 0, "the other caller reports zero");
assert_eq!(
metrics.counter("outbox_events_lost_total{reason=shutdown}"),
3,
"lost once, at shutdown"
);
assert_eq!(outbox.pending(), 0, "abandoned items are not pending");
assert_eq!(outbox.admission.staged_reserved.load(Ordering::Acquire), 0);
assert_eq!(outbox.admission.unreserved.load(Ordering::Acquire), 0);
assert_eq!(
outbox.admission.reserved.load(Ordering::Acquire),
1,
"only the unspent reservation still holds a slot"
);
drop(held);
assert_eq!(outbox.admission.reserved.load(Ordering::Acquire), 0);
}
#[tokio::test(flavor = "current_thread")]
async fn a_panicking_write_is_retried_once_in_order_without_loss() {
let metrics = TestMetrics::default();
let _metrics = metrics::set_default_local_recorder(&metrics);
let (outbox, state, started) = outbox_over(MemorySink::default());
let tenant = DataTenantId::new_v7();
state.lock().expect("unpoisoned").panicking.insert(tenant);
outbox.stage(tenant, 1_u32);
started.notified().await;
outbox.stage(tenant, 2_u32);
assert_eq!(outbox.pending(), 2, "the panicked batch stays pending");
assert_eq!(outbox.settle(within(10)).await, 0, "the retry drains");
assert_eq!(
state.lock().expect("unpoisoned").written.get(&tenant),
Some(&vec![1, 2]),
"the panicked batch commits once, ahead of the later item"
);
assert_eq!(metrics.counter("outbox_write_failures_total"), 1);
assert_eq!(metrics.counter("outbox_events_lost_total"), 0);
assert!(metrics.gauge("outbox_pending").abs() < f64::EPSILON);
}
#[tokio::test(flavor = "current_thread")]
async fn a_panicking_error_display_is_retried_once_in_order_without_loss() {
let metrics = TestMetrics::default();
let _metrics = metrics::set_default_local_recorder(&metrics);
let (outbox, state, started) = outbox_over(MemorySink::default());
let tenant = DataTenantId::new_v7();
state
.lock()
.expect("unpoisoned")
.display_panicking
.insert(tenant);
outbox.stage(tenant, 1_u32);
started.notified().await;
outbox.stage(tenant, 2_u32);
assert_eq!(outbox.pending(), 2, "the failed batch stays pending");
assert_eq!(outbox.settle(within(10)).await, 0, "the retry drains");
assert_eq!(
state.lock().expect("unpoisoned").written.get(&tenant),
Some(&vec![1, 2]),
"the failed batch commits once, ahead of the later item"
);
assert_eq!(metrics.counter("outbox_write_failures_total"), 1);
assert_eq!(metrics.counter("outbox_events_lost_total"), 0);
assert!(metrics.gauge("outbox_pending").abs() < f64::EPSILON);
}
#[tokio::test(flavor = "current_thread")]
async fn shutdown_refuses_items_staged_after_it_begins() {
let metrics = TestMetrics::default();
let _metrics = metrics::set_default_local_recorder(&metrics);
let sink = MemorySink::default();
let hang = Arc::clone(&sink.hang);
let (outbox, state, started) = outbox_over(sink);
let tenant = DataTenantId::new_v7();
state.lock().expect("unpoisoned").hanging.insert(tenant);
outbox.stage(tenant, 1_u32);
started.notified().await;
outbox.stage(tenant, 2_u32);
let mut shutdown = pin!(outbox.shutdown(within(10)));
let begun = poll_fn(|context| Poll::Ready(shutdown.as_mut().poll(context).is_pending()));
assert!(begun.await, "shutdown waits for the blocked write");
outbox.stage(tenant, 3_u32);
assert_eq!(outbox.pending(), 2, "the late item never entered the queue");
assert_eq!(
metrics.counter("outbox_events_lost_total{reason=shutdown}"),
1
);
state.lock().expect("unpoisoned").hanging.clear();
hang.notify_one();
assert_eq!(shutdown.await, 0, "pre-fence items drain");
assert_eq!(
state.lock().expect("unpoisoned").written.get(&tenant),
Some(&vec![1, 2]),
"only pre-fence items reach the sink"
);
assert_eq!(outbox.pending(), 0);
assert!(metrics.gauge("outbox_pending").abs() < f64::EPSILON);
}
#[tokio::test(flavor = "current_thread")]
async fn large_items_shed_on_bytes_before_items() {
let metrics = TestMetrics::default();
let _metrics = metrics::set_default_local_recorder(&metrics);
let sink = MemorySink::default();
let (hang, state) = (Arc::clone(&sink.hang), Arc::clone(&sink.state));
let outbox = Outbox::new(
sink,
4,
OutboxLimits {
items: 1_024,
bytes: 400,
},
);
let tenant = DataTenantId::new_v7();
state.lock().expect("unpoisoned").hanging.insert(tenant);
outbox.stage(tenant, 60_u32);
outbox.stage(tenant, 60_u32);
assert_eq!(
outbox.pending(),
1,
"the second large item is shed on bytes"
);
outbox.stage(tenant, 40_u32);
assert_eq!(outbox.pending(), 2, "a small item fits the bytes left");
assert_eq!(metrics.counter("outbox_events_lost_total{reason=shed}"), 1);
let reservation = outbox.reserve(2).expect("reserved bytes are free");
reservation.stage(tenant, 300_u32);
let full = outbox
.reserve(1)
.expect_err("reserved items hold the reserved byte share");
assert_eq!((full.admitted, full.limit, full.unit), (300, 300, "bytes"));
reservation.stage(tenant, 50_u32);
assert_eq!(
outbox.pending(),
4,
"a spent slot keeps its item past the share"
);
drop(reservation);
state.lock().expect("unpoisoned").hanging.clear();
hang.notify_one();
assert_eq!(outbox.settle(within(10)).await, 0);
assert_eq!(metrics.counter("outbox_items_written_total"), 4);
assert_eq!(metrics.counter("outbox_bytes_written_total"), 450);
assert_eq!(outbox.admission.unreserved_bytes.load(Ordering::Acquire), 0);
assert_eq!(outbox.admission.reserved_bytes.load(Ordering::Acquire), 0);
outbox.stage(tenant, 100_u32);
outbox
.reserve(1)
.expect("written items release their bytes");
assert_eq!(outbox.settle(within(10)).await, 0);
assert_eq!(
metrics.counter("outbox_events_lost_total"),
1,
"nothing else was lost"
);
}
#[derive(Default)]
struct ShedRecords {
counts: Mutex<Vec<u64>>,
}
struct ShedField<'a>(&'a mut Option<u64>);
impl tracing::field::Visit for ShedField<'_> {
fn record_u64(&mut self, field: &tracing::field::Field, value: u64) {
if field.name() == "shed" {
*self.0 = Some(value);
}
}
fn record_debug(&mut self, _: &tracing::field::Field, _: &dyn fmt::Debug) {}
}
impl tracing::Subscriber for ShedRecords {
fn enabled(&self, _: &tracing::Metadata<'_>) -> bool {
true
}
fn new_span(&self, _: &tracing::span::Attributes<'_>) -> tracing::span::Id {
tracing::span::Id::from_u64(1)
}
fn record(&self, _: &tracing::span::Id, _: &tracing::span::Record<'_>) {}
fn record_follows_from(&self, _: &tracing::span::Id, _: &tracing::span::Id) {}
fn event(&self, event: &tracing::Event<'_>) {
let mut shed = None;
event.record(&mut ShedField(&mut shed));
if let Some(shed) = shed {
self.counts.lock().expect("unpoisoned").push(shed);
}
}
fn enter(&self, _: &tracing::span::Id) {}
fn exit(&self, _: &tracing::span::Id) {}
}
#[tokio::test(flavor = "current_thread")]
async fn shed_records_are_aggregated_per_tenant_and_interval() {
let metrics = TestMetrics::default();
let _metrics = metrics::set_default_local_recorder(&metrics);
let records = Arc::new(ShedRecords::default());
let _records = tracing::subscriber::set_default(Arc::clone(&records));
let sink = MemorySink::default();
let state = Arc::clone(&sink.state);
let outbox = Outbox::new(sink, 4, items(4));
let (tenant, other) = (DataTenantId::new_v7(), DataTenantId::new_v7());
state.lock().expect("unpoisoned").hanging.insert(tenant);
outbox.stage(tenant, 1_u32);
for _ in 0..5 {
outbox.stage(tenant, 2_u32);
}
outbox.stage(other, 3_u32);
assert_eq!(
*records.counts.lock().expect("unpoisoned"),
vec![1, 1],
"one record per tenant within the interval"
);
outbox.shutdown(Instant::now()).await;
let counts = records.counts.lock().expect("unpoisoned").clone();
assert_eq!(counts, vec![1, 1, 4], "shutdown records the rest");
assert_eq!(
counts.iter().sum::<u64>(),
metrics.counter("outbox_events_lost_total{reason=shed}")
);
}
}