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; use aether_runtime_state::{RuntimeQueueEntry, RuntimeQueueStore}; use async_trait::async_trait; use tokio::sync::{mpsc, Notify}; use tracing::warn; use crate::event_capture_budget::{shared_capture_memory_budget, EventCaptureMemoryBudget}; use crate::executor::spawn_on_usage_background_runtime; use crate::keyed_lock::KeyedAsyncLockPool; use crate::queue::UsageDeadLetterOutcome; use crate::runtime::{ UsageBillingEventEnricher, UsageRuntimeAccess, UsageWorkerRecordConcurrencyGate, }; use crate::settlement::{ reconcile_usage_policy_cost_for_event_with_result, settle_usage_with_reconciled_cost, }; use crate::{ build_upsert_usage_record_from_event, UsageEvent, UsageEventType, UsageQueue, UsageRuntimeConfig, UsageSettlementWriter, }; const USAGE_WORKER_DB_PRESSURE_DEFER_MS: u64 = 10; const USAGE_WORKER_ACK_CHUNK_SIZE: usize = 100; enum EntryDisposition { NeedsAck, Complete, Deferred(DataLayerError), } #[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, ) -> 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 { 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 { None } async fn upsert_usage_record( &self, record: UpsertUsageRecord, ) -> Result, 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, ) -> Result<(), DataLayerError> { for record in records { self.upsert_first_byte_usage_record(record).await?; } Ok(()) } async fn upsert_pending_usage_records( &self, records: Vec, ) -> Result<(), DataLayerError> { for record in records { self.upsert_usage_record(record).await?; } Ok(()) } } pub struct UsageDataEventRecorder { data: Arc, record_gate: Option>, defer_for_database_pressure: bool, } impl UsageDataEventRecorder { pub fn new(data: Arc) -> Self { Self::with_record_gate(data, None) } pub(crate) fn with_record_gate( data: Arc, record_gate: Option>, ) -> Self { Self { data, record_gate, defer_for_database_pressure: false, } } pub(crate) fn with_record_gate_and_database_pressure_defer( data: Arc, record_gate: Option>, ) -> Self { Self { data, record_gate, defer_for_database_pressure: true, } } } #[async_trait] impl UsageEventRecorder for UsageDataEventRecorder 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> { static LOCKS: OnceLock = OnceLock::new(); LOCKS .get_or_init(KeyedAsyncLockPool::default) .lock_for(request_id) } pub struct UsageQueueWorker { queue: UsageQueue, recorder: Arc, consumer: String, worker_index: Option, control: Option, telemetry: Option>, config: UsageRuntimeConfig, capture_memory_budget: Arc, } #[derive(Debug, Clone, Default)] pub(crate) struct UsageWorkerControl { shutdown: Arc, shutdown_notify: Arc, } impl UsageWorkerControl { pub(crate) fn request_shutdown(&self) { self.shutdown.store(true, Ordering::Release); self.shutdown_notify.notify_waiters(); self.shutdown_notify.notify_one(); } fn should_shutdown(&self) -> bool { self.shutdown.load(Ordering::Acquire) } pub(crate) async fn wait_for_shutdown(&self) { loop { let notified = self.shutdown_notify.notified(); tokio::pin!(notified); notified.as_mut().enable(); if self.should_shutdown() { return; } notified.await; } } } #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub(crate) struct UsageWorkerObservation { pub worker_index: Option, 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, 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, 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, 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, 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) -> 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) -> 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) -> 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, } } } impl UsageQueueWorker { pub fn new( runner: Arc, recorder: Arc, config: UsageRuntimeConfig, worker_index: Option, ) -> Result { let queue = UsageQueue::new(runner, config.clone())?; let consumer = consumer_name(worker_index); Ok(Self { queue, recorder, consumer, worker_index, control: None, telemetry: None, config, capture_memory_budget: shared_capture_memory_budget(), }) } pub(crate) fn with_supervisor( mut self, control: UsageWorkerControl, telemetry: mpsc::Sender, ) -> 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 }) } pub(crate) fn with_shutdown(mut self, control: UsageWorkerControl) -> Self { self.control = Some(control); self } 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", 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; let mut reclaim_cursor = "0-0".to_string(); loop { if self.should_shutdown() { break; } let result = { let mut read_future = Box::pin(self.queue.read_group_reserved(&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(batch) => { self.report_read(batch.entries.len(), batch.requested_count); let reservation = batch.reservation; let result = self.process_entries(batch.entries).await; // The raw fields and decoded event must be gone, and ACK must finish, // before another worker can reuse this batch's receive allowance. drop(reservation); if let Err(err) = result { 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(&mut reclaim_cursor).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, cursor: &mut String) { let result = tokio::select! { biased; _ = self.wait_for_shutdown() => return, result = self.queue.claim_stale_page_reserved(&self.consumer, cursor) => result, }; match result { Ok(batch) => { let reservation = batch.reservation; let page = batch.page; // An empty page can still advance past a large active PEL prefix. // Failed writes remain pending and are revisited after the scan wraps. *cursor = page.next_start_id; self.report_reclaimed(page.entries.len()); let result = self.process_entries(page.entries).await; drop(page.deleted_ids); drop(reservation); if let Err(err) = result { 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" ); } } } fn should_shutdown(&self) -> bool { self.control .as_ref() .is_some_and(UsageWorkerControl::should_shutdown) } fn report_read(&self, entries_read: usize, requested_count: usize) { self.report(UsageWorkerObservation::read( self.worker_index, entries_read, requested_count, )); } 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) { let Some(telemetry) = &self.telemetry else { return; }; let _ = telemetry.try_send(observation); } async fn process_entries(&self, entries: Vec) -> Result<(), DataLayerError> { if entries.is_empty() { return Ok(()); } let mut ack_ids = Vec::new(); let mut deferred_error = None; for entry in entries { let id = entry.id.clone(); let result = self.process_entry(entry).await; match result { Ok(EntryDisposition::NeedsAck) => { ack_ids.push(id); if ack_ids.len() >= USAGE_WORKER_ACK_CHUNK_SIZE { self.acknowledge_entries(&ack_ids).await?; ack_ids.clear(); } } Ok(EntryDisposition::Complete) => {} Ok(EntryDisposition::Deferred(err)) => { // An entry that cannot fit the DLQ encoder must not indefinitely block // the healthy entries returned with it on every reclaim pass. deferred_error.get_or_insert(err); } Err(err) => { if !ack_ids.is_empty() { let _ = self.acknowledge_entries(&ack_ids).await; } return Err(err); } } } if !ack_ids.is_empty() { self.acknowledge_entries(&ack_ids).await?; } match deferred_error { Some(err) => Err(err), None => Ok(()), } } async fn acknowledge_entries(&self, ids: &[String]) -> Result<(), DataLayerError> { let acked = self.queue.ack_and_delete_counted(ids).await?; if acked > 0 { self.report_acked(acked); } Ok(()) } async fn dead_letter_entry( &self, entry: RuntimeQueueEntry, error: DataLayerError, event_name: &'static str, ) -> Result { let id = entry.id.clone(); let outcome = self .queue .transfer_dead_letter_owned(entry, error.to_string()) .await?; let (destination_id, disposition) = match outcome { UsageDeadLetterOutcome::Transferred { destination_id, acked, } => { if acked > 0 { self.report_acked(acked); } (destination_id, EntryDisposition::Complete) } UsageDeadLetterOutcome::Appended { destination_id } => { (destination_id, EntryDisposition::NeedsAck) } UsageDeadLetterOutcome::NotPending => { warn!( event_name = "usage_worker_dead_letter_source_not_pending", log_type = "ops", worker_consumer = %self.consumer, worker_group = %self.config.consumer_group, entry_id = %id, error = %error, "usage worker skipped dead letter transfer because the source is no longer pending" ); return Ok(EntryDisposition::Complete); } UsageDeadLetterOutcome::EncodingDeferred { error } => { warn!( event_name = "usage_worker_dead_letter_encoding_deferred", log_type = "ops", worker_consumer = %self.consumer, worker_group = %self.config.consumer_group, entry_id = %id, error = %error, "usage worker retained entry pending after dead letter encoding failed" ); return Ok(EntryDisposition::Deferred(error)); } }; self.report_dead_lettered(1); warn!( event_name, log_type = "ops", worker_consumer = %self.consumer, worker_group = %self.config.consumer_group, entry_id = %id, dead_letter_id = %destination_id, error = %error, "usage worker appended queue entry to dead letter" ); Ok(disposition) } async fn process_entry( &self, entry: RuntimeQueueEntry, ) -> Result { let event = match UsageEvent::from_stream_fields_with_capture_budget( &entry.fields, Arc::clone(&self.capture_memory_budget), ) { Ok(event) => event, Err(err) => { return self .dead_letter_entry(entry, err, "usage_worker_entry_decode_dead_lettered") .await; } }; match self.recorder.record_usage_event(&event).await { Ok(()) => Ok(EntryDisposition::NeedsAck), Err(err) if usage_event_record_error_is_permanent(&err) => { warn!( event_name = "usage_worker_entry_record_permanent_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 encountered a non-retryable usage record failure" ); drop(event); self.dead_letter_entry(entry, err, "usage_worker_entry_record_dead_lettered") .await } 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) } } } } 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( runner: Arc, data: Arc, config: UsageRuntimeConfig, worker_index: Option, ) -> Result 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( runner: Arc, data: Arc, config: UsageRuntimeConfig, record_gate: Option>, worker_index: Option, ) -> Result where T: UsageRuntimeAccess + 'static, { UsageQueueWorker::new( runner, Arc::new( UsageDataEventRecorder::with_record_gate_and_database_pressure_defer(data, record_gate), ), config, worker_index, ) } pub async fn write_event_record(data: &T, event: &UsageEvent) -> Result<(), DataLayerError> where T: UsageRecordWriter + UsageSettlementWriter + Send + Sync, { let reconciled = reconcile_usage_policy_cost_for_event_with_result(data, event).await?; let record = build_upsert_usage_record_from_event(event)?; if let Some(stored) = data.upsert_usage_record(record).await? { settle_usage_with_reconciled_cost(data, &stored, reconciled).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(data: &T, event: &mut UsageEvent) -> Result<(), DataLayerError> where T: UsageBillingEventEnricher + Send + Sync, { if !matches!( event.event_type, UsageEventType::Completed | UsageEventType::Failed | UsageEventType::Cancelled ) { return Ok(()); } 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" ); return Err(err); } Ok(()) } fn consumer_name(worker_index: Option) -> 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()); match worker_index { Some(worker_index) => format!("{host}:{}:{worker_index}", std::process::id()), None => format!("{host}:{}", std::process::id()), } } #[cfg(test)] mod tests { mod dead_letter_transfer { include!("worker_dead_letter_tests.rs"); } use std::collections::{BTreeMap, VecDeque}; use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; 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, UsageBodyCaptureState, }; use aether_data_contracts::DataLayerError; use aether_runtime_state::{ MemoryRuntimeStateConfig, RuntimeQueueEntry, RuntimeQueueReclaimConfig, RuntimeQueueReclaimPage, RuntimeQueueStats, RuntimeQueueStore, RuntimeState, }; use async_trait::async_trait; use tokio::sync::Notify; use super::{ build_usage_queue_worker_with_record_gate, usage_event_record_error_is_permanent, write_event_record, ManualProxyNodeCounter, UsageEventRecorder, UsageQueueWorker, UsageRecordWriter, UsageWorkerControl, }; use crate::dead_letter_encoding::DeadLetterEncodingBudget; use crate::event_capture_budget::EventCaptureMemoryBudget; use crate::queue_read_budget::QueueReadBudget; use crate::runtime::UsageWorkerRecordConcurrencyGate; use crate::UsageBillingEventEnricher; use crate::{ UsageEvent, UsageEventData, UsageEventType, UsageQueue, UsageRuntimeConfig, UsageSettlementWriter, }; #[derive(Default)] struct TestUsageStore { records: Mutex>, settlements: Mutex>, reconciliations: Mutex>, enrich_calls: Mutex>, enrich_outcomes: Mutex>, manual_proxy_counter_calls: AtomicUsize, } enum TestEnrichmentOutcome { TimedOut, Unpriced, Priced { listed: f64, actual: f64 }, } #[derive(Default)] struct ControlledRecorder { entered: Notify, release: Notify, calls: AtomicUsize, } #[async_trait] impl UsageEventRecorder for ControlledRecorder { async fn record_usage_event(&self, _event: &UsageEvent) -> Result<(), DataLayerError> { self.calls.fetch_add(1, Ordering::AcqRel); self.entered.notify_one(); self.release.notified().await; Ok(()) } } #[derive(Default)] struct SelectiveFailingRecorder { calls: Mutex>, } enum CaptureRecordOutcome { RetryOnce, PermanentFailure, Wait, } struct CaptureBudgetRecorder { budget: Arc, outcome: CaptureRecordOutcome, calls: AtomicUsize, entered: Notify, } #[async_trait] impl UsageEventRecorder for CaptureBudgetRecorder { async fn record_usage_event(&self, event: &UsageEvent) -> Result<(), DataLayerError> { let call = self.calls.fetch_add(1, Ordering::AcqRel); assert_eq!(event.data.input_tokens, Some(4)); assert_eq!(event.data.output_tokens, Some(6)); assert_eq!(event.data.total_tokens, Some(10)); assert_eq!(event.data.cache_read_input_tokens, Some(0)); assert_eq!(event.data.actual_total_cost_usd, Some(0.123)); let metadata = event.data.request_metadata.as_ref().expect("billing facts"); assert_eq!(metadata["requested_reasoning_effort"], "high"); assert_eq!(metadata["provider_service_tier"], "priority"); assert_eq!(metadata["provider_actual_service_tier"], "default"); if event.data.response_body.is_some() { let retained = self.budget.retained_bytes(); assert!(retained > 0); let recorder_copy = event.clone(); assert!(recorder_copy.data.response_body.is_some()); assert!(self.budget.retained_bytes() > retained); drop(recorder_copy); assert_eq!(self.budget.retained_bytes(), retained); } else { assert_eq!(self.budget.retained_bytes(), 0); assert_eq!( event.data.response_body_state, Some(UsageBodyCaptureState::Truncated) ); } self.entered.notify_one(); match self.outcome { CaptureRecordOutcome::RetryOnce if call == 0 => Err(DataLayerError::TimedOut( "retry test database write".to_string(), )), CaptureRecordOutcome::RetryOnce => Ok(()), CaptureRecordOutcome::PermanentFailure => Err(DataLayerError::UnexpectedValue( "permanent capture test error".to_string(), )), CaptureRecordOutcome::Wait => std::future::pending().await, } } } async fn capture_budget_worker( budget_bytes: usize, outcome: CaptureRecordOutcome, ) -> ( Arc, UsageQueueWorker, Arc, ) { let runner = Arc::new(RuntimeState::memory(MemoryRuntimeStateConfig::default())); let queue_runner: Arc = runner.clone(); let budget = Arc::new(EventCaptureMemoryBudget::new(budget_bytes)); let recorder = Arc::new(CaptureBudgetRecorder { budget: Arc::clone(&budget), outcome, calls: AtomicUsize::new(0), entered: Notify::new(), }); let config = UsageRuntimeConfig { enabled: true, stream_key: "usage:test:worker:capture".to_string(), consumer_group: "usage:test:worker:capture-group".to_string(), dlq_stream_key: "usage:test:worker:capture-dlq".to_string(), consumer_batch_size: 1, consumer_block_ms: 1, ..UsageRuntimeConfig::default() }; let mut worker = UsageQueueWorker::new(queue_runner, recorder.clone(), config, None) .expect("worker should build"); worker.capture_memory_budget = budget; worker.queue = worker .queue .with_dead_letter_encoding_budget(Arc::new(DeadLetterEncodingBudget::new( 64 * 1024 * 1024, 4, ))); worker .queue .ensure_consumer_group() .await .expect("consumer group"); (runner, worker, recorder) } fn captured_worker_event() -> UsageEvent { let mut event = sample_event(); event.data.model = "gpt-5.6-sol".to_string(); event.data.endpoint_api_format = Some("openai:responses".to_string()); event.data.request_body = Some(serde_json::json!({"reasoning": {"effort": "high"}})); event.data.provider_request_body = Some(serde_json::json!({ "model": "gpt-5.6-sol", "service_tier": "priority" })); event.data.response_body = Some(serde_json::json!({ "service_tier": "Default", "output": "x".repeat(1024) })); event.data.cache_read_input_tokens = Some(0); event.data.actual_total_cost_usd = Some(0.123); event } #[derive(Default)] struct SlowUsageStore { active: std::sync::atomic::AtomicUsize, max_active: std::sync::atomic::AtomicUsize, db_pressure: AtomicBool, records: Mutex>, } struct ReadReclaimRaceProbeQueue { entry: RuntimeQueueEntry, read_calls: AtomicUsize, first_read_cancelled: AtomicBool, read_completed: AtomicUsize, release_read: Notify, reclaim_calls: AtomicUsize, reclaim_pages: Mutex>>>, reclaim_cursors: Mutex>, acked: AtomicBool, requested_counts: Mutex>, block_ack: AtomicBool, ack_entered: Notify, release_ack: Notify, block_reclaim: AtomicBool, reclaim_cancelled: 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), reclaim_pages: Mutex::new(None), reclaim_cursors: Mutex::new(Vec::new()), acked: AtomicBool::new(false), requested_counts: Mutex::new(Vec::new()), block_ack: AtomicBool::new(false), ack_entered: Notify::new(), release_ack: Notify::new(), block_reclaim: AtomicBool::new(false), reclaim_cancelled: 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, _maxlen: Option, ) -> Result { Ok("0-0".to_string()) } async fn read_group( &self, _stream: &str, _group: &str, _consumer: &str, count: usize, _block_ms: Option, ) -> Result, DataLayerError> { let call_index = self.read_calls.fetch_add(1, Ordering::AcqRel); self.requested_counts .lock() .expect("requested counts lock") .push(count); 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, DataLayerError> { self.reclaim_calls.fetch_add(1, Ordering::AcqRel); Ok((!self.acked.load(Ordering::Acquire)) .then(|| self.entry.clone()) .into_iter() .collect()) } async fn claim_stale_page( &self, stream: &str, group: &str, consumer: &str, start_id: &str, config: RuntimeQueueReclaimConfig, ) -> Result { self.reclaim_cursors .lock() .expect("reclaim cursors lock") .push(start_id.to_string()); if self.block_reclaim.load(Ordering::Acquire) { let _guard = FirstReadDropGuard { cancelled: &self.reclaim_cancelled, completed: false, }; return std::future::pending().await; } let scripted = self .reclaim_pages .lock() .expect("reclaim pages lock") .as_mut() .map(|pages| pages.pop_front().expect("scripted reclaim page")); if let Some(page) = scripted { self.reclaim_calls.fetch_add(1, Ordering::AcqRel); return page; } Ok(RuntimeQueueReclaimPage { next_start_id: "0-0".to_string(), entries: self .claim_stale(stream, group, consumer, start_id, config) .await?, deleted_ids: Vec::new(), }) } async fn ack( &self, _stream: &str, _group: &str, ids: &[String], ) -> Result { if self.block_ack.load(Ordering::Acquire) { self.ack_entered.notify_one(); self.release_ack.notified().await; } 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 { Ok(ids.len()) } async fn stats( &self, _stream: &str, _group: Option<&str>, ) -> Result { Ok(RuntimeQueueStats::default()) } } #[async_trait] impl UsageRecordWriter for TestUsageStore { async fn upsert_usage_record( &self, record: UpsertUsageRecord, ) -> Result, 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, DataLayerError> { self.reconciliations .lock() .expect("reconciliations lock") .push(input); Ok(None) } async fn settle_usage( &self, input: UsageSettlementInput, ) -> Result, 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, ) -> 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()); match self .enrich_outcomes .lock() .expect("enrich outcomes lock") .pop_front() { Some(TestEnrichmentOutcome::TimedOut) => { return Err(DataLayerError::TimedOut("test pricing lookup".to_string())); } Some(TestEnrichmentOutcome::Unpriced) => return Ok(()), Some(TestEnrichmentOutcome::Priced { listed, actual }) => { event.data.total_cost_usd = Some(listed); event.data.actual_total_cost_usd = Some(actual); return Ok(()); } None => {} } 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> { None } } #[async_trait] impl UsageRecordWriter for SlowUsageStore { async fn upsert_usage_record( &self, record: UpsertUsageRecord, ) -> Result, 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, 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, ) -> 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> { None } fn usage_worker_should_defer_for_database_pressure(&self) -> bool { self.db_pressure.load(Ordering::Acquire) } } #[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"); } #[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_pricing_timeout_stays_pending_until_successful_reclaim() { let runner = Arc::new(RuntimeState::memory(MemoryRuntimeStateConfig::default())); let store = Arc::new(TestUsageStore::default()); store .enrich_outcomes .lock() .expect("enrich outcomes lock") .extend([ TestEnrichmentOutcome::TimedOut, TestEnrichmentOutcome::Priced { listed: 0.456, actual: 0.123, }, ]); let config = UsageRuntimeConfig { consumer_batch_size: 1, consumer_block_ms: 1, reclaim_count: 1, reclaim_idle_ms: 1, queue_payload_max_bytes: 4096, ..UsageRuntimeConfig::default() }; let budget = Arc::new(QueueReadBudget::new(4096, 4096)); let mut worker = build_usage_queue_worker_with_record_gate(runner, store.clone(), config, None, None) .expect("worker should build"); worker.queue = worker.queue.with_read_budget(Arc::clone(&budget)); worker.queue.ensure_consumer_group().await.expect("group"); let mut event = sample_event(); event.data.request_metadata = Some(serde_json::json!({ "plan_usage_reservation_token": "pricing-retry-reservation" })); worker.queue.enqueue(&event).await.expect("enqueue"); let batch = worker .queue .read_group_reserved(&worker.consumer) .await .expect("read"); let error = worker .process_entries(batch.entries) .await .expect_err("pricing timeout must reach the worker"); drop(batch.reservation); assert!(matches!(error, DataLayerError::TimedOut(_))); assert!(store.records.lock().expect("records lock").is_empty()); assert!(store .reconciliations .lock() .expect("reconciliations lock") .is_empty()); assert!(store .settlements .lock() .expect("settlements lock") .is_empty()); let stats = worker.queue.stats().await.expect("pending stats"); assert_eq!((stats.stream_length, stats.group_pending), (1, 1)); assert_eq!( worker .queue .dlq_stats() .await .expect("dlq stats") .stream_length, 0 ); assert_eq!(budget.snapshot().reserved_bytes, 0); let mut cursor = "0-0".to_string(); tokio::time::timeout(Duration::from_secs(1), async { loop { tokio::time::sleep(Duration::from_millis(1)).await; worker.reclaim_stale_entries(&mut cursor).await; if !store.records.lock().expect("records lock").is_empty() { break; } } }) .await .expect("pending entry should become reclaimable"); { let records = store.records.lock().expect("records lock"); assert_eq!(records.len(), 1); assert_eq!(records[0].total_cost_usd, Some(0.456)); assert_eq!(records[0].actual_total_cost_usd, Some(0.123)); assert_eq!(records[0].total_tokens, Some(10)); } { let reconciliations = store.reconciliations.lock().expect("reconciliations lock"); assert_eq!(reconciliations.len(), 1); assert_eq!(reconciliations[0].actual_cost_units, 12_300_000); assert_eq!( reconciliations[0].reservation_token, "pricing-retry-reservation" ); } assert_eq!(store.settlements.lock().expect("settlements lock").len(), 1); assert_eq!( store.enrich_calls.lock().expect("enrich calls lock").len(), 2 ); let stats = worker.queue.stats().await.expect("acked stats"); assert_eq!((stats.stream_length, stats.group_pending), (0, 0)); assert_eq!( worker .queue .dlq_stats() .await .expect("dlq stats") .stream_length, 0 ); assert_eq!(budget.snapshot().reserved_bytes, 0); } #[tokio::test] async fn data_event_recorder_successful_unpriced_enrichment_keeps_existing_write_behavior() { let store = Arc::new(TestUsageStore::default()); store .enrich_outcomes .lock() .expect("enrich outcomes lock") .push_back(TestEnrichmentOutcome::Unpriced); let recorder = super::UsageDataEventRecorder::new(Arc::clone(&store)); recorder .record_usage_event(&sample_event()) .await .expect("missing pricing is a successful enrichment outcome"); let records = store.records.lock().expect("records lock"); assert_eq!(records.len(), 1); assert_eq!(records[0].total_cost_usd, None); assert_eq!(records[0].actual_total_cost_usd, None); } #[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 = 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_reclaim_cursor_advances_on_empty_pages_and_retries_read_errors() { let event = sample_event(); let entry = RuntimeQueueEntry { id: "43-0".to_string(), fields: event.to_stream_fields().expect("event fields"), }; let runner = Arc::new(ReadReclaimRaceProbeQueue::new(entry.clone())); let page = |next_start_id: &str, entries, deleted_ids| RuntimeQueueReclaimPage { next_start_id: next_start_id.to_string(), entries, deleted_ids, }; *runner.reclaim_pages.lock().expect("reclaim pages lock") = Some(VecDeque::from([ Ok(page("11-0", Vec::new(), vec!["9-0".to_string()])), Err(DataLayerError::Redis("temporary reclaim error".to_string())), Ok(page("42-0", Vec::new(), Vec::new())), Ok(page("0-0", vec![entry], Vec::new())), Ok(page("0-0", Vec::new(), Vec::new())), ])); let recorder = Arc::new(SelectiveFailingRecorder::default()); let worker = UsageQueueWorker::new( runner.clone(), recorder.clone(), UsageRuntimeConfig::default(), None, ) .expect("cursor worker"); let mut cursor = "0-0".to_string(); for expected in ["11-0", "11-0", "42-0", "0-0", "0-0"] { worker.reclaim_stale_entries(&mut cursor).await; assert_eq!(cursor, expected); } assert_eq!( runner .reclaim_cursors .lock() .expect("reclaim cursors lock") .as_slice(), ["0-0", "11-0", "11-0", "42-0", "0-0"] ); assert_eq!( recorder.calls.lock().expect("calls lock").as_slice(), [event.request_id] ); assert!(runner.acked.load(Ordering::Acquire)); assert_eq!(runner.reclaim_calls.load(Ordering::Acquire), 5); } #[tokio::test] async fn usage_worker_reclaim_cursor_advances_after_write_failure_and_revisits_on_wrap() { struct FailOnceRecorder(AtomicUsize); #[async_trait] impl UsageEventRecorder for FailOnceRecorder { async fn record_usage_event(&self, _event: &UsageEvent) -> Result<(), DataLayerError> { if self.0.fetch_add(1, Ordering::AcqRel) == 0 { Err(DataLayerError::TimedOut( "temporary write failure".to_string(), )) } else { Ok(()) } } } let entry = RuntimeQueueEntry { id: "43-0".to_string(), fields: sample_event().to_stream_fields().expect("event fields"), }; let runner = Arc::new(ReadReclaimRaceProbeQueue::new(entry.clone())); *runner.reclaim_pages.lock().expect("reclaim pages lock") = Some(VecDeque::from([ Ok(RuntimeQueueReclaimPage { next_start_id: "50-0".to_string(), entries: vec![entry.clone()], deleted_ids: Vec::new(), }), Ok(RuntimeQueueReclaimPage { next_start_id: "0-0".to_string(), entries: Vec::new(), deleted_ids: Vec::new(), }), Ok(RuntimeQueueReclaimPage { next_start_id: "0-0".to_string(), entries: vec![entry], deleted_ids: Vec::new(), }), ])); let recorder = Arc::new(FailOnceRecorder(AtomicUsize::new(0))); let worker = UsageQueueWorker::new( runner.clone(), recorder.clone(), UsageRuntimeConfig::default(), None, ) .expect("cursor worker"); let mut cursor = "0-0".to_string(); worker.reclaim_stale_entries(&mut cursor).await; assert_eq!(cursor, "50-0"); assert!(!runner.acked.load(Ordering::Acquire)); worker.reclaim_stale_entries(&mut cursor).await; worker.reclaim_stale_entries(&mut cursor).await; assert_eq!(cursor, "0-0"); assert!(runner.acked.load(Ordering::Acquire)); assert_eq!(recorder.0.load(Ordering::Acquire), 2); assert_eq!( runner .reclaim_cursors .lock() .expect("reclaim cursors lock") .as_slice(), ["0-0", "50-0", "0-0"] ); } #[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 = 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 = 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 = 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 = 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); } fn receive_budget_worker_config() -> UsageRuntimeConfig { UsageRuntimeConfig { consumer_batch_size: 128, consumer_block_ms: 60_000, reclaim_interval_ms: 60_000, queue_payload_max_bytes: 4096, ..UsageRuntimeConfig::default() } } #[tokio::test] async fn usage_worker_shared_read_budget_waits_for_slow_recorder_and_releases_on_cancel() { let runner = Arc::new(RuntimeState::memory(MemoryRuntimeStateConfig::default())); let budget = Arc::new(QueueReadBudget::new(4096, 4096)); let config = receive_budget_worker_config(); let first_recorder = Arc::new(ControlledRecorder::default()); let second_recorder = Arc::new(ControlledRecorder::default()); let first_control = UsageWorkerControl::default(); let (telemetry, _observations) = tokio::sync::mpsc::channel(16); let mut first = UsageQueueWorker::new( runner.clone(), first_recorder.clone(), config.clone(), Some(0), ) .expect("first worker") .with_supervisor(first_control.clone(), telemetry); first.queue = first.queue.with_read_budget(Arc::clone(&budget)); let queue = first.queue.clone(); let mut second = UsageQueueWorker::new(runner, second_recorder.clone(), config, Some(1)) .expect("second worker"); second.queue = second.queue.with_read_budget(Arc::clone(&budget)); for index in 0..2 { let mut event = sample_event(); event.request_id = format!("receive-budget-{index}"); queue.enqueue(&event).await.expect("enqueue"); } let first_handle = tokio::spawn(first.run()); tokio::time::timeout(Duration::from_secs(1), first_recorder.entered.notified()) .await .expect("first recorder should hold a batch"); assert!(budget.snapshot().reserved_bytes > 0); let second_handle = tokio::spawn(second.run()); tokio::time::timeout(Duration::from_secs(1), async { while budget.snapshot().waiters == 0 { tokio::task::yield_now().await; } }) .await .expect("second worker should wait before reading"); assert_eq!(second_recorder.calls.load(Ordering::Acquire), 0); let stats = queue.stats().await.expect("blocked stats"); assert_eq!((stats.stream_length, stats.group_pending), (2, 1)); first_control.request_shutdown(); first_recorder.release.notify_one(); tokio::time::timeout(Duration::from_secs(1), first_handle) .await .expect("first worker should finish its acquired batch") .expect("first worker task"); tokio::time::timeout(Duration::from_secs(1), second_recorder.entered.notified()) .await .expect("second worker should acquire released allowance"); assert!(budget.snapshot().reserved_bytes > 0); second_handle.abort(); assert!(second_handle .await .expect_err("cancel worker") .is_cancelled()); assert_eq!(budget.snapshot().reserved_bytes, 0); assert_eq!(budget.snapshot().waiters, 0); let stats = queue.stats().await.expect("cancelled stats"); assert_eq!((stats.stream_length, stats.group_pending), (1, 1)); } #[tokio::test] async fn usage_worker_read_budget_survives_ack_and_reports_actual_requested_count() { let runner = Arc::new(ReadReclaimRaceProbeQueue::new(RuntimeQueueEntry { id: "7-0".to_string(), fields: sample_event().to_stream_fields().expect("event fields"), })); runner.block_ack.store(true, Ordering::Release); let budget = Arc::new(QueueReadBudget::new(4096, 4096)); let control = UsageWorkerControl::default(); let (telemetry, mut observations) = tokio::sync::mpsc::channel(16); let mut worker = UsageQueueWorker::new( runner.clone(), Arc::new(SelectiveFailingRecorder::default()), receive_budget_worker_config(), Some(3), ) .expect("worker") .with_supervisor(control.clone(), telemetry); worker.queue = worker.queue.with_read_budget(Arc::clone(&budget)); runner.release_read.notify_one(); let handle = tokio::spawn(worker.run()); tokio::time::timeout(Duration::from_secs(1), runner.ack_entered.notified()) .await .expect("worker should reach ACK"); assert!(budget.snapshot().reserved_bytes > 0); assert!(!runner.acked.load(Ordering::Acquire)); let observation = observations.recv().await.expect("read observation"); assert_eq!(observation.worker_index, Some(3)); assert_eq!(observation.entries_read, 1); assert_eq!(observation.batch_size, 1); assert_eq!( runner .requested_counts .lock() .expect("requested counts lock") .as_slice(), [1] ); control.request_shutdown(); runner.release_ack.notify_one(); tokio::time::timeout(Duration::from_secs(1), handle) .await .expect("worker should finish ACK before shutdown") .expect("worker task"); assert!(runner.acked.load(Ordering::Acquire)); assert_eq!(budget.snapshot().reserved_bytes, 0); } #[tokio::test] async fn usage_worker_shutdown_cancels_reclaim_budget_wait_without_reading() { use std::future::Future; use std::task::Poll; let runner = Arc::new(ReadReclaimRaceProbeQueue::new(RuntimeQueueEntry { id: "8-0".to_string(), fields: sample_event().to_stream_fields().expect("event fields"), })); let budget = Arc::new(QueueReadBudget::new(4096, 4096)); let (_, occupied) = budget.reserve(1, 4096).await.expect("occupy budget"); let control = UsageWorkerControl::default(); let (telemetry, _observations) = tokio::sync::mpsc::channel(16); let mut worker = UsageQueueWorker::new( runner.clone(), Arc::new(SelectiveFailingRecorder::default()), receive_budget_worker_config(), None, ) .expect("worker") .with_supervisor(control.clone(), telemetry); worker.queue = worker.queue.with_read_budget(Arc::clone(&budget)); let mut cursor = "8-0".to_string(); let mut reclaim = Box::pin(worker.reclaim_stale_entries(&mut cursor)); std::future::poll_fn(|cx| { assert!(reclaim.as_mut().poll(cx).is_pending()); Poll::Ready(()) }) .await; assert_eq!(budget.snapshot().waiters, 1); assert!(runner .reclaim_cursors .lock() .expect("cursors lock") .is_empty()); control.request_shutdown(); tokio::time::timeout(Duration::from_secs(1), &mut reclaim) .await .expect("shutdown should cancel budget wait"); drop(reclaim); assert_eq!(cursor, "8-0"); assert_eq!(budget.snapshot().waiters, 0); assert_eq!(budget.snapshot().reserved_bytes, 4096); drop(occupied); assert_eq!(budget.snapshot().reserved_bytes, 0); } #[tokio::test] async fn usage_worker_shutdown_cancels_inflight_reclaim_and_releases_budget() { use std::future::Future; use std::task::Poll; let runner = Arc::new(ReadReclaimRaceProbeQueue::new(RuntimeQueueEntry { id: "9-0".to_string(), fields: sample_event().to_stream_fields().expect("event fields"), })); runner.block_reclaim.store(true, Ordering::Release); let budget = Arc::new(QueueReadBudget::new(4096, 4096)); let control = UsageWorkerControl::default(); let (telemetry, _observations) = tokio::sync::mpsc::channel(16); let mut worker = UsageQueueWorker::new( runner.clone(), Arc::new(SelectiveFailingRecorder::default()), receive_budget_worker_config(), None, ) .expect("worker") .with_supervisor(control.clone(), telemetry); worker.queue = worker.queue.with_read_budget(Arc::clone(&budget)); let mut cursor = "9-0".to_string(); let mut reclaim = Box::pin(worker.reclaim_stale_entries(&mut cursor)); std::future::poll_fn(|cx| { assert!(reclaim.as_mut().poll(cx).is_pending()); Poll::Ready(()) }) .await; assert_eq!(budget.snapshot().reserved_bytes, 4096); assert_eq!( runner .reclaim_cursors .lock() .expect("cursors lock") .as_slice(), ["9-0"] ); control.request_shutdown(); tokio::time::timeout(Duration::from_secs(1), &mut reclaim) .await .expect("shutdown should cancel reclaim I/O"); drop(reclaim); assert_eq!(cursor, "9-0"); assert!(runner.reclaim_cancelled.load(Ordering::Acquire)); assert_eq!(budget.snapshot().reserved_bytes, 0); assert!(!runner.acked.load(Ordering::Acquire)); } #[tokio::test] async fn usage_worker_shutdown_finishes_acquired_reclaim_before_releasing_budget() { let runner = Arc::new(ReadReclaimRaceProbeQueue::new(RuntimeQueueEntry { id: "10-0".to_string(), fields: sample_event().to_stream_fields().expect("event fields"), })); let budget = Arc::new(QueueReadBudget::new(4096, 4096)); let recorder = Arc::new(ControlledRecorder::default()); let control = UsageWorkerControl::default(); let (telemetry, _observations) = tokio::sync::mpsc::channel(16); let mut worker = UsageQueueWorker::new( runner.clone(), recorder.clone(), receive_budget_worker_config(), None, ) .expect("worker") .with_supervisor(control.clone(), telemetry); worker.queue = worker.queue.with_read_budget(Arc::clone(&budget)); let handle = tokio::spawn(async move { let mut cursor = "10-0".to_string(); worker.reclaim_stale_entries(&mut cursor).await; cursor }); tokio::time::timeout(Duration::from_secs(1), recorder.entered.notified()) .await .expect("reclaimed entry should reach recorder"); control.request_shutdown(); tokio::task::yield_now().await; assert!(!handle.is_finished()); assert!(budget.snapshot().reserved_bytes > 0); assert!(!runner.acked.load(Ordering::Acquire)); recorder.release.notify_one(); let cursor = tokio::time::timeout(Duration::from_secs(1), handle) .await .expect("acquired page should finish processing") .expect("reclaim task"); assert_eq!(cursor, "0-0"); assert!(runner.acked.load(Ordering::Acquire)); assert_eq!(budget.snapshot().reserved_bytes, 0); } #[tokio::test] async fn usage_worker_read_budget_processes_oversized_historical_payload() { let runner = Arc::new(RuntimeState::memory(MemoryRuntimeStateConfig::default())); let store = Arc::new(TestUsageStore::default()); let budget = Arc::new(QueueReadBudget::new(4096, 4096)); let control = UsageWorkerControl::default(); let (telemetry, mut observations) = tokio::sync::mpsc::channel(16); let mut worker = build_usage_queue_worker_with_record_gate( runner.clone(), store.clone(), receive_budget_worker_config(), None, None, ) .expect("worker") .with_supervisor(control.clone(), telemetry); worker.queue = worker.queue.with_read_budget(Arc::clone(&budget)); worker.capture_memory_budget = Arc::new(EventCaptureMemoryBudget::new(128 * 1024)); let queue = worker.queue.clone(); let mut event = sample_event(); event.data.response_body = Some(serde_json::json!({"legacy": "x".repeat(8192)})); let fields = event.to_stream_fields().expect("historical envelope"); assert!(fields.values().map(String::len).sum::() > 4096); runner .append_fields_with_maxlen(&worker.config.stream_key, &fields, None) .await .expect("append pre-limit message"); let handle = tokio::spawn(worker.run()); tokio::time::timeout(Duration::from_secs(1), async { while observations .recv() .await .expect("worker observation") .acked_entries == 0 {} }) .await .expect("oversized message should be recorded and ACKed"); control.request_shutdown(); tokio::time::timeout(Duration::from_secs(1), handle) .await .expect("worker should stop") .expect("worker task"); { let records = store.records.lock().expect("records lock"); assert_eq!(records.len(), 1); assert_eq!(records[0].response_body, event.data.response_body); assert_eq!(records[0].total_tokens, Some(10)); } let snapshot = budget.snapshot(); assert_eq!(snapshot.reserved_bytes, 0); assert_eq!(snapshot.oversized_entries_total, 1); assert_eq!(snapshot.oversized_batches_total, 1); let stats = queue.stats().await.expect("acked stats"); assert_eq!((stats.stream_length, stats.group_pending), (0, 0)); assert_eq!(queue.dlq_stats().await.expect("dlq stats").stream_length, 0); } #[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 capture_budget_retry_releases_decoded_lease_and_preserves_pending_payload() { let (_runner, worker, recorder) = capture_budget_worker(64 * 1024, CaptureRecordOutcome::RetryOnce).await; worker .queue .enqueue(&captured_worker_event()) .await .expect("enqueue"); let entries = worker .queue .read_group(&worker.consumer) .await .expect("read event"); assert_eq!(entries.len(), 1); let retry_entries = entries.clone(); assert!(matches!( worker.process_entries(entries).await, Err(DataLayerError::TimedOut(_)) )); assert_eq!(recorder.budget.retained_bytes(), 0); let stats = worker.queue.stats().await.expect("pending stats"); assert_eq!(stats.stream_length, 1); assert_eq!(stats.group_pending, 1); assert_eq!( worker .queue .dlq_stats() .await .expect("dlq stats") .stream_length, 0 ); // Replay the same pending entry, as reclamation does, without a timing-dependent idle wait. worker .process_entries(retry_entries) .await .expect("retry should succeed"); assert_eq!(recorder.calls.load(Ordering::Acquire), 2); assert_eq!(recorder.budget.retained_bytes(), 0); let stats = worker.queue.stats().await.expect("ack stats"); assert_eq!(stats.stream_length, 0); assert_eq!(stats.group_pending, 0); } #[tokio::test] async fn capture_budget_downgrade_dead_letter_keeps_exact_original_fields() { let (runner, worker, recorder) = capture_budget_worker(0, CaptureRecordOutcome::PermanentFailure).await; let mut original = captured_worker_event() .to_stream_fields() .expect("wire serialization"); original.insert( "legacy_marker".to_string(), "preserve this field".to_string(), ); runner .append_fields_with_maxlen(&worker.config.stream_key, &original, None) .await .expect("enqueue raw fields"); let entries = worker .queue .read_group(&worker.consumer) .await .expect("read event"); worker .process_entries(entries) .await .expect("permanent failure should dead letter"); assert_eq!(recorder.calls.load(Ordering::Acquire), 1); assert_eq!(recorder.budget.retained_bytes(), 0); assert_eq!(recorder.budget.downgraded_total(), 1); let stats = worker.queue.stats().await.expect("ack stats"); assert_eq!(stats.stream_length, 0); assert_eq!(stats.group_pending, 0); runner .ensure_consumer_group( &worker.config.dlq_stream_key, "capture-dlq-inspection", "0-0", ) .await .expect("dlq group"); let dlq = runner .read_group( &worker.config.dlq_stream_key, "capture-dlq-inspection", "capture-inspector", 1, Some(1), ) .await .expect("read dlq"); assert_eq!(dlq.len(), 1); let payload: serde_json::Value = serde_json::from_str(&dlq[0].fields["payload"]).expect("dlq json"); assert_eq!( payload["fields"], serde_json::to_value(original).expect("original fields json") ); assert_eq!( payload["error"].as_str(), Some("unexpected database value: permanent capture test error") ); } #[tokio::test] async fn capture_budget_malformed_entry_dead_letters_without_recording() { let (runner, worker, recorder) = capture_budget_worker(1024, CaptureRecordOutcome::PermanentFailure).await; let original = BTreeMap::from([( "payload".to_string(), "malformed legacy payload".to_string(), )]); runner .append_fields_with_maxlen(&worker.config.stream_key, &original, None) .await .expect("enqueue raw fields"); let entries = worker .queue .read_group(&worker.consumer) .await .expect("read event"); worker .process_entries(entries) .await .expect("malformed event should dead letter"); assert_eq!(recorder.calls.load(Ordering::Acquire), 0); assert_eq!(recorder.budget.retained_bytes(), 0); assert_eq!(recorder.budget.downgraded_total(), 0); let stats = worker.queue.stats().await.expect("ack stats"); assert_eq!(stats.stream_length, 0); assert_eq!(stats.group_pending, 0); runner .ensure_consumer_group( &worker.config.dlq_stream_key, "capture-dlq-inspection", "0-0", ) .await .expect("dlq group"); let dlq = runner .read_group( &worker.config.dlq_stream_key, "capture-dlq-inspection", "capture-inspector", 1, Some(1), ) .await .expect("read dlq"); assert_eq!(dlq.len(), 1); let payload: serde_json::Value = serde_json::from_str(&dlq[0].fields["payload"]).expect("dlq json"); assert_eq!( payload["fields"], serde_json::to_value(original).expect("original fields json") ); } #[tokio::test] async fn capture_budget_cancelled_record_releases_lease_without_acknowledging() { let (_runner, worker, recorder) = capture_budget_worker(64 * 1024, CaptureRecordOutcome::Wait).await; worker .queue .enqueue(&captured_worker_event()) .await .expect("enqueue"); let entries = worker .queue .read_group(&worker.consumer) .await .expect("read event"); let worker = Arc::new(worker); let task_worker = Arc::clone(&worker); let task = tokio::spawn(async move { task_worker.process_entries(entries).await }); tokio::time::timeout(Duration::from_secs(1), recorder.entered.notified()) .await .expect("recorder should start"); assert!(recorder.budget.retained_bytes() > 0); task.abort(); assert!(task .await .expect_err("task should be cancelled") .is_cancelled()); assert_eq!(recorder.budget.retained_bytes(), 0); let stats = worker.queue.stats().await.expect("pending stats"); assert_eq!(stats.stream_length, 1); assert_eq!(stats.group_pending, 1); } #[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 = 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() }; let mut worker = UsageQueueWorker::new(queue_runner, recorder.clone(), config, None) .expect("worker should build"); worker.queue = worker .queue .with_dead_letter_encoding_budget(Arc::new(DeadLetterEncodingBudget::new( 64 * 1024 * 1024, 4, ))); 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")); } }