use std::collections::HashSet;
use std::future::Future;
use std::pin::Pin;
use std::sync::Arc;
use std::sync::atomic::{AtomicUsize, Ordering as AtomicOrdering};
use std::time::{Duration, Instant, SystemTime};
use async_trait::async_trait;
use tokio_postgres::{Client, Config, Transaction};
use super::{
Appended, Capacity, CapacityPolicy, Claim, ConsumerId, Delivery, DeliveryId, DeliveryMode,
JournalError, JournalStats, PoisonReason, UsageEvent, UsageJournal,
};
use crate::usage::{ObservedRecord, Status, UsageRecord};
const SCHEMA_DDL: &str = include_str!("../../../sql/usage_outbox_v1.sql");
const BACKEND: &str = "postgres";
const COUNT_REFRESH: Duration = Duration::from_secs(1);
const FLOOR_SETTLE_MARGIN: Duration = Duration::from_secs(300);
#[derive(Debug, Clone)]
pub struct PostgresJournalSettings {
pub schema: Option<String>,
pub create_schema: bool,
pub capacity: Capacity,
pub connect_timeout: Duration,
pub operation_timeout: Duration,
pub connections: usize,
}
impl Default for PostgresJournalSettings {
fn default() -> Self {
Self {
schema: None,
create_schema: false,
capacity: Capacity::BILLING_GRADE,
connect_timeout: Duration::from_secs(10),
operation_timeout: Duration::from_secs(10),
connections: 8,
}
}
}
pub struct PostgresJournal {
settings: PostgresJournalSettings,
pool: Pool,
stored: Arc<CapacityGate>,
}
impl std::fmt::Debug for PostgresJournal {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("PostgresJournal")
.field("schema", &self.settings.schema)
.field("capacity", &self.settings.capacity)
.finish_non_exhaustive()
}
}
impl PostgresJournal {
pub async fn connect(
dsn: &str,
settings: PostgresJournalSettings,
) -> Result<Self, JournalError> {
let mut config: Config = dsn
.parse()
.map_err(|error| {
backend(format!(
"the usage journal DSN could not be parsed: {error}"
))
})?;
config.connect_timeout(settings.connect_timeout);
config.application_name(crate::telemetry::SERVICE_NAME);
let search_path = settings
.schema
.as_deref()
.map(|schema| {
crate::usage::validate_table_name(schema).map_err(backend)?;
if schema.contains('.') {
return Err(backend(format!(
"`{schema}` is not a single unqualified schema name"
)));
}
Ok(schema.to_owned())
})
.transpose()?;
if settings.connections < 2 {
return Err(backend(
"a usage journal needs at least two connections: one is reserved for the \
delivery worker so its claims cannot stall the appends requests wait on",
));
}
let journal = Self {
pool: Pool::new(config, search_path, settings.connections),
settings,
stored: Arc::new(CapacityGate::new()),
};
if journal.settings.create_schema {
journal
.run("apply schema", Lane::Request, |client| {
Box::pin(async move {
client.batch_execute(SCHEMA_DDL).await?;
Ok(())
})
})
.await?;
}
journal.check_schema().await?;
Ok(journal)
}
async fn check_schema(&self) -> Result<(), JournalError> {
self.run("check schema", Lane::Request, |client| {
Box::pin(async move {
for (table, columns) in [
(
"axond_usage_outbox",
"position, request_id, schema_version, namespace, subject, record, \
observed_at, appended_at",
),
(
"axond_usage_outbox_consumer",
"consumer, registered_at, resolved_through",
),
(
"axond_usage_outbox_delivery",
"position, consumer, attempts, lease_expires_at, acknowledged_at, \
quarantined_at, poison_reason",
),
("axond_usage_outbox_loss", "id, dropped"),
] {
if let Err(error) = client
.query_opt(&format!("SELECT {columns} FROM {table} LIMIT 1"), &[])
.await
{
return Err(OpError::Journal(backend(format!(
"`{table}` is not readable with the columns this build needs, so it \
cannot own the usage outbox; apply (or re-apply) \
`ops/postgres/usage_outbox_v1.sql` (or set \
`[usage_journal] create_schema = true`): {error}"
))));
}
}
Ok(())
})
})
.await
}
async fn run<T, F>(&self, what: &'static str, lane: Lane, op: F) -> Result<T, JournalError>
where
T: Send,
F: for<'a> Fn(&'a mut Client) -> BoxFuture<'a, Result<T, OpError>> + Send + Sync,
{
match tokio::time::timeout(self.settings.operation_timeout, self.pool.run(&op, lane)).await
{
Ok(result) => result,
Err(_) => Err(backend(format!(
"`{what}` exceeded its {:?} bound",
self.settings.operation_timeout
))),
}
}
}
async fn stored_events(
tx: &Transaction<'_>,
gate: &CapacityGate,
max_events: u64,
) -> Result<u64, OpError> {
let span = tx.query_one(SPAN, &[]).await?.get::<_, i64>(0).max(0) as u64;
if span < max_events {
return Ok(span);
}
if let Some(estimate) = gate.estimate(span, max_events) {
return Ok(estimate);
}
let bound = i64::try_from(max_events.saturating_add(1)).unwrap_or(i64::MAX);
let row = tx
.query_one(
"SELECT count(*) FROM (SELECT 1 FROM axond_usage_outbox LIMIT $1) bounded",
&[&bound],
)
.await?;
let counted = row.get::<_, i64>(0).max(0) as u64;
gate.measured(span, counted, max_events);
Ok(counted)
}
const SPAN: &str = "SELECT COALESCE(max(position) - min(position) + 1, 0) FROM axond_usage_outbox";
#[async_trait]
impl UsageJournal for PostgresJournal {
fn name(&self) -> &'static str {
BACKEND
}
fn capacity(&self) -> Capacity {
self.settings.capacity
}
fn mode(&self) -> DeliveryMode {
DeliveryMode::BillingGrade
}
async fn append(&self, event: &UsageEvent) -> Result<Appended, JournalError> {
let idempotency_key = event.idempotency_key().clone();
let key = idempotency_key.as_str().to_owned();
let record = serde_json::to_value(event.record()).map_err(|error| {
backend(format!("the usage record could not be serialized: {error}"))
})?;
let capacity = self.settings.capacity;
let observed_at = event.observed_at();
let ordering = event.ordering_key().clone();
let version = i32::try_from(event.record().schema_version).unwrap_or(i32::MAX);
let gate = Arc::clone(&self.stored);
self.run("append", Lane::Request, move |client| {
let (key, record, ordering) = (key.clone(), record.clone(), ordering.clone());
let idempotency_key = idempotency_key.clone();
let gate = Arc::clone(&gate);
Box::pin(async move {
let tx = client.transaction().await?;
if let Some(row) = tx
.query_opt(
"SELECT position, record = $2::jsonb FROM axond_usage_outbox \
WHERE request_id = $1",
&[&key, &record],
)
.await?
{
let position = row.get::<_, i64>(0).max(0) as u64;
return if row.get::<_, bool>(1) {
Ok(Appended::AlreadyPresent { position })
} else {
Err(OpError::Journal(JournalError::Conflict {
key: idempotency_key,
}))
};
}
let mut dropped = 0;
let stored = stored_events(&tx, &gate, capacity.max_events).await?;
if stored >= capacity.max_events {
if gate.refusing() {
return Err(OpError::Journal(JournalError::AtCapacity {
pending: stored,
capacity,
}));
}
let mut surplus = surplus_cutoff(&tx, capacity.max_events).await?;
let mut reclaimed = 0;
if let Some(cutoff) = surplus {
reclaimed = reclaim_delivered(&tx, cutoff).await?;
if reclaimed > 0 {
surplus = surplus_cutoff(&tx, capacity.max_events).await?;
}
}
if let Some(cutoff) = surplus
&& capacity.policy == CapacityPolicy::DropOldest
{
dropped = drop_oldest(&tx, cutoff).await?;
if dropped > 0 {
surplus = surplus_cutoff(&tx, capacity.max_events).await?;
}
}
if reclaimed > 0 || dropped > 0 {
gate.invalidate();
}
if surplus.is_some() {
if reclaimed > 0 || dropped > 0 {
tx.commit().await?;
if dropped > 0 {
crate::telemetry::metrics::record_usage_journal_lost(
BACKEND,
"capacity_drop",
dropped,
);
}
} else {
gate.unreclaimable();
}
return Err(OpError::Journal(JournalError::AtCapacity {
pending: stored,
capacity,
}));
}
}
let row = tx
.query_one(
"INSERT INTO axond_usage_outbox \
(request_id, schema_version, namespace, subject, record, observed_at) \
VALUES ($1, $2, $3, $4, $5::jsonb, $6) \
RETURNING position",
&[
&key,
&version,
&ordering.namespace,
&ordering.subject,
&record,
&observed_at,
],
)
.await?;
tx.commit().await?;
if dropped > 0 {
crate::telemetry::metrics::record_usage_journal_lost(
BACKEND,
"capacity_drop",
dropped,
);
}
Ok(Appended::Accepted {
position: row.get::<_, i64>(0).max(0) as u64,
})
})
})
.await
}
async fn claim(
&self,
consumer: &ConsumerId,
claim: Claim,
) -> Result<Vec<Delivery>, JournalError> {
if claim.max_events == 0 {
return Ok(Vec::new());
}
let consumer = consumer.clone();
let name = consumer.as_str().to_owned();
let max_attempts =
i32::try_from(self.settings.capacity.max_delivery_attempts).unwrap_or(i32::MAX);
let lease_expires_at = claim.now + claim.lease;
let readable = i32::try_from(UsageRecord::SCHEMA_VERSION).unwrap_or(i32::MAX);
self.run("claim", Lane::Delivery, move |client| {
let (name, consumer) = (name.clone(), consumer.clone());
Box::pin(async move {
let tx = client.transaction().await?;
tx.execute(
"INSERT INTO axond_usage_outbox_consumer (consumer) VALUES ($1) \
ON CONFLICT (consumer) DO NOTHING",
&[&name],
)
.await?;
let floor: i64 = tx
.query_one(
"SELECT resolved_through FROM axond_usage_outbox_consumer \
WHERE consumer = $1",
&[&name],
)
.await?
.get(0);
let mut claimed: Vec<Delivery> = Vec::with_capacity(claim.max_events);
let mut condemnations: Vec<PoisonReason> = Vec::new();
let mut undeliverable = Undeliverable::default();
for _ in 0..claim.max_events {
let remaining =
i64::try_from(claim.max_events - claimed.len()).unwrap_or(i64::MAX);
let candidates = tx
.query(
"WITH open AS (
SELECT e.position, e.namespace, e.subject,
COALESCE(d.attempts, 0) AS attempts,
d.lease_expires_at
FROM axond_usage_outbox e
LEFT JOIN axond_usage_outbox_delivery d
ON d.position = e.position AND d.consumer = $1
AND d.position > $4
WHERE e.position > $4
AND d.acknowledged_at IS NULL AND d.quarantined_at IS NULL
),
head AS (
SELECT DISTINCT ON (namespace, subject)
position, attempts, lease_expires_at
FROM open
ORDER BY namespace, subject, position
)
SELECT h.position, h.attempts,
e.request_id, e.schema_version, e.record, e.observed_at
FROM head h
JOIN axond_usage_outbox e ON e.position = h.position
WHERE h.lease_expires_at IS NULL OR h.lease_expires_at <= $2
ORDER BY h.position
LIMIT $3
FOR UPDATE OF e SKIP LOCKED",
&[&name, &claim.now, &remaining, &floor],
)
.await?;
if candidates.is_empty() {
break;
}
let mut condemned = 0usize;
for row in candidates {
let position: i64 = row.get(0);
let attempt = row.get::<_, i32>(1).saturating_add(1);
let stored_version: i32 = row.get(3);
if stored_version > readable {
undeliverable.schema_ahead(position);
continue;
}
if attempt > max_attempts {
if condemn(
&tx,
position,
&name,
PoisonReason::AttemptsExhausted,
attempt,
)
.await?
{
condemnations.push(PoisonReason::AttemptsExhausted);
condemned += 1;
}
continue;
}
let event = match decode(&row) {
Ok(event) => event,
Err(reason) => {
tracing::error!(
position,
consumer = %name,
reason = %reason,
"usage outbox row could not be decoded; quarantining it"
);
undeliverable.corrupt();
if condemn(&tx, position, &name, PoisonReason::Malformed, attempt)
.await?
{
condemnations.push(PoisonReason::Malformed);
condemned += 1;
}
continue;
}
};
let taken = tx
.execute(
"INSERT INTO axond_usage_outbox_delivery
(position, consumer, attempts, lease_expires_at)
VALUES ($1, $2, $3, $4)
ON CONFLICT (position, consumer) DO UPDATE
SET attempts = $3, lease_expires_at = $4
WHERE axond_usage_outbox_delivery.acknowledged_at IS NULL
AND axond_usage_outbox_delivery.quarantined_at IS NULL
AND (axond_usage_outbox_delivery.lease_expires_at IS NULL
OR axond_usage_outbox_delivery.lease_expires_at <= $5)",
&[&position, &name, &attempt, &lease_expires_at, &claim.now],
)
.await?;
if taken == 0 {
continue;
}
claimed.push(Delivery {
id: DeliveryId {
consumer: consumer.clone(),
event: event.id(),
attempt: attempt.max(1) as u32,
},
event,
lease_expires_at,
});
}
if condemned == 0 || claimed.len() >= claim.max_events {
break;
}
}
tx.commit().await?;
for reason in condemnations {
crate::telemetry::metrics::record_usage_journal_quarantined(
BACKEND,
&name,
reason.as_str(),
);
}
for reason in undeliverable.reasons {
crate::telemetry::metrics::record_usage_journal_undeliverable(BACKEND, reason);
}
Ok(claimed)
})
})
.await
}
async fn ack(&self, delivery: &DeliveryId) -> Result<(), JournalError> {
self.verdict(delivery, None).await
}
async fn ack_all(&self, deliveries: &[DeliveryId]) -> Vec<Result<(), JournalError>> {
if deliveries.len() < 2 {
let mut verdicts = Vec::with_capacity(deliveries.len());
for delivery in deliveries {
verdicts.push(self.ack(delivery).await);
}
return verdicts;
}
let name = deliveries[0].consumer.as_str().to_owned();
let keys: Vec<String> = deliveries
.iter()
.filter(|delivery| delivery.consumer.as_str() == name)
.map(|delivery| delivery.event.to_string())
.collect();
let resolved = self
.run("ack_all", Lane::Delivery, {
let (name, keys) = (name.clone(), keys.clone());
move |client| {
let (name, keys) = (name.clone(), keys.clone());
Box::pin(async move {
let rows = client
.query(
"UPDATE axond_usage_outbox_delivery d
SET acknowledged_at = now(), lease_expires_at = NULL
FROM axond_usage_outbox e
WHERE e.position = d.position
AND d.consumer = $2
AND e.request_id = ANY($1)
AND d.attempts > 0
AND d.acknowledged_at IS NULL
AND d.quarantined_at IS NULL
RETURNING e.request_id",
&[&keys, &name],
)
.await?;
Ok(rows
.iter()
.map(|row| row.get::<_, String>(0))
.collect::<HashSet<String>>())
})
}
})
.await;
let resolved: HashSet<String> = resolved.unwrap_or_default();
let mut verdicts = Vec::with_capacity(deliveries.len());
for delivery in deliveries {
let resolved = delivery.consumer.as_str() == name
&& resolved.contains(&delivery.event.to_string());
if resolved {
verdicts.push(Ok(()));
} else {
verdicts.push(self.ack(delivery).await);
}
}
verdicts
}
async fn quarantine(
&self,
delivery: &DeliveryId,
reason: PoisonReason,
) -> Result<(), JournalError> {
self.verdict(delivery, Some(reason)).await
}
async fn relinquish(&self, delivery: &DeliveryId) -> Result<(), JournalError> {
let (name, key) = (
delivery.consumer.as_str().to_owned(),
delivery.event.to_string(),
);
let attempt = i32::try_from(delivery.attempt).unwrap_or(i32::MAX);
let refused = delivery.clone();
self.run("relinquish", Lane::Delivery, move |client| {
let (name, key, refused) = (name.clone(), key.clone(), refused.clone());
Box::pin(async move {
let refunded = client
.execute(
"UPDATE axond_usage_outbox_delivery d
SET attempts = d.attempts - 1
FROM axond_usage_outbox e
WHERE e.position = d.position AND e.request_id = $1
AND d.consumer = $2 AND d.attempts = $3
AND d.acknowledged_at IS NULL AND d.quarantined_at IS NULL",
&[&key, &name, &attempt],
)
.await?;
if refunded == 0 {
let exists = client
.query_opt(
"SELECT 1 FROM axond_usage_outbox e
JOIN axond_usage_outbox_delivery d
ON d.position = e.position AND d.consumer = $2
WHERE e.request_id = $1",
&[&key, &name],
)
.await?;
if exists.is_none()
&& client
.query_opt(
"SELECT 1 FROM axond_usage_outbox WHERE request_id = $1",
&[&key],
)
.await?
.is_some()
{
return Err(OpError::Journal(JournalError::NotOutstanding {
delivery: refused,
}));
}
}
Ok(())
})
})
.await
}
async fn stats(&self, consumer: &ConsumerId) -> Result<JournalStats, JournalError> {
let name = consumer.as_str().to_owned();
let now = SystemTime::now();
let capacity = self.settings.capacity;
self.run("stats", Lane::Delivery, move |client| {
let name = name.clone();
Box::pin(async move {
let row = client
.query_one(
"WITH floor AS (
SELECT COALESCE(
(SELECT resolved_through FROM axond_usage_outbox_consumer
WHERE consumer = $1), 0) AS position
),
state AS (
SELECT e.observed_at, d.acknowledged_at, d.quarantined_at,
d.lease_expires_at,
(d.acknowledged_at IS NULL AND d.quarantined_at IS NULL
AND (d.lease_expires_at IS NULL
OR d.lease_expires_at <= $2)) AS pending
FROM axond_usage_outbox e
LEFT JOIN axond_usage_outbox_delivery d
ON d.position = e.position AND d.consumer = $1
AND d.position > (SELECT position FROM floor)
WHERE e.position > (SELECT position FROM floor)
)
SELECT
count(*) FILTER (WHERE pending),
count(*) FILTER (WHERE acknowledged_at IS NULL
AND quarantined_at IS NULL
AND lease_expires_at > $2),
(SELECT count(*) FROM axond_usage_outbox_delivery
WHERE consumer = $1 AND quarantined_at IS NOT NULL),
min(observed_at) FILTER (WHERE pending),
(SELECT dropped FROM axond_usage_outbox_loss WHERE id)
FROM state",
&[&name, &now],
)
.await?;
let oldest: Option<SystemTime> = row.get(3);
Ok(JournalStats {
pending: row.get::<_, i64>(0).max(0) as u64,
in_flight: row.get::<_, i64>(1).max(0) as u64,
quarantined: row.get::<_, i64>(2).max(0) as u64,
oldest_pending_age: oldest.and_then(|oldest| now.duration_since(oldest).ok()),
dropped: row.get::<_, Option<i64>>(4).unwrap_or_default().max(0) as u64,
capacity,
})
})
})
.await
}
async fn maintain(&self, now: SystemTime) -> Result<u64, JournalError> {
let retain = self.settings.capacity.retain_acknowledged;
let settled = now
.checked_sub(FLOOR_SETTLE_MARGIN.max(self.settings.operation_timeout * 6))
.unwrap_or(now);
let pruned = self
.run("maintain", Lane::Delivery, move |client| {
Box::pin(async move {
let cutoff = now.checked_sub(retain).unwrap_or(now);
let pruned = client
.execute(
&format!(
"DELETE FROM axond_usage_outbox e WHERE e.observed_at <= $1 \
AND {DELIVERED}"
),
&[&cutoff],
)
.await?;
client
.execute(ADVANCE_RESOLVED_THROUGH, &[&settled])
.await?;
Ok(pruned)
})
})
.await?;
if pruned > 0 {
self.stored.invalidate();
}
Ok(pruned)
}
async fn consumers_besides(&self, mine: &ConsumerId) -> Result<Vec<String>, JournalError> {
let name = mine.as_str().to_owned();
self.run("consumers", Lane::Delivery, move |client| {
let name = name.clone();
Box::pin(async move {
let rows = client
.query(
"SELECT consumer FROM axond_usage_outbox_consumer WHERE consumer <> $1 \
ORDER BY consumer",
&[&name],
)
.await?;
Ok(rows.iter().map(|row| row.get::<_, String>(0)).collect())
})
})
.await
}
}
impl PostgresJournal {
async fn verdict(
&self,
delivery: &DeliveryId,
poison: Option<PoisonReason>,
) -> Result<(), JournalError> {
let (name, event) = (delivery.consumer.as_str().to_owned(), delivery.event);
let key = event.to_string();
let refused = delivery.clone();
self.run("verdict", Lane::Delivery, move |client| {
let (name, key, refused) = (name.clone(), key.clone(), refused.clone());
Box::pin(async move {
let tx = client.transaction().await?;
let state = tx
.query_opt(
"SELECT d.attempts, d.acknowledged_at IS NOT NULL,
d.quarantined_at IS NOT NULL
FROM axond_usage_outbox e
JOIN axond_usage_outbox_delivery d
ON d.position = e.position AND d.consumer = $2
WHERE e.request_id = $1
FOR UPDATE OF d",
&[&key, &name],
)
.await?;
let Some(state) = state.filter(|row| row.get::<_, i32>(0) > 0) else {
let vanished = tx
.query_opt(
"SELECT 1 FROM axond_usage_outbox WHERE request_id = $1",
&[&key],
)
.await?
.is_none();
if vanished && poison.is_none() {
return Ok(());
}
return Err(OpError::Journal(JournalError::NotOutstanding {
delivery: refused,
}));
};
let (acknowledged, quarantined) =
(state.get::<_, bool>(1), state.get::<_, bool>(2));
match poison {
None if acknowledged => return Ok(()),
None if quarantined => {
return Err(OpError::Journal(JournalError::Quarantined {
delivery: refused,
}));
}
Some(_) if quarantined => return Ok(()),
Some(_) if acknowledged => {
return Err(OpError::Journal(JournalError::AlreadyAcknowledged {
delivery: refused,
}));
}
_ => {}
}
let sql = match poison {
None => {
"UPDATE axond_usage_outbox_delivery d
SET acknowledged_at = now(), lease_expires_at = NULL
FROM axond_usage_outbox e
WHERE e.position = d.position AND e.request_id = $1 AND d.consumer = $2"
}
Some(_) => {
"UPDATE axond_usage_outbox_delivery d
SET quarantined_at = now(), lease_expires_at = NULL, poison_reason = $3
FROM axond_usage_outbox e
WHERE e.position = d.position AND e.request_id = $1 AND d.consumer = $2"
}
};
match poison {
None => tx.execute(sql, &[&key, &name]).await?,
Some(reason) => tx.execute(sql, &[&key, &name, &reason.as_str()]).await?,
};
tx.commit().await?;
Ok(())
})
})
.await
}
}
#[derive(Default)]
struct Undeliverable {
reasons: Vec<&'static str>,
ahead: HashSet<i64>,
}
impl Undeliverable {
fn schema_ahead(&mut self, position: i64) {
if self.ahead.insert(position) {
self.reasons.push("schema_ahead");
}
}
fn corrupt(&mut self) {
self.reasons.push("corrupt");
}
}
const ADVANCE_RESOLVED_THROUGH: &str = "UPDATE axond_usage_outbox_consumer c
SET resolved_through = GREATEST(
c.resolved_through,
LEAST(
COALESCE(
(SELECT e.position - 1
FROM axond_usage_outbox e
LEFT JOIN axond_usage_outbox_delivery d
ON d.position = e.position AND d.consumer = c.consumer
WHERE e.position > c.resolved_through
AND d.acknowledged_at IS NULL AND d.quarantined_at IS NULL
ORDER BY e.position
LIMIT 1),
(SELECT COALESCE(max(position), 0) FROM axond_usage_outbox)),
COALESCE(
(SELECT e.position FROM axond_usage_outbox e
WHERE e.appended_at <= $1
ORDER BY e.position DESC
LIMIT 1),
0)))";
const DELIVERED: &str = "EXISTS (SELECT 1 FROM axond_usage_outbox_consumer)
AND NOT EXISTS (
SELECT 1 FROM axond_usage_outbox_delivery d
WHERE d.position = e.position AND d.quarantined_at IS NOT NULL)
AND NOT EXISTS (
SELECT 1 FROM axond_usage_outbox_consumer c
WHERE NOT EXISTS (
SELECT 1 FROM axond_usage_outbox_delivery d
WHERE d.position = e.position AND d.consumer = c.consumer
AND d.acknowledged_at IS NOT NULL))";
async fn reclaim_delivered(tx: &Transaction<'_>, cutoff: i64) -> Result<u64, OpError> {
Ok(tx
.execute(
&format!(
"DELETE FROM axond_usage_outbox WHERE position IN (
SELECT e.position FROM axond_usage_outbox e
WHERE e.position <= $1 AND {DELIVERED})"
),
&[&cutoff],
)
.await?)
}
async fn surplus_cutoff(tx: &Transaction<'_>, max_events: u64) -> Result<Option<i64>, OpError> {
let keep = i64::try_from(max_events.saturating_sub(1)).unwrap_or(i64::MAX);
Ok(tx
.query_opt(
"SELECT position FROM axond_usage_outbox ORDER BY position DESC OFFSET $1 LIMIT 1",
&[&keep],
)
.await?
.map(|row| row.get::<_, i64>(0)))
}
async fn drop_oldest(tx: &Transaction<'_>, cutoff: i64) -> Result<u64, OpError> {
let dropped = tx
.execute(
"DELETE FROM axond_usage_outbox WHERE position IN (
SELECT e.position FROM axond_usage_outbox e
WHERE e.position <= $1 AND NOT EXISTS (
SELECT 1 FROM axond_usage_outbox_delivery d
WHERE d.position = e.position AND d.quarantined_at IS NOT NULL))",
&[&cutoff],
)
.await?;
if dropped > 0 {
let lost = i64::try_from(dropped).unwrap_or(i64::MAX);
tx.execute(
"UPDATE axond_usage_outbox_loss SET dropped = dropped + $1 WHERE id",
&[&lost],
)
.await?;
}
Ok(dropped)
}
async fn condemn(
tx: &Transaction<'_>,
position: i64,
consumer: &str,
reason: PoisonReason,
attempt: i32,
) -> Result<bool, OpError> {
let condemned = tx
.execute(
"INSERT INTO axond_usage_outbox_delivery
(position, consumer, attempts, quarantined_at, poison_reason)
VALUES ($1, $2, $3, now(), $4)
ON CONFLICT (position, consumer) DO UPDATE
SET attempts = $3, quarantined_at = now(), poison_reason = $4,
lease_expires_at = NULL
WHERE axond_usage_outbox_delivery.acknowledged_at IS NULL
AND axond_usage_outbox_delivery.quarantined_at IS NULL",
&[&position, &consumer, &attempt, &reason.as_str()],
)
.await?;
Ok(condemned == 1)
}
#[derive(serde::Deserialize)]
struct StoredRecord {
schema_version: u32,
request_id: String,
#[serde(default)]
trace_id: Option<String>,
namespace: String,
subject: String,
#[serde(default)]
signer_kid: Option<String>,
model: String,
target_provider: String,
target_model: String,
credential_source: String,
credential_id: String,
status: Status,
input_tokens: u64,
cache_read_tokens: u64,
cache_write_tokens: u64,
output_tokens: u64,
cost_microdollars: u64,
catalog_version: u64,
latency_ms: u64,
attempts: u32,
}
impl StoredRecord {
fn into_record(self) -> Result<UsageRecord, String> {
let credential_source = match self.credential_source.as_str() {
"platform" => "platform",
"byok" => "byok",
other => return Err(format!("`{other}` is not a credential source")),
};
Ok(UsageRecord {
schema_version: self.schema_version,
request_id: self.request_id,
trace_id: self.trace_id,
namespace: self.namespace,
subject: self.subject,
signer_kid: self.signer_kid,
model: self.model,
target_provider: self.target_provider,
target_model: self.target_model,
credential_source,
credential_id: self.credential_id,
status: self.status,
input_tokens: self.input_tokens,
cache_read_tokens: self.cache_read_tokens,
cache_write_tokens: self.cache_write_tokens,
output_tokens: self.output_tokens,
cost_microdollars: self.cost_microdollars,
catalog_version: self.catalog_version,
latency_ms: self.latency_ms,
attempts: self.attempts,
})
}
}
fn decode(row: &tokio_postgres::Row) -> Result<UsageEvent, String> {
let request_id: String = row.get(2);
let stored: StoredRecord = serde_json::from_value(row.get::<_, serde_json::Value>(4))
.map_err(|error| format!("the stored record is unreadable: {error}"))?;
if stored.request_id != request_id {
return Err(format!(
"the row's `request_id` and the record's identity disagree: `{request_id}` \
against `{}`",
stored.request_id
));
}
let record = stored.into_record()?;
UsageEvent::new(ObservedRecord {
record,
observed_at: row.get(5),
})
.map_err(|error| error.to_string())
}
fn backend(message: impl Into<String>) -> JournalError {
JournalError::Backend(message.into())
}
type BoxFuture<'a, T> = Pin<Box<dyn Future<Output = T> + Send + 'a>>;
enum OpError {
Db(tokio_postgres::Error),
Journal(JournalError),
}
impl From<tokio_postgres::Error> for OpError {
fn from(error: tokio_postgres::Error) -> Self {
Self::Db(error)
}
}
#[derive(Clone, Copy, PartialEq, Eq, Debug)]
enum Lane {
Request,
Delivery,
}
struct Pool {
config: Config,
search_path: Option<String>,
slots: Vec<tokio::sync::Mutex<Option<Client>>>,
next: AtomicUsize,
}
impl Pool {
fn new(config: Config, search_path: Option<String>, connections: usize) -> Self {
Self {
config,
search_path,
slots: (0..connections.max(2))
.map(|_| tokio::sync::Mutex::new(None))
.collect(),
next: AtomicUsize::new(0),
}
}
async fn run<T, F>(&self, op: &F, lane: Lane) -> Result<T, JournalError>
where
T: Send,
F: for<'a> Fn(&'a mut Client) -> BoxFuture<'a, Result<T, OpError>> + Send + Sync,
{
let mut guard = self.acquire(lane).await;
let mut last: Option<tokio_postgres::Error> = None;
for _ in 0..2 {
let mut client = match guard.take() {
Some(client) if !client.is_closed() => client,
_ => match self.connect().await {
Ok(client) => client,
Err(error) => {
last = Some(error);
continue;
}
},
};
match op(&mut client).await {
Ok(value) => {
*guard = Some(client);
return Ok(value);
}
Err(OpError::Journal(error)) => {
*guard = Some(client);
return Err(error);
}
Err(OpError::Db(error)) => last = Some(error),
}
}
Err(backend(last.map_or_else(
|| "the usage journal operation failed".to_owned(),
|error| error.to_string(),
)))
}
async fn acquire(&self, lane: Lane) -> tokio::sync::MutexGuard<'_, Option<Client>> {
let lane = match lane {
Lane::Delivery => &self.slots[self.slots.len() - 1..],
Lane::Request => &self.slots[..self.slots.len() - 1],
};
for slot in lane {
if let Ok(guard) = slot.try_lock() {
return guard;
}
}
let index = self.next.fetch_add(1, AtomicOrdering::Relaxed) % lane.len();
lane[index].lock().await
}
async fn connect(&self) -> Result<Client, tokio_postgres::Error> {
let (client, connection) = self.config.connect(crate::usage::tls_connector()).await?;
tokio::spawn(async move {
if let Err(error) = connection.await {
tracing::warn!(error = %error, "usage journal connection closed");
}
});
if let Some(schema) = self.search_path.as_deref() {
client
.batch_execute(&format!("SET search_path TO {schema}"))
.await?;
}
Ok(client)
}
}
struct CapacityGate {
state: std::sync::Mutex<GateState>,
}
#[derive(Default)]
struct GateState {
measured: Option<Measured>,
at: Option<Instant>,
unreclaimable: Option<Instant>,
}
#[derive(Clone, Copy)]
enum Measured {
Gaps { gaps: u64, span: u64 },
Over,
}
impl CapacityGate {
fn new() -> Self {
Self {
state: std::sync::Mutex::new(GateState::default()),
}
}
fn estimate(&self, span: u64, max_events: u64) -> Option<u64> {
let state = self.state.lock().expect("capacity gate");
if state.at?.elapsed() >= COUNT_REFRESH {
return None;
}
Some(match state.measured? {
Measured::Gaps {
gaps,
span: measured,
} if span >= measured => span.saturating_sub(gaps),
Measured::Gaps { .. } => return None,
Measured::Over => max_events.saturating_add(1),
})
}
fn measured(&self, span: u64, counted: u64, max_events: u64) {
let measured = if counted > max_events {
Measured::Over
} else {
Measured::Gaps {
gaps: span.saturating_sub(counted),
span,
}
};
let mut state = self.state.lock().expect("capacity gate");
state.measured = Some(measured);
state.at = Some(Instant::now());
}
fn refusing(&self) -> bool {
self.state
.lock()
.expect("capacity gate")
.unreclaimable
.is_some_and(|at| at.elapsed() < COUNT_REFRESH)
}
fn unreclaimable(&self) {
self.state.lock().expect("capacity gate").unreclaimable = Some(Instant::now());
}
fn invalidate(&self) {
let mut state = self.state.lock().expect("capacity gate");
state.at = None;
state.unreclaimable = None;
}
}
#[cfg(test)]
mod tests {
use super::super::tests::{consumer, event_for};
use super::*;
#[test]
fn a_stored_record_decodes_back_to_exactly_the_record_that_was_written() {
let record = crate::usage::tests::sample_record();
let stored: StoredRecord = serde_json::from_value(
serde_json::to_value(&record).expect("a usage record serializes"),
)
.expect("the stored mirror reads every field the record writes");
assert_eq!(
stored.into_record().expect("a known credential source"),
record
);
}
#[tokio::test]
async fn a_claim_in_flight_does_not_hold_a_connection_an_append_needs() {
let pool = Pool::new(
"host=127.0.0.1 user=nobody".parse().expect("a config"),
None,
2,
);
let claim = pool.acquire(Lane::Delivery).await;
let append = tokio::time::timeout(Duration::from_millis(50), pool.acquire(Lane::Request));
assert!(
append.await.is_ok(),
"an append waited on the worker's connection"
);
let second = tokio::time::timeout(Duration::from_millis(50), pool.acquire(Lane::Delivery));
assert!(
second.await.is_err(),
"a second claim took a connection reserved for requests"
);
drop(claim);
}
fn capacity(max_events: u64, policy: CapacityPolicy) -> Capacity {
Capacity {
max_events,
max_delivery_attempts: 3,
retain_acknowledged: Duration::from_secs(3600),
policy,
}
}
fn settings(schema: &str, create_schema: bool, capacity: Capacity) -> PostgresJournalSettings {
PostgresJournalSettings {
schema: Some(schema.to_owned()),
create_schema,
capacity,
..PostgresJournalSettings::default()
}
}
async fn client(dsn: &str, schema: Option<&str>) -> Client {
let (client, connection) = dsn
.parse::<Config>()
.expect("a test DSN")
.connect(crate::usage::tls_connector())
.await
.expect("connect");
tokio::spawn(async move {
let _ = connection.await;
});
if let Some(schema) = schema {
client
.batch_execute(&format!("SET search_path TO {schema}"))
.await
.expect("search_path");
}
client
}
async fn outbox(name: &str, capacity: Capacity) -> Option<(String, PostgresJournal)> {
let dsn = crate::test_services::postgres_dsn()?;
let schema = format!("axond_outbox_{name}");
let admin = client(&dsn, None).await;
admin
.batch_execute(&format!(
"DROP SCHEMA IF EXISTS {schema} CASCADE; CREATE SCHEMA {schema}"
))
.await
.expect("a schema of its own");
let journal = PostgresJournal::connect(&dsn, settings(&schema, true, capacity))
.await
.expect("connect");
Some((dsn, journal))
}
fn claim_at(max_events: usize, lease: Duration, now: SystemTime) -> Claim {
Claim {
max_events,
lease,
now,
}
}
#[tokio::test]
async fn an_outbox_that_is_not_there_refuses_to_boot() {
let Some(dsn) = crate::test_services::postgres_dsn() else {
return;
};
let schema = "axond_outbox_missing";
let admin = client(&dsn, None).await;
admin
.batch_execute(&format!(
"DROP SCHEMA IF EXISTS {schema} CASCADE; CREATE SCHEMA {schema}"
))
.await
.expect("an empty schema");
let error = PostgresJournal::connect(
&dsn,
settings(schema, false, capacity(16, CapacityPolicy::Refuse)),
)
.await
.expect_err("an outbox that is not there is a boot failure");
assert!(
matches!(&error, JournalError::Backend(message) if message.contains("axond_usage_outbox")),
"{error:?}"
);
}
#[tokio::test]
async fn a_consumer_table_from_before_the_claim_floor_is_migrated_not_served() {
let Some(dsn) = crate::test_services::postgres_dsn() else {
return;
};
let schema = "axond_outbox_migration";
let admin = client(&dsn, None).await;
admin
.batch_execute(&format!(
"DROP SCHEMA IF EXISTS {schema} CASCADE; CREATE SCHEMA {schema}"
))
.await
.expect("a schema of its own");
let old = client(&dsn, Some(schema)).await;
old.batch_execute(SCHEMA_DDL).await.expect("the schema");
old.batch_execute("ALTER TABLE axond_usage_outbox_consumer DROP COLUMN resolved_through")
.await
.expect("the shape an earlier copy created");
let error = PostgresJournal::connect(
&dsn,
settings(schema, false, capacity(16, CapacityPolicy::Refuse)),
)
.await
.expect_err("a consumer table this build cannot claim from is a boot failure");
assert!(
matches!(&error, JournalError::Backend(message)
if message.contains("axond_usage_outbox_consumer")),
"{error:?}"
);
let journal = PostgresJournal::connect(
&dsn,
settings(schema, true, capacity(16, CapacityPolicy::Refuse)),
)
.await
.expect("re-applying the DDL adopts the column");
let billing = consumer("billing");
let event = event_for("GW_INBOUND_ACME_KEY");
assert!(journal.append(&event).await.expect("append").is_new());
let claimed = journal
.claim(
&billing,
claim_at(4, Duration::from_secs(30), SystemTime::now()),
)
.await
.expect("a claim reads the migrated column");
assert_eq!(claimed.len(), 1);
}
#[tokio::test]
async fn an_appended_event_survives_the_process_that_appended_it() {
let Some((dsn, journal)) = outbox("restart", capacity(16, CapacityPolicy::Refuse)).await
else {
return;
};
let event = event_for("GW_INBOUND_ACME_KEY");
assert!(journal.append(&event).await.expect("append").is_new());
drop(journal);
let restarted = PostgresJournal::connect(
&dsn,
settings(
"axond_outbox_restart",
false,
capacity(16, CapacityPolicy::Refuse),
),
)
.await
.expect("reconnect");
let billing = consumer("billing");
let claimed = restarted
.claim(
&billing,
claim_at(8, Duration::from_secs(30), SystemTime::now()),
)
.await
.expect("claim");
assert_eq!(claimed.len(), 1, "{claimed:?}");
assert_eq!(claimed[0].event.id(), event.id());
assert_eq!(claimed[0].event.record(), event.record());
assert_eq!(claimed[0].id.attempt, 1);
restarted.ack(&claimed[0].id).await.expect("ack");
assert!(
restarted
.claim(
&billing,
claim_at(8, Duration::from_secs(30), SystemTime::now())
)
.await
.expect("claim")
.is_empty()
);
let stats = restarted.stats(&billing).await.expect("stats");
assert!(stats.is_drained(), "{stats:?}");
}
#[tokio::test]
async fn appending_the_same_event_twice_journals_it_once() {
let Some((_, journal)) = outbox("idempotent", capacity(16, CapacityPolicy::Refuse)).await
else {
return;
};
let event = event_for("GW_INBOUND_ACME_KEY");
let first = journal.append(&event).await.expect("append");
let again = UsageEvent::new(ObservedRecord {
record: event.record().clone(),
observed_at: event.observed_at() + Duration::from_secs(5),
})
.expect("the same fact");
let second = journal.append(&again).await.expect("append");
assert!(first.is_new());
assert!(!second.is_new(), "{second:?}");
assert_eq!(first.position(), second.position());
let mut different = event.record().clone();
different.cost_microdollars += 1;
let conflicting =
UsageEvent::new(ObservedRecord::now(different)).expect("a well-formed event");
let error = journal
.append(&conflicting)
.await
.expect_err("the same identity with different content is a conflict");
assert!(
matches!(&error, JournalError::Conflict { key } if key == event.idempotency_key()),
"{error:?}"
);
let claimed = journal
.claim(
&consumer("billing"),
claim_at(8, Duration::from_secs(30), SystemTime::now()),
)
.await
.expect("claim");
assert_eq!(claimed.len(), 1, "{claimed:?}");
assert_eq!(
claimed[0].event.record().cost_microdollars,
event.record().cost_microdollars
);
}
#[tokio::test]
async fn an_expired_lease_redelivers_the_event_as_a_new_attempt() {
let Some((_, journal)) = outbox("lease", capacity(16, CapacityPolicy::Refuse)).await else {
return;
};
let event = event_for("GW_INBOUND_ACME_KEY");
journal.append(&event).await.expect("append");
let billing = consumer("billing");
let now = SystemTime::now();
let first = journal
.claim(&billing, claim_at(8, Duration::from_secs(30), now))
.await
.expect("claim");
assert_eq!(first.len(), 1);
assert!(
journal
.claim(&billing, claim_at(8, Duration::from_secs(30), now))
.await
.expect("claim")
.is_empty()
);
let again = journal
.claim(
&billing,
claim_at(8, Duration::from_secs(30), now + Duration::from_secs(31)),
)
.await
.expect("claim");
assert_eq!(again.len(), 1, "{again:?}");
assert_eq!(again[0].event.id(), event.id());
assert_eq!(again[0].id.attempt, 2);
assert!(again[0].id.is_redelivery());
journal.ack(&first[0].id).await.expect("late ack");
journal.ack(&again[0].id).await.expect("idempotent ack");
}
#[tokio::test]
async fn one_callers_events_are_claimed_in_order_and_one_at_a_time() {
let Some((_, journal)) = outbox("ordering", capacity(16, CapacityPolicy::Refuse)).await
else {
return;
};
let first = event_for("GW_INBOUND_ACME_KEY");
let second = event_for("GW_INBOUND_ACME_KEY");
let other = event_for("GW_INBOUND_OTHER_KEY");
for event in [&first, &second, &other] {
journal.append(event).await.expect("append");
}
let billing = consumer("billing");
let now = SystemTime::now();
let claimed = journal
.claim(&billing, claim_at(8, Duration::from_secs(30), now))
.await
.expect("claim");
let ids: Vec<_> = claimed.iter().map(|d| d.event.id()).collect();
assert_eq!(ids, vec![first.id(), other.id()], "{ids:?}");
journal.ack(&claimed[0].id).await.expect("ack");
let next = journal
.claim(&billing, claim_at(8, Duration::from_secs(30), now))
.await
.expect("claim");
let ids: Vec<_> = next.iter().map(|d| d.event.id()).collect();
assert_eq!(ids, vec![second.id()], "{ids:?}");
}
#[tokio::test]
async fn a_delivery_that_was_never_claimed_cannot_be_acknowledged() {
let Some((_, journal)) = outbox("stray", capacity(16, CapacityPolicy::Refuse)).await else {
return;
};
let event = event_for("GW_INBOUND_ACME_KEY");
journal.append(&event).await.expect("append");
let stray = DeliveryId {
consumer: consumer("billing"),
event: event.id(),
attempt: 1,
};
let error = journal
.ack(&stray)
.await
.expect_err("a consumer that never claimed has nothing to acknowledge");
assert!(
matches!(error, JournalError::NotOutstanding { .. }),
"{error:?}"
);
let stats = journal.stats(&consumer("billing")).await.expect("stats");
assert_eq!(stats.pending, 1, "{stats:?}");
}
#[tokio::test]
async fn an_event_that_exhausts_its_attempts_is_quarantined_not_retried_forever() {
let Some((_, journal)) = outbox(
"attempts",
Capacity {
max_delivery_attempts: 2,
..capacity(16, CapacityPolicy::Refuse)
},
)
.await
else {
return;
};
let blocked = event_for("GW_INBOUND_ACME_KEY");
let behind = event_for("GW_INBOUND_ACME_KEY");
journal.append(&blocked).await.expect("append");
journal.append(&behind).await.expect("append");
let billing = consumer("billing");
let mut now = SystemTime::now();
for attempt in 1..=2 {
let claimed = journal
.claim(&billing, claim_at(8, Duration::from_secs(1), now))
.await
.expect("claim");
assert_eq!(claimed.len(), 1, "attempt {attempt}: {claimed:?}");
assert_eq!(claimed[0].event.id(), blocked.id());
now += Duration::from_secs(2);
}
let claimed = journal
.claim(&billing, claim_at(8, Duration::from_secs(1), now))
.await
.expect("claim");
assert_eq!(claimed.len(), 1, "{claimed:?}");
assert_eq!(claimed[0].event.id(), behind.id());
let stats = journal.stats(&billing).await.expect("stats");
assert_eq!(stats.quarantined, 1, "{stats:?}");
let error = journal
.ack(&DeliveryId {
consumer: billing.clone(),
event: blocked.id(),
attempt: 2,
})
.await
.expect_err("quarantine is terminal");
assert!(
matches!(error, JournalError::Quarantined { .. }),
"{error:?}"
);
}
#[tokio::test]
async fn a_full_outbox_refuses_the_append_rather_than_dropping_usage() {
let Some((_, journal)) = outbox("refuse", capacity(1, CapacityPolicy::Refuse)).await else {
return;
};
journal
.append(&event_for("GW_INBOUND_ACME_KEY"))
.await
.expect("append");
let error = journal
.append(&event_for("GW_INBOUND_ACME_KEY"))
.await
.expect_err("a full billing-grade outbox refuses");
assert!(
matches!(
error,
JournalError::AtCapacity { pending, capacity } if pending >= 1 && capacity.max_events == 1
),
"{error:?}"
);
let stats = journal.stats(&consumer("billing")).await.expect("stats");
assert_eq!(stats.dropped, 0, "a refusal is not a loss: {stats:?}");
}
#[tokio::test]
async fn drop_oldest_bounds_storage_and_counts_what_it_lost() {
let Some((_, journal)) =
outbox("drop_oldest", capacity(1, CapacityPolicy::DropOldest)).await
else {
return;
};
journal
.append(&event_for("GW_INBOUND_ACME_KEY"))
.await
.expect("append");
let kept = event_for("GW_INBOUND_ACME_KEY");
journal.append(&kept).await.expect("append");
let billing = consumer("billing");
let stats = journal.stats(&billing).await.expect("stats");
assert_eq!(stats.pending, 1, "{stats:?}");
assert_eq!(stats.dropped, 1, "{stats:?}");
let claimed = journal
.claim(
&billing,
claim_at(8, Duration::from_secs(30), SystemTime::now()),
)
.await
.expect("claim");
assert_eq!(claimed.len(), 1);
assert_eq!(claimed[0].event.id(), kept.id());
}
#[tokio::test]
async fn acknowledging_an_event_the_outbox_no_longer_holds_is_not_an_error() {
let Some((_, journal)) =
outbox("acked_gone", capacity(1, CapacityPolicy::DropOldest)).await
else {
return;
};
journal
.append(&event_for("GW_INBOUND_ACME_KEY"))
.await
.expect("append");
let billing = consumer("billing");
let claimed = journal
.claim(
&billing,
claim_at(1, Duration::from_secs(30), SystemTime::now()),
)
.await
.expect("claim");
journal
.append(&event_for("GW_INBOUND_ACME_KEY"))
.await
.expect("a lossy policy makes room rather than refusing");
journal
.ack(&claimed[0].id)
.await
.expect("nothing is left to redeliver, so this is not the worker's problem");
let error = journal
.quarantine(&claimed[0].id, PoisonReason::Malformed)
.await
.expect_err("a row that is gone cannot be set aside for an operator");
assert!(
matches!(error, JournalError::NotOutstanding { .. }),
"{error:?}"
);
}
#[tokio::test]
async fn acknowledging_a_claim_in_one_statement_still_answers_for_each_event() {
let Some((_, journal)) = outbox("ack_all", capacity(16, CapacityPolicy::Refuse)).await
else {
return;
};
let billing = consumer("billing");
for subject in ["acme", "globex", "initech", "umbrella"] {
journal.append(&event_for(subject)).await.expect("append");
}
let claimed = journal
.claim(
&billing,
claim_at(8, Duration::from_secs(30), SystemTime::now()),
)
.await
.expect("claim");
assert_eq!(claimed.len(), 4, "one event per ordering key");
journal
.quarantine(&claimed[3].id, PoisonReason::Malformed)
.await
.expect("quarantine");
let stray = DeliveryId {
consumer: consumer("audit"),
event: claimed[0].event.id(),
attempt: 1,
};
let ids: Vec<DeliveryId> = claimed
.iter()
.map(|delivery| delivery.id.clone())
.chain([stray])
.collect();
let verdicts = journal.ack_all(&ids).await;
assert!(verdicts[0].is_ok() && verdicts[1].is_ok() && verdicts[2].is_ok());
assert!(
matches!(verdicts[3], Err(JournalError::Quarantined { .. })),
"{:?}",
verdicts[3]
);
assert!(
matches!(verdicts[4], Err(JournalError::NotOutstanding { .. })),
"{:?}",
verdicts[4]
);
let repeated = journal.ack_all(&ids).await;
assert!(repeated[0].is_ok() && repeated[1].is_ok() && repeated[2].is_ok());
assert!(matches!(repeated[3], Err(JournalError::Quarantined { .. })));
let stats = journal.stats(&billing).await.expect("stats");
assert_eq!(stats.pending, 0, "{stats:?}");
assert_eq!(stats.quarantined, 1, "{stats:?}");
}
#[tokio::test]
async fn an_outbox_far_over_a_lowered_limit_is_brought_under_it_by_one_append() {
let Some((dsn, journal)) = outbox("surplus", capacity(8, CapacityPolicy::DropOldest)).await
else {
return;
};
for _ in 0..8 {
journal
.append(&event_for("GW_INBOUND_ACME_KEY"))
.await
.expect("append");
}
let lowered = PostgresJournal::connect(
&dsn,
settings(
"axond_outbox_surplus",
false,
capacity(2, CapacityPolicy::DropOldest),
),
)
.await
.expect("connect");
lowered
.append(&event_for("GW_INBOUND_ACME_KEY"))
.await
.expect("append");
let stored = client(&dsn, Some("axond_outbox_surplus"))
.await
.query_one("SELECT count(*) FROM axond_usage_outbox", &[])
.await
.expect("count")
.get::<_, i64>(0);
assert!(stored <= 2, "the outbox is inside its limit: {stored}");
}
#[tokio::test]
async fn a_delivered_event_yields_its_retention_window_before_an_append_is_refused() {
let Some((_, journal)) = outbox("reclaim", capacity(1, CapacityPolicy::Refuse)).await
else {
return;
};
let delivered = event_for("GW_INBOUND_ACME_KEY");
journal.append(&delivered).await.expect("append");
let billing = consumer("billing");
let claimed = journal
.claim(
&billing,
claim_at(8, Duration::from_secs(30), SystemTime::now()),
)
.await
.expect("claim");
journal.ack(&claimed[0].id).await.expect("ack");
let next = event_for("GW_INBOUND_ACME_KEY");
assert!(journal.append(&next).await.expect("append").is_new());
let stats = journal.stats(&billing).await.expect("stats");
assert_eq!(stats.pending, 1, "{stats:?}");
assert_eq!(stats.dropped, 0, "a reclaim is not a loss: {stats:?}");
}
#[test]
fn a_refusal_that_freed_nothing_is_remembered_until_something_is_deleted() {
let gate = CapacityGate::new();
assert!(!gate.refusing(), "an outbox with room refuses nothing");
gate.unreclaimable();
assert!(gate.refusing(), "the backlog cannot have changed by itself");
gate.invalidate();
assert!(
!gate.refusing(),
"a deletion made room, so the next append has to look again"
);
}
#[test]
fn a_gap_count_is_not_reused_against_a_span_that_has_since_collapsed() {
let gate = CapacityGate::new();
gate.measured(1_000, 900, 1_000);
assert_eq!(
gate.estimate(1_000, 1_000),
Some(900),
"the span it was measured against is exactly what it describes"
);
assert_eq!(
gate.estimate(1_100, 1_000),
Some(1_000),
"appends extend the span, and every one of them stored a row"
);
assert_eq!(
gate.estimate(899, 1_000),
None,
"the span collapsed below the measurement, so the gaps say nothing \
and the append has to count"
);
}
#[tokio::test]
async fn the_room_a_refused_append_made_is_kept_rather_than_rolled_back() {
let Some((dsn, journal)) =
outbox("refused_reclaim", capacity(8, CapacityPolicy::Refuse)).await
else {
return;
};
for _ in 0..8 {
journal
.append(&event_for("GW_INBOUND_ACME_KEY"))
.await
.expect("append");
}
let billing = consumer("billing");
let claimed = journal
.claim(
&billing,
claim_at(1, Duration::from_secs(30), SystemTime::now()),
)
.await
.expect("claim");
journal.ack(&claimed[0].id).await.expect("ack");
let lowered = PostgresJournal::connect(
&dsn,
settings(
"axond_outbox_refused_reclaim",
false,
capacity(2, CapacityPolicy::Refuse),
),
)
.await
.expect("connect");
let error = lowered
.append(&event_for("GW_INBOUND_ACME_KEY"))
.await
.expect_err("an outbox of undelivered events refuses");
assert!(
matches!(error, JournalError::AtCapacity { .. }),
"{error:?}"
);
let stored = client(&dsn, Some("axond_outbox_refused_reclaim"))
.await
.query_one("SELECT count(*) FROM axond_usage_outbox", &[])
.await
.expect("count")
.get::<_, i64>(0);
assert_eq!(
stored, 7,
"the reclaim rode out on the refusal instead of being rolled back \
for the next request to redo"
);
}
#[tokio::test]
async fn a_consumer_this_deployment_is_not_running_is_reported_not_deleted() {
let Some((dsn, journal)) = outbox("others", capacity(16, CapacityPolicy::Refuse)).await
else {
return;
};
let mine = consumer("billing");
assert!(
journal
.consumers_besides(&mine)
.await
.expect("consumers")
.is_empty(),
"nothing has claimed yet"
);
journal
.append(&event_for("GW_INBOUND_ACME_KEY"))
.await
.expect("append");
let now = SystemTime::now();
for name in ["billing", "retired"] {
journal
.claim(&consumer(name), claim_at(8, Duration::from_secs(30), now))
.await
.expect("claim");
}
assert_eq!(
journal.consumers_besides(&mine).await.expect("consumers"),
vec!["retired".to_owned()],
"the name retention is also waiting on"
);
let registered: i64 = client(&dsn, Some("axond_outbox_others"))
.await
.query_one("SELECT count(*) FROM axond_usage_outbox_consumer", &[])
.await
.expect("count")
.get(0);
assert_eq!(registered, 2);
}
#[tokio::test]
async fn retention_prunes_only_what_every_consumer_finished_with() {
let Some((dsn, journal)) = outbox(
"retention",
Capacity {
retain_acknowledged: Duration::ZERO,
..capacity(16, CapacityPolicy::Refuse)
},
)
.await
else {
return;
};
let acknowledged = event_for("GW_INBOUND_ACME_KEY");
let pending = event_for("GW_INBOUND_OTHER_KEY");
journal.append(&acknowledged).await.expect("append");
journal.append(&pending).await.expect("append");
let billing = consumer("billing");
let claimed = journal
.claim(
&billing,
claim_at(8, Duration::from_secs(30), SystemTime::now()),
)
.await
.expect("claim");
let first = claimed
.iter()
.find(|delivery| delivery.event.id() == acknowledged.id())
.expect("the acknowledged event was claimed");
journal.ack(&first.id).await.expect("ack");
assert_eq!(journal.maintain(SystemTime::now()).await.expect("prune"), 1);
let rows: i64 = client(&dsn, Some("axond_outbox_retention"))
.await
.query_one("SELECT count(*) FROM axond_usage_outbox", &[])
.await
.expect("count")
.get(0);
assert_eq!(rows, 1, "only the delivered event was pruned");
let stats = journal.stats(&billing).await.expect("stats");
assert_eq!(stats.in_flight + stats.pending, 1, "{stats:?}");
}
#[tokio::test]
async fn a_row_a_newer_build_wrote_is_left_for_that_build_to_deliver() {
let Some((dsn, journal)) =
outbox("schema_ahead", capacity(16, CapacityPolicy::Refuse)).await
else {
return;
};
let event = event_for("GW_INBOUND_ACME_KEY");
journal.append(&event).await.expect("append");
let ahead = i32::try_from(UsageRecord::SCHEMA_VERSION).expect("a small version") + 1;
client(&dsn, Some("axond_outbox_schema_ahead"))
.await
.execute(
"UPDATE axond_usage_outbox SET schema_version = $1",
&[&ahead],
)
.await
.expect("the row a rolling upgrade's newer replica wrote");
let billing = consumer("billing");
assert!(
journal
.claim(
&billing,
claim_at(8, Duration::from_secs(30), SystemTime::now())
)
.await
.expect("claim")
.is_empty()
);
let stats = journal.stats(&billing).await.expect("stats");
assert_eq!(stats.quarantined, 0, "{stats:?}");
assert_eq!(stats.pending, 1, "{stats:?}");
}
#[tokio::test]
async fn a_row_this_build_cannot_decode_is_quarantined_rather_than_blocking_its_key() {
let Some((dsn, journal)) = outbox("corrupt", capacity(16, CapacityPolicy::Refuse)).await
else {
return;
};
let corrupt = event_for("GW_INBOUND_ACME_KEY");
let behind = event_for("GW_INBOUND_ACME_KEY");
journal.append(&corrupt).await.expect("append");
journal.append(&behind).await.expect("append");
client(&dsn, Some("axond_outbox_corrupt"))
.await
.execute(
"UPDATE axond_usage_outbox \
SET record = jsonb_set(record, '{credential_source}', '\"nonsense\"') \
WHERE request_id = $1",
&[&corrupt.id().to_string()],
)
.await
.expect("corruption at this build's own version");
let billing = consumer("billing");
let claimed = journal
.claim(
&billing,
claim_at(8, Duration::from_secs(30), SystemTime::now()),
)
.await
.expect("claim");
assert_eq!(claimed.len(), 1, "{claimed:?}");
assert_eq!(claimed[0].event.id(), behind.id());
let stats = journal.stats(&billing).await.expect("stats");
assert_eq!(stats.quarantined, 1, "{stats:?}");
let rows: i64 = client(&dsn, Some("axond_outbox_corrupt"))
.await
.query_one("SELECT count(*) FROM axond_usage_outbox", &[])
.await
.expect("count")
.get(0);
assert_eq!(rows, 2);
}
#[tokio::test]
async fn consumers_acknowledge_independently() {
let Some((_, journal)) = outbox("consumers", capacity(16, CapacityPolicy::Refuse)).await
else {
return;
};
let event = event_for("GW_INBOUND_ACME_KEY");
journal.append(&event).await.expect("append");
let (billing, warehouse) = (consumer("billing"), consumer("warehouse"));
let now = SystemTime::now();
let claimed = journal
.claim(&billing, claim_at(8, Duration::from_secs(30), now))
.await
.expect("claim");
journal.ack(&claimed[0].id).await.expect("ack");
let other = journal
.claim(&warehouse, claim_at(8, Duration::from_secs(30), now))
.await
.expect("claim");
assert_eq!(other.len(), 1, "{other:?}");
assert_eq!(other[0].event.id(), event.id());
let stats = journal.stats(&billing).await.expect("stats");
assert!(stats.is_drained(), "{stats:?}");
let stats = journal.stats(&warehouse).await.expect("stats");
assert_eq!(stats.in_flight, 1, "{stats:?}");
}
#[tokio::test]
async fn stats_report_the_backlog_and_the_bound_it_is_measured_against() {
let Some((_, journal)) = outbox("stats", capacity(16, CapacityPolicy::Refuse)).await else {
return;
};
let mut record = event_for("GW_INBOUND_ACME_KEY").record().clone();
record.request_id = crate::usage::identity::next_request_id().to_string();
let old = UsageEvent::new(ObservedRecord {
record,
observed_at: SystemTime::now() - Duration::from_secs(120),
})
.expect("a well-formed event");
journal.append(&old).await.expect("append");
let stats = journal.stats(&consumer("billing")).await.expect("stats");
assert_eq!(stats.pending, 1, "{stats:?}");
assert_eq!(stats.capacity, journal.capacity());
let age = stats.oldest_pending_age.expect("an age");
assert!(age >= Duration::from_secs(100), "{age:?}");
}
#[test]
fn a_schema_ahead_row_is_reported_once_per_claim() {
let mut undeliverable = Undeliverable::default();
undeliverable.schema_ahead(7);
undeliverable.schema_ahead(7);
undeliverable.schema_ahead(9);
undeliverable.corrupt();
assert_eq!(
undeliverable.reasons,
vec!["schema_ahead", "schema_ahead", "corrupt"],
"one report per row, and corruption is never re-selected"
);
}
#[tokio::test]
async fn a_full_outbox_refuses_from_a_bounded_measurement_and_recovers() {
let Some((dsn, journal)) = outbox("near_full", capacity(2, CapacityPolicy::Refuse)).await
else {
return;
};
journal
.append(&event_for("GW_INBOUND_ACME_KEY"))
.await
.expect("append");
journal
.append(&event_for("GW_INBOUND_ACME_KEY"))
.await
.expect("the outbox is full only after this one");
for _ in 0..3 {
let error = journal
.append(&event_for("GW_INBOUND_ACME_KEY"))
.await
.expect_err("a full billing-grade outbox refuses");
assert!(
matches!(error, JournalError::AtCapacity { pending, .. } if pending >= 2),
"{error:?}"
);
}
client(&dsn, Some("axond_outbox_near_full"))
.await
.execute(
"DELETE FROM axond_usage_outbox WHERE position = \
(SELECT min(position) FROM axond_usage_outbox)",
&[],
)
.await
.expect("make room");
tokio::time::sleep(COUNT_REFRESH + Duration::from_millis(50)).await;
assert!(
journal
.append(&event_for("GW_INBOUND_ACME_KEY"))
.await
.expect("an outbox with room accepts")
.is_new()
);
}
#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn two_replicas_cannot_claim_the_same_event() {
let Some((dsn, first)) = outbox("concurrent", capacity(64, CapacityPolicy::Refuse)).await
else {
return;
};
let second = PostgresJournal::connect(
&dsn,
settings(
"axond_outbox_concurrent",
false,
capacity(64, CapacityPolicy::Refuse),
),
)
.await
.expect("a second replica");
let billing = consumer("billing");
let mut in_flight = 0u64;
for _ in 0..8 {
for key in ["GW_INBOUND_ACME_KEY", "GW_INBOUND_OTHER_KEY"] {
first.append(&event_for(key)).await.expect("append");
}
let now = SystemTime::now();
let (left, right) = tokio::join!(
first.claim(&billing, claim_at(8, Duration::from_secs(30), now)),
second.claim(&billing, claim_at(8, Duration::from_secs(30), now)),
);
let mut delivered: Vec<_> = left
.expect("claim")
.into_iter()
.chain(right.expect("claim"))
.map(|delivery| delivery.event.id())
.collect();
let claimed = delivered.len();
delivered.sort();
delivered.dedup();
assert_eq!(
delivered.len(),
claimed,
"an event was handed to both replicas at once"
);
in_flight += claimed as u64;
let stats = first.stats(&billing).await.expect("stats");
assert_eq!(stats.in_flight, in_flight, "{stats:?}");
}
}
#[tokio::test]
async fn maintenance_moves_the_claim_floor_past_the_acknowledged_prefix() {
let Some((dsn, journal)) = outbox("floor", capacity(64, CapacityPolicy::Refuse)).await
else {
return;
};
let billing = consumer("billing");
let now = SystemTime::now();
for _ in 0..4 {
journal
.append(&event_for("GW_INBOUND_ACME_KEY"))
.await
.expect("append");
let claimed = journal
.claim(&billing, claim_at(1, Duration::from_secs(30), now))
.await
.expect("claim");
journal.ack(&claimed[0].id).await.expect("ack");
}
let pending = event_for("GW_INBOUND_ACME_KEY");
journal.append(&pending).await.expect("append");
let settled = now + FLOOR_SETTLE_MARGIN + Duration::from_secs(1);
assert_eq!(journal.maintain(settled).await.expect("maintain"), 0);
let admin = client(&dsn, Some("axond_outbox_floor")).await;
let floor: i64 = admin
.query_one(
"SELECT resolved_through FROM axond_usage_outbox_consumer WHERE consumer = $1",
&[&billing.as_str()],
)
.await
.expect("the floor")
.get(0);
let first_open: i64 = admin
.query_one(
"SELECT min(e.position) FROM axond_usage_outbox e \
LEFT JOIN axond_usage_outbox_delivery d \
ON d.position = e.position AND d.consumer = $1 \
WHERE d.acknowledged_at IS NULL AND d.quarantined_at IS NULL",
&[&billing.as_str()],
)
.await
.expect("the first unresolved position")
.get(0);
assert_eq!(
floor,
first_open - 1,
"the floor sits just below the oldest event still to deliver"
);
let claimed = journal
.claim(&billing, claim_at(8, Duration::from_secs(30), now))
.await
.expect("claim");
assert_eq!(claimed.len(), 1, "{claimed:?}");
assert_eq!(claimed[0].event.id(), pending.id());
}
#[tokio::test]
async fn the_event_limit_holds_across_replicas() {
const LIMIT: u64 = 6;
let Some((dsn, first)) =
outbox("replica_capacity", capacity(LIMIT, CapacityPolicy::Refuse)).await
else {
return;
};
let second = PostgresJournal::connect(
&dsn,
settings(
"axond_outbox_replica_capacity",
false,
capacity(LIMIT, CapacityPolicy::Refuse),
),
)
.await
.expect("a second replica");
let mut accepted = 0u64;
let mut refused = 0u64;
for turn in 0..LIMIT * 2 {
let journal: &PostgresJournal = if turn % 2 == 0 { &first } else { &second };
match journal.append(&event_for("GW_INBOUND_ACME_KEY")).await {
Ok(_) => accepted += 1,
Err(JournalError::AtCapacity { .. }) => refused += 1,
Err(error) => panic!("unexpected append failure: {error}"),
}
}
assert_eq!(accepted, LIMIT, "the two replicas together overfilled it");
assert_eq!(refused, LIMIT, "the rest were refused, not lost");
let admin = client(&dsn, Some("axond_outbox_replica_capacity")).await;
let stored: i64 = admin
.query_one("SELECT count(*) FROM axond_usage_outbox", &[])
.await
.expect("the stored count")
.get(0);
assert_eq!(stored as u64, LIMIT, "the outbox holds more than its bound");
}
#[tokio::test]
async fn the_claim_floor_never_passes_an_append_that_has_not_committed() {
let Some((dsn, journal)) = outbox("floor_race", capacity(64, CapacityPolicy::Refuse)).await
else {
return;
};
let billing = consumer("billing");
let now = SystemTime::now();
let mut inflight = client(&dsn, Some("axond_outbox_floor_race")).await;
let held = inflight.transaction().await.expect("a held append");
let early = event_for("GW_INBOUND_ACME_KEY");
let record = serde_json::to_value(early.record()).expect("a record");
let position: i64 = held
.query_one(
"INSERT INTO axond_usage_outbox \
(request_id, schema_version, namespace, subject, record, observed_at) \
VALUES ($1, $2, $3, $4, $5::jsonb, $6) RETURNING position",
&[
&early.idempotency_key().as_str(),
&i32::try_from(early.record().schema_version).expect("a version"),
&early.ordering_key().namespace,
&early.ordering_key().subject,
&record,
&early.observed_at(),
],
)
.await
.expect("the held append takes a position")
.get(0);
let later = event_for("GW_INBOUND_OTHER_KEY");
journal.append(&later).await.expect("append");
let claimed = journal
.claim(&billing, claim_at(8, Duration::from_secs(30), now))
.await
.expect("claim");
assert_eq!(claimed.len(), 1, "{claimed:?}");
journal.ack(&claimed[0].id).await.expect("ack");
journal.maintain(now).await.expect("maintain");
let admin = client(&dsn, Some("axond_outbox_floor_race")).await;
let floor: i64 = admin
.query_one(
"SELECT resolved_through FROM axond_usage_outbox_consumer WHERE consumer = $1",
&[&billing.as_str()],
)
.await
.expect("the floor")
.get(0);
assert!(
floor < position,
"the floor ({floor}) passed a position ({position}) that had not committed"
);
held.commit().await.expect("the held append commits");
let claimed = journal
.claim(&billing, claim_at(8, Duration::from_secs(30), now))
.await
.expect("claim");
assert_eq!(
claimed.iter().map(|d| d.event.id()).collect::<Vec<_>>(),
vec![early.id()],
"the late-committing event is still claimable"
);
}
}