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
+11 -23
View File
@@ -61,28 +61,16 @@ ADMIN_USERNAME=admin123456
# docker compose 下 app 启动前自动执行 pending migration/backfill(默认 true
# AETHER_GATEWAY_AUTO_PREPARE_DATABASE=true
# PostgreSQL 连接池配置(默认适合单实例/小型部署;高并发可按需调大
# 推荐计算方式(单实例):
# MAX = CPU 核数 × 10(AI 网关偏 IO 等待,可激进些;纯 OLTP 用 × 4)
# MIN = MAX × 0.2(保留常驻连接应对突发流量,避免冷启动握手开销)
# 多实例部署时请按 实例数 × MAX 控制总和,PG 端 max_connections 至少为该总和 + 20 余量
# AETHER_GATEWAY_DATA_POSTGRES_MIN_CONNECTIONS=4
# AETHER_GATEWAY_DATA_POSTGRES_MAX_CONNECTIONS=20
# AETHER_GATEWAY_DATA_POSTGRES_STATEMENT_CACHE_CAPACITY=100
# AETHER_GATEWAY_DATA_POSTGRES_ACQUIRE_TIMEOUT_MS=3000
# PostgreSQL 连接池配置(默认按 CPU 自动计算;正式高并发环境可显式预算
# AETHER_GATEWAY_DATA_POSTGRES_MIN_CONNECTIONS=12
# AETHER_GATEWAY_DATA_POSTGRES_MAX_CONNECTIONS=80
# PostgreSQL 性能调优(默认值适合 2核4GB 机器,按实际配置覆盖)
# 参考:shared_buffers ≈ 可用内存 25%effective_cache_size ≈ 可用内存 50-75%
# POSTGRES_SHM_SIZE 控制 Docker 容器 /dev/shm;仪表盘统计等并行查询会使用它。
# work_mem 是每个连接每个排序操作的内存,不要设太大(并发数 × work_mem 是实际占用)
# | 系统内存 | shared_buffers | effective_cache_size | work_mem |
# | 2GB | 256MB | 768MB | 4MB |
# | 4GB | 1GB | 3GB | 16MB |
# | 8GB | 2GB | 6GB | 16MB |
# | 16GB | 4GB | 12GB | 32MB |
# | 32GB+ | 8GB | 24GB | 32MB |
# POSTGRES_SHARED_BUFFERS=1GB
# POSTGRES_EFFECTIVE_CACHE_SIZE=3GB
# POSTGRES_SHM_SIZE=512mb
# PostgreSQL 容器调优:docker-compose.yml 已内置通用默认值,通常不用配置。
# 只有在 Postgres 独占大内存、或压测显示 DB 缓存/排序/维护任务成为瓶颈时再覆盖。
# 内置默认:shared_buffers=1GB, effective_cache_size=3GB, shm_size=512mb,
# work_mem=16MB, maintenance_work_mem=256MB。
# POSTGRES_SHARED_BUFFERS=8GB
# POSTGRES_EFFECTIVE_CACHE_SIZE=24GB
# POSTGRES_SHM_SIZE=2gb
# POSTGRES_WORK_MEM=16MB
# POSTGRES_MAINTENANCE_WORK_MEM=256MB
# POSTGRES_MAINTENANCE_WORK_MEM=1GB
Generated
+4
View File
@@ -276,8 +276,10 @@ dependencies = [
"sha2",
"socket2 0.6.3",
"sqlx",
"sysinfo",
"tar",
"thiserror 2.0.18",
"tikv-jemalloc-sys",
"tikv-jemallocator",
"tokio",
"tokio-util",
@@ -477,10 +479,12 @@ dependencies = [
"reqwest",
"serde",
"serde_json",
"sha2",
"sqlx",
"sysinfo",
"tokio",
"tokio-tungstenite 0.28.0",
"uuid",
]
[[package]]
+8 -1
View File
@@ -8,7 +8,12 @@ description = "Rust ingress gateway for Aether phase 3a transparent proxy"
[features]
default = []
jemalloc = ["dep:tikv-jemallocator"]
jemalloc = [
"dep:tikv-jemallocator",
"dep:tikv-jemalloc-sys",
"tikv-jemallocator/stats",
"tikv-jemalloc-sys/stats",
]
testkit = []
[dependencies]
@@ -70,6 +75,7 @@ sha2 = { workspace = true, features = ["oid"] }
socket2.workspace = true
tar.workspace = true
sqlx.workspace = true
sysinfo = "0.32"
thiserror.workspace = true
tokio.workspace = true
tokio-util.workspace = true
@@ -85,6 +91,7 @@ zstd.workspace = true
[target.'cfg(not(target_env = "msvc"))'.dependencies]
tikv-jemallocator = { version = "0.6", optional = true }
tikv-jemalloc-sys = { version = "0.6", optional = true }
[dev-dependencies]
aether-testkit.workspace = true
@@ -0,0 +1,218 @@
use aether_runtime::{MetricKind, MetricSample};
pub(crate) fn gateway_allocator_metric_samples() -> Vec<MetricSample> {
match allocator_snapshot() {
Some(snapshot) => snapshot.to_metric_samples(),
None => unavailable_metric_samples(),
}
}
#[derive(Debug, Clone, Copy, Default)]
struct AllocatorSnapshot {
allocated_bytes: u64,
active_bytes: u64,
resident_bytes: u64,
mapped_bytes: u64,
retained_bytes: u64,
metadata_bytes: u64,
}
impl AllocatorSnapshot {
fn to_metric_samples(self) -> Vec<MetricSample> {
vec![
gauge(
"gateway_allocator_observability_available",
"Whether gateway allocator heap metrics were available for this scrape.",
1,
),
gauge(
"gateway_allocator_allocated_bytes",
"Bytes currently allocated by the gateway allocator.",
self.allocated_bytes,
),
gauge(
"gateway_allocator_active_bytes",
"Bytes in active pages managed by the gateway allocator.",
self.active_bytes,
),
gauge(
"gateway_allocator_resident_bytes",
"Bytes resident in physical memory for the gateway allocator.",
self.resident_bytes,
),
gauge(
"gateway_allocator_mapped_bytes",
"Bytes mapped by the gateway allocator.",
self.mapped_bytes,
),
gauge(
"gateway_allocator_retained_bytes",
"Bytes retained by the gateway allocator for future use.",
self.retained_bytes,
),
gauge(
"gateway_allocator_metadata_bytes",
"Bytes used for allocator metadata.",
self.metadata_bytes,
),
gauge(
"gateway_allocator_active_to_allocated_basis_points",
"Active allocator bytes divided by allocated bytes in basis points.",
ratio_basis_points(self.active_bytes, self.allocated_bytes),
),
gauge(
"gateway_allocator_resident_to_allocated_basis_points",
"Resident allocator bytes divided by allocated bytes in basis points.",
ratio_basis_points(self.resident_bytes, self.allocated_bytes),
),
]
}
}
fn unavailable_metric_samples() -> Vec<MetricSample> {
vec![
gauge(
"gateway_allocator_observability_available",
"Whether gateway allocator heap metrics were available for this scrape.",
0,
),
gauge(
"gateway_allocator_allocated_bytes",
"Bytes currently allocated by the gateway allocator.",
0,
),
gauge(
"gateway_allocator_active_bytes",
"Bytes in active pages managed by the gateway allocator.",
0,
),
gauge(
"gateway_allocator_resident_bytes",
"Bytes resident in physical memory for the gateway allocator.",
0,
),
gauge(
"gateway_allocator_mapped_bytes",
"Bytes mapped by the gateway allocator.",
0,
),
gauge(
"gateway_allocator_retained_bytes",
"Bytes retained by the gateway allocator for future use.",
0,
),
gauge(
"gateway_allocator_metadata_bytes",
"Bytes used for allocator metadata.",
0,
),
gauge(
"gateway_allocator_active_to_allocated_basis_points",
"Active allocator bytes divided by allocated bytes in basis points.",
0,
),
gauge(
"gateway_allocator_resident_to_allocated_basis_points",
"Resident allocator bytes divided by allocated bytes in basis points.",
0,
),
]
}
#[cfg(all(feature = "jemalloc", not(target_env = "msvc")))]
fn allocator_snapshot() -> Option<AllocatorSnapshot> {
refresh_jemalloc_epoch()?;
Some(AllocatorSnapshot {
allocated_bytes: read_jemalloc_stat("stats.allocated\0")?,
active_bytes: read_jemalloc_stat("stats.active\0")?,
resident_bytes: read_jemalloc_stat("stats.resident\0")?,
mapped_bytes: read_jemalloc_stat("stats.mapped\0")?,
retained_bytes: read_jemalloc_stat("stats.retained\0")?,
metadata_bytes: read_jemalloc_stat("stats.metadata\0")?,
})
}
#[cfg(not(all(feature = "jemalloc", not(target_env = "msvc"))))]
fn allocator_snapshot() -> Option<AllocatorSnapshot> {
None
}
#[cfg(all(feature = "jemalloc", not(target_env = "msvc")))]
fn refresh_jemalloc_epoch() -> Option<()> {
let mut epoch = 1_u64;
let result = unsafe {
tikv_jemalloc_sys::mallctl(
c"epoch".as_ptr(),
std::ptr::null_mut(),
std::ptr::null_mut(),
(&mut epoch as *mut u64).cast(),
std::mem::size_of::<u64>(),
)
};
if result == 0 {
Some(())
} else {
None
}
}
#[cfg(all(feature = "jemalloc", not(target_env = "msvc")))]
fn read_jemalloc_stat(name: &str) -> Option<u64> {
let mut value = 0_usize;
let mut size = std::mem::size_of::<usize>();
let result = unsafe {
tikv_jemalloc_sys::mallctl(
name.as_ptr().cast(),
(&mut value as *mut usize).cast(),
&mut size,
std::ptr::null_mut(),
0,
)
};
if result == 0 {
Some(u64_from_usize(value))
} else {
None
}
}
fn gauge(name: &'static str, help: &'static str, value: u64) -> MetricSample {
MetricSample::new(name, help, MetricKind::Gauge, value)
}
fn ratio_basis_points(numerator: u64, denominator: u64) -> u64 {
if denominator == 0 {
return 0;
}
numerator.saturating_mul(10_000) / denominator
}
fn u64_from_usize(value: usize) -> u64 {
u64::try_from(value).unwrap_or(u64::MAX)
}
#[cfg(test)]
mod tests {
use super::{gateway_allocator_metric_samples, ratio_basis_points};
#[test]
fn renders_allocator_metric_samples() {
let samples = gateway_allocator_metric_samples();
assert!(samples
.iter()
.any(|sample| sample.name == "gateway_allocator_observability_available"));
assert!(samples
.iter()
.any(|sample| sample.name == "gateway_allocator_allocated_bytes"));
assert!(samples
.iter()
.any(|sample| sample.name == "gateway_allocator_active_to_allocated_basis_points"));
}
#[test]
fn computes_ratio_basis_points() {
assert_eq!(ratio_basis_points(150, 100), 15_000);
assert_eq!(ratio_basis_points(1, 0), 0);
}
}
@@ -252,6 +252,12 @@ impl UsageRuntimeAccess for GatewayDataState {
GatewayDataState::usage_worker_queue(self)
}
fn usage_worker_should_defer_for_database_pressure(&self) -> bool {
self.database_pool_summary()
.as_ref()
.is_some_and(GatewayDataState::database_pool_summary_under_usage_worker_pressure)
}
async fn body_capture_policy(&self) -> Result<UsageBodyCapturePolicy, DataLayerError> {
let value = match GatewayDataState::find_system_config_value(self, REQUEST_RECORD_LEVEL_KEY)
.await?
@@ -170,6 +170,25 @@ impl GatewayDataState {
.and_then(|backends| backends.database_pool_summary())
}
pub(crate) async fn postgres_observability_snapshot(
&self,
) -> Result<Option<aether_data::DatabasePostgresObservabilitySnapshot>, DataLayerError> {
match &self.backends {
Some(backends) => backends.postgres_observability_snapshot().await,
None => Ok(None),
}
}
pub(crate) async fn postgres_activity_groups(
&self,
limit: i64,
) -> Result<Vec<aether_data::DatabasePostgresActivityGroup>, DataLayerError> {
match &self.backends {
Some(backends) => backends.postgres_activity_groups(limit).await,
None => Ok(Vec::new()),
}
}
pub(crate) fn database_pool_under_maintenance_pressure(&self) -> bool {
self.database_pool_summary()
.as_ref()
@@ -182,6 +201,14 @@ impl GatewayDataState {
summary.checked_out > 0 && summary.idle <= Self::maintenance_pool_idle_reserve(summary)
}
pub(crate) fn database_pool_summary_under_usage_worker_pressure(
summary: &aether_data::DatabasePoolSummary,
) -> bool {
summary.checked_out > 0
&& (summary.checked_out >= summary.max_connections as usize
|| summary.idle <= Self::usage_worker_pool_idle_reserve(summary))
}
pub(crate) fn maintenance_pool_idle_reserve(
summary: &aether_data::DatabasePoolSummary,
) -> usize {
@@ -201,6 +228,13 @@ impl GatewayDataState {
ten_percent_ceil.clamp(2, 10).min(max_connections)
}
fn usage_worker_pool_idle_reserve(summary: &aether_data::DatabasePoolSummary) -> usize {
if summary.max_connections <= 1 {
return 0;
}
1
}
pub(crate) fn should_defer_maintenance_for_database_pool_pressure(
&self,
deferred_since: &mut Option<Instant>,
+33
View File
@@ -97,6 +97,39 @@ fn maintenance_pool_pressure_keeps_idle_reserve_for_foreground_work() {
assert!(!GatewayDataState::database_pool_summary_under_maintenance_pressure(&idle));
}
#[test]
fn usage_worker_pool_pressure_only_defers_near_pool_exhaustion() {
let comfortable = aether_data::DatabasePoolSummary {
driver: DatabaseDriver::Postgres,
checked_out: 56,
pool_size: 64,
idle: 8,
max_connections: 64,
usage_rate: 87.5,
};
assert!(!GatewayDataState::database_pool_summary_under_usage_worker_pressure(&comfortable));
let last_idle_left = aether_data::DatabasePoolSummary {
driver: DatabaseDriver::Postgres,
checked_out: 63,
pool_size: 64,
idle: 1,
max_connections: 64,
usage_rate: 98.4375,
};
assert!(GatewayDataState::database_pool_summary_under_usage_worker_pressure(&last_idle_left));
let exhausted = aether_data::DatabasePoolSummary {
driver: DatabaseDriver::Postgres,
checked_out: 64,
pool_size: 64,
idle: 0,
max_connections: 64,
usage_rate: 100.0,
};
assert!(GatewayDataState::database_pool_summary_under_usage_worker_pressure(&exhausted));
}
#[test]
fn maintenance_pool_pressure_deferral_has_timeout() {
let mut deferred_since = None;
@@ -1420,9 +1420,10 @@ impl DirectPassthroughFinalizerCore {
}
if !self.pending_recorded {
self.pending_recorded = true;
let usage_data = self.state.data.as_ref().clone();
self.state
.usage_runtime
.record_pending_direct(self.state.data.as_ref(), self.lifecycle_seed.clone())
.record_pending_direct(&usage_data, self.lifecycle_seed.clone())
.await;
}
self.stream_started_recorded = true;
@@ -2016,9 +2017,10 @@ async fn record_stream_pending_lifecycle(
stage_trace: &mut RequestStageTrace,
) {
let usage_pending_started_at = Instant::now();
let usage_data = state.data.as_ref().clone();
state
.usage_runtime
.record_pending_direct(state.data.as_ref(), lifecycle_seed.clone())
.record_pending_direct(&usage_data, lifecycle_seed.clone())
.await;
observe_gateway_stage_trace_ms(
stage_trace,
@@ -5250,10 +5252,11 @@ async fn execute_stream_from_frame_stream(
telemetry.as_ref(),
&mut usage_stream_telemetry,
) {
let usage_data = state_for_report.data.as_ref().clone();
state_for_report
.usage_runtime
.record_stream_started_direct(
state_for_report.data.as_ref(),
&usage_data,
&lifecycle_seed_for_report,
status_code,
usage_stream_telemetry.as_ref(),
@@ -5458,10 +5461,11 @@ async fn execute_stream_from_frame_stream(
);
if should_refresh_stream_usage {
if usage_frame_telemetry.ttfb_ms.is_some() {
let usage_data = state_for_report.data.as_ref().clone();
state_for_report
.usage_runtime
.record_stream_started_direct(
state_for_report.data.as_ref(),
&usage_data,
&lifecycle_seed_for_report,
status_code,
Some(&usage_frame_telemetry),
@@ -1538,9 +1538,10 @@ async fn execute_execution_runtime_sync_impl(
.unwrap_or_else(|| "-".to_string());
let candidate_started_unix_secs = current_request_candidate_unix_ms();
let lifecycle_seed = build_lifecycle_usage_seed(&plan, report_context.as_ref());
let usage_data = state.data.as_ref().clone();
state
.usage_runtime
.record_pending_direct(state.data.as_ref(), lifecycle_seed)
.record_pending_direct(&usage_data, lifecycle_seed)
.await;
record_local_request_candidate_status(
state,
+3
View File
@@ -27,6 +27,7 @@
mod admin_api;
mod ai_serving;
mod allocator_metrics;
mod api;
mod async_task;
mod audit;
@@ -58,6 +59,7 @@ mod model_fetch;
mod oauth;
mod orchestration;
mod privacy;
mod process_metrics;
mod provider_key_auth;
mod provider_pool_demand;
pub(crate) use aether_provider_transport as provider_transport;
@@ -76,6 +78,7 @@ mod system_features;
mod task_runtime;
#[cfg(feature = "testkit")]
pub mod testkit;
mod tokio_metrics;
mod tunnel;
mod upstream_admission;
mod usage;
+144 -22
View File
@@ -250,6 +250,8 @@ const AUTO_USAGE_QUEUE_WORKERS_MIN: usize = 2;
const AUTO_USAGE_QUEUE_WORKERS_REQUESTS_PER_WORKER: usize = 128;
const AUTO_USAGE_QUEUE_WORKERS_DB_SHARE_ALL: usize = 4;
const AUTO_USAGE_QUEUE_WORKERS_DB_SHARE_BACKGROUND: usize = 2;
const AUTO_USAGE_WORKER_RECORD_DB_SHARE_ALL: usize = 8;
const AUTO_USAGE_WORKER_RECORD_DB_SHARE_BACKGROUND: usize = 4;
const MAX_USAGE_QUEUE_WORKERS: usize = 64;
const DEFAULT_GATEWAY_LISTEN_BACKLOG: i32 = 65_535;
const MIN_GATEWAY_LISTEN_BACKLOG: i32 = 128;
@@ -325,6 +327,29 @@ fn usage_queue_worker_database_cap(
.clamp(1, MAX_USAGE_QUEUE_WORKERS)
}
fn usage_worker_record_concurrency_database_cap(
node_role: NodeRoleArg,
database: Option<&SqlDatabaseConfig>,
) -> Option<usize> {
let database = database?;
if database.driver == DatabaseDriver::Sqlite {
return Some(1);
}
let divisor = if matches!(node_role, NodeRoleArg::Background) {
AUTO_USAGE_WORKER_RECORD_DB_SHARE_BACKGROUND
} else {
AUTO_USAGE_WORKER_RECORD_DB_SHARE_ALL
};
let max_connections = database.pool.max_connections.max(1) as usize;
Some(
max_connections
.checked_div(divisor.max(1))
.unwrap_or(1)
.clamp(1, MAX_USAGE_QUEUE_WORKERS),
)
}
fn automatic_usage_queue_workers_for_parallelism(
parallelism: usize,
node_role: NodeRoleArg,
@@ -631,10 +656,19 @@ struct GatewayUsageArgs {
#[arg(
long,
env = "AETHER_GATEWAY_USAGE_QUEUE_WORKER_MAX_COUNT",
value_name = "COUNT"
value_name = "COUNT",
default_value = "32"
)]
queue_worker_max_count: Option<usize>,
#[arg(
long,
env = "AETHER_GATEWAY_USAGE_WORKER_RECORD_CONCURRENCY_LIMIT",
value_name = "COUNT",
default_value = "32"
)]
worker_record_concurrency_limit: Option<usize>,
#[arg(
long,
env = "AETHER_GATEWAY_USAGE_QUEUE_WORKER_SCALE_INTERVAL_MS",
@@ -680,7 +714,7 @@ struct GatewayUsageArgs {
#[arg(
long,
env = "AETHER_GATEWAY_USAGE_QUEUE_BATCH_SIZE",
default_value_t = 500
default_value_t = 128
)]
queue_batch_size: usize,
@@ -694,14 +728,14 @@ struct GatewayUsageArgs {
#[arg(
long,
env = "AETHER_GATEWAY_USAGE_QUEUE_RECLAIM_IDLE_MS",
default_value_t = 30_000
default_value_t = 60_000
)]
queue_reclaim_idle_ms: u64,
#[arg(
long,
env = "AETHER_GATEWAY_USAGE_QUEUE_RECLAIM_COUNT",
default_value_t = 500
default_value_t = 128
)]
queue_reclaim_count: usize,
@@ -715,21 +749,28 @@ struct GatewayUsageArgs {
#[arg(
long,
env = "AETHER_GATEWAY_USAGE_TERMINAL_ENQUEUE_MAX_IN_FLIGHT",
default_value_t = 256
default_value_t = 1_024
)]
terminal_enqueue_max_in_flight: u64,
#[arg(
long,
env = "AETHER_GATEWAY_USAGE_LIFECYCLE_ENQUEUE_MAX_IN_FLIGHT",
default_value_t = 128
default_value_t = 512
)]
lifecycle_enqueue_max_in_flight: u64,
#[arg(
long,
env = "AETHER_GATEWAY_USAGE_LIFECYCLE_ENQUEUE_DELAY_MS",
default_value_t = 1_000
)]
lifecycle_enqueue_delay_ms: u64,
#[arg(
long,
env = "AETHER_GATEWAY_USAGE_RETRY_DEFERRED_LIFECYCLE_EVENTS",
default_value_t = false
default_value_t = true
)]
retry_deferred_lifecycle_events: bool,
@@ -743,7 +784,7 @@ struct GatewayUsageArgs {
#[arg(
long,
env = "AETHER_GATEWAY_USAGE_ENQUEUE_RETRY_WORKERS",
default_value_t = 4
default_value_t = 8
)]
enqueue_retry_workers: usize,
@@ -791,10 +832,12 @@ impl GatewayUsageArgs {
worker_count: usize,
) -> usize {
if !self.queue_worker_autoscale_enabled {
return worker_count.max(1).min(MAX_USAGE_QUEUE_WORKERS);
return worker_count.clamp(1, MAX_USAGE_QUEUE_WORKERS);
}
self.queue_worker_max_count
.unwrap_or_else(|| usage_queue_worker_database_cap(node_role, database))
.max(1)
.min(usage_queue_worker_database_cap(node_role, database))
.clamp(worker_count.max(1), MAX_USAGE_QUEUE_WORKERS)
}
@@ -813,7 +856,39 @@ impl GatewayUsageArgs {
Some(worker_max_count.clamp(1, MAX_USAGE_QUEUE_WORKERS))
}
fn to_config(&self, worker_count: usize, worker_max_count: usize) -> UsageRuntimeConfig {
fn effective_worker_record_concurrency_limit(
&self,
node_role: NodeRoleArg,
database: Option<&SqlDatabaseConfig>,
) -> Option<usize> {
if let Some(limit) = self.worker_record_concurrency_limit {
if limit == 0 {
return None;
}
return Some(
limit
.min(MAX_USAGE_QUEUE_WORKERS)
.min(
usage_worker_record_concurrency_database_cap(node_role, database)
.unwrap_or(MAX_USAGE_QUEUE_WORKERS),
)
.max(1),
);
}
if !node_role.spawns_background_tasks()
|| (!self.queue_terminal_events && !self.queue_lifecycle_events)
{
return None;
}
usage_worker_record_concurrency_database_cap(node_role, database)
}
fn to_config(
&self,
worker_count: usize,
worker_max_count: usize,
worker_record_concurrency_limit: Option<usize>,
) -> UsageRuntimeConfig {
UsageRuntimeConfig {
enabled: true,
queue_terminal_events: self.queue_terminal_events,
@@ -821,6 +896,7 @@ impl GatewayUsageArgs {
worker_count: worker_count.clamp(1, MAX_USAGE_QUEUE_WORKERS),
worker_autoscale_enabled: self.queue_worker_autoscale_enabled,
worker_max_count: worker_max_count.clamp(worker_count.max(1), MAX_USAGE_QUEUE_WORKERS),
worker_record_concurrency_limit,
worker_scale_interval_ms: self.queue_worker_scale_interval_ms.max(1),
worker_idle_scale_down_ticks: self.queue_worker_idle_scale_down_ticks.max(1),
stream_key: self.queue_stream_key.trim().to_string(),
@@ -834,6 +910,7 @@ impl GatewayUsageArgs {
reclaim_interval_ms: self.queue_reclaim_interval_ms.max(1),
terminal_enqueue_max_in_flight: self.terminal_enqueue_max_in_flight.max(1),
lifecycle_enqueue_max_in_flight: self.lifecycle_enqueue_max_in_flight.max(1),
lifecycle_enqueue_delay_ms: self.lifecycle_enqueue_delay_ms,
retry_deferred_lifecycle_events: self.retry_deferred_lifecycle_events,
enqueue_retry_buffer_capacity: self.enqueue_retry_buffer_capacity.max(1),
enqueue_retry_workers: self.enqueue_retry_workers.clamp(1, 64),
@@ -1593,9 +1670,14 @@ async fn run() -> Result<(), Box<dyn std::error::Error>> {
sql_database_config.as_ref(),
usage_queue_workers,
);
let usage_config = args
let usage_worker_record_concurrency_limit = args
.usage
.to_config(usage_queue_workers, usage_queue_worker_max_count);
.effective_worker_record_concurrency_limit(args.node_role, sql_database_config.as_ref());
let usage_config = args.usage.to_config(
usage_queue_workers,
usage_queue_worker_max_count,
usage_worker_record_concurrency_limit,
);
let usage_blocking_stream_lanes = args.usage.runtime_state_blocking_stream_lanes(
args.node_role,
sql_database_config.as_ref(),
@@ -1634,6 +1716,9 @@ async fn run() -> Result<(), Box<dyn std::error::Error>> {
usage_queue_workers = usage_config.worker_count,
usage_queue_worker_autoscale_enabled = usage_config.worker_autoscale_enabled,
usage_queue_worker_max_count = usage_config.worker_max_count,
usage_worker_record_concurrency_limit = usage_config
.worker_record_concurrency_limit
.unwrap_or_default(),
usage_queue_request_concurrency_hint =
usage_queue_request_concurrency_hint.unwrap_or_default(),
usage_queue_request_concurrency_hint_source = if usage_queue_request_concurrency_hint.is_some() {
@@ -1672,6 +1757,9 @@ async fn run() -> Result<(), Box<dyn std::error::Error>> {
},
usage_queue_worker_autoscale_enabled = usage_config.worker_autoscale_enabled,
usage_queue_worker_max_count = usage_config.worker_max_count,
usage_worker_record_concurrency_limit = usage_config
.worker_record_concurrency_limit
.unwrap_or_default(),
usage_queue_request_concurrency_hint =
usage_queue_request_concurrency_hint.unwrap_or_default(),
usage_queue_request_concurrency_hint_source =
@@ -2292,23 +2380,25 @@ mod tests {
queue_lifecycle_events: true,
queue_workers: Some(4),
queue_worker_autoscale_enabled: true,
queue_worker_max_count: None,
queue_worker_max_count: Some(32),
worker_record_concurrency_limit: Some(32),
queue_worker_scale_interval_ms: 1_000,
queue_worker_idle_scale_down_ticks: 30,
queue_stream_key: "usage:events".to_string(),
queue_group: "usage_consumers".to_string(),
queue_dlq_stream_key: "usage:events:dlq".to_string(),
queue_stream_maxlen: 200_000,
queue_batch_size: 500,
queue_batch_size: 128,
queue_block_ms: 500,
queue_reclaim_idle_ms: 30_000,
queue_reclaim_count: 500,
queue_reclaim_idle_ms: 60_000,
queue_reclaim_count: 128,
queue_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,
},
@@ -2526,7 +2616,7 @@ mod tests {
);
assert_eq!(workers, 64);
assert_eq!(args.usage.to_config(workers, 64).worker_count, 64);
assert_eq!(args.usage.to_config(workers, 64, Some(8)).worker_count, 64);
}
#[test]
@@ -2566,7 +2656,7 @@ mod tests {
let mut args = test_args();
args.usage.queue_workers = None;
args.usage.queue_worker_max_count = Some(32);
let database = test_database(DatabaseDriver::Postgres, 100);
let database = test_database(DatabaseDriver::Postgres, 200);
let workers =
args.usage
@@ -2579,6 +2669,38 @@ mod tests {
assert_eq!(max_workers, 32);
}
#[test]
fn gateway_usage_worker_record_concurrency_defaults_to_pool_reserve_share() {
let args = test_args();
let database = test_database(DatabaseDriver::Postgres, 64);
assert_eq!(
args.usage
.effective_worker_record_concurrency_limit(NodeRoleArg::All, Some(&database)),
Some(8)
);
assert_eq!(
args.usage.effective_worker_record_concurrency_limit(
NodeRoleArg::Background,
Some(&database)
),
Some(16)
);
}
#[test]
fn gateway_usage_worker_record_concurrency_can_be_explicitly_disabled() {
let mut args = test_args();
args.usage.worker_record_concurrency_limit = Some(0);
let database = test_database(DatabaseDriver::Postgres, 64);
assert_eq!(
args.usage
.effective_worker_record_concurrency_limit(NodeRoleArg::All, Some(&database)),
None
);
}
#[test]
fn gateway_usage_queue_blocking_stream_lanes_only_expand_when_worker_can_spawn() {
let database = test_database(DatabaseDriver::Postgres, 100);
+2 -1
View File
@@ -30,5 +30,6 @@ pub(crate) use runtime::{
ProxyUpgradeRolloutCancelSummary, ProxyUpgradeRolloutConflictClearSummary,
ProxyUpgradeRolloutNodeActionSummary, ProxyUpgradeRolloutProbeConfig,
ProxyUpgradeRolloutSkippedRestoreSummary, ProxyUpgradeRolloutStatus,
ProxyUpgradeRolloutTrackedNodeState,
ProxyUpgradeRolloutTrackedNodeState, UsageCounterFlushRuntimeMetrics,
UsageCounterFlushWorkerConfig,
};
@@ -116,6 +116,9 @@ pub(crate) use usage_cleanup::{
preview_manual_usage_cleanup, ManualUsageCleanupMode, ManualUsageCleanupOptions,
};
use usage_counter_flush::*;
pub(crate) use usage_counter_flush::{
UsageCounterFlushRuntimeMetrics, UsageCounterFlushWorkerConfig,
};
use wallet_daily_usage::*;
pub(crate) use workers::*;
@@ -136,12 +139,6 @@ const PROXY_NODE_STALE_MIN_GRACE_SECS: u64 = 15;
const PROXY_NODE_STALE_MISSED_HEARTBEATS: u64 = 3;
const POOL_MONITOR_INTERVAL: Duration = Duration::from_secs(5 * 60);
const OAUTH_TOKEN_REFRESH_INTERVAL: Duration = Duration::from_secs(60);
const USAGE_COUNTER_FLUSH_INTERVAL: Duration = Duration::from_secs(1);
const USAGE_COUNTER_FLUSH_BATCH_SIZE: usize = 1_000;
const USAGE_COUNTER_FLUSH_CATCH_UP_BURST_LIMIT: usize = 20;
const USAGE_COUNTER_DELTA_CLEANUP_INTERVAL: Duration = Duration::from_secs(60);
const USAGE_COUNTER_DELTA_CLEANUP_BATCH_SIZE: usize = 5_000;
const USAGE_COUNTER_DELTA_RETENTION_SECS: u64 = 7 * 24 * 60 * 60;
const PROVIDER_CHECKIN_CONCURRENCY: usize = 3;
const PROVIDER_QUOTA_ALERT_CONCURRENCY: usize = 3;
const PROVIDER_QUOTA_ALERT_INTERVAL: Duration = Duration::from_secs(5);
@@ -1,6 +1,300 @@
use std::sync::atomic::{AtomicU64, Ordering};
use std::time::Duration;
use crate::data::GatewayDataState;
use aether_data::DataLayerError;
use aether_data_contracts::repository::usage::UsageCounterFlushSummary;
use aether_runtime::{MetricKind, MetricLabel, MetricSample};
const USAGE_COUNTER_FLUSH_INTERVAL_MS_ENV: &str = "AETHER_GATEWAY_USAGE_COUNTER_FLUSH_INTERVAL_MS";
const USAGE_COUNTER_FLUSH_BATCH_SIZE_ENV: &str = "AETHER_GATEWAY_USAGE_COUNTER_FLUSH_BATCH_SIZE";
const USAGE_COUNTER_FLUSH_CATCH_UP_BURST_LIMIT_ENV: &str =
"AETHER_GATEWAY_USAGE_COUNTER_FLUSH_CATCH_UP_BURST_LIMIT";
const USAGE_COUNTER_DELTA_CLEANUP_INTERVAL_MS_ENV: &str =
"AETHER_GATEWAY_USAGE_COUNTER_DELTA_CLEANUP_INTERVAL_MS";
const USAGE_COUNTER_DELTA_CLEANUP_BATCH_SIZE_ENV: &str =
"AETHER_GATEWAY_USAGE_COUNTER_DELTA_CLEANUP_BATCH_SIZE";
const USAGE_COUNTER_DELTA_RETENTION_SECS_ENV: &str =
"AETHER_GATEWAY_USAGE_COUNTER_DELTA_RETENTION_SECS";
const DEFAULT_USAGE_COUNTER_FLUSH_INTERVAL_MS: u64 = 1_000;
const DEFAULT_USAGE_COUNTER_FLUSH_BATCH_SIZE: usize = 1_000;
const DEFAULT_USAGE_COUNTER_FLUSH_CATCH_UP_BURST_LIMIT: usize = 20;
const DEFAULT_USAGE_COUNTER_DELTA_CLEANUP_INTERVAL_MS: u64 = 60_000;
const DEFAULT_USAGE_COUNTER_DELTA_CLEANUP_BATCH_SIZE: usize = 5_000;
const DEFAULT_USAGE_COUNTER_DELTA_RETENTION_SECS: u64 = 7 * 24 * 60 * 60;
#[derive(Debug, Clone, PartialEq, Eq)]
pub(crate) struct UsageCounterFlushWorkerConfig {
pub(crate) flush_interval: Duration,
pub(crate) flush_batch_size: usize,
pub(crate) flush_catch_up_burst_limit: usize,
pub(crate) cleanup_interval: Duration,
pub(crate) cleanup_batch_size: usize,
pub(crate) delta_retention_secs: u64,
}
impl Default for UsageCounterFlushWorkerConfig {
fn default() -> Self {
Self {
flush_interval: Duration::from_millis(DEFAULT_USAGE_COUNTER_FLUSH_INTERVAL_MS),
flush_batch_size: DEFAULT_USAGE_COUNTER_FLUSH_BATCH_SIZE,
flush_catch_up_burst_limit: DEFAULT_USAGE_COUNTER_FLUSH_CATCH_UP_BURST_LIMIT,
cleanup_interval: Duration::from_millis(
DEFAULT_USAGE_COUNTER_DELTA_CLEANUP_INTERVAL_MS,
),
cleanup_batch_size: DEFAULT_USAGE_COUNTER_DELTA_CLEANUP_BATCH_SIZE,
delta_retention_secs: DEFAULT_USAGE_COUNTER_DELTA_RETENTION_SECS,
}
}
}
impl UsageCounterFlushWorkerConfig {
pub(crate) fn from_env() -> Self {
let defaults = Self::default();
Self {
flush_interval: Duration::from_millis(env_u64(
USAGE_COUNTER_FLUSH_INTERVAL_MS_ENV,
duration_millis_u64(defaults.flush_interval),
)),
flush_batch_size: env_usize(
USAGE_COUNTER_FLUSH_BATCH_SIZE_ENV,
defaults.flush_batch_size,
),
flush_catch_up_burst_limit: env_usize(
USAGE_COUNTER_FLUSH_CATCH_UP_BURST_LIMIT_ENV,
defaults.flush_catch_up_burst_limit,
),
cleanup_interval: Duration::from_millis(env_u64(
USAGE_COUNTER_DELTA_CLEANUP_INTERVAL_MS_ENV,
duration_millis_u64(defaults.cleanup_interval),
)),
cleanup_batch_size: env_usize(
USAGE_COUNTER_DELTA_CLEANUP_BATCH_SIZE_ENV,
defaults.cleanup_batch_size,
),
delta_retention_secs: env_u64(
USAGE_COUNTER_DELTA_RETENTION_SECS_ENV,
defaults.delta_retention_secs,
),
}
.normalized()
}
fn normalized(mut self) -> Self {
self.flush_interval = self.flush_interval.max(Duration::from_millis(1));
self.flush_batch_size = self.flush_batch_size.max(1);
self.flush_catch_up_burst_limit = self.flush_catch_up_burst_limit.max(1);
self.cleanup_interval = self.cleanup_interval.max(Duration::from_millis(1));
self.cleanup_batch_size = self.cleanup_batch_size.max(1);
self
}
}
#[derive(Debug, Default)]
pub(crate) struct UsageCounterFlushRuntimeMetrics {
flush_batches_total: AtomicU64,
flush_empty_batches_total: AtomicU64,
flush_rows_claimed_total: AtomicU64,
flush_api_key_targets_total: AtomicU64,
flush_provider_api_key_targets_total: AtomicU64,
flush_model_targets_total: AtomicU64,
flush_provider_monthly_targets_total: AtomicU64,
flush_proxy_node_targets_total: AtomicU64,
flush_management_token_targets_total: AtomicU64,
flush_api_key_last_used_targets_total: AtomicU64,
flush_failed_batches_total: AtomicU64,
flush_deferred_total: AtomicU64,
cleanup_batches_total: AtomicU64,
cleanup_rows_total: AtomicU64,
cleanup_failed_batches_total: AtomicU64,
cleanup_deferred_total: AtomicU64,
}
impl UsageCounterFlushRuntimeMetrics {
pub(crate) fn record_flush_success(&self, summary: &UsageCounterFlushSummary) {
if summary.rows_claimed > 0 {
self.flush_batches_total.fetch_add(1, Ordering::AcqRel);
} else {
self.flush_empty_batches_total
.fetch_add(1, Ordering::AcqRel);
}
self.flush_rows_claimed_total
.fetch_add(usize_to_u64(summary.rows_claimed), Ordering::AcqRel);
self.flush_api_key_targets_total
.fetch_add(usize_to_u64(summary.api_key_targets), Ordering::AcqRel);
self.flush_provider_api_key_targets_total.fetch_add(
usize_to_u64(summary.provider_api_key_targets),
Ordering::AcqRel,
);
self.flush_model_targets_total
.fetch_add(usize_to_u64(summary.model_targets), Ordering::AcqRel);
self.flush_provider_monthly_targets_total.fetch_add(
usize_to_u64(summary.provider_monthly_targets),
Ordering::AcqRel,
);
self.flush_proxy_node_targets_total
.fetch_add(usize_to_u64(summary.proxy_node_targets), Ordering::AcqRel);
self.flush_management_token_targets_total.fetch_add(
usize_to_u64(summary.management_token_targets),
Ordering::AcqRel,
);
self.flush_api_key_last_used_targets_total.fetch_add(
usize_to_u64(summary.api_key_last_used_targets),
Ordering::AcqRel,
);
}
pub(crate) fn record_flush_failed(&self) {
self.flush_failed_batches_total
.fetch_add(1, Ordering::AcqRel);
}
pub(crate) fn record_flush_deferred(&self) {
self.flush_deferred_total.fetch_add(1, Ordering::AcqRel);
}
pub(crate) fn record_cleanup_success(&self, rows_deleted: usize) {
self.cleanup_batches_total.fetch_add(1, Ordering::AcqRel);
self.cleanup_rows_total
.fetch_add(usize_to_u64(rows_deleted), Ordering::AcqRel);
}
pub(crate) fn record_cleanup_failed(&self) {
self.cleanup_failed_batches_total
.fetch_add(1, Ordering::AcqRel);
}
pub(crate) fn record_cleanup_deferred(&self) {
self.cleanup_deferred_total.fetch_add(1, Ordering::AcqRel);
}
pub(crate) fn metric_samples(&self) -> Vec<MetricSample> {
let mut samples = vec![
MetricSample::new(
"usage_counter_outbox_flush_batches_total",
"Total non-empty usage counter outbox flush batches completed by this gateway process.",
MetricKind::Counter,
self.flush_batches_total.load(Ordering::Acquire),
),
MetricSample::new(
"usage_counter_outbox_flush_empty_batches_total",
"Total successful usage counter outbox flush checks that found no rows.",
MetricKind::Counter,
self.flush_empty_batches_total.load(Ordering::Acquire),
),
MetricSample::new(
"usage_counter_outbox_flush_rows_claimed_total",
"Total usage counter outbox rows claimed and processed by this gateway process.",
MetricKind::Counter,
self.flush_rows_claimed_total.load(Ordering::Acquire),
),
MetricSample::new(
"usage_counter_outbox_flush_failed_batches_total",
"Total usage counter outbox flush batches that failed in this gateway process.",
MetricKind::Counter,
self.flush_failed_batches_total.load(Ordering::Acquire),
),
MetricSample::new(
"usage_counter_outbox_flush_deferred_total",
"Total usage counter outbox flush ticks deferred because of database pool pressure.",
MetricKind::Counter,
self.flush_deferred_total.load(Ordering::Acquire),
),
MetricSample::new(
"usage_counter_outbox_cleanup_batches_total",
"Total usage counter outbox cleanup batches completed by this gateway process.",
MetricKind::Counter,
self.cleanup_batches_total.load(Ordering::Acquire),
),
MetricSample::new(
"usage_counter_outbox_cleanup_rows_total",
"Total processed usage counter outbox rows deleted by cleanup in this gateway process.",
MetricKind::Counter,
self.cleanup_rows_total.load(Ordering::Acquire),
),
MetricSample::new(
"usage_counter_outbox_cleanup_failed_batches_total",
"Total usage counter outbox cleanup batches that failed in this gateway process.",
MetricKind::Counter,
self.cleanup_failed_batches_total.load(Ordering::Acquire),
),
MetricSample::new(
"usage_counter_outbox_cleanup_deferred_total",
"Total usage counter outbox cleanup ticks deferred because of database pool pressure.",
MetricKind::Counter,
self.cleanup_deferred_total.load(Ordering::Acquire),
),
];
for (kind, value) in [
(
"api_key",
self.flush_api_key_targets_total.load(Ordering::Acquire),
),
(
"provider_api_key",
self.flush_provider_api_key_targets_total
.load(Ordering::Acquire),
),
(
"model",
self.flush_model_targets_total.load(Ordering::Acquire),
),
(
"provider_monthly",
self.flush_provider_monthly_targets_total
.load(Ordering::Acquire),
),
(
"proxy_node",
self.flush_proxy_node_targets_total.load(Ordering::Acquire),
),
(
"management_token",
self.flush_management_token_targets_total
.load(Ordering::Acquire),
),
(
"api_key_last_used",
self.flush_api_key_last_used_targets_total
.load(Ordering::Acquire),
),
] {
samples.push(
MetricSample::new(
"usage_counter_outbox_flush_targets_total",
"Total usage counter outbox aggregate targets updated by target kind.",
MetricKind::Counter,
value,
)
.with_labels(vec![MetricLabel::new("kind", kind)]),
);
}
samples
}
}
fn usize_to_u64(value: usize) -> u64 {
u64::try_from(value).unwrap_or(u64::MAX)
}
fn duration_millis_u64(duration: Duration) -> u64 {
u64::try_from(duration.as_millis()).unwrap_or(u64::MAX)
}
fn env_u64(name: &str, default_value: u64) -> u64 {
std::env::var(name)
.ok()
.and_then(|value| value.trim().parse::<u64>().ok())
.unwrap_or(default_value)
}
fn env_usize(name: &str, default_value: usize) -> usize {
std::env::var(name)
.ok()
.and_then(|value| value.trim().parse::<usize>().ok())
.unwrap_or(default_value)
}
pub(crate) async fn run_usage_counter_flush_once(
data: &GatewayDataState,
@@ -19,3 +313,186 @@ pub(crate) async fn cleanup_processed_usage_counter_deltas_once(
data.cleanup_processed_usage_counter_deltas(cutoff, batch_size)
.await
}
#[cfg(test)]
mod tests {
use std::sync::{Mutex, OnceLock};
use std::time::Duration;
use super::{
UsageCounterFlushRuntimeMetrics, UsageCounterFlushSummary, UsageCounterFlushWorkerConfig,
USAGE_COUNTER_DELTA_CLEANUP_BATCH_SIZE_ENV, USAGE_COUNTER_DELTA_CLEANUP_INTERVAL_MS_ENV,
USAGE_COUNTER_DELTA_RETENTION_SECS_ENV, USAGE_COUNTER_FLUSH_BATCH_SIZE_ENV,
USAGE_COUNTER_FLUSH_CATCH_UP_BURST_LIMIT_ENV, USAGE_COUNTER_FLUSH_INTERVAL_MS_ENV,
};
const CONFIG_ENV_KEYS: &[&str] = &[
USAGE_COUNTER_FLUSH_INTERVAL_MS_ENV,
USAGE_COUNTER_FLUSH_BATCH_SIZE_ENV,
USAGE_COUNTER_FLUSH_CATCH_UP_BURST_LIMIT_ENV,
USAGE_COUNTER_DELTA_CLEANUP_INTERVAL_MS_ENV,
USAGE_COUNTER_DELTA_CLEANUP_BATCH_SIZE_ENV,
USAGE_COUNTER_DELTA_RETENTION_SECS_ENV,
];
fn env_lock() -> &'static Mutex<()> {
static LOCK: OnceLock<Mutex<()>> = OnceLock::new();
LOCK.get_or_init(|| Mutex::new(()))
}
struct EnvVarGuard {
values: Vec<(&'static str, Option<String>)>,
}
impl EnvVarGuard {
fn new(keys: &[&'static str]) -> Self {
let values = keys
.iter()
.map(|key| (*key, std::env::var(key).ok()))
.collect();
Self { values }
}
}
impl Drop for EnvVarGuard {
fn drop(&mut self) {
for (key, value) in &self.values {
match value {
Some(value) => std::env::set_var(key, value),
None => std::env::remove_var(key),
}
}
}
}
fn sample_value(samples: &[aether_runtime::MetricSample], name: &str) -> u64 {
samples
.iter()
.find(|sample| sample.name == name)
.map(|sample| sample.value)
.expect("sample should exist")
}
fn target_value(samples: &[aether_runtime::MetricSample], kind: &str) -> u64 {
samples
.iter()
.find(|sample| {
sample.name == "usage_counter_outbox_flush_targets_total"
&& sample
.labels
.iter()
.any(|label| label.key == "kind" && label.value == kind)
})
.map(|sample| sample.value)
.expect("target sample should exist")
}
#[test]
fn usage_counter_flush_worker_config_reads_env() {
let _lock = env_lock().lock().expect("env lock should not be poisoned");
let _guard = EnvVarGuard::new(CONFIG_ENV_KEYS);
std::env::set_var(USAGE_COUNTER_FLUSH_INTERVAL_MS_ENV, "100");
std::env::set_var(USAGE_COUNTER_FLUSH_BATCH_SIZE_ENV, "2000");
std::env::set_var(USAGE_COUNTER_FLUSH_CATCH_UP_BURST_LIMIT_ENV, "50");
std::env::set_var(USAGE_COUNTER_DELTA_CLEANUP_INTERVAL_MS_ENV, "250");
std::env::set_var(USAGE_COUNTER_DELTA_CLEANUP_BATCH_SIZE_ENV, "7000");
std::env::set_var(USAGE_COUNTER_DELTA_RETENTION_SECS_ENV, "0");
let config = UsageCounterFlushWorkerConfig::from_env();
assert_eq!(config.flush_interval, Duration::from_millis(100));
assert_eq!(config.flush_batch_size, 2000);
assert_eq!(config.flush_catch_up_burst_limit, 50);
assert_eq!(config.cleanup_interval, Duration::from_millis(250));
assert_eq!(config.cleanup_batch_size, 7000);
assert_eq!(config.delta_retention_secs, 0);
}
#[test]
fn usage_counter_flush_worker_config_normalizes_zero_values() {
let _lock = env_lock().lock().expect("env lock should not be poisoned");
let _guard = EnvVarGuard::new(CONFIG_ENV_KEYS);
for key in CONFIG_ENV_KEYS {
std::env::set_var(key, "0");
}
let config = UsageCounterFlushWorkerConfig::from_env();
assert_eq!(config.flush_interval, Duration::from_millis(1));
assert_eq!(config.flush_batch_size, 1);
assert_eq!(config.flush_catch_up_burst_limit, 1);
assert_eq!(config.cleanup_interval, Duration::from_millis(1));
assert_eq!(config.cleanup_batch_size, 1);
assert_eq!(config.delta_retention_secs, 0);
}
#[test]
fn usage_counter_flush_runtime_metrics_record_success_and_failures() {
let metrics = UsageCounterFlushRuntimeMetrics::default();
metrics.record_flush_success(&UsageCounterFlushSummary {
rows_claimed: 3,
api_key_targets: 1,
provider_api_key_targets: 2,
model_targets: 3,
provider_monthly_targets: 4,
proxy_node_targets: 5,
management_token_targets: 6,
api_key_last_used_targets: 7,
});
metrics.record_flush_success(&UsageCounterFlushSummary::default());
metrics.record_flush_failed();
metrics.record_flush_deferred();
metrics.record_cleanup_success(11);
metrics.record_cleanup_failed();
metrics.record_cleanup_deferred();
let samples = metrics.metric_samples();
assert_eq!(
sample_value(&samples, "usage_counter_outbox_flush_batches_total"),
1
);
assert_eq!(
sample_value(&samples, "usage_counter_outbox_flush_empty_batches_total"),
1
);
assert_eq!(
sample_value(&samples, "usage_counter_outbox_flush_rows_claimed_total"),
3
);
assert_eq!(
sample_value(&samples, "usage_counter_outbox_flush_failed_batches_total"),
1
);
assert_eq!(
sample_value(&samples, "usage_counter_outbox_flush_deferred_total"),
1
);
assert_eq!(
sample_value(&samples, "usage_counter_outbox_cleanup_batches_total"),
1
);
assert_eq!(
sample_value(&samples, "usage_counter_outbox_cleanup_rows_total"),
11
);
assert_eq!(
sample_value(
&samples,
"usage_counter_outbox_cleanup_failed_batches_total"
),
1
);
assert_eq!(
sample_value(&samples, "usage_counter_outbox_cleanup_deferred_total"),
1
);
assert_eq!(target_value(&samples, "api_key"), 1);
assert_eq!(target_value(&samples, "provider_api_key"), 2);
assert_eq!(target_value(&samples, "model"), 3);
assert_eq!(target_value(&samples, "provider_monthly"), 4);
assert_eq!(target_value(&samples, "proxy_node"), 5);
assert_eq!(target_value(&samples, "management_token"), 6);
assert_eq!(target_value(&samples, "api_key_last_used"), 7);
}
}
@@ -23,11 +23,9 @@ use super::{
PROXY_NODE_METRICS_CLEANUP_HOUR, PROXY_NODE_METRICS_CLEANUP_MINUTE,
PROXY_NODE_STALE_SWEEP_INTERVAL, PROXY_UPGRADE_ROLLOUT_INTERVAL,
REQUEST_CANDIDATE_CLEANUP_INTERVAL, USAGE_CLEANUP_HOUR, USAGE_CLEANUP_MINUTE,
USAGE_COUNTER_DELTA_CLEANUP_BATCH_SIZE, USAGE_COUNTER_DELTA_CLEANUP_INTERVAL,
USAGE_COUNTER_DELTA_RETENTION_SECS, USAGE_COUNTER_FLUSH_BATCH_SIZE,
USAGE_COUNTER_FLUSH_CATCH_UP_BURST_LIMIT, USAGE_COUNTER_FLUSH_INTERVAL,
WALLET_DAILY_USAGE_AGGREGATION_HOUR, WALLET_DAILY_USAGE_AGGREGATION_MINUTE,
};
use super::{UsageCounterFlushRuntimeMetrics, UsageCounterFlushWorkerConfig};
const STATS_DAILY_CATCH_UP_BURST_LIMIT: usize = 14;
const STATS_HOURLY_CATCH_UP_BURST_LIMIT: usize = 72;
@@ -249,13 +247,26 @@ pub(crate) fn spawn_usage_cleanup_worker(
pub(crate) fn spawn_usage_counter_flush_worker(
data: Arc<GatewayDataState>,
metrics: Arc<UsageCounterFlushRuntimeMetrics>,
) -> Option<tokio::task::JoinHandle<()>> {
spawn_usage_counter_flush_worker_with_config(
data,
metrics,
UsageCounterFlushWorkerConfig::from_env(),
)
}
pub(crate) fn spawn_usage_counter_flush_worker_with_config(
data: Arc<GatewayDataState>,
metrics: Arc<UsageCounterFlushRuntimeMetrics>,
config: UsageCounterFlushWorkerConfig,
) -> Option<tokio::task::JoinHandle<()>> {
if !data.has_usage_counter_flush_backend() {
return None;
}
Some(tokio::spawn(async move {
let mut interval = tokio::time::interval(USAGE_COUNTER_FLUSH_INTERVAL);
let mut interval = tokio::time::interval(config.flush_interval);
interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
interval.tick().await;
let mut last_delta_cleanup = tokio::time::Instant::now();
@@ -268,47 +279,66 @@ pub(crate) fn spawn_usage_counter_flush_worker(
"usage_counter_flush",
&mut usage_counter_flush_deferred_since,
) {
metrics.record_flush_deferred();
interval.tick().await;
continue;
}
let mut batches = 0_usize;
while batches < USAGE_COUNTER_FLUSH_CATCH_UP_BURST_LIMIT {
match run_usage_counter_flush_once(&data, USAGE_COUNTER_FLUSH_BATCH_SIZE).await {
Ok(summary) if summary.rows_claimed > 0 => batches += 1,
Ok(_) => break,
while batches < config.flush_catch_up_burst_limit {
match run_usage_counter_flush_once(&data, config.flush_batch_size).await {
Ok(summary) if summary.rows_claimed > 0 => {
metrics.record_flush_success(&summary);
batches += 1;
}
Ok(summary) => {
metrics.record_flush_success(&summary);
break;
}
Err(err) => {
metrics.record_flush_failed();
log_maintenance_worker_failure("usage_counter_flush", "tick", &err);
break;
}
}
}
if batches >= USAGE_COUNTER_FLUSH_CATCH_UP_BURST_LIMIT {
if batches >= config.flush_catch_up_burst_limit {
tokio::task::yield_now().await;
continue;
}
if last_delta_cleanup.elapsed() >= USAGE_COUNTER_DELTA_CLEANUP_INTERVAL {
if last_delta_cleanup.elapsed() >= config.cleanup_interval {
if should_defer_for_database_pressure(
&data,
"usage_counter_delta_cleanup",
&mut usage_counter_delta_cleanup_deferred_since,
) {
metrics.record_cleanup_deferred();
debug!(
event_name = "maintenance_worker_deferred",
log_type = "ops",
worker = "usage_counter_delta_cleanup",
"gateway maintenance worker deferred cleanup under database pressure"
);
} else if let Err(err) = cleanup_processed_usage_counter_deltas_once(
&data,
USAGE_COUNTER_DELTA_RETENTION_SECS,
USAGE_COUNTER_DELTA_CLEANUP_BATCH_SIZE,
)
.await
{
log_maintenance_worker_failure("usage_counter_delta_cleanup", "tick", &err);
} else {
match cleanup_processed_usage_counter_deltas_once(
&data,
config.delta_retention_secs,
config.cleanup_batch_size,
)
.await
{
Ok(rows_deleted) => metrics.record_cleanup_success(rows_deleted),
Err(err) => {
metrics.record_cleanup_failed();
log_maintenance_worker_failure(
"usage_counter_delta_cleanup",
"tick",
&err,
);
}
}
}
last_delta_cleanup = tokio::time::Instant::now();
}
+787
View File
@@ -0,0 +1,787 @@
use std::collections::HashSet;
use std::sync::Mutex;
use std::time::{SystemTime, UNIX_EPOCH};
use aether_runtime::{MetricKind, MetricSample};
use sysinfo::{get_current_pid, Networks, Pid, ProcessesToUpdate, System};
#[derive(Debug, Clone, Copy, Default)]
pub(crate) struct GatewayProcessResourceSnapshot {
pub(crate) sampled_at_unix_secs: u64,
pub(crate) system_cpu_usage_basis_points: u64,
pub(crate) process_cpu_usage_basis_points: u64,
pub(crate) memory_total_bytes: u64,
pub(crate) memory_used_bytes: u64,
pub(crate) memory_available_bytes: u64,
pub(crate) memory_used_basis_points: u64,
pub(crate) process_memory_bytes: u64,
pub(crate) process_virtual_memory_bytes: u64,
pub(crate) process_memory_basis_points: u64,
pub(crate) process_uptime_secs: Option<u64>,
pub(crate) process_threads: u64,
pub(crate) fd_open_count: u64,
pub(crate) fd_limit: u64,
pub(crate) fd_usage_basis_points: u64,
pub(crate) network_observability_available: u64,
pub(crate) network_interface_count: u64,
pub(crate) network_received_bytes_total: u64,
pub(crate) network_transmitted_bytes_total: u64,
pub(crate) network_received_packets_total: u64,
pub(crate) network_transmitted_packets_total: u64,
pub(crate) network_receive_errors_total: u64,
pub(crate) network_transmit_errors_total: u64,
pub(crate) network_receive_dropped_total: u64,
pub(crate) network_transmit_dropped_total: u64,
pub(crate) process_socket_fds: u64,
pub(crate) tcp_state_observability_available: u64,
pub(crate) host_tcp_connections: u64,
pub(crate) host_tcp_established_connections: u64,
pub(crate) host_tcp_listen_connections: u64,
pub(crate) host_tcp_time_wait_connections: u64,
pub(crate) host_tcp_syn_sent_connections: u64,
pub(crate) host_tcp_syn_recv_connections: u64,
pub(crate) host_tcp_close_wait_connections: u64,
pub(crate) process_tcp_connections: u64,
pub(crate) process_tcp_established_connections: u64,
pub(crate) process_tcp_listen_connections: u64,
pub(crate) process_tcp_time_wait_connections: u64,
pub(crate) process_tcp_syn_sent_connections: u64,
pub(crate) process_tcp_syn_recv_connections: u64,
pub(crate) process_tcp_close_wait_connections: u64,
}
impl GatewayProcessResourceSnapshot {
pub(crate) fn to_metric_samples(self) -> Vec<MetricSample> {
let mut samples = vec![
MetricSample::new(
"gateway_process_sampled_at_unix_secs",
"Unix timestamp of the current gateway process resource sample.",
MetricKind::Gauge,
self.sampled_at_unix_secs,
),
MetricSample::new(
"gateway_system_cpu_usage_basis_points",
"Host CPU usage in basis points of percent, where 10000 means 100 percent.",
MetricKind::Gauge,
self.system_cpu_usage_basis_points,
),
MetricSample::new(
"gateway_process_cpu_usage_basis_points",
"Gateway process CPU usage in basis points of percent, where 10000 means 100 percent.",
MetricKind::Gauge,
self.process_cpu_usage_basis_points,
),
MetricSample::new(
"gateway_system_memory_total_bytes",
"Total host memory visible to the gateway process.",
MetricKind::Gauge,
self.memory_total_bytes,
),
MetricSample::new(
"gateway_system_memory_used_bytes",
"Used host memory visible to the gateway process.",
MetricKind::Gauge,
self.memory_used_bytes,
),
MetricSample::new(
"gateway_system_memory_available_bytes",
"Available host memory visible to the gateway process.",
MetricKind::Gauge,
self.memory_available_bytes,
),
MetricSample::new(
"gateway_system_memory_usage_basis_points",
"Host memory usage in basis points, where 10000 means 100 percent.",
MetricKind::Gauge,
self.memory_used_basis_points,
),
MetricSample::new(
"gateway_process_memory_bytes",
"Gateway process resident memory bytes.",
MetricKind::Gauge,
self.process_memory_bytes,
),
MetricSample::new(
"gateway_process_virtual_memory_bytes",
"Gateway process virtual memory bytes.",
MetricKind::Gauge,
self.process_virtual_memory_bytes,
),
MetricSample::new(
"gateway_process_memory_basis_points",
"Gateway process resident memory as basis points of host memory.",
MetricKind::Gauge,
self.process_memory_basis_points,
),
MetricSample::new(
"gateway_process_threads",
"Current number of threads owned by the gateway process where available.",
MetricKind::Gauge,
self.process_threads,
),
MetricSample::new(
"gateway_process_open_fds",
"Current number of file descriptors opened by the gateway process.",
MetricKind::Gauge,
self.fd_open_count,
),
MetricSample::new(
"gateway_process_fd_limit",
"Current soft file descriptor limit for the gateway process.",
MetricKind::Gauge,
self.fd_limit,
),
MetricSample::new(
"gateway_process_fd_usage_basis_points",
"Gateway process file descriptor usage in basis points, where 10000 means 100 percent.",
MetricKind::Gauge,
self.fd_usage_basis_points,
),
MetricSample::new(
"gateway_network_observability_available",
"Whether host network interface counters are available.",
MetricKind::Gauge,
self.network_observability_available,
),
MetricSample::new(
"gateway_network_interfaces",
"Number of host network interfaces visible to the gateway process.",
MetricKind::Gauge,
self.network_interface_count,
),
MetricSample::new(
"gateway_network_received_bytes_total",
"Total host network bytes received across visible interfaces.",
MetricKind::Counter,
self.network_received_bytes_total,
),
MetricSample::new(
"gateway_network_transmitted_bytes_total",
"Total host network bytes transmitted across visible interfaces.",
MetricKind::Counter,
self.network_transmitted_bytes_total,
),
MetricSample::new(
"gateway_network_received_packets_total",
"Total host network packets received across visible interfaces.",
MetricKind::Counter,
self.network_received_packets_total,
),
MetricSample::new(
"gateway_network_transmitted_packets_total",
"Total host network packets transmitted across visible interfaces.",
MetricKind::Counter,
self.network_transmitted_packets_total,
),
MetricSample::new(
"gateway_network_receive_errors_total",
"Total host network receive errors across visible interfaces.",
MetricKind::Counter,
self.network_receive_errors_total,
),
MetricSample::new(
"gateway_network_transmit_errors_total",
"Total host network transmit errors across visible interfaces.",
MetricKind::Counter,
self.network_transmit_errors_total,
),
MetricSample::new(
"gateway_network_receive_dropped_total",
"Total host network receive drops across visible interfaces where available.",
MetricKind::Counter,
self.network_receive_dropped_total,
),
MetricSample::new(
"gateway_network_transmit_dropped_total",
"Total host network transmit drops across visible interfaces where available.",
MetricKind::Counter,
self.network_transmit_dropped_total,
),
MetricSample::new(
"gateway_process_socket_fds",
"Current number of socket file descriptors opened by the gateway process.",
MetricKind::Gauge,
self.process_socket_fds,
),
MetricSample::new(
"gateway_tcp_state_observability_available",
"Whether Linux TCP state counters are available from procfs.",
MetricKind::Gauge,
self.tcp_state_observability_available,
),
MetricSample::new(
"gateway_host_tcp_connections",
"Current host TCP connections visible in procfs.",
MetricKind::Gauge,
self.host_tcp_connections,
),
MetricSample::new(
"gateway_host_tcp_established_connections",
"Current host TCP connections in ESTABLISHED state visible in procfs.",
MetricKind::Gauge,
self.host_tcp_established_connections,
),
MetricSample::new(
"gateway_host_tcp_listen_connections",
"Current host TCP sockets in LISTEN state visible in procfs.",
MetricKind::Gauge,
self.host_tcp_listen_connections,
),
MetricSample::new(
"gateway_host_tcp_time_wait_connections",
"Current host TCP connections in TIME_WAIT state visible in procfs.",
MetricKind::Gauge,
self.host_tcp_time_wait_connections,
),
MetricSample::new(
"gateway_host_tcp_syn_sent_connections",
"Current host TCP connections in SYN_SENT state visible in procfs.",
MetricKind::Gauge,
self.host_tcp_syn_sent_connections,
),
MetricSample::new(
"gateway_host_tcp_syn_recv_connections",
"Current host TCP connections in SYN_RECV state visible in procfs.",
MetricKind::Gauge,
self.host_tcp_syn_recv_connections,
),
MetricSample::new(
"gateway_host_tcp_close_wait_connections",
"Current host TCP connections in CLOSE_WAIT state visible in procfs.",
MetricKind::Gauge,
self.host_tcp_close_wait_connections,
),
MetricSample::new(
"gateway_process_tcp_connections",
"Current gateway process TCP connections visible in procfs.",
MetricKind::Gauge,
self.process_tcp_connections,
),
MetricSample::new(
"gateway_process_tcp_established_connections",
"Current gateway process TCP connections in ESTABLISHED state visible in procfs.",
MetricKind::Gauge,
self.process_tcp_established_connections,
),
MetricSample::new(
"gateway_process_tcp_listen_connections",
"Current gateway process TCP sockets in LISTEN state visible in procfs.",
MetricKind::Gauge,
self.process_tcp_listen_connections,
),
MetricSample::new(
"gateway_process_tcp_time_wait_connections",
"Current gateway process TCP connections in TIME_WAIT state visible in procfs.",
MetricKind::Gauge,
self.process_tcp_time_wait_connections,
),
MetricSample::new(
"gateway_process_tcp_syn_sent_connections",
"Current gateway process TCP connections in SYN_SENT state visible in procfs.",
MetricKind::Gauge,
self.process_tcp_syn_sent_connections,
),
MetricSample::new(
"gateway_process_tcp_syn_recv_connections",
"Current gateway process TCP connections in SYN_RECV state visible in procfs.",
MetricKind::Gauge,
self.process_tcp_syn_recv_connections,
),
MetricSample::new(
"gateway_process_tcp_close_wait_connections",
"Current gateway process TCP connections in CLOSE_WAIT state visible in procfs.",
MetricKind::Gauge,
self.process_tcp_close_wait_connections,
),
];
if let Some(process_uptime_secs) = self.process_uptime_secs {
samples.push(MetricSample::new(
"gateway_process_uptime_seconds",
"Gateway process uptime in seconds.",
MetricKind::Gauge,
process_uptime_secs,
));
}
samples
}
}
pub(crate) struct GatewayProcessResourceMonitor {
system: Mutex<System>,
networks: Mutex<Networks>,
current_pid: Option<Pid>,
}
impl GatewayProcessResourceMonitor {
pub(crate) fn new() -> Self {
let mut system = System::new_all();
let mut networks = Networks::new_with_refreshed_list();
let current_pid = get_current_pid().ok();
if let Some(pid) = current_pid {
system.refresh_processes(ProcessesToUpdate::Some(&[pid]), true);
}
system.refresh_cpu_usage();
system.refresh_memory();
networks.refresh();
Self {
system: Mutex::new(system),
networks: Mutex::new(networks),
current_pid,
}
}
pub(crate) fn snapshot(&self) -> GatewayProcessResourceSnapshot {
let mut system = match self.system.lock() {
Ok(guard) => guard,
Err(poisoned) => poisoned.into_inner(),
};
system.refresh_cpu_usage();
system.refresh_memory();
if let Some(pid) = self.current_pid {
system.refresh_processes(ProcessesToUpdate::Some(&[pid]), true);
}
let memory_total_bytes = system.total_memory();
let memory_used_bytes = system.used_memory();
let memory_available_bytes = system.available_memory();
let (
process_cpu_usage_basis_points,
process_memory_bytes,
process_virtual_memory_bytes,
process_uptime_secs,
) = self
.current_pid
.and_then(|pid| system.process(pid))
.map(|process| {
(
percent_to_basis_points(process.cpu_usage() as f64),
process.memory(),
process.virtual_memory(),
Some(process.run_time()),
)
})
.unwrap_or((0, 0, 0, None));
let fd_open_count = open_file_descriptors().unwrap_or(0);
let fd_limit = file_descriptor_limit();
let network = self.network_snapshot();
let sockets = socket_snapshot().unwrap_or_default();
GatewayProcessResourceSnapshot {
sampled_at_unix_secs: current_unix_secs(),
system_cpu_usage_basis_points: percent_to_basis_points(system.global_cpu_usage() as f64),
process_cpu_usage_basis_points,
memory_total_bytes,
memory_used_bytes,
memory_available_bytes,
memory_used_basis_points: ratio_to_basis_points(memory_used_bytes, memory_total_bytes),
process_memory_bytes,
process_virtual_memory_bytes,
process_memory_basis_points: ratio_to_basis_points(
process_memory_bytes,
memory_total_bytes,
),
process_uptime_secs,
process_threads: process_thread_count().unwrap_or(0),
fd_open_count,
fd_limit,
fd_usage_basis_points: ratio_to_basis_points(fd_open_count, fd_limit),
network_observability_available: network.observability_available,
network_interface_count: network.interface_count,
network_received_bytes_total: network.received_bytes_total,
network_transmitted_bytes_total: network.transmitted_bytes_total,
network_received_packets_total: network.received_packets_total,
network_transmitted_packets_total: network.transmitted_packets_total,
network_receive_errors_total: network.receive_errors_total,
network_transmit_errors_total: network.transmit_errors_total,
network_receive_dropped_total: network.receive_dropped_total,
network_transmit_dropped_total: network.transmit_dropped_total,
process_socket_fds: sockets.process_socket_fds,
tcp_state_observability_available: sockets.tcp_state_observability_available,
host_tcp_connections: sockets.host_tcp.total,
host_tcp_established_connections: sockets.host_tcp.established,
host_tcp_listen_connections: sockets.host_tcp.listen,
host_tcp_time_wait_connections: sockets.host_tcp.time_wait,
host_tcp_syn_sent_connections: sockets.host_tcp.syn_sent,
host_tcp_syn_recv_connections: sockets.host_tcp.syn_recv,
host_tcp_close_wait_connections: sockets.host_tcp.close_wait,
process_tcp_connections: sockets.process_tcp.total,
process_tcp_established_connections: sockets.process_tcp.established,
process_tcp_listen_connections: sockets.process_tcp.listen,
process_tcp_time_wait_connections: sockets.process_tcp.time_wait,
process_tcp_syn_sent_connections: sockets.process_tcp.syn_sent,
process_tcp_syn_recv_connections: sockets.process_tcp.syn_recv,
process_tcp_close_wait_connections: sockets.process_tcp.close_wait,
}
}
pub(crate) fn metric_samples(&self) -> Vec<MetricSample> {
self.snapshot().to_metric_samples()
}
fn network_snapshot(&self) -> GatewayNetworkSnapshot {
let mut networks = match self.networks.lock() {
Ok(guard) => guard,
Err(poisoned) => poisoned.into_inner(),
};
networks.refresh();
let mut snapshot = GatewayNetworkSnapshot {
observability_available: u64::from(!networks.list().is_empty()),
interface_count: networks.list().len() as u64,
..GatewayNetworkSnapshot::default()
};
for network in networks.list().values() {
snapshot.received_bytes_total = snapshot
.received_bytes_total
.saturating_add(network.total_received());
snapshot.transmitted_bytes_total = snapshot
.transmitted_bytes_total
.saturating_add(network.total_transmitted());
snapshot.received_packets_total = snapshot
.received_packets_total
.saturating_add(network.total_packets_received());
snapshot.transmitted_packets_total = snapshot
.transmitted_packets_total
.saturating_add(network.total_packets_transmitted());
snapshot.receive_errors_total = snapshot
.receive_errors_total
.saturating_add(network.total_errors_on_received());
snapshot.transmit_errors_total = snapshot
.transmit_errors_total
.saturating_add(network.total_errors_on_transmitted());
}
let drops = network_drop_totals();
snapshot.receive_dropped_total = drops.receive_dropped_total;
snapshot.transmit_dropped_total = drops.transmit_dropped_total;
snapshot
}
}
impl Default for GatewayProcessResourceMonitor {
fn default() -> Self {
Self::new()
}
}
impl std::fmt::Debug for GatewayProcessResourceMonitor {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("GatewayProcessResourceMonitor")
.field("current_pid", &self.current_pid)
.finish_non_exhaustive()
}
}
fn current_unix_secs() -> u64 {
SystemTime::now()
.duration_since(UNIX_EPOCH)
.map(|duration| duration.as_secs())
.unwrap_or(0)
}
fn percent_to_basis_points(value: f64) -> u64 {
if !value.is_finite() || value.is_sign_negative() {
0
} else {
(value * 100.0).round().clamp(0.0, u64::MAX as f64) as u64
}
}
fn ratio_to_basis_points(value: u64, total: u64) -> u64 {
value.saturating_mul(10_000).checked_div(total).unwrap_or(0)
}
fn open_file_descriptors() -> Option<u64> {
#[cfg(unix)]
{
for dir in ["/proc/self/fd", "/dev/fd"] {
if let Ok(entries) = std::fs::read_dir(dir) {
return Some(entries.count() as u64);
}
}
}
None
}
fn file_descriptor_limit() -> u64 {
#[cfg(unix)]
{
let mut limit = libc::rlimit {
rlim_cur: 0,
rlim_max: 0,
};
let result = unsafe { libc::getrlimit(libc::RLIMIT_NOFILE, &mut limit) };
if result == 0 {
return limit.rlim_cur;
}
}
0
}
fn process_thread_count() -> Option<u64> {
#[cfg(target_os = "linux")]
{
return std::fs::read_to_string("/proc/self/status")
.ok()
.and_then(|raw| parse_linux_process_thread_count(&raw));
}
#[allow(unreachable_code)]
None
}
#[cfg(target_os = "linux")]
fn parse_linux_process_thread_count(raw: &str) -> Option<u64> {
raw.lines().find_map(|line| {
let value = line.strip_prefix("Threads:")?.trim();
value.parse::<u64>().ok()
})
}
#[derive(Debug, Clone, Copy, Default)]
struct GatewayNetworkSnapshot {
observability_available: u64,
interface_count: u64,
received_bytes_total: u64,
transmitted_bytes_total: u64,
received_packets_total: u64,
transmitted_packets_total: u64,
receive_errors_total: u64,
transmit_errors_total: u64,
receive_dropped_total: u64,
transmit_dropped_total: u64,
}
#[derive(Debug, Clone, Copy, Default)]
struct NetworkDropTotals {
receive_dropped_total: u64,
transmit_dropped_total: u64,
}
fn network_drop_totals() -> NetworkDropTotals {
#[cfg(target_os = "linux")]
{
return std::fs::read_to_string("/proc/net/dev")
.ok()
.map(|raw| parse_linux_network_drop_totals(&raw))
.unwrap_or_default();
}
#[allow(unreachable_code)]
NetworkDropTotals::default()
}
fn parse_linux_network_drop_totals(raw: &str) -> NetworkDropTotals {
let mut totals = NetworkDropTotals::default();
for line in raw.lines().skip(2) {
let Some((_, counters)) = line.split_once(':') else {
continue;
};
let fields: Vec<&str> = counters.split_whitespace().collect();
if fields.len() < 12 {
continue;
}
totals.receive_dropped_total = totals
.receive_dropped_total
.saturating_add(fields[3].parse::<u64>().unwrap_or_default());
totals.transmit_dropped_total = totals
.transmit_dropped_total
.saturating_add(fields[11].parse::<u64>().unwrap_or_default());
}
totals
}
#[derive(Debug, Clone, Copy, Default)]
struct SocketSnapshot {
process_socket_fds: u64,
tcp_state_observability_available: u64,
host_tcp: TcpStateCounts,
process_tcp: TcpStateCounts,
}
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
struct TcpStateCounts {
total: u64,
established: u64,
listen: u64,
time_wait: u64,
syn_sent: u64,
syn_recv: u64,
close_wait: u64,
}
impl TcpStateCounts {
fn observe(&mut self, state: &str) {
self.total = self.total.saturating_add(1);
match state {
"01" => self.established = self.established.saturating_add(1),
"02" => self.syn_sent = self.syn_sent.saturating_add(1),
"03" => self.syn_recv = self.syn_recv.saturating_add(1),
"06" => self.time_wait = self.time_wait.saturating_add(1),
"08" => self.close_wait = self.close_wait.saturating_add(1),
"0A" | "0a" => self.listen = self.listen.saturating_add(1),
_ => {}
}
}
}
fn socket_snapshot() -> Option<SocketSnapshot> {
#[cfg(target_os = "linux")]
{
let process_inodes = process_socket_inodes().unwrap_or_default();
let mut snapshot = SocketSnapshot {
process_socket_fds: process_inodes.len() as u64,
..SocketSnapshot::default()
};
let mut observed_tcp_table = false;
for path in ["/proc/net/tcp", "/proc/net/tcp6"] {
let Ok(raw) = std::fs::read_to_string(path) else {
continue;
};
observed_tcp_table = true;
observe_linux_tcp_table(&raw, &process_inodes, &mut snapshot);
}
if observed_tcp_table {
snapshot.tcp_state_observability_available = 1;
return Some(snapshot);
}
}
#[allow(unreachable_code)]
None
}
#[cfg(target_os = "linux")]
fn process_socket_inodes() -> Option<HashSet<u64>> {
let mut inodes = HashSet::new();
for entry in std::fs::read_dir("/proc/self/fd").ok()? {
let Ok(entry) = entry else {
continue;
};
let Ok(target) = std::fs::read_link(entry.path()) else {
continue;
};
if let Some(inode) = parse_socket_inode(&target.to_string_lossy()) {
inodes.insert(inode);
}
}
Some(inodes)
}
fn parse_socket_inode(target: &str) -> Option<u64> {
target
.strip_prefix("socket:[")
.and_then(|value| value.strip_suffix(']'))
.and_then(|value| value.parse::<u64>().ok())
}
fn observe_linux_tcp_table(
raw: &str,
process_inodes: &HashSet<u64>,
snapshot: &mut SocketSnapshot,
) {
for line in raw.lines().skip(1) {
let fields: Vec<&str> = line.split_whitespace().collect();
if fields.len() <= 9 {
continue;
}
let state = fields[3];
snapshot.host_tcp.observe(state);
let inode = fields[9].parse::<u64>().unwrap_or_default();
if process_inodes.contains(&inode) {
snapshot.process_tcp.observe(state);
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn parses_socket_inode_targets() {
assert_eq!(parse_socket_inode("socket:[12345]"), Some(12345));
assert_eq!(parse_socket_inode("anon_inode:[eventpoll]"), None);
assert_eq!(parse_socket_inode("socket:12345"), None);
}
#[test]
fn parses_linux_tcp_table_state_counts() {
let raw = "\
sl local_address rem_address st tx_queue rx_queue tr tm->when retrnsmt uid timeout inode
0: 0100007F:1F90 00000000:0000 0A 00000000:00000000 00:00000000 00000000 501 0 111 1 0000000000000000 100 0 0 10 0
1: 0100007F:9C40 0100007F:1F90 01 00000000:00000000 00:00000000 00000000 501 0 222 1 0000000000000000 20 4 30 10 -1
2: 0100007F:9C41 0100007F:1F90 08 00000000:00000000 00:00000000 00000000 501 0 333 1 0000000000000000 20 4 30 10 -1
";
let mut inodes = HashSet::new();
inodes.insert(111);
inodes.insert(333);
let mut snapshot = SocketSnapshot::default();
observe_linux_tcp_table(raw, &inodes, &mut snapshot);
assert_eq!(snapshot.host_tcp.total, 3);
assert_eq!(snapshot.host_tcp.listen, 1);
assert_eq!(snapshot.host_tcp.established, 1);
assert_eq!(snapshot.host_tcp.close_wait, 1);
assert_eq!(snapshot.process_tcp.total, 2);
assert_eq!(snapshot.process_tcp.listen, 1);
assert_eq!(snapshot.process_tcp.close_wait, 1);
}
#[test]
fn parses_linux_network_drop_totals() {
let raw = "\
Inter-| Receive | Transmit
face |bytes packets errs drop fifo frame compressed multicast|bytes packets errs drop fifo colls carrier compressed
lo: 100 1 0 2 0 0 0 0 200 2 0 3 0 0 0 0
eth0: 300 3 0 5 0 0 0 0 400 4 0 7 0 0 0 0
";
let totals = parse_linux_network_drop_totals(raw);
assert_eq!(totals.receive_dropped_total, 7);
assert_eq!(totals.transmit_dropped_total, 10);
}
#[test]
fn process_resource_monitor_renders_gateway_metrics() {
let samples = GatewayProcessResourceMonitor::new().metric_samples();
assert!(samples
.iter()
.any(|sample| sample.name == "gateway_process_memory_bytes"));
assert!(samples
.iter()
.any(|sample| sample.name == "gateway_process_open_fds"));
assert!(samples
.iter()
.any(|sample| sample.name == "gateway_process_fd_usage_basis_points"));
assert!(samples
.iter()
.any(|sample| sample.name == "gateway_process_threads"));
assert!(samples
.iter()
.any(|sample| sample.name == "gateway_process_socket_fds"));
assert!(samples
.iter()
.any(|sample| sample.name == "gateway_network_observability_available"));
}
#[cfg(target_os = "linux")]
#[test]
fn parses_linux_process_thread_count() {
let raw = "\
Name:\taether-gateway
State:\tS (sleeping)
Threads:\t42
";
assert_eq!(parse_linux_process_thread_count(raw), Some(42));
assert_eq!(parse_linux_process_thread_count("Name:\ttest\n"), None);
}
}
@@ -17,9 +17,13 @@ const BATCH_SIZE_ENV: &str = "AETHER_GATEWAY_REQUEST_CANDIDATE_QUEUE_BATCH_SIZE"
const FLUSH_INTERVAL_MS_ENV: &str = "AETHER_GATEWAY_REQUEST_CANDIDATE_QUEUE_FLUSH_INTERVAL_MS";
const WORKERS_ENV: &str = "AETHER_GATEWAY_REQUEST_CANDIDATE_QUEUE_WORKERS";
const QUEUE_FULL_ENV: &str = "AETHER_GATEWAY_REQUEST_CANDIDATE_QUEUE_FULL";
const DB_WRITE_CONCURRENCY_LIMIT_ENV: &str =
"AETHER_GATEWAY_REQUEST_CANDIDATE_DB_WRITE_CONCURRENCY_LIMIT";
const DB_BATCH_SIZE_ENV: &str = "AETHER_GATEWAY_REQUEST_CANDIDATE_DB_BATCH_SIZE";
const DEFAULT_QUEUE_CAPACITY: usize = 65_536;
const DEFAULT_BATCH_SIZE: usize = 512;
const DEFAULT_DB_BATCH_SIZE: usize = 128;
const DEFAULT_FLUSH_INTERVAL_MS: u64 = 50;
const DEFAULT_WORKERS: usize = 2;
const FAILED_FLUSH_RETRY_DELAY_MS: u64 = 25;
@@ -41,8 +45,10 @@ pub(crate) struct RequestCandidateQueueConfig {
pub(crate) mode: RequestCandidateWriteMode,
pub(crate) capacity: usize,
pub(crate) batch_size: usize,
pub(crate) db_batch_size: usize,
pub(crate) flush_interval: Duration,
pub(crate) workers: usize,
pub(crate) db_write_concurrency_limit: Option<usize>,
pub(crate) full_policy: RequestCandidateQueueFullPolicy,
}
@@ -52,8 +58,10 @@ impl Default for RequestCandidateQueueConfig {
mode: RequestCandidateWriteMode::Sync,
capacity: DEFAULT_QUEUE_CAPACITY,
batch_size: DEFAULT_BATCH_SIZE,
db_batch_size: DEFAULT_DB_BATCH_SIZE,
flush_interval: Duration::from_millis(DEFAULT_FLUSH_INTERVAL_MS),
workers: DEFAULT_WORKERS,
db_write_concurrency_limit: None,
full_policy: RequestCandidateQueueFullPolicy::Sync,
}
}
@@ -69,9 +77,12 @@ impl RequestCandidateQueueConfig {
};
config.capacity = env_usize(QUEUE_CAPACITY_ENV, DEFAULT_QUEUE_CAPACITY).max(1);
config.batch_size = env_usize(BATCH_SIZE_ENV, DEFAULT_BATCH_SIZE).max(1);
config.db_batch_size = env_usize(DB_BATCH_SIZE_ENV, DEFAULT_DB_BATCH_SIZE).max(1);
config.flush_interval =
Duration::from_millis(env_u64(FLUSH_INTERVAL_MS_ENV, DEFAULT_FLUSH_INTERVAL_MS).max(1));
config.workers = env_usize(WORKERS_ENV, DEFAULT_WORKERS).clamp(1, 32);
config.db_write_concurrency_limit =
env_optional_usize(DB_WRITE_CONCURRENCY_LIMIT_ENV).map(|limit| limit.clamp(1, 32));
config.full_policy = match env_string(QUEUE_FULL_ENV).as_deref() {
Some("drop") | Some("best_effort") | Some("best-effort") => {
RequestCandidateQueueFullPolicy::Drop
@@ -100,6 +111,9 @@ struct RequestCandidateQueueMetrics {
flush_batches_total: AtomicU64,
flush_sql_ops_total: AtomicU64,
flush_sql_records_total: AtomicU64,
db_write_in_flight: AtomicUsize,
db_write_max_in_flight: AtomicUsize,
db_write_wait_total: AtomicU64,
compacted_total: AtomicU64,
sync_fallback_total: AtomicU64,
}
@@ -110,6 +124,7 @@ pub(crate) struct RequestCandidateQueueRuntime {
repository: Arc<dyn RequestCandidateWriteRepository>,
config: RequestCandidateQueueConfig,
metrics: Arc<RequestCandidateQueueMetrics>,
db_write_gate: Option<Arc<RequestCandidateDbWriteGate>>,
}
impl std::fmt::Debug for RequestCandidateQueueRuntime {
@@ -138,11 +153,16 @@ impl RequestCandidateQueueRuntime {
senders.push(sender);
receivers.push(receiver);
}
let db_write_gate = config
.db_write_concurrency_limit
.map(RequestCandidateDbWriteGate::new)
.map(Arc::new);
let runtime = Arc::new(Self {
senders,
repository,
config,
metrics: Arc::new(RequestCandidateQueueMetrics::default()),
db_write_gate,
});
runtime.spawn_workers(receivers);
runtime
@@ -264,6 +284,36 @@ impl RequestCandidateQueueRuntime {
MetricKind::Counter,
self.metrics.flush_sql_records_total.load(Ordering::Acquire),
),
MetricSample::new(
"request_candidate_queue_db_batch_size",
"Maximum request candidate records submitted in one async DB batch upsert after compaction.",
MetricKind::Gauge,
self.config.db_batch_size as u64,
),
MetricSample::new(
"request_candidate_queue_db_write_concurrency_limit",
"Maximum concurrent request candidate async DB write batches; zero means unlimited.",
MetricKind::Gauge,
self.config.db_write_concurrency_limit.unwrap_or_default() as u64,
),
MetricSample::new(
"request_candidate_queue_db_write_in_flight",
"Current request candidate async DB write batches in flight.",
MetricKind::Gauge,
self.metrics.db_write_in_flight.load(Ordering::Acquire) as u64,
),
MetricSample::new(
"request_candidate_queue_db_write_max_in_flight",
"Maximum observed request candidate async DB write batches in flight.",
MetricKind::Gauge,
self.metrics.db_write_max_in_flight.load(Ordering::Acquire) as u64,
),
MetricSample::new(
"request_candidate_queue_db_write_wait_total",
"Total request candidate async DB write batches that had to wait for the DB write gate.",
MetricKind::Counter,
self.metrics.db_write_wait_total.load(Ordering::Acquire),
),
MetricSample::new(
"request_candidate_queue_compacted_total",
"Total request candidate records compacted before async persistence because a later queued record covered the same request candidate slot.",
@@ -287,8 +337,17 @@ impl RequestCandidateQueueRuntime {
let repository = Arc::clone(&self.repository);
let config = self.config.clone();
let metrics = Arc::clone(&self.metrics);
let db_write_gate = self.db_write_gate.clone();
tokio::spawn(async move {
run_worker(repository, config, metrics, worker_index, receiver).await;
run_worker(
repository,
config,
metrics,
db_write_gate,
worker_index,
receiver,
)
.await;
});
}
}
@@ -306,6 +365,7 @@ async fn run_worker(
repository: Arc<dyn RequestCandidateWriteRepository>,
config: RequestCandidateQueueConfig,
metrics: Arc<RequestCandidateQueueMetrics>,
db_write_gate: Option<Arc<RequestCandidateDbWriteGate>>,
worker_index: usize,
mut receiver: mpsc::Receiver<UpsertRequestCandidateRecord>,
) {
@@ -317,7 +377,14 @@ async fn run_worker(
tokio::select! {
_ = ticker.tick() => {
if !batch.is_empty() {
flush_batch(&repository, &metrics, worker_index, &mut batch).await;
flush_batch(
&repository,
&config,
&metrics,
db_write_gate.as_ref(),
worker_index,
&mut batch,
).await;
}
}
received = receiver.recv() => {
@@ -326,12 +393,26 @@ async fn run_worker(
decrement_atomic_usize(&metrics.queued_current);
batch.push(record);
if batch.len() >= config.batch_size {
flush_batch(&repository, &metrics, worker_index, &mut batch).await;
flush_batch(
&repository,
&config,
&metrics,
db_write_gate.as_ref(),
worker_index,
&mut batch,
).await;
}
}
None => {
if !batch.is_empty() {
flush_batch(&repository, &metrics, worker_index, &mut batch).await;
flush_batch(
&repository,
&config,
&metrics,
db_write_gate.as_ref(),
worker_index,
&mut batch,
).await;
}
break;
}
@@ -343,7 +424,9 @@ async fn run_worker(
async fn flush_batch(
repository: &Arc<dyn RequestCandidateWriteRepository>,
config: &RequestCandidateQueueConfig,
metrics: &RequestCandidateQueueMetrics,
db_write_gate: Option<&Arc<RequestCandidateDbWriteGate>>,
worker_index: usize,
batch: &mut Vec<UpsertRequestCandidateRecord>,
) {
@@ -360,42 +443,49 @@ async fn flush_batch(
.fetch_add(compacted as u64, Ordering::AcqRel);
}
metrics.flush_batches_total.fetch_add(1, Ordering::AcqRel);
let source_count = records
.iter()
.map(|record| record.source_count)
.sum::<usize>();
let record_count = records.len();
let upsert_records = records
.into_iter()
.map(|record| record.record)
.collect::<Vec<_>>();
metrics.flush_sql_ops_total.fetch_add(1, Ordering::AcqRel);
metrics
.flush_sql_records_total
.fetch_add(record_count as u64, Ordering::AcqRel);
let mut failed = 0_u64;
let mut retry_records = Vec::new();
if let Err(err) = repository.upsert_many(upsert_records.clone()).await {
failed = source_count as u64;
decrement_atomic_usize_by(
&metrics.pending_current,
source_count.saturating_sub(record_count),
);
warn!(
event_name = "request_candidate_async_flush_failed",
log_type = "event",
worker_index,
record_count,
source_count,
error = ?err,
"gateway failed to asynchronously persist request candidate batch"
);
retry_records = upsert_records;
} else {
for chunk in records.chunks(config.db_batch_size) {
let source_count = chunk
.iter()
.map(|record| record.source_count)
.sum::<usize>();
let record_count = chunk.len();
let upsert_records = chunk
.iter()
.map(|record| record.record.clone())
.collect::<Vec<_>>();
metrics.flush_sql_ops_total.fetch_add(1, Ordering::AcqRel);
metrics
.flushed_total
.fetch_add(source_count as u64, Ordering::AcqRel);
decrement_atomic_usize_by(&metrics.pending_current, source_count);
.flush_sql_records_total
.fetch_add(record_count as u64, Ordering::AcqRel);
let _db_write_permit = match db_write_gate {
Some(gate) => Some(gate.acquire(metrics).await),
None => None,
};
if let Err(err) = repository.upsert_many(upsert_records.clone()).await {
failed = failed.saturating_add(source_count as u64);
decrement_atomic_usize_by(
&metrics.pending_current,
source_count.saturating_sub(record_count),
);
warn!(
event_name = "request_candidate_async_flush_failed",
log_type = "event",
worker_index,
record_count,
source_count,
error = ?err,
"gateway failed to asynchronously persist request candidate DB batch"
);
retry_records.extend(upsert_records);
} else {
metrics
.flushed_total
.fetch_add(source_count as u64, Ordering::AcqRel);
decrement_atomic_usize_by(&metrics.pending_current, source_count);
}
}
if failed > 0 {
metrics
@@ -476,6 +566,54 @@ fn worker_queue_capacity(total_capacity: usize, workers: usize, worker_index: us
(base + usize::from(worker_index < remainder)).max(1)
}
#[derive(Debug)]
struct RequestCandidateDbWriteGate {
semaphore: tokio::sync::Semaphore,
}
impl RequestCandidateDbWriteGate {
fn new(limit: usize) -> Self {
Self {
semaphore: tokio::sync::Semaphore::new(limit.max(1)),
}
}
async fn acquire<'a>(
&'a self,
metrics: &'a RequestCandidateQueueMetrics,
) -> RequestCandidateDbWritePermit<'a> {
if self.semaphore.available_permits() == 0 {
metrics.db_write_wait_total.fetch_add(1, Ordering::AcqRel);
}
let permit = self
.semaphore
.acquire()
.await
.expect("request candidate DB write gate semaphore should not be closed");
let in_flight = metrics.db_write_in_flight.fetch_add(1, Ordering::AcqRel) + 1;
metrics
.db_write_max_in_flight
.fetch_max(in_flight, Ordering::AcqRel);
RequestCandidateDbWritePermit {
metrics,
_permit: permit,
}
}
}
struct RequestCandidateDbWritePermit<'a> {
metrics: &'a RequestCandidateQueueMetrics,
_permit: tokio::sync::SemaphorePermit<'a>,
}
impl Drop for RequestCandidateDbWritePermit<'_> {
fn drop(&mut self) {
self.metrics
.db_write_in_flight
.fetch_sub(1, Ordering::AcqRel);
}
}
fn merge_request_candidate_record(
target: &mut UpsertRequestCandidateRecord,
incoming: UpsertRequestCandidateRecord,
@@ -601,6 +739,13 @@ fn env_usize(key: &str, default: usize) -> usize {
.unwrap_or(default)
}
fn env_optional_usize(key: &str) -> Option<usize> {
std::env::var(key)
.ok()
.and_then(|value| value.trim().parse::<usize>().ok())
.filter(|value| *value > 0)
}
fn env_u64(key: &str, default: u64) -> u64 {
std::env::var(key)
.ok()
@@ -666,6 +811,8 @@ mod tests {
inner: InMemoryRequestCandidateRepository,
upsert_calls: AtomicUsize,
upsert_many_calls: AtomicUsize,
active_upsert_many: AtomicUsize,
max_active_upsert_many: AtomicUsize,
}
#[async_trait::async_trait]
@@ -683,10 +830,15 @@ mod tests {
candidates: Vec<UpsertRequestCandidateRecord>,
) -> Result<usize, DataLayerError> {
self.upsert_many_calls.fetch_add(1, Ordering::AcqRel);
let active = self.active_upsert_many.fetch_add(1, Ordering::AcqRel) + 1;
self.max_active_upsert_many
.fetch_max(active, Ordering::AcqRel);
tokio::time::sleep(Duration::from_millis(30)).await;
let count = candidates.len();
for candidate in candidates {
self.inner.upsert(candidate).await?;
}
self.active_upsert_many.fetch_sub(1, Ordering::AcqRel);
Ok(count)
}
@@ -826,8 +978,10 @@ mod tests {
mode: super::RequestCandidateWriteMode::Async,
capacity: 16,
batch_size: 2,
db_batch_size: 128,
flush_interval: Duration::from_millis(10),
workers: 1,
db_write_concurrency_limit: None,
full_policy: super::RequestCandidateQueueFullPolicy::Drop,
},
);
@@ -858,8 +1012,10 @@ mod tests {
mode: super::RequestCandidateWriteMode::Async,
capacity: 16,
batch_size: 4,
db_batch_size: 128,
flush_interval: Duration::from_millis(100),
workers: 1,
db_write_concurrency_limit: None,
full_policy: super::RequestCandidateQueueFullPolicy::Drop,
},
);
@@ -899,6 +1055,109 @@ mod tests {
panic!("async request candidate queue did not finish batch flush in time");
}
#[tokio::test]
async fn async_queue_splits_compacted_flush_into_db_batches() {
let repository = Arc::new(CountingBatchRequestCandidateRepository::default());
let runtime = RequestCandidateQueueRuntime::spawn(
repository.clone(),
RequestCandidateQueueConfig {
mode: super::RequestCandidateWriteMode::Async,
capacity: 16,
batch_size: 5,
db_batch_size: 2,
flush_interval: Duration::from_millis(100),
workers: 1,
db_write_concurrency_limit: None,
full_policy: super::RequestCandidateQueueFullPolicy::Drop,
},
);
for index in 0..5 {
runtime
.enqueue_or_fallback(record(
"req-db-batch",
index,
0,
RequestCandidateStatus::Success,
))
.await
.unwrap();
}
for _ in 0..50 {
if runtime.metrics.pending_current.load(Ordering::Acquire) == 0 {
assert_eq!(repository.upsert_many_calls.load(Ordering::Acquire), 3);
assert_eq!(
runtime.metrics.flush_batches_total.load(Ordering::Acquire),
1
);
assert_eq!(
runtime.metrics.flush_sql_ops_total.load(Ordering::Acquire),
3
);
assert_eq!(
runtime
.metrics
.flush_sql_records_total
.load(Ordering::Acquire),
5
);
return;
}
tokio::time::sleep(Duration::from_millis(10)).await;
}
panic!("async request candidate queue did not split DB batches in time");
}
#[tokio::test]
async fn async_queue_db_write_gate_limits_concurrent_batch_writes() {
let repository = Arc::new(CountingBatchRequestCandidateRepository::default());
let runtime = RequestCandidateQueueRuntime::spawn(
repository.clone(),
RequestCandidateQueueConfig {
mode: super::RequestCandidateWriteMode::Async,
capacity: 64,
batch_size: 1,
db_batch_size: 128,
flush_interval: Duration::from_millis(100),
workers: 4,
db_write_concurrency_limit: Some(2),
full_policy: super::RequestCandidateQueueFullPolicy::Drop,
},
);
for index in 0..8 {
runtime
.enqueue_or_fallback(record(
&format!("req-gate-{index}"),
0,
0,
RequestCandidateStatus::Success,
))
.await
.unwrap();
}
for _ in 0..100 {
if runtime.metrics.pending_current.load(Ordering::Acquire) == 0 {
assert_eq!(repository.max_active_upsert_many.load(Ordering::Acquire), 2);
assert_eq!(
runtime
.metrics
.db_write_max_in_flight
.load(Ordering::Acquire),
2
);
assert!(runtime.metrics.db_write_wait_total.load(Ordering::Acquire) > 0);
return;
}
tokio::time::sleep(Duration::from_millis(10)).await;
}
panic!("async request candidate queue did not finish gated writes in time");
}
#[tokio::test]
async fn async_queue_preserves_same_slot_order_with_multiple_workers() {
let repository = Arc::new(DelayedPendingRequestCandidateRepository::default());
@@ -908,8 +1167,10 @@ mod tests {
mode: super::RequestCandidateWriteMode::Async,
capacity: 16,
batch_size: 1,
db_batch_size: 128,
flush_interval: Duration::from_millis(100),
workers: 2,
db_write_concurrency_limit: None,
full_policy: super::RequestCandidateQueueFullPolicy::Drop,
},
);
+39
View File
@@ -20,8 +20,10 @@ use super::super::cache::{
};
use super::super::data::GatewayDataState;
use super::super::fallback_metrics;
use super::super::maintenance::UsageCounterFlushRuntimeMetrics;
use super::super::rate_limit::FrontdoorUserRpmLimiter;
use super::super::request_candidate_queue::RequestCandidateQueueRuntime;
use super::super::task_runtime::TaskSupervisorMetrics;
use super::super::{provider_transport, usage};
use super::{
AdminBillingCollectorRecord, AdminBillingRuleRecord, AdminPaymentCallbackRecord,
@@ -40,18 +42,22 @@ const MIN_LOCAL_EXECUTION_PLANNING_TIMEOUT_MS: u64 = 500;
const MAX_LOCAL_EXECUTION_PLANNING_TIMEOUT_MS: u64 = 120_000;
const LOCAL_EXECUTION_PLANNING_TIMEOUT_MS_ENV: &str =
"AETHER_GATEWAY_LOCAL_EXECUTION_PLANNING_TIMEOUT_MS";
const DEFAULT_AUTH_SNAPSHOT_LOAD_GATE_LIMIT: usize = 64;
const DEFAULT_CANDIDATE_PLANNING_GATE_LIMIT: usize = 1024;
const DEFAULT_UPSTREAM_EXECUTION_GATE_LIMIT: usize = 10_000;
const DEFAULT_UPSTREAM_TARGET_GATE_LIMIT: usize = 10_000;
const MAX_AUTH_SNAPSHOT_LOAD_GATE_LIMIT: usize = 1024;
const MAX_CANDIDATE_PLANNING_GATE_LIMIT: usize = 8192;
const MAX_UPSTREAM_EXECUTION_GATE_LIMIT: usize = 16_384;
const MAX_UPSTREAM_TARGET_GATE_LIMIT: usize = 16_384;
const AUTH_SNAPSHOT_LOAD_GATE_LIMIT_PER_CPU: usize = 16;
const CANDIDATE_PLANNING_GATE_LIMIT_PER_CPU: usize = 256;
const UPSTREAM_EXECUTION_GATE_LIMIT_PER_CPU: usize = 1024;
const UPSTREAM_TARGET_GATE_LIMIT_PER_CPU: usize = 1024;
const GATE_LIMIT_FD_RESERVE: usize = 128;
const DEFAULT_INTERNAL_GATE_QUEUE_BUDGET_MS: u64 = 250;
const MAX_INTERNAL_GATE_QUEUE_BUDGET_MS: u64 = 5_000;
const AUTH_SNAPSHOT_LOAD_GATE_LIMIT_ENV: &str = "AETHER_GATEWAY_AUTH_SNAPSHOT_LOAD_GATE_LIMIT";
const CANDIDATE_PLANNING_GATE_LIMIT_ENV: &str = "AETHER_GATEWAY_CANDIDATE_PLANNING_GATE_LIMIT";
const UPSTREAM_EXECUTION_GATE_LIMIT_ENV: &str = "AETHER_GATEWAY_UPSTREAM_EXECUTION_GATE_LIMIT";
const UPSTREAM_TARGET_GATE_LIMIT_ENV: &str = "AETHER_GATEWAY_UPSTREAM_TARGET_GATE_LIMIT";
@@ -87,6 +93,7 @@ pub(crate) struct FrontdoorRuntimeGuardConfig {
pub(crate) local_execution_planning_timeout: Duration,
pub(crate) internal_gate_queue_budget: Duration,
pub(crate) auth_capacity_cache_ttl: Duration,
pub(crate) auth_snapshot_load_gate_limit: Option<usize>,
pub(crate) candidate_planning_gate_limit: Option<usize>,
pub(crate) upstream_execution_gate_limit: Option<usize>,
pub(crate) upstream_target_gate_limit: Option<usize>,
@@ -119,6 +126,7 @@ impl FrontdoorRuntimeGuardConfig {
MIN_AUTH_CAPACITY_CACHE_TTL_MS,
MAX_AUTH_CAPACITY_CACHE_TTL_MS,
),
auth_snapshot_load_gate_limit: auth_snapshot_load_gate_limit_from_env(),
candidate_planning_gate_limit: candidate_planning_gate_limit_from_env(),
upstream_execution_gate_limit: upstream_execution_gate_limit_from_env(),
upstream_target_gate_limit: upstream_target_gate_limit_from_env(),
@@ -137,6 +145,7 @@ impl FrontdoorRuntimeGuardConfig {
DEFAULT_INTERNAL_GATE_QUEUE_BUDGET_MS,
),
auth_capacity_cache_ttl: Duration::from_millis(DEFAULT_AUTH_CAPACITY_CACHE_TTL_MS),
auth_snapshot_load_gate_limit: Some(DEFAULT_AUTH_SNAPSHOT_LOAD_GATE_LIMIT),
candidate_planning_gate_limit: Some(DEFAULT_CANDIDATE_PLANNING_GATE_LIMIT),
upstream_execution_gate_limit: Some(DEFAULT_UPSTREAM_EXECUTION_GATE_LIMIT),
upstream_target_gate_limit: Some(DEFAULT_UPSTREAM_TARGET_GATE_LIMIT),
@@ -192,6 +201,13 @@ const CANDIDATE_PLANNING_GATE_AUTO_PROFILE: GateAutoProfile = GateAutoProfile {
fd_divisor: None,
};
const AUTH_SNAPSHOT_LOAD_GATE_AUTO_PROFILE: GateAutoProfile = GateAutoProfile {
floor: DEFAULT_AUTH_SNAPSHOT_LOAD_GATE_LIMIT,
cap: MAX_AUTH_SNAPSHOT_LOAD_GATE_LIMIT,
per_cpu: AUTH_SNAPSHOT_LOAD_GATE_LIMIT_PER_CPU,
fd_divisor: None,
};
const UPSTREAM_EXECUTION_GATE_AUTO_PROFILE: GateAutoProfile = GateAutoProfile {
floor: DEFAULT_UPSTREAM_EXECUTION_GATE_LIMIT,
cap: MAX_UPSTREAM_EXECUTION_GATE_LIMIT,
@@ -213,6 +229,13 @@ fn candidate_planning_gate_limit_from_env() -> Option<usize> {
)
}
fn auth_snapshot_load_gate_limit_from_env() -> Option<usize> {
env_gate_limit(
AUTH_SNAPSHOT_LOAD_GATE_LIMIT_ENV,
AUTH_SNAPSHOT_LOAD_GATE_AUTO_PROFILE,
)
}
fn upstream_execution_gate_limit_from_env() -> Option<usize> {
env_gate_limit(
UPSTREAM_EXECUTION_GATE_LIMIT_ENV,
@@ -315,6 +338,7 @@ pub struct AppState {
pub(crate) video_task_poller: Option<VideoTaskPollerConfig>,
pub(crate) frontdoor_runtime_guards: Arc<FrontdoorRuntimeGuardConfig>,
pub(crate) request_gate: Option<Arc<ConcurrencyGate>>,
pub(crate) auth_snapshot_load_gate: Option<Arc<ConcurrencyGate>>,
pub(crate) candidate_planning_gate: Option<Arc<ConcurrencyGate>>,
pub(crate) upstream_execution_gate: Option<Arc<ConcurrencyGate>>,
pub(crate) upstream_target_admission: Arc<crate::upstream_admission::UpstreamTargetAdmission>,
@@ -349,6 +373,9 @@ pub struct AppState {
pub(crate) chat_pii_redaction_runtime_config_cache:
crate::privacy::ChatPiiRedactionRuntimeConfigCacheHandle,
pub(crate) fallback_metrics: Arc<fallback_metrics::GatewayFallbackMetrics>,
pub(crate) usage_counter_flush_metrics: Arc<UsageCounterFlushRuntimeMetrics>,
pub(crate) task_supervisor_metrics: TaskSupervisorMetrics,
pub(crate) process_resource_monitor: Arc<crate::process_metrics::GatewayProcessResourceMonitor>,
pub(crate) request_candidate_queue: Option<Arc<RequestCandidateQueueRuntime>>,
pub(crate) frontdoor_cors: Option<Arc<FrontdoorCorsConfig>>,
pub(crate) frontdoor_user_rpm: Arc<FrontdoorUserRpmLimiter>,
@@ -454,6 +481,18 @@ mod tests {
);
}
#[test]
fn auth_snapshot_load_gate_auto_limit_uses_smaller_frontdoor_profile() {
assert_eq!(
parse_gate_limit_value(
Some("auto"),
AUTH_SNAPSHOT_LOAD_GATE_AUTO_PROFILE,
TEST_CAPACITY
),
Some(192)
);
}
#[test]
fn gate_limit_parser_accepts_fixed_numbers() {
assert_eq!(
File diff suppressed because it is too large Load Diff
@@ -13,6 +13,18 @@ const AUTH_API_KEY_SNAPSHOT_RUNTIME_CACHE_TTL: Duration = Duration::from_secs(30
use super::super::super::{AUTH_API_KEY_LAST_USED_MAX_ENTRIES, AUTH_API_KEY_LAST_USED_TTL};
impl AppState {
async fn acquire_auth_snapshot_load_gate(
&self,
) -> Result<Option<aether_runtime::ConcurrencyPermit>, GatewayError> {
let Some(gate) = self.auth_snapshot_load_gate.as_ref() else {
return Ok(None);
};
gate.acquire()
.await
.map(Some)
.map_err(|err| GatewayError::Internal(err.to_string()))
}
pub(crate) async fn read_cached_auth_api_key_snapshot(
&self,
user_id: &str,
@@ -28,6 +40,7 @@ impl AppState {
cache_key,
AUTH_API_KEY_SNAPSHOT_RUNTIME_CACHE_TTL,
|| async move {
let _permit = self.acquire_auth_snapshot_load_gate().await?;
self.data
.read_auth_api_key_snapshot(user_id, api_key_id, now_unix_secs)
.await
@@ -52,6 +65,7 @@ impl AppState {
cache_key.clone(),
AUTH_API_KEY_SNAPSHOT_RUNTIME_CACHE_TTL,
|| async move {
let _permit = self.acquire_auth_snapshot_load_gate().await?;
self.data
.read_auth_api_key_snapshot_by_key_hash(key_hash, now_unix_secs)
.await
+1 -1
View File
@@ -6,8 +6,8 @@ use aether_data_contracts::repository::background_tasks::{
UpsertBackgroundTaskRun,
};
use aether_runtime::task::spawn_named;
pub(crate) use aether_task_runtime::TaskSupervisor;
use aether_task_runtime::{RetryPolicy, TaskDefinition, TaskKind};
pub(crate) use aether_task_runtime::{TaskSupervisor, TaskSupervisorMetrics};
use serde_json::Value;
use tokio::task::JoinHandle;
use tracing::warn;
@@ -343,6 +343,63 @@ async fn gateway_exposes_request_concurrency_metrics_impl() {
assert!(body.contains("tunnel_proxy_connections 0"));
assert!(body.contains("tunnel_nodes 0"));
assert!(body.contains("tunnel_active_streams 0"));
assert!(body.contains("gateway_process_cpu_usage_basis_points "));
assert!(body.contains("gateway_process_memory_bytes "));
assert!(body.contains("gateway_process_threads "));
assert!(body.contains("gateway_process_open_fds "));
assert!(body.contains("gateway_process_fd_limit "));
assert!(body.contains("gateway_process_socket_fds "));
assert!(body.contains("gateway_allocator_observability_available "));
assert!(body.contains("gateway_allocator_allocated_bytes "));
assert!(body.contains("gateway_allocator_active_bytes "));
assert!(body.contains("gateway_allocator_resident_bytes "));
assert!(body.contains("gateway_allocator_active_to_allocated_basis_points "));
assert!(body.contains("gateway_network_observability_available "));
assert!(body.contains("gateway_network_received_bytes_total "));
assert!(body.contains("gateway_tcp_state_observability_available "));
assert!(body.contains("gateway_host_tcp_established_connections "));
assert!(body.contains("gateway_process_tcp_established_connections "));
assert!(body.contains("postgres_observability_available{driver=\"postgres\"} 0"));
assert!(body.contains("postgres_observability_unavailable{driver=\"postgres\"} 0"));
assert!(body.contains("postgres_lock_waiting_connections{driver=\"postgres\"} 0"));
assert!(body.contains("postgres_oldest_active_query_age_ms{driver=\"postgres\"} 0"));
assert!(body.contains("postgres_oldest_transaction_age_ms{driver=\"postgres\"} 0"));
assert!(body.contains("redis_runtime_enabled{backend=\"redis\"} 0"));
assert!(body.contains("redis_runtime_health_unavailable{backend=\"redis\"} 0"));
assert!(body.contains("redis_runtime_connected_clients{backend=\"redis\"} 0"));
assert!(body.contains("redis_runtime_used_memory_bytes{backend=\"redis\"} 0"));
assert!(body.contains("usage_runtime_queue_worker_read_batches_total 0"));
assert!(body.contains("usage_runtime_queue_worker_read_entries_total 0"));
assert!(body.contains("usage_runtime_queue_worker_reclaimed_entries_total 0"));
assert!(body.contains("usage_runtime_queue_worker_acked_entries_total 0"));
assert!(body.contains("usage_runtime_queue_worker_dead_lettered_entries_total 0"));
assert!(body.contains("usage_runtime_queue_worker_process_failures_total 0"));
assert!(body.contains("usage_runtime_queue_worker_read_failures_total 0"));
assert!(body.contains("usage_runtime_queue_worker_reclaim_failures_total 0"));
assert!(body.contains("usage_queue_health_unavailable 0"));
assert!(
body.contains("usage_queue_enabled{stream=\"usage:events\",group=\"usage_consumers\"} 0")
);
assert!(body
.contains("usage_queue_configured{stream=\"usage:events\",group=\"usage_consumers\"} 0"));
assert!(body.contains("usage_queue_dlq_length{stream=\"usage:events:dlq\"} 0"));
assert!(body.contains("usage_counter_health_unavailable 0"));
assert!(body.contains("usage_counter_outbox_pending_rows 0"));
assert!(body.contains("usage_counter_outbox_oldest_pending_age_seconds 0"));
assert!(body.contains("usage_counter_outbox_flush_batches_total 0"));
assert!(body.contains("usage_counter_outbox_flush_rows_claimed_total 0"));
assert!(body.contains("usage_counter_outbox_flush_failed_batches_total 0"));
assert!(body.contains("usage_counter_outbox_cleanup_rows_total 0"));
assert!(body.contains("usage_counter_outbox_cleanup_failed_batches_total 0"));
assert!(body.contains("gateway_background_tasks_active 0"));
assert!(body.contains("gateway_background_tasks_supervised_total 0"));
assert!(body.contains("gateway_background_tasks_unexpected_exits_total 0"));
assert!(body.contains("gateway_background_tasks_panicked_total 0"));
assert!(body.contains("gateway_background_tasks_aborted_total 0"));
assert!(body.contains("gateway_tokio_runtime_observability_available 1"));
assert!(body.contains("gateway_tokio_runtime_workers "));
assert!(body.contains("gateway_tokio_runtime_alive_tasks "));
assert!(body.contains("gateway_tokio_runtime_global_queue_depth "));
gateway_handle.abort();
}
+89
View File
@@ -0,0 +1,89 @@
use aether_runtime::{MetricKind, MetricSample};
pub(crate) fn gateway_tokio_runtime_metric_samples() -> Vec<MetricSample> {
let Ok(handle) = tokio::runtime::Handle::try_current() else {
return vec![
availability_sample(0),
gauge(
"gateway_tokio_runtime_workers",
TOKIO_RUNTIME_WORKERS_HELP,
0,
),
gauge(
"gateway_tokio_runtime_alive_tasks",
TOKIO_RUNTIME_ALIVE_TASKS_HELP,
0,
),
gauge(
"gateway_tokio_runtime_global_queue_depth",
TOKIO_RUNTIME_GLOBAL_QUEUE_DEPTH_HELP,
0,
),
];
};
let metrics = handle.metrics();
vec![
availability_sample(1),
gauge(
"gateway_tokio_runtime_workers",
TOKIO_RUNTIME_WORKERS_HELP,
u64_from_usize(metrics.num_workers()),
),
gauge(
"gateway_tokio_runtime_alive_tasks",
TOKIO_RUNTIME_ALIVE_TASKS_HELP,
u64_from_usize(metrics.num_alive_tasks()),
),
gauge(
"gateway_tokio_runtime_global_queue_depth",
TOKIO_RUNTIME_GLOBAL_QUEUE_DEPTH_HELP,
u64_from_usize(metrics.global_queue_depth()),
),
]
}
const TOKIO_RUNTIME_WORKERS_HELP: &str =
"Number of worker threads configured for the gateway Tokio runtime.";
const TOKIO_RUNTIME_ALIVE_TASKS_HELP: &str =
"Current number of alive tasks tracked by the gateway Tokio runtime.";
const TOKIO_RUNTIME_GLOBAL_QUEUE_DEPTH_HELP: &str =
"Current number of tasks waiting in the gateway Tokio runtime global queue.";
fn availability_sample(value: u64) -> MetricSample {
gauge(
"gateway_tokio_runtime_observability_available",
"Whether gateway Tokio runtime metrics were available for this scrape.",
value,
)
}
fn gauge(name: &'static str, help: &'static str, value: u64) -> MetricSample {
MetricSample::new(name, help, MetricKind::Gauge, value)
}
fn u64_from_usize(value: usize) -> u64 {
u64::try_from(value).unwrap_or(u64::MAX)
}
#[cfg(test)]
mod tests {
use super::gateway_tokio_runtime_metric_samples;
#[tokio::test]
async fn renders_tokio_runtime_metrics_inside_runtime() {
let samples = gateway_tokio_runtime_metric_samples();
assert!(samples.iter().any(|sample| {
sample.name == "gateway_tokio_runtime_observability_available" && sample.value == 1
}));
assert!(samples
.iter()
.any(|sample| sample.name == "gateway_tokio_runtime_workers"));
assert!(samples
.iter()
.any(|sample| sample.name == "gateway_tokio_runtime_alive_tasks"));
assert!(samples
.iter()
.any(|sample| sample.name == "gateway_tokio_runtime_global_queue_depth"));
}
}
+1 -1
View File
@@ -5,11 +5,11 @@ pub(crate) mod write;
pub(crate) use aether_usage_runtime::UsageRuntime;
pub use aether_usage_runtime::UsageRuntimeConfig;
pub(crate) use aether_usage_runtime::UsageRuntimeMetricsSnapshot;
pub(crate) use aether_usage_runtime::{
now_ms, UsageEvent, UsageEventData, UsageEventType, UsageQueue, UsageRequestRecordLevel,
USAGE_EVENT_VERSION,
};
pub(crate) use aether_usage_runtime::{UsageQueueHealthSnapshot, UsageRuntimeMetricsSnapshot};
pub(crate) use reporting::{
spawn_sync_report, submit_stream_report, submit_sync_report, GatewayStreamReportRequest,
GatewaySyncReportRequest,
+649 -1
View File
@@ -3,7 +3,8 @@ use super::{
};
use crate::error::{SqlResultExt, SqlxResultExt};
use crate::maintenance::{
DatabaseMaintenanceSummary, DatabasePoolSummary, StatsDailyAggregationInput,
DatabaseMaintenanceSummary, DatabasePoolSummary, DatabasePostgresActivityGroup,
DatabasePostgresObservabilitySnapshot, StatsDailyAggregationInput,
StatsDailyAggregationSummary, StatsHourlyAggregationInput, StatsHourlyAggregationSummary,
WalletDailyUsageAggregationInput, WalletDailyUsageAggregationResult,
};
@@ -14,6 +15,7 @@ use crate::repository::system::{
};
use crate::DataLayerError;
use sqlx::migrate::MigrateError;
use sqlx::Row;
fn maintenance_identifier(value: &str) -> Result<&str, DataLayerError> {
let valid = !value.is_empty()
@@ -109,6 +111,25 @@ impl DataBackends {
self.sql_backend().map(SqlBackendRef::database_pool_summary)
}
pub async fn postgres_observability_snapshot(
&self,
) -> Result<Option<DatabasePostgresObservabilitySnapshot>, DataLayerError> {
match self.postgres() {
Some(postgres) => postgres.postgres_observability_snapshot().await.map(Some),
None => Ok(None),
}
}
pub async fn postgres_activity_groups(
&self,
limit: i64,
) -> Result<Vec<DatabasePostgresActivityGroup>, DataLayerError> {
match self.postgres() {
Some(postgres) => postgres.postgres_activity_groups(limit).await,
None => Ok(Vec::new()),
}
}
pub async fn aggregate_wallet_daily_usage(
&self,
input: &WalletDailyUsageAggregationInput,
@@ -260,6 +281,633 @@ impl PostgresBackend {
}
Ok(summary)
}
pub async fn postgres_observability_snapshot(
&self,
) -> Result<DatabasePostgresObservabilitySnapshot, DataLayerError> {
const ACTIVITY_SQL: &str = r#"
SELECT
COUNT(*) FILTER (WHERE state = 'active')::BIGINT AS active_connections,
COUNT(*) FILTER (WHERE state = 'idle')::BIGINT AS idle_connections,
COUNT(*) FILTER (WHERE state = 'idle in transaction')::BIGINT AS idle_in_transaction_connections,
COUNT(*) FILTER (WHERE state = 'active' AND wait_event_type IS NOT NULL)::BIGINT AS waiting_connections,
COUNT(*) FILTER (WHERE state = 'active' AND wait_event_type = 'Lock')::BIGINT AS lock_waiting_connections,
COALESCE(MAX(EXTRACT(EPOCH FROM now() - query_start) * 1000) FILTER (WHERE state = 'active' AND query_start IS NOT NULL), 0)::BIGINT AS oldest_active_query_age_ms,
COALESCE(MAX(EXTRACT(EPOCH FROM now() - xact_start) * 1000) FILTER (WHERE xact_start IS NOT NULL), 0)::BIGINT AS oldest_transaction_age_ms
FROM pg_stat_activity
WHERE datname = current_database()
AND pid <> pg_backend_pid()
"#;
const DEADLOCKS_SQL: &str = r#"
SELECT
COALESCE(SUM(deadlocks), 0)::BIGINT AS deadlocks_total,
COALESCE(SUM(blks_read), 0)::BIGINT AS block_read_total,
COALESCE(SUM(blks_hit), 0)::BIGINT AS block_hit_total,
COALESCE(SUM(temp_files), 0)::BIGINT AS temp_files_total,
COALESCE(SUM(temp_bytes), 0)::BIGINT AS temp_bytes_total,
COALESCE(SUM(xact_commit), 0)::BIGINT AS xact_commit_total,
COALESCE(SUM(xact_rollback), 0)::BIGINT AS xact_rollback_total
FROM pg_stat_database
WHERE datname = current_database()
"#;
let activity = sqlx::query(ACTIVITY_SQL)
.fetch_one(self.pool())
.await
.map_postgres_err()?;
let database = sqlx::query(DEADLOCKS_SQL)
.fetch_one(self.pool())
.await
.map_postgres_err()?;
let wal = self.postgres_wal_observability_snapshot().await;
let checkpoint = self.postgres_checkpoint_observability_snapshot().await;
let statements = self.postgres_statement_observability_snapshot().await;
let block_read_total = row_u64(&database, "block_read_total")?;
let block_hit_total = row_u64(&database, "block_hit_total")?;
Ok(DatabasePostgresObservabilitySnapshot {
active_connections: row_u64(&activity, "active_connections")?,
idle_connections: row_u64(&activity, "idle_connections")?,
idle_in_transaction_connections: row_u64(&activity, "idle_in_transaction_connections")?,
waiting_connections: row_u64(&activity, "waiting_connections")?,
lock_waiting_connections: row_u64(&activity, "lock_waiting_connections")?,
oldest_active_query_age_ms: row_u64(&activity, "oldest_active_query_age_ms")?,
oldest_transaction_age_ms: row_u64(&activity, "oldest_transaction_age_ms")?,
deadlocks_total: row_u64(&database, "deadlocks_total")?,
block_read_total,
block_hit_total,
block_cache_hit_rate_basis_points: ratio_to_basis_points(
block_hit_total,
block_read_total.saturating_add(block_hit_total),
),
temp_files_total: row_u64(&database, "temp_files_total")?,
temp_bytes_total: row_u64(&database, "temp_bytes_total")?,
xact_commit_total: row_u64(&database, "xact_commit_total")?,
xact_rollback_total: row_u64(&database, "xact_rollback_total")?,
wal_observability_available: wal.available,
wal_observability_unavailable: wal.unavailable,
wal_records_total: wal.records_total,
wal_fpi_total: wal.fpi_total,
wal_bytes_total: wal.bytes_total,
wal_buffers_full_total: wal.buffers_full_total,
wal_write_total: wal.write_total,
wal_sync_total: wal.sync_total,
wal_write_time_ms_total: wal.write_time_ms_total,
wal_sync_time_ms_total: wal.sync_time_ms_total,
checkpoint_observability_available: checkpoint.available,
checkpoint_observability_unavailable: checkpoint.unavailable,
checkpoints_timed_total: checkpoint.timed_total,
checkpoints_requested_total: checkpoint.requested_total,
checkpoint_write_time_ms_total: checkpoint.write_time_ms_total,
checkpoint_sync_time_ms_total: checkpoint.sync_time_ms_total,
buffers_checkpoint_total: checkpoint.buffers_checkpoint_total,
buffers_backend_total: checkpoint.buffers_backend_total,
statement_observability_available: statements.available,
statement_observability_unavailable: statements.unavailable,
statement_top_calls_total: statements.top_calls_total,
statement_top_exec_time_ms_total: statements.top_exec_time_ms_total,
statement_top_max_mean_exec_time_ms: statements.top_max_mean_exec_time_ms,
statement_top_max_exec_time_ms: statements.top_max_exec_time_ms,
statement_top_shared_blks_read_total: statements.top_shared_blks_read_total,
statement_top_shared_blks_hit_total: statements.top_shared_blks_hit_total,
statement_top_temp_blks_total: statements.top_temp_blks_total,
})
}
pub async fn postgres_activity_groups(
&self,
limit: i64,
) -> Result<Vec<DatabasePostgresActivityGroup>, DataLayerError> {
const ACTIVITY_GROUP_SQL: &str = r#"
WITH normalized_activity AS (
SELECT
COALESCE(NULLIF(state, ''), 'unknown') AS state,
COALESCE(NULLIF(wait_event_type, ''), 'none') AS wait_event_type,
COALESCE(NULLIF(wait_event, ''), 'none') AS wait_event,
LEFT(
regexp_replace(
regexp_replace(
COALESCE(NULLIF(query, ''), '<empty>'),
'\s+',
' ',
'g'
),
'([0-9a-fA-F]{8,}|[0-9]+)',
'?',
'g'
),
160
) AS query_prefix,
COALESCE(EXTRACT(EPOCH FROM now() - query_start) * 1000, 0)::BIGINT AS query_age_ms,
COALESCE(EXTRACT(EPOCH FROM now() - xact_start) * 1000, 0)::BIGINT AS transaction_age_ms
FROM pg_stat_activity
WHERE datname = current_database()
AND pid <> pg_backend_pid()
)
SELECT
state,
wait_event_type,
wait_event,
query_prefix,
COUNT(*)::BIGINT AS connections,
COALESCE(MAX(query_age_ms), 0)::BIGINT AS max_query_age_ms,
COALESCE(MAX(transaction_age_ms), 0)::BIGINT AS max_transaction_age_ms
FROM normalized_activity
GROUP BY state, wait_event_type, wait_event, query_prefix
ORDER BY connections DESC, max_transaction_age_ms DESC, max_query_age_ms DESC
LIMIT $1
"#;
let rows = sqlx::query(ACTIVITY_GROUP_SQL)
.bind(limit.clamp(1, 20))
.fetch_all(self.pool())
.await
.map_postgres_err()?;
rows.into_iter()
.map(|row| {
Ok(DatabasePostgresActivityGroup {
state: row.try_get::<String, _>("state").map_postgres_err()?,
wait_event_type: row
.try_get::<String, _>("wait_event_type")
.map_postgres_err()?,
wait_event: row.try_get::<String, _>("wait_event").map_postgres_err()?,
query_prefix: row
.try_get::<String, _>("query_prefix")
.map_postgres_err()?,
connections: row_u64(&row, "connections")?,
max_query_age_ms: row_u64(&row, "max_query_age_ms")?,
max_transaction_age_ms: row_u64(&row, "max_transaction_age_ms")?,
})
})
.collect()
}
async fn postgres_wal_observability_snapshot(&self) -> PostgresWalObservabilitySnapshot {
if !self
.postgres_catalog_relation_has_columns(
"pg_catalog.pg_stat_wal",
&["wal_records", "wal_fpi", "wal_bytes", "wal_buffers_full"],
)
.await
{
return PostgresWalObservabilitySnapshot::default();
}
const WAL_SQL: &str = r#"
SELECT
COALESCE(SUM(wal_records), 0)::BIGINT AS records_total,
COALESCE(SUM(wal_fpi), 0)::BIGINT AS fpi_total,
COALESCE(SUM(wal_bytes), 0)::BIGINT AS bytes_total,
COALESCE(SUM(wal_buffers_full), 0)::BIGINT AS buffers_full_total
FROM pg_stat_wal
"#;
match sqlx::query(WAL_SQL).fetch_one(self.pool()).await {
Ok(row) => {
let io = self.postgres_wal_io_observability_snapshot().await;
PostgresWalObservabilitySnapshot {
available: 1,
records_total: row_u64(&row, "records_total").unwrap_or_default(),
fpi_total: row_u64(&row, "fpi_total").unwrap_or_default(),
bytes_total: row_u64(&row, "bytes_total").unwrap_or_default(),
buffers_full_total: row_u64(&row, "buffers_full_total").unwrap_or_default(),
write_total: io.write_total,
sync_total: io.sync_total,
write_time_ms_total: io.write_time_ms_total,
sync_time_ms_total: io.sync_time_ms_total,
..PostgresWalObservabilitySnapshot::default()
}
}
Err(_) => PostgresWalObservabilitySnapshot {
unavailable: 1,
..PostgresWalObservabilitySnapshot::default()
},
}
}
async fn postgres_wal_io_observability_snapshot(&self) -> PostgresWalIoObservabilitySnapshot {
if self
.postgres_catalog_relation_has_columns(
"pg_catalog.pg_stat_wal",
&["wal_write", "wal_sync", "wal_write_time", "wal_sync_time"],
)
.await
{
return self.postgres_wal_legacy_io_observability_snapshot().await;
}
if !self
.postgres_catalog_relation_has_columns(
"pg_catalog.pg_stat_io",
&["object", "writes", "fsyncs", "write_time", "fsync_time"],
)
.await
{
return PostgresWalIoObservabilitySnapshot::default();
}
const WAL_IO_SQL: &str = r#"
SELECT
COALESCE(SUM(writes), 0)::BIGINT AS write_total,
COALESCE(SUM(fsyncs), 0)::BIGINT AS sync_total,
COALESCE(SUM(write_time), 0)::BIGINT AS write_time_ms_total,
COALESCE(SUM(fsync_time), 0)::BIGINT AS sync_time_ms_total
FROM pg_stat_io
WHERE object = 'wal'
"#;
match sqlx::query(WAL_IO_SQL).fetch_one(self.pool()).await {
Ok(row) => PostgresWalIoObservabilitySnapshot {
write_total: row_u64(&row, "write_total").unwrap_or_default(),
sync_total: row_u64(&row, "sync_total").unwrap_or_default(),
write_time_ms_total: row_u64(&row, "write_time_ms_total").unwrap_or_default(),
sync_time_ms_total: row_u64(&row, "sync_time_ms_total").unwrap_or_default(),
},
Err(_) => PostgresWalIoObservabilitySnapshot::default(),
}
}
async fn postgres_wal_legacy_io_observability_snapshot(
&self,
) -> PostgresWalIoObservabilitySnapshot {
const WAL_IO_SQL: &str = r#"
SELECT
COALESCE(SUM(wal_write), 0)::BIGINT AS write_total,
COALESCE(SUM(wal_sync), 0)::BIGINT AS sync_total,
COALESCE(SUM(wal_write_time), 0)::BIGINT AS write_time_ms_total,
COALESCE(SUM(wal_sync_time), 0)::BIGINT AS sync_time_ms_total
FROM pg_stat_wal
"#;
match sqlx::query(WAL_IO_SQL).fetch_one(self.pool()).await {
Ok(row) => PostgresWalIoObservabilitySnapshot {
write_total: row_u64(&row, "write_total").unwrap_or_default(),
sync_total: row_u64(&row, "sync_total").unwrap_or_default(),
write_time_ms_total: row_u64(&row, "write_time_ms_total").unwrap_or_default(),
sync_time_ms_total: row_u64(&row, "sync_time_ms_total").unwrap_or_default(),
},
Err(_) => PostgresWalIoObservabilitySnapshot::default(),
}
}
async fn postgres_checkpoint_observability_snapshot(
&self,
) -> PostgresCheckpointObservabilitySnapshot {
if self
.postgres_catalog_relation_has_columns(
"pg_catalog.pg_stat_checkpointer",
&[
"num_timed",
"num_requested",
"write_time",
"sync_time",
"buffers_written",
],
)
.await
{
return self
.postgres_checkpoint_observability_snapshot_from_checkpointer()
.await;
}
if !self
.postgres_catalog_relation_has_columns(
"pg_catalog.pg_stat_bgwriter",
&[
"checkpoints_timed",
"checkpoints_req",
"checkpoint_write_time",
"checkpoint_sync_time",
"buffers_checkpoint",
"buffers_backend",
],
)
.await
{
return PostgresCheckpointObservabilitySnapshot::default();
}
const CHECKPOINT_SQL: &str = r#"
SELECT
COALESCE(SUM(checkpoints_timed), 0)::BIGINT AS timed_total,
COALESCE(SUM(checkpoints_req), 0)::BIGINT AS requested_total,
COALESCE(SUM(checkpoint_write_time), 0)::BIGINT AS write_time_ms_total,
COALESCE(SUM(checkpoint_sync_time), 0)::BIGINT AS sync_time_ms_total,
COALESCE(SUM(buffers_checkpoint), 0)::BIGINT AS buffers_checkpoint_total,
COALESCE(SUM(buffers_backend), 0)::BIGINT AS buffers_backend_total
FROM pg_stat_bgwriter
"#;
match sqlx::query(CHECKPOINT_SQL).fetch_one(self.pool()).await {
Ok(row) => PostgresCheckpointObservabilitySnapshot {
available: 1,
timed_total: row_u64(&row, "timed_total").unwrap_or_default(),
requested_total: row_u64(&row, "requested_total").unwrap_or_default(),
write_time_ms_total: row_u64(&row, "write_time_ms_total").unwrap_or_default(),
sync_time_ms_total: row_u64(&row, "sync_time_ms_total").unwrap_or_default(),
buffers_checkpoint_total: row_u64(&row, "buffers_checkpoint_total")
.unwrap_or_default(),
buffers_backend_total: row_u64(&row, "buffers_backend_total").unwrap_or_default(),
..PostgresCheckpointObservabilitySnapshot::default()
},
Err(_) => PostgresCheckpointObservabilitySnapshot {
unavailable: 1,
..PostgresCheckpointObservabilitySnapshot::default()
},
}
}
async fn postgres_checkpoint_observability_snapshot_from_checkpointer(
&self,
) -> PostgresCheckpointObservabilitySnapshot {
const CHECKPOINT_SQL: &str = r#"
SELECT
COALESCE(SUM(num_timed), 0)::BIGINT AS timed_total,
COALESCE(SUM(num_requested), 0)::BIGINT AS requested_total,
COALESCE(SUM(write_time), 0)::BIGINT AS write_time_ms_total,
COALESCE(SUM(sync_time), 0)::BIGINT AS sync_time_ms_total,
COALESCE(SUM(buffers_written), 0)::BIGINT AS buffers_checkpoint_total
FROM pg_stat_checkpointer
"#;
match sqlx::query(CHECKPOINT_SQL).fetch_one(self.pool()).await {
Ok(row) => PostgresCheckpointObservabilitySnapshot {
available: 1,
timed_total: row_u64(&row, "timed_total").unwrap_or_default(),
requested_total: row_u64(&row, "requested_total").unwrap_or_default(),
write_time_ms_total: row_u64(&row, "write_time_ms_total").unwrap_or_default(),
sync_time_ms_total: row_u64(&row, "sync_time_ms_total").unwrap_or_default(),
buffers_checkpoint_total: row_u64(&row, "buffers_checkpoint_total")
.unwrap_or_default(),
..PostgresCheckpointObservabilitySnapshot::default()
},
Err(_) => PostgresCheckpointObservabilitySnapshot {
unavailable: 1,
..PostgresCheckpointObservabilitySnapshot::default()
},
}
}
async fn postgres_statement_observability_snapshot(
&self,
) -> PostgresStatementObservabilitySnapshot {
let extension_installed = sqlx::query_scalar::<_, bool>(
"SELECT EXISTS (SELECT 1 FROM pg_extension WHERE extname = 'pg_stat_statements')",
)
.fetch_one(self.pool())
.await
.unwrap_or(false);
if !extension_installed {
return PostgresStatementObservabilitySnapshot::default();
}
if self
.postgres_catalog_relation_has_columns(
"pg_stat_statements",
&[
"calls",
"total_exec_time",
"mean_exec_time",
"max_exec_time",
"shared_blks_read",
"shared_blks_hit",
"temp_blks_read",
"temp_blks_written",
"dbid",
],
)
.await
{
return self
.postgres_statement_observability_snapshot_with_exec_time()
.await;
}
if self
.postgres_catalog_relation_has_columns(
"pg_stat_statements",
&[
"calls",
"total_time",
"mean_time",
"max_time",
"shared_blks_read",
"shared_blks_hit",
"temp_blks_read",
"temp_blks_written",
"dbid",
],
)
.await
{
return self
.postgres_statement_observability_snapshot_with_total_time()
.await;
}
PostgresStatementObservabilitySnapshot::default()
}
async fn postgres_statement_observability_snapshot_with_exec_time(
&self,
) -> PostgresStatementObservabilitySnapshot {
const STATEMENTS_SQL: &str = r#"
SELECT
COALESCE(SUM(calls), 0)::BIGINT AS top_calls_total,
COALESCE(SUM(total_exec_time), 0)::BIGINT AS top_exec_time_ms_total,
COALESCE(MAX(mean_exec_time), 0)::BIGINT AS top_max_mean_exec_time_ms,
COALESCE(MAX(max_exec_time), 0)::BIGINT AS top_max_exec_time_ms,
COALESCE(SUM(shared_blks_read), 0)::BIGINT AS top_shared_blks_read_total,
COALESCE(SUM(shared_blks_hit), 0)::BIGINT AS top_shared_blks_hit_total,
COALESCE(SUM(temp_blks_read + temp_blks_written), 0)::BIGINT AS top_temp_blks_total
FROM (
SELECT
calls,
total_exec_time,
mean_exec_time,
max_exec_time,
shared_blks_read,
shared_blks_hit,
temp_blks_read,
temp_blks_written
FROM pg_stat_statements
WHERE dbid = (SELECT oid FROM pg_database WHERE datname = current_database())
ORDER BY total_exec_time DESC
LIMIT 20
) top_statements
"#;
match sqlx::query(STATEMENTS_SQL).fetch_one(self.pool()).await {
Ok(row) => PostgresStatementObservabilitySnapshot {
available: 1,
top_calls_total: row_u64(&row, "top_calls_total").unwrap_or_default(),
top_exec_time_ms_total: row_u64(&row, "top_exec_time_ms_total").unwrap_or_default(),
top_max_mean_exec_time_ms: row_u64(&row, "top_max_mean_exec_time_ms")
.unwrap_or_default(),
top_max_exec_time_ms: row_u64(&row, "top_max_exec_time_ms").unwrap_or_default(),
top_shared_blks_read_total: row_u64(&row, "top_shared_blks_read_total")
.unwrap_or_default(),
top_shared_blks_hit_total: row_u64(&row, "top_shared_blks_hit_total")
.unwrap_or_default(),
top_temp_blks_total: row_u64(&row, "top_temp_blks_total").unwrap_or_default(),
..PostgresStatementObservabilitySnapshot::default()
},
Err(_) => PostgresStatementObservabilitySnapshot {
unavailable: 1,
..PostgresStatementObservabilitySnapshot::default()
},
}
}
async fn postgres_statement_observability_snapshot_with_total_time(
&self,
) -> PostgresStatementObservabilitySnapshot {
const STATEMENTS_SQL: &str = r#"
SELECT
COALESCE(SUM(calls), 0)::BIGINT AS top_calls_total,
COALESCE(SUM(total_time), 0)::BIGINT AS top_exec_time_ms_total,
COALESCE(MAX(mean_time), 0)::BIGINT AS top_max_mean_exec_time_ms,
COALESCE(MAX(max_time), 0)::BIGINT AS top_max_exec_time_ms,
COALESCE(SUM(shared_blks_read), 0)::BIGINT AS top_shared_blks_read_total,
COALESCE(SUM(shared_blks_hit), 0)::BIGINT AS top_shared_blks_hit_total,
COALESCE(SUM(temp_blks_read + temp_blks_written), 0)::BIGINT AS top_temp_blks_total
FROM (
SELECT
calls,
total_time,
mean_time,
max_time,
shared_blks_read,
shared_blks_hit,
temp_blks_read,
temp_blks_written
FROM pg_stat_statements
WHERE dbid = (SELECT oid FROM pg_database WHERE datname = current_database())
ORDER BY total_time DESC
LIMIT 20
) top_statements
"#;
match sqlx::query(STATEMENTS_SQL).fetch_one(self.pool()).await {
Ok(row) => PostgresStatementObservabilitySnapshot {
available: 1,
top_calls_total: row_u64(&row, "top_calls_total").unwrap_or_default(),
top_exec_time_ms_total: row_u64(&row, "top_exec_time_ms_total").unwrap_or_default(),
top_max_mean_exec_time_ms: row_u64(&row, "top_max_mean_exec_time_ms")
.unwrap_or_default(),
top_max_exec_time_ms: row_u64(&row, "top_max_exec_time_ms").unwrap_or_default(),
top_shared_blks_read_total: row_u64(&row, "top_shared_blks_read_total")
.unwrap_or_default(),
top_shared_blks_hit_total: row_u64(&row, "top_shared_blks_hit_total")
.unwrap_or_default(),
top_temp_blks_total: row_u64(&row, "top_temp_blks_total").unwrap_or_default(),
..PostgresStatementObservabilitySnapshot::default()
},
Err(_) => PostgresStatementObservabilitySnapshot {
unavailable: 1,
..PostgresStatementObservabilitySnapshot::default()
},
}
}
async fn postgres_catalog_relation_has_columns(
&self,
relation: &str,
columns: &[&str],
) -> bool {
if !self.postgres_catalog_relation_exists(relation).await {
return false;
}
for column in columns {
if !self.postgres_catalog_column_exists(relation, column).await {
return false;
}
}
true
}
async fn postgres_catalog_column_exists(&self, relation: &str, column: &str) -> bool {
sqlx::query_scalar::<_, bool>(
r#"
SELECT EXISTS (
SELECT 1
FROM pg_attribute
WHERE attrelid = to_regclass($1)
AND attname = $2
AND NOT attisdropped
)
"#,
)
.bind(relation)
.bind(column)
.fetch_one(self.pool())
.await
.unwrap_or(false)
}
async fn postgres_catalog_relation_exists(&self, relation: &str) -> bool {
sqlx::query_scalar::<_, Option<String>>("SELECT to_regclass($1)::TEXT")
.bind(relation)
.fetch_one(self.pool())
.await
.ok()
.flatten()
.is_some()
}
}
fn row_u64(row: &sqlx::postgres::PgRow, name: &str) -> Result<u64, DataLayerError> {
row.try_get::<i64, _>(name)
.map(u64_from_i64)
.map_postgres_err()
}
fn u64_from_i64(value: i64) -> u64 {
u64::try_from(value).unwrap_or_default()
}
fn ratio_to_basis_points(value: u64, total: u64) -> u64 {
value.saturating_mul(10_000).checked_div(total).unwrap_or(0)
}
#[derive(Debug, Clone, Copy, Default)]
struct PostgresWalObservabilitySnapshot {
available: u64,
unavailable: u64,
records_total: u64,
fpi_total: u64,
bytes_total: u64,
buffers_full_total: u64,
write_total: u64,
sync_total: u64,
write_time_ms_total: u64,
sync_time_ms_total: u64,
}
#[derive(Debug, Clone, Copy, Default)]
struct PostgresWalIoObservabilitySnapshot {
write_total: u64,
sync_total: u64,
write_time_ms_total: u64,
sync_time_ms_total: u64,
}
#[derive(Debug, Clone, Copy, Default)]
struct PostgresCheckpointObservabilitySnapshot {
available: u64,
unavailable: u64,
timed_total: u64,
requested_total: u64,
write_time_ms_total: u64,
sync_time_ms_total: u64,
buffers_checkpoint_total: u64,
buffers_backend_total: u64,
}
#[derive(Debug, Clone, Copy, Default)]
struct PostgresStatementObservabilitySnapshot {
available: u64,
unavailable: u64,
top_calls_total: u64,
top_exec_time_ms_total: u64,
top_max_mean_exec_time_ms: u64,
top_max_exec_time_ms: u64,
top_shared_blks_read_total: u64,
top_shared_blks_hit_total: u64,
top_temp_blks_total: u64,
}
impl MysqlBackend {
+2 -1
View File
@@ -24,7 +24,8 @@ pub use config::DataLayerConfig;
pub use database::{DatabaseDriver, SqlDatabaseConfig, SqlPoolConfig, DEFAULT_SQLITE_DATABASE_URL};
pub use error::DataLayerError;
pub use maintenance::{
DatabaseMaintenanceSummary, DatabasePoolSummary, StatsDailyAggregationInput,
DatabaseMaintenanceSummary, DatabasePoolSummary, DatabasePostgresActivityGroup,
DatabasePostgresObservabilitySnapshot, StatsDailyAggregationInput,
StatsDailyAggregationSummary, StatsHourlyAggregationInput, StatsHourlyAggregationSummary,
WalletDailyUsageAggregationInput, WalletDailyUsageAggregationResult,
};
+57
View File
@@ -21,6 +21,63 @@ pub struct DatabasePoolSummary {
pub usage_rate: f64,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub struct DatabasePostgresObservabilitySnapshot {
pub active_connections: u64,
pub idle_connections: u64,
pub idle_in_transaction_connections: u64,
pub waiting_connections: u64,
pub lock_waiting_connections: u64,
pub oldest_active_query_age_ms: u64,
pub oldest_transaction_age_ms: u64,
pub deadlocks_total: u64,
pub block_read_total: u64,
pub block_hit_total: u64,
pub block_cache_hit_rate_basis_points: u64,
pub temp_files_total: u64,
pub temp_bytes_total: u64,
pub xact_commit_total: u64,
pub xact_rollback_total: u64,
pub wal_observability_available: u64,
pub wal_observability_unavailable: u64,
pub wal_records_total: u64,
pub wal_fpi_total: u64,
pub wal_bytes_total: u64,
pub wal_buffers_full_total: u64,
pub wal_write_total: u64,
pub wal_sync_total: u64,
pub wal_write_time_ms_total: u64,
pub wal_sync_time_ms_total: u64,
pub checkpoint_observability_available: u64,
pub checkpoint_observability_unavailable: u64,
pub checkpoints_timed_total: u64,
pub checkpoints_requested_total: u64,
pub checkpoint_write_time_ms_total: u64,
pub checkpoint_sync_time_ms_total: u64,
pub buffers_checkpoint_total: u64,
pub buffers_backend_total: u64,
pub statement_observability_available: u64,
pub statement_observability_unavailable: u64,
pub statement_top_calls_total: u64,
pub statement_top_exec_time_ms_total: u64,
pub statement_top_max_mean_exec_time_ms: u64,
pub statement_top_max_exec_time_ms: u64,
pub statement_top_shared_blks_read_total: u64,
pub statement_top_shared_blks_hit_total: u64,
pub statement_top_temp_blks_total: u64,
}
#[derive(Debug, Clone, PartialEq, Eq, Default)]
pub struct DatabasePostgresActivityGroup {
pub state: String,
pub wait_event_type: String,
pub wait_event: String,
pub query_prefix: String,
pub connections: u64,
pub max_query_age_ms: u64,
pub max_transaction_age_ms: u64,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct WalletDailyUsageAggregationInput {
pub billing_date: String,
@@ -374,6 +374,17 @@ fn merge_candidate(
candidate: UpsertRequestCandidateRecord,
existing: Option<StoredRequestCandidate>,
) -> Result<StoredRequestCandidate, DataLayerError> {
let preserve_existing_lifecycle = existing.as_ref().is_some_and(|value| {
request_candidate_lifecycle_would_regress(value.status, candidate.status)
});
let merged_status = if preserve_existing_lifecycle {
existing
.as_ref()
.map(|value| value.status)
.unwrap_or(candidate.status)
} else {
candidate.status
};
let created_at_unix_ms = candidate
.created_at_unix_ms
.filter(|value| *value > 1000)
@@ -426,7 +437,7 @@ fn merge_candidate(
candidate
.key_id
.or_else(|| existing.as_ref().and_then(|value| value.key_id.clone())),
candidate.status,
merged_status,
candidate.skip_reason.or_else(|| {
existing
.as_ref()
@@ -435,25 +446,48 @@ fn merge_candidate(
candidate
.is_cached
.unwrap_or_else(|| existing.as_ref().is_some_and(|value| value.is_cached)),
candidate.status_code.map(i32::from).or_else(|| {
if preserve_existing_lifecycle {
existing
.as_ref()
.and_then(|value| value.status_code.map(i32::from))
}),
candidate
.error_type
.or_else(|| existing.as_ref().and_then(|value| value.error_type.clone())),
candidate.error_message.or_else(|| {
} else {
candidate.status_code.map(i32::from).or_else(|| {
existing
.as_ref()
.and_then(|value| value.status_code.map(i32::from))
})
},
if preserve_existing_lifecycle {
existing.as_ref().and_then(|value| value.error_type.clone())
} else {
candidate
.error_type
.or_else(|| existing.as_ref().and_then(|value| value.error_type.clone()))
},
if preserve_existing_lifecycle {
existing
.as_ref()
.and_then(|value| value.error_message.clone())
}),
candidate.latency_ms.map(to_i32_u64).transpose()?.or(
} else {
candidate.error_message.or_else(|| {
existing
.as_ref()
.and_then(|value| value.error_message.clone())
})
},
if preserve_existing_lifecycle {
match existing.as_ref().and_then(|value| value.latency_ms) {
Some(value) => Some(to_i32_u64(value)?),
None => None,
},
),
}
} else {
candidate.latency_ms.map(to_i32_u64).transpose()?.or(
match existing.as_ref().and_then(|value| value.latency_ms) {
Some(value) => Some(to_i32_u64(value)?),
None => None,
},
)
},
candidate.concurrent_requests.map(to_i32).transpose()?.or(
match existing
.as_ref()
@@ -475,18 +509,42 @@ fn merge_candidate(
.or_else(|| existing.as_ref().and_then(|value| value.started_at_unix_ms))
.map(|value| u64_to_i64(value, "request candidate started_at"))
.transpose()?,
candidate
.finished_at_unix_ms
.or_else(|| {
if preserve_existing_lifecycle {
existing
.as_ref()
.and_then(|value| value.finished_at_unix_ms)
} else {
candidate.finished_at_unix_ms.or_else(|| {
existing
.as_ref()
.and_then(|value| value.finished_at_unix_ms)
})
.map(|value| u64_to_i64(value, "request candidate finished_at"))
.transpose()?,
}
.map(|value| u64_to_i64(value, "request candidate finished_at"))
.transpose()?,
)
}
fn request_candidate_lifecycle_would_regress(
existing: RequestCandidateStatus,
incoming: RequestCandidateStatus,
) -> bool {
matches!(
existing,
RequestCandidateStatus::Success
| RequestCandidateStatus::Failed
| RequestCandidateStatus::Cancelled
| RequestCandidateStatus::Skipped
) && matches!(
incoming,
RequestCandidateStatus::Available
| RequestCandidateStatus::Unused
| RequestCandidateStatus::Pending
| RequestCandidateStatus::Streaming
) || existing == RequestCandidateStatus::Streaming
&& incoming == RequestCandidateStatus::Pending
}
fn aggregate_timeline(
candidates: Vec<StoredRequestCandidate>,
since_unix_secs: u64,
@@ -681,6 +739,9 @@ fn optional_u64_to_i64(value: Option<u64>, name: &str) -> Result<Option<i64>, Da
#[cfg(test)]
mod tests {
use super::MysqlRequestCandidateRepository;
use crate::repository::candidates::{
RequestCandidateStatus, StoredRequestCandidate, UpsertRequestCandidateRecord,
};
#[tokio::test]
async fn repository_builds_from_lazy_pool() {
@@ -692,4 +753,75 @@ mod tests {
let _repository = MysqlRequestCandidateRepository::new(pool);
}
#[test]
fn merge_candidate_keeps_terminal_status_when_streaming_arrives_late() {
let existing = StoredRequestCandidate::new(
"candidate-1".to_string(),
"request-1".to_string(),
Some("user-1".to_string()),
Some("key-1".to_string()),
None,
None,
0,
0,
Some("provider-1".to_string()),
Some("endpoint-1".to_string()),
Some("provider-key-1".to_string()),
RequestCandidateStatus::Success,
None,
false,
Some(200),
None,
None,
Some(123),
None,
Some(serde_json::json!({"terminal": true})),
None,
1_000,
Some(1_001),
Some(1_123),
)
.expect("existing candidate should build");
let merged = super::merge_candidate(
UpsertRequestCandidateRecord {
id: "candidate-late".to_string(),
request_id: "request-1".to_string(),
user_id: Some("user-1".to_string()),
api_key_id: Some("key-1".to_string()),
username: None,
api_key_name: None,
candidate_index: 0,
retry_index: 0,
provider_id: Some("provider-1".to_string()),
endpoint_id: Some("endpoint-1".to_string()),
key_id: Some("provider-key-1".to_string()),
status: RequestCandidateStatus::Streaming,
skip_reason: None,
is_cached: Some(false),
status_code: Some(200),
error_type: None,
error_message: None,
latency_ms: Some(9_999),
concurrent_requests: None,
extra_data: Some(serde_json::json!({"late": true})),
required_capabilities: None,
created_at_unix_ms: Some(1_050),
started_at_unix_ms: Some(1_051),
finished_at_unix_ms: None,
},
Some(existing),
)
.expect("candidate should merge");
assert_eq!(merged.id, "candidate-1");
assert_eq!(merged.status, RequestCandidateStatus::Success);
assert_eq!(merged.latency_ms, Some(123));
assert_eq!(merged.finished_at_unix_ms, Some(1_123));
assert_eq!(
merged.extra_data,
Some(serde_json::json!({"terminal": true, "late": true}))
);
}
}
@@ -132,13 +132,48 @@ DO UPDATE SET
provider_id = COALESCE(EXCLUDED.provider_id, request_candidates.provider_id),
endpoint_id = COALESCE(EXCLUDED.endpoint_id, request_candidates.endpoint_id),
key_id = COALESCE(EXCLUDED.key_id, request_candidates.key_id),
status = EXCLUDED.status,
status = CASE
WHEN request_candidates.status IN ('success', 'failed', 'cancelled', 'skipped')
AND EXCLUDED.status IN ('available', 'unused', 'pending', 'streaming')
THEN request_candidates.status
WHEN request_candidates.status = 'streaming' AND EXCLUDED.status = 'pending'
THEN request_candidates.status
ELSE EXCLUDED.status
END,
skip_reason = COALESCE(EXCLUDED.skip_reason, request_candidates.skip_reason),
is_cached = COALESCE($14, request_candidates.is_cached),
status_code = COALESCE(EXCLUDED.status_code, request_candidates.status_code),
error_type = COALESCE(EXCLUDED.error_type, request_candidates.error_type),
error_message = COALESCE(EXCLUDED.error_message, request_candidates.error_message),
latency_ms = COALESCE(EXCLUDED.latency_ms, request_candidates.latency_ms),
status_code = CASE
WHEN request_candidates.status IN ('success', 'failed', 'cancelled', 'skipped')
AND EXCLUDED.status IN ('available', 'unused', 'pending', 'streaming')
THEN request_candidates.status_code
WHEN request_candidates.status = 'streaming' AND EXCLUDED.status = 'pending'
THEN request_candidates.status_code
ELSE COALESCE(EXCLUDED.status_code, request_candidates.status_code)
END,
error_type = CASE
WHEN request_candidates.status IN ('success', 'failed', 'cancelled', 'skipped')
AND EXCLUDED.status IN ('available', 'unused', 'pending', 'streaming')
THEN request_candidates.error_type
WHEN request_candidates.status = 'streaming' AND EXCLUDED.status = 'pending'
THEN request_candidates.error_type
ELSE COALESCE(EXCLUDED.error_type, request_candidates.error_type)
END,
error_message = CASE
WHEN request_candidates.status IN ('success', 'failed', 'cancelled', 'skipped')
AND EXCLUDED.status IN ('available', 'unused', 'pending', 'streaming')
THEN request_candidates.error_message
WHEN request_candidates.status = 'streaming' AND EXCLUDED.status = 'pending'
THEN request_candidates.error_message
ELSE COALESCE(EXCLUDED.error_message, request_candidates.error_message)
END,
latency_ms = CASE
WHEN request_candidates.status IN ('success', 'failed', 'cancelled', 'skipped')
AND EXCLUDED.status IN ('available', 'unused', 'pending', 'streaming')
THEN request_candidates.latency_ms
WHEN request_candidates.status = 'streaming' AND EXCLUDED.status = 'pending'
THEN request_candidates.latency_ms
ELSE COALESCE(EXCLUDED.latency_ms, request_candidates.latency_ms)
END,
concurrent_requests = COALESCE(EXCLUDED.concurrent_requests, request_candidates.concurrent_requests),
extra_data = CASE
WHEN request_candidates.extra_data IS NULL THEN EXCLUDED.extra_data
@@ -155,7 +190,14 @@ DO UPDATE SET
ELSE request_candidates.created_at
END,
started_at = COALESCE(EXCLUDED.started_at, request_candidates.started_at),
finished_at = COALESCE(EXCLUDED.finished_at, request_candidates.finished_at)
finished_at = CASE
WHEN request_candidates.status IN ('success', 'failed', 'cancelled', 'skipped')
AND EXCLUDED.status IN ('available', 'unused', 'pending', 'streaming')
THEN request_candidates.finished_at
WHEN request_candidates.status = 'streaming' AND EXCLUDED.status = 'pending'
THEN request_candidates.finished_at
ELSE COALESCE(EXCLUDED.finished_at, request_candidates.finished_at)
END
RETURNING
id,
request_id,
@@ -193,13 +235,48 @@ DO UPDATE SET
provider_id = COALESCE(EXCLUDED.provider_id, request_candidates.provider_id),
endpoint_id = COALESCE(EXCLUDED.endpoint_id, request_candidates.endpoint_id),
key_id = COALESCE(EXCLUDED.key_id, request_candidates.key_id),
status = EXCLUDED.status,
status = CASE
WHEN request_candidates.status IN ('success', 'failed', 'cancelled', 'skipped')
AND EXCLUDED.status IN ('available', 'unused', 'pending', 'streaming')
THEN request_candidates.status
WHEN request_candidates.status = 'streaming' AND EXCLUDED.status = 'pending'
THEN request_candidates.status
ELSE EXCLUDED.status
END,
skip_reason = COALESCE(EXCLUDED.skip_reason, request_candidates.skip_reason),
is_cached = COALESCE(EXCLUDED.is_cached, request_candidates.is_cached),
status_code = COALESCE(EXCLUDED.status_code, request_candidates.status_code),
error_type = COALESCE(EXCLUDED.error_type, request_candidates.error_type),
error_message = COALESCE(EXCLUDED.error_message, request_candidates.error_message),
latency_ms = COALESCE(EXCLUDED.latency_ms, request_candidates.latency_ms),
status_code = CASE
WHEN request_candidates.status IN ('success', 'failed', 'cancelled', 'skipped')
AND EXCLUDED.status IN ('available', 'unused', 'pending', 'streaming')
THEN request_candidates.status_code
WHEN request_candidates.status = 'streaming' AND EXCLUDED.status = 'pending'
THEN request_candidates.status_code
ELSE COALESCE(EXCLUDED.status_code, request_candidates.status_code)
END,
error_type = CASE
WHEN request_candidates.status IN ('success', 'failed', 'cancelled', 'skipped')
AND EXCLUDED.status IN ('available', 'unused', 'pending', 'streaming')
THEN request_candidates.error_type
WHEN request_candidates.status = 'streaming' AND EXCLUDED.status = 'pending'
THEN request_candidates.error_type
ELSE COALESCE(EXCLUDED.error_type, request_candidates.error_type)
END,
error_message = CASE
WHEN request_candidates.status IN ('success', 'failed', 'cancelled', 'skipped')
AND EXCLUDED.status IN ('available', 'unused', 'pending', 'streaming')
THEN request_candidates.error_message
WHEN request_candidates.status = 'streaming' AND EXCLUDED.status = 'pending'
THEN request_candidates.error_message
ELSE COALESCE(EXCLUDED.error_message, request_candidates.error_message)
END,
latency_ms = CASE
WHEN request_candidates.status IN ('success', 'failed', 'cancelled', 'skipped')
AND EXCLUDED.status IN ('available', 'unused', 'pending', 'streaming')
THEN request_candidates.latency_ms
WHEN request_candidates.status = 'streaming' AND EXCLUDED.status = 'pending'
THEN request_candidates.latency_ms
ELSE COALESCE(EXCLUDED.latency_ms, request_candidates.latency_ms)
END,
concurrent_requests = COALESCE(EXCLUDED.concurrent_requests, request_candidates.concurrent_requests),
extra_data = CASE
WHEN request_candidates.extra_data IS NULL THEN EXCLUDED.extra_data
@@ -216,7 +293,14 @@ DO UPDATE SET
ELSE request_candidates.created_at
END,
started_at = COALESCE(EXCLUDED.started_at, request_candidates.started_at),
finished_at = COALESCE(EXCLUDED.finished_at, request_candidates.finished_at)
finished_at = CASE
WHEN request_candidates.status IN ('success', 'failed', 'cancelled', 'skipped')
AND EXCLUDED.status IN ('available', 'unused', 'pending', 'streaming')
THEN request_candidates.finished_at
WHEN request_candidates.status = 'streaming' AND EXCLUDED.status = 'pending'
THEN request_candidates.finished_at
ELSE COALESCE(EXCLUDED.finished_at, request_candidates.finished_at)
END
"#;
const UPSERT_CONFLICT_INHERIT_IS_CACHED_SQL: &str = r#"
@@ -229,13 +313,48 @@ DO UPDATE SET
provider_id = COALESCE(EXCLUDED.provider_id, request_candidates.provider_id),
endpoint_id = COALESCE(EXCLUDED.endpoint_id, request_candidates.endpoint_id),
key_id = COALESCE(EXCLUDED.key_id, request_candidates.key_id),
status = EXCLUDED.status,
status = CASE
WHEN request_candidates.status IN ('success', 'failed', 'cancelled', 'skipped')
AND EXCLUDED.status IN ('available', 'unused', 'pending', 'streaming')
THEN request_candidates.status
WHEN request_candidates.status = 'streaming' AND EXCLUDED.status = 'pending'
THEN request_candidates.status
ELSE EXCLUDED.status
END,
skip_reason = COALESCE(EXCLUDED.skip_reason, request_candidates.skip_reason),
is_cached = request_candidates.is_cached,
status_code = COALESCE(EXCLUDED.status_code, request_candidates.status_code),
error_type = COALESCE(EXCLUDED.error_type, request_candidates.error_type),
error_message = COALESCE(EXCLUDED.error_message, request_candidates.error_message),
latency_ms = COALESCE(EXCLUDED.latency_ms, request_candidates.latency_ms),
status_code = CASE
WHEN request_candidates.status IN ('success', 'failed', 'cancelled', 'skipped')
AND EXCLUDED.status IN ('available', 'unused', 'pending', 'streaming')
THEN request_candidates.status_code
WHEN request_candidates.status = 'streaming' AND EXCLUDED.status = 'pending'
THEN request_candidates.status_code
ELSE COALESCE(EXCLUDED.status_code, request_candidates.status_code)
END,
error_type = CASE
WHEN request_candidates.status IN ('success', 'failed', 'cancelled', 'skipped')
AND EXCLUDED.status IN ('available', 'unused', 'pending', 'streaming')
THEN request_candidates.error_type
WHEN request_candidates.status = 'streaming' AND EXCLUDED.status = 'pending'
THEN request_candidates.error_type
ELSE COALESCE(EXCLUDED.error_type, request_candidates.error_type)
END,
error_message = CASE
WHEN request_candidates.status IN ('success', 'failed', 'cancelled', 'skipped')
AND EXCLUDED.status IN ('available', 'unused', 'pending', 'streaming')
THEN request_candidates.error_message
WHEN request_candidates.status = 'streaming' AND EXCLUDED.status = 'pending'
THEN request_candidates.error_message
ELSE COALESCE(EXCLUDED.error_message, request_candidates.error_message)
END,
latency_ms = CASE
WHEN request_candidates.status IN ('success', 'failed', 'cancelled', 'skipped')
AND EXCLUDED.status IN ('available', 'unused', 'pending', 'streaming')
THEN request_candidates.latency_ms
WHEN request_candidates.status = 'streaming' AND EXCLUDED.status = 'pending'
THEN request_candidates.latency_ms
ELSE COALESCE(EXCLUDED.latency_ms, request_candidates.latency_ms)
END,
concurrent_requests = COALESCE(EXCLUDED.concurrent_requests, request_candidates.concurrent_requests),
extra_data = CASE
WHEN request_candidates.extra_data IS NULL THEN EXCLUDED.extra_data
@@ -252,7 +371,14 @@ DO UPDATE SET
ELSE request_candidates.created_at
END,
started_at = COALESCE(EXCLUDED.started_at, request_candidates.started_at),
finished_at = COALESCE(EXCLUDED.finished_at, request_candidates.finished_at)
finished_at = CASE
WHEN request_candidates.status IN ('success', 'failed', 'cancelled', 'skipped')
AND EXCLUDED.status IN ('available', 'unused', 'pending', 'streaming')
THEN request_candidates.finished_at
WHEN request_candidates.status = 'streaming' AND EXCLUDED.status = 'pending'
THEN request_candidates.finished_at
ELSE COALESCE(EXCLUDED.finished_at, request_candidates.finished_at)
END
"#;
const UPSERT_MANY_PREFIX_SQL: &str = r#"
@@ -1015,7 +1141,10 @@ fn to_i32_u64(value: u64) -> Result<i32, DataLayerError> {
#[cfg(test)]
mod tests {
use super::{SqlxRequestCandidateReadRepository, UPSERT_SQL};
use super::{
SqlxRequestCandidateReadRepository, UPSERT_CONFLICT_INHERIT_IS_CACHED_SQL,
UPSERT_CONFLICT_SQL, UPSERT_SQL,
};
use crate::driver::postgres::{PostgresPoolConfig, PostgresPoolFactory};
#[test]
@@ -1030,6 +1159,27 @@ mod tests {
assert!(UPSERT_SQL.contains("THEN EXCLUDED.created_at"));
}
#[test]
fn upsert_sql_keeps_terminal_candidate_state_when_lifecycle_events_arrive_late() {
for sql in [
UPSERT_SQL,
UPSERT_CONFLICT_SQL,
UPSERT_CONFLICT_INHERIT_IS_CACHED_SQL,
] {
assert!(sql.contains(
"request_candidates.status IN ('success', 'failed', 'cancelled', 'skipped')"
));
assert!(
sql.contains("EXCLUDED.status IN ('available', 'unused', 'pending', 'streaming')")
);
assert!(sql.contains(
"request_candidates.status = 'streaming' AND EXCLUDED.status = 'pending'"
));
assert!(sql.contains("THEN request_candidates.status"));
assert!(sql.contains("THEN request_candidates.latency_ms"));
}
}
#[tokio::test]
async fn repository_constructs_from_lazy_pool() {
let factory = PostgresPoolFactory::new(PostgresPoolConfig {
@@ -361,6 +361,17 @@ fn merge_candidate(
candidate: UpsertRequestCandidateRecord,
existing: Option<StoredRequestCandidate>,
) -> Result<StoredRequestCandidate, DataLayerError> {
let preserve_existing_lifecycle = existing.as_ref().is_some_and(|value| {
request_candidate_lifecycle_would_regress(value.status, candidate.status)
});
let merged_status = if preserve_existing_lifecycle {
existing
.as_ref()
.map(|value| value.status)
.unwrap_or(candidate.status)
} else {
candidate.status
};
let created_at_unix_ms = candidate
.created_at_unix_ms
.filter(|value| *value > 1000)
@@ -413,7 +424,7 @@ fn merge_candidate(
candidate
.key_id
.or_else(|| existing.as_ref().and_then(|value| value.key_id.clone())),
candidate.status,
merged_status,
candidate.skip_reason.or_else(|| {
existing
.as_ref()
@@ -422,25 +433,48 @@ fn merge_candidate(
candidate
.is_cached
.unwrap_or_else(|| existing.as_ref().is_some_and(|value| value.is_cached)),
candidate.status_code.map(i32::from).or_else(|| {
if preserve_existing_lifecycle {
existing
.as_ref()
.and_then(|value| value.status_code.map(i32::from))
}),
candidate
.error_type
.or_else(|| existing.as_ref().and_then(|value| value.error_type.clone())),
candidate.error_message.or_else(|| {
} else {
candidate.status_code.map(i32::from).or_else(|| {
existing
.as_ref()
.and_then(|value| value.status_code.map(i32::from))
})
},
if preserve_existing_lifecycle {
existing.as_ref().and_then(|value| value.error_type.clone())
} else {
candidate
.error_type
.or_else(|| existing.as_ref().and_then(|value| value.error_type.clone()))
},
if preserve_existing_lifecycle {
existing
.as_ref()
.and_then(|value| value.error_message.clone())
}),
candidate.latency_ms.map(to_i32_u64).transpose()?.or(
} else {
candidate.error_message.or_else(|| {
existing
.as_ref()
.and_then(|value| value.error_message.clone())
})
},
if preserve_existing_lifecycle {
match existing.as_ref().and_then(|value| value.latency_ms) {
Some(value) => Some(to_i32_u64(value)?),
None => None,
},
),
}
} else {
candidate.latency_ms.map(to_i32_u64).transpose()?.or(
match existing.as_ref().and_then(|value| value.latency_ms) {
Some(value) => Some(to_i32_u64(value)?),
None => None,
},
)
},
candidate.concurrent_requests.map(to_i32).transpose()?.or(
match existing
.as_ref()
@@ -462,18 +496,42 @@ fn merge_candidate(
.or_else(|| existing.as_ref().and_then(|value| value.started_at_unix_ms))
.map(|value| u64_to_i64(value, "request candidate started_at"))
.transpose()?,
candidate
.finished_at_unix_ms
.or_else(|| {
if preserve_existing_lifecycle {
existing
.as_ref()
.and_then(|value| value.finished_at_unix_ms)
} else {
candidate.finished_at_unix_ms.or_else(|| {
existing
.as_ref()
.and_then(|value| value.finished_at_unix_ms)
})
.map(|value| u64_to_i64(value, "request candidate finished_at"))
.transpose()?,
}
.map(|value| u64_to_i64(value, "request candidate finished_at"))
.transpose()?,
)
}
fn request_candidate_lifecycle_would_regress(
existing: RequestCandidateStatus,
incoming: RequestCandidateStatus,
) -> bool {
matches!(
existing,
RequestCandidateStatus::Success
| RequestCandidateStatus::Failed
| RequestCandidateStatus::Cancelled
| RequestCandidateStatus::Skipped
) && matches!(
incoming,
RequestCandidateStatus::Available
| RequestCandidateStatus::Unused
| RequestCandidateStatus::Pending
| RequestCandidateStatus::Streaming
) || existing == RequestCandidateStatus::Streaming
&& incoming == RequestCandidateStatus::Pending
}
fn aggregate_timeline(
candidates: Vec<StoredRequestCandidate>,
since_unix_secs: u64,
@@ -710,6 +768,23 @@ mod tests {
assert_eq!(updated.id, "candidate-1");
assert_eq!(updated.extra_data, Some(json!({"a": 1, "b": 2})));
let late_streaming = repository
.upsert(sample_upsert(
"candidate-late-streaming",
RequestCandidateStatus::Streaming,
Some(json!({"late": true})),
1_000_250,
))
.await
.expect("late streaming candidate should not regress terminal status");
assert_eq!(late_streaming.id, "candidate-1");
assert_eq!(late_streaming.status, RequestCandidateStatus::Success);
assert_eq!(late_streaming.finished_at_unix_ms, Some(1_000_502));
assert_eq!(
late_streaming.extra_data,
Some(json!({"a": 1, "b": 2, "late": true}))
);
assert_eq!(
repository
.list_by_request_id("request-1")
@@ -151,33 +151,85 @@ ON DUPLICATE KEY UPDATE
has_format_conversion = VALUES(has_format_conversion),
is_stream = VALUES(is_stream),
upstream_is_stream = VALUES(upstream_is_stream),
input_tokens = VALUES(input_tokens),
output_tokens = VALUES(output_tokens),
total_tokens = VALUES(total_tokens),
cache_creation_input_tokens = VALUES(cache_creation_input_tokens),
cache_creation_ephemeral_5m_input_tokens = VALUES(cache_creation_ephemeral_5m_input_tokens),
cache_creation_ephemeral_1h_input_tokens = VALUES(cache_creation_ephemeral_1h_input_tokens),
cache_read_input_tokens = VALUES(cache_read_input_tokens),
cache_creation_cost_usd = VALUES(cache_creation_cost_usd),
cache_read_cost_usd = VALUES(cache_read_cost_usd),
output_price_per_1m = VALUES(output_price_per_1m),
total_cost_usd = VALUES(total_cost_usd),
actual_total_cost_usd = VALUES(actual_total_cost_usd),
status_code = VALUES(status_code),
error_message = VALUES(error_message),
error_category = VALUES(error_category),
input_tokens = CASE
WHEN status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming') THEN input_tokens
ELSE VALUES(input_tokens)
END,
output_tokens = CASE
WHEN status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming') THEN output_tokens
ELSE VALUES(output_tokens)
END,
total_tokens = CASE
WHEN status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming') THEN total_tokens
ELSE VALUES(total_tokens)
END,
cache_creation_input_tokens = CASE
WHEN status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming') THEN cache_creation_input_tokens
ELSE VALUES(cache_creation_input_tokens)
END,
cache_creation_ephemeral_5m_input_tokens = CASE
WHEN status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming') THEN cache_creation_ephemeral_5m_input_tokens
ELSE VALUES(cache_creation_ephemeral_5m_input_tokens)
END,
cache_creation_ephemeral_1h_input_tokens = CASE
WHEN status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming') THEN cache_creation_ephemeral_1h_input_tokens
ELSE VALUES(cache_creation_ephemeral_1h_input_tokens)
END,
cache_read_input_tokens = CASE
WHEN status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming') THEN cache_read_input_tokens
ELSE VALUES(cache_read_input_tokens)
END,
cache_creation_cost_usd = CASE
WHEN status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming') THEN cache_creation_cost_usd
ELSE VALUES(cache_creation_cost_usd)
END,
cache_read_cost_usd = CASE
WHEN status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming') THEN cache_read_cost_usd
ELSE VALUES(cache_read_cost_usd)
END,
output_price_per_1m = CASE
WHEN status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming') THEN output_price_per_1m
ELSE VALUES(output_price_per_1m)
END,
total_cost_usd = CASE
WHEN status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming') THEN total_cost_usd
ELSE VALUES(total_cost_usd)
END,
actual_total_cost_usd = CASE
WHEN status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming') THEN actual_total_cost_usd
ELSE VALUES(actual_total_cost_usd)
END,
status_code = CASE
WHEN status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming') THEN status_code
WHEN status = 'streaming' AND VALUES(status) = 'pending' THEN status_code
ELSE VALUES(status_code)
END,
error_message = CASE
WHEN status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming') THEN error_message
WHEN status = 'streaming' AND VALUES(status) = 'pending' THEN error_message
ELSE VALUES(error_message)
END,
error_category = CASE
WHEN status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming') THEN error_category
WHEN status = 'streaming' AND VALUES(status) = 'pending' THEN error_category
ELSE VALUES(error_category)
END,
response_time_ms = CASE
WHEN status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming') THEN response_time_ms
WHEN VALUES(response_time_ms) IS NULL OR VALUES(response_time_ms) = 0
THEN COALESCE(response_time_ms, VALUES(response_time_ms))
ELSE VALUES(response_time_ms)
END,
first_byte_time_ms = CASE
WHEN status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming') THEN first_byte_time_ms
WHEN VALUES(first_byte_time_ms) IS NULL OR VALUES(first_byte_time_ms) = 0
THEN COALESCE(first_byte_time_ms, VALUES(first_byte_time_ms))
ELSE VALUES(first_byte_time_ms)
END,
status = VALUES(status),
billing_status = VALUES(billing_status),
billing_status = CASE
WHEN status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming') THEN billing_status
ELSE VALUES(billing_status)
END,
request_metadata = VALUES(request_metadata),
candidate_id = VALUES(candidate_id),
candidate_index = VALUES(candidate_index),
@@ -187,8 +239,19 @@ ON DUPLICATE KEY UPDATE
route_kind = VALUES(route_kind),
execution_path = VALUES(execution_path),
local_execution_runtime_miss_reason = VALUES(local_execution_runtime_miss_reason),
finalized_at = VALUES(finalized_at),
updated_at_unix_secs = VALUES(updated_at_unix_secs)
finalized_at = CASE
WHEN status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming') THEN finalized_at
ELSE VALUES(finalized_at)
END,
updated_at_unix_secs = CASE
WHEN status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming') THEN updated_at_unix_secs
ELSE VALUES(updated_at_unix_secs)
END,
status = CASE
WHEN status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming') THEN status
WHEN status = 'streaming' AND VALUES(status) = 'pending' THEN status
ELSE VALUES(status)
END
"#;
const SELECT_STALE_PENDING_USAGE_BATCH_SQL: &str = r#"
@@ -1550,6 +1613,20 @@ mod tests {
assert!(source.contains("CAST(COALESCE(SUM(total_requests), 0) AS SIGNED) AS requests"));
}
#[test]
fn mysql_usage_upsert_keeps_terminal_state_when_streaming_arrives_late() {
assert!(super::UPSERT_USAGE_SQL.contains(
"status IN ('completed', 'failed', 'cancelled') AND VALUES(status) IN ('pending', 'streaming')"
));
assert!(super::UPSERT_USAGE_SQL.contains("input_tokens = CASE"));
assert!(super::UPSERT_USAGE_SQL.contains("status_code = CASE"));
assert!(super::UPSERT_USAGE_SQL.contains("billing_status = CASE"));
assert!(super::UPSERT_USAGE_SQL.contains("finalized_at = CASE"));
assert!(super::UPSERT_USAGE_SQL.contains("updated_at_unix_secs = CASE"));
assert!(super::UPSERT_USAGE_SQL
.contains("WHEN status = 'streaming' AND VALUES(status) = 'pending' THEN status"));
}
#[tokio::test]
async fn mysql_usage_write_repository_upserts_when_url_is_set() {
let Some(database_url) = std::env::var("AETHER_TEST_MYSQL_URL")
@@ -8029,8 +8029,26 @@ ORDER BY "usage".user_id ASC
) -> Result<StoredRequestUsageAudit, DataLayerError> {
usage.validate()?;
let usage = strip_deprecated_usage_display_fields(usage);
let prepared = prepare_usage_upsert_context(&usage)?;
self.tx_runner
.run_read_write(|tx| {
let PreparedUsageUpsert {
request_headers_json,
provider_request_headers_json,
response_headers_json,
client_response_headers_json,
request_body_storage,
provider_request_body_storage,
response_body_storage,
client_response_body_storage,
http_audit_refs,
http_audit_states,
http_audit_capture_mode,
routing_snapshot,
settlement_pricing_snapshot,
request_metadata_value,
request_metadata_json,
} = prepared;
Box::pin(async move {
lock_usage_request_id_in_tx(tx, &usage.request_id).await?;
@@ -8052,102 +8070,6 @@ ORDER BY "usage".user_id ASC
let previous_usage =
find_usage_by_request_id_in_tx(tx, &usage.request_id).await?;
let request_headers_json = json_bind_text(usage.request_headers.as_ref())?;
let request_body_storage =
prepare_usage_body_storage(usage.request_body.as_ref())?;
let provider_request_headers_json =
json_bind_text(usage.provider_request_headers.as_ref())?;
let provider_request_body_storage =
prepare_usage_body_storage(usage.provider_request_body.as_ref())?;
let response_headers_json = json_bind_text(usage.response_headers.as_ref())?;
let response_body_storage =
prepare_usage_body_storage(usage.response_body.as_ref())?;
let client_response_headers_json =
json_bind_text(usage.client_response_headers.as_ref())?;
let client_response_body_storage =
prepare_usage_body_storage(usage.client_response_body.as_ref())?;
let http_audit_refs = UsageHttpAuditRefs {
request_body_ref: resolved_write_usage_body_ref(
usage.request_body_ref.as_deref(),
&usage.request_id,
UsageBodyField::RequestBody,
request_body_storage.has_detached_blob(),
None,
),
provider_request_body_ref: resolved_write_usage_body_ref(
usage.provider_request_body_ref.as_deref(),
&usage.request_id,
UsageBodyField::ProviderRequestBody,
provider_request_body_storage.has_detached_blob(),
None,
),
response_body_ref: resolved_write_usage_body_ref(
usage.response_body_ref.as_deref(),
&usage.request_id,
UsageBodyField::ResponseBody,
response_body_storage.has_detached_blob(),
None,
),
client_response_body_ref: resolved_write_usage_body_ref(
usage.client_response_body_ref.as_deref(),
&usage.request_id,
UsageBodyField::ClientResponseBody,
client_response_body_storage.has_detached_blob(),
None,
),
};
let http_audit_states = UsageHttpAuditStates {
request_body_state: usage.request_body_state,
provider_request_body_state: usage.provider_request_body_state,
response_body_state: usage.response_body_state,
client_response_body_state: usage.client_response_body_state,
};
let request_metadata_value = prepare_request_metadata_for_body_storage(
usage.request_metadata.clone(),
[
(
UsageBodyField::RequestBody,
&request_body_storage,
usage.request_body.as_ref(),
usage.request_body_ref.as_deref(),
),
(
UsageBodyField::ProviderRequestBody,
&provider_request_body_storage,
usage.provider_request_body.as_ref(),
usage.provider_request_body_ref.as_deref(),
),
(
UsageBodyField::ResponseBody,
&response_body_storage,
usage.response_body.as_ref(),
usage.response_body_ref.as_deref(),
),
(
UsageBodyField::ClientResponseBody,
&client_response_body_storage,
usage.client_response_body.as_ref(),
usage.client_response_body_ref.as_deref(),
),
],
);
let http_audit_capture_mode = usage_http_audit_capture_mode(
&http_audit_refs,
[
usage.request_body.as_ref(),
usage.provider_request_body.as_ref(),
usage.response_body.as_ref(),
usage.client_response_body.as_ref(),
],
);
let routing_snapshot =
usage_routing_snapshot_from_usage(&usage, request_metadata_value.as_ref());
let settlement_pricing_snapshot = usage_settlement_pricing_snapshot_from_usage(
&usage,
request_metadata_value.as_ref(),
)?;
let request_metadata_json = json_bind_text(request_metadata_value.as_ref())?;
let _row = sqlx::query(UPSERT_SQL)
.bind(Uuid::new_v4().to_string())
.bind(&usage.request_id)
@@ -10486,6 +10408,25 @@ impl UsageSettlementPricingSnapshot {
}
}
#[derive(Debug)]
struct PreparedUsageUpsert {
request_headers_json: Option<String>,
provider_request_headers_json: Option<String>,
response_headers_json: Option<String>,
client_response_headers_json: Option<String>,
request_body_storage: UsageBodyStorage,
provider_request_body_storage: UsageBodyStorage,
response_body_storage: UsageBodyStorage,
client_response_body_storage: UsageBodyStorage,
http_audit_refs: UsageHttpAuditRefs,
http_audit_states: UsageHttpAuditStates,
http_audit_capture_mode: &'static str,
routing_snapshot: UsageRoutingSnapshot,
settlement_pricing_snapshot: UsageSettlementPricingSnapshot,
request_metadata_value: Option<Value>,
request_metadata_json: Option<String>,
}
fn prepare_usage_body_storage(value: Option<&Value>) -> Result<UsageBodyStorage, DataLayerError> {
let Some(value) = value else {
return Ok(UsageBodyStorage {
@@ -10530,6 +10471,118 @@ fn json_bind_text(value: Option<&Value>) -> Result<Option<String>, DataLayerErro
.transpose()
}
fn prepare_usage_upsert_context(
usage: &UpsertUsageRecord,
) -> Result<PreparedUsageUpsert, DataLayerError> {
let request_headers_json = json_bind_text(usage.request_headers.as_ref())?;
let request_body_storage = prepare_usage_body_storage(usage.request_body.as_ref())?;
let provider_request_headers_json = json_bind_text(usage.provider_request_headers.as_ref())?;
let provider_request_body_storage =
prepare_usage_body_storage(usage.provider_request_body.as_ref())?;
let response_headers_json = json_bind_text(usage.response_headers.as_ref())?;
let response_body_storage = prepare_usage_body_storage(usage.response_body.as_ref())?;
let client_response_headers_json = json_bind_text(usage.client_response_headers.as_ref())?;
let client_response_body_storage =
prepare_usage_body_storage(usage.client_response_body.as_ref())?;
let http_audit_refs = UsageHttpAuditRefs {
request_body_ref: resolved_write_usage_body_ref(
usage.request_body_ref.as_deref(),
&usage.request_id,
UsageBodyField::RequestBody,
request_body_storage.has_detached_blob(),
None,
),
provider_request_body_ref: resolved_write_usage_body_ref(
usage.provider_request_body_ref.as_deref(),
&usage.request_id,
UsageBodyField::ProviderRequestBody,
provider_request_body_storage.has_detached_blob(),
None,
),
response_body_ref: resolved_write_usage_body_ref(
usage.response_body_ref.as_deref(),
&usage.request_id,
UsageBodyField::ResponseBody,
response_body_storage.has_detached_blob(),
None,
),
client_response_body_ref: resolved_write_usage_body_ref(
usage.client_response_body_ref.as_deref(),
&usage.request_id,
UsageBodyField::ClientResponseBody,
client_response_body_storage.has_detached_blob(),
None,
),
};
let http_audit_states = UsageHttpAuditStates {
request_body_state: usage.request_body_state,
provider_request_body_state: usage.provider_request_body_state,
response_body_state: usage.response_body_state,
client_response_body_state: usage.client_response_body_state,
};
let request_metadata_value = prepare_request_metadata_for_body_storage(
usage.request_metadata.clone(),
[
(
UsageBodyField::RequestBody,
&request_body_storage,
usage.request_body.as_ref(),
usage.request_body_ref.as_deref(),
),
(
UsageBodyField::ProviderRequestBody,
&provider_request_body_storage,
usage.provider_request_body.as_ref(),
usage.provider_request_body_ref.as_deref(),
),
(
UsageBodyField::ResponseBody,
&response_body_storage,
usage.response_body.as_ref(),
usage.response_body_ref.as_deref(),
),
(
UsageBodyField::ClientResponseBody,
&client_response_body_storage,
usage.client_response_body.as_ref(),
usage.client_response_body_ref.as_deref(),
),
],
);
let http_audit_capture_mode = usage_http_audit_capture_mode(
&http_audit_refs,
[
usage.request_body.as_ref(),
usage.provider_request_body.as_ref(),
usage.response_body.as_ref(),
usage.client_response_body.as_ref(),
],
);
let routing_snapshot =
usage_routing_snapshot_from_usage(usage, request_metadata_value.as_ref());
let settlement_pricing_snapshot =
usage_settlement_pricing_snapshot_from_usage(usage, request_metadata_value.as_ref())?;
let request_metadata_json = json_bind_text(request_metadata_value.as_ref())?;
Ok(PreparedUsageUpsert {
request_headers_json,
provider_request_headers_json,
response_headers_json,
client_response_headers_json,
request_body_storage,
provider_request_body_storage,
response_body_storage,
client_response_body_storage,
http_audit_refs,
http_audit_states,
http_audit_capture_mode,
routing_snapshot,
settlement_pricing_snapshot,
request_metadata_value,
request_metadata_json,
})
}
fn usage_body_capture_state_bind_text(
value: Option<UsageBodyCaptureState>,
) -> Option<&'static str> {
+107 -15
View File
@@ -172,18 +172,54 @@ ON CONFLICT (request_id) DO UPDATE SET
has_format_conversion = excluded.has_format_conversion,
is_stream = excluded.is_stream,
upstream_is_stream = excluded.upstream_is_stream,
input_tokens = excluded.input_tokens,
output_tokens = excluded.output_tokens,
total_tokens = excluded.total_tokens,
cache_creation_input_tokens = excluded.cache_creation_input_tokens,
cache_creation_ephemeral_5m_input_tokens = excluded.cache_creation_ephemeral_5m_input_tokens,
cache_creation_ephemeral_1h_input_tokens = excluded.cache_creation_ephemeral_1h_input_tokens,
cache_read_input_tokens = excluded.cache_read_input_tokens,
cache_creation_cost_usd = excluded.cache_creation_cost_usd,
cache_read_cost_usd = excluded.cache_read_cost_usd,
output_price_per_1m = excluded.output_price_per_1m,
total_cost_usd = excluded.total_cost_usd,
actual_total_cost_usd = excluded.actual_total_cost_usd,
input_tokens = CASE
WHEN "usage".status IN ('completed', 'failed', 'cancelled') AND excluded.status IN ('pending', 'streaming') THEN "usage".input_tokens
ELSE excluded.input_tokens
END,
output_tokens = CASE
WHEN "usage".status IN ('completed', 'failed', 'cancelled') AND excluded.status IN ('pending', 'streaming') THEN "usage".output_tokens
ELSE excluded.output_tokens
END,
total_tokens = CASE
WHEN "usage".status IN ('completed', 'failed', 'cancelled') AND excluded.status IN ('pending', 'streaming') THEN "usage".total_tokens
ELSE excluded.total_tokens
END,
cache_creation_input_tokens = CASE
WHEN "usage".status IN ('completed', 'failed', 'cancelled') AND excluded.status IN ('pending', 'streaming') THEN "usage".cache_creation_input_tokens
ELSE excluded.cache_creation_input_tokens
END,
cache_creation_ephemeral_5m_input_tokens = CASE
WHEN "usage".status IN ('completed', 'failed', 'cancelled') AND excluded.status IN ('pending', 'streaming') THEN "usage".cache_creation_ephemeral_5m_input_tokens
ELSE excluded.cache_creation_ephemeral_5m_input_tokens
END,
cache_creation_ephemeral_1h_input_tokens = CASE
WHEN "usage".status IN ('completed', 'failed', 'cancelled') AND excluded.status IN ('pending', 'streaming') THEN "usage".cache_creation_ephemeral_1h_input_tokens
ELSE excluded.cache_creation_ephemeral_1h_input_tokens
END,
cache_read_input_tokens = CASE
WHEN "usage".status IN ('completed', 'failed', 'cancelled') AND excluded.status IN ('pending', 'streaming') THEN "usage".cache_read_input_tokens
ELSE excluded.cache_read_input_tokens
END,
cache_creation_cost_usd = CASE
WHEN "usage".status IN ('completed', 'failed', 'cancelled') AND excluded.status IN ('pending', 'streaming') THEN "usage".cache_creation_cost_usd
ELSE excluded.cache_creation_cost_usd
END,
cache_read_cost_usd = CASE
WHEN "usage".status IN ('completed', 'failed', 'cancelled') AND excluded.status IN ('pending', 'streaming') THEN "usage".cache_read_cost_usd
ELSE excluded.cache_read_cost_usd
END,
output_price_per_1m = CASE
WHEN "usage".status IN ('completed', 'failed', 'cancelled') AND excluded.status IN ('pending', 'streaming') THEN "usage".output_price_per_1m
ELSE excluded.output_price_per_1m
END,
total_cost_usd = CASE
WHEN "usage".status IN ('completed', 'failed', 'cancelled') AND excluded.status IN ('pending', 'streaming') THEN "usage".total_cost_usd
ELSE excluded.total_cost_usd
END,
actual_total_cost_usd = CASE
WHEN "usage".status IN ('completed', 'failed', 'cancelled') AND excluded.status IN ('pending', 'streaming') THEN "usage".actual_total_cost_usd
ELSE excluded.actual_total_cost_usd
END,
status_code = CASE
WHEN "usage".status IN ('completed', 'failed', 'cancelled') AND excluded.status IN ('pending', 'streaming') THEN "usage".status_code
WHEN "usage".status = 'streaming' AND excluded.status = 'pending' THEN "usage".status_code
@@ -214,7 +250,10 @@ ON CONFLICT (request_id) DO UPDATE SET
WHEN "usage".status = 'streaming' AND excluded.status = 'pending' THEN "usage".status
ELSE excluded.status
END,
billing_status = excluded.billing_status,
billing_status = CASE
WHEN "usage".status IN ('completed', 'failed', 'cancelled') AND excluded.status IN ('pending', 'streaming') THEN "usage".billing_status
ELSE excluded.billing_status
END,
request_metadata = excluded.request_metadata,
candidate_id = COALESCE(excluded.candidate_id, "usage".candidate_id),
candidate_index = COALESCE(excluded.candidate_index, "usage".candidate_index),
@@ -224,8 +263,14 @@ ON CONFLICT (request_id) DO UPDATE SET
route_kind = excluded.route_kind,
execution_path = excluded.execution_path,
local_execution_runtime_miss_reason = excluded.local_execution_runtime_miss_reason,
finalized_at = excluded.finalized_at,
updated_at_unix_secs = excluded.updated_at_unix_secs
finalized_at = CASE
WHEN "usage".status IN ('completed', 'failed', 'cancelled') AND excluded.status IN ('pending', 'streaming') THEN "usage".finalized_at
ELSE excluded.finalized_at
END,
updated_at_unix_secs = CASE
WHEN "usage".status IN ('completed', 'failed', 'cancelled') AND excluded.status IN ('pending', 'streaming') THEN "usage".updated_at_unix_secs
ELSE excluded.updated_at_unix_secs
END
"#;
const SELECT_STALE_PENDING_USAGE_BATCH_SQL: &str = r#"
@@ -4520,6 +4565,53 @@ mod tests {
assert_eq!(existing.updated_at_unix_secs, 1_000);
}
#[tokio::test]
async fn sqlite_usage_write_repository_does_not_regress_terminal_usage_from_late_streaming() {
let pool = sqlx::sqlite::SqlitePoolOptions::new()
.max_connections(1)
.connect("sqlite::memory:")
.await
.expect("sqlite pool should connect");
run_sqlite_migrations(&pool)
.await
.expect("sqlite migrations should run");
seed_stats_targets(&pool).await;
let repository = SqliteUsageWriteRepository::new(pool);
repository
.upsert(sample_usage("request-1", "completed", "pending", 1_000))
.await
.expect("terminal usage should upsert");
let mut late_streaming = sample_usage("request-1", "streaming", "pending", 1_001);
late_streaming.input_tokens = Some(0);
late_streaming.output_tokens = Some(0);
late_streaming.total_tokens = Some(0);
late_streaming.cache_read_input_tokens = Some(0);
late_streaming.cache_read_cost_usd = Some(0.0);
late_streaming.total_cost_usd = Some(0.0);
late_streaming.actual_total_cost_usd = Some(0.0);
late_streaming.response_time_ms = Some(9_999);
late_streaming.first_byte_time_ms = Some(9_999);
late_streaming.finalized_at_unix_secs = None;
let current = repository
.upsert(late_streaming)
.await
.expect("late streaming usage should not regress terminal usage");
assert_eq!(current.status, "completed");
assert_eq!(current.billing_status, "pending");
assert_eq!(current.total_tokens, 7);
assert_eq!(current.cache_read_input_tokens, 2);
assert_eq!(current.total_cost_usd, 0.5);
assert_eq!(current.actual_total_cost_usd, 0.4);
assert_eq!(current.response_time_ms, Some(42));
assert_eq!(current.first_byte_time_ms, Some(12));
assert_eq!(current.finalized_at_unix_secs, Some(1_000));
assert_eq!(current.updated_at_unix_secs, 1_000);
}
#[tokio::test]
async fn sqlite_usage_write_repository_cleans_stale_pending_requests() {
let pool = sqlx::sqlite::SqlitePoolOptions::new()
+56
View File
@@ -748,6 +748,14 @@ pub struct RuntimeQueueEntry {
pub fields: BTreeMap<String, String>,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub struct RuntimeQueueStats {
pub stream_length: u64,
pub group_pending: u64,
pub group_lag: Option<u64>,
pub oldest_pending_idle_ms: Option<u64>,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct RuntimeQueueReclaimConfig {
pub min_idle_ms: u64,
@@ -817,6 +825,12 @@ pub trait RuntimeQueueStore: Send + Sync {
-> Result<usize, DataLayerError>;
async fn delete(&self, stream: &str, ids: &[String]) -> Result<usize, DataLayerError>;
async fn stats(
&self,
stream: &str,
group: Option<&str>,
) -> Result<RuntimeQueueStats, DataLayerError>;
}
#[async_trait]
@@ -999,6 +1013,31 @@ impl RuntimeQueueStore for RuntimeState {
}
}
}
async fn stats(
&self,
stream: &str,
group: Option<&str>,
) -> Result<RuntimeQueueStats, DataLayerError> {
validate_runtime_queue_name(stream, "runtime queue stream")?;
if let Some(group) = group {
validate_runtime_queue_name(group, "runtime queue group")?;
}
match self.backend.as_ref() {
RuntimeStateBackend::Memory(memory) => Ok(memory.queue_stats(stream, group).await),
RuntimeStateBackend::Redis(redis) => {
redis
.stream
.stats(
&RedisStreamName(stream.to_string()),
group
.map(|value| RedisConsumerGroup(value.to_string()))
.as_ref(),
)
.await
}
}
}
}
#[async_trait]
@@ -1921,6 +1960,12 @@ mod tests {
.collect::<Vec<_>>(),
vec![first.clone(), second.clone()]
);
let stats = RuntimeQueueStore::stats(runtime, "contract:stream", Some("workers"))
.await
.expect("queue stats");
assert_eq!(stats.stream_length, 2);
assert_eq!(stats.group_pending, 2);
assert_eq!(stats.group_lag, Some(0));
assert!(RuntimeQueueStore::read_group(
runtime,
"contract:stream",
@@ -1951,6 +1996,11 @@ mod tests {
.map(|entry| entry.id.clone())
.collect::<Vec<_>>();
assert_eq!(ids, vec![first.clone(), second.clone()]);
let stats = RuntimeQueueStore::stats(runtime, "contract:stream", Some("workers"))
.await
.expect("queue stats after claim");
assert_eq!(stats.group_pending, 2);
assert!(stats.oldest_pending_idle_ms.unwrap_or_default() <= 5000);
assert_eq!(
RuntimeQueueStore::ack(runtime, "contract:stream", "workers", &ids)
.await
@@ -1963,6 +2013,12 @@ mod tests {
.expect("delete"),
2
);
let stats = RuntimeQueueStore::stats(runtime, "contract:stream", Some("workers"))
.await
.expect("queue stats after delete");
assert_eq!(stats.stream_length, 0);
assert_eq!(stats.group_pending, 0);
assert_eq!(stats.group_lag, Some(0));
assert!(RuntimeQueueStore::read_group(
runtime,
"contract:stream",
+43 -1
View File
@@ -5,7 +5,7 @@ use std::time::{Duration, Instant};
use tokio::sync::Mutex;
use crate::{DataLayerError, RuntimeQueueEntry, RuntimeQueueReclaimConfig};
use crate::{DataLayerError, RuntimeQueueEntry, RuntimeQueueReclaimConfig, RuntimeQueueStats};
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct MemoryRuntimeStateConfig {
@@ -801,6 +801,48 @@ impl MemoryRuntimeBackend {
before.saturating_sub(stream_state.entries.len())
}
pub(crate) async fn queue_stats(&self, stream: &str, group: Option<&str>) -> RuntimeQueueStats {
let mut queues = self.queues.lock().await;
let now = Instant::now();
prune_memory_key(&mut queues, stream, now);
let Some(stream_state) = queues.get(stream) else {
return RuntimeQueueStats::default();
};
let stream_length = stream_state.entries.len() as u64;
let Some(group_name) = group else {
return RuntimeQueueStats {
stream_length,
..RuntimeQueueStats::default()
};
};
let Some(group_state) = stream_state.groups.get(group_name) else {
return RuntimeQueueStats {
stream_length,
..RuntimeQueueStats::default()
};
};
let group_lag = stream_state
.entries
.iter()
.filter(|entry| entry.sequence > group_state.last_delivered_sequence)
.count() as u64;
let oldest_pending_idle_ms = group_state
.pending
.values()
.map(|entry| {
now.saturating_duration_since(entry.delivered_at)
.as_millis() as u64
})
.max();
RuntimeQueueStats {
stream_length,
group_pending: group_state.pending.len() as u64,
group_lag: Some(group_lag),
oldest_pending_idle_ms,
}
}
pub(crate) async fn lock_try_acquire(
&self,
key: &str,
+120 -2
View File
@@ -12,6 +12,9 @@ pub(crate) type RedisManagedConnection = redis::aio::ConnectionManager;
const DEFAULT_BLOCKING_STREAM_LANES_FALLBACK: usize = 4;
const DEFAULT_BLOCKING_STREAM_LANES_CAP: usize = 16;
const MAX_BLOCKING_STREAM_LANES_CAP: usize = 64;
pub(crate) const REDIS_COMMAND_LATENCY_BUCKETS_MS: [u64; 12] =
[1, 5, 10, 25, 50, 100, 250, 500, 1_000, 2_500, 5_000, 10_000];
const REDIS_COMMAND_LATENCY_BUCKET_COUNT: usize = REDIS_COMMAND_LATENCY_BUCKETS_MS.len() + 1;
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize, PartialEq, Eq)]
pub struct RedisClientConfig {
@@ -190,6 +193,10 @@ impl RedisConnectionRouter {
.fetch_add(1, Ordering::Relaxed);
}
pub(crate) fn record_latency(&self, lane: RedisConnectionLane, elapsed: Duration) {
self.metrics.for_lane(lane).record_latency(elapsed);
}
pub(crate) fn lane_diagnostics(&self) -> Vec<RedisLaneDiagnostics> {
[
RedisConnectionLane::Fast,
@@ -204,6 +211,10 @@ impl RedisConnectionRouter {
lane: lane.as_str(),
command_errors: metrics.errors.load(Ordering::Relaxed),
command_timeouts: metrics.timeouts.load(Ordering::Relaxed),
command_count: metrics.command_count.load(Ordering::Relaxed),
command_latency_total_ms: metrics.latency_total_ms.load(Ordering::Relaxed),
command_latency_max_ms: metrics.latency_max_ms.load(Ordering::Relaxed),
command_latency_buckets: metrics.latency_buckets(),
}
})
.collect()
@@ -215,6 +226,16 @@ pub struct RedisLaneDiagnostics {
pub lane: &'static str,
pub command_errors: u64,
pub command_timeouts: u64,
pub command_count: u64,
pub command_latency_total_ms: u64,
pub command_latency_max_ms: u64,
pub command_latency_buckets: Vec<RedisCommandLatencyBucket>,
}
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize)]
pub struct RedisCommandLatencyBucket {
pub le_ms: Option<u64>,
pub count: u64,
}
#[derive(Default)]
@@ -236,10 +257,74 @@ impl RedisConnectionMetrics {
}
}
#[derive(Default)]
struct RedisLaneMetrics {
errors: AtomicU64,
timeouts: AtomicU64,
command_count: AtomicU64,
latency_total_ms: AtomicU64,
latency_max_ms: AtomicU64,
latency_bucket_counts: [AtomicU64; REDIS_COMMAND_LATENCY_BUCKET_COUNT],
}
impl Default for RedisLaneMetrics {
fn default() -> Self {
Self {
errors: AtomicU64::new(0),
timeouts: AtomicU64::new(0),
command_count: AtomicU64::new(0),
latency_total_ms: AtomicU64::new(0),
latency_max_ms: AtomicU64::new(0),
latency_bucket_counts: std::array::from_fn(|_| AtomicU64::new(0)),
}
}
}
impl RedisLaneMetrics {
fn record_latency(&self, elapsed: Duration) {
let elapsed_ms = u64::try_from(elapsed.as_millis()).unwrap_or(u64::MAX);
self.command_count.fetch_add(1, Ordering::Relaxed);
self.latency_total_ms
.fetch_add(elapsed_ms, Ordering::Relaxed);
update_atomic_max(&self.latency_max_ms, elapsed_ms);
let bucket_index = REDIS_COMMAND_LATENCY_BUCKETS_MS
.iter()
.position(|upper_bound_ms| elapsed_ms <= *upper_bound_ms)
.unwrap_or(REDIS_COMMAND_LATENCY_BUCKETS_MS.len());
self.latency_bucket_counts[bucket_index].fetch_add(1, Ordering::Relaxed);
}
fn latency_buckets(&self) -> Vec<RedisCommandLatencyBucket> {
let mut cumulative = 0u64;
let mut buckets = Vec::with_capacity(REDIS_COMMAND_LATENCY_BUCKET_COUNT);
for (index, upper_bound_ms) in REDIS_COMMAND_LATENCY_BUCKETS_MS.iter().enumerate() {
cumulative = cumulative
.saturating_add(self.latency_bucket_counts[index].load(Ordering::Relaxed));
buckets.push(RedisCommandLatencyBucket {
le_ms: Some(*upper_bound_ms),
count: cumulative,
});
}
cumulative = cumulative.saturating_add(
self.latency_bucket_counts[REDIS_COMMAND_LATENCY_BUCKETS_MS.len()]
.load(Ordering::Relaxed),
);
buckets.push(RedisCommandLatencyBucket {
le_ms: None,
count: cumulative,
});
buckets
}
}
fn update_atomic_max(target: &AtomicU64, value: u64) {
let mut current = target.load(Ordering::Relaxed);
while value > current {
match target.compare_exchange_weak(current, value, Ordering::Relaxed, Ordering::Relaxed) {
Ok(_) => break,
Err(next) => current = next,
}
}
}
fn connection_manager_config(
@@ -330,8 +415,11 @@ async fn connect_lane(
mod tests {
use super::{
blocking_stream_lane_count, default_blocking_stream_lane_count, RedisClientConfig,
RedisClientFactory, MAX_BLOCKING_STREAM_LANES_CAP,
RedisClientFactory, RedisLaneMetrics, MAX_BLOCKING_STREAM_LANES_CAP,
REDIS_COMMAND_LATENCY_BUCKETS_MS,
};
use std::sync::atomic::Ordering;
use std::time::Duration;
#[test]
fn factory_builds_lazy_client_from_valid_config() {
@@ -370,4 +458,34 @@ mod tests {
);
assert!(blocking_stream_lane_count(Some(0)).is_err());
}
#[test]
fn lane_metrics_record_cumulative_latency_buckets() {
let metrics = RedisLaneMetrics::default();
metrics.record_latency(Duration::from_millis(0));
metrics.record_latency(Duration::from_millis(12));
metrics.record_latency(Duration::from_millis(12_345));
assert_eq!(metrics.command_count.load(Ordering::Relaxed), 3);
assert_eq!(metrics.latency_total_ms.load(Ordering::Relaxed), 12_357);
assert_eq!(metrics.latency_max_ms.load(Ordering::Relaxed), 12_345);
let buckets = metrics.latency_buckets();
let le_1 = buckets
.iter()
.find(|bucket| bucket.le_ms == Some(1))
.expect("1ms bucket");
let le_25 = buckets
.iter()
.find(|bucket| bucket.le_ms == Some(25))
.expect("25ms bucket");
let plus_inf = buckets.last().expect("+Inf bucket");
assert_eq!(buckets.len(), REDIS_COMMAND_LATENCY_BUCKETS_MS.len() + 1);
assert_eq!(le_1.count, 1);
assert_eq!(le_25.count, 2);
assert_eq!(plus_inf.le_ms, None);
assert_eq!(plus_inf.count, 3);
}
}
@@ -39,11 +39,13 @@ pub(crate) async fn run_lane_with_timeout<T, F>(
where
F: std::future::Future<Output = Result<T, crate::DataLayerError>>,
{
let started = std::time::Instant::now();
let result = if let Some(timeout_ms) = timeout_ms {
match tokio::time::timeout(std::time::Duration::from_millis(timeout_ms), future).await {
Ok(result) => result,
Err(_) => {
connections.record_timeout(lane);
connections.record_latency(lane, started.elapsed());
return Err(crate::DataLayerError::TimedOut(format!(
"{operation} exceeded {timeout_ms}ms timeout"
)));
@@ -52,6 +54,7 @@ where
} else {
future.await
};
connections.record_latency(lane, started.elapsed());
if result.is_err() {
connections.record_error(lane);
}
@@ -54,7 +54,19 @@ return {1, 0, 0, remaining}
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize)]
pub struct RedisRuntimeDiagnostics {
pub connected_clients: Option<u64>,
pub blocked_clients: Option<u64>,
pub total_connections_received: Option<u64>,
pub rejected_connections: Option<u64>,
pub total_commands_processed: Option<u64>,
pub instantaneous_ops_per_sec: Option<u64>,
pub total_error_replies: Option<u64>,
pub expired_keys: Option<u64>,
pub evicted_keys: Option<u64>,
pub keyspace_hits: Option<u64>,
pub keyspace_misses: Option<u64>,
pub used_memory_bytes: Option<u64>,
pub maxmemory_bytes: Option<u64>,
pub memory_fragmentation_ratio_basis_points: Option<u64>,
pub lanes: Vec<RedisLaneDiagnostics>,
}
@@ -646,7 +658,22 @@ impl RedisRuntimeRunner {
fn parse_diagnostics(info: &str, lanes: Vec<RedisLaneDiagnostics>) -> RedisRuntimeDiagnostics {
RedisRuntimeDiagnostics {
connected_clients: parse_info_u64(info, "connected_clients"),
blocked_clients: parse_info_u64(info, "blocked_clients"),
total_connections_received: parse_info_u64(info, "total_connections_received"),
rejected_connections: parse_info_u64(info, "rejected_connections"),
total_commands_processed: parse_info_u64(info, "total_commands_processed"),
instantaneous_ops_per_sec: parse_info_u64(info, "instantaneous_ops_per_sec"),
total_error_replies: parse_info_u64(info, "total_error_replies"),
expired_keys: parse_info_u64(info, "expired_keys"),
evicted_keys: parse_info_u64(info, "evicted_keys"),
keyspace_hits: parse_info_u64(info, "keyspace_hits"),
keyspace_misses: parse_info_u64(info, "keyspace_misses"),
used_memory_bytes: parse_info_u64(info, "used_memory"),
maxmemory_bytes: parse_info_u64(info, "maxmemory"),
memory_fragmentation_ratio_basis_points: parse_info_f64_basis_points(
info,
"mem_fragmentation_ratio",
),
lanes,
}
}
@@ -660,6 +687,17 @@ fn parse_info_u64(info: &str, key: &str) -> Option<u64> {
})
}
fn parse_info_f64_basis_points(info: &str, key: &str) -> Option<u64> {
info.lines().find_map(|line| {
let (name, value) = line.split_once(':')?;
if name != key {
return None;
}
let parsed = value.trim().parse::<f64>().ok()?;
(parsed.is_finite() && parsed >= 0.0).then(|| (parsed * 10_000.0).round() as u64)
})
}
fn key_belongs_to_prefix(key: &str, prefix: &str) -> bool {
prefix.is_empty()
|| key == prefix
@@ -675,7 +713,7 @@ mod tests {
#[test]
fn parses_runtime_diagnostics_from_info() {
let parsed = parse_diagnostics(
"# Clients\r\nconnected_clients:5\r\n# Stats\r\ntotal_connections_received:42\r\n",
"# Clients\r\nconnected_clients:5\r\nblocked_clients:2\r\n# Memory\r\nused_memory:1048576\r\nmaxmemory:8388608\r\nmem_fragmentation_ratio:1.25\r\n# Stats\r\ntotal_connections_received:42\r\nrejected_connections:0\r\ntotal_commands_processed:99\r\ninstantaneous_ops_per_sec:7\r\ntotal_error_replies:1\r\nexpired_keys:3\r\nevicted_keys:4\r\nkeyspace_hits:10\r\nkeyspace_misses:2\r\n",
Vec::new(),
);
@@ -683,7 +721,19 @@ mod tests {
parsed,
RedisRuntimeDiagnostics {
connected_clients: Some(5),
blocked_clients: Some(2),
total_connections_received: Some(42),
rejected_connections: Some(0),
total_commands_processed: Some(99),
instantaneous_ops_per_sec: Some(7),
total_error_replies: Some(1),
expired_keys: Some(3),
evicted_keys: Some(4),
keyspace_hits: Some(10),
keyspace_misses: Some(2),
used_memory_bytes: Some(1_048_576),
maxmemory_bytes: Some(8_388_608),
memory_fragmentation_ratio_basis_points: Some(12_500),
lanes: Vec::new(),
}
);
+328 -2
View File
@@ -10,7 +10,7 @@ use crate::redis::{
run_lane_with_timeout, RedisClientConfig, RedisClientFactory, RedisConnectionLane,
RedisConnectionRouter, RedisKeyspace,
};
use crate::DataLayerError;
use crate::{DataLayerError, RuntimeQueueStats};
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub struct RedisStreamName(pub String);
@@ -383,6 +383,87 @@ impl RedisStreamRunner {
.await
}
pub async fn stats(
&self,
stream: &RedisStreamName,
group: Option<&RedisConsumerGroup>,
) -> Result<RuntimeQueueStats, DataLayerError> {
validate_stream_name(stream)?;
if let Some(group) = group {
validate_group(group)?;
}
let stream_length = self
.run_with_timeout(RedisConnectionLane::Stream, "redis stream xlen", async {
let mut connection = self.connections.connection(RedisConnectionLane::Stream);
redis::cmd("XLEN")
.arg(&stream.0)
.query_async::<u64>(&mut connection)
.await
.map_redis_err()
})
.await?;
let Some(group) = group else {
return Ok(RuntimeQueueStats {
stream_length,
..RuntimeQueueStats::default()
});
};
let groups = self
.run_with_timeout(
RedisConnectionLane::Stream,
"redis stream xinfo groups",
async {
let mut connection = self.connections.connection(RedisConnectionLane::Stream);
let reply = match redis::cmd("XINFO")
.arg("GROUPS")
.arg(&stream.0)
.query_async::<RedisValue>(&mut connection)
.await
{
Ok(reply) => reply,
Err(err) if redis_stream_stats_missing_stream(&err) => {
return Ok(RedisXInfoGroupStats::default());
}
Err(err) => return Err(redis_error(err)),
};
parse_xinfo_group_stats(reply, &group.0)
},
)
.await?;
let oldest_pending_idle_ms = self
.run_with_timeout(
RedisConnectionLane::Stream,
"redis stream xpending summary",
async {
let mut connection = self.connections.connection(RedisConnectionLane::Stream);
let reply = match redis::cmd("XPENDING")
.arg(&stream.0)
.arg(&group.0)
.query_async::<RedisValue>(&mut connection)
.await
{
Ok(reply) => reply,
Err(err) if redis_stream_stats_missing_stream_or_group(&err) => {
return Ok(None);
}
Err(err) => return Err(redis_error(err)),
};
parse_xpending_oldest_idle_ms(reply)
},
)
.await?;
Ok(RuntimeQueueStats {
stream_length,
group_pending: groups.pending.unwrap_or_default(),
group_lag: groups.lag,
oldest_pending_idle_ms,
})
}
async fn run_with_timeout<T, F>(
&self,
lane: RedisConnectionLane,
@@ -403,6 +484,17 @@ impl RedisStreamRunner {
}
}
fn redis_stream_stats_missing_stream(error: &redis::RedisError) -> bool {
error.code() == Some("ERR") && error.to_string().contains("no such key")
}
fn redis_stream_stats_missing_stream_or_group(error: &redis::RedisError) -> bool {
let message = error.to_string();
redis_stream_stats_missing_stream(error)
|| error.code() == Some("NOGROUP")
|| (message.contains("NOGROUP") && message.contains("consumer group"))
}
fn validate_stream_name(stream: &RedisStreamName) -> Result<(), DataLayerError> {
if stream.0.trim().is_empty() {
return Err(DataLayerError::InvalidInput(
@@ -577,12 +669,133 @@ fn parse_string_value(value: &RedisValue, context: &str) -> Result<String, DataL
})
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
struct RedisXInfoGroupStats {
pending: Option<u64>,
lag: Option<u64>,
}
fn parse_xinfo_group_stats(
value: RedisValue,
group_name: &str,
) -> Result<RedisXInfoGroupStats, DataLayerError> {
let RedisValue::Array(groups) = value else {
return Err(DataLayerError::UnexpectedValue(
"redis xinfo groups returned non-array payload".to_string(),
));
};
for group in groups {
let fields = parse_info_fields(&group, "redis xinfo group")?;
if fields
.get("name")
.is_some_and(|value| value.as_str() == group_name)
{
return Ok(RedisXInfoGroupStats {
pending: fields
.get("pending")
.and_then(|value| value.parse::<u64>().ok()),
lag: fields
.get("lag")
.and_then(|value| value.parse::<u64>().ok()),
});
}
}
Ok(RedisXInfoGroupStats::default())
}
fn parse_xpending_oldest_idle_ms(value: RedisValue) -> Result<Option<u64>, DataLayerError> {
let RedisValue::Array(parts) = value else {
return Err(DataLayerError::UnexpectedValue(
"redis xpending summary returned non-array payload".to_string(),
));
};
let Some(total) = parts
.first()
.and_then(|value| parse_u64_value(value, "redis xpending pending count").ok())
else {
return Ok(None);
};
if total == 0 {
return Ok(None);
}
let Some(consumers) = parts.get(3) else {
return Ok(None);
};
let RedisValue::Array(consumers) = consumers else {
return Ok(None);
};
let mut oldest: Option<u64> = None;
for consumer in consumers {
let RedisValue::Array(fields) = consumer else {
continue;
};
let Some(idle_ms) = fields
.get(2)
.and_then(|value| parse_u64_value(value, "redis xpending consumer idle").ok())
else {
continue;
};
oldest = Some(oldest.map_or(idle_ms, |current| current.max(idle_ms)));
}
Ok(oldest)
}
fn parse_info_fields(
value: &RedisValue,
context: &str,
) -> Result<BTreeMap<String, String>, DataLayerError> {
match value {
RedisValue::Array(values) => {
if values.len() % 2 != 0 {
return Err(DataLayerError::UnexpectedValue(format!(
"{context} expected an even number of field elements, got {}",
values.len()
)));
}
let mut fields = BTreeMap::new();
for pair in values.chunks(2) {
if matches!(pair[1], RedisValue::Nil) {
continue;
}
fields.insert(
parse_string_value(&pair[0], context)?,
parse_string_value(&pair[1], context)?,
);
}
Ok(fields)
}
RedisValue::Map(entries) => entries
.iter()
.map(|(key, value)| {
Ok((
parse_string_value(key, context)?,
parse_string_value(value, context)?,
))
})
.collect(),
RedisValue::Nil => Ok(BTreeMap::new()),
_ => Err(DataLayerError::UnexpectedValue(format!(
"{context} expected a redis array/map payload"
))),
}
}
fn parse_u64_value(value: &RedisValue, context: &str) -> Result<u64, DataLayerError> {
from_redis_value::<u64>(value).map_err(|err| {
DataLayerError::UnexpectedValue(format!(
"{context} was not an unsigned integer-compatible redis value: {err}"
))
})
}
#[cfg(test)]
mod tests {
use std::collections::BTreeMap;
use super::{
parse_reclaim_result, parse_stream_read_entries, validate_consumer, validate_group,
parse_reclaim_result, parse_stream_read_entries, parse_xinfo_group_stats,
parse_xpending_oldest_idle_ms, redis_stream_stats_missing_stream,
redis_stream_stats_missing_stream_or_group, validate_consumer, validate_group,
validate_stream_name, validate_stream_position, RedisConsumerName, RedisStreamName,
RedisStreamReclaimConfig, RedisStreamReclaimResult, RedisStreamRunnerConfig,
};
@@ -706,4 +919,117 @@ mod tests {
}
);
}
#[test]
fn parses_xinfo_group_stats() {
let parsed = parse_xinfo_group_stats(
RedisValue::Array(vec![RedisValue::Array(vec![
RedisValue::BulkString(b"name".to_vec()),
RedisValue::BulkString(b"usage_consumers".to_vec()),
RedisValue::BulkString(b"consumers".to_vec()),
RedisValue::Int(4),
RedisValue::BulkString(b"pending".to_vec()),
RedisValue::Int(7),
RedisValue::BulkString(b"last-delivered-id".to_vec()),
RedisValue::BulkString(b"1710000000000-0".to_vec()),
RedisValue::BulkString(b"entries-read".to_vec()),
RedisValue::Int(12),
RedisValue::BulkString(b"lag".to_vec()),
RedisValue::Int(3),
])]),
"usage_consumers",
)
.expect("xinfo groups should parse");
assert_eq!(parsed.pending, Some(7));
assert_eq!(parsed.lag, Some(3));
}
#[test]
fn parses_xinfo_group_stats_with_nil_lag() {
let parsed = parse_xinfo_group_stats(
RedisValue::Array(vec![RedisValue::Array(vec![
RedisValue::BulkString(b"name".to_vec()),
RedisValue::BulkString(b"usage_consumers".to_vec()),
RedisValue::BulkString(b"consumers".to_vec()),
RedisValue::Int(4),
RedisValue::BulkString(b"pending".to_vec()),
RedisValue::Int(7),
RedisValue::BulkString(b"last-delivered-id".to_vec()),
RedisValue::BulkString(b"1710000000000-0".to_vec()),
RedisValue::BulkString(b"entries-read".to_vec()),
RedisValue::Int(12),
RedisValue::BulkString(b"lag".to_vec()),
RedisValue::Nil,
])]),
"usage_consumers",
)
.expect("xinfo groups should parse nil lag");
assert_eq!(parsed.pending, Some(7));
assert_eq!(parsed.lag, None);
}
#[test]
fn parses_xpending_oldest_idle_ms_from_summary() {
let parsed = parse_xpending_oldest_idle_ms(RedisValue::Array(vec![
RedisValue::Int(7),
RedisValue::BulkString(b"1710000000000-0".to_vec()),
RedisValue::BulkString(b"1710000000006-0".to_vec()),
RedisValue::Array(vec![
RedisValue::Array(vec![
RedisValue::BulkString(b"consumer-a".to_vec()),
RedisValue::Int(4),
RedisValue::Int(1200),
]),
RedisValue::Array(vec![
RedisValue::BulkString(b"consumer-b".to_vec()),
RedisValue::Int(3),
RedisValue::Int(3400),
]),
]),
]))
.expect("xpending should parse");
assert_eq!(parsed, Some(3400));
}
#[test]
fn parses_empty_xpending_as_none() {
let parsed = parse_xpending_oldest_idle_ms(RedisValue::Array(vec![
RedisValue::Int(0),
RedisValue::Nil,
RedisValue::Nil,
RedisValue::Array(vec![]),
]))
.expect("empty xpending should parse");
assert_eq!(parsed, None);
}
#[test]
fn classifies_missing_stream_stats_errors() {
let missing_stream = redis::RedisError::from((
redis::ErrorKind::ResponseError,
"ERR",
"no such key".to_string(),
));
let missing_group = redis::RedisError::from((
redis::ErrorKind::ResponseError,
"ERR",
"NOGROUP No such key 'usage:events' or consumer group 'usage_consumers'".to_string(),
));
let other_error = redis::RedisError::from((
redis::ErrorKind::ResponseError,
"ERR",
"wrong type".to_string(),
));
assert!(redis_stream_stats_missing_stream(&missing_stream));
assert!(redis_stream_stats_missing_stream_or_group(&missing_stream));
assert!(!redis_stream_stats_missing_stream(&missing_group));
assert!(redis_stream_stats_missing_stream_or_group(&missing_group));
assert!(!redis_stream_stats_missing_stream(&other_error));
assert!(!redis_stream_stats_missing_stream_or_group(&other_error));
}
}
+213 -4
View File
@@ -1,4 +1,6 @@
use std::collections::BTreeMap;
use std::future::Future;
use std::sync::{Arc, Mutex};
use aether_runtime::task::spawn_named;
use tokio::task::{JoinHandle, JoinSet};
@@ -94,6 +96,125 @@ impl TaskDefinition {
}
}
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub struct TaskSupervisorTaskSnapshot {
pub task_name: &'static str,
pub active_tasks: u64,
pub supervised_total: u64,
pub completed_total: u64,
pub panicked_total: u64,
pub aborted_total: u64,
pub cancelled_total: u64,
}
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct TaskSupervisorMetricsSnapshot {
pub active_tasks: u64,
pub supervised_total: u64,
pub completed_total: u64,
pub panicked_total: u64,
pub aborted_total: u64,
pub cancelled_total: u64,
pub tasks: Vec<TaskSupervisorTaskSnapshot>,
}
#[derive(Debug, Clone, Default)]
pub struct TaskSupervisorMetrics {
inner: Arc<Mutex<BTreeMap<&'static str, TaskSupervisorTaskCounters>>>,
}
#[derive(Debug, Clone, Copy, Default)]
struct TaskSupervisorTaskCounters {
active_tasks: u64,
supervised_total: u64,
completed_total: u64,
panicked_total: u64,
aborted_total: u64,
cancelled_total: u64,
}
impl TaskSupervisorMetrics {
pub fn record_supervised(&self, task_name: &'static str) {
self.with_task_counters(task_name, |counters| {
counters.supervised_total = counters.supervised_total.saturating_add(1);
counters.active_tasks = counters.active_tasks.saturating_add(1);
});
}
pub fn record_completed(&self, task_name: &'static str) {
self.with_task_counters(task_name, |counters| {
counters.active_tasks = counters.active_tasks.saturating_sub(1);
counters.completed_total = counters.completed_total.saturating_add(1);
});
}
pub fn record_panicked(&self, task_name: &'static str) {
self.with_task_counters(task_name, |counters| {
counters.active_tasks = counters.active_tasks.saturating_sub(1);
counters.panicked_total = counters.panicked_total.saturating_add(1);
});
}
pub fn record_aborted(&self, task_name: &'static str) {
self.with_task_counters(task_name, |counters| {
counters.active_tasks = counters.active_tasks.saturating_sub(1);
counters.aborted_total = counters.aborted_total.saturating_add(1);
});
}
pub fn record_cancelled(&self, task_name: &'static str) {
self.with_task_counters(task_name, |counters| {
counters.active_tasks = counters.active_tasks.saturating_sub(1);
counters.cancelled_total = counters.cancelled_total.saturating_add(1);
});
}
pub fn snapshot(&self) -> TaskSupervisorMetricsSnapshot {
let Ok(guard) = self.inner.lock() else {
return TaskSupervisorMetricsSnapshot::default();
};
let mut snapshot = TaskSupervisorMetricsSnapshot::default();
for (task_name, counters) in guard.iter() {
snapshot.active_tasks = snapshot.active_tasks.saturating_add(counters.active_tasks);
snapshot.supervised_total = snapshot
.supervised_total
.saturating_add(counters.supervised_total);
snapshot.completed_total = snapshot
.completed_total
.saturating_add(counters.completed_total);
snapshot.panicked_total = snapshot
.panicked_total
.saturating_add(counters.panicked_total);
snapshot.aborted_total = snapshot
.aborted_total
.saturating_add(counters.aborted_total);
snapshot.cancelled_total = snapshot
.cancelled_total
.saturating_add(counters.cancelled_total);
snapshot.tasks.push(TaskSupervisorTaskSnapshot {
task_name,
active_tasks: counters.active_tasks,
supervised_total: counters.supervised_total,
completed_total: counters.completed_total,
panicked_total: counters.panicked_total,
aborted_total: counters.aborted_total,
cancelled_total: counters.cancelled_total,
});
}
snapshot
}
fn with_task_counters(
&self,
task_name: &'static str,
update: impl FnOnce(&mut TaskSupervisorTaskCounters),
) {
if let Ok(mut guard) = self.inner.lock() {
update(guard.entry(task_name).or_default());
}
}
}
#[derive(Debug, Clone)]
pub struct TaskContext<TPayload = serde_json::Value> {
run_id: String,
@@ -146,14 +267,20 @@ impl<TPayload> TaskContext<TPayload> {
pub struct TaskSupervisor {
cancellation_token: CancellationToken,
join_set: JoinSet<()>,
metrics: TaskSupervisorMetrics,
supervised_task_count: usize,
}
impl TaskSupervisor {
pub fn new() -> Self {
Self::with_metrics(TaskSupervisorMetrics::default())
}
pub fn with_metrics(metrics: TaskSupervisorMetrics) -> Self {
Self {
cancellation_token: CancellationToken::new(),
join_set: JoinSet::new(),
metrics,
supervised_task_count: 0,
}
}
@@ -162,22 +289,37 @@ impl TaskSupervisor {
self.cancellation_token.clone()
}
pub fn metrics(&self) -> TaskSupervisorMetrics {
self.metrics.clone()
}
pub fn spawn_named<F>(&mut self, task_name: &'static str, future: F)
where
F: Future<Output = ()> + Send + 'static,
{
self.supervised_task_count = self.supervised_task_count.saturating_add(1);
let cancellation_token = self.cancellation_token.clone();
let metrics = self.metrics.clone();
metrics.record_supervised(task_name);
self.join_set.spawn(async move {
let mut handle = spawn_named(task_name, future);
tokio::select! {
_ = cancellation_token.cancelled() => {
handle.abort();
let _ = handle.await;
metrics.record_cancelled(task_name);
}
result = &mut handle => {
if let Err(error) = result {
warn!(task = task_name, error = ?error, "supervised task failed");
match result {
Ok(()) => metrics.record_completed(task_name),
Err(error) => {
if error.is_panic() {
metrics.record_panicked(task_name);
} else {
metrics.record_aborted(task_name);
}
warn!(task = task_name, error = ?error, "supervised task failed");
}
}
}
}
@@ -187,15 +329,26 @@ impl TaskSupervisor {
pub fn supervise_handle(&mut self, task_name: &'static str, mut handle: JoinHandle<()>) {
self.supervised_task_count = self.supervised_task_count.saturating_add(1);
let cancellation_token = self.cancellation_token.clone();
let metrics = self.metrics.clone();
metrics.record_supervised(task_name);
self.join_set.spawn(async move {
tokio::select! {
_ = cancellation_token.cancelled() => {
handle.abort();
let _ = handle.await;
metrics.record_cancelled(task_name);
}
result = &mut handle => {
if let Err(error) = result {
warn!(task = task_name, error = ?error, "supervised task failed");
match result {
Ok(()) => metrics.record_completed(task_name),
Err(error) => {
if error.is_panic() {
metrics.record_panicked(task_name);
} else {
metrics.record_aborted(task_name);
}
warn!(task = task_name, error = ?error, "supervised task failed");
}
}
}
}
@@ -225,3 +378,59 @@ impl Default for TaskSupervisor {
Self::new()
}
}
#[cfg(test)]
mod tests {
use super::{TaskSupervisor, TaskSupervisorMetrics};
use tokio::sync::oneshot;
#[tokio::test]
async fn supervisor_metrics_record_completion_and_cancellation() {
let metrics = TaskSupervisorMetrics::default();
let (tx, rx) = oneshot::channel::<()>();
let mut supervisor = TaskSupervisor::with_metrics(metrics.clone());
supervisor.spawn_named("test.completed", async {});
supervisor.spawn_named("test.cancelled", async move {
let _ = rx.await;
});
tokio::time::sleep(std::time::Duration::from_millis(10)).await;
let snapshot = metrics.snapshot();
assert_eq!(snapshot.supervised_total, 2);
assert_eq!(snapshot.completed_total, 1);
assert_eq!(snapshot.active_tasks, 1);
supervisor.shutdown().await;
drop(tx);
let snapshot = metrics.snapshot();
assert_eq!(snapshot.supervised_total, 2);
assert_eq!(snapshot.completed_total, 1);
assert_eq!(snapshot.cancelled_total, 1);
assert_eq!(snapshot.active_tasks, 0);
}
#[tokio::test]
async fn supervisor_metrics_record_panics_and_external_aborts() {
let metrics = TaskSupervisorMetrics::default();
let mut supervisor = TaskSupervisor::with_metrics(metrics.clone());
let handle = tokio::spawn(async {});
handle.abort();
supervisor.supervise_handle("test.aborted", handle);
supervisor.spawn_named("test.panicked", async {
panic!("intentional task panic");
});
tokio::time::sleep(std::time::Duration::from_millis(10)).await;
let snapshot = metrics.snapshot();
assert_eq!(snapshot.supervised_total, 2);
assert_eq!(snapshot.aborted_total, 1);
assert_eq!(snapshot.panicked_total, 1);
assert_eq!(snapshot.active_tasks, 0);
supervisor.shutdown().await;
}
}
+2
View File
@@ -22,8 +22,10 @@ futures-util.workspace = true
reqwest.workspace = true
serde.workspace = true
serde_json.workspace = true
sha2.workspace = true
sqlx.workspace = true
libc = "0.2"
sysinfo = "0.32"
tokio.workspace = true
tokio-tungstenite = { version = "0.28", features = ["rustls-tls-webpki-roots"] }
uuid.workspace = true
@@ -386,6 +386,7 @@ fn load_config(
warmup_url: Some(gateway_health_url.to_string()),
method: Method::POST,
headers,
header_sets: Vec::new(),
body: Some(
serde_json::json!({
"model": "gpt-5",
File diff suppressed because it is too large Load Diff
@@ -0,0 +1,756 @@
use std::env;
use std::fs;
use std::path::PathBuf;
use aether_data::repository::auth::CreateStandaloneApiKeyRecord;
use aether_data::repository::wallet::WalletLookupKey;
use aether_data::{
DataBackends, DataLayerConfig, DatabaseDriver, SqlDatabaseConfig, SqlPoolConfig,
};
use aether_data_contracts::repository::global_models::{
CreateAdminGlobalModelRecord, UpdateAdminGlobalModelRecord, UpsertAdminProviderModelRecord,
};
use aether_data_contracts::repository::provider_catalog::{
StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogProvider,
};
use serde_json::json;
use sha2::Digest;
const DEFAULT_POSTGRES_URL: &str = "postgresql://postgres:aether@127.0.0.1:5432/aether";
const DEFAULT_OUTPUT_ENV_PATH: &str = "/tmp/aether_local_env.sh";
const DEFAULT_OUTPUT_KEY_PATH: &str = "/tmp/aether_fullchain_api_key";
const DEFAULT_OUTPUT_KEY_LIST_PATH: &str = "/tmp/aether_fullchain_api_keys";
const DEFAULT_PROVIDER_ID: &str = "provider-local-pressure-openai";
const DEFAULT_ENDPOINT_ID: &str = "endpoint-local-pressure-openai-chat";
const DEFAULT_PROVIDER_KEY_ID: &str = "provider-key-local-pressure-openai";
const DEFAULT_GLOBAL_MODEL_ID: &str = "gm-local-pressure-gpt-5-mini";
const DEFAULT_MODEL_ID: &str = "model-local-pressure-gpt-5-mini";
const DEFAULT_API_KEY_ID: &str = "api-key-local-pressure";
const DEFAULT_OPERATOR_ID: &str = "pressure-local";
const DEFAULT_MODEL: &str = "gpt-5-mini";
const DEFAULT_MOCK_UPSTREAM_BASE_URL: &str = "http://127.0.0.1:18181/v1";
const DEFAULT_GATEWAY_BASE_URL: &str = "http://127.0.0.1:8084";
const DEFAULT_API_KEY: &str = "sk-aether-local-pressure";
#[derive(Debug, Clone)]
struct Config {
database_url: String,
output_env_path: PathBuf,
output_key_path: PathBuf,
output_key_list_path: PathBuf,
provider_id: String,
endpoint_id: String,
provider_key_id: String,
global_model_id: String,
model_id: String,
api_key_id: String,
operator_id: String,
model: String,
mock_upstream_base_url: String,
gateway_base_url: String,
api_key: String,
api_key_count: usize,
provider_api_key: String,
postgres_min_connections: u32,
postgres_max_connections: u32,
}
impl Config {
fn from_env_and_args() -> Result<Self, String> {
let mut config = Self {
database_url: env_value("DATABASE_URL")
.or_else(|| env_value("AETHER_DATABASE_URL"))
.or_else(|| env_value("AETHER_GATEWAY_DATA_POSTGRES_URL"))
.unwrap_or_else(|| DEFAULT_POSTGRES_URL.to_string()),
output_env_path: PathBuf::from(
env_value("OUTPUT_ENV_PATH").unwrap_or_else(|| DEFAULT_OUTPUT_ENV_PATH.to_string()),
),
output_key_path: PathBuf::from(
env_value("OUTPUT_KEY_PATH").unwrap_or_else(|| DEFAULT_OUTPUT_KEY_PATH.to_string()),
),
output_key_list_path: PathBuf::from(
env_value("OUTPUT_KEY_LIST_PATH")
.or_else(|| env_value("PRESSURE_API_KEY_LIST_FILE"))
.unwrap_or_else(|| DEFAULT_OUTPUT_KEY_LIST_PATH.to_string()),
),
provider_id: env_value("PRESSURE_PROVIDER_ID")
.unwrap_or_else(|| DEFAULT_PROVIDER_ID.to_string()),
endpoint_id: env_value("PRESSURE_ENDPOINT_ID")
.unwrap_or_else(|| DEFAULT_ENDPOINT_ID.to_string()),
provider_key_id: env_value("PRESSURE_PROVIDER_KEY_ID")
.unwrap_or_else(|| DEFAULT_PROVIDER_KEY_ID.to_string()),
global_model_id: env_value("PRESSURE_GLOBAL_MODEL_ID")
.unwrap_or_else(|| DEFAULT_GLOBAL_MODEL_ID.to_string()),
model_id: env_value("PRESSURE_MODEL_ID")
.unwrap_or_else(|| DEFAULT_MODEL_ID.to_string()),
api_key_id: env_value("PRESSURE_API_KEY_ID")
.unwrap_or_else(|| DEFAULT_API_KEY_ID.to_string()),
operator_id: env_value("PRESSURE_OPERATOR_ID")
.unwrap_or_else(|| DEFAULT_OPERATOR_ID.to_string()),
model: env_value("PRESSURE_MODEL").unwrap_or_else(|| DEFAULT_MODEL.to_string()),
mock_upstream_base_url: env_value("PRESSURE_MOCK_UPSTREAM_BASE_URL")
.unwrap_or_else(|| DEFAULT_MOCK_UPSTREAM_BASE_URL.to_string()),
gateway_base_url: env_value("GATEWAY_BASE_URL")
.unwrap_or_else(|| DEFAULT_GATEWAY_BASE_URL.to_string()),
api_key: env_value("AETHER_API_KEY").unwrap_or_else(|| DEFAULT_API_KEY.to_string()),
api_key_count: env_value("PRESSURE_API_KEY_COUNT")
.and_then(|value| value.parse::<usize>().ok())
.unwrap_or(1),
provider_api_key: env_value("PRESSURE_PROVIDER_API_KEY")
.unwrap_or_else(|| "dummy-local-pressure-provider-key".to_string()),
postgres_min_connections: 1,
postgres_max_connections: 8,
};
let args = env::args().skip(1).collect::<Vec<_>>();
let mut index = 0;
while index < args.len() {
let arg = &args[index];
match arg.as_str() {
"--database-url" => config.database_url = arg_value(&args, &mut index, arg)?,
"--output-env" => {
config.output_env_path = PathBuf::from(arg_value(&args, &mut index, arg)?)
}
"--output-key" => {
config.output_key_path = PathBuf::from(arg_value(&args, &mut index, arg)?)
}
"--output-key-list" => {
config.output_key_list_path = PathBuf::from(arg_value(&args, &mut index, arg)?)
}
"--provider-id" => config.provider_id = arg_value(&args, &mut index, arg)?,
"--endpoint-id" => config.endpoint_id = arg_value(&args, &mut index, arg)?,
"--provider-key-id" => config.provider_key_id = arg_value(&args, &mut index, arg)?,
"--global-model-id" => config.global_model_id = arg_value(&args, &mut index, arg)?,
"--model-id" => config.model_id = arg_value(&args, &mut index, arg)?,
"--api-key-id" => config.api_key_id = arg_value(&args, &mut index, arg)?,
"--operator-id" => config.operator_id = arg_value(&args, &mut index, arg)?,
"--model" => config.model = arg_value(&args, &mut index, arg)?,
"--mock-upstream-base-url" => {
config.mock_upstream_base_url = arg_value(&args, &mut index, arg)?
}
"--gateway-base-url" => {
config.gateway_base_url = arg_value(&args, &mut index, arg)?
}
"--api-key" => config.api_key = arg_value(&args, &mut index, arg)?,
"--api-key-count" => {
config.api_key_count = parse_usize(&arg_value(&args, &mut index, arg)?, arg)?
}
"--provider-api-key" => {
config.provider_api_key = arg_value(&args, &mut index, arg)?
}
"--postgres-min-connections" => {
config.postgres_min_connections =
parse_u32(&arg_value(&args, &mut index, arg)?, arg)?
}
"--postgres-max-connections" => {
config.postgres_max_connections =
parse_u32(&arg_value(&args, &mut index, arg)?, arg)?
}
"--help" | "-h" => {
print_help();
std::process::exit(0);
}
_ => return Err(format!("unknown argument: {arg}")),
}
index += 1;
}
config.validate()?;
Ok(config)
}
fn validate(&self) -> Result<(), String> {
for (name, value) in [
("database_url", &self.database_url),
("provider_id", &self.provider_id),
("endpoint_id", &self.endpoint_id),
("provider_key_id", &self.provider_key_id),
("global_model_id", &self.global_model_id),
("model_id", &self.model_id),
("api_key_id", &self.api_key_id),
("operator_id", &self.operator_id),
("model", &self.model),
("mock_upstream_base_url", &self.mock_upstream_base_url),
("gateway_base_url", &self.gateway_base_url),
("api_key", &self.api_key),
] {
if value.trim().is_empty() {
return Err(format!("{name} cannot be empty"));
}
}
if self.postgres_min_connections > self.postgres_max_connections {
return Err("postgres min connections cannot exceed max connections".to_string());
}
if self.api_key_count == 0 {
return Err("api_key_count must be positive".to_string());
}
Ok(())
}
}
#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
let config = Config::from_env_and_args().map_err(|err| format!("invalid config: {err}"))?;
let backends = DataBackends::from_config(DataLayerConfig::from_database(SqlDatabaseConfig {
driver: DatabaseDriver::Postgres,
url: config.database_url.clone(),
pool: SqlPoolConfig {
min_connections: config.postgres_min_connections,
max_connections: config.postgres_max_connections,
acquire_timeout_ms: 5_000,
idle_timeout_ms: 30_000,
max_lifetime_ms: 300_000,
statement_cache_capacity: 128,
require_ssl: false,
},
}))?;
seed_provider_catalog(&backends, &config).await?;
seed_models(&backends, &config).await?;
let operator_user_id = seed_operator_user(&backends, &config).await?;
seed_api_keys(&backends, &config, &operator_user_id).await?;
verify_candidate_selection(&backends, &config).await?;
write_outputs(&config)?;
println!("gateway pressure seed complete");
println!("provider_id={}", config.provider_id);
println!("endpoint_id={}", config.endpoint_id);
println!("provider_key_id={}", config.provider_key_id);
println!("model={}", config.model);
println!("api_key_id={}", config.api_key_id);
println!("api_key_count={}", config.api_key_count);
println!("env written to {}", config.output_env_path.display());
println!("api key written to {}", config.output_key_path.display());
println!(
"api key list written to {}",
config.output_key_list_path.display()
);
Ok(())
}
async fn seed_provider_catalog(
backends: &DataBackends,
config: &Config,
) -> Result<(), Box<dyn std::error::Error>> {
let reader = backends
.read()
.provider_catalog()
.ok_or("provider catalog reader unavailable")?;
let writer = backends
.write()
.provider_catalog()
.ok_or("provider catalog writer unavailable")?;
let provider = StoredProviderCatalogProvider::new(
config.provider_id.clone(),
"Local pressure OpenAI mock".to_string(),
Some("http://127.0.0.1:18181".to_string()),
"openai".to_string(),
)?
.with_transport_fields(
true,
false,
false,
None,
Some(0),
None,
Some(120.0),
Some(30.0),
None,
)
.with_routing_fields(0)
.with_description(Some(
"Local OpenAI-compatible mock provider for gateway pressure tests".to_string(),
));
if reader
.list_providers_by_ids(std::slice::from_ref(&config.provider_id))
.await?
.is_empty()
{
writer.create_provider(&provider, None).await?;
} else {
writer.update_provider(&provider).await?;
}
let endpoint = StoredProviderCatalogEndpoint::new(
config.endpoint_id.clone(),
config.provider_id.clone(),
"openai:chat".to_string(),
Some("openai".to_string()),
Some("chat_completions".to_string()),
true,
)?
.with_transport_fields(
config.mock_upstream_base_url.clone(),
None,
None,
Some(0),
None,
None,
None,
None,
)?
.with_health_score(1.0);
if reader
.list_endpoints_by_ids(std::slice::from_ref(&config.endpoint_id))
.await?
.is_empty()
{
writer.create_endpoint(&endpoint).await?;
} else {
writer.update_endpoint(&endpoint).await?;
}
let provider_key = StoredProviderCatalogKey::new(
config.provider_key_id.clone(),
config.provider_id.clone(),
"Local pressure mock key".to_string(),
"api_key".to_string(),
Some(json!({"streaming": true})),
true,
)?
.with_transport_fields(
Some(json!(["openai:chat"])),
Some(config.provider_api_key.clone()),
None,
None,
None,
Some(json!([config.model.clone()])),
None,
None,
None,
)?
.with_rate_limit_fields(None, None, None, None, None, None, None, None, None)
.with_health_fields(
Some(json!({"openai:chat": {"status": "healthy"}})),
Some(json!({"openai:chat": {"state": "closed"}})),
);
if reader
.list_keys_by_ids(std::slice::from_ref(&config.provider_key_id))
.await?
.is_empty()
{
writer.create_key(&provider_key).await?;
} else {
writer.update_key(&provider_key).await?;
}
Ok(())
}
async fn seed_models(
backends: &DataBackends,
config: &Config,
) -> Result<(), Box<dyn std::error::Error>> {
let reader = backends
.read()
.global_models()
.ok_or("global model reader unavailable")?;
let writer = backends
.write()
.global_models()
.ok_or("global model writer unavailable")?;
let capabilities = Some(json!({
"streaming": true,
"chat": true
}));
let global_config = Some(json!({
"model_mappings": [config.model],
"pressure_seed": true
}));
if reader
.get_admin_global_model_by_id(&config.global_model_id)
.await?
.is_some()
{
writer
.update_admin_global_model(&UpdateAdminGlobalModelRecord::new(
config.global_model_id.clone(),
config.model.clone(),
true,
Some(0.0),
None,
capabilities.clone(),
global_config.clone(),
)?)
.await?;
} else {
writer
.create_admin_global_model(&CreateAdminGlobalModelRecord::new(
config.global_model_id.clone(),
config.model.clone(),
config.model.clone(),
true,
Some(0.0),
None,
capabilities.clone(),
global_config.clone(),
)?)
.await?;
}
let provider_model_mappings = Some(json!([
{
"name": config.model,
"priority": 0,
"api_formats": ["openai:chat"],
"endpoint_ids": [config.endpoint_id]
}
]));
let provider_model = UpsertAdminProviderModelRecord::new(
config.model_id.clone(),
config.provider_id.clone(),
config.global_model_id.clone(),
config.model.clone(),
provider_model_mappings,
Some(0.0),
None,
Some(false),
Some(false),
Some(true),
Some(false),
Some(false),
true,
true,
Some(json!({"pressure_seed": true})),
)?;
if reader
.list_admin_provider_models_by_global_model_id(&config.global_model_id)
.await?
.iter()
.any(|model| model.id == config.model_id)
{
writer.update_admin_provider_model(&provider_model).await?;
} else {
writer.create_admin_provider_model(&provider_model).await?;
}
Ok(())
}
async fn seed_api_keys(
backends: &DataBackends,
config: &Config,
operator_user_id: &str,
) -> Result<(), Box<dyn std::error::Error>> {
for index in 0..config.api_key_count {
seed_api_key(backends, config, operator_user_id, index).await?;
}
Ok(())
}
async fn seed_api_key(
backends: &DataBackends,
config: &Config,
operator_user_id: &str,
key_index: usize,
) -> Result<(), Box<dyn std::error::Error>> {
let auth_reader = backends
.read()
.auth_api_keys()
.ok_or("auth api key reader unavailable")?;
let auth_writer = backends
.write()
.auth_api_keys()
.ok_or("auth api key writer unavailable")?;
let wallet_reader = backends
.read()
.wallets()
.ok_or("wallet reader unavailable")?;
let api_key_id = pressure_api_key_id(config, key_index);
let api_key_value = pressure_api_key_value(config, key_index);
let existing = auth_reader
.find_export_standalone_api_key_by_id(&api_key_id)
.await?;
if existing.is_none() {
auth_writer
.create_standalone_api_key(CreateStandaloneApiKeyRecord {
user_id: operator_user_id.to_string(),
api_key_id: api_key_id.clone(),
key_hash: sha256_hex(&api_key_value),
key_encrypted: Some(api_key_value),
name: Some(format!("Local pressure API key {}", key_index + 1)),
allowed_providers: Some(vec![config.provider_id.clone()]),
allowed_api_formats: Some(vec!["openai:chat".to_string()]),
allowed_models: Some(vec![config.model.clone()]),
ip_rules: None,
rate_limit: Some(0),
concurrent_limit: None,
force_capabilities: None,
is_active: true,
expires_at_unix_secs: None,
auto_delete_on_expiry: false,
total_requests: 0,
total_tokens: 0,
total_cost_usd: 0.0,
})
.await?;
} else {
auth_writer
.update_standalone_api_key_basic(
aether_data::repository::auth::UpdateStandaloneApiKeyBasicRecord {
api_key_id: api_key_id.clone(),
name: Some(format!("Local pressure API key {}", key_index + 1)),
rate_limit_present: true,
rate_limit: Some(0),
concurrent_limit_present: true,
concurrent_limit: None,
allowed_providers: Some(Some(vec![config.provider_id.clone()])),
allowed_api_formats: Some(Some(vec!["openai:chat".to_string()])),
allowed_models: Some(Some(vec![config.model.clone()])),
ip_rules: Some(None),
expires_at_present: true,
expires_at_unix_secs: None,
auto_delete_on_expiry_present: true,
auto_delete_on_expiry: false,
},
)
.await?;
auth_writer
.set_standalone_api_key_active(&api_key_id, true)
.await?;
}
if wallet_reader
.find(WalletLookupKey::ApiKeyId(&api_key_id))
.await?
.is_none()
{
wallet_reader
.initialize_auth_api_key_wallet(&api_key_id, 0.0, true)
.await?;
} else {
wallet_reader
.update_auth_api_key_wallet_limit_mode(&api_key_id, "unlimited")
.await?;
}
Ok(())
}
async fn seed_operator_user(
backends: &DataBackends,
config: &Config,
) -> Result<String, Box<dyn std::error::Error>> {
let user_reader = backends.read().users().ok_or("user reader unavailable")?;
let wallet_reader = backends
.read()
.wallets()
.ok_or("wallet reader unavailable")?;
let username = format!("{}-user", config.operator_id);
let user = match user_reader.find_user_auth_by_username(&username).await? {
Some(user) => user_reader
.update_local_auth_user_admin_fields(
&user.id,
Some("admin".to_string()),
true,
Some(vec![config.provider_id.clone()]),
true,
Some(vec!["openai:chat".to_string()]),
true,
Some(vec![config.model.clone()]),
true,
None,
Some(true),
)
.await?
.unwrap_or(user),
None => user_reader
.create_local_auth_user_with_settings(
Some(format!("{}@local.pressure", config.operator_id)),
true,
username,
"local-pressure-password-disabled".to_string(),
"admin".to_string(),
Some(vec![config.provider_id.clone()]),
Some(vec!["openai:chat".to_string()]),
Some(vec![config.model.clone()]),
None,
)
.await?
.ok_or("failed to create pressure operator user")?,
};
if wallet_reader
.find(WalletLookupKey::UserId(&user.id))
.await?
.is_none()
{
wallet_reader
.initialize_auth_user_wallet(&user.id, 0.0, true)
.await?;
} else {
wallet_reader
.update_auth_user_wallet_limit_mode(&user.id, "unlimited")
.await?;
}
Ok(user.id)
}
async fn verify_candidate_selection(
backends: &DataBackends,
config: &Config,
) -> Result<(), Box<dyn std::error::Error>> {
let reader = backends
.read()
.minimal_candidate_selection()
.ok_or("candidate selection reader unavailable")?;
let rows = reader
.list_for_exact_api_format_and_requested_model("openai:chat", &config.model)
.await?;
let has_pressure_row = rows.iter().any(|row| {
row.provider_id == config.provider_id
&& row.endpoint_id == config.endpoint_id
&& row.key_id == config.provider_key_id
&& row.global_model_id == config.global_model_id
&& row.model_id == config.model_id
&& row.provider_is_active
&& row.endpoint_is_active
&& row.key_is_active
&& row.model_is_active
&& row.model_is_available
});
if !has_pressure_row {
return Err(format!(
"seeded candidate not visible for model {} and openai:chat",
config.model
)
.into());
}
Ok(())
}
fn write_outputs(config: &Config) -> Result<(), Box<dyn std::error::Error>> {
if let Some(parent) = config.output_env_path.parent() {
fs::create_dir_all(parent)?;
}
if let Some(parent) = config.output_key_path.parent() {
fs::create_dir_all(parent)?;
}
if let Some(parent) = config.output_key_list_path.parent() {
fs::create_dir_all(parent)?;
}
fs::write(&config.output_key_path, format!("{}\n", config.api_key))?;
let key_list = (0..config.api_key_count)
.map(|index| pressure_api_key_value(config, index))
.collect::<Vec<_>>()
.join("\n");
fs::write(&config.output_key_list_path, format!("{key_list}\n"))?;
let env_content = format!(
concat!(
"export AETHER_API_KEY_FILE={key_path}\n",
"export AETHER_API_KEY_LIST_FILE={key_list_path}\n",
"export AETHER_API_KEY=$(cat {key_path})\n",
"export GATEWAY_BASE_URL={gateway_base_url}\n",
"export TARGET_URL={gateway_base_url}/v1/chat/completions\n",
"export METRICS_URL={gateway_base_url}/_gateway/metrics\n",
"export PRESSURE_MODEL={model}\n",
"export PRESSURE_MOCK_UPSTREAM_BASE_URL={mock_upstream_base_url}\n",
"export PRESSURE_MOCK_UPSTREAM_METRICS_URL=http://127.0.0.1:18181/metrics\n"
),
key_path = shell_escape(&config.output_key_path.display().to_string()),
key_list_path = shell_escape(&config.output_key_list_path.display().to_string()),
gateway_base_url = shell_escape(config.gateway_base_url.trim_end_matches('/')),
model = shell_escape(&config.model),
mock_upstream_base_url = shell_escape(&config.mock_upstream_base_url),
);
fs::write(&config.output_env_path, env_content)?;
Ok(())
}
fn sha256_hex(value: &str) -> String {
let mut hasher = sha2::Sha256::new();
hasher.update(value.as_bytes());
format!("{:x}", hasher.finalize())
}
fn pressure_api_key_id(config: &Config, index: usize) -> String {
if index == 0 {
config.api_key_id.clone()
} else {
format!("{}-{}", config.api_key_id, index + 1)
}
}
fn pressure_api_key_value(config: &Config, index: usize) -> String {
if index == 0 {
config.api_key.clone()
} else {
format!("{}-{}", config.api_key, index + 1)
}
}
fn shell_escape(value: &str) -> String {
if value
.chars()
.all(|ch| ch.is_ascii_alphanumeric() || matches!(ch, '/' | '.' | ':' | '_' | '-'))
{
return value.to_string();
}
format!("'{}'", value.replace('\'', "'\\''"))
}
fn env_value(name: &str) -> Option<String> {
env::var(name)
.ok()
.map(|value| value.trim().to_string())
.filter(|value| !value.is_empty())
}
fn arg_value(args: &[String], index: &mut usize, name: &str) -> Result<String, String> {
*index += 1;
args.get(*index)
.filter(|value| !value.starts_with("--"))
.cloned()
.ok_or_else(|| format!("{name} requires a value"))
}
fn parse_u32(value: &str, name: &str) -> Result<u32, String> {
value
.parse::<u32>()
.map_err(|_| format!("{name} must be an unsigned integer"))
}
fn parse_usize(value: &str, name: &str) -> Result<usize, String> {
value
.parse::<usize>()
.map_err(|_| format!("{name} must be an unsigned integer"))
}
fn print_help() {
println!(
"Usage: cargo run -p aether-testkit --bin gateway_pressure_seed -- [options]\n\
\n\
Options:\n\
--database-url URL\n\
--output-env PATH\n\
--output-key PATH\n\
--output-key-list PATH\n\
--gateway-base-url URL\n\
--mock-upstream-base-url URL\n\
--model NAME\n\
--api-key VALUE\n\
--api-key-count N\n\
--provider-api-key VALUE\n\
--provider-id ID\n\
--endpoint-id ID\n\
--provider-key-id ID\n\
--global-model-id ID\n\
--model-id ID\n\
--api-key-id ID\n"
);
}
+125 -8
View File
@@ -11,6 +11,8 @@ use tokio::sync::Mutex;
use crate::runtime::{BenchmarkRuntimeSampler, BenchmarkRuntimeSnapshot};
const MAX_ERROR_SAMPLES: usize = 32;
const MAX_STATUS_SAMPLES: usize = 32;
const MAX_STATUS_SAMPLE_BODY_CHARS: usize = 512;
const FIRST_BODY_BACKGROUND_DRAIN_CHUNKS_ENV: &str =
"AETHER_TESTKIT_FIRST_BODY_BACKGROUND_DRAIN_CHUNKS";
const FIRST_BODY_BACKGROUND_DRAIN_MS_ENV: &str = "AETHER_TESTKIT_FIRST_BODY_BACKGROUND_DRAIN_MS";
@@ -34,6 +36,7 @@ pub struct HttpLoadProbeConfig {
pub warmup_url: Option<String>,
pub method: Method,
pub headers: BTreeMap<String, String>,
pub header_sets: Vec<BTreeMap<String, String>>,
pub body: Option<Vec<u8>>,
pub total_requests: usize,
pub concurrency: usize,
@@ -56,6 +59,7 @@ impl Default for HttpLoadProbeConfig {
warmup_url: None,
method: Method::GET,
headers: BTreeMap::new(),
header_sets: Vec::new(),
body: None,
total_requests: 100,
concurrency: 10,
@@ -93,6 +97,11 @@ impl HttpLoadProbeConfig {
if self.client_shards == 0 {
return Err("load probe client_shards must be positive".to_string());
}
for (index, headers) in self.header_sets.iter().enumerate() {
if headers.is_empty() {
return Err(format!("load probe header_sets[{index}] cannot be empty"));
}
}
if self.http1_only && self.http2_prior_knowledge {
return Err(
"load probe cannot enable both http1_only and http2_prior_knowledge".to_string(),
@@ -114,6 +123,15 @@ pub struct HttpLoadProbeErrorSample {
pub source: Option<String>,
}
#[derive(Debug, Clone, serde::Serialize, PartialEq, Eq)]
pub struct HttpLoadProbeStatusSample {
pub request_index: usize,
pub url: String,
pub status: u16,
pub elapsed_ms: u64,
pub body: String,
}
#[derive(Debug, Clone, serde::Serialize, PartialEq, Eq)]
pub struct HttpLoadProbeResult {
pub url: String,
@@ -159,6 +177,8 @@ pub struct HttpLoadProbeResult {
pub error_counts: BTreeMap<String, usize>,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub error_samples: Vec<HttpLoadProbeErrorSample>,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub non_success_status_samples: Vec<HttpLoadProbeStatusSample>,
}
#[derive(Debug, Clone, serde::Serialize, PartialEq, Eq)]
@@ -207,6 +227,8 @@ pub struct MultiUrlHttpLoadProbeResult {
pub error_counts: BTreeMap<String, usize>,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub error_samples: Vec<HttpLoadProbeErrorSample>,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub non_success_status_samples: Vec<HttpLoadProbeStatusSample>,
}
pub async fn run_http_load_probe(
@@ -253,6 +275,7 @@ pub async fn run_http_load_probe(
status_counts: result.status_counts,
error_counts: result.error_counts,
error_samples: result.error_samples,
non_success_status_samples: result.non_success_status_samples,
})
}
@@ -274,7 +297,7 @@ async fn run_http_load_probe_against_urls(
let effective_client_shards = effective_probe_client_shards(config);
let clients = Arc::new(build_probe_clients(config, effective_client_shards)?);
let total_requests = config.total_requests;
let request_headers = build_headers(&config.headers)?;
let request_headers = build_request_header_sets(config)?;
let request_body = config.body.clone().map(Arc::new);
let response_mode = config.response_mode;
let first_body_hold = config.first_body_hold;
@@ -290,6 +313,7 @@ async fn run_http_load_probe_against_urls(
let status_counts = Arc::new(Mutex::new(BTreeMap::<u16, usize>::new()));
let error_counts = Arc::new(Mutex::new(BTreeMap::<String, usize>::new()));
let error_samples = Arc::new(Mutex::new(Vec::<HttpLoadProbeErrorSample>::new()));
let non_success_status_samples = Arc::new(Mutex::new(Vec::<HttpLoadProbeStatusSample>::new()));
let target_request_counts = Arc::new(Mutex::new(BTreeMap::<String, usize>::new()));
let failed_requests = Arc::new(AtomicUsize::new(0));
let completed_requests = Arc::new(AtomicUsize::new(0));
@@ -304,12 +328,13 @@ async fn run_http_load_probe_against_urls(
let status_counts = Arc::clone(&status_counts);
let error_counts = Arc::clone(&error_counts);
let error_samples = Arc::clone(&error_samples);
let non_success_status_samples = Arc::clone(&non_success_status_samples);
let target_request_counts = Arc::clone(&target_request_counts);
let failed_requests = Arc::clone(&failed_requests);
let completed_requests = Arc::clone(&completed_requests);
let method = config.method.clone();
let urls = urls.to_vec();
let request_headers = request_headers.clone();
let request_headers = Arc::clone(&request_headers);
let request_body = request_body.clone();
let start_delay = worker_start_delay(start_ramp, worker_index, config.concurrency);
@@ -326,7 +351,8 @@ async fn run_http_load_probe_against_urls(
let started_at = Instant::now();
let url = urls[current % urls.len()].clone();
let mut request = client.request(method.clone(), &url);
for (name, value) in request_headers.iter() {
let headers = &request_headers[current % request_headers.len()];
for (name, value) in headers.iter() {
request = request.header(name, value);
}
if let Some(body) = request_body.as_ref() {
@@ -357,6 +383,17 @@ async fn run_http_load_probe_against_urls(
.await;
}
Ok(observation) => {
if !(200..300).contains(&status) {
record_non_success_status_sample(
&non_success_status_samples,
current,
&url,
status,
started_at.elapsed().as_millis() as u64,
observation.body_sample.as_deref().unwrap_or_default(),
)
.await;
}
let mut counts = status_counts.lock().await;
*counts.entry(status).or_insert(0) += 1;
drop(counts);
@@ -404,6 +441,7 @@ async fn run_http_load_probe_against_urls(
let status_counts = status_counts.lock().await.clone();
let error_counts = error_counts.lock().await.clone();
let error_samples = error_samples.lock().await.clone();
let non_success_status_samples = non_success_status_samples.lock().await.clone();
let target_request_counts = target_request_counts.lock().await.clone();
let mut latencies = latencies_ms.lock().await.clone();
let mut header_latencies = header_latencies_ms.lock().await.clone();
@@ -460,6 +498,7 @@ async fn run_http_load_probe_against_urls(
status_counts,
error_counts,
error_samples,
non_success_status_samples,
})
}
@@ -613,6 +652,26 @@ async fn record_load_error(
}
}
async fn record_non_success_status_sample(
non_success_status_samples: &Arc<Mutex<Vec<HttpLoadProbeStatusSample>>>,
request_index: usize,
url: &str,
status: u16,
elapsed_ms: u64,
body: &str,
) {
let mut samples = non_success_status_samples.lock().await;
if samples.len() < MAX_STATUS_SAMPLES {
samples.push(HttpLoadProbeStatusSample {
request_index,
url: url.to_string(),
status,
elapsed_ms,
body: compact_error_text(body, MAX_STATUS_SAMPLE_BODY_CHARS),
});
}
}
fn classify_reqwest_error(phase: &str, err: &reqwest::Error) -> ClassifiedLoadError {
let kind = if err.is_timeout() && err.is_connect() {
"connect_timeout"
@@ -689,17 +748,19 @@ async fn observe_response_body(
)
})?;
let first_body_latency_ms = started_at.elapsed().as_millis() as u64;
drop(first);
let body_sample = compact_bytes_sample(&first, MAX_STATUS_SAMPLE_BODY_CHARS);
if !first_body_hold.is_zero() {
tokio::time::sleep(first_body_hold).await;
}
drain_first_body_response_tail(response).await?;
Ok(BodyObservation {
first_body_latency_ms: Some(first_body_latency_ms),
body_sample: Some(body_sample),
})
}
HttpLoadProbeResponseMode::FullBody => {
let mut first_body_latency_ms = None;
let mut body_sample = Vec::new();
while let Some(chunk) = response
.chunk()
.await
@@ -708,7 +769,7 @@ async fn observe_response_body(
if first_body_latency_ms.is_none() {
first_body_latency_ms = Some(started_at.elapsed().as_millis() as u64);
}
drop(chunk);
append_body_sample(&mut body_sample, &chunk, MAX_STATUS_SAMPLE_BODY_CHARS);
}
if first_body_latency_ms.is_none() {
return Err(ClassifiedLoadError::static_body(
@@ -718,11 +779,24 @@ async fn observe_response_body(
}
Ok(BodyObservation {
first_body_latency_ms,
body_sample: Some(String::from_utf8_lossy(&body_sample).into_owned()),
})
}
}
}
fn compact_bytes_sample(bytes: &[u8], max_chars: usize) -> String {
compact_error_text(String::from_utf8_lossy(bytes), max_chars)
}
fn append_body_sample(target: &mut Vec<u8>, chunk: &[u8], max_chars: usize) {
if target.len() >= max_chars {
return;
}
let remaining = max_chars.saturating_sub(target.len());
target.extend_from_slice(&chunk[..chunk.len().min(remaining)]);
}
async fn drain_first_body_response_tail(
mut response: reqwest::Response,
) -> Result<(), ClassifiedLoadError> {
@@ -773,9 +847,26 @@ fn env_u64(key: &str, default_value: u64) -> u64 {
.unwrap_or(default_value)
}
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
#[derive(Debug, Clone, Default, PartialEq, Eq)]
struct BodyObservation {
first_body_latency_ms: Option<u64>,
body_sample: Option<String>,
}
fn build_request_header_sets(config: &HttpLoadProbeConfig) -> Result<Arc<Vec<HeaderMap>>, String> {
let raw_sets = if config.header_sets.is_empty() {
vec![config.headers.clone()]
} else {
config.header_sets.clone()
};
let mut sets = Vec::with_capacity(raw_sets.len().max(1));
for headers in raw_sets {
sets.push(build_headers(&headers)?);
}
if sets.is_empty() {
sets.push(HeaderMap::new());
}
Ok(Arc::new(sets))
}
fn build_headers(headers: &BTreeMap<String, String>) -> Result<HeaderMap, String> {
@@ -815,8 +906,8 @@ fn percentile(latencies: &[u64], percentile: u8) -> u64 {
#[cfg(test)]
mod tests {
use super::{
build_headers, summarize_latencies, worker_start_delay, HttpLoadProbeConfig,
HttpLoadProbeResponseMode,
build_headers, build_request_header_sets, summarize_latencies, worker_start_delay,
HttpLoadProbeConfig, HttpLoadProbeResponseMode,
};
use reqwest::Method;
use std::collections::BTreeMap;
@@ -886,6 +977,7 @@ mod tests {
assert_eq!(config.method, Method::GET);
assert!(config.warmup_url.is_none());
assert!(config.headers.is_empty());
assert!(config.header_sets.is_empty());
assert!(config.body.is_none());
assert_eq!(config.total_requests, 100);
assert_eq!(config.concurrency, 10);
@@ -916,6 +1008,31 @@ mod tests {
assert!(build_headers(&invalid).is_err());
}
#[test]
fn validates_header_sets() {
let mut config = HttpLoadProbeConfig {
url: "http://127.0.0.1/".to_string(),
header_sets: vec![BTreeMap::new()],
..HttpLoadProbeConfig::default()
};
assert!(config.validate().is_err());
config.header_sets = vec![BTreeMap::from([(
"authorization".to_string(),
"Bearer test".to_string(),
)])];
assert!(config.validate().is_ok());
let sets = build_request_header_sets(&config).expect("header set should build");
assert_eq!(sets.len(), 1);
assert_eq!(
sets[0]
.get("authorization")
.and_then(|value| value.to_str().ok()),
Some("Bearer test")
);
}
#[test]
fn spreads_worker_start_delay_across_ramp() {
assert_eq!(
+88 -7
View File
@@ -85,22 +85,69 @@ fn parse_prometheus_line(line: &str) -> Option<PrometheusSample> {
fn parse_labels(raw: &str) -> BTreeMap<String, String> {
let mut labels = BTreeMap::new();
for pair in raw.split(',').filter(|pair| !pair.is_empty()) {
for pair in split_label_pairs(raw) {
if let Some((key, value)) = pair.split_once('=') {
labels.insert(
key.trim().to_string(),
value
.trim()
.trim_matches('"')
.replace("\\\"", "\"")
.replace("\\n", "\n")
.replace("\\\\", "\\"),
unescape_label_value(value.trim().trim_matches('"')),
);
}
}
labels
}
fn split_label_pairs(raw: &str) -> Vec<&str> {
let mut pairs = Vec::new();
let mut start = 0;
let mut in_string = false;
let mut escaped = false;
for (index, ch) in raw.char_indices() {
if escaped {
escaped = false;
continue;
}
match ch {
'\\' if in_string => escaped = true,
'"' => in_string = !in_string,
',' if !in_string => {
let pair = raw[start..index].trim();
if !pair.is_empty() {
pairs.push(pair);
}
start = index + ch.len_utf8();
}
_ => {}
}
}
let pair = raw[start..].trim();
if !pair.is_empty() {
pairs.push(pair);
}
pairs
}
fn unescape_label_value(value: &str) -> String {
let mut output = String::with_capacity(value.len());
let mut chars = value.chars();
while let Some(ch) = chars.next() {
if ch != '\\' {
output.push(ch);
continue;
}
match chars.next() {
Some('n') => output.push('\n'),
Some('\\') => output.push('\\'),
Some('"') => output.push('"'),
Some(next) => {
output.push('\\');
output.push(next);
}
None => output.push('\\'),
}
}
output
}
#[cfg(test)]
mod tests {
use super::{find_metric_value_u64, parse_prometheus_samples};
@@ -133,4 +180,38 @@ aether_gateway_concurrency_rejected_total{gate="gateway_requests"} 12
Some(12)
);
}
#[test]
fn parses_quoted_label_values_containing_commas() {
let samples = parse_prometheus_samples(
r#"
metric_with_sql{rank="1",query_prefix="SELECT id, name, created_at FROM request_candidates",state="active"} 2
"#,
);
assert_eq!(samples.len(), 1);
assert_eq!(samples[0].labels.get("rank").map(String::as_str), Some("1"));
assert_eq!(
samples[0].labels.get("query_prefix").map(String::as_str),
Some("SELECT id, name, created_at FROM request_candidates")
);
assert_eq!(
samples[0].labels.get("state").map(String::as_str),
Some("active")
);
}
#[test]
fn unescapes_quoted_label_values() {
let samples = parse_prometheus_samples(
r#"
metric_with_escape{message="bad\"line\nx\\y"} 1
"#,
);
assert_eq!(
samples[0].labels.get("message").map(String::as_str),
Some("bad\"line\nx\\y")
);
}
}
+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(
-116
View File
@@ -1,116 +0,0 @@
# Aether Data Schema Inventory
This inventory is the maintenance map for the three SQL drivers. The executable
`sqlx` migrations remain under `crates/aether-data/migrations/{postgres,mysql,sqlite}`.
Do not split already-shipped migration files without also deciding how to handle
existing `_sqlx_migrations` rows.
The maintainable schema source is under `crates/aether-data/schema`. Manifests
there compose back into the executable SQL files and are checked by tests, so
the split source is not just documentation.
The schema directory separates human-maintained sources from generated/runtime
outputs:
| Layer | Purpose |
|---|---|
| `schema/logical/*.toml` | Human-maintained long-term table-structure source. |
| `schema/drivers/{postgres,mysql,sqlite}/**` | Human-maintained executable-SQL fragments while generation is being promoted incrementally. |
| `schema/bootstrap/postgres/**` | Human-maintained source fragments for the Postgres empty-database bootstrap snapshot. |
| `schema/generated/**` | Machine-written SQL from logical schema; checked in for audit and drift detection only. |
| `migrations/**` | Runtime SQL artifacts composed from manifests. |
Generated SQL is not hand-maintained and runtime code does not load it. It
exists to prove that the logical schema can emit driver SQL and to provide the
candidate replacement for handwritten fragments.
## Logical Type Map
| Logical type | Postgres | MySQL | SQLite | Notes |
|---|---|---|---|---|
| `id` | `varchar/text` | `varchar` | `text` | Repository DTOs treat ids as strings. |
| `bool` | `boolean` | `tinyint(1)/boolean` | `integer` | Repositories normalize to Rust `bool`. |
| `time_unix` | `bigint` or legacy `timestamptz` | `bigint` | `integer` | New cross-driver paths prefer unix seconds/ms. |
| `json` | `json/jsonb` | `text/json-compatible` | `text` | Application parses through `serde_json::Value`. |
| `decimal_money` | `numeric` or `double precision` legacy | `double` | `real` | Wallet precision should be reviewed before new money tables. |
| `blob` | `bytea` | `longblob/blob` | `blob` | Used for compressed body payloads. |
| `enum` | `enum` or `varchar` legacy | `varchar` | `text` | Repository contracts own allowed values. |
## Baseline Source Plan
The current executable SQL files are intentionally kept stable for runtime
compatibility. Their maintainable sources are:
| Driver | Executable SQL | Source manifest |
|---|---|---|
| Postgres baseline | `migrations/postgres/20260403000000_baseline.sql` | `schema/drivers/postgres/baseline/manifest.txt` |
| Postgres empty-database snapshot | `aether-data` build output (`OUT_DIR/empty_database_snapshot.sql`) | `schema/bootstrap/postgres/manifest.txt` |
| MySQL baseline | `migrations/mysql/20260403000000_baseline.sql` | `schema/drivers/mysql/baseline/manifest.txt` |
| SQLite baseline | `migrations/sqlite/20260403000000_baseline.sql` | `schema/drivers/sqlite/baseline/manifest.txt` |
All driver manifests are kept as a small set of numbered SQL fragments. Postgres
uses execution-phase fragments (`001_types_and_tables.sql`,
`002_defaults.sql`, `003_constraints.sql`, `004_indexes.sql`,
`005_foreign_keys.sql`, `006_footer.sql`) so pg_dump ordering remains stable
when composed. MySQL and SQLite use similarly numbered domain fragments. After
editing fragments, run:
```bash
bash crates/aether-data/schema/compose_schema.sh compose
bash crates/aether-data/schema/compose_schema.sh check
```
Use `schema/logical/*.toml` for new table structure first; handwritten driver
fragments remain for executable migration compatibility and generator gaps.
`crates/aether-data/schema/logical/*.toml` is the single-maintenance source for
table structure. `aether-data-schema` renders it to
`schema/generated/{postgres,mysql,sqlite}/baseline`, and
`compose_schema.sh check` verifies that the generated SQL is current and that
required executable SQL tables are represented in logical schema. The generated
directory carries its own machine-generated README and per-file `Do not edit`
headers; changes there should come only from `compose_schema.sh generate`.
## Table Inventory
| Area | Tables | Owner | Generation target |
|---|---|---|---|
| Identity/auth | `users`, `api_keys`, `management_tokens`, `user_preferences`, `user_sessions`, `user_oauth_links` | `repository/users`, `repository/auth`, `repository/management_tokens`, auth modules | Good first candidate for schema manifest/query helper generation. |
| Provider catalog | `providers`, `provider_api_keys`, `provider_endpoints`, `models`, `global_models`, `api_key_provider_mappings`, `provider_usage_tracking` | `repository/provider_catalog`, `repository/global_models`, scheduler read paths | Keep complex selection SQL handwritten; generate basic CRUD only. |
| Auth config | `auth_modules`, `oauth_providers`, `ldap_configs` | `repository/auth_modules`, `repository/oauth_providers`, `repository/users` | Good candidate for generated CRUD. |
| Proxy nodes | `proxy_nodes`, `proxy_node_events` | `repository/proxy_nodes` | Good candidate for generated CRUD plus handwritten heartbeat update. |
| Wallet/billing | `wallets`, `wallet_transactions`, `wallet_daily_usage_ledgers`, `payment_orders`, `payment_callbacks`, `refund_requests`, `redeem_code_batches`, `redeem_codes`, `billing_rules`, `dimension_collectors` | `repository/wallet`, `repository/billing`, `repository/settlement` | Keep settlement/ledger math explicit; generate table definitions and simple reads. |
| Usage/audit | `usage`, `usage_counter_deltas`, `usage_body_blobs`, `usage_http_audits`, `usage_routing_snapshots`, `usage_settlement_snapshots`, `request_candidates`, `audit_logs` | `repository/usage`, `repository/candidates`, `repository/audit` | Keep core write/audit queries handwritten. |
| Runtime tasks | `video_tasks`, `gemini_file_mappings`, `announcements`, `announcement_reads` | `repository/video_tasks`, `repository/gemini_file_mappings`, `repository/announcements` | Good candidate for generated CRUD except polling claim logic. |
| Stats | `stats_*`, `schema_backfills` | backend aggregation modules | Keep aggregation SQL per-driver; generate table/index definitions only. |
| System | `system_configs` | `repository/system` through backend dispatch | Good candidate for generated CRUD. |
## Logical Schema Coverage
Logical schema currently covers the clean baseline table set plus portable
MySQL/SQLite table-creation migrations. Postgres-only historical follow-up
migrations remain driver-specific until their schema is normalized or promoted
as explicit generated/override fragments.
## Maintenance Rules
1. Keep driver-specific SQL inside driver-specific migration/repository files.
2. Use logical type names in docs and future schema manifests, not raw database
type names.
3. Keep `jsonb` only in Postgres migrations/repositories/tests.
4. Prefer generated helpers for simple CRUD first; do not rewrite complex usage,
billing, stats, or candidate-selection queries until contract tests cover the
behavior.
5. When adding a new table, update this inventory and add it to the export domain
plan if it must move across databases.
6. If a baseline fragment changes, run `compose_schema.sh compose` before tests
so the executable SQL artifact is regenerated from the source manifest.
## Performance Notes
- Shared hot counters must use `bigint` and be updated by durable outbox flush
workers, not request transactions.
- `usage_counter_deltas` is append-only until processed; processed rows are
retained briefly for audit/replay and then batch-deleted by maintenance.
- Candidate selection joins stay handwritten but are protected in the gateway by
a short TTL, single-flight cache invalidated by provider/routing writes.
@@ -1,411 +0,0 @@
# Gemini API Endpoint Routing Design
**状态:** implementation design
**最后更新:** 2026-05-18
**目标:** 把 Gemini Developer API 和 Vertex AI 的端点语义在 Aether 内部做成明确、可测试、可审计的一等路由语义,根治 `generativelanguage.googleapis.com``aiplatform.googleapis.com` 混用、批量 embedding 伪成功、provider 能力声明不完整等问题。
---
## 速查结论
Aether 里同一个 `api_format` 只描述请求/响应数据形态,不等于实际 Google 后端产品面。
| Aether 语义 | 默认后端产品面 | 官方 host | 主要认证形态 | 说明 |
| --- | --- | --- | --- | --- |
| Google / Gemini Developer API | Gemini Developer API, 也就是 AI Studio 这条 Gemini API | `generativelanguage.googleapis.com` | API key | 默认 Gemini provider 应走这里 |
| Vertex AI | Vertex AI Gemini API | `aiplatform.googleapis.com``{region}-aiplatform.googleapis.com` | service account / Vertex API key | `provider_type = vertex_ai` 应走这里 |
Aether 还必须区分 Google 官方的 OpenAI-compatible 表面。它们使用 OpenAI request/response schema,但不等于 native `generateContent` / `embedContent` endpoint
| OpenAI-compatible 表面 | 后端产品面 | 官方 API root | 主要认证形态 | Aether 处理原则 |
| --- | --- | --- | --- | --- |
| Gemini Developer API OpenAI compatibility | Gemini Developer API / AI Studio | `https://generativelanguage.googleapis.com/v1beta/openai` | Gemini API key as Bearer | 只在 provider format 是 `openai:*` 且显式配置该 root 时使用 |
| Vertex AI OpenAI compatibility | Vertex AI / Google Cloud | `https://aiplatform.googleapis.com/v1/projects/{project}/locations/{location}/endpoints/openapi` | Google Cloud access token / service account | 只在显式 OpenAI-compatible endpoint 上使用;不得替代 native Vertex provider 主链 |
端点动作必须按后端产品面区分:
| 能力 | Gemini Developer API | Vertex AI | Aether 处理原则 |
| --- | --- | --- | --- |
| Generate Content | `models/{model}:generateContent` | `projects/{project}/locations/{location}/publishers/google/models/{model}:generateContent` | 两边都支持,但 URL 构造不同 |
| Stream Generate Content | `models/{model}:streamGenerateContent?alt=sse` | `projects/{project}/locations/{location}/publishers/google/models/{model}:streamGenerateContent?alt=sse` | 两边都支持,但 URL 构造不同 |
| Single Embedding | `models/{model}:embedContent` | `projects/{project}/locations/{location}/publishers/google/models/{model}:predict` | Vertex 文本 embedding 使用 Predict contract`instances[]` + `parameters` |
| Batch Embedding | `models/{model}:batchEmbedContents` | 同一个 `:predict`,由 `instances[]` 表达多输入 | Aether 不切到 Developer API;模型自身的批量限制由 Vertex 明确返回 |
工程不变量:
1. 默认 Gemini provider 只能生成 Gemini Developer API URL,不得因为模型名是 Gemini 就走 Vertex。
2. `provider_type = vertex_ai` 或明确的 Vertex auth/host 只能生成 Vertex URL,不得回退到 Gemini Developer API URL。
3. Vertex embedding 请求必须使用 Vertex Predict contract,不得把 Developer API 的 `model/content/requests` body 原样发给 `:predict`
4. 任何“不支持”的情况必须在调度/URL 构造阶段显式暴露为不可用,不能伪成功。
5. Provider 模板、runtime policy、URL builder、conversion policy、测试连接、live DB reconciliation 必须消费同一个语义模型。
6. Google 官方 OpenAI-compatible root 已经包含 API rootAether 不得额外拼接 `/v1`,否则会生成 `.../openai/v1/...``.../endpoints/openapi/v1/...` 这类错误 URL。
7. Native Gemini endpoint 与 Google OpenAI-compatible endpoint 不是互相 fallback 的关系。显式配置 `openai:*` 才能走 OpenAI-compatible;显式配置 `gemini:*` 才能走 native Gemini REST。
---
## 官方资料依据
本节只记录影响工程设计的官方事实。实现前必须以这些来源为真源,而不是以旧代码行为为真源。
### Gemini Developer API / AI Studio
官方 Gemini API 文档把 Developer API 作为可直接用 API key 调用的产品面。其 REST API host 是 `generativelanguage.googleapis.com`,常见路径是 `/v1beta/models/{model}:...`
关键资料:
- Gemini API reference: <https://ai.google.dev/api>
- Gemini API Generate Content: <https://ai.google.dev/api/generate-content>
- Gemini API Embeddings guide: <https://ai.google.dev/gemini-api/docs/embeddings>
- Gemini API embeddings reference: <https://ai.google.dev/api/embeddings>
- Gemini API OpenAI compatibility: <https://ai.google.dev/gemini-api/docs/openai>
- Gemini API migrate to cloud / Vertex AI: <https://ai.google.dev/gemini-api/docs/migrate-to-cloud>
工程含义:
- `generateContent``streamGenerateContent` 可以走 Developer API host。
- `embedContent` 是单条 embedding。
- `batchEmbedContents` 是 Developer API 的批量 embedding 方法;批量 body 形态是顶层 `requests[]`,每项包含 `model``content`
- Developer API key 不应被拼进 pathAether URL builder 应继续过滤或独立处理 `key` query,避免 query 重复或泄露。
- Developer API 的 OpenAI-compatible root 是 `/v1beta/openai`,其 chat / embedding path 是 `/chat/completions``/embeddings`,不是 `/v1/chat/completions``/v1/embeddings`
### Vertex AI Gemini API
Vertex AI 的 Gemini API REST reference 使用 `aiplatform.googleapis.com` 或 region host,路径包含 GCP project 与 location。
关键资料:
- Vertex AI Generate Content REST: <https://docs.cloud.google.com/vertex-ai/generative-ai/docs/reference/rest/v1/projects.locations.publishers.models/generateContent>
- Vertex AI Stream Generate Content REST: <https://docs.cloud.google.com/vertex-ai/generative-ai/docs/reference/rest/v1/projects.locations.publishers.models/streamGenerateContent>
- Vertex AI Embed Content REST: <https://docs.cloud.google.com/vertex-ai/generative-ai/docs/reference/rest/v1/projects.locations.publishers.models/embedContent>
- Vertex AI Predict REST: <https://docs.cloud.google.com/vertex-ai/generative-ai/docs/reference/rest/v1/projects.locations.publishers.models/predict>
- Vertex AI REST resources: <https://docs.cloud.google.com/vertex-ai/generative-ai/docs/reference/rest/v1/projects.locations.publishers.models>
- Vertex AI Model Garden publisher model list: <https://docs.cloud.google.com/vertex-ai/docs/reference/rest/v1beta1/publishers.models/list>
- Vertex AI text embeddings API: <https://cloud.google.com/vertex-ai/generative-ai/docs/model-reference/text-embeddings-api>
- Vertex AI OpenAI compatibility: <https://cloud.google.com/vertex-ai/generative-ai/docs/start/openai>
工程含义:
- Vertex service account 路径必须包含 project 和 location
- `https://{region}-aiplatform.googleapis.com/v1/projects/{project}/locations/{region}/publishers/google/models/{model}:{action}`
- 对 `global` location,可使用 `https://aiplatform.googleapis.com/v1/projects/{project}/locations/global/...`
- Vertex API key 路径可走:
- `https://aiplatform.googleapis.com/v1/publishers/google/models/{model}:{action}?key=...`
- Vertex 模型目录拉取不是推理请求,必须走 Model Garden publisher list
- `https://aiplatform.googleapis.com/v1beta1/publishers/{publisher}/models`
- 不得使用 `projects/{project}/locations/{location}/publishers/{publisher}/models``projects.locations.publishers.models` 资源没有 list 方法,只有 generate / stream / predict / embed 等动作。
- Vertex 文本 embedding API 文档使用 `:predict`,请求体是 `instances[]`,可选参数在 `parameters` 下;响应是 `predictions[].embeddings.values`
- Vertex REST reference 也列出 `embedContent`,但 Aether 当前 text embedding 主链使用 text embeddings guide 和 Predict API 的 contract。
- Vertex `instances[]` 是在线 Predict 请求体,不等同于异步 batch prediction job。模型级输入数量限制由 Vertex 返回;Aether 不把超出限制的请求静默改走其他产品面。
- Vertex OpenAI-compatible root 是 `/v1/projects/{project}/locations/{location}/endpoints/openapi`,其 OpenAI path 直接挂在这个 root 之后。
- 自定义 Vertex OpenAI-compatible endpoint 可以使用 service account token 刷新,但只有 base URL 明确落在 `/endpoints/openapi` 时才能启用该 Vertex auth 语义。普通 `aiplatform.googleapis.com` + `openai:*` 不能被误判成 Vertex OpenAI compatibility。
### Google Gen AI SDK 的后端切换语义
官方 SDK 同时支持 Gemini Developer API 与 Vertex AI,但二者需要显式选择后端。SDK 层面的 `vertexai=true` / `GOOGLE_GENAI_USE_VERTEXAI=true` 说明:这不是“同一 URL 自动兼容”的关系,而是同一 SDK 下的两个后端产品面。
关键资料:
- Google Gen AI SDK docs: <https://googleapis.github.io/python-genai/>
- Vertex AI SDK overview: <https://cloud.google.com/vertex-ai/generative-ai/docs/sdks/overview>
- Gemini API migrate to cloud / Vertex AI: <https://ai.google.dev/gemini-api/docs/migrate-to-cloud>
工程含义:
- Aether 也应把“选择 Gemini Developer API 还是 Vertex AI”作为显式路由语义,而不是让 URL builder 通过零散 host 字符串猜测。
- `api_format = gemini:generate_content``api_format = gemini:embedding` 只是数据格式。真正的后端产品面由 provider family / auth / endpoint host 决定。
---
## Aether 当前相关链路
这一节说明在 Aether 内部,哪些对象共同决定一次 Gemini 请求实际打到哪里。
| 层 | 代表文件 | 当前职责 | 设计要求 |
| --- | --- | --- | --- |
| 请求格式转换 | `crates/aether-ai-formats/src/formats/...` | OpenAI / Gemini / Claude 等格式互转 | 只负责 body 形态,不决定 Google 后端产品面 |
| Provider 类型模板 | `crates/aether-provider-transport/src/provider_types.rs` | 固定 provider 默认 endpoint、runtime policy | Vertex 模板必须声明 generate + embedding 能力 |
| Runtime policy | `crates/aether-provider-transport/src/provider_types.rs` 和 provider policy | 判断 provider 是否本地可消费 | Vertex embedding 必须进入支持矩阵 |
| URL builder | `crates/aether-provider-transport/src/request_url/mod.rs` | 把 transport + mapped_model + api_format 转成 upstream URL | 必须按后端产品面构造 URL |
| Vertex helpers | `crates/aether-provider-transport/src/vertex/url.rs` | 构造 Vertex 特有 URL | 必须覆盖 generate / stream / embedding |
| Conversion policy | `crates/aether-provider-transport/src/conversion.rs` | 判定跨格式请求能否走某个 transport | OpenAI embedding -> Gemini embedding 在 Vertex 上必须可判定、可认证、可 URL |
| Gateway 测试连接 | `apps/aether-gateway/src/handlers/public/support/test_connection/route.rs` | 测试 provider endpoint 是否可用 | 不得用过低 token 或伪成功规则误判 Gemini 3 |
| Live provider reconciliation | gateway admin/provider 初始化与 DB | 把固定模板同步到 live DB | 新 endpoint 不应只存在源码里,必须进入 live provider/endpoints |
---
## 目标语义模型
新增或显式固化一个内部概念:`GeminiEndpointFamily`
```rust
enum GeminiEndpointFamily {
DeveloperApi,
VertexAi,
}
```
该概念不一定必须以公开 enum 落地,但所有相关函数必须在行为上遵守同一判定:
| 判定输入 | 结果 | 备注 |
| --- | --- | --- |
| `provider_type == "vertex_ai"` | `VertexAi` | 固定 provider 主判据 |
| endpoint host 看起来是 `aiplatform.googleapis.com``{region}-aiplatform.googleapis.com` | `VertexAi` | 支持自定义 Vertex provider,但不可反客为主覆盖固定 provider |
| Vertex service account auth 可解析 | `VertexAi` | service account 是 Vertex 强语义 |
| Vertex API key query auth 可解析 | `VertexAi` | Vertex API key 仍是 Vertex 后端 |
| 普通 Google/Gemini provider + `generativelanguage.googleapis.com` | `DeveloperApi` | 默认 Gemini API |
禁止规则:
- 不得因为 `api_format``gemini:*` 就默认走 Vertex。
- 不得因为 Vertex 缺少某个 endpoint 就回退到 Developer API。
- 不得在 URL builder 里用“host 像谁就算谁”覆盖固定 provider 的 provider_type。
- 不得在 body converter 里偷偷决定 endpoint familybody converter 只能做数据形态转换。
---
## URL 构造矩阵
### Developer API URL
| Aether api_format | stream | batch | URL 形态 |
| --- | --- | --- | --- |
| `gemini:generate_content` | false | 不适用 | `/v1beta/models/{model}:generateContent` |
| `gemini:generate_content` | true | 不适用 | `/v1beta/models/{model}:streamGenerateContent?alt=sse` |
| `gemini:embedding` | false | false | `/v1beta/models/{model}:embedContent` |
| `gemini:embedding` | false | true | `/v1beta/models/{model}:batchEmbedContents` |
Developer API 的批量 embedding 支持顶层 `requests[]`。Aether 可以继续用 body 检测来决定单条还是批量 URL,但该检测只允许影响 Developer API URL。
### Vertex AI URL
| Aether api_format | stream | batch | URL 形态 |
| --- | --- | --- | --- |
| `gemini:generate_content` | false | 不适用 | `/v1/projects/{project}/locations/{location}/publishers/google/models/{model}:generateContent` |
| `gemini:generate_content` | true | 不适用 | `/v1/projects/{project}/locations/{location}/publishers/google/models/{model}:streamGenerateContent?alt=sse` |
| `gemini:embedding` | false | false | `/v1/projects/{project}/locations/{location}/publishers/google/models/{model}:predict` |
| `gemini:embedding` | false | true | `/v1/projects/{project}/locations/{location}/publishers/google/models/{model}:predict` |
Vertex text embedding 的模型由 URL path 承载,body 不得重复携带顶层 `model` 字段,否则会触发 Vertex `oneof field '_model' is already set` 一类错误。Aether 在 Vertex transport context 下必须把 Gemini Developer API embedding body 转成 Predict body
```json
{
"instances": [
{ "content": "text", "task_type": "RETRIEVAL_QUERY", "title": "optional" }
],
"parameters": {
"outputDimensionality": 768,
"autoTruncate": true
}
}
```
如果输入已经是 Predict bodyAether 只移除重复的顶层 `model`。如果输入仍是 OpenAI body 或无法确定可转换,调度阶段必须显式失败,不得把未转换 body 发到 Vertex native endpoint。
### Google OpenAI-Compatible URL
OpenAI-compatible URL 属于显式 passthrough root,不参与 native Gemini URL builder。
| Aether api_format | Gemini Developer API OpenAI compatibility | Vertex AI OpenAI compatibility | Aether 处理原则 |
| --- | --- | --- | --- |
| `openai:chat` | `/v1beta/openai/chat/completions` | `/v1/projects/{project}/locations/{location}/endpoints/openapi/chat/completions` | root 已含 API 版本,不再补 `/v1` |
| `openai:embedding` | `/v1beta/openai/embeddings` | `/v1/projects/{project}/locations/{location}/endpoints/openapi/embeddings` | root 已含 API 版本,不再补 `/v1` |
这条链路的关键边界:
1. `openai:*` provider format 走 OpenAI-compatible schema,不做 OpenAI -> Gemini native body 转换。
2. `gemini:*` provider format 走 native Gemini schema,不因目标是 Google provider 就切到 OpenAI-compatible endpoint。
3. 如果用户请求 `openai:*`、provider endpoint 是 `gemini:*`Aether 走格式转换后打 native Gemini endpoint。
4. 如果用户请求 `openai:*`、provider endpoint 也是 `openai:*`Aether 保持 OpenAI schema 并打显式 OpenAI-compatible root。
5. 以上两条都是正式主链,不能互相静默顶替。
---
## 请求体转换边界
`aether-ai-formats` 中的 Gemini embedding converter 当前负责:
- 单条 input -> `embedContent` body
- 多条 input -> Developer API `batchEmbedContents` body
- `dimensions` -> `outputDimensionality`
- embedding task -> Gemini `taskType`
设计要求:
1. 该 converter 可以继续生成 Gemini Developer API 的 `embedContent` / `batchEmbedContents` body。
2. Transport 层必须在 Vertex context 下把该 body 转成 Predict body,并在无法转换时 fail closed。
3. 所有 taskType / outputDimensionality 必须保持显式传递;不得默认注入会改变语义的 task 或维度。
4. Vertex Predict 的 `instances[]` 只能表示在线 Predict 请求的一次调用;它不是异步 batch prediction job,也不是 Developer API `batchEmbedContents` 的静默替身。
5. Developer API 单条 embedding body 可以保留 `model`Vertex 单条 embedding 在 gateway transport 语义层必须删除顶层 `model`,因为 Vertex 模型已在 path 中指定。
### 格式转换矩阵
端点族和格式转换是两层语义:
- 端点族决定请求发往 `generativelanguage.googleapis.com` 还是 `aiplatform.googleapis.com`
- 格式转换决定客户端传入的 body 如何变成 provider 所需 body,以及 provider response 如何变回客户端期望 body。
`gemini:generate_content` 在 Developer API 与 Vertex AI 上使用同一 Gemini generate-content body 形态,因此格式转换器不应区分这两个产品面。产品面差异只留给 URL/auth 层处理。
`gemini:embedding` 在格式层仍先表达为 Gemini Developer API 的 embedding body。进入 Vertex transport context 时,transport 层再把它收敛到 Vertex Predict body。这样格式转换器不需要知道认证方式,URL/body transport 也不会把 Developer API body 原样发给 Vertex。
| 客户端格式 | Provider 格式 | Developer API | Vertex AI | 处理要求 |
| --- | --- | --- | --- | --- |
| `openai:chat` | `gemini:generate_content` | 支持 | 支持 | OpenAI chat -> Gemini contents / generationConfig |
| `gemini:generate_content` | `openai:chat` | 支持 | 支持 | Gemini contents -> OpenAI messages |
| `openai:embedding` | `gemini:embedding` 单条 | 支持 | 支持 | OpenAI input string 或单项数组 -> Gemini `embedContent` bodyVertex transport 再转 `instances[]` |
| `openai:embedding` | `gemini:embedding` 多条 | 支持 | 支持于 transport 层 | Developer API -> `batchEmbedContents`Vertex transport -> Predict `instances[]`,模型限制由 Vertex 返回 |
| `gemini:embedding` 单条 | `openai:embedding` | 支持 | 支持 | Gemini `content.parts[].text` -> OpenAI `input` string |
| `gemini:embedding` 批量 | `openai:embedding` | 支持 | 支持 | Gemini `requests[]` -> OpenAI `input[]`Vertex Predict response 同样可转 |
| `gemini:embedding` response | `openai:embedding` response | 支持 | 支持 | Gemini `embedding.values` / `embeddings[].values` / Vertex `predictions[].embeddings.values` -> OpenAI `data[].embedding` |
| `openai:embedding` response | `gemini:embedding` response | 支持 | 支持于格式层 | OpenAI `data[]` -> Gemini single `embedding` 或 batch `embeddings[]` |
| `openai:chat` | `openai:chat` on Google OpenAI-compatible root | 支持 | 支持 | passthrough OpenAI schema,不做 native Gemini 转换 |
| `openai:embedding` | `openai:embedding` on Google OpenAI-compatible root | 支持 | 支持 | passthrough OpenAI schema,不做 native Gemini 转换 |
这张矩阵的关键点:
1. 格式层必须能双向理解 Gemini native embedding request/response 与 OpenAI embedding request/response。
2. Vertex text embedding 使用 Predict contract;多输入由 `instances[]` 表达,不构造不存在的 `:batchEmbedContents`
3. 一旦 provider family 是 Vertex,任何 embedding 请求都不能借格式转换之名回退到 Developer API。
4. 对 OpenAI embedding 单项数组,转换器必须生成 Gemini 单条 body,避免把“单条业务请求”误判成 Vertex batch。
5. Google OpenAI-compatible passthrough 与 OpenAI -> Gemini native conversion 是两条显式路径。管理员通过 provider endpoint format 选择路径,Aether 不得自动“择优”改路。
---
## Provider 能力声明与调度
Vertex provider 的固定模板必须包含:
- `gemini:generate_content`
- `gemini:embedding`
- `claude:messages`,如果当前上游 Vertex Claude 支持仍保留
Runtime policy 必须表达:
- Vertex 能本地消费 Gemini generate content。
- Vertex 能本地消费 Gemini embedding,并在 transport 层生成 Predict URL/body。
- Vertex text embedding 的模型级输入数量限制由 Vertex 返回;Aether 不静默拆分、不静默降级到 Developer API。
- 全局模型名与 Vertex 实际 provider 模型名必须可以分离。例如客户端继续请求全局 `gemini-embedding-2-preview` 时,Vertex provider model 可以映射到官方可用的 `gemini-embedding-2`;调度、key allowed_models、URL builder 必须消费映射后的 provider model,不得拿全局 preview 名直打 Vertex。
调度与 conversion policy 必须表达:
- `openai:embedding -> gemini:embedding` 可以被 Vertex provider 接收,transport 层负责把 Gemini Developer API body 转成 Vertex Predict body。
- 对批量 input,不能生成 Vertex `:batchEmbedContents`,也不能回退到 `generativelanguage.googleapis.com`
- `request_pair_direct_auth` 对 Vertex API key 必须返回 `key` query authservice account auth 由 OAuth refresh path 处理,不能伪造成普通 bearer key。
---
## 测试设计
必须覆盖这些测试面:
1. Developer API generate URL:
- non-stream -> `generativelanguage.googleapis.com/...:generateContent`
- stream -> `...:streamGenerateContent?alt=sse`
2. Developer API embedding URL:
- 单条 body -> `...:embedContent`
- 多条 body -> `...:batchEmbedContents`
3. Vertex generate URL:
- API key auth -> `aiplatform.googleapis.com/v1/publishers/google/models/...`
- service account -> project/location path
4. Vertex embedding URL:
- API key auth -> `...:predict?key=...`
- service account -> project/location `...:predict`
5. Vertex embedding body:
- 单条 `model/content` body -> `instances[]`,且移除顶层 `model`
- body 含顶层 `requests[]` 时 -> `instances[]`
- body 已经是 `instances[]` 时只清理重复 `model`
- 无法映射的 body 在调度/模型测试阶段显式失败
- 不得生成 `generativelanguage.googleapis.com`
- 不得生成 `aiplatform.googleapis.com/...:batchEmbedContents`
6. Provider template:
- Vertex fixed template 包含 `gemini:embedding`
- provider embedding support 矩阵包含 Vertex -> Gemini embedding
7. Conversion:
- OpenAI embedding 可以被转换到 Gemini embedding provider format
- Vertex embedding transport 可通过支持检查
- Vertex embedding execution plan 的 URL 使用 mapped provider modelbody 不含顶层 `model`
- Vertex embedding execution plan 的 body 使用 `instances[]` / `parameters`
8. Gateway test connection:
- Gemini generate content 测试不能强制 `maxOutputTokens = 5`
- Google OpenAI-compatible `openai:chat` 测试不能强制 `max_tokens = 5`,否则 Gemini thinking 模型仍可能只返回 thought token / 空 visible content
- Gemini 3 / thinking 模型返回 HTTP 200 但无 visible content 时必须判失败,不能写成成功
9. Google OpenAI-compatible roots:
- Developer API OpenAI root `.../v1beta/openai``openai:chat` 时生成 `.../chat/completions`,不得生成 `.../openai/v1/chat/completions`
- Developer API OpenAI root `.../v1beta/openai``openai:embedding` 时生成 `.../embeddings`,不得生成 `.../openai/v1/embeddings`
- Vertex OpenAI root `.../endpoints/openapi``openai:chat` / `openai:embedding` 时直接挂对应 OpenAI path
- 自定义 Vertex OpenAI-compatible service account endpoint 必须进入 Vertex auth refresh 上下文;普通 `aiplatform.googleapis.com` + `openai:*` 不得被误认成 Vertex OpenAI-compatible
10. Admin write validation:
- Vertex API key key formats 允许 `gemini:generate_content``gemini:embedding`
- Vertex service account key formats 允许 `claude:messages``gemini:generate_content``gemini:embedding`
测试断言必须检查具体 URL、具体 action、具体 body contract 和具体失败原因,不能只检查 `Some(url)` 或状态码。
---
## Live 迁移与验证
上线后必须做四类验证:
1. 源码测试:
- `cargo test -p aether-provider-transport --lib`
- 必要时补 `cargo test -p aether-ai-formats --lib`
- 必要时补 gateway 相关 test
2. Live DB reconciliation
- `vertex_ai` provider 的 endpoints 中必须出现 `gemini:embedding`
- Google/Gemini provider 的 embedding endpoint 仍指向 Developer API,不被 Vertex 改写
3. Live HTTP smoke
- Developer API embedding 单条可用
- Developer API embedding 批量可用
- Vertex generate content 返回 visible content 才算成功
- Vertex embedding 单输入可用
- Vertex embedding 多输入如果模型拒绝,必须暴露 Vertex 原始失败,不能降级成 Developer API 成功
4. 接入方地址核验:
- astrbot plugin ltm
- codex cli config
- 其它容器中引用 Aether 的配置
接入方默认应使用容器网络内稳定地址:
```text
http://aether-app:8084/v1
```
只有在调用方不在 `edge-stack-aether-internal` 这类 Docker 内网、或需要从宿主机/外网访问时,才使用宿主机映射地址或域名。
---
## 明确不做的事
1. 不把 Vertex embedding 多输入写成隐藏循环。隐藏 fan-out 会改变成本、延迟、断路器行为和重试语义,必须另开设计。
2. 不为了测试通过把 Vertex 请求降级到 Developer API。
3. 不为了让 HTTP 200 看起来成功而接受空 candidate / MAX_TOKENS 无 visible content。
4. 不改前端视觉定制、字体、品牌名、landing page 设计。
5. 不用旧 provider endpoint 继续承担新主链。
6. 不把 Google OpenAI-compatible endpoint 当成 native Gemini endpoint 的隐藏 fallback;它只能通过显式 `openai:*` endpoint 进入。
---
## 施工顺序
1. 固化 endpoint family 判定与 URL helper。
2. 为 Vertex `gemini:embedding` 补齐 provider template、runtime policy、conversion policy。
3. 让 request URL builder 对 Vertex embedding 走 Vertex Predict helper。
4. 让 transport body semantics 对 Vertex embedding 生成 `instances[]` / `parameters`,无法转换时 fail closed。
5. 移除测试连接中对 Gemini generate content 的过低 `maxOutputTokens` 硬编码,防止 Gemini 3 thinking 被预算挤空。
6. 移除公开 test-connection 对 OpenAI-compatible chat 的过低 `max_tokens` 硬编码,防止 Google OpenAI-compatible root 复发同类空输出。
7. 跑 red/green 测试。
8. 部署 live。
9. 校验 live DB provider endpoints 与外部接入方地址。
---
## 后续可选增强
如果 7 天 embedding 重算必须在 Vertex 上高吞吐完成,建议后续单独实现 `VertexEmbeddingFanoutExecutor`
- 输入 OpenAI embedding 数组。
- 按配置分片,每片发单条或有限并发 Vertex `predict`
- 合并为 OpenAI embedding response。
- 将每个子请求的失败、重试、成本、断路器状态独立记录。
这项增强不能混入本次 endpoint 语义修复,否则会扩大风险面。
-279
View File
@@ -1,279 +0,0 @@
# Usage Counter Root Fix
This worktree implements the first production cut of the root fix for
usage-related lock contention and hot-row pressure.
## Problem
The current Postgres write path mixes:
- request facts (`usage`, audit snapshots, settlement state)
- shared counters (`api_keys`, `provider_api_keys`, `global_models`)
- provider quota window JSON updates
- wallet settlement writes
That means a single request can hold a transaction while touching multiple shared rows, and high-frequency traffic can serialize on the same `api_key_id` / `provider_api_key_id` / wallet rows.
Relevant code paths:
- `crates/aether-data/src/repository/usage/postgres/mod.rs`
- `crates/aether-data/src/repository/settlement/postgres.rs`
- `crates/aether-usage-runtime/src/runtime.rs`
## Goal
Remove shared-hot-row updates from the request path without losing correctness.
The request path must become:
1. write immutable request facts
2. write durable delta records
3. commit
All shared counters must be derived later by a worker.
## Target Architecture
### 1. Facts first
Keep `usage` as the source of truth for per-request facts:
- request identity
- status transitions
- billing state
- token / cost / latency payloads
- audit snapshots
### 2. Durable counter outbox
Add a new append-only delta table, for example:
`usage_counter_deltas`
Each row should represent one logical contribution:
- `kind` (`api_key`, `provider_api_key`, `model`, `provider_monthly`,
`proxy_node`, `management_token`, `api_key_last_used`, `window`)
- `target_id`
- delta fields
- request id / revision
- `processed_at`
This table is the bridge between request facts and derived counters.
### 3. Flush worker
Add a background worker that:
1. reads unprocessed deltas with `FOR UPDATE SKIP LOCKED`
2. aggregates them in memory by `(kind, target_id)`
3. applies one UPDATE per target
4. marks the source delta rows processed in the same transaction
This keeps memory useful as a buffer, but never as the only source of truth.
### 4. Read models
Move high-read counters to separate tables or compact materialized read models:
- `api_key_usage_counters`
- `provider_api_key_usage_counters`
- `model_usage_counters`
- `provider_monthly_usage_counters`
- `provider_api_key_window_usage_counters`
Keep the original business tables (`api_keys`, `provider_api_keys`, `global_models`) for configuration and compatibility only.
## Implementation Boundary
This branch should land the fix in phases, but the direction must not change:
1. Postgres gets the durable outbox and counter flush path first, because the
reported production lock wait is on Postgres row locks.
2. MySQL and SQLite keep their current usage write behavior until their smaller
deployment paths are moved to the same contract.
3. Request-path Postgres writes may still write `usage`, audit blobs, routing
snapshots, and settlement pricing snapshots. They must not directly update
shared aggregate rows.
4. Compatibility mirror columns may be updated by the flush worker only, never
by the request transaction.
5. Dashboard/statistics read paths should be migrated after the write pressure
is removed, otherwise we risk mixing a large read refactor with the lock fix.
The first executable migration creates the outbox. The first code patch makes
`upsert_usage_record` enqueue deltas in the same transaction as the usage fact
write, then a background worker batches those deltas with
`FOR UPDATE SKIP LOCKED`. Dedicated counter read-model tables remain a follow-up
after the request-path lock pressure is removed.
## Locking Model
### Keep
- advisory lock per `request_id` for idempotent request transitions
- wallet row lock only inside settlement, where correctness depends on it
### Remove from request path
- direct `UPDATE api_keys`
- direct `UPDATE provider_api_keys`
- direct `UPDATE global_models`
- direct `UPDATE providers.monthly_used_usd`
- direct `FOR UPDATE` on provider quota JSON for request-level usage windows
## Memory Cache Rules
Allowed:
- short TTL snapshot cache for read-only admin UI data
- short TTL + single-flight cache for provider/model candidate selection rows
- worker-side delta aggregation buffer
- per-key last-used-at max tracking before flush
Not allowed:
- using memory as the only accounting source
- using memory as the only settlement source
- depending on a clipped Redis stream as the only record of usage
## Rollout Plan
1. Stop counting `pending` / `streaming` as shared counter contributions.
2. Introduce the delta outbox and worker.
3. Redirect request path to facts + outbox only.
4. Migrate reads to the new counter tables.
5. Decommission synchronous hot-row updates.
6. Move provider quota windows out of request transactions.
## Implemented In This Branch
- `usage_counter_deltas` durable outbox for api key, provider api key, model,
provider monthly, proxy node, management token, and api key last-used counters.
- Postgres usage upsert writes request facts plus outbox rows, not shared counter
rows.
- Postgres settlement enqueues provider monthly usage deltas instead of updating
`providers.monthly_used_usd` in the request transaction.
- Gateway request-adjacent proxy node, management token, and api key last-used
writes enqueue durable deltas and fall back to direct writes only when the
usage writer is unavailable.
- Gateway provider/model candidate selection reads use a 5 second in-memory TTL
cache with per-key single-flight. Provider/routing catalog writes invalidate
this cache alongside provider transport and scheduler affinity caches.
- Gateway maintenance worker flushes deltas in batches and aggregates in memory
inside the worker before applying compatibility counter updates.
- Daily quota lookup index on `(user_entitlement_id, usage_date)` removes the
avoidable aggregate scan in quota settlement checks.
- Hot counter columns are widened to `bigint` where old bootstrap schemas still
used `integer`, preventing long-running counter overflow.
## Integration Pressure Tests
The hotspot benchmarks start a managed local Postgres instance, run the
migrations, seed a single hot target, then monitor `pg_stat_activity` while the
load is running. Use a separate target directory when the main worktree target
lock is not writable.
Usage write path, one hot `api_key` / `provider_api_key` / `global_model`:
```sh
CARGO_TARGET_DIR=/tmp/aether-rootfix-target \
cargo run -p aether-testkit --bin usage_counter_hotspot_baseline -- \
--requests 5000 \
--concurrency 200 \
--flush-interval-ms 50 \
--monitor-interval-ms 20 \
--output /tmp/usage_counter_hotspot_after_5000.json
```
Settlement path, one hot provider monthly counter:
```sh
CARGO_TARGET_DIR=/tmp/aether-rootfix-target \
cargo run -p aether-testkit --bin usage_settlement_hotspot_baseline -- \
--requests 5000 \
--concurrency 200 \
--flush-interval-ms 50 \
--monitor-interval-ms 20 \
--output /tmp/usage_settlement_hotspot_after_5000.json
```
Auxiliary hot counters, one hot proxy node / management token / api key
last-used target:
```sh
CARGO_TARGET_DIR=/tmp/aether-rootfix-target \
cargo run -p aether-testkit --bin usage_aux_counter_hotspot_baseline -- \
--requests 5000 \
--concurrency 200 \
--flush-interval-ms 50 \
--monitor-interval-ms 20 \
--output /tmp/usage_aux_counter_hotspot_after_5000.json
```
Latest local run on this worktree:
- usage hotspot: 5000 requests, 200 concurrency, p95 173 ms, 0 failures,
15000 outbox rows processed, 0 pending rows, 0 `api_keys` /
`provider_api_keys` / `global_models` update waiters.
- settlement hotspot: 5000 requests, 200 concurrency, p95 39 ms, 0 failures,
5000 provider monthly deltas processed, `providers.monthly_used_usd = 5.0`,
0 provider update waiters.
Run the auxiliary counter hotspot after changing outbox schemas or gateway
fallback routing; it should drain all pending outbox rows and report zero
request-path waiters for `proxy_nodes`, `management_tokens`, and `api_keys`.
## Runtime Observability
The same outbox health signals used by the pressure tools are exposed through
admin runtime endpoints:
- `GET /api/admin/system/stats`
- `GET /api/admin/monitoring/system-status`
- `GET /api/admin/stats/performance/providers`
These responses include `usage_counter`:
- `status`: `idle`, `catching_up`, or `backlogged`
- `outbox_pending_rows`
- `outbox_processed_rows`
- `oldest_pending_created_at_unix_secs`
- `oldest_pending_age_secs`
- `latest_processed_at_unix_secs`
- `pending_by_kind`
Operational alerting should page when `status = backlogged`, when pending rows
continue growing across several flush intervals, or when the oldest pending age
stays above one minute. A transient non-zero backlog is acceptable during catch
up bursts.
## Remaining Correctness Locks
Wallet debit settlement and daily quota consumption still use database locks
because they protect money/quota correctness, not derived counters. Removing
those locks safely requires a separate wallet debit ledger/reservation worker:
1. request settlement writes an immutable debit intent keyed by `request_id`
2. a per-wallet worker claims intents with `FOR UPDATE SKIP LOCKED`
3. the worker applies balance changes and writes final settlement snapshots
4. request-facing APIs read `pending/settled/insufficient_quota` from the
settlement snapshot
Do not replace this with memory-only balance caches. A cache may accelerate
read-side availability estimates, but the durable ledger must remain the source
of truth.
## Acceptance Criteria
- request transactions no longer update shared counter rows
- counter updates become batchable and replayable
- lock wait time on `api_keys` / `provider_api_keys` drops sharply under concurrency
- wallet settlement remains correct and isolated
- dashboard/statistics reads stay fast from dedicated read models
## Open Decisions
- table names for counter read models
- whether provider window counters live in Postgres only or are dual-written into Redis for display latency
- flush cadence and batch size defaults
- whether to keep compatibility snapshot columns as low-frequency mirrors
-271
View File
@@ -1,271 +0,0 @@
# Aether Gateway DB pressure testing
目标:验证 6k+ 并发长连接/流式请求时,gateway 不把请求并发线性放大成 DB 连接并发,并确认 DB pool、usage queue、后台维护任务不会成为瓶颈。
## 1. 预设环境
生产/压测环境建议显式设置:
```bash
export AETHER_GATEWAY_DATA_POSTGRES_MAX_CONNECTIONS=80
export AETHER_GATEWAY_DATA_POSTGRES_MIN_CONNECTIONS=12
export AETHER_GATEWAY_MAINTENANCE_POOL_IDLE_RESERVE=8
export AETHER_GATEWAY_USAGE_QUEUE_TERMINAL_EVENTS=true
export AETHER_GATEWAY_USAGE_QUEUE_LIFECYCLE_EVENTS=true
export AETHER_GATEWAY_USAGE_QUEUE_STREAM_MAXLEN=200000
export AETHER_GATEWAY_USAGE_QUEUE_BATCH_SIZE=500
export AETHER_GATEWAY_USAGE_QUEUE_RECLAIM_COUNT=500
# Pool score DB feedback is rate-limited per provider key to avoid one
# synchronous score UPDATE per successful request.
export AETHER_GATEWAY_POOL_SCORE_SUCCESS_FEEDBACK_MIN_INTERVAL_SECS=5
export AETHER_GATEWAY_POOL_SCORE_FAILURE_FEEDBACK_MIN_INTERVAL_SECS=1
# Invalid API keys are cached briefly so repeated bad credentials cannot
# linearly amplify into DB lookups. Set 0 to disable during auth debugging.
export AETHER_GATEWAY_AUTH_CONTEXT_NEGATIVE_CACHE_TTL_SECS=10
```
SQLite 不适合 6k 并发压测;请使用 Postgres/MySQL 和 Redis runtime backend。
## 2. Gateway HTTP 压测
准备请求体:
```bash
cat >/tmp/aether-pressure-request.json <<'JSON'
{"model":"gpt-5-mini","messages":[{"role":"user","content":"ping"}],"stream":true}
JSON
```
运行 6k 并发(高并发结论必须使用 release;debug 构建会把 CPU/调试开销误判成 planning timeout):
```bash
TARGET_URL=http://127.0.0.1:18080/v1/chat/completions \
METRICS_URL=http://127.0.0.1:18080/_gateway/metrics \
PRESSURE_METHOD=POST \
PRESSURE_REQUESTS=60000 \
PRESSURE_CONCURRENCY=6000 \
PRESSURE_TIMEOUT_MS=120000 \
PRESSURE_BODY_FILE=/tmp/aether-pressure-request.json \
AUTH_HEADER='Authorization: Bearer <api-key>' \
EXTRA_HEADERS='Content-Type: application/json' \
PRESSURE_RESPONSE_MODE=full \
PRESSURE_CARGO_PROFILE=release \
OUTPUT=/tmp/aether_gateway_pressure_6k.json \
tools/pressure/run_gateway_6k_pressure.sh
```
如果压测流式长连接,建议让 probe 读完整响应体,否则客户端拿到 headers 后会立刻断开:
```bash
PRESSURE_RESPONSE_MODE=full
```
## 3. 本地 mock upstream
没有可承载 6k 并发的真实上游 key 时,先用 testkit 启一个 OpenAI-compatible mock upstream
```bash
cargo run --release -p aether-testkit --bin mock_openai_upstream -- \
--bind 127.0.0.1:18181 \
--chunks 8 \
--first-byte-delay-ms 0 \
--chunk-delay-ms 20 \
--payload-bytes 32
```
可直接压 mock 网络栈:
```bash
cat >/tmp/aether-mock-request.json <<'JSON'
{"model":"mock-model","messages":[{"role":"user","content":"ping"}],"stream":true}
JSON
cargo run --release -p aether-testkit --bin http_load_probe -- \
--url http://127.0.0.1:18181/v1/chat/completions \
--method POST \
--requests 60000 \
--concurrency 6000 \
--timeout-ms 120000 \
--header 'Content-Type: application/json' \
--body-file /tmp/aether-mock-request.json \
--response-mode full
```
要做 gateway 端到端压测,把一个本地压测 provider/endpoint 指到
`http://127.0.0.1:18181/v1`provider key 用 dummy 值即可;gateway 侧仍需一个
本地 Aether API key,但不再消耗真实上游额度。
报告重点字段:
```json
{
"load": {
"throughput_rps": 0,
"failed_requests": 0,
"error_counts": {},
"p95_ms": 0,
"p99_ms": 0
},
"metrics": {
"db_pool_max_checked_out": 0,
"db_pool_min_idle": 0,
"db_pool_max_usage_basis_points": 0,
"db_pool_pressure_samples": 0,
"gateway_requests_max_rejected_total": 0
}
}
```
### 本地全链路 release 基线(mock upstream
本地全链路压测固定为:release gateway + release mock upstream + 本地 Postgres/Redis + 临时 Aether API key/provider/model。
不要把 6k 长连接理解为 6k DB 连接;目标是 6k 前端连接下 DB pool 维持在几十级。
最近可复现基线:
| 场景 | 结果 | throughput | p50 | p95 | p99 | DB pool | mock in-flight |
| --- | --- | ---: | ---: | ---: | ---: | --- | ---: |
| 1000 req / 100 conc | 1000x 200 / 0 fail | 211 rps | 209ms | 729ms | 1171ms | max checked out 33/48 | - |
| 6000 req / 6000 conc, sync terminal candidate | 6000x 200 / 0 fail, `error_counts={}` | 262 rps | 20198ms | 22496ms | 22748ms | max checked out 48/48, pressure samples 7 | 396 |
| 6000 req / 6000 conc, async candidate queue | 6000x 200 / 0 fail, `error_counts={}` | 367 rps | 13678ms | 16060ms | 16245ms | max checked out 48/48, pressure samples 3 | - |
| 6000 req / 6000 conc, async queue + slot compaction | 6000x 200 / 0 fail, `error_counts={}` | 344 rps | 14613ms | 17143ms | 17315ms | max checked out 48/48, pressure samples 3 | - |
6000/6000 下实际打开约 6k FD,说明客户端长连接链路有效。terminal 模式下同一请求可能产生 2 条 candidate 状态写入;
async queue 会把这部分从前台请求路径移出,slot compaction 会把同 slot/同状态的重复记录合并后再落库。
如需测纯转发极限,可临时把 `AETHER_GATEWAY_REQUEST_CANDIDATE_PERSISTENCE=none` 与 terminal 模式对照。
常用本地命令(不要打印 API key):
```bash
source /tmp/aether_local_env.sh
KEY=$(cat /tmp/aether_fullchain_api_key)
CARGO_TARGET_DIR=/tmp/aether-release-pressure \
TARGET_URL=http://127.0.0.1:8088/v1/chat/completions \
METRICS_URL=http://127.0.0.1:8088/_gateway/metrics \
PRESSURE_METHOD=POST \
PRESSURE_REQUESTS=6000 \
PRESSURE_CONCURRENCY=6000 \
PRESSURE_TIMEOUT_MS=120000 \
PRESSURE_SAMPLE_INTERVAL_MS=250 \
PRESSURE_BODY_FILE=/tmp/aether-mock-request.json \
PRESSURE_RESPONSE_MODE=full \
PRESSURE_CARGO_PROFILE=release \
AUTH_HEADER="Authorization: Bearer ${KEY}" \
EXTRA_HEADERS='Content-Type: application/json' \
OUTPUT=/tmp/aether_gateway_release_6000_6000_full.json \
tools/pressure/run_gateway_6k_pressure.sh
```
```bash
docker compose exec -T -e PGPASSWORD="$DB_PASSWORD" postgres psql -U postgres -d aether -P pager=off -c "
SELECT calls, round(total_exec_time::numeric,1) total_ms,
round(mean_exec_time::numeric,3) mean_ms, rows,
left(regexp_replace(query, '\s+', ' ', 'g'), 280) query
FROM pg_stat_statements
ORDER BY calls DESC
LIMIT 60;"
```
### Request candidate async persistence
`request_candidates` 是当前 6k 全链路下最明显的同步 DB 写入点。生产/压测可以先保留
`AETHER_GATEWAY_REQUEST_CANDIDATE_PERSISTENCE=terminal`,再把 terminal 写入从前台 await 改为异步队列:
```bash
export AETHER_GATEWAY_REQUEST_CANDIDATE_PERSISTENCE=terminal
export AETHER_GATEWAY_REQUEST_CANDIDATE_WRITE_MODE=async
export AETHER_GATEWAY_REQUEST_CANDIDATE_QUEUE_CAPACITY=65536
export AETHER_GATEWAY_REQUEST_CANDIDATE_QUEUE_BATCH_SIZE=512
export AETHER_GATEWAY_REQUEST_CANDIDATE_QUEUE_FLUSH_INTERVAL_MS=50
export AETHER_GATEWAY_REQUEST_CANDIDATE_QUEUE_WORKERS=2
# 队列满时默认 drop trace,保护前台请求;需要强一致审计时可设 sync。
export AETHER_GATEWAY_REQUEST_CANDIDATE_QUEUE_FULL=drop
```
新增 metrics
- `request_candidate_queue_depth`
- `request_candidate_queue_pending_depth`
- `request_candidate_queue_capacity`
- `request_candidate_queue_enqueued_total`
- `request_candidate_queue_dropped_total`
- `request_candidate_queue_flushed_total`
- `request_candidate_queue_flush_failed_total`
- `request_candidate_queue_flush_batches_total`
- `request_candidate_queue_flush_sql_ops_total`
- `request_candidate_queue_compacted_total`
- `request_candidate_queue_sync_fallback_total`
判定:6k/6k 下 `dropped_total=0``flush_failed_total=0`,压测结束后
`queue_depth``pending_depth` 应回到 0。
`flush_sql_ops_total` 应低于 `flushed_total`;差值体现在 `compacted_total`
用于确认 terminal candidate 的重复状态写入已在队列侧合并。
当前本地 6000/6000 合并版观测:
- `enqueued_total=12000`
- `flushed_total=12000`
- `flush_sql_ops_total=6406`
- `compacted_total=5594`
- `dropped_total=0`
- `flush_failed_total=0`
- `request_candidates` 最终 `6000 rows / 6000 request_id`
## 4. 判定标准
优先看这些信号:
- `failed_requests == 0` 或仅包含预期的上游错误。
- `db_pool_max_checked_out` 不应接近 `AETHER_GATEWAY_DATA_POSTGRES_MAX_CONNECTIONS`
- `db_pool_min_idle` 在大部分采样中应高于 idle reserve。
- `db_pool_max_usage_basis_points < 8000` 比较健康;持续 `>9000` 表示 DB pool 或 SQL 写入已接近瓶颈。
- `db_pool_pressure_samples` 可以短暂出现,但不应贯穿压测全程。
- `gateway_requests_max_rejected_total` 不应增长,除非有意测试 admission limit。
## 5. DB 热点写入专项压测
这些 testkit 场景会启动临时 Postgres,适合回归验证 counter/settlement 热点锁竞争:
```bash
cargo run -p aether-testkit --bin usage_counter_hotspot_baseline -- \
--requests 20000 --concurrency 1000 \
--flush-interval-ms 50 --monitor-interval-ms 20 \
--output /tmp/aether_usage_counter_20000_1000.json
cargo run -p aether-testkit --bin usage_settlement_hotspot_baseline -- \
--requests 20000 --concurrency 1000 \
--flush-interval-ms 50 --monitor-interval-ms 20 \
--output /tmp/aether_usage_settlement_20000_1000.json
cargo run -p aether-testkit --bin usage_aux_counter_hotspot_baseline -- \
--requests 20000 --concurrency 1000 \
--flush-interval-ms 50 --monitor-interval-ms 20 \
--output /tmp/aether_usage_aux_counter_20000_1000.json
```
关注:
- `failed_requests`
- `throughput_rps`
- `p95_ms`
- `lock_monitor.max_*_update_waiters`
- `lock_monitor.max_oldest_lock_wait_ms`
定向 update waiter 持续大于 0,说明某类 counter/settlement 仍有热点行锁竞争,需要继续分桶或延迟聚合。
## 6. 本地回归基线
最近一次本地临时 Postgres 回归(`requests=60000, concurrency=6000, max_connections=64`):
| suite | throughput | failed | p95 | 定向 update waiters |
| --- | ---: | ---: | ---: | ---: |
| `usage_settlement_hotspot_baseline` | 5452 rps | 0 | 740ms | usage 10 / wallet 0 / provider 0 |
| `usage_counter_hotspot_baseline` | 1358 rps | 0 | 1865ms | api_key 0 / provider_key 0 / model 0 / provider 0 |
| `usage_aux_counter_hotspot_baseline` | 1899 rps | 0 | 785ms | proxy 0 / management_token 0 / api_key 0 |
这些是 DB 热点写入专项结果,不等价于完整 gateway 端到端 6k 流式压测;完整压测仍需使用第 2 节的 gateway URL、真实 API key、Redis runtime backend 与目标模型请求体。
@@ -0,0 +1,348 @@
import { describe, expect, it } from 'vitest'
import { buildGatewayMetricsSummary } from '../monitoring'
describe('buildGatewayMetricsSummary', () => {
it('parses gateway process resource metrics with namespace prefixes', () => {
const summary = buildGatewayMetricsSummary(`
aether_gateway_service_up{service="aether-gateway"} 1
aether_gateway_gateway_process_cpu_usage_basis_points 1234
aether_gateway_gateway_process_memory_bytes 268435456
aether_gateway_gateway_process_memory_basis_points 250
aether_gateway_gateway_process_threads 64
aether_gateway_gateway_process_open_fds 2048
aether_gateway_gateway_process_fd_limit 500000
aether_gateway_gateway_process_fd_usage_basis_points 40
aether_gateway_gateway_process_socket_fds 1800
aether_gateway_gateway_network_observability_available 1
aether_gateway_gateway_network_interfaces 4
aether_gateway_gateway_network_received_bytes_total 123456789
aether_gateway_gateway_network_transmitted_bytes_total 987654321
aether_gateway_gateway_network_received_packets_total 12345
aether_gateway_gateway_network_transmitted_packets_total 23456
aether_gateway_gateway_network_receive_errors_total 1
aether_gateway_gateway_network_transmit_errors_total 2
aether_gateway_gateway_network_receive_dropped_total 3
aether_gateway_gateway_network_transmit_dropped_total 4
aether_gateway_gateway_tcp_state_observability_available 1
aether_gateway_gateway_host_tcp_connections 2200
aether_gateway_gateway_host_tcp_established_connections 1801
aether_gateway_gateway_host_tcp_listen_connections 12
aether_gateway_gateway_host_tcp_time_wait_connections 210
aether_gateway_gateway_host_tcp_syn_sent_connections 5
aether_gateway_gateway_host_tcp_syn_recv_connections 6
aether_gateway_gateway_host_tcp_close_wait_connections 7
aether_gateway_gateway_process_tcp_connections 1812
aether_gateway_gateway_process_tcp_established_connections 1800
aether_gateway_gateway_process_tcp_listen_connections 2
aether_gateway_gateway_process_tcp_time_wait_connections 3
aether_gateway_gateway_process_tcp_syn_sent_connections 4
aether_gateway_gateway_process_tcp_syn_recv_connections 5
aether_gateway_gateway_process_tcp_close_wait_connections 0
aether_gateway_gateway_allocator_observability_available 1
aether_gateway_gateway_allocator_allocated_bytes 67108864
aether_gateway_gateway_allocator_active_bytes 83886080
aether_gateway_gateway_allocator_resident_bytes 100663296
aether_gateway_gateway_allocator_mapped_bytes 134217728
aether_gateway_gateway_allocator_retained_bytes 33554432
aether_gateway_gateway_allocator_metadata_bytes 4194304
aether_gateway_gateway_allocator_active_to_allocated_basis_points 12500
aether_gateway_gateway_allocator_resident_to_allocated_basis_points 15000
aether_gateway_gateway_background_tasks_active 18
aether_gateway_gateway_background_tasks_supervised_total 18
aether_gateway_gateway_background_tasks_unexpected_exits_total 1
aether_gateway_gateway_background_tasks_completed_total 1
aether_gateway_gateway_background_tasks_panicked_total 0
aether_gateway_gateway_background_tasks_aborted_total 0
aether_gateway_gateway_background_tasks_cancelled_total 2
aether_gateway_gateway_tokio_runtime_observability_available 1
aether_gateway_gateway_tokio_runtime_workers 16
aether_gateway_gateway_tokio_runtime_alive_tasks 123
aether_gateway_gateway_tokio_runtime_global_queue_depth 4
aether_gateway_postgres_observability_available{driver="postgres"} 1
aether_gateway_postgres_observability_unavailable{driver="postgres"} 0
aether_gateway_postgres_active_connections{driver="postgres"} 8
aether_gateway_postgres_idle_connections{driver="postgres"} 12
aether_gateway_postgres_idle_in_transaction_connections{driver="postgres"} 0
aether_gateway_postgres_waiting_connections{driver="postgres"} 1
aether_gateway_postgres_lock_waiting_connections{driver="postgres"} 0
aether_gateway_postgres_oldest_active_query_age_ms{driver="postgres"} 123
aether_gateway_postgres_oldest_transaction_age_ms{driver="postgres"} 456
aether_gateway_postgres_deadlocks_total{driver="postgres"} 2
aether_gateway_postgres_block_read_total{driver="postgres"} 100
aether_gateway_postgres_block_hit_total{driver="postgres"} 9900
aether_gateway_postgres_block_cache_hit_rate_basis_points{driver="postgres"} 9900
aether_gateway_postgres_temp_files_total{driver="postgres"} 3
aether_gateway_postgres_temp_bytes_total{driver="postgres"} 4096
aether_gateway_postgres_xact_commit_total{driver="postgres"} 1000
aether_gateway_postgres_xact_rollback_total{driver="postgres"} 4
aether_gateway_postgres_wal_observability_available{driver="postgres"} 1
aether_gateway_postgres_wal_observability_unavailable{driver="postgres"} 0
aether_gateway_postgres_wal_records_total{driver="postgres"} 5000
aether_gateway_postgres_wal_fpi_total{driver="postgres"} 6
aether_gateway_postgres_wal_bytes_total{driver="postgres"} 1048576
aether_gateway_postgres_wal_buffers_full_total{driver="postgres"} 0
aether_gateway_postgres_wal_write_total{driver="postgres"} 50
aether_gateway_postgres_wal_sync_total{driver="postgres"} 40
aether_gateway_postgres_wal_write_time_ms_total{driver="postgres"} 700
aether_gateway_postgres_wal_sync_time_ms_total{driver="postgres"} 80
aether_gateway_postgres_checkpoint_observability_available{driver="postgres"} 1
aether_gateway_postgres_checkpoint_observability_unavailable{driver="postgres"} 0
aether_gateway_postgres_checkpoints_timed_total{driver="postgres"} 2
aether_gateway_postgres_checkpoints_requested_total{driver="postgres"} 1
aether_gateway_postgres_checkpoint_write_time_ms_total{driver="postgres"} 900
aether_gateway_postgres_checkpoint_sync_time_ms_total{driver="postgres"} 120
aether_gateway_postgres_buffers_checkpoint_total{driver="postgres"} 123
aether_gateway_postgres_buffers_backend_total{driver="postgres"} 45
aether_gateway_postgres_statement_observability_available{driver="postgres"} 1
aether_gateway_postgres_statement_observability_unavailable{driver="postgres"} 0
aether_gateway_postgres_statement_top_calls_total{driver="postgres"} 42
aether_gateway_postgres_statement_top_exec_time_ms_total{driver="postgres"} 2345
aether_gateway_postgres_statement_top_max_mean_exec_time_ms{driver="postgres"} 12
aether_gateway_postgres_statement_top_max_exec_time_ms{driver="postgres"} 345
aether_gateway_postgres_statement_top_shared_blks_read_total{driver="postgres"} 22
aether_gateway_postgres_statement_top_shared_blks_hit_total{driver="postgres"} 900
aether_gateway_postgres_statement_top_temp_blks_total{driver="postgres"} 7
aether_gateway_redis_runtime_enabled{backend="redis"} 1
aether_gateway_redis_runtime_health_unavailable{backend="redis"} 0
aether_gateway_redis_runtime_connected_clients{backend="redis"} 9
aether_gateway_redis_runtime_blocked_clients{backend="redis"} 2
aether_gateway_redis_runtime_total_connections_received{backend="redis"} 40
aether_gateway_redis_runtime_rejected_connections_total{backend="redis"} 0
aether_gateway_redis_runtime_total_commands_processed{backend="redis"} 100
aether_gateway_redis_runtime_instantaneous_ops_per_sec{backend="redis"} 17
aether_gateway_redis_runtime_total_error_replies{backend="redis"} 1
aether_gateway_redis_runtime_expired_keys_total{backend="redis"} 3
aether_gateway_redis_runtime_evicted_keys_total{backend="redis"} 0
aether_gateway_redis_runtime_keyspace_hits_total{backend="redis"} 20
aether_gateway_redis_runtime_keyspace_misses_total{backend="redis"} 5
aether_gateway_redis_runtime_keyspace_hit_rate_basis_points{backend="redis"} 8000
aether_gateway_redis_runtime_used_memory_bytes{backend="redis"} 1048576
aether_gateway_redis_runtime_maxmemory_bytes{backend="redis"} 8388608
aether_gateway_redis_runtime_memory_usage_basis_points{backend="redis"} 1250
aether_gateway_redis_runtime_memory_fragmentation_ratio_basis_points{backend="redis"} 12500
aether_gateway_redis_runtime_lane_command_errors_total{backend="redis",lane="fast"} 1
aether_gateway_redis_runtime_lane_command_errors_total{backend="redis",lane="stream"} 2
aether_gateway_redis_runtime_lane_command_timeouts_total{backend="redis",lane="fast"} 3
aether_gateway_redis_runtime_lane_command_timeouts_total{backend="redis",lane="stream"} 4
aether_gateway_redis_runtime_lane_command_count_total{backend="redis",lane="fast"} 10
aether_gateway_redis_runtime_lane_command_count_total{backend="redis",lane="stream"} 5
aether_gateway_redis_runtime_lane_command_latency_ms_sum{backend="redis",lane="fast"} 70
aether_gateway_redis_runtime_lane_command_latency_ms_sum{backend="redis",lane="stream"} 55
aether_gateway_redis_runtime_lane_command_latency_ms_count{backend="redis",lane="fast"} 10
aether_gateway_redis_runtime_lane_command_latency_ms_count{backend="redis",lane="stream"} 5
aether_gateway_redis_runtime_lane_command_latency_ms_max{backend="redis",lane="fast"} 12
aether_gateway_redis_runtime_lane_command_latency_ms_max{backend="redis",lane="stream"} 23
aether_gateway_redis_runtime_lane_command_latency_ms_max{backend="redis",lane="blocking_stream"} 1001
aether_gateway_redis_runtime_lane_command_latency_ms_bucket{backend="redis",lane="fast",le="1"} 2
aether_gateway_redis_runtime_lane_command_latency_ms_bucket{backend="redis",lane="fast",le="+Inf"} 10
aether_gateway_usage_runtime_enabled 1
aether_gateway_usage_runtime_queue_terminal_events_enabled 1
aether_gateway_usage_runtime_queue_lifecycle_events_enabled 1
aether_gateway_usage_runtime_queue_worker_count 2
aether_gateway_usage_runtime_queue_worker_autoscale_enabled 1
aether_gateway_usage_runtime_queue_worker_active_count 3
aether_gateway_usage_runtime_queue_worker_desired_count 4
aether_gateway_usage_runtime_queue_worker_max_count 8
aether_gateway_usage_runtime_queue_worker_read_batches_total 10
aether_gateway_usage_runtime_queue_worker_read_entries_total 20
aether_gateway_usage_runtime_queue_worker_reclaimed_entries_total 2
aether_gateway_usage_runtime_queue_worker_acked_entries_total 18
aether_gateway_usage_runtime_queue_worker_dead_lettered_entries_total 1
aether_gateway_usage_runtime_queue_worker_process_failures_total 2
aether_gateway_usage_runtime_queue_worker_read_failures_total 3
aether_gateway_usage_runtime_queue_worker_reclaim_failures_total 4
aether_gateway_usage_runtime_terminal_enqueue_failed_total 5
aether_gateway_usage_runtime_lifecycle_enqueue_failed_total 6
aether_gateway_request_candidate_queue_depth 5
aether_gateway_request_candidate_queue_pending_depth 4
aether_gateway_request_candidate_queue_capacity 1024
aether_gateway_request_candidate_queue_enqueued_total 30
aether_gateway_request_candidate_queue_dropped_total 1
aether_gateway_request_candidate_queue_flushed_total 28
aether_gateway_request_candidate_queue_flush_failed_total 2
aether_gateway_request_candidate_queue_flush_batches_total 7
aether_gateway_request_candidate_queue_flush_sql_ops_total 8
aether_gateway_request_candidate_queue_flush_sql_records_total 24
aether_gateway_request_candidate_queue_compacted_total 4
aether_gateway_request_candidate_queue_sync_fallback_total 3
aether_gateway_usage_queue_health_unavailable 0
aether_gateway_usage_queue_enabled{stream="usage:events",group="usage_consumers"} 1
aether_gateway_usage_queue_configured{stream="usage:events",group="usage_consumers"} 1
aether_gateway_usage_queue_stream_length{stream="usage:events",group="usage_consumers"} 12
aether_gateway_usage_queue_group_pending{stream="usage:events",group="usage_consumers"} 2
aether_gateway_usage_queue_group_lag{stream="usage:events",group="usage_consumers"} 3
aether_gateway_usage_queue_oldest_pending_idle_ms{stream="usage:events",group="usage_consumers"} 4500
aether_gateway_usage_queue_dlq_length{stream="usage:events:dlq"} 1
aether_gateway_usage_counter_health_unavailable 0
aether_gateway_usage_counter_outbox_pending_rows 3
aether_gateway_usage_counter_outbox_processed_rows 42
aether_gateway_usage_counter_outbox_oldest_pending_age_seconds 12
aether_gateway_usage_counter_outbox_oldest_pending_created_at_unix_secs 1800000000
aether_gateway_usage_counter_outbox_latest_processed_at_unix_secs 1800000012
aether_gateway_usage_counter_outbox_pending_rows_by_kind{kind="api_key"} 2
aether_gateway_usage_counter_outbox_pending_rows_by_kind{kind="model"} 1
aether_gateway_usage_counter_outbox_flush_batches_total 7
aether_gateway_usage_counter_outbox_flush_rows_claimed_total 11
aether_gateway_usage_counter_outbox_flush_targets_total{kind="api_key"} 3
aether_gateway_usage_counter_outbox_flush_targets_total{kind="model"} 4
aether_gateway_usage_counter_outbox_flush_failed_batches_total 1
aether_gateway_usage_counter_outbox_cleanup_rows_total 5
aether_gateway_usage_counter_outbox_cleanup_failed_batches_total 2
`)
expect(summary.process.processCpuUsageBasisPoints).toBe(1234)
expect(summary.process.processMemoryBytes).toBe(268435456)
expect(summary.process.processMemoryBasisPoints).toBe(250)
expect(summary.process.processThreads).toBe(64)
expect(summary.process.openFds).toBe(2048)
expect(summary.process.fdLimit).toBe(500000)
expect(summary.process.fdUsageBasisPoints).toBe(40)
expect(summary.process.socketFds).toBe(1800)
expect(summary.process.networkAvailable).toBe(true)
expect(summary.process.networkInterfaces).toBe(4)
expect(summary.process.networkReceivedBytesTotal).toBe(123456789)
expect(summary.process.networkTransmittedBytesTotal).toBe(987654321)
expect(summary.process.networkReceiveErrorsTotal).toBe(1)
expect(summary.process.networkTransmitErrorsTotal).toBe(2)
expect(summary.process.networkReceiveDroppedTotal).toBe(3)
expect(summary.process.networkTransmitDroppedTotal).toBe(4)
expect(summary.process.tcpStateAvailable).toBe(true)
expect(summary.process.hostTcpEstablishedConnections).toBe(1801)
expect(summary.process.hostTcpTimeWaitConnections).toBe(210)
expect(summary.process.hostTcpCloseWaitConnections).toBe(7)
expect(summary.process.processTcpConnections).toBe(1812)
expect(summary.process.processTcpEstablishedConnections).toBe(1800)
expect(summary.process.processTcpCloseWaitConnections).toBe(0)
expect(summary.allocator.available).toBe(true)
expect(summary.allocator.allocatedBytes).toBe(67108864)
expect(summary.allocator.activeBytes).toBe(83886080)
expect(summary.allocator.residentBytes).toBe(100663296)
expect(summary.allocator.retainedBytes).toBe(33554432)
expect(summary.allocator.activeToAllocatedBasisPoints).toBe(12500)
expect(summary.allocator.residentToAllocatedBasisPoints).toBe(15000)
expect(summary.backgroundTasks.active).toBe(18)
expect(summary.backgroundTasks.supervisedTotal).toBe(18)
expect(summary.backgroundTasks.unexpectedExitsTotal).toBe(1)
expect(summary.backgroundTasks.completedTotal).toBe(1)
expect(summary.backgroundTasks.cancelledTotal).toBe(2)
expect(summary.tokioRuntime.available).toBe(true)
expect(summary.tokioRuntime.workers).toBe(16)
expect(summary.tokioRuntime.aliveTasks).toBe(123)
expect(summary.tokioRuntime.globalQueueDepth).toBe(4)
expect(summary.postgres.driver).toBe('postgres')
expect(summary.postgres.available).toBe(true)
expect(summary.postgres.unavailable).toBe(false)
expect(summary.postgres.activeConnections).toBe(8)
expect(summary.postgres.idleConnections).toBe(12)
expect(summary.postgres.idleInTransactionConnections).toBe(0)
expect(summary.postgres.waitingConnections).toBe(1)
expect(summary.postgres.lockWaitingConnections).toBe(0)
expect(summary.postgres.oldestActiveQueryAgeMs).toBe(123)
expect(summary.postgres.oldestTransactionAgeMs).toBe(456)
expect(summary.postgres.deadlocksTotal).toBe(2)
expect(summary.postgres.blockReadTotal).toBe(100)
expect(summary.postgres.blockHitTotal).toBe(9900)
expect(summary.postgres.blockCacheHitRateBasisPoints).toBe(9900)
expect(summary.postgres.tempBytesTotal).toBe(4096)
expect(summary.postgres.xactRollbackTotal).toBe(4)
expect(summary.postgres.walAvailable).toBe(true)
expect(summary.postgres.walUnavailable).toBe(false)
expect(summary.postgres.walBytesTotal).toBe(1048576)
expect(summary.postgres.walWriteTimeMsTotal).toBe(700)
expect(summary.postgres.walSyncTimeMsTotal).toBe(80)
expect(summary.postgres.checkpointAvailable).toBe(true)
expect(summary.postgres.checkpointUnavailable).toBe(false)
expect(summary.postgres.checkpointWriteTimeMsTotal).toBe(900)
expect(summary.postgres.checkpointSyncTimeMsTotal).toBe(120)
expect(summary.postgres.buffersCheckpointTotal).toBe(123)
expect(summary.postgres.buffersBackendTotal).toBe(45)
expect(summary.postgres.statementAvailable).toBe(true)
expect(summary.postgres.statementUnavailable).toBe(false)
expect(summary.postgres.statementTopExecTimeMsTotal).toBe(2345)
expect(summary.postgres.statementTopMaxMeanExecTimeMs).toBe(12)
expect(summary.postgres.statementTopMaxExecTimeMs).toBe(345)
expect(summary.postgres.statementTopTempBlksTotal).toBe(7)
expect(summary.redisRuntime.enabled).toBe(true)
expect(summary.redisRuntime.unavailable).toBe(false)
expect(summary.redisRuntime.connectedClients).toBe(9)
expect(summary.redisRuntime.blockedClients).toBe(2)
expect(summary.redisRuntime.totalConnectionsReceived).toBe(40)
expect(summary.redisRuntime.rejectedConnectionsTotal).toBe(0)
expect(summary.redisRuntime.totalCommandsProcessed).toBe(100)
expect(summary.redisRuntime.instantaneousOpsPerSec).toBe(17)
expect(summary.redisRuntime.totalErrorReplies).toBe(1)
expect(summary.redisRuntime.expiredKeysTotal).toBe(3)
expect(summary.redisRuntime.evictedKeysTotal).toBe(0)
expect(summary.redisRuntime.keyspaceHitsTotal).toBe(20)
expect(summary.redisRuntime.keyspaceMissesTotal).toBe(5)
expect(summary.redisRuntime.keyspaceHitRateBasisPoints).toBe(8000)
expect(summary.redisRuntime.usedMemoryBytes).toBe(1048576)
expect(summary.redisRuntime.maxmemoryBytes).toBe(8388608)
expect(summary.redisRuntime.memoryUsageBasisPoints).toBe(1250)
expect(summary.redisRuntime.memoryFragmentationRatioBasisPoints).toBe(12500)
expect(summary.redisRuntime.laneCommandErrorsTotal).toBe(3)
expect(summary.redisRuntime.laneCommandTimeoutsTotal).toBe(7)
expect(summary.redisRuntime.laneCommandCountTotal).toBe(15)
expect(summary.redisRuntime.commandLatencyTotalMs).toBe(125)
expect(summary.redisRuntime.commandLatencyObservationCount).toBe(15)
expect(summary.redisRuntime.commandLatencyMaxMs).toBe(1001)
expect(summary.redisRuntime.nonblockingCommandLatencyMaxMs).toBe(23)
expect(summary.usageRuntime.enabled).toBe(true)
expect(summary.usageRuntime.terminalQueueEnabled).toBe(true)
expect(summary.usageRuntime.lifecycleQueueEnabled).toBe(true)
expect(summary.usageRuntime.workerCount).toBe(2)
expect(summary.usageRuntime.workerAutoscaleEnabled).toBe(true)
expect(summary.usageRuntime.workerActiveCount).toBe(3)
expect(summary.usageRuntime.workerDesiredCount).toBe(4)
expect(summary.usageRuntime.workerMaxCount).toBe(8)
expect(summary.usageRuntime.workerReadBatchesTotal).toBe(10)
expect(summary.usageRuntime.workerReadEntriesTotal).toBe(20)
expect(summary.usageRuntime.workerReclaimedEntriesTotal).toBe(2)
expect(summary.usageRuntime.workerAckedEntriesTotal).toBe(18)
expect(summary.usageRuntime.workerDeadLetteredEntriesTotal).toBe(1)
expect(summary.usageRuntime.workerProcessFailuresTotal).toBe(2)
expect(summary.usageRuntime.workerReadFailuresTotal).toBe(3)
expect(summary.usageRuntime.workerReclaimFailuresTotal).toBe(4)
expect(summary.usageRuntime.terminalEnqueueFailedTotal).toBe(5)
expect(summary.usageRuntime.lifecycleEnqueueFailedTotal).toBe(6)
expect(summary.requestCandidateQueue.depth).toBe(5)
expect(summary.requestCandidateQueue.pendingDepth).toBe(4)
expect(summary.requestCandidateQueue.capacity).toBe(1024)
expect(summary.requestCandidateQueue.enqueuedTotal).toBe(30)
expect(summary.requestCandidateQueue.droppedTotal).toBe(1)
expect(summary.requestCandidateQueue.flushedTotal).toBe(28)
expect(summary.requestCandidateQueue.flushFailedTotal).toBe(2)
expect(summary.requestCandidateQueue.flushBatchesTotal).toBe(7)
expect(summary.requestCandidateQueue.flushSqlOpsTotal).toBe(8)
expect(summary.requestCandidateQueue.flushSqlRecordsTotal).toBe(24)
expect(summary.requestCandidateQueue.compactedTotal).toBe(4)
expect(summary.requestCandidateQueue.syncFallbackTotal).toBe(3)
expect(summary.usageQueue.unavailable).toBe(false)
expect(summary.usageQueue.enabled).toBe(true)
expect(summary.usageQueue.configured).toBe(true)
expect(summary.usageQueue.stream).toBe('usage:events')
expect(summary.usageQueue.group).toBe('usage_consumers')
expect(summary.usageQueue.streamLength).toBe(12)
expect(summary.usageQueue.groupPending).toBe(2)
expect(summary.usageQueue.groupLag).toBe(3)
expect(summary.usageQueue.oldestPendingIdleMs).toBe(4500)
expect(summary.usageQueue.dlqStream).toBe('usage:events:dlq')
expect(summary.usageQueue.dlqLength).toBe(1)
expect(summary.usageCounter.unavailable).toBe(false)
expect(summary.usageCounter.pendingRows).toBe(3)
expect(summary.usageCounter.processedRows).toBe(42)
expect(summary.usageCounter.oldestPendingAgeSeconds).toBe(12)
expect(summary.usageCounter.oldestPendingCreatedAtUnixSecs).toBe(1800000000)
expect(summary.usageCounter.latestProcessedAtUnixSecs).toBe(1800000012)
expect(summary.usageCounter.flushBatchesTotal).toBe(7)
expect(summary.usageCounter.flushRowsClaimedTotal).toBe(11)
expect(summary.usageCounter.flushTargetsTotal).toBe(7)
expect(summary.usageCounter.flushFailedBatchesTotal).toBe(1)
expect(summary.usageCounter.cleanupRowsTotal).toBe(5)
expect(summary.usageCounter.cleanupFailedBatchesTotal).toBe(2)
expect(summary.usageCounter.pendingByKind).toEqual([
{ kind: 'api_key', pendingRows: 2 },
{ kind: 'model', pendingRows: 1 },
])
})
})
+667 -1
View File
@@ -1,8 +1,10 @@
import apiClient from './client'
import {
findMetricSamples,
findMetricValueNumber,
parsePrometheusSamples,
sumMetricValues,
type PrometheusSample,
} from '@/utils/prometheus'
export interface AdminMonitoringSystemStatus {
@@ -116,6 +118,21 @@ export interface GatewayMetricsSummary {
serviceUp: number | null
local: GatewayGateMetrics
distributed: GatewayGateMetrics
candidatePlanning: GatewayGateMetrics
upstreamExecution: GatewayGateMetrics
databasePool: GatewayDatabasePoolMetrics
postgres: GatewayPostgresObservabilityMetrics
redisRuntime: GatewayRedisRuntimeMetrics
process: GatewayProcessResourceMetrics
allocator: GatewayAllocatorMetrics
backgroundTasks: GatewayBackgroundTaskMetrics
tokioRuntime: GatewayTokioRuntimeMetrics
usageRuntime: GatewayUsageRuntimeMetrics
usageQueue: GatewayUsageQueueMetrics
usageCounter: GatewayUsageCounterMetrics
requestCandidateQueue: GatewayRequestCandidateQueueMetrics
upstreamTargets: GatewayUpstreamTargetMetrics
stageLatency: GatewayStageLatencyMetrics
tunnel: {
proxyConnections: number | null
availableProxyConnections: number | null
@@ -138,6 +155,276 @@ export interface GatewayMetricsSummary {
fallbacks: GatewayFallbackMetricSummary[]
}
export interface GatewayDatabasePoolMetrics {
driver: string | null
checkedOut: number | null
idle: number | null
size: number | null
max: number | null
usageBasisPoints: number | null
idleReserve: number | null
underMaintenancePressure: boolean | null
}
export interface GatewayPostgresObservabilityMetrics {
driver: string | null
available: boolean | null
unavailable: boolean | null
activeConnections: number | null
idleConnections: number | null
idleInTransactionConnections: number | null
waitingConnections: number | null
lockWaitingConnections: number | null
oldestActiveQueryAgeMs: number | null
oldestTransactionAgeMs: number | null
deadlocksTotal: number | null
blockReadTotal: number | null
blockHitTotal: number | null
blockCacheHitRateBasisPoints: number | null
tempFilesTotal: number | null
tempBytesTotal: number | null
xactCommitTotal: number | null
xactRollbackTotal: number | null
walAvailable: boolean | null
walUnavailable: boolean | null
walRecordsTotal: number | null
walFpiTotal: number | null
walBytesTotal: number | null
walBuffersFullTotal: number | null
walWriteTotal: number | null
walSyncTotal: number | null
walWriteTimeMsTotal: number | null
walSyncTimeMsTotal: number | null
checkpointAvailable: boolean | null
checkpointUnavailable: boolean | null
checkpointsTimedTotal: number | null
checkpointsRequestedTotal: number | null
checkpointWriteTimeMsTotal: number | null
checkpointSyncTimeMsTotal: number | null
buffersCheckpointTotal: number | null
buffersBackendTotal: number | null
statementAvailable: boolean | null
statementUnavailable: boolean | null
statementTopCallsTotal: number | null
statementTopExecTimeMsTotal: number | null
statementTopMaxMeanExecTimeMs: number | null
statementTopMaxExecTimeMs: number | null
statementTopSharedBlksReadTotal: number | null
statementTopSharedBlksHitTotal: number | null
statementTopTempBlksTotal: number | null
}
export interface GatewayRedisRuntimeMetrics {
enabled: boolean | null
unavailable: boolean | null
connectedClients: number | null
blockedClients: number | null
totalConnectionsReceived: number | null
rejectedConnectionsTotal: number | null
totalCommandsProcessed: number | null
instantaneousOpsPerSec: number | null
totalErrorReplies: number | null
expiredKeysTotal: number | null
evictedKeysTotal: number | null
keyspaceHitsTotal: number | null
keyspaceMissesTotal: number | null
keyspaceHitRateBasisPoints: number | null
usedMemoryBytes: number | null
maxmemoryBytes: number | null
memoryUsageBasisPoints: number | null
memoryFragmentationRatioBasisPoints: number | null
laneCommandErrorsTotal: number
laneCommandTimeoutsTotal: number
laneCommandCountTotal: number
commandLatencyTotalMs: number
commandLatencyObservationCount: number
commandLatencyMaxMs: number | null
nonblockingCommandLatencyMaxMs: number | null
}
export interface GatewayProcessResourceMetrics {
sampledAtUnixSecs: number | null
systemCpuUsageBasisPoints: number | null
processCpuUsageBasisPoints: number | null
systemMemoryTotalBytes: number | null
systemMemoryUsedBytes: number | null
systemMemoryAvailableBytes: number | null
systemMemoryUsageBasisPoints: number | null
processMemoryBytes: number | null
processVirtualMemoryBytes: number | null
processMemoryBasisPoints: number | null
processUptimeSeconds: number | null
processThreads: number | null
openFds: number | null
fdLimit: number | null
fdUsageBasisPoints: number | null
socketFds: number | null
networkAvailable: boolean | null
networkInterfaces: number | null
networkReceivedBytesTotal: number | null
networkTransmittedBytesTotal: number | null
networkReceivedPacketsTotal: number | null
networkTransmittedPacketsTotal: number | null
networkReceiveErrorsTotal: number | null
networkTransmitErrorsTotal: number | null
networkReceiveDroppedTotal: number | null
networkTransmitDroppedTotal: number | null
tcpStateAvailable: boolean | null
hostTcpConnections: number | null
hostTcpEstablishedConnections: number | null
hostTcpListenConnections: number | null
hostTcpTimeWaitConnections: number | null
hostTcpSynSentConnections: number | null
hostTcpSynRecvConnections: number | null
hostTcpCloseWaitConnections: number | null
processTcpConnections: number | null
processTcpEstablishedConnections: number | null
processTcpListenConnections: number | null
processTcpTimeWaitConnections: number | null
processTcpSynSentConnections: number | null
processTcpSynRecvConnections: number | null
processTcpCloseWaitConnections: number | null
}
export interface GatewayAllocatorMetrics {
available: boolean | null
allocatedBytes: number | null
activeBytes: number | null
residentBytes: number | null
mappedBytes: number | null
retainedBytes: number | null
metadataBytes: number | null
activeToAllocatedBasisPoints: number | null
residentToAllocatedBasisPoints: number | null
}
export interface GatewayBackgroundTaskMetrics {
active: number | null
supervisedTotal: number | null
unexpectedExitsTotal: number | null
completedTotal: number | null
panickedTotal: number | null
abortedTotal: number | null
cancelledTotal: number | null
}
export interface GatewayTokioRuntimeMetrics {
available: boolean | null
workers: number | null
aliveTasks: number | null
globalQueueDepth: number | null
}
export interface GatewayUsageRuntimeMetrics {
enabled: boolean | null
terminalQueueEnabled: boolean | null
lifecycleQueueEnabled: boolean | null
workerCount: number | null
workerAutoscaleEnabled: boolean | null
workerActiveCount: number | null
workerDesiredCount: number | null
workerMaxCount: number | null
workerReadBatchesTotal: number | null
workerReadEntriesTotal: number | null
workerReclaimedEntriesTotal: number | null
workerAckedEntriesTotal: number | null
workerDeadLetteredEntriesTotal: number | null
workerProcessFailuresTotal: number | null
workerReadFailuresTotal: number | null
workerReclaimFailuresTotal: number | null
terminalEnqueueInFlight: number | null
terminalEnqueueDeferredTotal: number | null
terminalEnqueueDeferredRetryTotal: number | null
terminalEnqueueFailedTotal: number | null
lifecycleEnqueueInFlight: number | null
lifecycleEnqueueDeferredTotal: number | null
lifecycleEnqueueDeferredDroppedTotal: number | null
lifecycleEnqueueDeferredRetryTotal: number | null
lifecycleEnqueueFailedTotal: number | null
enqueueRetryScheduledTotal: number | null
}
export interface GatewayUsageQueueMetrics {
unavailable: boolean | null
enabled: boolean | null
configured: boolean | null
stream: string | null
group: string | null
streamLength: number | null
groupPending: number | null
groupLag: number | null
oldestPendingIdleMs: number | null
dlqStream: string | null
dlqLength: number | null
}
export interface GatewayUsageCounterKindMetrics {
kind: string
pendingRows: number | null
}
export interface GatewayUsageCounterMetrics {
unavailable: boolean | null
pendingRows: number | null
processedRows: number | null
oldestPendingAgeSeconds: number | null
oldestPendingCreatedAtUnixSecs: number | null
latestProcessedAtUnixSecs: number | null
flushBatchesTotal: number | null
flushRowsClaimedTotal: number | null
flushTargetsTotal: number | null
flushFailedBatchesTotal: number | null
cleanupRowsTotal: number | null
cleanupFailedBatchesTotal: number | null
pendingByKind: GatewayUsageCounterKindMetrics[]
}
export interface GatewayRequestCandidateQueueMetrics {
depth: number | null
pendingDepth: number | null
capacity: number | null
enqueuedTotal: number | null
droppedTotal: number | null
flushedTotal: number | null
flushFailedTotal: number | null
flushBatchesTotal: number | null
flushSqlOpsTotal: number | null
flushSqlRecordsTotal: number | null
compactedTotal: number | null
syncFallbackTotal: number | null
}
export interface GatewayUpstreamTargetRow {
target: string
inFlight: number | null
availablePermits: number | null
highWatermark: number | null
rejectedTotal: number | null
selectedTotal: number | null
saturatedTotal: number | null
}
export interface GatewayUpstreamTargetMetrics {
activeTargets: number | null
limit: number | null
selectedTotal: number
saturatedTotal: number
rejectedTotal: number
rows: GatewayUpstreamTargetRow[]
}
export interface GatewayStageLatencyRow {
stage: string
label: string
count: number | null
avgMs: number | null
maxMs: number | null
}
export interface GatewayStageLatencyMetrics {
rows: GatewayStageLatencyRow[]
}
const FALLBACK_METRICS: Array<{ name: string; label: string }> = [
{ name: 'decision_remote_total', label: '远端决策回退' },
{ name: 'plan_fallback_total', label: 'Plan 回退' },
@@ -146,8 +433,24 @@ const FALLBACK_METRICS: Array<{ name: string; label: string }> = [
{ name: 'local_execution_runtime_miss_total', label: '本地运行时缺失' },
]
const CAPACITY_STAGES: Array<{ stage: string; label: string }> = [
{ stage: 'frontdoor_handler_queue', label: '入口排队' },
{ stage: 'frontdoor_admission', label: '入口准入' },
{ stage: 'candidate_planning_gate_wait', label: '候选规划 Gate' },
{ stage: 'candidate_page_load', label: '候选分页读取' },
{ stage: 'candidate_page_resolve', label: '候选分页解析' },
{ stage: 'upstream_execution_gate_wait', label: '上游执行 Gate' },
{ stage: 'stream_upstream_target_admission', label: 'Target 准入' },
{ stage: 'stream_total', label: '流式总耗时' },
]
function metricBoolean(value: number | null): boolean | null {
if (value == null) return null
return value === 1
}
function buildGateMetrics(
samples: ReturnType<typeof parsePrometheusSamples>,
samples: PrometheusSample[],
gate: string
): GatewayGateMetrics {
return {
@@ -159,6 +462,354 @@ function buildGateMetrics(
}
}
function buildDatabasePoolMetrics(samples: PrometheusSample[]): GatewayDatabasePoolMetrics {
const maxSample = findMetricSamples(samples, 'database_pool_max_connections')[0]
return {
driver: maxSample?.labels.driver ?? null,
checkedOut: findMetricValueNumber(samples, 'database_pool_checked_out_connections'),
idle: findMetricValueNumber(samples, 'database_pool_idle_connections'),
size: findMetricValueNumber(samples, 'database_pool_size_connections'),
max: findMetricValueNumber(samples, 'database_pool_max_connections'),
usageBasisPoints: findMetricValueNumber(samples, 'database_pool_usage_basis_points'),
idleReserve: findMetricValueNumber(samples, 'database_pool_idle_reserve_connections'),
underMaintenancePressure: metricBoolean(findMetricValueNumber(samples, 'database_pool_under_maintenance_pressure')),
}
}
function buildPostgresObservabilityMetrics(samples: PrometheusSample[]): GatewayPostgresObservabilityMetrics {
const availabilitySample = findMetricSamples(samples, 'postgres_observability_available')[0]
return {
driver: availabilitySample?.labels.driver ?? null,
available: metricBoolean(findMetricValueNumber(samples, 'postgres_observability_available')),
unavailable: metricBoolean(findMetricValueNumber(samples, 'postgres_observability_unavailable')),
activeConnections: findMetricValueNumber(samples, 'postgres_active_connections'),
idleConnections: findMetricValueNumber(samples, 'postgres_idle_connections'),
idleInTransactionConnections: findMetricValueNumber(samples, 'postgres_idle_in_transaction_connections'),
waitingConnections: findMetricValueNumber(samples, 'postgres_waiting_connections'),
lockWaitingConnections: findMetricValueNumber(samples, 'postgres_lock_waiting_connections'),
oldestActiveQueryAgeMs: findMetricValueNumber(samples, 'postgres_oldest_active_query_age_ms'),
oldestTransactionAgeMs: findMetricValueNumber(samples, 'postgres_oldest_transaction_age_ms'),
deadlocksTotal: findMetricValueNumber(samples, 'postgres_deadlocks_total'),
blockReadTotal: findMetricValueNumber(samples, 'postgres_block_read_total'),
blockHitTotal: findMetricValueNumber(samples, 'postgres_block_hit_total'),
blockCacheHitRateBasisPoints: findMetricValueNumber(samples, 'postgres_block_cache_hit_rate_basis_points'),
tempFilesTotal: findMetricValueNumber(samples, 'postgres_temp_files_total'),
tempBytesTotal: findMetricValueNumber(samples, 'postgres_temp_bytes_total'),
xactCommitTotal: findMetricValueNumber(samples, 'postgres_xact_commit_total'),
xactRollbackTotal: findMetricValueNumber(samples, 'postgres_xact_rollback_total'),
walAvailable: metricBoolean(findMetricValueNumber(samples, 'postgres_wal_observability_available')),
walUnavailable: metricBoolean(findMetricValueNumber(samples, 'postgres_wal_observability_unavailable')),
walRecordsTotal: findMetricValueNumber(samples, 'postgres_wal_records_total'),
walFpiTotal: findMetricValueNumber(samples, 'postgres_wal_fpi_total'),
walBytesTotal: findMetricValueNumber(samples, 'postgres_wal_bytes_total'),
walBuffersFullTotal: findMetricValueNumber(samples, 'postgres_wal_buffers_full_total'),
walWriteTotal: findMetricValueNumber(samples, 'postgres_wal_write_total'),
walSyncTotal: findMetricValueNumber(samples, 'postgres_wal_sync_total'),
walWriteTimeMsTotal: findMetricValueNumber(samples, 'postgres_wal_write_time_ms_total'),
walSyncTimeMsTotal: findMetricValueNumber(samples, 'postgres_wal_sync_time_ms_total'),
checkpointAvailable: metricBoolean(findMetricValueNumber(samples, 'postgres_checkpoint_observability_available')),
checkpointUnavailable: metricBoolean(findMetricValueNumber(samples, 'postgres_checkpoint_observability_unavailable')),
checkpointsTimedTotal: findMetricValueNumber(samples, 'postgres_checkpoints_timed_total'),
checkpointsRequestedTotal: findMetricValueNumber(samples, 'postgres_checkpoints_requested_total'),
checkpointWriteTimeMsTotal: findMetricValueNumber(samples, 'postgres_checkpoint_write_time_ms_total'),
checkpointSyncTimeMsTotal: findMetricValueNumber(samples, 'postgres_checkpoint_sync_time_ms_total'),
buffersCheckpointTotal: findMetricValueNumber(samples, 'postgres_buffers_checkpoint_total'),
buffersBackendTotal: findMetricValueNumber(samples, 'postgres_buffers_backend_total'),
statementAvailable: metricBoolean(findMetricValueNumber(samples, 'postgres_statement_observability_available')),
statementUnavailable: metricBoolean(findMetricValueNumber(samples, 'postgres_statement_observability_unavailable')),
statementTopCallsTotal: findMetricValueNumber(samples, 'postgres_statement_top_calls_total'),
statementTopExecTimeMsTotal: findMetricValueNumber(samples, 'postgres_statement_top_exec_time_ms_total'),
statementTopMaxMeanExecTimeMs: findMetricValueNumber(samples, 'postgres_statement_top_max_mean_exec_time_ms'),
statementTopMaxExecTimeMs: findMetricValueNumber(samples, 'postgres_statement_top_max_exec_time_ms'),
statementTopSharedBlksReadTotal: findMetricValueNumber(samples, 'postgres_statement_top_shared_blks_read_total'),
statementTopSharedBlksHitTotal: findMetricValueNumber(samples, 'postgres_statement_top_shared_blks_hit_total'),
statementTopTempBlksTotal: findMetricValueNumber(samples, 'postgres_statement_top_temp_blks_total'),
}
}
function buildRedisRuntimeMetrics(samples: PrometheusSample[]): GatewayRedisRuntimeMetrics {
const commandLatencyMaxMs = maxMetricValue(samples, 'redis_runtime_lane_command_latency_ms_max')
const nonblockingLatencyValues = [
findMetricValueNumber(samples, 'redis_runtime_lane_command_latency_ms_max', { lane: 'fast' }),
findMetricValueNumber(samples, 'redis_runtime_lane_command_latency_ms_max', { lane: 'stream' }),
findMetricValueNumber(samples, 'redis_runtime_lane_command_latency_ms_max', { lane: 'admin' }),
].filter((value): value is number => value != null && Number.isFinite(value))
const nonblockingCommandLatencyMaxMs = nonblockingLatencyValues.length > 0
? Math.max(...nonblockingLatencyValues)
: null
return {
enabled: metricBoolean(findMetricValueNumber(samples, 'redis_runtime_enabled')),
unavailable: metricBoolean(findMetricValueNumber(samples, 'redis_runtime_health_unavailable')),
connectedClients: findMetricValueNumber(samples, 'redis_runtime_connected_clients'),
blockedClients: findMetricValueNumber(samples, 'redis_runtime_blocked_clients'),
totalConnectionsReceived: findMetricValueNumber(samples, 'redis_runtime_total_connections_received'),
rejectedConnectionsTotal: findMetricValueNumber(samples, 'redis_runtime_rejected_connections_total'),
totalCommandsProcessed: findMetricValueNumber(samples, 'redis_runtime_total_commands_processed'),
instantaneousOpsPerSec: findMetricValueNumber(samples, 'redis_runtime_instantaneous_ops_per_sec'),
totalErrorReplies: findMetricValueNumber(samples, 'redis_runtime_total_error_replies'),
expiredKeysTotal: findMetricValueNumber(samples, 'redis_runtime_expired_keys_total'),
evictedKeysTotal: findMetricValueNumber(samples, 'redis_runtime_evicted_keys_total'),
keyspaceHitsTotal: findMetricValueNumber(samples, 'redis_runtime_keyspace_hits_total'),
keyspaceMissesTotal: findMetricValueNumber(samples, 'redis_runtime_keyspace_misses_total'),
keyspaceHitRateBasisPoints: findMetricValueNumber(samples, 'redis_runtime_keyspace_hit_rate_basis_points'),
usedMemoryBytes: findMetricValueNumber(samples, 'redis_runtime_used_memory_bytes'),
maxmemoryBytes: findMetricValueNumber(samples, 'redis_runtime_maxmemory_bytes'),
memoryUsageBasisPoints: findMetricValueNumber(samples, 'redis_runtime_memory_usage_basis_points'),
memoryFragmentationRatioBasisPoints: findMetricValueNumber(samples, 'redis_runtime_memory_fragmentation_ratio_basis_points'),
laneCommandErrorsTotal: sumMetricValues(samples, 'redis_runtime_lane_command_errors_total'),
laneCommandTimeoutsTotal: sumMetricValues(samples, 'redis_runtime_lane_command_timeouts_total'),
laneCommandCountTotal: sumMetricValues(samples, 'redis_runtime_lane_command_count_total'),
commandLatencyTotalMs: sumMetricValues(samples, 'redis_runtime_lane_command_latency_ms_sum'),
commandLatencyObservationCount: sumMetricValues(samples, 'redis_runtime_lane_command_latency_ms_count'),
commandLatencyMaxMs,
nonblockingCommandLatencyMaxMs,
}
}
function maxMetricValue(samples: PrometheusSample[], metricName: string): number | null {
const values = findMetricSamples(samples, metricName)
.map((sample) => Number(sample.value))
.filter((value) => Number.isFinite(value))
if (values.length === 0) return null
return Math.max(...values)
}
function buildProcessResourceMetrics(samples: PrometheusSample[]): GatewayProcessResourceMetrics {
return {
sampledAtUnixSecs: findMetricValueNumber(samples, 'gateway_process_sampled_at_unix_secs'),
systemCpuUsageBasisPoints: findMetricValueNumber(samples, 'gateway_system_cpu_usage_basis_points'),
processCpuUsageBasisPoints: findMetricValueNumber(samples, 'gateway_process_cpu_usage_basis_points'),
systemMemoryTotalBytes: findMetricValueNumber(samples, 'gateway_system_memory_total_bytes'),
systemMemoryUsedBytes: findMetricValueNumber(samples, 'gateway_system_memory_used_bytes'),
systemMemoryAvailableBytes: findMetricValueNumber(samples, 'gateway_system_memory_available_bytes'),
systemMemoryUsageBasisPoints: findMetricValueNumber(samples, 'gateway_system_memory_usage_basis_points'),
processMemoryBytes: findMetricValueNumber(samples, 'gateway_process_memory_bytes'),
processVirtualMemoryBytes: findMetricValueNumber(samples, 'gateway_process_virtual_memory_bytes'),
processMemoryBasisPoints: findMetricValueNumber(samples, 'gateway_process_memory_basis_points'),
processUptimeSeconds: findMetricValueNumber(samples, 'gateway_process_uptime_seconds'),
processThreads: findMetricValueNumber(samples, 'gateway_process_threads'),
openFds: findMetricValueNumber(samples, 'gateway_process_open_fds'),
fdLimit: findMetricValueNumber(samples, 'gateway_process_fd_limit'),
fdUsageBasisPoints: findMetricValueNumber(samples, 'gateway_process_fd_usage_basis_points'),
socketFds: findMetricValueNumber(samples, 'gateway_process_socket_fds'),
networkAvailable: metricBoolean(findMetricValueNumber(samples, 'gateway_network_observability_available')),
networkInterfaces: findMetricValueNumber(samples, 'gateway_network_interfaces'),
networkReceivedBytesTotal: findMetricValueNumber(samples, 'gateway_network_received_bytes_total'),
networkTransmittedBytesTotal: findMetricValueNumber(samples, 'gateway_network_transmitted_bytes_total'),
networkReceivedPacketsTotal: findMetricValueNumber(samples, 'gateway_network_received_packets_total'),
networkTransmittedPacketsTotal: findMetricValueNumber(samples, 'gateway_network_transmitted_packets_total'),
networkReceiveErrorsTotal: findMetricValueNumber(samples, 'gateway_network_receive_errors_total'),
networkTransmitErrorsTotal: findMetricValueNumber(samples, 'gateway_network_transmit_errors_total'),
networkReceiveDroppedTotal: findMetricValueNumber(samples, 'gateway_network_receive_dropped_total'),
networkTransmitDroppedTotal: findMetricValueNumber(samples, 'gateway_network_transmit_dropped_total'),
tcpStateAvailable: metricBoolean(findMetricValueNumber(samples, 'gateway_tcp_state_observability_available')),
hostTcpConnections: findMetricValueNumber(samples, 'gateway_host_tcp_connections'),
hostTcpEstablishedConnections: findMetricValueNumber(samples, 'gateway_host_tcp_established_connections'),
hostTcpListenConnections: findMetricValueNumber(samples, 'gateway_host_tcp_listen_connections'),
hostTcpTimeWaitConnections: findMetricValueNumber(samples, 'gateway_host_tcp_time_wait_connections'),
hostTcpSynSentConnections: findMetricValueNumber(samples, 'gateway_host_tcp_syn_sent_connections'),
hostTcpSynRecvConnections: findMetricValueNumber(samples, 'gateway_host_tcp_syn_recv_connections'),
hostTcpCloseWaitConnections: findMetricValueNumber(samples, 'gateway_host_tcp_close_wait_connections'),
processTcpConnections: findMetricValueNumber(samples, 'gateway_process_tcp_connections'),
processTcpEstablishedConnections: findMetricValueNumber(samples, 'gateway_process_tcp_established_connections'),
processTcpListenConnections: findMetricValueNumber(samples, 'gateway_process_tcp_listen_connections'),
processTcpTimeWaitConnections: findMetricValueNumber(samples, 'gateway_process_tcp_time_wait_connections'),
processTcpSynSentConnections: findMetricValueNumber(samples, 'gateway_process_tcp_syn_sent_connections'),
processTcpSynRecvConnections: findMetricValueNumber(samples, 'gateway_process_tcp_syn_recv_connections'),
processTcpCloseWaitConnections: findMetricValueNumber(samples, 'gateway_process_tcp_close_wait_connections'),
}
}
function buildAllocatorMetrics(samples: PrometheusSample[]): GatewayAllocatorMetrics {
return {
available: metricBoolean(findMetricValueNumber(samples, 'gateway_allocator_observability_available')),
allocatedBytes: findMetricValueNumber(samples, 'gateway_allocator_allocated_bytes'),
activeBytes: findMetricValueNumber(samples, 'gateway_allocator_active_bytes'),
residentBytes: findMetricValueNumber(samples, 'gateway_allocator_resident_bytes'),
mappedBytes: findMetricValueNumber(samples, 'gateway_allocator_mapped_bytes'),
retainedBytes: findMetricValueNumber(samples, 'gateway_allocator_retained_bytes'),
metadataBytes: findMetricValueNumber(samples, 'gateway_allocator_metadata_bytes'),
activeToAllocatedBasisPoints: findMetricValueNumber(samples, 'gateway_allocator_active_to_allocated_basis_points'),
residentToAllocatedBasisPoints: findMetricValueNumber(samples, 'gateway_allocator_resident_to_allocated_basis_points'),
}
}
function buildBackgroundTaskMetrics(samples: PrometheusSample[]): GatewayBackgroundTaskMetrics {
return {
active: findMetricValueNumber(samples, 'gateway_background_tasks_active'),
supervisedTotal: findMetricValueNumber(samples, 'gateway_background_tasks_supervised_total'),
unexpectedExitsTotal: findMetricValueNumber(samples, 'gateway_background_tasks_unexpected_exits_total'),
completedTotal: findMetricValueNumber(samples, 'gateway_background_tasks_completed_total'),
panickedTotal: findMetricValueNumber(samples, 'gateway_background_tasks_panicked_total'),
abortedTotal: findMetricValueNumber(samples, 'gateway_background_tasks_aborted_total'),
cancelledTotal: findMetricValueNumber(samples, 'gateway_background_tasks_cancelled_total'),
}
}
function buildTokioRuntimeMetrics(samples: PrometheusSample[]): GatewayTokioRuntimeMetrics {
return {
available: metricBoolean(findMetricValueNumber(samples, 'gateway_tokio_runtime_observability_available')),
workers: findMetricValueNumber(samples, 'gateway_tokio_runtime_workers'),
aliveTasks: findMetricValueNumber(samples, 'gateway_tokio_runtime_alive_tasks'),
globalQueueDepth: findMetricValueNumber(samples, 'gateway_tokio_runtime_global_queue_depth'),
}
}
function buildUsageRuntimeMetrics(samples: PrometheusSample[]): GatewayUsageRuntimeMetrics {
return {
enabled: metricBoolean(findMetricValueNumber(samples, 'usage_runtime_enabled')),
terminalQueueEnabled: metricBoolean(findMetricValueNumber(samples, 'usage_runtime_queue_terminal_events_enabled')),
lifecycleQueueEnabled: metricBoolean(findMetricValueNumber(samples, 'usage_runtime_queue_lifecycle_events_enabled')),
workerCount: findMetricValueNumber(samples, 'usage_runtime_queue_worker_count'),
workerAutoscaleEnabled: metricBoolean(findMetricValueNumber(samples, 'usage_runtime_queue_worker_autoscale_enabled')),
workerActiveCount: findMetricValueNumber(samples, 'usage_runtime_queue_worker_active_count'),
workerDesiredCount: findMetricValueNumber(samples, 'usage_runtime_queue_worker_desired_count'),
workerMaxCount: findMetricValueNumber(samples, 'usage_runtime_queue_worker_max_count'),
workerReadBatchesTotal: findMetricValueNumber(samples, 'usage_runtime_queue_worker_read_batches_total'),
workerReadEntriesTotal: findMetricValueNumber(samples, 'usage_runtime_queue_worker_read_entries_total'),
workerReclaimedEntriesTotal: findMetricValueNumber(samples, 'usage_runtime_queue_worker_reclaimed_entries_total'),
workerAckedEntriesTotal: findMetricValueNumber(samples, 'usage_runtime_queue_worker_acked_entries_total'),
workerDeadLetteredEntriesTotal: findMetricValueNumber(samples, 'usage_runtime_queue_worker_dead_lettered_entries_total'),
workerProcessFailuresTotal: findMetricValueNumber(samples, 'usage_runtime_queue_worker_process_failures_total'),
workerReadFailuresTotal: findMetricValueNumber(samples, 'usage_runtime_queue_worker_read_failures_total'),
workerReclaimFailuresTotal: findMetricValueNumber(samples, 'usage_runtime_queue_worker_reclaim_failures_total'),
terminalEnqueueInFlight: findMetricValueNumber(samples, 'usage_runtime_terminal_enqueue_in_flight'),
terminalEnqueueDeferredTotal: findMetricValueNumber(samples, 'usage_runtime_terminal_enqueue_deferred_total'),
terminalEnqueueDeferredRetryTotal: findMetricValueNumber(samples, 'usage_runtime_terminal_enqueue_deferred_retry_total'),
terminalEnqueueFailedTotal: findMetricValueNumber(samples, 'usage_runtime_terminal_enqueue_failed_total'),
lifecycleEnqueueInFlight: findMetricValueNumber(samples, 'usage_runtime_lifecycle_enqueue_in_flight'),
lifecycleEnqueueDeferredTotal: findMetricValueNumber(samples, 'usage_runtime_lifecycle_enqueue_deferred_total'),
lifecycleEnqueueDeferredDroppedTotal: findMetricValueNumber(samples, 'usage_runtime_lifecycle_enqueue_deferred_dropped_total'),
lifecycleEnqueueDeferredRetryTotal: findMetricValueNumber(samples, 'usage_runtime_lifecycle_enqueue_deferred_retry_total'),
lifecycleEnqueueFailedTotal: findMetricValueNumber(samples, 'usage_runtime_lifecycle_enqueue_failed_total'),
enqueueRetryScheduledTotal: findMetricValueNumber(samples, 'usage_runtime_enqueue_retry_scheduled_total'),
}
}
function buildUsageQueueMetrics(samples: PrometheusSample[]): GatewayUsageQueueMetrics {
const streamSample = findMetricSamples(samples, 'usage_queue_stream_length')[0]
const dlqSample = findMetricSamples(samples, 'usage_queue_dlq_length')[0]
return {
unavailable: metricBoolean(findMetricValueNumber(samples, 'usage_queue_health_unavailable')),
enabled: metricBoolean(findMetricValueNumber(samples, 'usage_queue_enabled')),
configured: metricBoolean(findMetricValueNumber(samples, 'usage_queue_configured')),
stream: streamSample?.labels.stream ?? null,
group: streamSample?.labels.group ?? null,
streamLength: findMetricValueNumber(samples, 'usage_queue_stream_length'),
groupPending: findMetricValueNumber(samples, 'usage_queue_group_pending'),
groupLag: findMetricValueNumber(samples, 'usage_queue_group_lag'),
oldestPendingIdleMs: findMetricValueNumber(samples, 'usage_queue_oldest_pending_idle_ms'),
dlqStream: dlqSample?.labels.stream ?? null,
dlqLength: findMetricValueNumber(samples, 'usage_queue_dlq_length'),
}
}
function buildUsageCounterMetrics(samples: PrometheusSample[]): GatewayUsageCounterMetrics {
const pendingByKind = findMetricSamples(samples, 'usage_counter_outbox_pending_rows_by_kind')
.map(sample => ({
kind: sample.labels.kind || 'unknown',
pendingRows: Number.isFinite(Number(sample.value)) ? Number(sample.value) : null,
}))
.sort((left, right) => (right.pendingRows ?? 0) - (left.pendingRows ?? 0))
return {
unavailable: metricBoolean(findMetricValueNumber(samples, 'usage_counter_health_unavailable')),
pendingRows: findMetricValueNumber(samples, 'usage_counter_outbox_pending_rows'),
processedRows: findMetricValueNumber(samples, 'usage_counter_outbox_processed_rows'),
oldestPendingAgeSeconds: findMetricValueNumber(samples, 'usage_counter_outbox_oldest_pending_age_seconds'),
oldestPendingCreatedAtUnixSecs: findMetricValueNumber(samples, 'usage_counter_outbox_oldest_pending_created_at_unix_secs'),
latestProcessedAtUnixSecs: findMetricValueNumber(samples, 'usage_counter_outbox_latest_processed_at_unix_secs'),
flushBatchesTotal: findMetricValueNumber(samples, 'usage_counter_outbox_flush_batches_total'),
flushRowsClaimedTotal: findMetricValueNumber(samples, 'usage_counter_outbox_flush_rows_claimed_total'),
flushTargetsTotal: sumMetricValues(samples, 'usage_counter_outbox_flush_targets_total'),
flushFailedBatchesTotal: findMetricValueNumber(samples, 'usage_counter_outbox_flush_failed_batches_total'),
cleanupRowsTotal: findMetricValueNumber(samples, 'usage_counter_outbox_cleanup_rows_total'),
cleanupFailedBatchesTotal: findMetricValueNumber(samples, 'usage_counter_outbox_cleanup_failed_batches_total'),
pendingByKind,
}
}
function buildRequestCandidateQueueMetrics(samples: PrometheusSample[]): GatewayRequestCandidateQueueMetrics {
return {
depth: findMetricValueNumber(samples, 'request_candidate_queue_depth'),
pendingDepth: findMetricValueNumber(samples, 'request_candidate_queue_pending_depth'),
capacity: findMetricValueNumber(samples, 'request_candidate_queue_capacity'),
enqueuedTotal: findMetricValueNumber(samples, 'request_candidate_queue_enqueued_total'),
droppedTotal: findMetricValueNumber(samples, 'request_candidate_queue_dropped_total'),
flushedTotal: findMetricValueNumber(samples, 'request_candidate_queue_flushed_total'),
flushFailedTotal: findMetricValueNumber(samples, 'request_candidate_queue_flush_failed_total'),
flushBatchesTotal: findMetricValueNumber(samples, 'request_candidate_queue_flush_batches_total'),
flushSqlOpsTotal: findMetricValueNumber(samples, 'request_candidate_queue_flush_sql_ops_total'),
flushSqlRecordsTotal: findMetricValueNumber(samples, 'request_candidate_queue_flush_sql_records_total'),
compactedTotal: findMetricValueNumber(samples, 'request_candidate_queue_compacted_total'),
syncFallbackTotal: findMetricValueNumber(samples, 'request_candidate_queue_sync_fallback_total'),
}
}
function collectUpstreamTargets(samples: PrometheusSample[]): string[] {
const targets = new Set<string>()
const metricNames = [
'upstream_target_gate_in_flight',
'upstream_target_gate_available_permits',
'upstream_target_gate_high_watermark',
'upstream_target_gate_rejected_total',
'upstream_target_selected_total',
'upstream_target_saturated_total',
]
for (const metricName of metricNames) {
for (const sample of findMetricSamples(samples, metricName)) {
if (sample.labels.target) {
targets.add(sample.labels.target)
}
}
}
return Array.from(targets)
}
function buildUpstreamTargetMetrics(samples: PrometheusSample[]): GatewayUpstreamTargetMetrics {
const rows = collectUpstreamTargets(samples)
.map(target => ({
target,
inFlight: findMetricValueNumber(samples, 'upstream_target_gate_in_flight', { target }),
availablePermits: findMetricValueNumber(samples, 'upstream_target_gate_available_permits', { target }),
highWatermark: findMetricValueNumber(samples, 'upstream_target_gate_high_watermark', { target }),
rejectedTotal: findMetricValueNumber(samples, 'upstream_target_gate_rejected_total', { target }),
selectedTotal: findMetricValueNumber(samples, 'upstream_target_selected_total', { target }),
saturatedTotal: findMetricValueNumber(samples, 'upstream_target_saturated_total', { target }),
}))
.sort((left, right) => (
(right.inFlight ?? 0) - (left.inFlight ?? 0)
|| (right.saturatedTotal ?? 0) - (left.saturatedTotal ?? 0)
|| (right.selectedTotal ?? 0) - (left.selectedTotal ?? 0)
))
.slice(0, 8)
return {
activeTargets: findMetricValueNumber(samples, 'upstream_target_gate_active_targets'),
limit: findMetricValueNumber(samples, 'upstream_target_gate_limit'),
selectedTotal: sumMetricValues(samples, 'upstream_target_selected_total'),
saturatedTotal: sumMetricValues(samples, 'upstream_target_saturated_total'),
rejectedTotal: sumMetricValues(samples, 'upstream_target_gate_rejected_total'),
rows,
}
}
function buildStageLatencyMetrics(samples: PrometheusSample[]): GatewayStageLatencyMetrics {
return {
rows: CAPACITY_STAGES.map(({ stage, label }) => {
const count = findMetricValueNumber(samples, 'gateway_stage_latency_count', { stage })
const sumMs = findMetricValueNumber(samples, 'gateway_stage_latency_sum_ms', { stage })
return {
stage,
label,
count,
avgMs: count != null && count > 0 && sumMs != null ? sumMs / count : null,
maxMs: findMetricValueNumber(samples, 'gateway_stage_latency_max_ms', { stage }),
}
}),
}
}
export function buildGatewayMetricsSummary(text: string): GatewayMetricsSummary {
const samples = parsePrometheusSamples(text)
const fallbacks = FALLBACK_METRICS.map(item => ({
@@ -170,6 +821,21 @@ export function buildGatewayMetricsSummary(text: string): GatewayMetricsSummary
serviceUp: findMetricValueNumber(samples, 'service_up', { service: 'aether-gateway' }),
local: buildGateMetrics(samples, 'gateway_requests'),
distributed: buildGateMetrics(samples, 'gateway_requests_distributed'),
candidatePlanning: buildGateMetrics(samples, 'gateway_candidate_planning'),
upstreamExecution: buildGateMetrics(samples, 'gateway_upstream_execution'),
databasePool: buildDatabasePoolMetrics(samples),
postgres: buildPostgresObservabilityMetrics(samples),
redisRuntime: buildRedisRuntimeMetrics(samples),
process: buildProcessResourceMetrics(samples),
allocator: buildAllocatorMetrics(samples),
backgroundTasks: buildBackgroundTaskMetrics(samples),
tokioRuntime: buildTokioRuntimeMetrics(samples),
usageRuntime: buildUsageRuntimeMetrics(samples),
usageQueue: buildUsageQueueMetrics(samples),
usageCounter: buildUsageCounterMetrics(samples),
requestCandidateQueue: buildRequestCandidateQueueMetrics(samples),
upstreamTargets: buildUpstreamTargetMetrics(samples),
stageLatency: buildStageLatencyMetrics(samples),
tunnel: {
proxyConnections: findMetricValueNumber(samples, 'tunnel_proxy_connections'),
availableProxyConnections: findMetricValueNumber(samples, 'tunnel_proxy_connections_available'),
@@ -1,5 +1,6 @@
import { describe, expect, it } from 'vitest'
import {
findMetricSamples,
findMetricValueNumber,
parsePrometheusSamples,
sumMetricValues,
@@ -34,4 +35,15 @@ decision_remote_total{route_kind="responses",reason="remote_decision_miss"} 3
expect(sumMetricValues(samples, 'decision_remote_total')).toBe(5)
})
it('finds all matching samples by full or suffix metric name', () => {
const samples = parsePrometheusSamples(`
aether_gateway_upstream_target_selected_total{target="openai"} 4
aether_gateway_upstream_target_selected_total{target="azure"} 6
aether_gateway_upstream_target_saturated_total{target="openai"} 1
`)
expect(findMetricSamples(samples, 'upstream_target_selected_total')).toHaveLength(2)
expect(sumMetricValues(samples, 'upstream_target_selected_total')).toBe(10)
})
})
+7
View File
@@ -44,6 +44,13 @@ export function sumMetricValues(
}, 0)
}
export function findMetricSamples(
samples: PrometheusSample[],
metricName: string
): PrometheusSample[] {
return samples.filter(sample => metricNameMatches(sample.name, metricName))
}
function metricNameMatches(actual: string, expected: string): boolean {
return actual === expected || actual.split('_').pop() === expected || actual.endsWith(`_${expected}`)
}
File diff suppressed because it is too large Load Diff
+11
View File
@@ -0,0 +1,11 @@
#!/usr/bin/env node
const path = require('node:path')
const { spawnSync } = require('node:child_process')
const reportPath = process.argv[2] || '/tmp/aether_gateway_pressure_s1_1k.json'
const checker = path.join(__dirname, 'check_gateway_stage_report.js')
const result = spawnSync(process.execPath, [checker, '--stage', 'S1', reportPath], {
stdio: 'inherit',
})
process.exit(result.status ?? 1)
+452
View File
@@ -0,0 +1,452 @@
#!/usr/bin/env node
const http = require('node:http')
const https = require('node:https')
const DEFAULT_STAGE = 'S1'
const DEFAULT_GATEWAY_BASE_URL = 'http://127.0.0.1:8084'
const DEFAULT_MOCK_UPSTREAM_METRICS_URL = 'http://127.0.0.1:18181/metrics'
const REQUIRED_M4_METRICS = [
'gateway_process_open_fds',
'gateway_process_fd_limit',
'gateway_process_fd_usage_basis_points',
'gateway_process_threads',
'gateway_process_socket_fds',
'gateway_process_tcp_established_connections',
'gateway_host_tcp_established_connections',
'gateway_network_observability_available',
'gateway_network_received_bytes_total',
'gateway_background_tasks_active',
'gateway_background_tasks_unexpected_exits_total',
'gateway_tokio_runtime_observability_available',
'gateway_tokio_runtime_workers',
'gateway_allocator_observability_available',
'postgres_observability_available',
'postgres_wal_observability_available',
'postgres_checkpoint_observability_available',
'postgres_statement_observability_available',
'redis_runtime_enabled',
'redis_runtime_lane_command_latency_ms_max',
'usage_runtime_queue_worker_read_batches_total',
'usage_runtime_queue_worker_acked_entries_total',
'usage_counter_outbox_flush_batches_total',
'usage_counter_outbox_cleanup_rows_total',
'request_candidate_queue_flush_batches_total',
'request_candidate_queue_flush_sql_ops_total',
]
function parseArgs(argv) {
const gatewayBaseUrl = process.env.GATEWAY_BASE_URL || DEFAULT_GATEWAY_BASE_URL
const options = {
stage: process.env.PRESSURE_STAGE || DEFAULT_STAGE,
gatewayBaseUrl,
healthUrl: `${gatewayBaseUrl.replace(/\/$/, '')}/_gateway/health`,
metricsUrl: process.env.METRICS_URL || `${gatewayBaseUrl.replace(/\/$/, '')}/_gateway/metrics`,
targetUrl: process.env.TARGET_URL || `${gatewayBaseUrl.replace(/\/$/, '')}/v1/chat/completions`,
mockUpstreamMetricsUrl:
process.env.PRESSURE_MOCK_UPSTREAM_METRICS_URL || DEFAULT_MOCK_UPSTREAM_METRICS_URL,
timeoutMs: numberEnv('PRESSURE_PREFLIGHT_TIMEOUT_MS') ?? 5000,
requireAuth: boolEnv('PRESSURE_REQUIRE_AUTH', true),
requireM4Metrics: boolEnv('PRESSURE_REQUIRE_M4_METRICS', true),
requireMockUpstream: boolEnv('PRESSURE_REQUIRE_MOCK_UPSTREAM', true),
apiKeyFile:
process.env.AETHER_API_KEY_FILE ||
process.env.API_KEY_FILE ||
process.env.PRESSURE_API_KEY_FILE ||
'',
apiKeyListFile:
process.env.AETHER_API_KEY_LIST_FILE ||
process.env.API_KEY_LIST_FILE ||
process.env.PRESSURE_API_KEY_LIST_FILE ||
'',
}
for (let index = 0; index < argv.length; index += 1) {
const arg = argv[index]
switch (arg) {
case '--stage':
options.stage = requireValue(argv, ++index, arg)
break
case '--gateway-base-url':
options.gatewayBaseUrl = requireValue(argv, ++index, arg)
options.healthUrl = `${options.gatewayBaseUrl.replace(/\/$/, '')}/_gateway/health`
if (!process.env.METRICS_URL) {
options.metricsUrl = `${options.gatewayBaseUrl.replace(/\/$/, '')}/_gateway/metrics`
}
if (!process.env.TARGET_URL) {
options.targetUrl = `${options.gatewayBaseUrl.replace(/\/$/, '')}/v1/chat/completions`
}
break
case '--health-url':
options.healthUrl = requireValue(argv, ++index, arg)
break
case '--metrics-url':
options.metricsUrl = requireValue(argv, ++index, arg)
break
case '--target-url':
options.targetUrl = requireValue(argv, ++index, arg)
break
case '--mock-upstream-metrics-url':
options.mockUpstreamMetricsUrl = requireValue(argv, ++index, arg)
break
case '--timeout-ms':
options.timeoutMs = parsePositiveInteger(requireValue(argv, ++index, arg), arg)
break
case '--require-auth':
options.requireAuth = true
break
case '--skip-auth':
options.requireAuth = false
break
case '--api-key-file':
options.apiKeyFile = requireValue(argv, ++index, arg)
break
case '--api-key-list-file':
options.apiKeyListFile = requireValue(argv, ++index, arg)
break
case '--require-m4-metrics':
options.requireM4Metrics = true
break
case '--skip-m4-metrics':
options.requireM4Metrics = false
break
case '--require-mock-upstream':
options.requireMockUpstream = true
break
case '--skip-mock-upstream':
options.requireMockUpstream = false
break
case '--help':
case '-h':
printHelp()
process.exit(0)
default:
throw new Error(`unknown option: ${arg}`)
}
}
options.stage = String(options.stage).trim().toUpperCase()
return options
}
function numberEnv(name) {
const value = process.env[name]
if (value == null || value === '') {
return undefined
}
return parsePositiveInteger(value, name)
}
function boolEnv(name, fallback) {
const value = process.env[name]
if (value == null || value === '') {
return fallback
}
return ['1', 'true', 'yes', 'on'].includes(value.toLowerCase())
}
function requireValue(argv, index, option) {
const value = argv[index]
if (!value || value.startsWith('--')) {
throw new Error(`${option} requires a value`)
}
return value
}
function parsePositiveInteger(value, name) {
const number = Number(value)
if (!Number.isInteger(number) || number <= 0) {
throw new Error(`${name} must be a positive integer, got ${value}`)
}
return number
}
function printHelp() {
console.log(`Usage: tools/pressure/check_gateway_stage_preflight.js [options]
Checks whether a gateway is ready to run staged mock streaming pressure tests.
The script never prints auth header or API key values.
Options:
--stage S1|S2|S3|S4|S5
--gateway-base-url URL
--health-url URL
--metrics-url URL
--target-url URL
--mock-upstream-metrics-url URL
--timeout-ms N
--skip-auth
--api-key-file PATH
--api-key-list-file PATH
--skip-mock-upstream
--skip-m4-metrics
`)
}
async function main() {
const options = parseArgs(process.argv.slice(2))
const ok = []
const failures = []
if (options.requireAuth) {
if (authConfigured(options)) {
ok.push('auth configured')
} else {
failures.push(
'missing auth: set AUTH_HEADER, AETHER_API_KEY, API_KEY, AETHER_API_KEY_FILE, or AETHER_API_KEY_LIST_FILE',
)
}
}
if (!isHttpUrl(options.targetUrl)) {
failures.push(`target URL is not http(s): ${options.targetUrl}`)
} else {
ok.push(`target configured: ${options.targetUrl}`)
}
let metricsText = ''
await checkHttpText('gateway health', options.healthUrl, options.timeoutMs, ok, failures, (body) => {
const health = parseJson(body)
if (!health) {
failures.push('gateway health did not return JSON')
return
}
if (health.status !== 'ok') {
failures.push(`gateway health status=${health.status ?? 'missing'}, expected ok`)
return
}
ok.push('gateway health status ok')
})
await checkHttpText('gateway metrics', options.metricsUrl, options.timeoutMs, ok, failures, (body) => {
metricsText = body
const metricNames = parseMetricNames(body)
if (metricNames.size === 0) {
failures.push('gateway metrics response did not contain Prometheus samples')
return
}
ok.push(`gateway metrics samples available (${metricNames.size} metric names)`)
if (options.requireM4Metrics) {
const missing = REQUIRED_M4_METRICS.filter((name) => !metricNames.has(name))
if (missing.length > 0) {
failures.push(`gateway metrics missing M4 required metrics: ${missing.join(', ')}`)
} else {
ok.push('gateway M4 metrics present')
}
const tokioAvailable = metricMax(body, 'gateway_tokio_runtime_observability_available')
if (tokioAvailable !== null && tokioAvailable !== 1) {
failures.push(`gateway_tokio_runtime_observability_available=${tokioAvailable}, expected 1`)
}
}
})
if (options.requireMockUpstream) {
await checkHttpText(
'mock upstream metrics',
options.mockUpstreamMetricsUrl,
options.timeoutMs,
ok,
failures,
(body) => {
if (!body.trim()) {
failures.push('mock upstream metrics response was empty')
return
}
ok.push('mock upstream metrics available')
},
)
}
console.log(`gateway staged pressure preflight: ${options.stage}`)
ok.forEach((line) => console.log(`OK ${line}`))
if (failures.length > 0) {
console.error('FAIL preflight checks failed:')
failures.forEach((line) => console.error(`FAIL ${line}`))
process.exit(1)
}
if (metricsText) {
const dbPoolMax = metricMax(metricsText, 'database_pool_max_connections')
const upstreamPermits = metricMax(metricsText, 'concurrency_available_permits', {
gate: 'gateway_upstream_execution',
})
if (dbPoolMax !== null) {
console.log(`OK database_pool_max_connections=${dbPoolMax}`)
}
if (upstreamPermits !== null) {
console.log(`OK gateway_upstream_execution_available_permits=${upstreamPermits}`)
}
}
console.log('PASS gateway staged pressure preflight')
}
function authConfigured(options) {
if (['AUTH_HEADER', 'AETHER_API_KEY', 'API_KEY'].some((name) => {
const value = process.env[name]
return typeof value === 'string' && value.trim().length > 0
})) {
return true
}
return fileHasSecret(options.apiKeyFile) || fileHasSecret(options.apiKeyListFile)
}
function fileHasSecret(path) {
if (!path || !path.trim()) {
return false
}
try {
const fs = require('node:fs')
return fs.readFileSync(path, 'utf8').trim().length > 0
} catch (_error) {
return false
}
}
async function checkHttpText(label, url, timeoutMs, ok, failures, inspect) {
if (!isHttpUrl(url)) {
failures.push(`${label} URL is not http(s): ${url}`)
return
}
try {
const response = await requestText(url, timeoutMs)
if (response.statusCode < 200 || response.statusCode >= 300) {
failures.push(`${label} returned HTTP ${response.statusCode}`)
return
}
inspect(response.body)
} catch (error) {
failures.push(`${label} unreachable at ${url}: ${error.message}`)
}
}
function isHttpUrl(value) {
try {
const url = new URL(value)
return url.protocol === 'http:' || url.protocol === 'https:'
} catch (_error) {
return false
}
}
function requestText(urlString, timeoutMs) {
return new Promise((resolve, reject) => {
const url = new URL(urlString)
const client = url.protocol === 'https:' ? https : http
const request = client.get(
url,
{
headers: {
accept: 'text/plain, application/json;q=0.9, */*;q=0.1',
},
},
(response) => {
response.setEncoding('utf8')
let body = ''
response.on('data', (chunk) => {
body += chunk
})
response.on('end', () => {
resolve({
statusCode: response.statusCode ?? 0,
body,
})
})
},
)
request.setTimeout(timeoutMs, () => {
request.destroy(new Error(`timed out after ${timeoutMs}ms`))
})
request.on('error', reject)
})
}
function parseJson(value) {
try {
return JSON.parse(value)
} catch (_error) {
return null
}
}
function parseMetricNames(text) {
const names = new Set()
for (const line of text.split(/\r?\n/)) {
if (!line || line.startsWith('#')) {
continue
}
const match = line.match(/^([a-zA-Z_:][a-zA-Z0-9_:.-]*)(?:\{|[\s])/)
if (!match) {
continue
}
addMetricName(names, match[1])
}
return names
}
function addMetricName(names, name) {
names.add(name)
if (name.startsWith('aether-gateway_')) {
names.add(name.slice('aether-gateway_'.length))
}
}
function metricMax(text, metricName, labels = {}) {
let max = null
for (const line of text.split(/\r?\n/)) {
if (!line || line.startsWith('#')) {
continue
}
const parsed = parseMetricSample(line)
if (!parsed) {
continue
}
if (parsed.name !== metricName && parsed.name !== `aether-gateway_${metricName}`) {
continue
}
if (!labelsMatch(parsed.labels, labels)) {
continue
}
max = max === null ? parsed.value : Math.max(max, parsed.value)
}
return max
}
function parseMetricSample(line) {
const match = line.match(/^([a-zA-Z_:][a-zA-Z0-9_:.-]*)(\{[^}]*\})?\s+(-?(?:\d+\.?\d*|\d*\.\d+)(?:[eE][+-]?\d+)?)\s*$/)
if (!match) {
return null
}
return {
name: match[1],
labels: parseLabels(match[2]),
value: Number(match[3]),
}
}
function parseLabels(labelText) {
if (!labelText) {
return {}
}
const labels = {}
const inner = labelText.slice(1, -1)
const pattern = /([a-zA-Z_][a-zA-Z0-9_]*)="((?:\\.|[^"\\])*)"/g
let match
while ((match = pattern.exec(inner)) !== null) {
labels[match[1]] = match[2].replace(/\\"/g, '"').replace(/\\\\/g, '\\')
}
return labels
}
function labelsMatch(actual, expected) {
return Object.entries(expected).every(([name, value]) => actual[name] === value)
}
main().catch((error) => {
console.error(`FAIL ${error.message}`)
process.exit(1)
})
File diff suppressed because it is too large Load Diff
+197
View File
@@ -0,0 +1,197 @@
#!/usr/bin/env node
const assert = require('node:assert/strict')
const fs = require('node:fs')
const os = require('node:os')
const path = require('node:path')
const test = require('node:test')
const { spawnSync } = require('node:child_process')
const checker = path.join(__dirname, 'check_gateway_stage_report.js')
test('realistic-stream report passes full-chain latency and throughput checks', () => {
const report = reportFor({
totalRequests: 1000,
concurrency: 1000,
throughputRps: 180,
headersP95Ms: 120,
firstBodyP95Ms: 350,
p95Ms: 6000,
p99Ms: 9000,
})
const result = runChecker('--stage', 'realistic-stream', writeReport(report))
assert.equal(result.status, 0, result.stderr)
assert.match(result.stdout, /REALISTIC_STREAM PASS/)
assert.match(result.stdout, /throughput_rps=180/)
})
test('tps report passes completed request throughput checks', () => {
const report = reportFor({
totalRequests: 20000,
concurrency: 500,
throughputRps: 750,
headersP95Ms: 90,
firstBodyP95Ms: 180,
p95Ms: 900,
p99Ms: 1600,
})
const result = runChecker('--stage', 'TPS', writeReport(report))
assert.equal(result.status, 0, result.stderr)
assert.match(result.stdout, /TPS PASS/)
})
test('tps report fails when throughput is below the acceptance threshold', () => {
const report = reportFor({
totalRequests: 20000,
concurrency: 500,
throughputRps: 499,
headersP95Ms: 90,
firstBodyP95Ms: 180,
p95Ms: 900,
p99Ms: 1600,
})
const result = runChecker('--stage', 'tps', writeReport(report))
assert.equal(result.status, 1)
assert.match(result.stderr, /TPS FAIL: load\.throughput_rps=499, expected >= 500/)
})
test('tps report fails when settle drain does not complete', () => {
const report = reportFor({
totalRequests: 20000,
concurrency: 500,
throughputRps: 750,
headersP95Ms: 90,
firstBodyP95Ms: 180,
p95Ms: 900,
p99Ms: 1600,
})
report.settle_drain_completed = false
const result = runChecker('--stage', 'tps', writeReport(report))
assert.equal(result.status, 1)
assert.match(result.stderr, /TPS FAIL: settle_drain_completed=false after 5000ms/)
})
test('tps report fails when lifecycle enqueue drops deferred events', () => {
const report = reportFor({
totalRequests: 20000,
concurrency: 500,
throughputRps: 750,
headersP95Ms: 90,
firstBodyP95Ms: 180,
p95Ms: 900,
p99Ms: 1600,
})
report.metrics.usage_runtime_max_lifecycle_enqueue_deferred_dropped_total = 1
const result = runChecker('--stage', 'tps', writeReport(report))
assert.equal(result.status, 1)
assert.match(
result.stderr,
/TPS FAIL: usage_runtime_max_lifecycle_enqueue_deferred_dropped_total=1/,
)
})
test('tps report ignores shared redis error replies when gateway lane errors are clean', () => {
const report = reportFor({
totalRequests: 20000,
concurrency: 500,
throughputRps: 750,
headersP95Ms: 90,
firstBodyP95Ms: 180,
p95Ms: 900,
p99Ms: 1600,
})
report.metrics.redis_runtime_total_error_replies_delta = 12
report.metrics.redis_runtime_lane_command_errors_total_delta = 0
report.metrics.redis_runtime_lane_command_timeouts_total_delta = 0
const result = runChecker('--stage', 'tps', writeReport(report))
assert.equal(result.status, 0, result.stderr)
})
test('tps report falls back to redis error replies when lane metrics are absent', () => {
const report = reportFor({
totalRequests: 20000,
concurrency: 500,
throughputRps: 750,
headersP95Ms: 90,
firstBodyP95Ms: 180,
p95Ms: 900,
p99Ms: 1600,
})
report.metrics.redis_runtime_total_error_replies_delta = 1
delete report.metrics.redis_runtime_lane_command_errors_total_delta
delete report.metrics.redis_runtime_lane_command_timeouts_total_delta
const result = runChecker('--stage', 'tps', writeReport(report))
assert.equal(result.status, 1)
assert.match(result.stderr, /TPS FAIL: redis_runtime_total_error_replies_delta=1/)
})
function runChecker(...args) {
return spawnSync(process.execPath, [checker, ...args], {
cwd: path.resolve(__dirname, '../..'),
encoding: 'utf8',
})
}
function writeReport(report) {
const file = path.join(
fs.mkdtempSync(path.join(os.tmpdir(), 'aether-stage-report-')),
'report.json',
)
fs.writeFileSync(file, `${JSON.stringify(report)}\n`)
return file
}
function reportFor({
totalRequests,
concurrency,
throughputRps,
headersP95Ms,
firstBodyP95Ms,
p95Ms,
p99Ms,
}) {
return {
settle_after_ms: 5000,
settle_drain_completed: true,
settle_drain_elapsed_ms: 250,
load: {
response_mode: 'FullBody',
total_requests: totalRequests,
completed_requests: totalRequests,
failed_requests: 0,
concurrency,
first_body_hold_ms: 0,
throughput_rps: throughputRps,
headers_p95_ms: headersP95Ms,
first_body_p95_ms: firstBodyP95Ms,
p95_ms: p95Ms,
p99_ms: p99Ms,
status_counts: { 200: totalRequests },
error_counts: {},
},
metrics: {
samples: 10,
db_pool_max_usage_basis_points: 2500,
db_pool_pressure_samples: 0,
gateway_requests_max_rejected_total: 0,
gateway_requests_distributed_max_rejected_total: 0,
request_candidate_queue_final_depth: 0,
request_candidate_queue_final_pending_depth: 0,
request_candidate_queue_max_flush_failed_total: 0,
request_candidate_queue_max_dropped_total: 0,
request_candidate_queue_max_sync_fallback_total: 0,
usage_runtime_max_terminal_enqueue_failed_total: 0,
usage_runtime_max_lifecycle_enqueue_failed_total: 0,
usage_runtime_max_lifecycle_enqueue_deferred_dropped_total: 0,
redis_runtime_lane_command_errors_total_delta: 0,
redis_runtime_lane_command_timeouts_total_delta: 0,
upstream_target_max_rejected_total: 0,
},
}
}
+5
View File
@@ -0,0 +1,5 @@
#!/usr/bin/env bash
set -euo pipefail
PRESSURE_STAGE="${PRESSURE_STAGE:-S1}"
exec "$(dirname "$0")/run_gateway_mock_streaming_stage.sh"
+238
View File
@@ -0,0 +1,238 @@
#!/usr/bin/env bash
set -euo pipefail
# Gateway staged mock streaming pressure probe.
#
# Required auth:
# AETHER_API_KEY_FILE=/path/to/api-key
# or:
# AUTH_HEADER='Authorization: Bearer <aether-api-key>'
# or:
# AETHER_API_KEY='<aether-api-key>'
#
# Common settings:
# PRESSURE_STAGE=S1|S2|S3|S4|S5
# GATEWAY_BASE_URL=http://127.0.0.1:8084
# TARGET_URL=http://127.0.0.1:8084/v1/chat/completions
# METRICS_URL=http://127.0.0.1:8084/_gateway/metrics
# PRESSURE_MODEL=gpt-5-mini
# PRESSURE_RESPONSE_MODE=first-body-byte
# PRESSURE_CARGO_PROFILE=release
script_dir="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)"
repo_root="$(cd -- "$script_dir/../.." && pwd)"
PRESSURE_STAGE="${PRESSURE_STAGE:-S1}"
PRESSURE_STAGE="$(printf '%s' "$PRESSURE_STAGE" | tr '[:lower:]' '[:upper:]')"
case "$PRESSURE_STAGE" in
S1)
default_requests=1000
default_concurrency=1000
default_hold_ms=600000
default_timeout_ms=720000
default_start_ramp_ms=10000
default_output=/tmp/aether_gateway_pressure_s1_1k.json
;;
S2)
default_requests=3000
default_concurrency=3000
default_hold_ms=900000
default_timeout_ms=1080000
default_start_ramp_ms=30000
default_output=/tmp/aether_gateway_pressure_s2_3k.json
;;
S3)
default_requests=6000
default_concurrency=6000
default_hold_ms=1800000
default_timeout_ms=1980000
default_start_ramp_ms=60000
default_output=/tmp/aether_gateway_pressure_s3_6k.json
;;
S4)
default_requests=10000
default_concurrency=10000
default_hold_ms=1800000
default_timeout_ms=2100000
default_start_ramp_ms=90000
default_output=/tmp/aether_gateway_pressure_s4_10k.json
;;
S5)
default_requests=10000
default_concurrency=10000
default_hold_ms=7200000
default_timeout_ms=7500000
default_start_ramp_ms=120000
default_output=/tmp/aether_gateway_pressure_s5_10k_soak.json
;;
*)
echo "unsupported PRESSURE_STAGE=$PRESSURE_STAGE; expected S1, S2, S3, S4, or S5" >&2
exit 2
;;
esac
GATEWAY_BASE_URL="${GATEWAY_BASE_URL:-http://127.0.0.1:8084}"
TARGET_URL="${TARGET_URL:-${GATEWAY_BASE_URL%/}/v1/chat/completions}"
METRICS_URL="${METRICS_URL:-${GATEWAY_BASE_URL%/}/_gateway/metrics}"
PRESSURE_REQUESTS="${PRESSURE_REQUESTS:-$default_requests}"
PRESSURE_CONCURRENCY="${PRESSURE_CONCURRENCY:-$default_concurrency}"
PRESSURE_TIMEOUT_MS="${PRESSURE_TIMEOUT_MS:-$default_timeout_ms}"
PRESSURE_CONNECT_TIMEOUT_MS="${PRESSURE_CONNECT_TIMEOUT_MS:-30000}"
PRESSURE_SAMPLE_INTERVAL_MS="${PRESSURE_SAMPLE_INTERVAL_MS:-500}"
PRESSURE_SETTLE_AFTER_MS="${PRESSURE_SETTLE_AFTER_MS:-2000}"
PRESSURE_START_RAMP_MS="${PRESSURE_START_RAMP_MS:-$default_start_ramp_ms}"
PRESSURE_FIRST_BODY_HOLD_MS="${PRESSURE_FIRST_BODY_HOLD_MS:-$default_hold_ms}"
PRESSURE_METHOD="${PRESSURE_METHOD:-POST}"
PRESSURE_RESPONSE_MODE="${PRESSURE_RESPONSE_MODE:-first-body-byte}"
PRESSURE_CARGO_PROFILE="${PRESSURE_CARGO_PROFILE:-release}"
PRESSURE_MODEL="${PRESSURE_MODEL:-gpt-5-mini}"
OUTPUT="${OUTPUT:-$default_output}"
api_key_file="${AETHER_API_KEY_FILE:-${API_KEY_FILE:-${PRESSURE_API_KEY_FILE:-}}}"
stage_lower="$(printf '%s' "$PRESSURE_STAGE" | tr '[:upper:]' '[:lower:]')"
PRESSURE_BODY_FILE="${PRESSURE_BODY_FILE:-/tmp/aether-pressure-${stage_lower}-mock-streaming-request.json}"
if [[ -z "${AUTH_HEADER:-}" ]]; then
if [[ -n "$api_key_file" && -s "$api_key_file" ]]; then
:
elif [[ -n "${AETHER_API_KEY:-}" ]]; then
AUTH_HEADER="Authorization: Bearer ${AETHER_API_KEY}"
elif [[ -n "${API_KEY:-}" ]]; then
AUTH_HEADER="Authorization: Bearer ${API_KEY}"
else
echo "missing auth: set AETHER_API_KEY_FILE, AUTH_HEADER, or AETHER_API_KEY before running gateway staged pressure" >&2
exit 2
fi
fi
if [[ -z "${PRESSURE_BODY:-}" && ! -s "$PRESSURE_BODY_FILE" ]]; then
cat >"$PRESSURE_BODY_FILE" <<JSON
{"model":"${PRESSURE_MODEL}","messages":[{"role":"user","content":"ping"}],"stream":true}
JSON
fi
args=(run)
case "$PRESSURE_CARGO_PROFILE" in
release)
args+=(--release)
;;
debug)
;;
*)
echo "unsupported PRESSURE_CARGO_PROFILE=$PRESSURE_CARGO_PROFILE; expected release or debug" >&2
exit 2
;;
esac
args+=(
-p aether-testkit --bin gateway_pressure_probe --
--url "$TARGET_URL"
--metrics-url "$METRICS_URL"
--requests "$PRESSURE_REQUESTS"
--concurrency "$PRESSURE_CONCURRENCY"
--timeout-ms "$PRESSURE_TIMEOUT_MS"
--connect-timeout-ms "$PRESSURE_CONNECT_TIMEOUT_MS"
--sample-interval-ms "$PRESSURE_SAMPLE_INTERVAL_MS"
--settle-after-ms "$PRESSURE_SETTLE_AFTER_MS"
--start-ramp-ms "$PRESSURE_START_RAMP_MS"
--first-body-hold-ms "$PRESSURE_FIRST_BODY_HOLD_MS"
--method "$PRESSURE_METHOD"
--response-mode "$PRESSURE_RESPONSE_MODE"
--output "$OUTPUT"
)
if [[ -n "$api_key_file" && -s "$api_key_file" ]]; then
args+=(--api-key-file "$api_key_file")
else
args+=(--header "$AUTH_HEADER")
fi
if [[ -n "${EXTRA_HEADERS:-}" ]]; then
while IFS= read -r header; do
[[ -z "$header" ]] && continue
args+=(--header "$header")
done <<< "$EXTRA_HEADERS"
else
args+=(--header "Content-Type: application/json")
fi
if [[ -n "${PRESSURE_CLIENT_SHARDS:-}" ]]; then
args+=(--client-shards "$PRESSURE_CLIENT_SHARDS")
fi
if [[ -n "${PRESSURE_POOL_MAX_IDLE_PER_HOST:-}" ]]; then
args+=(--pool-max-idle-per-host "$PRESSURE_POOL_MAX_IDLE_PER_HOST")
fi
if [[ -n "${PRESSURE_WARMUP_CONNECTIONS:-}" ]]; then
args+=(--warmup-connections "$PRESSURE_WARMUP_CONNECTIONS")
fi
if [[ -n "${PRESSURE_WARMUP_URL:-}" ]]; then
args+=(--warmup-url "$PRESSURE_WARMUP_URL")
fi
if [[ "${PRESSURE_HTTP1_ONLY:-false}" == "true" ]]; then
args+=(--http1-only)
fi
if [[ "${PRESSURE_HTTP2_PRIOR_KNOWLEDGE:-false}" == "true" ]]; then
args+=(--http2-prior-knowledge)
fi
if [[ -n "${PRESSURE_BODY:-}" ]]; then
args+=(--body "$PRESSURE_BODY")
else
args+=(--body-file "$PRESSURE_BODY_FILE")
fi
metrics_before="${OUTPUT%.json}.metrics.before.prom"
metrics_after="${OUTPUT%.json}.metrics.after.prom"
if [[ "${PRESSURE_PREFLIGHT:-true}" == "true" ]]; then
preflight_args=(
--stage "$PRESSURE_STAGE"
--gateway-base-url "$GATEWAY_BASE_URL"
--target-url "$TARGET_URL"
--metrics-url "$METRICS_URL"
)
if [[ -n "$api_key_file" && -s "$api_key_file" ]]; then
preflight_args+=(--api-key-file "$api_key_file")
fi
"$script_dir/check_gateway_stage_preflight.js" \
"${preflight_args[@]}"
fi
echo "running $PRESSURE_STAGE gateway mock streaming pressure probe"
echo " target: $TARGET_URL"
echo " metrics: $METRICS_URL"
echo " requests: $PRESSURE_REQUESTS"
echo " concurrency: $PRESSURE_CONCURRENCY"
echo " hold ms: $PRESSURE_FIRST_BODY_HOLD_MS"
echo " ramp ms: $PRESSURE_START_RAMP_MS"
echo " settle ms: $PRESSURE_SETTLE_AFTER_MS"
echo " response mode: $PRESSURE_RESPONSE_MODE"
echo " cargo: $PRESSURE_CARGO_PROFILE"
echo " output: $OUTPUT"
if [[ "${PRESSURE_CAPTURE_METRICS_SNAPSHOTS:-true}" == "true" ]]; then
curl -fsS "$METRICS_URL" >"$metrics_before" || true
fi
# Use quiet cargo output so sensitive header values are not echoed back as part
# of Cargo's `Running ...` command line.
(cd "$repo_root" && cargo -q "${args[@]}")
if [[ "${PRESSURE_CAPTURE_METRICS_SNAPSHOTS:-true}" == "true" ]]; then
curl -fsS "$METRICS_URL" >"$metrics_after" || true
echo "metrics snapshots written to:"
echo " before: $metrics_before"
echo " after: $metrics_after"
fi
echo
echo "$PRESSURE_STAGE pressure report written to $OUTPUT"
if [[ "${PRESSURE_CHECK_REPORT:-true}" == "true" ]]; then
"$script_dir/check_gateway_stage_report.js" --stage "$PRESSURE_STAGE" "$OUTPUT"
fi
+241
View File
@@ -0,0 +1,241 @@
#!/usr/bin/env bash
set -euo pipefail
# Gateway realistic profile pressure probe.
#
# Profiles:
# realistic-stream: full-body streaming; use with mock upstream chunks/delay/payload set to realistic values.
# tps: no artificial hold; measures completed request throughput through auth + DB/Redis + usage/counter paths.
#
# Suggested mock upstream for realistic-stream:
# cargo run --release -p aether-testkit --bin mock_openai_upstream -- \
# --bind 127.0.0.1:18181 --chunks 80 --first-byte-delay-ms 150 --chunk-delay-ms 50 --payload-bytes 128
#
# Suggested mock upstream for tps:
# cargo run --release -p aether-testkit --bin mock_openai_upstream -- \
# --bind 127.0.0.1:18181 --chunks 8 --first-byte-delay-ms 20 --chunk-delay-ms 5 --payload-bytes 64
#
# Required auth:
# AETHER_API_KEY_FILE=/path/to/api-key
# or:
# AUTH_HEADER='Authorization: Bearer <aether-api-key>'
# or:
# AETHER_API_KEY='<aether-api-key>'
script_dir="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)"
repo_root="$(cd -- "$script_dir/../.." && pwd)"
PROFILE="${PRESSURE_PROFILE:-${1:-realistic-stream}}"
PROFILE="$(printf '%s' "$PROFILE" | tr '[:upper:]' '[:lower:]')"
case "$PROFILE" in
realistic-stream)
default_requests=1000
default_concurrency=1000
default_timeout_ms=300000
default_start_ramp_ms=10000
default_response_mode=full
default_settle_after_ms=5000
default_output=/tmp/aether_gateway_realistic_stream_1k.json
;;
tps)
default_requests=20000
default_concurrency=500
default_timeout_ms=180000
default_start_ramp_ms=5000
default_response_mode=full
default_settle_after_ms=5000
default_output=/tmp/aether_gateway_tps_20k_c500.json
;;
*)
echo "unsupported PRESSURE_PROFILE=$PROFILE; expected realistic-stream or tps" >&2
exit 2
;;
esac
GATEWAY_BASE_URL="${GATEWAY_BASE_URL:-http://127.0.0.1:8084}"
TARGET_URL="${TARGET_URL:-${GATEWAY_BASE_URL%/}/v1/chat/completions}"
METRICS_URL="${METRICS_URL:-${GATEWAY_BASE_URL%/}/_gateway/metrics}"
PRESSURE_REQUESTS="${PRESSURE_REQUESTS:-$default_requests}"
PRESSURE_CONCURRENCY="${PRESSURE_CONCURRENCY:-$default_concurrency}"
PRESSURE_TIMEOUT_MS="${PRESSURE_TIMEOUT_MS:-$default_timeout_ms}"
PRESSURE_CONNECT_TIMEOUT_MS="${PRESSURE_CONNECT_TIMEOUT_MS:-30000}"
PRESSURE_SAMPLE_INTERVAL_MS="${PRESSURE_SAMPLE_INTERVAL_MS:-500}"
PRESSURE_SETTLE_AFTER_MS="${PRESSURE_SETTLE_AFTER_MS:-$default_settle_after_ms}"
PRESSURE_START_RAMP_MS="${PRESSURE_START_RAMP_MS:-$default_start_ramp_ms}"
PRESSURE_FIRST_BODY_HOLD_MS="${PRESSURE_FIRST_BODY_HOLD_MS:-0}"
PRESSURE_METHOD="${PRESSURE_METHOD:-POST}"
PRESSURE_RESPONSE_MODE="${PRESSURE_RESPONSE_MODE:-$default_response_mode}"
PRESSURE_CARGO_PROFILE="${PRESSURE_CARGO_PROFILE:-release}"
PRESSURE_MODEL="${PRESSURE_MODEL:-gpt-5-mini}"
OUTPUT="${OUTPUT:-$default_output}"
api_key_file="${AETHER_API_KEY_FILE:-${API_KEY_FILE:-${PRESSURE_API_KEY_FILE:-}}}"
api_key_list_file="${AETHER_API_KEY_LIST_FILE:-${API_KEY_LIST_FILE:-${PRESSURE_API_KEY_LIST_FILE:-}}}"
PRESSURE_BODY_FILE="${PRESSURE_BODY_FILE:-/tmp/aether-pressure-${PROFILE}-request.json}"
if [[ -z "${AUTH_HEADER:-}" ]]; then
if [[ -n "$api_key_list_file" && -s "$api_key_list_file" ]]; then
:
elif [[ -n "$api_key_file" && -s "$api_key_file" ]]; then
:
elif [[ -n "${AETHER_API_KEY:-}" ]]; then
AUTH_HEADER="Authorization: Bearer ${AETHER_API_KEY}"
elif [[ -n "${API_KEY:-}" ]]; then
AUTH_HEADER="Authorization: Bearer ${API_KEY}"
else
echo "missing auth: set AETHER_API_KEY_FILE, AUTH_HEADER, or AETHER_API_KEY before running gateway realistic pressure" >&2
exit 2
fi
fi
if [[ -z "${PRESSURE_BODY:-}" && ! -s "$PRESSURE_BODY_FILE" ]]; then
cat >"$PRESSURE_BODY_FILE" <<JSON
{"model":"${PRESSURE_MODEL}","messages":[{"role":"system","content":"You are a concise assistant."},{"role":"user","content":"Write a practical deployment checklist for a high-concurrency API gateway. Include authentication, billing, observability, rollout, rollback, and incident handling."}],"stream":true}
JSON
fi
args=(run)
case "$PRESSURE_CARGO_PROFILE" in
release)
args+=(--release)
;;
debug)
;;
*)
echo "unsupported PRESSURE_CARGO_PROFILE=$PRESSURE_CARGO_PROFILE; expected release or debug" >&2
exit 2
;;
esac
args+=(
-p aether-testkit --bin gateway_pressure_probe --
--url "$TARGET_URL"
--metrics-url "$METRICS_URL"
--requests "$PRESSURE_REQUESTS"
--concurrency "$PRESSURE_CONCURRENCY"
--timeout-ms "$PRESSURE_TIMEOUT_MS"
--connect-timeout-ms "$PRESSURE_CONNECT_TIMEOUT_MS"
--sample-interval-ms "$PRESSURE_SAMPLE_INTERVAL_MS"
--settle-after-ms "$PRESSURE_SETTLE_AFTER_MS"
--start-ramp-ms "$PRESSURE_START_RAMP_MS"
--first-body-hold-ms "$PRESSURE_FIRST_BODY_HOLD_MS"
--method "$PRESSURE_METHOD"
--response-mode "$PRESSURE_RESPONSE_MODE"
--output "$OUTPUT"
)
if [[ -n "$api_key_list_file" && -s "$api_key_list_file" ]]; then
args+=(--api-key-list-file "$api_key_list_file")
elif [[ -n "$api_key_file" && -s "$api_key_file" ]]; then
args+=(--api-key-file "$api_key_file")
else
args+=(--header "$AUTH_HEADER")
fi
if [[ -n "${EXTRA_HEADERS:-}" ]]; then
while IFS= read -r header; do
[[ -z "$header" ]] && continue
args+=(--header "$header")
done <<< "$EXTRA_HEADERS"
else
args+=(--header "Content-Type: application/json")
fi
if [[ -n "${PRESSURE_CLIENT_SHARDS:-}" ]]; then
args+=(--client-shards "$PRESSURE_CLIENT_SHARDS")
fi
if [[ -n "${PRESSURE_POOL_MAX_IDLE_PER_HOST:-}" ]]; then
args+=(--pool-max-idle-per-host "$PRESSURE_POOL_MAX_IDLE_PER_HOST")
fi
if [[ -n "${PRESSURE_WARMUP_CONNECTIONS:-}" ]]; then
args+=(--warmup-connections "$PRESSURE_WARMUP_CONNECTIONS")
fi
if [[ -n "${PRESSURE_WARMUP_URL:-}" ]]; then
args+=(--warmup-url "$PRESSURE_WARMUP_URL")
fi
if [[ "${PRESSURE_HTTP1_ONLY:-false}" == "true" ]]; then
args+=(--http1-only)
fi
if [[ "${PRESSURE_HTTP2_PRIOR_KNOWLEDGE:-false}" == "true" ]]; then
args+=(--http2-prior-knowledge)
fi
if [[ -n "${PRESSURE_BODY:-}" ]]; then
args+=(--body "$PRESSURE_BODY")
else
args+=(--body-file "$PRESSURE_BODY_FILE")
fi
metrics_before="${OUTPUT%.json}.metrics.before.prom"
metrics_after="${OUTPUT%.json}.metrics.after.prom"
if [[ "${PRESSURE_PREFLIGHT:-true}" == "true" ]]; then
preflight_args=(
--stage "$PROFILE"
--gateway-base-url "$GATEWAY_BASE_URL"
--target-url "$TARGET_URL"
--metrics-url "$METRICS_URL"
)
if [[ -n "$api_key_list_file" && -s "$api_key_list_file" ]]; then
preflight_args+=(--api-key-list-file "$api_key_list_file")
elif [[ -n "$api_key_file" && -s "$api_key_file" ]]; then
preflight_args+=(--api-key-file "$api_key_file")
fi
"$script_dir/check_gateway_stage_preflight.js" "${preflight_args[@]}"
fi
echo "running $PROFILE gateway realistic pressure probe"
echo " target: $TARGET_URL"
echo " metrics: $METRICS_URL"
echo " requests: $PRESSURE_REQUESTS"
echo " concurrency: $PRESSURE_CONCURRENCY"
echo " response mode: $PRESSURE_RESPONSE_MODE"
echo " hold ms: $PRESSURE_FIRST_BODY_HOLD_MS"
echo " ramp ms: $PRESSURE_START_RAMP_MS"
echo " settle ms: $PRESSURE_SETTLE_AFTER_MS"
echo " cargo: $PRESSURE_CARGO_PROFILE"
echo " output: $OUTPUT"
if [[ "${PRESSURE_CAPTURE_METRICS_SNAPSHOTS:-true}" == "true" ]]; then
curl -fsS "$METRICS_URL" >"$metrics_before" || true
fi
(cd "$repo_root" && cargo -q "${args[@]}")
if [[ "${PRESSURE_CAPTURE_METRICS_SNAPSHOTS:-true}" == "true" ]]; then
curl -fsS "$METRICS_URL" >"$metrics_after" || true
echo "metrics snapshots written to:"
echo " before: $metrics_before"
echo " after: $metrics_after"
fi
echo
echo "$PROFILE pressure report written to $OUTPUT"
if [[ "${PRESSURE_CHECK_REPORT:-true}" == "true" ]]; then
check_args=(--stage "$PROFILE")
if [[ -n "${PRESSURE_MIN_THROUGHPUT_RPS:-}" ]]; then
check_args+=(--min-throughput-rps "$PRESSURE_MIN_THROUGHPUT_RPS")
fi
if [[ -n "${PRESSURE_MAX_HEADERS_P95_MS:-}" ]]; then
check_args+=(--max-headers-p95-ms "$PRESSURE_MAX_HEADERS_P95_MS")
fi
if [[ -n "${PRESSURE_MAX_FIRST_BODY_P95_MS:-}" ]]; then
check_args+=(--max-first-body-p95-ms "$PRESSURE_MAX_FIRST_BODY_P95_MS")
fi
if [[ -n "${PRESSURE_MAX_P95_MS:-}" ]]; then
check_args+=(--max-p95-ms "$PRESSURE_MAX_P95_MS")
fi
if [[ -n "${PRESSURE_MAX_P99_MS:-}" ]]; then
check_args+=(--max-p99-ms "$PRESSURE_MAX_P99_MS")
fi
if [[ -n "${PRESSURE_MAX_FIRST_BODY_HOLD_MS:-}" ]]; then
check_args+=(--max-first-body-hold-ms "$PRESSURE_MAX_FIRST_BODY_HOLD_MS")
fi
"$script_dir/check_gateway_stage_report.js" "${check_args[@]}" "$OUTPUT"
fi