mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-12 06:00:20 +08:00
chore: update gateway pressure observability
This commit is contained in:
@@ -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(),
|
||||
|
||||
@@ -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);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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,
|
||||
};
|
||||
|
||||
@@ -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);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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(
|
||||
|
||||
Reference in New Issue
Block a user