chore: update gateway pressure observability

This commit is contained in:
elky
2026-06-30 17:01:39 +08:00
parent 974def5fef
commit f179ee72f9
69 changed files with 14843 additions and 1609 deletions
+21 -8
View File
@@ -8,6 +8,7 @@ pub struct UsageRuntimeConfig {
pub worker_count: usize,
pub worker_autoscale_enabled: bool,
pub worker_max_count: usize,
pub worker_record_concurrency_limit: Option<usize>,
pub worker_scale_interval_ms: u64,
pub worker_idle_scale_down_ticks: u64,
pub stream_key: String,
@@ -21,6 +22,7 @@ pub struct UsageRuntimeConfig {
pub reclaim_interval_ms: u64,
pub terminal_enqueue_max_in_flight: u64,
pub lifecycle_enqueue_max_in_flight: u64,
pub lifecycle_enqueue_delay_ms: u64,
pub retry_deferred_lifecycle_events: bool,
pub enqueue_retry_buffer_capacity: usize,
pub enqueue_retry_workers: usize,
@@ -36,23 +38,25 @@ impl Default for UsageRuntimeConfig {
queue_lifecycle_events: false,
worker_count: 4,
worker_autoscale_enabled: true,
worker_max_count: 64,
worker_max_count: 32,
worker_record_concurrency_limit: Some(32),
worker_scale_interval_ms: 1_000,
worker_idle_scale_down_ticks: 30,
stream_key: "usage:events".to_string(),
consumer_group: "usage_consumers".to_string(),
dlq_stream_key: "usage:events:dlq".to_string(),
stream_maxlen: 200_000,
consumer_batch_size: 500,
consumer_batch_size: 128,
consumer_block_ms: 500,
reclaim_idle_ms: 30_000,
reclaim_count: 500,
reclaim_idle_ms: 60_000,
reclaim_count: 128,
reclaim_interval_ms: 5_000,
terminal_enqueue_max_in_flight: 256,
lifecycle_enqueue_max_in_flight: 128,
retry_deferred_lifecycle_events: false,
terminal_enqueue_max_in_flight: 1_024,
lifecycle_enqueue_max_in_flight: 512,
lifecycle_enqueue_delay_ms: 1_000,
retry_deferred_lifecycle_events: true,
enqueue_retry_buffer_capacity: 131_072,
enqueue_retry_workers: 4,
enqueue_retry_workers: 8,
enqueue_retry_initial_backoff_ms: 3_000,
enqueue_retry_max_backoff_ms: 10_000,
}
@@ -94,6 +98,15 @@ impl UsageRuntimeConfig {
"usage runtime worker_max_count must be positive".to_string(),
));
}
if self
.worker_record_concurrency_limit
.is_some_and(|limit| limit == 0)
{
return Err(DataLayerError::InvalidConfiguration(
"usage runtime worker_record_concurrency_limit must be positive when set"
.to_string(),
));
}
if self.worker_scale_interval_ms == 0 {
return Err(DataLayerError::InvalidConfiguration(
"usage runtime worker_scale_interval_ms must be positive".to_string(),
+35 -3
View File
@@ -1,7 +1,10 @@
use std::future::Future;
use std::sync::OnceLock;
const USAGE_BACKGROUND_RUNTIME_THREADS: usize = 2;
const DEFAULT_USAGE_BACKGROUND_RUNTIME_THREADS: usize = 8;
const MAX_USAGE_BACKGROUND_RUNTIME_THREADS: usize = 64;
const GATEWAY_USAGE_BACKGROUND_RUNTIME_THREADS_ENV: &str = "AETHER_GATEWAY_USAGE_RUNTIME_THREADS";
const USAGE_BACKGROUND_RUNTIME_THREADS_ENV: &str = "AETHER_USAGE_RUNTIME_THREADS";
const USAGE_BACKGROUND_RUNTIME_STACK_BYTES: usize = 8 * 1024 * 1024;
const USAGE_BACKGROUND_RUNTIME_THREAD_NAME: &str = "aether-usage-runtime";
@@ -19,7 +22,7 @@ fn usage_background_runtime() -> &'static tokio::runtime::Runtime {
RUNTIME.get_or_init(|| {
let runtime = tokio::runtime::Builder::new_multi_thread()
.enable_all()
.worker_threads(USAGE_BACKGROUND_RUNTIME_THREADS)
.worker_threads(usage_background_runtime_threads())
.thread_name(USAGE_BACKGROUND_RUNTIME_THREAD_NAME)
.thread_stack_size(USAGE_BACKGROUND_RUNTIME_STACK_BYTES)
.build()
@@ -28,9 +31,26 @@ fn usage_background_runtime() -> &'static tokio::runtime::Runtime {
})
}
fn usage_background_runtime_threads() -> usize {
parse_usage_background_runtime_threads(
std::env::var(GATEWAY_USAGE_BACKGROUND_RUNTIME_THREADS_ENV)
.ok()
.or_else(|| std::env::var(USAGE_BACKGROUND_RUNTIME_THREADS_ENV).ok())
.as_deref(),
)
}
fn parse_usage_background_runtime_threads(value: Option<&str>) -> usize {
value
.and_then(|value| value.trim().parse::<usize>().ok())
.filter(|threads| *threads > 0)
.unwrap_or(DEFAULT_USAGE_BACKGROUND_RUNTIME_THREADS)
.clamp(1, MAX_USAGE_BACKGROUND_RUNTIME_THREADS)
}
#[cfg(test)]
mod tests {
use super::spawn_on_usage_background_runtime;
use super::{parse_usage_background_runtime_threads, spawn_on_usage_background_runtime};
#[tokio::test]
async fn usage_background_runtime_runs_on_dedicated_named_threads() {
@@ -45,4 +65,16 @@ mod tests {
assert_eq!(thread_name, "aether-usage-runtime");
}
#[test]
fn usage_background_runtime_threads_are_configurable() {
assert_eq!(parse_usage_background_runtime_threads(None), 8);
assert_eq!(parse_usage_background_runtime_threads(Some("12")), 12);
assert_eq!(parse_usage_background_runtime_threads(Some("0")), 8);
assert_eq!(
parse_usage_background_runtime_threads(Some("not-a-number")),
8
);
assert_eq!(parse_usage_background_runtime_threads(Some("999")), 64);
}
}
+2 -2
View File
@@ -40,8 +40,8 @@ pub use report_context::{
build_locally_actionable_report_context_from_video_task, report_context_is_locally_actionable,
};
pub use runtime::{
UsageBillingEventEnricher, UsageBodyCapturePolicy, UsageRequestRecordLevel, UsageRuntime,
UsageRuntimeAccess, UsageRuntimeMetricsSnapshot,
UsageBillingEventEnricher, UsageBodyCapturePolicy, UsageQueueHealthSnapshot,
UsageRequestRecordLevel, UsageRuntime, UsageRuntimeAccess, UsageRuntimeMetricsSnapshot,
DEFAULT_USAGE_REQUEST_BODY_CAPTURE_LIMIT_BYTES,
DEFAULT_USAGE_RESPONSE_BODY_CAPTURE_LIMIT_BYTES,
};
+11 -1
View File
@@ -3,7 +3,9 @@ use std::sync::Arc;
use serde_json::json;
use aether_data_contracts::DataLayerError;
use aether_runtime_state::{RuntimeQueueEntry, RuntimeQueueReclaimConfig, RuntimeQueueStore};
use aether_runtime_state::{
RuntimeQueueEntry, RuntimeQueueReclaimConfig, RuntimeQueueStats, RuntimeQueueStore,
};
use super::config::UsageRuntimeConfig;
use super::event::UsageEvent;
@@ -103,6 +105,14 @@ impl UsageQueue {
.append_fields_with_maxlen(&self.dlq_stream, &fields, None)
.await
}
pub async fn stats(&self) -> Result<RuntimeQueueStats, DataLayerError> {
self.runner.stats(&self.stream, Some(&self.group)).await
}
pub async fn dlq_stats(&self) -> Result<RuntimeQueueStats, DataLayerError> {
self.runner.stats(&self.dlq_stream, None).await
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
File diff suppressed because it is too large Load Diff
@@ -1,3 +1,7 @@
use std::collections::hash_map::DefaultHasher;
use std::hash::{Hash, Hasher};
use std::sync::OnceLock;
use aether_data_contracts::repository::settlement::{StoredUsageSettlement, UsageSettlementInput};
use aether_data_contracts::repository::usage::StoredRequestUsageAudit;
use aether_data_contracts::{DataLayerError, DataLayerError::InvalidInput};
@@ -39,10 +43,50 @@ pub async fn settle_usage_if_needed(
actual_total_cost_usd: finite_cost(usage.actual_total_cost_usd)?,
finalized_at_unix_secs,
};
let settlement_key = usage_settlement_lock_key(&input);
let _guard = usage_settlement_lock(&settlement_key).lock().await;
let _ = writer.settle_usage(input).await?;
Ok(())
}
const USAGE_SETTLEMENT_LOCK_SHARDS: usize = 4096;
fn usage_settlement_lock(key: &str) -> &'static tokio::sync::Mutex<()> {
static LOCKS: OnceLock<Vec<tokio::sync::Mutex<()>>> = OnceLock::new();
let locks = LOCKS.get_or_init(|| {
(0..USAGE_SETTLEMENT_LOCK_SHARDS)
.map(|_| tokio::sync::Mutex::new(()))
.collect()
});
&locks[usage_settlement_lock_shard(key)]
}
fn usage_settlement_lock_shard(key: &str) -> usize {
let mut hasher = DefaultHasher::new();
key.hash(&mut hasher);
(hasher.finish() as usize) % USAGE_SETTLEMENT_LOCK_SHARDS
}
fn usage_settlement_lock_key(input: &UsageSettlementInput) -> String {
if input.api_key_is_standalone {
if let Some(api_key_id) = input.api_key_id.as_deref().and_then(non_empty_trimmed) {
return format!("api-key:{api_key_id}");
}
}
if let Some(user_id) = input.user_id.as_deref().and_then(non_empty_trimmed) {
return format!("user:{user_id}");
}
if let Some(api_key_id) = input.api_key_id.as_deref().and_then(non_empty_trimmed) {
return format!("api-key:{api_key_id}");
}
format!("request:{}", input.request_id.trim())
}
fn non_empty_trimmed(value: &str) -> Option<&str> {
let value = value.trim();
(!value.is_empty()).then_some(value)
}
fn usage_api_key_is_standalone(usage: &StoredRequestUsageAudit) -> bool {
usage
.request_metadata
@@ -64,7 +108,9 @@ fn finite_cost(value: f64) -> Result<f64, DataLayerError> {
#[cfg(test)]
mod tests {
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::Mutex;
use std::time::Duration;
use super::{settle_usage_if_needed, UsageSettlementWriter};
use aether_data_contracts::repository::settlement::UsageSettlementInput;
@@ -78,6 +124,13 @@ mod tests {
inputs: Mutex<Vec<UsageSettlementInput>>,
}
#[derive(Default)]
struct SlowSettlementWriter {
active: AtomicUsize,
max_active: AtomicUsize,
inputs: Mutex<Vec<UsageSettlementInput>>,
}
#[async_trait]
impl UsageSettlementWriter for TestSettlementWriter {
fn has_usage_settlement_writer(&self) -> bool {
@@ -99,6 +152,31 @@ mod tests {
}
}
#[async_trait]
impl UsageSettlementWriter for SlowSettlementWriter {
fn has_usage_settlement_writer(&self) -> bool {
true
}
async fn settle_usage(
&self,
input: UsageSettlementInput,
) -> Result<
Option<aether_data_contracts::repository::settlement::StoredUsageSettlement>,
aether_data_contracts::DataLayerError,
> {
let active = self.active.fetch_add(1, Ordering::AcqRel) + 1;
self.max_active.fetch_max(active, Ordering::AcqRel);
tokio::time::sleep(Duration::from_millis(30)).await;
self.inputs
.lock()
.expect("settlement inputs lock")
.push(input);
self.active.fetch_sub(1, Ordering::AcqRel);
Ok(None)
}
}
fn sample_usage() -> StoredRequestUsageAudit {
StoredRequestUsageAudit::new(
"usage-1".to_string(),
@@ -238,4 +316,22 @@ mod tests {
let inputs = writer.inputs.lock().expect("settlement inputs lock");
assert!(inputs.is_empty());
}
#[tokio::test]
async fn serializes_settlements_for_same_billing_subject() {
let writer = SlowSettlementWriter::default();
let mut first = sample_usage();
first.request_id = "req-same-subject-1".to_string();
let mut second = sample_usage();
second.request_id = "req-same-subject-2".to_string();
tokio::try_join!(
settle_usage_if_needed(&writer, &first),
settle_usage_if_needed(&writer, &second)
)
.expect("settlements should succeed");
assert_eq!(writer.max_active.load(Ordering::Acquire), 1);
assert_eq!(writer.inputs.lock().expect("inputs lock").len(), 2);
}
}
+471 -33
View File
@@ -1,5 +1,7 @@
use std::collections::hash_map::DefaultHasher;
use std::hash::{Hash, Hasher};
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::Arc;
use std::sync::{Arc, OnceLock};
use std::time::Duration;
use aether_data_contracts::repository::usage::{StoredRequestUsageAudit, UpsertUsageRecord};
@@ -10,12 +12,17 @@ use tokio::sync::mpsc;
use tracing::warn;
use crate::executor::spawn_on_usage_background_runtime;
use crate::runtime::UsageBillingEventEnricher;
use crate::runtime::{
UsageBillingEventEnricher, UsageRuntimeAccess, UsageWorkerRecordConcurrencyGate,
};
use crate::{
build_upsert_usage_record_from_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>;
@@ -42,31 +49,81 @@ pub trait UsageRecordWriter: Send + Sync {
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 { data }
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: UsageRecordWriter
+ UsageSettlementWriter
+ UsageBillingEventEnricher
+ ManualProxyNodeCounter
+ Send
+ Sync,
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 _guard = usage_request_lock(&event.request_id).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
}
}
const USAGE_REQUEST_LOCK_SHARDS: usize = 4096;
fn usage_request_lock(request_id: &str) -> &'static tokio::sync::Mutex<()> {
static LOCKS: OnceLock<Vec<tokio::sync::Mutex<()>>> = OnceLock::new();
let locks = LOCKS.get_or_init(|| {
(0..USAGE_REQUEST_LOCK_SHARDS)
.map(|_| tokio::sync::Mutex::new(()))
.collect()
});
&locks[usage_request_lock_shard(request_id)]
}
fn usage_request_lock_shard(request_id: &str) -> usize {
let mut hasher = DefaultHasher::new();
request_id.hash(&mut hasher);
(hasher.finish() as usize) % USAGE_REQUEST_LOCK_SHARDS
}
pub struct UsageQueueWorker {
queue: UsageQueue,
recorder: Arc<dyn UsageEventRecorder>,
@@ -97,6 +154,112 @@ 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,
}
}
}
impl UsageQueueWorker {
@@ -163,7 +326,9 @@ impl UsageQueueWorker {
_ = reclaim_interval.tick() => {
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",
@@ -174,14 +339,17 @@ impl UsageQueueWorker {
);
}
}
Err(err) => 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"
),
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"
);
}
}
}
result = self.queue.read_group(&self.consumer) => {
@@ -192,6 +360,7 @@ impl UsageQueueWorker {
break;
}
if let Err(err) = self.process_entries(entries).await {
self.report_process_failed();
warn!(
event_name = "usage_worker_process_failed",
log_type = "ops",
@@ -207,6 +376,7 @@ impl UsageQueueWorker {
}
}
Err(err) => {
self.report_read_failed();
warn!(
event_name = "usage_worker_read_failed",
log_type = "ops",
@@ -230,14 +400,51 @@ impl UsageQueueWorker {
}
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) {
let Some(telemetry) = &self.telemetry else {
return;
};
let _ = telemetry.try_send(UsageWorkerObservation {
worker_index: self.worker_index,
entries_read,
batch_size: self.config.consumer_batch_size.max(1),
});
let _ = telemetry.try_send(observation);
}
async fn process_entries(&self, entries: Vec<RuntimeQueueEntry>) -> Result<(), DataLayerError> {
@@ -251,6 +458,11 @@ impl UsageQueueWorker {
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) => {
@@ -264,6 +476,7 @@ impl UsageQueueWorker {
if !ack_ids.is_empty() {
self.queue.ack_and_delete(&ack_ids).await?;
self.report_acked(ack_ids.len());
}
Ok(())
@@ -283,6 +496,7 @@ impl UsageQueueWorker {
"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);
}
};
@@ -308,6 +522,7 @@ impl UsageQueueWorker {
"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);
Ok(true)
}
Err(err) => {
@@ -357,17 +572,26 @@ pub fn build_usage_queue_worker<T>(
worker_index: Option<usize>,
) -> Result<UsageQueueWorker, DataLayerError>
where
T: UsageRecordWriter
+ UsageSettlementWriter
+ UsageBillingEventEnricher
+ ManualProxyNodeCounter
+ Send
+ Sync
+ 'static,
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,
{
UsageQueueWorker::new(
runner,
Arc::new(UsageDataEventRecorder::new(data)),
Arc::new(
UsageDataEventRecorder::with_record_gate_and_database_pressure_defer(data, record_gate),
),
config,
worker_index,
)
@@ -472,7 +696,9 @@ fn consumer_name(worker_index: Option<usize>) -> String {
#[cfg(test)]
mod tests {
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{Arc, Mutex};
use std::time::Duration;
use aether_data_contracts::repository::settlement::{
StoredUsageSettlement, UsageSettlementInput,
@@ -483,12 +709,15 @@ mod tests {
use async_trait::async_trait;
use super::{
usage_event_record_error_is_permanent, write_event_record, ManualProxyNodeCounter,
UsageEventRecorder, UsageQueueWorker, UsageRecordWriter,
build_usage_queue_worker_with_record_gate, usage_event_record_error_is_permanent,
write_event_record, ManualProxyNodeCounter, UsageEventRecorder, UsageQueueWorker,
UsageRecordWriter,
};
use crate::runtime::UsageWorkerRecordConcurrencyGate;
use crate::UsageBillingEventEnricher;
use crate::{
UsageEvent, UsageEventData, UsageEventType, UsageRuntimeConfig, UsageSettlementWriter,
UsageEvent, UsageEventData, UsageEventType, UsageQueue, UsageRuntimeConfig,
UsageSettlementWriter,
};
#[derive(Default)]
@@ -503,6 +732,14 @@ mod tests {
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>>,
}
#[async_trait]
impl UsageRecordWriter for TestUsageStore {
async fn upsert_usage_record(
@@ -603,6 +840,95 @@ mod tests {
}
}
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)
}
}
#[async_trait]
impl UsageEventRecorder for SelectiveFailingRecorder {
async fn record_usage_event(&self, event: &UsageEvent) -> Result<(), DataLayerError> {
@@ -710,6 +1036,118 @@ mod tests {
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);
}
#[test]
fn usage_event_record_error_classifies_permanent_failures() {
assert!(usage_event_record_error_is_permanent(