use std::any::Any;
use std::collections::HashMap;
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};
use std::sync::{Arc, 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 wyrd_spec::DataTenantId;
pub const INITIAL_BACKOFF: Duration = Duration::from_millis(50);
pub const MAX_BACKOFF: Duration = Duration::from_secs(5);
pub trait OutboxSink: Send + Sync + 'static {
type Item: Send + 'static;
type Error: Display + Send + 'static;
const NAME: &'static str;
fn write(
&self,
tenant: DataTenantId,
items: &[Self::Item],
) -> impl Future<Output = Result<(), Self::Error>> + Send;
}
type QueueSender<S> = mpsc::UnboundedSender<(DataTenantId, <S as OutboxSink>::Item)>;
pub struct Outbox<S: OutboxSink> {
queue: RwLock<Option<QueueSender<S>>>,
pending: Arc<AtomicUsize>,
idle: Arc<Notify>,
abandon: CancellationToken,
writer: TaskTracker,
sink: Arc<S>,
}
impl<S: OutboxSink> Outbox<S> {
#[must_use]
pub fn new(sink: S, concurrency: usize) -> Arc<Self> {
let (queue, requests) = mpsc::unbounded_channel();
let pending = Arc::new(AtomicUsize::new(0));
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(),
failed: HashMap::new(),
writing: JoinSet::new(),
in_flight: HashMap::new(),
retry_at: HashMap::new(),
concurrency: concurrency.max(1),
pending: Arc::clone(&pending),
idle: Arc::clone(&idle),
abandon: abandon.clone(),
}
.run(),
);
writer.close();
Arc::new(Self {
queue: RwLock::new(Some(queue)),
pending,
idle,
abandon,
writer,
sink,
})
}
#[must_use]
pub fn sink(&self) -> &S {
&self.sink
}
pub fn stage(&self, tenant: DataTenantId, item: impl Into<S::Item>) {
let queue = self.queue.read().unwrap_or_else(PoisonError::into_inner);
if let Some(queue) = queue.as_ref() {
let gauge = metrics::gauge!("outbox_pending", "outbox" => S::NAME);
self.pending.fetch_add(1, Ordering::AcqRel);
gauge.increment(1.0);
if queue.send((tenant, item.into())).is_ok() {
return;
}
self.pending.fetch_sub(1, Ordering::AcqRel);
gauge.decrement(1.0);
}
count_lost::<S>(1);
tracing::error!(outbox = S::NAME, %tenant, "outbox item staged after shutdown was lost");
}
#[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(),
);
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 {
count_lost::<S>(lost);
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
}
}
type Written<S> = (
DataTenantId,
Vec<<S as OutboxSink>::Item>,
Result<(), String>,
);
struct OutboxWriter<S: OutboxSink> {
sink: Arc<S>,
requests: mpsc::UnboundedReceiver<(DataTenantId, S::Item)>,
waiting: HashMap<DataTenantId, Vec<S::Item>>,
failed: HashMap<DataTenantId, Vec<S::Item>>,
writing: JoinSet<Written<S>>,
in_flight: HashMap<Id, (DataTenantId, usize)>,
retry_at: HashMap<DataTenantId, (Instant, Duration)>,
concurrency: usize,
pending: Arc<AtomicUsize>,
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 {
self.dispatch();
if !open && self.waiting.is_empty() && self.failed.is_empty() && self.writing.is_empty()
{
return;
}
let retry = self.next_retry();
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;
for (tenant, item) in received.drain(..) {
self.waiting.entry(tenant).or_default().push(item);
}
}
() = async {
match retry {
Some(at) => sleep_until(at.into()).await,
None => pending().await,
}
} => {}
}
}
}
fn dispatch(&mut self) {
let now = Instant::now();
let ready: Vec<DataTenantId> = self
.waiting
.keys()
.chain(
self.failed
.keys()
.filter(|tenant| !self.waiting.contains_key(*tenant)),
)
.filter(|tenant| self.is_ready(**tenant, now))
.copied()
.collect();
for tenant in ready {
if self.writing.len() >= self.concurrency {
return;
}
let Some(items) = self
.failed
.remove(&tenant)
.or_else(|| self.waiting.remove(&tenant))
else {
continue;
};
let count = items.len();
let sink = Arc::clone(&self.sink);
let task = self.writing.spawn(async move {
let outcome = {
let write = catch_unwind(AssertUnwindSafe(|| sink.write(tenant, &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)
))
});
(tenant, items, result)
});
self.in_flight.insert(task.id(), (tenant, count));
}
}
fn is_ready(&self, tenant: DataTenantId, now: Instant) -> bool {
!self.in_flight.values().any(|(busy, _)| *busy == tenant)
&& self.retry_at.get(&tenant).is_none_or(|(at, _)| *at <= now)
}
fn next_retry(&self) -> Option<Instant> {
if self.writing.len() >= self.concurrency {
return None;
}
let now = Instant::now();
self.retry_at
.iter()
.filter(|(tenant, (at, _))| *at > now && self.failed.contains_key(*tenant))
.map(|(_, (at, _))| *at)
.min()
}
fn finish(&mut self, done: Result<(Id, Written<S>), JoinError>) {
let (tenant, items, result) = match done {
Ok((id, written)) => {
self.in_flight.remove(&id);
written
}
Err(error) => {
if let Some((tenant, count)) = self.in_flight.remove(&error.id()) {
count_lost::<S>(count);
tracing::error!(outbox = S::NAME, %tenant, %error, lost = count, "outbox write task failed");
self.release(count);
}
return;
}
};
let error = match result {
Ok(()) => {
self.retry_at.remove(&tenant);
self.release(items.len());
return;
}
Err(error) => error,
};
let backoff = self
.retry_at
.get(&tenant)
.map_or(INITIAL_BACKOFF, |(_, last)| (*last * 2).min(MAX_BACKOFF));
self.retry_at
.insert(tenant, (Instant::now() + backoff, backoff));
metrics::counter!("outbox_write_failures_total", "outbox" => S::NAME).increment(1);
tracing::error!(
outbox = S::NAME,
%tenant,
%error,
items = items.len(),
retry_in_ms = backoff.as_millis(),
"outbox write failed; retrying"
);
self.failed.insert(tenant, items);
}
fn release(&self, count: usize) {
metrics::gauge!("outbox_pending", "outbox" => S::NAME).decrement(count as f64);
if self.pending.fetch_sub(count, Ordering::AcqRel) == count {
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) {
metrics::counter!("outbox_events_lost_total", "outbox" => S::NAME).increment(count as u64);
}
#[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 wyrd_spec::DataTenantId;
use super::{Outbox, 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, 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";
async fn write(&self, tenant: DataTenantId, 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, 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")
}
}
fn outbox(concurrency: usize) -> (Arc<Outbox<MemorySink>>, Arc<Mutex<MemoryState>>) {
let sink = MemorySink::default();
let state = Arc::clone(&sink.state);
(Outbox::new(sink, concurrency), 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), state, started)
}
#[derive(Default)]
struct TestMetrics {
values: Mutex<HashMap<String, Arc<AtomicU64>>>,
}
impl TestMetrics {
fn cell(&self, key: &Key) -> Arc<AtomicU64> {
let mut values = self.values.lock().expect("test metrics are never poisoned");
Arc::clone(values.entry(key.name().to_owned()).or_default())
}
fn raw(&self, name: &str) -> u64 {
self.values
.lock()
.expect("test metrics are never poisoned")
.get(name)
.map_or(0, |value| value.load(Ordering::Acquire))
}
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, &(tenant, vec![3]), "the later item is its own write");
assert!(retries.len() >= 3, "two failures and one successful retry");
assert!(
retries
.iter()
.all(|attempt| attempt == &(tenant, vec![1, 2])),
"every retry is the identical failed slice: {retries:?}"
);
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 the_queue_has_no_count_limit() {
let (outbox, state) = outbox(4);
let tenant = DataTenantId::new_v7();
for item in 0..50_000_u32 {
outbox.stage(tenant, item);
}
assert_eq!(outbox.settle(within(30)).await, 0);
let written = state.lock().expect("unpoisoned").written[&tenant].clone();
assert_eq!(written, (0..50_000_u32).collect::<Vec<_>>());
}
#[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"), 3, "lost once");
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 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"), 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);
}
}