mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-08 20:20:19 +08:00
chore: update gateway pressure observability
This commit is contained in:
+11
-23
@@ -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
@@ -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,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>,
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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
@@ -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);
|
||||
|
||||
@@ -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();
|
||||
}
|
||||
|
||||
@@ -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,
|
||||
},
|
||||
);
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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();
|
||||
}
|
||||
|
||||
@@ -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"));
|
||||
}
|
||||
}
|
||||
@@ -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,
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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,
|
||||
};
|
||||
|
||||
@@ -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> {
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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(),
|
||||
}
|
||||
);
|
||||
|
||||
@@ -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));
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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;
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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"
|
||||
);
|
||||
}
|
||||
@@ -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!(
|
||||
|
||||
@@ -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")
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -8,6 +8,7 @@ pub struct UsageRuntimeConfig {
|
||||
pub worker_count: usize,
|
||||
pub worker_autoscale_enabled: bool,
|
||||
pub worker_max_count: usize,
|
||||
pub worker_record_concurrency_limit: Option<usize>,
|
||||
pub worker_scale_interval_ms: u64,
|
||||
pub worker_idle_scale_down_ticks: u64,
|
||||
pub stream_key: String,
|
||||
@@ -21,6 +22,7 @@ pub struct UsageRuntimeConfig {
|
||||
pub reclaim_interval_ms: u64,
|
||||
pub terminal_enqueue_max_in_flight: u64,
|
||||
pub lifecycle_enqueue_max_in_flight: u64,
|
||||
pub lifecycle_enqueue_delay_ms: u64,
|
||||
pub retry_deferred_lifecycle_events: bool,
|
||||
pub enqueue_retry_buffer_capacity: usize,
|
||||
pub enqueue_retry_workers: usize,
|
||||
@@ -36,23 +38,25 @@ impl Default for UsageRuntimeConfig {
|
||||
queue_lifecycle_events: false,
|
||||
worker_count: 4,
|
||||
worker_autoscale_enabled: true,
|
||||
worker_max_count: 64,
|
||||
worker_max_count: 32,
|
||||
worker_record_concurrency_limit: Some(32),
|
||||
worker_scale_interval_ms: 1_000,
|
||||
worker_idle_scale_down_ticks: 30,
|
||||
stream_key: "usage:events".to_string(),
|
||||
consumer_group: "usage_consumers".to_string(),
|
||||
dlq_stream_key: "usage:events:dlq".to_string(),
|
||||
stream_maxlen: 200_000,
|
||||
consumer_batch_size: 500,
|
||||
consumer_batch_size: 128,
|
||||
consumer_block_ms: 500,
|
||||
reclaim_idle_ms: 30_000,
|
||||
reclaim_count: 500,
|
||||
reclaim_idle_ms: 60_000,
|
||||
reclaim_count: 128,
|
||||
reclaim_interval_ms: 5_000,
|
||||
terminal_enqueue_max_in_flight: 256,
|
||||
lifecycle_enqueue_max_in_flight: 128,
|
||||
retry_deferred_lifecycle_events: false,
|
||||
terminal_enqueue_max_in_flight: 1_024,
|
||||
lifecycle_enqueue_max_in_flight: 512,
|
||||
lifecycle_enqueue_delay_ms: 1_000,
|
||||
retry_deferred_lifecycle_events: true,
|
||||
enqueue_retry_buffer_capacity: 131_072,
|
||||
enqueue_retry_workers: 4,
|
||||
enqueue_retry_workers: 8,
|
||||
enqueue_retry_initial_backoff_ms: 3_000,
|
||||
enqueue_retry_max_backoff_ms: 10_000,
|
||||
}
|
||||
@@ -94,6 +98,15 @@ impl UsageRuntimeConfig {
|
||||
"usage runtime worker_max_count must be positive".to_string(),
|
||||
));
|
||||
}
|
||||
if self
|
||||
.worker_record_concurrency_limit
|
||||
.is_some_and(|limit| limit == 0)
|
||||
{
|
||||
return Err(DataLayerError::InvalidConfiguration(
|
||||
"usage runtime worker_record_concurrency_limit must be positive when set"
|
||||
.to_string(),
|
||||
));
|
||||
}
|
||||
if self.worker_scale_interval_ms == 0 {
|
||||
return Err(DataLayerError::InvalidConfiguration(
|
||||
"usage runtime worker_scale_interval_ms must be positive".to_string(),
|
||||
|
||||
@@ -1,7 +1,10 @@
|
||||
use std::future::Future;
|
||||
use std::sync::OnceLock;
|
||||
|
||||
const USAGE_BACKGROUND_RUNTIME_THREADS: usize = 2;
|
||||
const DEFAULT_USAGE_BACKGROUND_RUNTIME_THREADS: usize = 8;
|
||||
const MAX_USAGE_BACKGROUND_RUNTIME_THREADS: usize = 64;
|
||||
const GATEWAY_USAGE_BACKGROUND_RUNTIME_THREADS_ENV: &str = "AETHER_GATEWAY_USAGE_RUNTIME_THREADS";
|
||||
const USAGE_BACKGROUND_RUNTIME_THREADS_ENV: &str = "AETHER_USAGE_RUNTIME_THREADS";
|
||||
const USAGE_BACKGROUND_RUNTIME_STACK_BYTES: usize = 8 * 1024 * 1024;
|
||||
const USAGE_BACKGROUND_RUNTIME_THREAD_NAME: &str = "aether-usage-runtime";
|
||||
|
||||
@@ -19,7 +22,7 @@ fn usage_background_runtime() -> &'static tokio::runtime::Runtime {
|
||||
RUNTIME.get_or_init(|| {
|
||||
let runtime = tokio::runtime::Builder::new_multi_thread()
|
||||
.enable_all()
|
||||
.worker_threads(USAGE_BACKGROUND_RUNTIME_THREADS)
|
||||
.worker_threads(usage_background_runtime_threads())
|
||||
.thread_name(USAGE_BACKGROUND_RUNTIME_THREAD_NAME)
|
||||
.thread_stack_size(USAGE_BACKGROUND_RUNTIME_STACK_BYTES)
|
||||
.build()
|
||||
@@ -28,9 +31,26 @@ fn usage_background_runtime() -> &'static tokio::runtime::Runtime {
|
||||
})
|
||||
}
|
||||
|
||||
fn usage_background_runtime_threads() -> usize {
|
||||
parse_usage_background_runtime_threads(
|
||||
std::env::var(GATEWAY_USAGE_BACKGROUND_RUNTIME_THREADS_ENV)
|
||||
.ok()
|
||||
.or_else(|| std::env::var(USAGE_BACKGROUND_RUNTIME_THREADS_ENV).ok())
|
||||
.as_deref(),
|
||||
)
|
||||
}
|
||||
|
||||
fn parse_usage_background_runtime_threads(value: Option<&str>) -> usize {
|
||||
value
|
||||
.and_then(|value| value.trim().parse::<usize>().ok())
|
||||
.filter(|threads| *threads > 0)
|
||||
.unwrap_or(DEFAULT_USAGE_BACKGROUND_RUNTIME_THREADS)
|
||||
.clamp(1, MAX_USAGE_BACKGROUND_RUNTIME_THREADS)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::spawn_on_usage_background_runtime;
|
||||
use super::{parse_usage_background_runtime_threads, spawn_on_usage_background_runtime};
|
||||
|
||||
#[tokio::test]
|
||||
async fn usage_background_runtime_runs_on_dedicated_named_threads() {
|
||||
@@ -45,4 +65,16 @@ mod tests {
|
||||
|
||||
assert_eq!(thread_name, "aether-usage-runtime");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn usage_background_runtime_threads_are_configurable() {
|
||||
assert_eq!(parse_usage_background_runtime_threads(None), 8);
|
||||
assert_eq!(parse_usage_background_runtime_threads(Some("12")), 12);
|
||||
assert_eq!(parse_usage_background_runtime_threads(Some("0")), 8);
|
||||
assert_eq!(
|
||||
parse_usage_background_runtime_threads(Some("not-a-number")),
|
||||
8
|
||||
);
|
||||
assert_eq!(parse_usage_background_runtime_threads(Some("999")), 64);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -40,8 +40,8 @@ pub use report_context::{
|
||||
build_locally_actionable_report_context_from_video_task, report_context_is_locally_actionable,
|
||||
};
|
||||
pub use runtime::{
|
||||
UsageBillingEventEnricher, UsageBodyCapturePolicy, UsageRequestRecordLevel, UsageRuntime,
|
||||
UsageRuntimeAccess, UsageRuntimeMetricsSnapshot,
|
||||
UsageBillingEventEnricher, UsageBodyCapturePolicy, UsageQueueHealthSnapshot,
|
||||
UsageRequestRecordLevel, UsageRuntime, UsageRuntimeAccess, UsageRuntimeMetricsSnapshot,
|
||||
DEFAULT_USAGE_REQUEST_BODY_CAPTURE_LIMIT_BYTES,
|
||||
DEFAULT_USAGE_RESPONSE_BODY_CAPTURE_LIMIT_BYTES,
|
||||
};
|
||||
|
||||
@@ -3,7 +3,9 @@ use std::sync::Arc;
|
||||
use serde_json::json;
|
||||
|
||||
use aether_data_contracts::DataLayerError;
|
||||
use aether_runtime_state::{RuntimeQueueEntry, RuntimeQueueReclaimConfig, RuntimeQueueStore};
|
||||
use aether_runtime_state::{
|
||||
RuntimeQueueEntry, RuntimeQueueReclaimConfig, RuntimeQueueStats, RuntimeQueueStore,
|
||||
};
|
||||
|
||||
use super::config::UsageRuntimeConfig;
|
||||
use super::event::UsageEvent;
|
||||
@@ -103,6 +105,14 @@ impl UsageQueue {
|
||||
.append_fields_with_maxlen(&self.dlq_stream, &fields, None)
|
||||
.await
|
||||
}
|
||||
|
||||
pub async fn stats(&self) -> Result<RuntimeQueueStats, DataLayerError> {
|
||||
self.runner.stats(&self.stream, Some(&self.group)).await
|
||||
}
|
||||
|
||||
pub async fn dlq_stats(&self) -> Result<RuntimeQueueStats, DataLayerError> {
|
||||
self.runner.stats(&self.dlq_stream, None).await
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -1,3 +1,7 @@
|
||||
use std::collections::hash_map::DefaultHasher;
|
||||
use std::hash::{Hash, Hasher};
|
||||
use std::sync::OnceLock;
|
||||
|
||||
use aether_data_contracts::repository::settlement::{StoredUsageSettlement, UsageSettlementInput};
|
||||
use aether_data_contracts::repository::usage::StoredRequestUsageAudit;
|
||||
use aether_data_contracts::{DataLayerError, DataLayerError::InvalidInput};
|
||||
@@ -39,10 +43,50 @@ pub async fn settle_usage_if_needed(
|
||||
actual_total_cost_usd: finite_cost(usage.actual_total_cost_usd)?,
|
||||
finalized_at_unix_secs,
|
||||
};
|
||||
let settlement_key = usage_settlement_lock_key(&input);
|
||||
let _guard = usage_settlement_lock(&settlement_key).lock().await;
|
||||
let _ = writer.settle_usage(input).await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
const USAGE_SETTLEMENT_LOCK_SHARDS: usize = 4096;
|
||||
|
||||
fn usage_settlement_lock(key: &str) -> &'static tokio::sync::Mutex<()> {
|
||||
static LOCKS: OnceLock<Vec<tokio::sync::Mutex<()>>> = OnceLock::new();
|
||||
let locks = LOCKS.get_or_init(|| {
|
||||
(0..USAGE_SETTLEMENT_LOCK_SHARDS)
|
||||
.map(|_| tokio::sync::Mutex::new(()))
|
||||
.collect()
|
||||
});
|
||||
&locks[usage_settlement_lock_shard(key)]
|
||||
}
|
||||
|
||||
fn usage_settlement_lock_shard(key: &str) -> usize {
|
||||
let mut hasher = DefaultHasher::new();
|
||||
key.hash(&mut hasher);
|
||||
(hasher.finish() as usize) % USAGE_SETTLEMENT_LOCK_SHARDS
|
||||
}
|
||||
|
||||
fn usage_settlement_lock_key(input: &UsageSettlementInput) -> String {
|
||||
if input.api_key_is_standalone {
|
||||
if let Some(api_key_id) = input.api_key_id.as_deref().and_then(non_empty_trimmed) {
|
||||
return format!("api-key:{api_key_id}");
|
||||
}
|
||||
}
|
||||
if let Some(user_id) = input.user_id.as_deref().and_then(non_empty_trimmed) {
|
||||
return format!("user:{user_id}");
|
||||
}
|
||||
if let Some(api_key_id) = input.api_key_id.as_deref().and_then(non_empty_trimmed) {
|
||||
return format!("api-key:{api_key_id}");
|
||||
}
|
||||
format!("request:{}", input.request_id.trim())
|
||||
}
|
||||
|
||||
fn non_empty_trimmed(value: &str) -> Option<&str> {
|
||||
let value = value.trim();
|
||||
(!value.is_empty()).then_some(value)
|
||||
}
|
||||
|
||||
fn usage_api_key_is_standalone(usage: &StoredRequestUsageAudit) -> bool {
|
||||
usage
|
||||
.request_metadata
|
||||
@@ -64,7 +108,9 @@ fn finite_cost(value: f64) -> Result<f64, DataLayerError> {
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::sync::atomic::{AtomicUsize, Ordering};
|
||||
use std::sync::Mutex;
|
||||
use std::time::Duration;
|
||||
|
||||
use super::{settle_usage_if_needed, UsageSettlementWriter};
|
||||
use aether_data_contracts::repository::settlement::UsageSettlementInput;
|
||||
@@ -78,6 +124,13 @@ mod tests {
|
||||
inputs: Mutex<Vec<UsageSettlementInput>>,
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
struct SlowSettlementWriter {
|
||||
active: AtomicUsize,
|
||||
max_active: AtomicUsize,
|
||||
inputs: Mutex<Vec<UsageSettlementInput>>,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl UsageSettlementWriter for TestSettlementWriter {
|
||||
fn has_usage_settlement_writer(&self) -> bool {
|
||||
@@ -99,6 +152,31 @@ mod tests {
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl UsageSettlementWriter for SlowSettlementWriter {
|
||||
fn has_usage_settlement_writer(&self) -> bool {
|
||||
true
|
||||
}
|
||||
|
||||
async fn settle_usage(
|
||||
&self,
|
||||
input: UsageSettlementInput,
|
||||
) -> Result<
|
||||
Option<aether_data_contracts::repository::settlement::StoredUsageSettlement>,
|
||||
aether_data_contracts::DataLayerError,
|
||||
> {
|
||||
let active = self.active.fetch_add(1, Ordering::AcqRel) + 1;
|
||||
self.max_active.fetch_max(active, Ordering::AcqRel);
|
||||
tokio::time::sleep(Duration::from_millis(30)).await;
|
||||
self.inputs
|
||||
.lock()
|
||||
.expect("settlement inputs lock")
|
||||
.push(input);
|
||||
self.active.fetch_sub(1, Ordering::AcqRel);
|
||||
Ok(None)
|
||||
}
|
||||
}
|
||||
|
||||
fn sample_usage() -> StoredRequestUsageAudit {
|
||||
StoredRequestUsageAudit::new(
|
||||
"usage-1".to_string(),
|
||||
@@ -238,4 +316,22 @@ mod tests {
|
||||
let inputs = writer.inputs.lock().expect("settlement inputs lock");
|
||||
assert!(inputs.is_empty());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn serializes_settlements_for_same_billing_subject() {
|
||||
let writer = SlowSettlementWriter::default();
|
||||
let mut first = sample_usage();
|
||||
first.request_id = "req-same-subject-1".to_string();
|
||||
let mut second = sample_usage();
|
||||
second.request_id = "req-same-subject-2".to_string();
|
||||
|
||||
tokio::try_join!(
|
||||
settle_usage_if_needed(&writer, &first),
|
||||
settle_usage_if_needed(&writer, &second)
|
||||
)
|
||||
.expect("settlements should succeed");
|
||||
|
||||
assert_eq!(writer.max_active.load(Ordering::Acquire), 1);
|
||||
assert_eq!(writer.inputs.lock().expect("inputs lock").len(), 2);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,5 +1,7 @@
|
||||
use std::collections::hash_map::DefaultHasher;
|
||||
use std::hash::{Hash, Hasher};
|
||||
use std::sync::atomic::{AtomicBool, Ordering};
|
||||
use std::sync::Arc;
|
||||
use std::sync::{Arc, OnceLock};
|
||||
use std::time::Duration;
|
||||
|
||||
use aether_data_contracts::repository::usage::{StoredRequestUsageAudit, UpsertUsageRecord};
|
||||
@@ -10,12 +12,17 @@ use tokio::sync::mpsc;
|
||||
use tracing::warn;
|
||||
|
||||
use crate::executor::spawn_on_usage_background_runtime;
|
||||
use crate::runtime::UsageBillingEventEnricher;
|
||||
use crate::runtime::{
|
||||
UsageBillingEventEnricher, UsageRuntimeAccess, UsageWorkerRecordConcurrencyGate,
|
||||
};
|
||||
use crate::{
|
||||
build_upsert_usage_record_from_event, settle_usage_if_needed, UsageEvent, UsageEventType,
|
||||
UsageQueue, UsageRuntimeConfig, UsageSettlementWriter,
|
||||
};
|
||||
|
||||
const USAGE_WORKER_DB_PRESSURE_DEFER_MS: u64 = 10;
|
||||
const USAGE_WORKER_ACK_CHUNK_SIZE: usize = 100;
|
||||
|
||||
#[async_trait]
|
||||
pub trait UsageEventRecorder: Send + Sync {
|
||||
async fn record_usage_event(&self, event: &UsageEvent) -> Result<(), DataLayerError>;
|
||||
@@ -42,31 +49,81 @@ pub trait UsageRecordWriter: Send + Sync {
|
||||
|
||||
pub struct UsageDataEventRecorder<T> {
|
||||
data: Arc<T>,
|
||||
record_gate: Option<Arc<UsageWorkerRecordConcurrencyGate>>,
|
||||
defer_for_database_pressure: bool,
|
||||
}
|
||||
|
||||
impl<T> UsageDataEventRecorder<T> {
|
||||
pub fn new(data: Arc<T>) -> Self {
|
||||
Self { data }
|
||||
Self::with_record_gate(data, None)
|
||||
}
|
||||
|
||||
pub(crate) fn with_record_gate(
|
||||
data: Arc<T>,
|
||||
record_gate: Option<Arc<UsageWorkerRecordConcurrencyGate>>,
|
||||
) -> Self {
|
||||
Self {
|
||||
data,
|
||||
record_gate,
|
||||
defer_for_database_pressure: false,
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn with_record_gate_and_database_pressure_defer(
|
||||
data: Arc<T>,
|
||||
record_gate: Option<Arc<UsageWorkerRecordConcurrencyGate>>,
|
||||
) -> Self {
|
||||
Self {
|
||||
data,
|
||||
record_gate,
|
||||
defer_for_database_pressure: true,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl<T> UsageEventRecorder for UsageDataEventRecorder<T>
|
||||
where
|
||||
T: UsageRecordWriter
|
||||
+ UsageSettlementWriter
|
||||
+ UsageBillingEventEnricher
|
||||
+ ManualProxyNodeCounter
|
||||
+ Send
|
||||
+ Sync,
|
||||
T: UsageRuntimeAccess,
|
||||
{
|
||||
async fn record_usage_event(&self, event: &UsageEvent) -> Result<(), DataLayerError> {
|
||||
if self.defer_for_database_pressure
|
||||
&& self.data.usage_worker_should_defer_for_database_pressure()
|
||||
{
|
||||
if let Some(gate) = self.record_gate.as_ref() {
|
||||
gate.record_deferred();
|
||||
}
|
||||
tokio::time::sleep(Duration::from_millis(USAGE_WORKER_DB_PRESSURE_DEFER_MS)).await;
|
||||
}
|
||||
let _record_gate_permit = match self.record_gate.as_ref() {
|
||||
Some(gate) => Some(gate.acquire().await),
|
||||
None => None,
|
||||
};
|
||||
let _guard = usage_request_lock(&event.request_id).lock().await;
|
||||
let mut event = event.clone();
|
||||
enrich_terminal_event(self.data.as_ref(), &mut event).await;
|
||||
write_event_record(self.data.as_ref(), &event).await
|
||||
}
|
||||
}
|
||||
|
||||
const USAGE_REQUEST_LOCK_SHARDS: usize = 4096;
|
||||
|
||||
fn usage_request_lock(request_id: &str) -> &'static tokio::sync::Mutex<()> {
|
||||
static LOCKS: OnceLock<Vec<tokio::sync::Mutex<()>>> = OnceLock::new();
|
||||
let locks = LOCKS.get_or_init(|| {
|
||||
(0..USAGE_REQUEST_LOCK_SHARDS)
|
||||
.map(|_| tokio::sync::Mutex::new(()))
|
||||
.collect()
|
||||
});
|
||||
&locks[usage_request_lock_shard(request_id)]
|
||||
}
|
||||
|
||||
fn usage_request_lock_shard(request_id: &str) -> usize {
|
||||
let mut hasher = DefaultHasher::new();
|
||||
request_id.hash(&mut hasher);
|
||||
(hasher.finish() as usize) % USAGE_REQUEST_LOCK_SHARDS
|
||||
}
|
||||
|
||||
pub struct UsageQueueWorker {
|
||||
queue: UsageQueue,
|
||||
recorder: Arc<dyn UsageEventRecorder>,
|
||||
@@ -97,6 +154,112 @@ pub(crate) struct UsageWorkerObservation {
|
||||
pub worker_index: Option<usize>,
|
||||
pub entries_read: usize,
|
||||
pub batch_size: usize,
|
||||
pub reclaimed_entries: usize,
|
||||
pub acked_entries: usize,
|
||||
pub dead_lettered_entries: usize,
|
||||
pub process_failures: usize,
|
||||
pub read_failures: usize,
|
||||
pub reclaim_failures: usize,
|
||||
}
|
||||
|
||||
impl UsageWorkerObservation {
|
||||
fn read(worker_index: Option<usize>, entries_read: usize, batch_size: usize) -> Self {
|
||||
Self {
|
||||
worker_index,
|
||||
entries_read,
|
||||
batch_size,
|
||||
reclaimed_entries: 0,
|
||||
acked_entries: 0,
|
||||
dead_lettered_entries: 0,
|
||||
process_failures: 0,
|
||||
read_failures: 0,
|
||||
reclaim_failures: 0,
|
||||
}
|
||||
}
|
||||
|
||||
fn reclaimed(worker_index: Option<usize>, reclaimed_entries: usize) -> Self {
|
||||
Self {
|
||||
worker_index,
|
||||
entries_read: 0,
|
||||
batch_size: 0,
|
||||
reclaimed_entries,
|
||||
acked_entries: 0,
|
||||
dead_lettered_entries: 0,
|
||||
process_failures: 0,
|
||||
read_failures: 0,
|
||||
reclaim_failures: 0,
|
||||
}
|
||||
}
|
||||
|
||||
fn acked(worker_index: Option<usize>, acked_entries: usize) -> Self {
|
||||
Self {
|
||||
worker_index,
|
||||
entries_read: 0,
|
||||
batch_size: 0,
|
||||
reclaimed_entries: 0,
|
||||
acked_entries,
|
||||
dead_lettered_entries: 0,
|
||||
process_failures: 0,
|
||||
read_failures: 0,
|
||||
reclaim_failures: 0,
|
||||
}
|
||||
}
|
||||
|
||||
fn dead_lettered(worker_index: Option<usize>, dead_lettered_entries: usize) -> Self {
|
||||
Self {
|
||||
worker_index,
|
||||
entries_read: 0,
|
||||
batch_size: 0,
|
||||
reclaimed_entries: 0,
|
||||
acked_entries: 0,
|
||||
dead_lettered_entries,
|
||||
process_failures: 0,
|
||||
read_failures: 0,
|
||||
reclaim_failures: 0,
|
||||
}
|
||||
}
|
||||
|
||||
fn process_failed(worker_index: Option<usize>) -> Self {
|
||||
Self {
|
||||
worker_index,
|
||||
entries_read: 0,
|
||||
batch_size: 0,
|
||||
reclaimed_entries: 0,
|
||||
acked_entries: 0,
|
||||
dead_lettered_entries: 0,
|
||||
process_failures: 1,
|
||||
read_failures: 0,
|
||||
reclaim_failures: 0,
|
||||
}
|
||||
}
|
||||
|
||||
fn read_failed(worker_index: Option<usize>) -> Self {
|
||||
Self {
|
||||
worker_index,
|
||||
entries_read: 0,
|
||||
batch_size: 0,
|
||||
reclaimed_entries: 0,
|
||||
acked_entries: 0,
|
||||
dead_lettered_entries: 0,
|
||||
process_failures: 0,
|
||||
read_failures: 1,
|
||||
reclaim_failures: 0,
|
||||
}
|
||||
}
|
||||
|
||||
fn reclaim_failed(worker_index: Option<usize>) -> Self {
|
||||
Self {
|
||||
worker_index,
|
||||
entries_read: 0,
|
||||
batch_size: 0,
|
||||
reclaimed_entries: 0,
|
||||
acked_entries: 0,
|
||||
dead_lettered_entries: 0,
|
||||
process_failures: 0,
|
||||
read_failures: 0,
|
||||
reclaim_failures: 1,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl UsageQueueWorker {
|
||||
@@ -163,7 +326,9 @@ impl UsageQueueWorker {
|
||||
_ = reclaim_interval.tick() => {
|
||||
match self.queue.claim_stale(&self.consumer, "0-0").await {
|
||||
Ok(entries) => {
|
||||
self.report_reclaimed(entries.len());
|
||||
if let Err(err) = self.process_entries(entries).await {
|
||||
self.report_process_failed();
|
||||
warn!(
|
||||
event_name = "usage_worker_reclaim_process_failed",
|
||||
log_type = "ops",
|
||||
@@ -174,14 +339,17 @@ impl UsageQueueWorker {
|
||||
);
|
||||
}
|
||||
}
|
||||
Err(err) => warn!(
|
||||
event_name = "usage_worker_reclaim_failed",
|
||||
log_type = "ops",
|
||||
worker_consumer = %self.consumer,
|
||||
worker_group = %self.config.consumer_group,
|
||||
error = %err,
|
||||
"usage worker failed to reclaim stale entries"
|
||||
),
|
||||
Err(err) => {
|
||||
self.report_reclaim_failed();
|
||||
warn!(
|
||||
event_name = "usage_worker_reclaim_failed",
|
||||
log_type = "ops",
|
||||
worker_consumer = %self.consumer,
|
||||
worker_group = %self.config.consumer_group,
|
||||
error = %err,
|
||||
"usage worker failed to reclaim stale entries"
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
result = self.queue.read_group(&self.consumer) => {
|
||||
@@ -192,6 +360,7 @@ impl UsageQueueWorker {
|
||||
break;
|
||||
}
|
||||
if let Err(err) = self.process_entries(entries).await {
|
||||
self.report_process_failed();
|
||||
warn!(
|
||||
event_name = "usage_worker_process_failed",
|
||||
log_type = "ops",
|
||||
@@ -207,6 +376,7 @@ impl UsageQueueWorker {
|
||||
}
|
||||
}
|
||||
Err(err) => {
|
||||
self.report_read_failed();
|
||||
warn!(
|
||||
event_name = "usage_worker_read_failed",
|
||||
log_type = "ops",
|
||||
@@ -230,14 +400,51 @@ impl UsageQueueWorker {
|
||||
}
|
||||
|
||||
fn report_read(&self, entries_read: usize) {
|
||||
self.report(UsageWorkerObservation::read(
|
||||
self.worker_index,
|
||||
entries_read,
|
||||
self.config.consumer_batch_size.max(1),
|
||||
));
|
||||
}
|
||||
|
||||
fn report_reclaimed(&self, reclaimed_entries: usize) {
|
||||
self.report(UsageWorkerObservation::reclaimed(
|
||||
self.worker_index,
|
||||
reclaimed_entries,
|
||||
));
|
||||
}
|
||||
|
||||
fn report_acked(&self, acked_entries: usize) {
|
||||
self.report(UsageWorkerObservation::acked(
|
||||
self.worker_index,
|
||||
acked_entries,
|
||||
));
|
||||
}
|
||||
|
||||
fn report_dead_lettered(&self, dead_lettered_entries: usize) {
|
||||
self.report(UsageWorkerObservation::dead_lettered(
|
||||
self.worker_index,
|
||||
dead_lettered_entries,
|
||||
));
|
||||
}
|
||||
|
||||
fn report_process_failed(&self) {
|
||||
self.report(UsageWorkerObservation::process_failed(self.worker_index));
|
||||
}
|
||||
|
||||
fn report_read_failed(&self) {
|
||||
self.report(UsageWorkerObservation::read_failed(self.worker_index));
|
||||
}
|
||||
|
||||
fn report_reclaim_failed(&self) {
|
||||
self.report(UsageWorkerObservation::reclaim_failed(self.worker_index));
|
||||
}
|
||||
|
||||
fn report(&self, observation: UsageWorkerObservation) {
|
||||
let Some(telemetry) = &self.telemetry else {
|
||||
return;
|
||||
};
|
||||
let _ = telemetry.try_send(UsageWorkerObservation {
|
||||
worker_index: self.worker_index,
|
||||
entries_read,
|
||||
batch_size: self.config.consumer_batch_size.max(1),
|
||||
});
|
||||
let _ = telemetry.try_send(observation);
|
||||
}
|
||||
|
||||
async fn process_entries(&self, entries: Vec<RuntimeQueueEntry>) -> Result<(), DataLayerError> {
|
||||
@@ -251,6 +458,11 @@ impl UsageQueueWorker {
|
||||
Ok(should_ack) => {
|
||||
if should_ack {
|
||||
ack_ids.push(entry.id.clone());
|
||||
if ack_ids.len() >= USAGE_WORKER_ACK_CHUNK_SIZE {
|
||||
self.queue.ack_and_delete(&ack_ids).await?;
|
||||
self.report_acked(ack_ids.len());
|
||||
ack_ids.clear();
|
||||
}
|
||||
}
|
||||
}
|
||||
Err(err) => {
|
||||
@@ -264,6 +476,7 @@ impl UsageQueueWorker {
|
||||
|
||||
if !ack_ids.is_empty() {
|
||||
self.queue.ack_and_delete(&ack_ids).await?;
|
||||
self.report_acked(ack_ids.len());
|
||||
}
|
||||
|
||||
Ok(())
|
||||
@@ -283,6 +496,7 @@ impl UsageQueueWorker {
|
||||
"usage worker moved malformed queue entry to dead letter"
|
||||
);
|
||||
self.queue.push_dead_letter(entry, &err.to_string()).await?;
|
||||
self.report_dead_lettered(1);
|
||||
return Ok(true);
|
||||
}
|
||||
};
|
||||
@@ -308,6 +522,7 @@ impl UsageQueueWorker {
|
||||
"usage worker moved non-retryable usage event to dead letter"
|
||||
);
|
||||
self.queue.push_dead_letter(entry, &err.to_string()).await?;
|
||||
self.report_dead_lettered(1);
|
||||
Ok(true)
|
||||
}
|
||||
Err(err) => {
|
||||
@@ -357,17 +572,26 @@ pub fn build_usage_queue_worker<T>(
|
||||
worker_index: Option<usize>,
|
||||
) -> Result<UsageQueueWorker, DataLayerError>
|
||||
where
|
||||
T: UsageRecordWriter
|
||||
+ UsageSettlementWriter
|
||||
+ UsageBillingEventEnricher
|
||||
+ ManualProxyNodeCounter
|
||||
+ Send
|
||||
+ Sync
|
||||
+ 'static,
|
||||
T: UsageRuntimeAccess + 'static,
|
||||
{
|
||||
build_usage_queue_worker_with_record_gate(runner, data, config, None, worker_index)
|
||||
}
|
||||
|
||||
pub(crate) fn build_usage_queue_worker_with_record_gate<T>(
|
||||
runner: Arc<dyn RuntimeQueueStore>,
|
||||
data: Arc<T>,
|
||||
config: UsageRuntimeConfig,
|
||||
record_gate: Option<Arc<UsageWorkerRecordConcurrencyGate>>,
|
||||
worker_index: Option<usize>,
|
||||
) -> Result<UsageQueueWorker, DataLayerError>
|
||||
where
|
||||
T: UsageRuntimeAccess + 'static,
|
||||
{
|
||||
UsageQueueWorker::new(
|
||||
runner,
|
||||
Arc::new(UsageDataEventRecorder::new(data)),
|
||||
Arc::new(
|
||||
UsageDataEventRecorder::with_record_gate_and_database_pressure_defer(data, record_gate),
|
||||
),
|
||||
config,
|
||||
worker_index,
|
||||
)
|
||||
@@ -472,7 +696,9 @@ fn consumer_name(worker_index: Option<usize>) -> String {
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::sync::atomic::{AtomicBool, Ordering};
|
||||
use std::sync::{Arc, Mutex};
|
||||
use std::time::Duration;
|
||||
|
||||
use aether_data_contracts::repository::settlement::{
|
||||
StoredUsageSettlement, UsageSettlementInput,
|
||||
@@ -483,12 +709,15 @@ mod tests {
|
||||
use async_trait::async_trait;
|
||||
|
||||
use super::{
|
||||
usage_event_record_error_is_permanent, write_event_record, ManualProxyNodeCounter,
|
||||
UsageEventRecorder, UsageQueueWorker, UsageRecordWriter,
|
||||
build_usage_queue_worker_with_record_gate, usage_event_record_error_is_permanent,
|
||||
write_event_record, ManualProxyNodeCounter, UsageEventRecorder, UsageQueueWorker,
|
||||
UsageRecordWriter,
|
||||
};
|
||||
use crate::runtime::UsageWorkerRecordConcurrencyGate;
|
||||
use crate::UsageBillingEventEnricher;
|
||||
use crate::{
|
||||
UsageEvent, UsageEventData, UsageEventType, UsageRuntimeConfig, UsageSettlementWriter,
|
||||
UsageEvent, UsageEventData, UsageEventType, UsageQueue, UsageRuntimeConfig,
|
||||
UsageSettlementWriter,
|
||||
};
|
||||
|
||||
#[derive(Default)]
|
||||
@@ -503,6 +732,14 @@ mod tests {
|
||||
calls: Mutex<Vec<String>>,
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
struct SlowUsageStore {
|
||||
active: std::sync::atomic::AtomicUsize,
|
||||
max_active: std::sync::atomic::AtomicUsize,
|
||||
db_pressure: AtomicBool,
|
||||
records: Mutex<Vec<String>>,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl UsageRecordWriter for TestUsageStore {
|
||||
async fn upsert_usage_record(
|
||||
@@ -603,6 +840,95 @@ mod tests {
|
||||
}
|
||||
}
|
||||
|
||||
impl crate::runtime::UsageRuntimeAccess for TestUsageStore {
|
||||
fn has_usage_writer(&self) -> bool {
|
||||
true
|
||||
}
|
||||
|
||||
fn has_usage_worker_queue(&self) -> bool {
|
||||
true
|
||||
}
|
||||
|
||||
fn usage_worker_queue(&self) -> Option<Arc<dyn RuntimeQueueStore>> {
|
||||
None
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl UsageRecordWriter for SlowUsageStore {
|
||||
async fn upsert_usage_record(
|
||||
&self,
|
||||
record: UpsertUsageRecord,
|
||||
) -> Result<Option<StoredRequestUsageAudit>, DataLayerError> {
|
||||
let active = self
|
||||
.active
|
||||
.fetch_add(1, std::sync::atomic::Ordering::AcqRel)
|
||||
+ 1;
|
||||
self.max_active
|
||||
.fetch_max(active, std::sync::atomic::Ordering::AcqRel);
|
||||
tokio::time::sleep(Duration::from_millis(30)).await;
|
||||
self.records
|
||||
.lock()
|
||||
.expect("records lock")
|
||||
.push(record.request_id.clone());
|
||||
self.active
|
||||
.fetch_sub(1, std::sync::atomic::Ordering::AcqRel);
|
||||
Ok(None)
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl UsageSettlementWriter for SlowUsageStore {
|
||||
fn has_usage_settlement_writer(&self) -> bool {
|
||||
false
|
||||
}
|
||||
|
||||
async fn settle_usage(
|
||||
&self,
|
||||
_input: UsageSettlementInput,
|
||||
) -> Result<Option<StoredUsageSettlement>, DataLayerError> {
|
||||
Ok(None)
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl ManualProxyNodeCounter for SlowUsageStore {
|
||||
async fn increment_manual_proxy_node_requests(
|
||||
&self,
|
||||
_node_id: &str,
|
||||
_total_delta: i64,
|
||||
_failed_delta: i64,
|
||||
_latency_ms: Option<i64>,
|
||||
) -> Result<(), DataLayerError> {
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl UsageBillingEventEnricher for SlowUsageStore {
|
||||
async fn enrich_usage_event(&self, _event: &mut UsageEvent) -> Result<(), DataLayerError> {
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
impl crate::runtime::UsageRuntimeAccess for SlowUsageStore {
|
||||
fn has_usage_writer(&self) -> bool {
|
||||
true
|
||||
}
|
||||
|
||||
fn has_usage_worker_queue(&self) -> bool {
|
||||
true
|
||||
}
|
||||
|
||||
fn usage_worker_queue(&self) -> Option<Arc<dyn RuntimeQueueStore>> {
|
||||
None
|
||||
}
|
||||
|
||||
fn usage_worker_should_defer_for_database_pressure(&self) -> bool {
|
||||
self.db_pressure.load(Ordering::Acquire)
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl UsageEventRecorder for SelectiveFailingRecorder {
|
||||
async fn record_usage_event(&self, event: &UsageEvent) -> Result<(), DataLayerError> {
|
||||
@@ -710,6 +1036,118 @@ mod tests {
|
||||
assert_eq!(records[0].total_cost_usd, None);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn data_event_recorder_serializes_same_request_id_writes() {
|
||||
let store = Arc::new(SlowUsageStore::default());
|
||||
let recorder = Arc::new(super::UsageDataEventRecorder::new(Arc::clone(&store)));
|
||||
let mut first = sample_event();
|
||||
first.request_id = "req-same".to_string();
|
||||
first.event_type = UsageEventType::Pending;
|
||||
let mut second = sample_event();
|
||||
second.request_id = "req-same".to_string();
|
||||
second.event_type = UsageEventType::Completed;
|
||||
|
||||
let first_recorder = Arc::clone(&recorder);
|
||||
let second_recorder = Arc::clone(&recorder);
|
||||
tokio::try_join!(
|
||||
async move { first_recorder.record_usage_event(&first).await },
|
||||
async move { second_recorder.record_usage_event(&second).await }
|
||||
)
|
||||
.expect("same request writes should both succeed");
|
||||
|
||||
assert_eq!(
|
||||
store.max_active.load(std::sync::atomic::Ordering::Acquire),
|
||||
1
|
||||
);
|
||||
assert_eq!(store.records.lock().expect("records lock").len(), 2);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn data_event_recorder_defers_when_database_pool_is_under_pressure() {
|
||||
let store = Arc::new(SlowUsageStore::default());
|
||||
store.db_pressure.store(true, Ordering::Release);
|
||||
let gate = Arc::new(UsageWorkerRecordConcurrencyGate::new(1));
|
||||
let recorder = super::UsageDataEventRecorder::with_record_gate_and_database_pressure_defer(
|
||||
Arc::clone(&store),
|
||||
Some(Arc::clone(&gate)),
|
||||
);
|
||||
|
||||
recorder
|
||||
.record_usage_event(&sample_event())
|
||||
.await
|
||||
.expect("recorder should write after brief defer");
|
||||
|
||||
assert_eq!(gate.deferred_total(), 1);
|
||||
assert_eq!(store.records.lock().expect("records lock").len(), 1);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn usage_worker_record_gate_limits_concurrent_record_writes() {
|
||||
let runner = Arc::new(RuntimeState::memory(MemoryRuntimeStateConfig::default()));
|
||||
let queue_runner: Arc<dyn RuntimeQueueStore> = runner.clone();
|
||||
let store = Arc::new(SlowUsageStore::default());
|
||||
let gate = Arc::new(UsageWorkerRecordConcurrencyGate::new(2));
|
||||
let config = UsageRuntimeConfig {
|
||||
enabled: true,
|
||||
stream_key: "usage:test:worker:record-gate".to_string(),
|
||||
consumer_group: "usage:test:worker:record-gate-group".to_string(),
|
||||
dlq_stream_key: "usage:test:worker:record-gate-dlq".to_string(),
|
||||
consumer_batch_size: 1,
|
||||
consumer_block_ms: 1,
|
||||
worker_record_concurrency_limit: Some(2),
|
||||
..UsageRuntimeConfig::default()
|
||||
};
|
||||
let mut handles = Vec::new();
|
||||
for worker_index in 0..4 {
|
||||
let worker = build_usage_queue_worker_with_record_gate(
|
||||
Arc::clone(&queue_runner),
|
||||
Arc::clone(&store),
|
||||
config.clone(),
|
||||
Some(Arc::clone(&gate)),
|
||||
Some(worker_index),
|
||||
)
|
||||
.expect("worker should build");
|
||||
worker
|
||||
.queue
|
||||
.ensure_consumer_group()
|
||||
.await
|
||||
.expect("group should initialize");
|
||||
handles.push(tokio::spawn(async move {
|
||||
let entries = worker
|
||||
.queue
|
||||
.read_group(&worker.consumer)
|
||||
.await
|
||||
.expect("event should read");
|
||||
worker
|
||||
.process_entries(entries)
|
||||
.await
|
||||
.expect("event should process");
|
||||
}));
|
||||
}
|
||||
|
||||
for index in 0..4 {
|
||||
let mut event = sample_event();
|
||||
event.request_id = format!("req-record-gate-{index}");
|
||||
UsageQueue::new(queue_runner.clone(), config.clone())
|
||||
.expect("queue should build")
|
||||
.enqueue(&event)
|
||||
.await
|
||||
.expect("event should enqueue");
|
||||
}
|
||||
|
||||
for handle in handles {
|
||||
handle.await.expect("worker should complete");
|
||||
}
|
||||
|
||||
assert_eq!(
|
||||
store.max_active.load(std::sync::atomic::Ordering::Acquire),
|
||||
2
|
||||
);
|
||||
assert_eq!(gate.max_in_flight(), 2);
|
||||
assert!(gate.wait_total() > 0);
|
||||
assert_eq!(store.records.lock().expect("records lock").len(), 4);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn usage_event_record_error_classifies_permanent_failures() {
|
||||
assert!(usage_event_record_error_is_permanent(
|
||||
|
||||
@@ -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 root,Aether 不得额外拼接 `/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 不应被拼进 path;Aether 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 family;body 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 body,Aether 只移除重复的顶层 `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` body;Vertex 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 auth;service 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 model,body 不含顶层 `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 语义修复,否则会扩大风险面。
|
||||
@@ -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
|
||||
@@ -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 },
|
||||
])
|
||||
})
|
||||
})
|
||||
@@ -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)
|
||||
})
|
||||
})
|
||||
|
||||
@@ -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
Executable
+11
@@ -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
@@ -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)
|
||||
})
|
||||
Executable
+1252
File diff suppressed because it is too large
Load Diff
+197
@@ -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
@@ -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
@@ -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
@@ -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
|
||||
Reference in New Issue
Block a user