Files
Aether/crates/aether-usage/runtime/src/worker.rs
T

1600 lines
55 KiB
Rust
Raw Normal View History

2026-06-26 14:02:57 +08:00
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{Arc, OnceLock};
use std::time::Duration;
use aether_data_contracts::repository::usage::{StoredRequestUsageAudit, UpsertUsageRecord};
use aether_data_contracts::DataLayerError;
2026-05-08 00:18:12 +08:00
use aether_runtime_state::{RuntimeQueueEntry, RuntimeQueueStore};
use async_trait::async_trait;
use tokio::sync::{mpsc, Notify};
use tracing::warn;
use crate::executor::spawn_on_usage_background_runtime;
use crate::keyed_lock::KeyedAsyncLockPool;
use crate::runtime::{
UsageBillingEventEnricher, UsageRuntimeAccess, UsageWorkerRecordConcurrencyGate,
};
use crate::{
build_upsert_usage_record_from_event, reconcile_usage_policy_cost_for_event,
settle_usage_if_needed, UsageEvent, UsageEventType, UsageQueue, UsageRuntimeConfig,
UsageSettlementWriter,
};
const USAGE_WORKER_DB_PRESSURE_DEFER_MS: u64 = 10;
const USAGE_WORKER_ACK_CHUNK_SIZE: usize = 100;
#[async_trait]
pub trait UsageEventRecorder: Send + Sync {
async fn record_usage_event(&self, event: &UsageEvent) -> Result<(), DataLayerError>;
}
#[async_trait]
pub trait ManualProxyNodeCounter: Send + Sync {
async fn increment_manual_proxy_node_requests(
&self,
node_id: &str,
total_delta: i64,
failed_delta: i64,
latency_ms: Option<i64>,
) -> Result<(), DataLayerError>;
}
#[async_trait]
pub trait UsageRecordWriter: Send + Sync {
/// Native batch support is opt-in; the default preserves one-row writes for other backends.
fn supports_first_byte_usage_batch(&self) -> bool {
false
}
/// Stable identity for the underlying first-byte writer. Implementations that opt into
/// batching must return the same value for clones backed by the same repository.
fn first_byte_usage_writer_identity(&self) -> Option<usize> {
None
}
/// Native pending batching is opt-in because it must retain the complete usage audit write
/// contract, not just the base lifecycle row.
fn supports_pending_usage_batch(&self) -> bool {
false
}
/// Stable identity for clones backed by the same pending usage repository.
fn pending_usage_writer_identity(&self) -> Option<usize> {
None
}
async fn upsert_usage_record(
&self,
record: UpsertUsageRecord,
) -> Result<Option<StoredRequestUsageAudit>, DataLayerError>;
async fn upsert_first_byte_usage_record(
&self,
record: UpsertUsageRecord,
) -> Result<(), DataLayerError> {
self.upsert_usage_record(record).await.map(|_| ())
}
async fn upsert_first_byte_usage_records(
&self,
records: Vec<UpsertUsageRecord>,
) -> Result<(), DataLayerError> {
for record in records {
self.upsert_first_byte_usage_record(record).await?;
}
Ok(())
}
async fn upsert_pending_usage_records(
&self,
records: Vec<UpsertUsageRecord>,
) -> Result<(), DataLayerError> {
for record in records {
self.upsert_usage_record(record).await?;
}
Ok(())
}
}
pub struct UsageDataEventRecorder<T> {
data: Arc<T>,
record_gate: Option<Arc<UsageWorkerRecordConcurrencyGate>>,
defer_for_database_pressure: bool,
}
impl<T> UsageDataEventRecorder<T> {
pub fn new(data: Arc<T>) -> Self {
Self::with_record_gate(data, None)
}
pub(crate) fn with_record_gate(
data: Arc<T>,
record_gate: Option<Arc<UsageWorkerRecordConcurrencyGate>>,
) -> Self {
Self {
data,
record_gate,
defer_for_database_pressure: false,
}
}
pub(crate) fn with_record_gate_and_database_pressure_defer(
data: Arc<T>,
record_gate: Option<Arc<UsageWorkerRecordConcurrencyGate>>,
) -> Self {
Self {
data,
record_gate,
defer_for_database_pressure: true,
}
}
}
#[async_trait]
impl<T> UsageEventRecorder for UsageDataEventRecorder<T>
where
T: UsageRuntimeAccess,
{
async fn record_usage_event(&self, event: &UsageEvent) -> Result<(), DataLayerError> {
if self.defer_for_database_pressure
&& self.data.usage_worker_should_defer_for_database_pressure()
{
if let Some(gate) = self.record_gate.as_ref() {
gate.record_deferred();
}
tokio::time::sleep(Duration::from_millis(USAGE_WORKER_DB_PRESSURE_DEFER_MS)).await;
}
let _record_gate_permit = match self.record_gate.as_ref() {
Some(gate) => Some(gate.acquire().await),
None => None,
};
let request_lock = usage_request_lock(&event.request_id);
let _guard = request_lock.lock().await;
let mut event = event.clone();
enrich_terminal_event(self.data.as_ref(), &mut event).await;
write_event_record(self.data.as_ref(), &event).await
}
}
fn usage_request_lock(request_id: &str) -> Arc<tokio::sync::Mutex<()>> {
static LOCKS: OnceLock<KeyedAsyncLockPool> = OnceLock::new();
LOCKS
.get_or_init(KeyedAsyncLockPool::default)
.lock_for(request_id)
}
pub struct UsageQueueWorker {
queue: UsageQueue,
recorder: Arc<dyn UsageEventRecorder>,
2026-05-08 00:18:12 +08:00
consumer: String,
2026-06-26 14:02:57 +08:00
worker_index: Option<usize>,
control: Option<UsageWorkerControl>,
telemetry: Option<mpsc::Sender<UsageWorkerObservation>>,
config: UsageRuntimeConfig,
}
2026-06-26 14:02:57 +08:00
#[derive(Clone, Default)]
pub(crate) struct UsageWorkerControl {
shutdown: Arc<AtomicBool>,
shutdown_notify: Arc<Notify>,
2026-06-26 14:02:57 +08:00
}
impl UsageWorkerControl {
pub(crate) fn request_shutdown(&self) {
self.shutdown.store(true, Ordering::Release);
self.shutdown_notify.notify_one();
2026-06-26 14:02:57 +08:00
}
fn should_shutdown(&self) -> bool {
self.shutdown.load(Ordering::Acquire)
}
async fn wait_for_shutdown(&self) {
while !self.should_shutdown() {
self.shutdown_notify.notified().await;
}
}
2026-06-26 14:02:57 +08:00
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) struct UsageWorkerObservation {
pub worker_index: Option<usize>,
pub entries_read: usize,
pub batch_size: usize,
pub reclaimed_entries: usize,
pub acked_entries: usize,
pub dead_lettered_entries: usize,
pub process_failures: usize,
pub read_failures: usize,
pub reclaim_failures: usize,
}
impl UsageWorkerObservation {
fn read(worker_index: Option<usize>, entries_read: usize, batch_size: usize) -> Self {
Self {
worker_index,
entries_read,
batch_size,
reclaimed_entries: 0,
acked_entries: 0,
dead_lettered_entries: 0,
process_failures: 0,
read_failures: 0,
reclaim_failures: 0,
}
}
fn reclaimed(worker_index: Option<usize>, reclaimed_entries: usize) -> Self {
Self {
worker_index,
entries_read: 0,
batch_size: 0,
reclaimed_entries,
acked_entries: 0,
dead_lettered_entries: 0,
process_failures: 0,
read_failures: 0,
reclaim_failures: 0,
}
}
fn acked(worker_index: Option<usize>, acked_entries: usize) -> Self {
Self {
worker_index,
entries_read: 0,
batch_size: 0,
reclaimed_entries: 0,
acked_entries,
dead_lettered_entries: 0,
process_failures: 0,
read_failures: 0,
reclaim_failures: 0,
}
}
fn dead_lettered(worker_index: Option<usize>, dead_lettered_entries: usize) -> Self {
Self {
worker_index,
entries_read: 0,
batch_size: 0,
reclaimed_entries: 0,
acked_entries: 0,
dead_lettered_entries,
process_failures: 0,
read_failures: 0,
reclaim_failures: 0,
}
}
fn process_failed(worker_index: Option<usize>) -> Self {
Self {
worker_index,
entries_read: 0,
batch_size: 0,
reclaimed_entries: 0,
acked_entries: 0,
dead_lettered_entries: 0,
process_failures: 1,
read_failures: 0,
reclaim_failures: 0,
}
}
fn read_failed(worker_index: Option<usize>) -> Self {
Self {
worker_index,
entries_read: 0,
batch_size: 0,
reclaimed_entries: 0,
acked_entries: 0,
dead_lettered_entries: 0,
process_failures: 0,
read_failures: 1,
reclaim_failures: 0,
}
}
fn reclaim_failed(worker_index: Option<usize>) -> Self {
Self {
worker_index,
entries_read: 0,
batch_size: 0,
reclaimed_entries: 0,
acked_entries: 0,
dead_lettered_entries: 0,
process_failures: 0,
read_failures: 0,
reclaim_failures: 1,
}
}
2026-06-26 14:02:57 +08:00
}
impl UsageQueueWorker {
pub fn new(
2026-05-08 00:18:12 +08:00
runner: Arc<dyn RuntimeQueueStore>,
recorder: Arc<dyn UsageEventRecorder>,
config: UsageRuntimeConfig,
2026-06-26 14:02:57 +08:00
worker_index: Option<usize>,
) -> Result<Self, DataLayerError> {
let queue = UsageQueue::new(runner, config.clone())?;
2026-06-26 14:02:57 +08:00
let consumer = consumer_name(worker_index);
Ok(Self {
queue,
recorder,
consumer,
2026-06-26 14:02:57 +08:00
worker_index,
control: None,
telemetry: None,
config,
})
}
2026-06-26 14:02:57 +08:00
pub(crate) fn with_supervisor(
mut self,
control: UsageWorkerControl,
telemetry: mpsc::Sender<UsageWorkerObservation>,
) -> Self {
self.control = Some(control);
self.telemetry = Some(telemetry);
self
}
pub fn spawn(self) -> tokio::task::JoinHandle<()> {
spawn_on_usage_background_runtime(async move { self.run_forever().await })
}
2026-06-26 14:02:57 +08:00
pub(crate) async fn run(self) {
self.run_forever().await;
}
async fn run_forever(self) {
if let Err(err) = self.queue.ensure_consumer_group().await {
warn!(
event_name = "usage_worker_consumer_group_failed",
log_type = "ops",
2026-05-08 00:18:12 +08:00
worker_consumer = %self.consumer,
worker_group = %self.config.consumer_group,
error = %err,
"usage worker failed to ensure consumer group"
);
return;
}
let mut reclaim_interval =
tokio::time::interval(Duration::from_millis(self.config.reclaim_interval_ms));
reclaim_interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
reclaim_interval.tick().await;
let mut reclaim_due = false;
loop {
2026-06-26 14:02:57 +08:00
if self.should_shutdown() {
break;
}
let result = {
let mut read_future = Box::pin(self.queue.read_group(&self.consumer));
loop {
tokio::select! {
biased;
// A command already delivered by Redis remains in the PEL and is recovered
// by a subsequent worker reclaim after the idle period.
_ = self.wait_for_shutdown() => return,
_ = reclaim_interval.tick(), if !reclaim_due => {
// Do not reclaim while XREADGROUP is in flight. Redis can add an entry
// to this consumer's PEL before delivering the response; claiming it in
// that window would make both paths process the same stream entry.
reclaim_due = true;
}
result = &mut read_future => break result,
}
}
};
match result {
Ok(entries) => {
self.report_read(entries.len());
if let Err(err) = self.process_entries(entries).await {
self.report_process_failed();
warn!(
event_name = "usage_worker_process_failed",
log_type = "ops",
worker_consumer = %self.consumer,
worker_group = %self.config.consumer_group,
error = %err,
"usage worker failed to process queue entries"
);
tokio::time::sleep(Duration::from_millis(250)).await;
}
}
Err(err) => {
self.report_read_failed();
warn!(
event_name = "usage_worker_read_failed",
log_type = "ops",
worker_consumer = %self.consumer,
worker_group = %self.config.consumer_group,
error = %err,
"usage worker failed to read queue"
);
tokio::time::sleep(Duration::from_millis(500)).await;
}
}
if self.should_shutdown() {
break;
}
if reclaim_due {
reclaim_due = false;
self.reclaim_stale_entries().await;
}
}
}
async fn wait_for_shutdown(&self) {
match self.control.as_ref() {
Some(control) => control.wait_for_shutdown().await,
None => std::future::pending().await,
}
}
async fn reclaim_stale_entries(&self) {
match self.queue.claim_stale(&self.consumer, "0-0").await {
Ok(entries) => {
self.report_reclaimed(entries.len());
if let Err(err) = self.process_entries(entries).await {
self.report_process_failed();
warn!(
event_name = "usage_worker_reclaim_process_failed",
log_type = "ops",
worker_consumer = %self.consumer,
worker_group = %self.config.consumer_group,
error = %err,
"usage worker failed while reclaiming stale entries"
);
}
}
Err(err) => {
self.report_reclaim_failed();
warn!(
event_name = "usage_worker_reclaim_failed",
log_type = "ops",
worker_consumer = %self.consumer,
worker_group = %self.config.consumer_group,
error = %err,
"usage worker failed to reclaim stale entries"
);
}
}
}
2026-06-26 14:02:57 +08:00
fn should_shutdown(&self) -> bool {
self.control
.as_ref()
.is_some_and(UsageWorkerControl::should_shutdown)
}
fn report_read(&self, entries_read: usize) {
self.report(UsageWorkerObservation::read(
self.worker_index,
entries_read,
self.config.consumer_batch_size.max(1),
));
}
fn report_reclaimed(&self, reclaimed_entries: usize) {
self.report(UsageWorkerObservation::reclaimed(
self.worker_index,
reclaimed_entries,
));
}
fn report_acked(&self, acked_entries: usize) {
self.report(UsageWorkerObservation::acked(
self.worker_index,
acked_entries,
));
}
fn report_dead_lettered(&self, dead_lettered_entries: usize) {
self.report(UsageWorkerObservation::dead_lettered(
self.worker_index,
dead_lettered_entries,
));
}
fn report_process_failed(&self) {
self.report(UsageWorkerObservation::process_failed(self.worker_index));
}
fn report_read_failed(&self) {
self.report(UsageWorkerObservation::read_failed(self.worker_index));
}
fn report_reclaim_failed(&self) {
self.report(UsageWorkerObservation::reclaim_failed(self.worker_index));
}
fn report(&self, observation: UsageWorkerObservation) {
2026-06-26 14:02:57 +08:00
let Some(telemetry) = &self.telemetry else {
return;
};
let _ = telemetry.try_send(observation);
2026-06-26 14:02:57 +08:00
}
2026-05-08 00:18:12 +08:00
async fn process_entries(&self, entries: Vec<RuntimeQueueEntry>) -> Result<(), DataLayerError> {
if entries.is_empty() {
return Ok(());
}
let mut ack_ids = Vec::new();
for entry in entries {
match self.process_entry(&entry).await {
Ok(should_ack) => {
if should_ack {
ack_ids.push(entry.id.clone());
if ack_ids.len() >= USAGE_WORKER_ACK_CHUNK_SIZE {
self.queue.ack_and_delete(&ack_ids).await?;
self.report_acked(ack_ids.len());
ack_ids.clear();
}
}
}
Err(err) => {
if !ack_ids.is_empty() {
let _ = self.queue.ack_and_delete(&ack_ids).await;
}
return Err(err);
}
}
}
if !ack_ids.is_empty() {
self.queue.ack_and_delete(&ack_ids).await?;
self.report_acked(ack_ids.len());
}
Ok(())
}
2026-05-08 00:18:12 +08:00
async fn process_entry(&self, entry: &RuntimeQueueEntry) -> Result<bool, DataLayerError> {
let event = match UsageEvent::from_stream_fields(&entry.fields) {
Ok(event) => event,
Err(err) => {
2026-05-22 23:57:38 +08:00
warn!(
event_name = "usage_worker_entry_decode_dead_lettered",
log_type = "ops",
worker_consumer = %self.consumer,
worker_group = %self.config.consumer_group,
entry_id = %entry.id,
error = %err,
"usage worker moved malformed queue entry to dead letter"
);
self.queue.push_dead_letter(entry, &err.to_string()).await?;
self.report_dead_lettered(1);
return Ok(true);
}
};
2026-05-22 23:57:38 +08:00
match self.recorder.record_usage_event(&event).await {
Ok(()) => Ok(true),
Err(err) if usage_event_record_error_is_permanent(&err) => {
warn!(
event_name = "usage_worker_entry_record_dead_lettered",
log_type = "ops",
worker_consumer = %self.consumer,
worker_group = %self.config.consumer_group,
entry_id = %entry.id,
request_id = %event.request_id,
event_type = ?event.event_type,
provider_name = %event.data.provider_name,
model = %event.data.model,
api_format = event.data.api_format.as_deref().unwrap_or(""),
provider_id = event.data.provider_id.as_deref().unwrap_or(""),
provider_endpoint_id = event.data.provider_endpoint_id.as_deref().unwrap_or(""),
provider_api_key_id = event.data.provider_api_key_id.as_deref().unwrap_or(""),
error = %err,
"usage worker moved non-retryable usage event to dead letter"
);
self.queue.push_dead_letter(entry, &err.to_string()).await?;
self.report_dead_lettered(1);
2026-05-22 23:57:38 +08:00
Ok(true)
}
Err(err) => {
warn!(
event_name = "usage_worker_entry_record_retryable_failed",
log_type = "ops",
worker_consumer = %self.consumer,
worker_group = %self.config.consumer_group,
entry_id = %entry.id,
request_id = %event.request_id,
event_type = ?event.event_type,
provider_name = %event.data.provider_name,
model = %event.data.model,
api_format = event.data.api_format.as_deref().unwrap_or(""),
provider_id = event.data.provider_id.as_deref().unwrap_or(""),
provider_endpoint_id = event.data.provider_endpoint_id.as_deref().unwrap_or(""),
provider_api_key_id = event.data.provider_api_key_id.as_deref().unwrap_or(""),
error = %err,
"usage worker will retry usage event after record failure"
);
Err(err)
}
}
}
}
2026-05-22 23:57:38 +08:00
fn usage_event_record_error_is_permanent(err: &DataLayerError) -> bool {
match err {
DataLayerError::InvalidConfiguration(_)
| DataLayerError::InvalidInput(_)
| DataLayerError::UnexpectedValue(_) => true,
DataLayerError::Postgres(message) | DataLayerError::Sql(message) => {
database_error_is_known_permanent(message)
}
DataLayerError::Redis(_) | DataLayerError::TimedOut(_) => false,
}
}
fn database_error_is_known_permanent(message: &str) -> bool {
message.contains("SQLSTATE 23503") || message.contains("violates foreign key constraint")
}
pub fn build_usage_queue_worker<T>(
2026-05-08 00:18:12 +08:00
runner: Arc<dyn RuntimeQueueStore>,
data: Arc<T>,
config: UsageRuntimeConfig,
2026-06-26 14:02:57 +08:00
worker_index: Option<usize>,
) -> Result<UsageQueueWorker, DataLayerError>
where
T: UsageRuntimeAccess + 'static,
{
build_usage_queue_worker_with_record_gate(runner, data, config, None, worker_index)
}
pub(crate) fn build_usage_queue_worker_with_record_gate<T>(
runner: Arc<dyn RuntimeQueueStore>,
data: Arc<T>,
config: UsageRuntimeConfig,
record_gate: Option<Arc<UsageWorkerRecordConcurrencyGate>>,
worker_index: Option<usize>,
) -> Result<UsageQueueWorker, DataLayerError>
where
T: UsageRuntimeAccess + 'static,
{
2026-06-26 14:02:57 +08:00
UsageQueueWorker::new(
runner,
Arc::new(
UsageDataEventRecorder::with_record_gate_and_database_pressure_defer(data, record_gate),
),
2026-06-26 14:02:57 +08:00
config,
worker_index,
)
}
pub async fn write_event_record<T>(data: &T, event: &UsageEvent) -> Result<(), DataLayerError>
where
T: UsageRecordWriter + UsageSettlementWriter + Send + Sync,
{
reconcile_usage_policy_cost_for_event(data, event).await?;
let record = build_upsert_usage_record_from_event(event)?;
if let Some(stored) = data.upsert_usage_record(record).await? {
settle_usage_if_needed(data, &stored).await?;
}
// Manual proxy traffic is counted at the actual transport-attempt boundary. Usage events are
// replayable, so emitting that side effect here would count normal requests and reclaims twice.
Ok(())
}
async fn enrich_terminal_event<T>(data: &T, event: &mut UsageEvent)
where
T: UsageBillingEventEnricher + Send + Sync,
{
if !matches!(
event.event_type,
UsageEventType::Completed | UsageEventType::Failed | UsageEventType::Cancelled
) {
return;
}
if let Err(err) = data.enrich_usage_event(event).await {
warn!(
event_name = "usage_worker_billing_enrichment_failed",
log_type = "event",
request_id = %event.request_id,
event_type = ?event.event_type,
error = %err,
"usage worker failed to enrich terminal usage event with billing"
);
}
}
2026-06-26 14:02:57 +08:00
fn consumer_name(worker_index: Option<usize>) -> String {
let host = std::env::var("HOSTNAME")
.ok()
.map(|value| value.trim().to_string())
.filter(|value| !value.is_empty())
.unwrap_or_else(|| "aether-gateway".to_string());
2026-06-26 14:02:57 +08:00
match worker_index {
Some(worker_index) => format!("{host}:{}:{worker_index}", std::process::id()),
None => format!("{host}:{}", std::process::id()),
}
}
#[cfg(test)]
mod tests {
use std::collections::BTreeMap;
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
2026-05-22 23:57:38 +08:00
use std::sync::{Arc, Mutex};
use std::time::Duration;
use aether_data_contracts::repository::settlement::{
ReconcileUsagePolicyCostInput, StoredUsagePolicyCostReservation, StoredUsageSettlement,
UsageSettlementInput,
};
use aether_data_contracts::repository::usage::{StoredRequestUsageAudit, UpsertUsageRecord};
2026-05-22 23:57:38 +08:00
use aether_data_contracts::DataLayerError;
use aether_runtime_state::{
MemoryRuntimeStateConfig, RuntimeQueueEntry, RuntimeQueueReclaimConfig, RuntimeQueueStats,
RuntimeQueueStore, RuntimeState,
};
use async_trait::async_trait;
use tokio::sync::Notify;
2026-05-22 23:57:38 +08:00
use super::{
build_usage_queue_worker_with_record_gate, usage_event_record_error_is_permanent,
write_event_record, ManualProxyNodeCounter, UsageEventRecorder, UsageQueueWorker,
UsageRecordWriter, UsageWorkerControl,
2026-05-22 23:57:38 +08:00
};
use crate::runtime::UsageWorkerRecordConcurrencyGate;
use crate::UsageBillingEventEnricher;
2026-05-22 23:57:38 +08:00
use crate::{
UsageEvent, UsageEventData, UsageEventType, UsageQueue, UsageRuntimeConfig,
UsageSettlementWriter,
2026-05-22 23:57:38 +08:00
};
#[derive(Default)]
struct TestUsageStore {
records: Mutex<Vec<UpsertUsageRecord>>,
settlements: Mutex<Vec<UsageSettlementInput>>,
reconciliations: Mutex<Vec<ReconcileUsagePolicyCostInput>>,
enrich_calls: Mutex<Vec<String>>,
manual_proxy_counter_calls: AtomicUsize,
}
2026-05-22 23:57:38 +08:00
#[derive(Default)]
struct SelectiveFailingRecorder {
calls: Mutex<Vec<String>>,
}
#[derive(Default)]
struct SlowUsageStore {
active: std::sync::atomic::AtomicUsize,
max_active: std::sync::atomic::AtomicUsize,
db_pressure: AtomicBool,
records: Mutex<Vec<String>>,
}
struct ReadReclaimRaceProbeQueue {
entry: RuntimeQueueEntry,
read_calls: AtomicUsize,
first_read_cancelled: AtomicBool,
read_completed: AtomicUsize,
release_read: Notify,
reclaim_calls: AtomicUsize,
acked: AtomicBool,
}
impl ReadReclaimRaceProbeQueue {
fn new(entry: RuntimeQueueEntry) -> Self {
Self {
entry,
read_calls: AtomicUsize::new(0),
first_read_cancelled: AtomicBool::new(false),
read_completed: AtomicUsize::new(0),
release_read: Notify::new(),
reclaim_calls: AtomicUsize::new(0),
acked: AtomicBool::new(false),
}
}
}
struct FirstReadDropGuard<'a> {
cancelled: &'a AtomicBool,
completed: bool,
}
impl Drop for FirstReadDropGuard<'_> {
fn drop(&mut self) {
if !self.completed {
self.cancelled.store(true, Ordering::Release);
}
}
}
#[async_trait]
impl RuntimeQueueStore for ReadReclaimRaceProbeQueue {
async fn ensure_consumer_group(
&self,
_stream: &str,
_group: &str,
_start_id: &str,
) -> Result<(), DataLayerError> {
Ok(())
}
async fn append_fields_with_maxlen(
&self,
_stream: &str,
_fields: &BTreeMap<String, String>,
_maxlen: Option<usize>,
) -> Result<String, DataLayerError> {
Ok("0-0".to_string())
}
async fn read_group(
&self,
_stream: &str,
_group: &str,
_consumer: &str,
_count: usize,
_block_ms: Option<u64>,
) -> Result<Vec<RuntimeQueueEntry>, DataLayerError> {
let call_index = self.read_calls.fetch_add(1, Ordering::AcqRel);
let mut first_read_guard = (call_index == 0).then(|| FirstReadDropGuard {
cancelled: &self.first_read_cancelled,
completed: false,
});
self.release_read.notified().await;
if let Some(guard) = first_read_guard.as_mut() {
guard.completed = true;
}
self.read_completed.fetch_add(1, Ordering::AcqRel);
Ok((call_index == 0)
.then(|| self.entry.clone())
.into_iter()
.collect())
}
async fn claim_stale(
&self,
_stream: &str,
_group: &str,
_consumer: &str,
_start_id: &str,
_config: RuntimeQueueReclaimConfig,
) -> Result<Vec<RuntimeQueueEntry>, DataLayerError> {
self.reclaim_calls.fetch_add(1, Ordering::AcqRel);
Ok((!self.acked.load(Ordering::Acquire))
.then(|| self.entry.clone())
.into_iter()
.collect())
}
async fn ack(
&self,
_stream: &str,
_group: &str,
ids: &[String],
) -> Result<usize, DataLayerError> {
if ids.iter().any(|id| id == &self.entry.id) {
self.acked.store(true, Ordering::Release);
Ok(1)
} else {
Ok(0)
}
}
async fn delete(&self, _stream: &str, ids: &[String]) -> Result<usize, DataLayerError> {
Ok(ids.len())
}
async fn stats(
&self,
_stream: &str,
_group: Option<&str>,
) -> Result<RuntimeQueueStats, DataLayerError> {
Ok(RuntimeQueueStats::default())
}
}
#[async_trait]
impl UsageRecordWriter for TestUsageStore {
async fn upsert_usage_record(
&self,
record: UpsertUsageRecord,
) -> Result<Option<StoredRequestUsageAudit>, aether_data_contracts::DataLayerError>
{
self.records
.lock()
.expect("records lock")
.push(record.clone());
Ok(Some(
StoredRequestUsageAudit::new(
"usage-1".to_string(),
record.request_id,
record.user_id,
record.api_key_id,
record.username,
record.api_key_name,
record.provider_name,
record.model,
record.target_model,
record.provider_id,
record.provider_endpoint_id,
record.provider_api_key_id,
record.request_type,
record.api_format,
record.api_family,
record.endpoint_kind,
record.endpoint_api_format,
record.provider_api_family,
record.provider_endpoint_kind,
record.has_format_conversion.unwrap_or(false),
record.is_stream.unwrap_or(false),
record.input_tokens.unwrap_or_default() as i32,
record.output_tokens.unwrap_or_default() as i32,
record.total_tokens.unwrap_or_default() as i32,
record.total_cost_usd.unwrap_or_default(),
record.actual_total_cost_usd.unwrap_or_default(),
record.status_code.map(i32::from),
record.error_message,
record.error_category,
record.response_time_ms.map(|value| value as i32),
record.first_byte_time_ms.map(|value| value as i32),
record.status,
record.billing_status,
record
.created_at_unix_ms
.unwrap_or(record.updated_at_unix_secs) as i64,
record.updated_at_unix_secs as i64,
record.finalized_at_unix_secs.map(|value| value as i64),
)
.expect("stored usage should build"),
))
}
}
#[async_trait]
impl UsageSettlementWriter for TestUsageStore {
fn has_usage_settlement_writer(&self) -> bool {
true
}
async fn reconcile_usage_policy_cost(
&self,
input: ReconcileUsagePolicyCostInput,
) -> Result<Option<StoredUsagePolicyCostReservation>, DataLayerError> {
self.reconciliations
.lock()
.expect("reconciliations lock")
.push(input);
Ok(None)
}
async fn settle_usage(
&self,
input: UsageSettlementInput,
) -> Result<Option<StoredUsageSettlement>, aether_data_contracts::DataLayerError> {
self.settlements
.lock()
.expect("settlements lock")
.push(input);
Ok(None)
}
}
#[async_trait]
impl ManualProxyNodeCounter for TestUsageStore {
async fn increment_manual_proxy_node_requests(
&self,
_node_id: &str,
_total_delta: i64,
_failed_delta: i64,
_latency_ms: Option<i64>,
) -> Result<(), aether_data_contracts::DataLayerError> {
self.manual_proxy_counter_calls
.fetch_add(1, Ordering::AcqRel);
Ok(())
}
}
#[async_trait]
impl UsageBillingEventEnricher for TestUsageStore {
async fn enrich_usage_event(&self, event: &mut UsageEvent) -> Result<(), DataLayerError> {
self.enrich_calls
.lock()
.expect("enrich calls lock")
.push(event.request_id.clone());
event.data.total_cost_usd = Some(0.456);
Ok(())
}
}
impl crate::runtime::UsageRuntimeAccess for TestUsageStore {
fn has_usage_writer(&self) -> bool {
true
}
fn has_usage_worker_queue(&self) -> bool {
true
}
fn usage_worker_queue(&self) -> Option<Arc<dyn RuntimeQueueStore>> {
None
}
}
#[async_trait]
impl UsageRecordWriter for SlowUsageStore {
async fn upsert_usage_record(
&self,
record: UpsertUsageRecord,
) -> Result<Option<StoredRequestUsageAudit>, DataLayerError> {
let active = self
.active
.fetch_add(1, std::sync::atomic::Ordering::AcqRel)
+ 1;
self.max_active
.fetch_max(active, std::sync::atomic::Ordering::AcqRel);
tokio::time::sleep(Duration::from_millis(30)).await;
self.records
.lock()
.expect("records lock")
.push(record.request_id.clone());
self.active
.fetch_sub(1, std::sync::atomic::Ordering::AcqRel);
Ok(None)
}
}
#[async_trait]
impl UsageSettlementWriter for SlowUsageStore {
fn has_usage_settlement_writer(&self) -> bool {
false
}
async fn settle_usage(
&self,
_input: UsageSettlementInput,
) -> Result<Option<StoredUsageSettlement>, DataLayerError> {
Ok(None)
}
}
#[async_trait]
impl ManualProxyNodeCounter for SlowUsageStore {
async fn increment_manual_proxy_node_requests(
&self,
_node_id: &str,
_total_delta: i64,
_failed_delta: i64,
_latency_ms: Option<i64>,
) -> Result<(), DataLayerError> {
Ok(())
}
}
#[async_trait]
impl UsageBillingEventEnricher for SlowUsageStore {
async fn enrich_usage_event(&self, _event: &mut UsageEvent) -> Result<(), DataLayerError> {
Ok(())
}
}
impl crate::runtime::UsageRuntimeAccess for SlowUsageStore {
fn has_usage_writer(&self) -> bool {
true
}
fn has_usage_worker_queue(&self) -> bool {
true
}
fn usage_worker_queue(&self) -> Option<Arc<dyn RuntimeQueueStore>> {
None
}
fn usage_worker_should_defer_for_database_pressure(&self) -> bool {
self.db_pressure.load(Ordering::Acquire)
}
}
2026-05-22 23:57:38 +08:00
#[async_trait]
impl UsageEventRecorder for SelectiveFailingRecorder {
async fn record_usage_event(&self, event: &UsageEvent) -> Result<(), DataLayerError> {
self.calls
.lock()
.expect("calls lock")
.push(event.request_id.clone());
if event.request_id == "req-worker-poison" {
return Err(DataLayerError::UnexpectedValue(
"permanent test error".to_string(),
));
}
Ok(())
}
}
fn sample_event() -> UsageEvent {
UsageEvent::new(
UsageEventType::Completed,
"req-worker-123".to_string(),
UsageEventData {
user_id: Some("user-worker-123".to_string()),
api_key_id: Some("api-key-worker-123".to_string()),
provider_name: "openai".to_string(),
provider_id: Some("provider-worker-123".to_string()),
provider_endpoint_id: Some("endpoint-worker-123".to_string()),
provider_api_key_id: Some("provider-key-worker-123".to_string()),
model: "gpt-5".to_string(),
api_format: Some("openai:chat".to_string()),
endpoint_api_format: Some("openai:chat".to_string()),
is_stream: Some(false),
status_code: Some(200),
input_tokens: Some(4),
output_tokens: Some(6),
total_tokens: Some(10),
response_time_ms: Some(52),
..UsageEventData::default()
},
)
}
#[tokio::test]
async fn write_event_record_persists_usage_and_triggers_settlement() {
let store = TestUsageStore::default();
let event = sample_event();
write_event_record(&store, &event)
.await
.expect("worker should write usage record");
let records = store.records.lock().expect("records lock");
assert_eq!(records.len(), 1);
assert_eq!(records[0].request_id, "req-worker-123");
assert_eq!(records[0].status, "completed");
drop(records);
let settlements = store.settlements.lock().expect("settlements lock");
assert_eq!(settlements.len(), 1);
assert_eq!(settlements[0].request_id, "req-worker-123");
}
2026-05-22 23:57:38 +08:00
#[tokio::test]
async fn same_request_id_terminal_events_reconcile_each_reservation_token_before_upsert() {
let store = TestUsageStore::default();
let mut first = sample_event();
first.request_id = "shared-client-trace".to_string();
first.data.actual_total_cost_usd = Some(0.25);
first.data.request_metadata = Some(serde_json::json!({
"plan_usage_reservation_token": "server-token-a"
}));
let mut second = first.clone();
second.data.actual_total_cost_usd = Some(0.75);
second.data.request_metadata = Some(serde_json::json!({
"plan_usage_reservation_token": "server-token-b"
}));
write_event_record(&store, &first)
.await
.expect("first terminal event");
write_event_record(&store, &second)
.await
.expect("second terminal event");
let reconciliations = store.reconciliations.lock().expect("reconciliations lock");
assert_eq!(reconciliations.len(), 2);
assert_eq!(reconciliations[0].request_id, "shared-client-trace");
assert_eq!(reconciliations[0].reservation_token, "server-token-a");
assert_eq!(reconciliations[0].actual_cost_units, 25_000_000);
assert_eq!(reconciliations[1].request_id, "shared-client-trace");
assert_eq!(reconciliations[1].reservation_token, "server-token-b");
assert_eq!(reconciliations[1].actual_cost_units, 75_000_000);
assert_eq!(store.records.lock().expect("records lock").len(), 2);
}
#[tokio::test]
async fn replayable_usage_write_does_not_duplicate_transport_owned_proxy_counter() {
let store = TestUsageStore::default();
let mut event = sample_event();
event.data.request_metadata = Some(serde_json::json!({
"proxy": {"mode": "manual", "node_id": "manual-node-1"}
}));
write_event_record(&store, &event)
.await
.expect("first usage write should succeed");
write_event_record(&store, &event)
.await
.expect("replayed usage write should succeed");
assert_eq!(
store.manual_proxy_counter_calls.load(Ordering::Acquire),
0,
"proxy traffic belongs to the transport attempt, not the replayable usage worker"
);
}
#[tokio::test]
async fn data_event_recorder_enriches_terminal_event_before_write() {
let store = Arc::new(TestUsageStore::default());
let recorder = super::UsageDataEventRecorder::new(Arc::clone(&store));
let event = sample_event();
recorder
.record_usage_event(&event)
.await
.expect("recorder should enrich and write usage");
assert_eq!(
store
.enrich_calls
.lock()
.expect("enrich calls lock")
.as_slice(),
["req-worker-123"]
);
let records = store.records.lock().expect("records lock");
assert_eq!(records.len(), 1);
assert_eq!(records[0].total_cost_usd, Some(0.456));
}
#[tokio::test]
async fn data_event_recorder_skips_enrichment_for_lifecycle_event() {
let store = Arc::new(TestUsageStore::default());
let recorder = super::UsageDataEventRecorder::new(Arc::clone(&store));
let mut event = sample_event();
event.event_type = UsageEventType::Pending;
recorder
.record_usage_event(&event)
.await
.expect("recorder should write lifecycle usage");
assert!(store
.enrich_calls
.lock()
.expect("enrich calls lock")
.is_empty());
let records = store.records.lock().expect("records lock");
assert_eq!(records.len(), 1);
assert_eq!(records[0].total_cost_usd, None);
}
#[tokio::test]
async fn data_event_recorder_serializes_same_request_id_writes() {
let store = Arc::new(SlowUsageStore::default());
let recorder = Arc::new(super::UsageDataEventRecorder::new(Arc::clone(&store)));
let mut first = sample_event();
first.request_id = "req-same".to_string();
first.event_type = UsageEventType::Pending;
let mut second = sample_event();
second.request_id = "req-same".to_string();
second.event_type = UsageEventType::Completed;
let first_recorder = Arc::clone(&recorder);
let second_recorder = Arc::clone(&recorder);
tokio::try_join!(
async move { first_recorder.record_usage_event(&first).await },
async move { second_recorder.record_usage_event(&second).await }
)
.expect("same request writes should both succeed");
assert_eq!(
store.max_active.load(std::sync::atomic::Ordering::Acquire),
1
);
assert_eq!(store.records.lock().expect("records lock").len(), 2);
}
#[tokio::test]
async fn data_event_recorder_defers_when_database_pool_is_under_pressure() {
let store = Arc::new(SlowUsageStore::default());
store.db_pressure.store(true, Ordering::Release);
let gate = Arc::new(UsageWorkerRecordConcurrencyGate::new(1));
let recorder = super::UsageDataEventRecorder::with_record_gate_and_database_pressure_defer(
Arc::clone(&store),
Some(Arc::clone(&gate)),
);
recorder
.record_usage_event(&sample_event())
.await
.expect("recorder should write after brief defer");
assert_eq!(gate.deferred_total(), 1);
assert_eq!(store.records.lock().expect("records lock").len(), 1);
}
#[tokio::test]
async fn usage_worker_record_gate_limits_concurrent_record_writes() {
let runner = Arc::new(RuntimeState::memory(MemoryRuntimeStateConfig::default()));
let queue_runner: Arc<dyn RuntimeQueueStore> = runner.clone();
let store = Arc::new(SlowUsageStore::default());
let gate = Arc::new(UsageWorkerRecordConcurrencyGate::new(2));
let config = UsageRuntimeConfig {
enabled: true,
stream_key: "usage:test:worker:record-gate".to_string(),
consumer_group: "usage:test:worker:record-gate-group".to_string(),
dlq_stream_key: "usage:test:worker:record-gate-dlq".to_string(),
consumer_batch_size: 1,
consumer_block_ms: 1,
worker_record_concurrency_limit: Some(2),
..UsageRuntimeConfig::default()
};
let mut handles = Vec::new();
for worker_index in 0..4 {
let worker = build_usage_queue_worker_with_record_gate(
Arc::clone(&queue_runner),
Arc::clone(&store),
config.clone(),
Some(Arc::clone(&gate)),
Some(worker_index),
)
.expect("worker should build");
worker
.queue
.ensure_consumer_group()
.await
.expect("group should initialize");
handles.push(tokio::spawn(async move {
let entries = worker
.queue
.read_group(&worker.consumer)
.await
.expect("event should read");
worker
.process_entries(entries)
.await
.expect("event should process");
}));
}
for index in 0..4 {
let mut event = sample_event();
event.request_id = format!("req-record-gate-{index}");
UsageQueue::new(queue_runner.clone(), config.clone())
.expect("queue should build")
.enqueue(&event)
.await
.expect("event should enqueue");
}
for handle in handles {
handle.await.expect("worker should complete");
}
assert_eq!(
store.max_active.load(std::sync::atomic::Ordering::Acquire),
2
);
assert_eq!(gate.max_in_flight(), 2);
assert!(gate.wait_total() > 0);
assert_eq!(store.records.lock().expect("records lock").len(), 4);
}
#[tokio::test]
async fn usage_worker_defers_reclaim_until_inflight_read_is_processed() {
let event = sample_event();
let queue = Arc::new(ReadReclaimRaceProbeQueue::new(RuntimeQueueEntry {
id: "1-0".to_string(),
fields: event
.to_stream_fields()
.expect("usage event should serialize"),
}));
let queue_runner: Arc<dyn RuntimeQueueStore> = queue.clone();
let config = UsageRuntimeConfig {
enabled: true,
stream_key: "usage:test:worker:read-reclaim-race".to_string(),
consumer_group: "usage:test:worker:read-reclaim-race-group".to_string(),
dlq_stream_key: "usage:test:worker:read-reclaim-race-dlq".to_string(),
consumer_batch_size: 1,
consumer_block_ms: 1_000,
reclaim_interval_ms: 10,
reclaim_idle_ms: 1,
reclaim_count: 1,
..UsageRuntimeConfig::default()
};
let recorder = Arc::new(SelectiveFailingRecorder::default());
let worker_recorder: Arc<dyn UsageEventRecorder> = recorder.clone();
let control = UsageWorkerControl::default();
let (telemetry_tx, _telemetry_rx) = tokio::sync::mpsc::channel(8);
let worker = UsageQueueWorker::new(queue_runner, worker_recorder, config, None)
.expect("worker should build")
.with_supervisor(control.clone(), telemetry_tx);
let handle = tokio::spawn(worker.run());
tokio::time::timeout(Duration::from_secs(1), async {
while queue.read_calls.load(Ordering::Acquire) == 0 {
tokio::task::yield_now().await;
}
})
.await
.expect("worker should start the blocking read");
tokio::time::sleep(Duration::from_millis(50)).await;
assert_eq!(
queue.reclaim_calls.load(Ordering::Acquire),
0,
"reclaim must not run while XREADGROUP can still return the same PEL entry"
);
assert!(recorder.calls.lock().expect("calls lock").is_empty());
queue.release_read.notify_one();
tokio::time::timeout(Duration::from_secs(1), async {
while !queue.acked.load(Ordering::Acquire)
|| queue.reclaim_calls.load(Ordering::Acquire) == 0
{
tokio::task::yield_now().await;
}
})
.await
.expect("read entry should be processed before deferred reclaim runs");
assert_eq!(
recorder.calls.lock().expect("calls lock").as_slice(),
[event.request_id.as_str()],
"the stream entry must be recorded exactly once"
);
assert_eq!(queue.read_completed.load(Ordering::Acquire), 1);
assert!(!queue.first_read_cancelled.load(Ordering::Acquire));
control.request_shutdown();
tokio::time::timeout(Duration::from_secs(1), handle)
.await
.expect("worker should stop promptly")
.expect("worker task should not panic");
}
#[tokio::test]
async fn usage_worker_shutdown_cancels_blocking_read_promptly() {
let event = sample_event();
let queue = Arc::new(ReadReclaimRaceProbeQueue::new(RuntimeQueueEntry {
id: "2-0".to_string(),
fields: event
.to_stream_fields()
.expect("usage event should serialize"),
}));
let queue_runner: Arc<dyn RuntimeQueueStore> = queue.clone();
let config = UsageRuntimeConfig {
enabled: true,
stream_key: "usage:test:worker:shutdown-read".to_string(),
consumer_group: "usage:test:worker:shutdown-read-group".to_string(),
dlq_stream_key: "usage:test:worker:shutdown-read-dlq".to_string(),
consumer_batch_size: 1,
consumer_block_ms: 60_000,
reclaim_interval_ms: 10,
reclaim_idle_ms: 1,
reclaim_count: 1,
..UsageRuntimeConfig::default()
};
let recorder: Arc<dyn UsageEventRecorder> = Arc::new(SelectiveFailingRecorder::default());
let control = UsageWorkerControl::default();
let (telemetry_tx, _telemetry_rx) = tokio::sync::mpsc::channel(8);
let worker = UsageQueueWorker::new(queue_runner, recorder, config, None)
.expect("worker should build")
.with_supervisor(control.clone(), telemetry_tx);
let handle = tokio::spawn(worker.run());
tokio::time::timeout(Duration::from_secs(1), async {
while queue.read_calls.load(Ordering::Acquire) == 0 {
tokio::task::yield_now().await;
}
})
.await
.expect("worker should start the blocking read");
control.request_shutdown();
tokio::time::timeout(Duration::from_secs(1), handle)
.await
.expect("shutdown should interrupt the blocking read")
.expect("worker task should not panic");
assert!(queue.first_read_cancelled.load(Ordering::Acquire));
assert_eq!(queue.read_completed.load(Ordering::Acquire), 0);
assert_eq!(queue.reclaim_calls.load(Ordering::Acquire), 0);
}
2026-05-22 23:57:38 +08:00
#[test]
fn usage_event_record_error_classifies_permanent_failures() {
assert!(usage_event_record_error_is_permanent(
&DataLayerError::UnexpectedValue("bad payload".to_string())
));
assert!(usage_event_record_error_is_permanent(
&DataLayerError::Postgres(
"error returned from database: violates foreign key constraint (SQLSTATE 23503)"
.to_string()
)
));
assert!(!usage_event_record_error_is_permanent(
&DataLayerError::Redis("connection refused".to_string())
));
assert!(!usage_event_record_error_is_permanent(
&DataLayerError::TimedOut("postgres acquire".to_string())
));
}
#[tokio::test]
async fn process_entries_dead_letters_permanent_record_error_and_continues() {
let runner = Arc::new(RuntimeState::memory(MemoryRuntimeStateConfig::default()));
let queue_runner: Arc<dyn RuntimeQueueStore> = runner.clone();
let recorder = Arc::new(SelectiveFailingRecorder::default());
let config = UsageRuntimeConfig {
enabled: true,
stream_key: "usage:test:worker:events".to_string(),
consumer_group: "usage:test:worker:group".to_string(),
dlq_stream_key: "usage:test:worker:dlq".to_string(),
consumer_batch_size: 10,
consumer_block_ms: 1,
..UsageRuntimeConfig::default()
};
2026-06-26 14:02:57 +08:00
let worker = UsageQueueWorker::new(queue_runner, recorder.clone(), config, None)
2026-05-22 23:57:38 +08:00
.expect("worker should build");
worker
.queue
.ensure_consumer_group()
.await
.expect("group should initialize");
let mut poison = sample_event();
poison.request_id = "req-worker-poison".to_string();
let mut ok = sample_event();
ok.request_id = "req-worker-ok".to_string();
worker
.queue
.enqueue(&poison)
.await
.expect("poison event should enqueue");
worker
.queue
.enqueue(&ok)
.await
.expect("ok event should enqueue");
let entries = worker
.queue
.read_group(&worker.consumer)
.await
.expect("events should read");
assert_eq!(entries.len(), 2);
worker
.process_entries(entries)
.await
.expect("permanent failure should not block batch");
assert_eq!(
recorder.calls.lock().expect("calls lock").as_slice(),
["req-worker-poison", "req-worker-ok"]
);
runner
.ensure_consumer_group(
"usage:test:worker:dlq",
"usage:test:worker:dlq-group",
"0-0",
)
.await
.expect("dlq group should initialize");
let dlq_entries = runner
.read_group(
"usage:test:worker:dlq",
"usage:test:worker:dlq-group",
"usage-test-dlq-consumer",
10,
Some(1),
)
.await
.expect("dlq should read");
assert_eq!(dlq_entries.len(), 1);
let payload = dlq_entries[0]
.fields
.get("payload")
.expect("dlq payload should exist");
assert!(payload.contains("req-worker-poison"));
assert!(payload.contains("permanent test error"));
}
}