i/Aether
1
0
Fork 0
mirror of https://github.com/fawney19/Aether.git synced 2026-09-12 22:20:19 +08:00
Code Issues Packages Projects Releases Wiki Activity

Add usage queue worker autoscaling

Browse Source
This commit is contained in:
elky
2026-06-26 14:02:57 +08:00
parent 6c5e70ccb1
commit 7e9424008f
21 changed files with 1801 additions and 85 deletions
Download Patch File Download Diff File Expand all files Collapse all files
apps/aether-gateway/src/handlers/admin/provider/oauth/quota/codex/mod.rs
+29 -3
View File
@@ -15,7 +15,8 @@ use self::plan::{build_codex_quota_request_spec, execute_codex_quota_plan};
use super::shared::{
build_quota_snapshot_payload, extract_execution_error_message,
oauth_refresh_auto_removed_result, persist_provider_quota_refresh_state,
provider_auto_remove_banned_keys, quota_key_auto_removed, quota_refresh_success_invalid_state,
provider_auto_remove_banned_keys, provider_auto_remove_quota_exhausted_keys,
quota_key_auto_removed, quota_refresh_success_invalid_state,
should_auto_remove_structured_reason, ProviderQuotaExecutionOutcome,
};
use crate::handlers::admin::request::AdminAppState;
@@ -61,6 +62,8 @@ pub(crate) async fn refresh_codex_provider_quota_locally(
proxy_override: Option<ProxySnapshot>,
) -> Result<Option<serde_json::Value>, GatewayError> {
let auto_remove_abnormal_keys = provider_auto_remove_banned_keys(provider.config.as_ref());
let auto_remove_quota_exhausted_keys =
provider_auto_remove_quota_exhausted_keys(provider.config.as_ref());
let mut results = Vec::new();
let mut success_count = 0usize;
let mut failed_count = 0usize;
@@ -312,7 +315,7 @@ pub(crate) async fn refresh_codex_provider_quota_locally(
}));
continue;
}
let auto_removed = if auto_remove_candidate {
let auto_removed_hard_banned = if auto_remove_candidate {
state
.cleanup_provider_catalog_key_if_current(provider, &key.id, |latest_key| {
should_auto_remove_structured_reason(latest_key.oauth_invalid_reason.as_deref())
@@ -321,10 +324,28 @@ pub(crate) async fn refresh_codex_provider_quota_locally(
} else {
false
};
if auto_removed {
if auto_removed_hard_banned {
auto_removed_count += 1;
auto_removed_hard_banned_count += 1;
}
let auto_removed_quota_exhausted =
if !auto_removed_hard_banned && auto_remove_quota_exhausted_keys {
state
.cleanup_provider_catalog_key_if_current(provider, &key.id, |latest_key| {
aether_admin::provider::pool::admin_pool_key_account_quota_exhausted(
latest_key,
provider.provider_type.as_str(),
)
})
.await?
} else {
false
};
if auto_removed_quota_exhausted {
auto_removed_count += 1;
status = "quota_exhausted".to_string();
}
let auto_removed = auto_removed_hard_banned || auto_removed_quota_exhausted;
let refresh_fixed =
status == "success" && had_oauth_refresh_issue && oauth_invalid_reason.is_none();
if refresh_fixed {
@@ -370,8 +391,13 @@ pub(crate) async fn refresh_codex_provider_quota_locally(
}
if auto_removed {
payload.insert("auto_removed".to_string(), json!(true));
}
if auto_removed_hard_banned {
payload.insert("auto_removed_hard_banned".to_string(), json!(true));
}
if auto_removed_quota_exhausted {
payload.insert("auto_removed_quota_exhausted".to_string(), json!(true));
}
if refresh_fixed {
payload.insert("refresh_fixed".to_string(), json!(true));
}
apps/aether-gateway/src/handlers/admin/provider/oauth/quota/kiro/mod.rs
+25 -2
View File
@@ -6,8 +6,8 @@ use self::plan::execute_kiro_quota_plan;
use super::shared::{
build_quota_snapshot_payload, extract_execution_error_message,
oauth_refresh_auto_removed_result, persist_provider_quota_refresh_state,
persist_quota_oauth_refresh_failure_state, quota_refresh_success_invalid_state,
ProviderQuotaExecutionOutcome,
persist_quota_oauth_refresh_failure_state, provider_auto_remove_quota_exhausted_keys,
quota_refresh_success_invalid_state, ProviderQuotaExecutionOutcome,
};
use crate::handlers::admin::request::{AdminAppState, AdminLocalOAuthRefreshError};
use crate::GatewayError;
@@ -97,6 +97,8 @@ pub(crate) async fn refresh_kiro_provider_quota_locally(
let mut success_count = 0usize;
let mut failed_count = 0usize;
let mut auto_removed_count = 0usize;
let auto_remove_quota_exhausted_keys =
provider_auto_remove_quota_exhausted_keys(provider.config.as_ref());
for key in keys {
let transport = match state
@@ -304,6 +306,23 @@ pub(crate) async fn refresh_kiro_provider_quota_locally(
continue;
}
let auto_removed_quota_exhausted = if auto_remove_quota_exhausted_keys {
state
.cleanup_provider_catalog_key_if_current(provider, &key.id, |latest_key| {
aether_admin::provider::pool::admin_pool_key_account_quota_exhausted(
latest_key,
provider.provider_type.as_str(),
)
})
.await?
} else {
false
};
if auto_removed_quota_exhausted {
auto_removed_count += 1;
status = "quota_exhausted".to_string();
}
if status == "success" {
success_count += 1;
} else {
@@ -331,6 +350,10 @@ pub(crate) async fn refresh_kiro_provider_quota_locally(
) {
payload.insert("quota_snapshot".to_string(), quota_snapshot);
}
if auto_removed_quota_exhausted {
payload.insert("auto_removed".to_string(), json!(true));
payload.insert("auto_removed_quota_exhausted".to_string(), json!(true));
}
results.push(serde_json::Value::Object(payload));
}
apps/aether-gateway/src/handlers/admin/provider/oauth/quota/shared.rs
+6
View File
@@ -67,6 +67,12 @@ pub(crate) fn provider_auto_remove_banned_keys(config: Option<&serde_json::Value
admin_provider_quota_pure::provider_auto_remove_banned_keys(config)
}
pub(crate) fn provider_auto_remove_quota_exhausted_keys(
config: Option<&serde_json::Value>,
) -> bool {
admin_provider_quota_pure::provider_auto_remove_quota_exhausted_keys(config)
}
pub(super) fn should_auto_remove_structured_reason(reason: Option<&str>) -> bool {
admin_provider_quota_pure::should_auto_remove_structured_reason(reason)
}
apps/aether-gateway/src/handlers/admin/provider/write/provider/template.rs
+93 -2
View File
@@ -16,6 +16,8 @@ use std::collections::{BTreeMap, BTreeSet};
const FIXED_PROVIDER_TEMPLATE_METADATA_KEY: &str = "_aether_fixed_provider_template";
const OVERRIDE_BODY_RULES: &str = "body_rules";
const OVERRIDE_FORMAT_ACCEPTANCE_CONFIG: &str = "format_acceptance_config";
const OVERRIDE_BASE_URL: &str = "base_url";
const OVERRIDE_CUSTOM_PATH: &str = "custom_path";
const OVERRIDE_HEADER_RULES: &str = "header_rules";
const OVERRIDE_IS_ACTIVE: &str = "is_active";
const OVERRIDE_MAX_RETRIES: &str = "max_retries";
@@ -158,6 +160,20 @@ pub(crate) fn apply_admin_fixed_provider_endpoint_template_overrides(
.unwrap_or_else(|| managed_fixed_provider_endpoint_metadata(template, endpoint_template));
let mut overrides = metadata.overrides.clone();
sync_override_if_changed(
&mut overrides,
OVERRIDE_BASE_URL,
&existing_endpoint.base_url,
&updated_endpoint.base_url,
&defaults.base_url,
);
sync_override_if_changed(
&mut overrides,
OVERRIDE_CUSTOM_PATH,
&existing_endpoint.custom_path,
&updated_endpoint.custom_path,
&defaults.custom_path,
);
sync_override_if_changed(
&mut overrides,
OVERRIDE_HEADER_RULES,
@@ -249,9 +265,13 @@ fn reconcile_fixed_provider_endpoint(
updated.api_format = defaults.api_format.clone();
updated.api_family = Some(defaults.api_family.clone());
updated.endpoint_kind = Some(defaults.endpoint_kind.clone());
updated.base_url = defaults.base_url;
updated.custom_path = defaults.custom_path;
if !metadata.overrides.contains(OVERRIDE_BASE_URL) {
updated.base_url = defaults.base_url;
}
if !metadata.overrides.contains(OVERRIDE_CUSTOM_PATH) {
updated.custom_path = defaults.custom_path;
}
if !metadata.overrides.contains(OVERRIDE_HEADER_RULES) {
updated.header_rules = defaults.header_rules;
}
@@ -534,3 +554,74 @@ fn sync_override_if_changed<T>(
}
sync_override(overrides, key, actual, desired);
}
#[cfg(test)]
mod tests {
use super::{
apply_admin_fixed_provider_endpoint_template_overrides, fixed_provider_endpoint_metadata,
reconcile_fixed_provider_endpoint,
};
use aether_data_contracts::repository::provider_catalog::{
StoredProviderCatalogEndpoint, StoredProviderCatalogProvider,
};
use aether_provider_transport::provider_types::fixed_provider_template;
fn sample_codex_provider() -> StoredProviderCatalogProvider {
StoredProviderCatalogProvider::new(
"provider-codex".to_string(),
"Codex".to_string(),
None,
"codex".to_string(),
)
.expect("provider should build")
}
fn sample_codex_endpoint(base_url: &str) -> StoredProviderCatalogEndpoint {
StoredProviderCatalogEndpoint::new(
"endpoint-codex-responses".to_string(),
"provider-codex".to_string(),
"openai:responses".to_string(),
Some("openai".to_string()),
Some("responses".to_string()),
true,
)
.expect("endpoint should build")
.with_transport_fields(
base_url.to_string(),
None,
None,
Some(2),
None,
None,
None,
None,
)
.expect("endpoint transport should build")
}
#[test]
fn fixed_provider_endpoint_reconcile_preserves_base_url_override() {
let provider = sample_codex_provider();
let template = fixed_provider_template("codex").expect("codex template should exist");
let endpoint_template = template
.endpoints
.iter()
.find(|endpoint| endpoint.api_format == "openai:responses")
.expect("responses endpoint template should exist");
let existing = sample_codex_endpoint("https://chatgpt.com/backend-api/codex");
let mut updated = existing.clone();
updated.base_url = "http://127.0.0.1:18181/v1".to_string();
apply_admin_fixed_provider_endpoint_template_overrides(&provider, &existing, &mut updated)
.expect("override metadata should apply");
let metadata = fixed_provider_endpoint_metadata(&updated)
.expect("fixed provider metadata should exist");
assert!(metadata.overrides.contains("base_url"));
let reconciled =
reconcile_fixed_provider_endpoint(&provider, &updated, template, endpoint_template)
.expect("endpoint should reconcile");
assert_eq!(reconciled.base_url, "http://127.0.0.1:18181/v1");
}
}
apps/aether-gateway/src/handlers/admin/request/provider/builders.rs
-8
View File
@@ -315,14 +315,6 @@ impl<'a> AdminAppState<'a> {
return Err("Gemini CLI Endpoint 由系统固定管理，不允许修改".to_string());
}
if self.provider_type_is_fixed(&provider.provider_type)
&& (fields.contains("base_url") || fields.contains("custom_path"))
{
return Err(
"固定类型 Provider 的 Endpoint 不允许修改 base_url/custom_path".to_string(),
);
}
let mut update_fields = admin_provider_endpoints_pure::AdminProviderEndpointUpdateFields {
base_url: payload.base_url,
custom_path: payload.custom_path,
apps/aether-gateway/src/handlers/admin/request/provider/tasks.rs
+98
View File
@@ -1,4 +1,8 @@
use super::*;
use crate::ai_serving::provider_key_pool_score_scope;
use aether_data_contracts::repository::pool_scores::{
ListPoolMemberScoresQuery, PoolMemberHardState, POOL_KIND_PROVIDER_KEY_POOL,
};
use aether_data_contracts::repository::provider_catalog::{
StoredProviderCatalogKey, StoredProviderCatalogProvider,
};
@@ -266,6 +270,100 @@ impl<'a> AdminAppState<'a> {
Ok(affected)
}
pub(crate) async fn cleanup_quota_exhausted_provider_catalog_keys(
&self,
provider: &StoredProviderCatalogProvider,
provider_type: &str,
) -> Result<usize, GatewayError> {
use aether_admin::provider::pool as admin_provider_pool_pure;
let keys = self
.list_provider_catalog_keys_by_provider_ids(std::slice::from_ref(&provider.id))
.await?;
if keys.is_empty() {
return Ok(0);
}
let known_key_ids = keys
.iter()
.map(|key| key.id.as_str())
.collect::<std::collections::BTreeSet<_>>();
let mut exhausted_key_ids = keys
.iter()
.filter(|key| {
admin_provider_pool_pure::admin_pool_key_account_quota_exhausted(key, provider_type)
})
.map(|key| key.id.clone())
.collect::<std::collections::BTreeSet<_>>();
if self.app().data.has_pool_score_reader() {
let scope = provider_key_pool_score_scope();
let page_size = 10_000usize;
let mut offset = 0usize;
loop {
let scores = self
.app()
.data
.list_pool_member_scores(&ListPoolMemberScoresQuery {
pool_kind: POOL_KIND_PROVIDER_KEY_POOL.to_string(),
pool_id: provider.id.clone(),
capability: Some(scope.capability.clone()),
scope_kind: Some(scope.scope_kind.clone()),
scope_id: scope.scope_id.clone(),
hard_states: vec![PoolMemberHardState::QuotaExhausted],
probe_statuses: None,
offset,
limit: page_size,
})
.await
.map_err(|err| GatewayError::Internal(err.to_string()))?;
if scores.is_empty() {
break;
}
let page_len = scores.len();
for score in scores {
if known_key_ids.contains(score.member_id.as_str()) {
exhausted_key_ids.insert(score.member_id);
}
}
if page_len < page_size {
break;
}
offset = offset.saturating_add(page_size);
}
}
let exhausted_keys = keys
.iter()
.filter(|key| exhausted_key_ids.contains(&key.id))
.collect::<Vec<_>>();
if exhausted_keys.is_empty() {
return Ok(0);
}
let deleted_key_ids = exhausted_keys
.iter()
.map(|key| key.id.clone())
.collect::<Vec<_>>();
for key in exhausted_keys {
self.clear_admin_provider_pool_cooldown(&provider.id, &key.id)
.await;
self.reset_admin_provider_pool_cost(&provider.id, &key.id)
.await;
}
let mut affected = 0usize;
for key_id in &deleted_key_ids {
if self.delete_provider_catalog_key(key_id).await? {
affected += 1;
}
}
self.cleanup_deleted_provider_catalog_refs(&provider.id, false, &[], &deleted_key_ids)
.await?;
Ok(affected)
}
pub(crate) async fn cleanup_provider_catalog_key_if_current<F>(
&self,
provider: &StoredProviderCatalogProvider,
apps/aether-gateway/src/main.rs
+450 -19
View File
@@ -245,6 +245,12 @@ const AUTO_SERVER_SQL_POOL_MIN_CONNECTIONS_FLOOR: u32 = 4;
const AUTO_SERVER_SQL_POOL_MIN_CONNECTIONS_CAP: u32 = 16;
const AUTO_SERVER_SQL_POOL_MAX_CONNECTIONS_FLOOR: u32 = 20;
const AUTO_SERVER_SQL_POOL_MAX_CONNECTIONS_CAP: u32 = 100;
const DEFAULT_USAGE_QUEUE_WORKERS_CAP: usize = 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 MAX_USAGE_QUEUE_WORKERS: usize = 64;
const DEFAULT_GATEWAY_LISTEN_BACKLOG: i32 = 65_535;
const MIN_GATEWAY_LISTEN_BACKLOG: i32 = 128;
const MAX_GATEWAY_LISTEN_BACKLOG: i32 = 65_535;
@@ -261,12 +267,99 @@ fn env_var_trimmed(name: &str) -> Option<String> {
}
fn available_parallelism_u32() -> u32 {
std::thread::available_parallelism()
.map(|value| u32::try_from(value.get()).unwrap_or(u32::MAX))
.unwrap_or(AUTO_SERVER_SQL_POOL_MIN_CONNECTIONS_FLOOR)
u32::try_from(available_parallelism_usize())
.unwrap_or(u32::MAX)
.max(1)
}
fn available_parallelism_usize() -> usize {
std::thread::available_parallelism()
.map(|value| value.get())
.unwrap_or(AUTO_SERVER_SQL_POOL_MIN_CONNECTIONS_FLOOR as usize)
.max(1)
}
fn usage_queue_request_concurrency_hint(
max_in_flight_requests: Option<usize>,
distributed_request_limit: Option<usize>,
) -> Option<usize> {
match (
max_in_flight_requests.filter(|limit| *limit > 0),
distributed_request_limit.filter(|limit| *limit > 0),
) {
(Some(local), Some(distributed)) => Some(local.min(distributed)),
(Some(local), None) => Some(local),
(None, Some(distributed)) => Some(distributed),
(None, None) => None,
}
}
fn usage_queue_workers_for_request_concurrency(request_concurrency: usize) -> usize {
let workers = request_concurrency
.saturating_add(AUTO_USAGE_QUEUE_WORKERS_REQUESTS_PER_WORKER - 1)
/ AUTO_USAGE_QUEUE_WORKERS_REQUESTS_PER_WORKER;
workers.clamp(AUTO_USAGE_QUEUE_WORKERS_MIN, MAX_USAGE_QUEUE_WORKERS)
}
fn usage_queue_worker_database_cap(
node_role: NodeRoleArg,
database: Option<&SqlDatabaseConfig>,
) -> usize {
let Some(database) = database else {
return MAX_USAGE_QUEUE_WORKERS;
};
if database.driver == DatabaseDriver::Sqlite {
return 1;
}
let divisor = if matches!(node_role, NodeRoleArg::Background) {
AUTO_USAGE_QUEUE_WORKERS_DB_SHARE_BACKGROUND
} else {
AUTO_USAGE_QUEUE_WORKERS_DB_SHARE_ALL
};
let max_connections = database.pool.max_connections.max(1) as usize;
max_connections
.saturating_add(divisor - 1)
.checked_div(divisor)
.unwrap_or(1)
.clamp(1, MAX_USAGE_QUEUE_WORKERS)
}
fn automatic_usage_queue_workers_for_parallelism(
parallelism: usize,
node_role: NodeRoleArg,
max_in_flight_requests: Option<usize>,
distributed_request_limit: Option<usize>,
database: Option<&SqlDatabaseConfig>,
) -> usize {
let cpu_default = parallelism.max(1).clamp(
AUTO_USAGE_QUEUE_WORKERS_MIN,
DEFAULT_USAGE_QUEUE_WORKERS_CAP,
);
let requested =
usage_queue_request_concurrency_hint(max_in_flight_requests, distributed_request_limit)
.map(usage_queue_workers_for_request_concurrency)
.unwrap_or(cpu_default);
requested
.min(usage_queue_worker_database_cap(node_role, database))
.clamp(1, MAX_USAGE_QUEUE_WORKERS)
}
fn automatic_usage_queue_workers(
node_role: NodeRoleArg,
max_in_flight_requests: Option<usize>,
distributed_request_limit: Option<usize>,
database: Option<&SqlDatabaseConfig>,
) -> usize {
automatic_usage_queue_workers_for_parallelism(
available_parallelism_usize(),
node_role,
max_in_flight_requests,
distributed_request_limit,
database,
)
}
fn automatic_sql_pool_config(driver: DatabaseDriver) -> SqlPoolConfig {
automatic_sql_pool_config_for_parallelism(driver, available_parallelism_u32())
}
@@ -525,6 +618,37 @@ struct GatewayUsageArgs {
)]
queue_lifecycle_events: bool,
#[arg(long, env = "AETHER_GATEWAY_USAGE_QUEUE_WORKERS", value_name = "COUNT")]
queue_workers: Option<usize>,
#[arg(
long,
env = "AETHER_GATEWAY_USAGE_QUEUE_WORKER_AUTOSCALE_ENABLED",
default_value_t = true
)]
queue_worker_autoscale_enabled: bool,
#[arg(
long,
env = "AETHER_GATEWAY_USAGE_QUEUE_WORKER_MAX_COUNT",
value_name = "COUNT"
)]
queue_worker_max_count: Option<usize>,
#[arg(
long,
env = "AETHER_GATEWAY_USAGE_QUEUE_WORKER_SCALE_INTERVAL_MS",
default_value_t = 1_000
)]
queue_worker_scale_interval_ms: u64,
#[arg(
long,
env = "AETHER_GATEWAY_USAGE_QUEUE_WORKER_IDLE_SCALE_DOWN_TICKS",
default_value_t = 30
)]
queue_worker_idle_scale_down_ticks: u64,
#[arg(
long,
env = "AETHER_GATEWAY_USAGE_QUEUE_STREAM_KEY",
@@ -639,11 +763,66 @@ struct GatewayUsageArgs {
}
impl GatewayUsageArgs {
fn to_config(&self) -> UsageRuntimeConfig {
fn effective_queue_workers(
&self,
node_role: NodeRoleArg,
max_in_flight_requests: Option<usize>,
distributed_request_limit: Option<usize>,
database: Option<&SqlDatabaseConfig>,
) -> usize {
if let Some(queue_workers) = self.queue_workers {
return queue_workers.clamp(1, MAX_USAGE_QUEUE_WORKERS);
}
if !self.queue_terminal_events && !self.queue_lifecycle_events {
return 1;
}
automatic_usage_queue_workers(
node_role,
max_in_flight_requests,
distributed_request_limit,
database,
)
}
fn effective_queue_worker_max_count(
&self,
node_role: NodeRoleArg,
database: Option<&SqlDatabaseConfig>,
worker_count: usize,
) -> usize {
if !self.queue_worker_autoscale_enabled {
return worker_count.max(1).min(MAX_USAGE_QUEUE_WORKERS);
}
self.queue_worker_max_count
.unwrap_or_else(|| usage_queue_worker_database_cap(node_role, database))
.clamp(worker_count.max(1), MAX_USAGE_QUEUE_WORKERS)
}
fn runtime_state_blocking_stream_lanes(
&self,
node_role: NodeRoleArg,
database: Option<&SqlDatabaseConfig>,
worker_max_count: usize,
) -> Option<usize> {
if !node_role.spawns_background_tasks()
|| (!self.queue_terminal_events && !self.queue_lifecycle_events)
|| database.is_none()
{
return None;
}
Some(worker_max_count.clamp(1, MAX_USAGE_QUEUE_WORKERS))
}
fn to_config(&self, worker_count: usize, worker_max_count: usize) -> UsageRuntimeConfig {
UsageRuntimeConfig {
enabled: true,
queue_terminal_events: self.queue_terminal_events,
queue_lifecycle_events: self.queue_lifecycle_events,
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_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(),
consumer_group: self.queue_group.trim().to_string(),
dlq_stream_key: self.queue_dlq_stream_key.trim().to_string(),
@@ -1041,6 +1220,7 @@ impl Args {
&self,
runtime_backend: RuntimeBackendArg,
data_redis_url: Option<&str>,
blocking_stream_lanes: Option<usize>,
) -> RuntimeStateConfig {
let redis = self
.effective_runtime_redis_url(data_redis_url)
@@ -1052,6 +1232,7 @@ impl Args {
backend: runtime_backend.to_runtime_state_backend(),
redis,
command_timeout_ms: Some(self.runtime_command_timeout_ms.max(1)),
blocking_stream_lanes,
..RuntimeStateConfig::default()
}
}
@@ -1397,10 +1578,35 @@ async fn run() -> Result<(), Box<dyn std::error::Error>> {
runtime_redis_url.as_deref(),
runtime_backend,
)?;
let usage_queue_request_concurrency_hint = usage_queue_request_concurrency_hint(
args.max_in_flight_requests,
args.distributed_request_limit,
);
let usage_queue_workers = args.usage.effective_queue_workers(
args.node_role,
args.max_in_flight_requests,
args.distributed_request_limit,
sql_database_config.as_ref(),
);
let usage_queue_worker_max_count = args.usage.effective_queue_worker_max_count(
args.node_role,
sql_database_config.as_ref(),
usage_queue_workers,
);
let usage_config = args
.usage
.to_config(usage_queue_workers, usage_queue_worker_max_count);
let usage_blocking_stream_lanes = args.usage.runtime_state_blocking_stream_lanes(
args.node_role,
sql_database_config.as_ref(),
usage_config.worker_max_count,
);
let runtime_state = Arc::new(
RuntimeState::from_config(
args.runtime_state_config(runtime_backend, data_redis_url.as_deref()),
)
RuntimeState::from_config(args.runtime_state_config(
runtime_backend,
data_redis_url.as_deref(),
usage_blocking_stream_lanes,
))
.await
.map_err(|err| std::io::Error::new(std::io::ErrorKind::InvalidInput, err.to_string()))?,
);
@@ -1425,6 +1631,16 @@ async fn run() -> Result<(), Box<dyn std::error::Error>> {
deployment_topology = args.deployment_topology.as_str(),
node_role = args.node_role.as_str(),
runtime_backend = runtime_backend.as_str(),
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_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() {
"explicit"
} else {
"none"
},
frontdoor_mode = "compatibility_frontdoor",
log_format = ?args.logging.log_format,
log_destination = args.logging.log_destination.as_str(),
@@ -1448,6 +1664,22 @@ async fn run() -> Result<(), Box<dyn std::error::Error>> {
video_task_poller_interval_ms = args.video_task_poller_interval_ms,
video_task_poller_batch_size = args.video_task_poller_batch_size,
video_task_store_path = args.video_task_store_path.as_deref().unwrap_or("-"),
usage_queue_workers = usage_config.worker_count,
usage_queue_workers_source = if args.usage.queue_workers.is_some() {
"explicit"
} else {
"auto"
},
usage_queue_worker_autoscale_enabled = usage_config.worker_autoscale_enabled,
usage_queue_worker_max_count = usage_config.worker_max_count,
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() {
"explicit"
} else {
"none"
},
max_in_flight_requests = args.max_in_flight_requests.unwrap_or_default(),
distributed_request_limit = args.distributed_request_limit.unwrap_or_default(),
distributed_request_redis_configured = args
@@ -1479,7 +1711,7 @@ async fn run() -> Result<(), Box<dyn std::error::Error>> {
let mut state = AppState::new()?
.with_runtime_state(runtime_state)
.with_data_config(data_config)?
.with_usage_runtime_config(args.usage.to_config())?
.with_usage_runtime_config(usage_config)?
.with_video_task_truth_source_mode(args.video_task_truth_source_mode.into());
if let Some(cors_config) = args.frontdoor.cors_config() {
state = state.with_frontdoor_cors_config(cors_config);
@@ -1977,9 +2209,7 @@ fn pending_schema_error(
) -> std::io::Error {
std::io::Error::other(format!(
"database schema is behind by {} migration(s); next pending migration is {} ({})\nrun `aether-gateway --migrate` before starting the service",
pending_count,
next_version,
next_description
pending_count, next_version, next_description
))
}
@@ -1990,9 +2220,7 @@ fn pending_backfills_error(
) -> std::io::Error {
std::io::Error::other(format!(
"database backfills are behind by {} backfill(s); next pending backfill is {} ({})\nrun `aether-gateway --apply-backfills` before starting the service",
pending_count,
next_version,
next_description
pending_count, next_version, next_description
))
}
@@ -2000,11 +2228,11 @@ fn pending_backfills_error(
mod tests {
use super::{
automatic_sql_pool_config, automatic_sql_pool_config_for_parallelism,
ensure_database_backfills_are_current, ensure_database_schema_is_current,
pending_backfills_error, pending_schema_error, resolve_healthcheck_url, Args,
DatabaseDriverArg, DeploymentTopologyArg, GatewayDataArgs, GatewayFrontdoorArgs,
GatewayLogDestinationArg, GatewayLogFormatArg, GatewayLogRotationArg, GatewayLoggingArgs,
GatewayRateLimitArgs, GatewayUsageArgs, NodeRoleArg, RuntimeBackendArg,
automatic_usage_queue_workers_for_parallelism, ensure_database_backfills_are_current,
ensure_database_schema_is_current, pending_backfills_error, pending_schema_error,
resolve_healthcheck_url, Args, DatabaseDriverArg, DeploymentTopologyArg, GatewayDataArgs,
GatewayFrontdoorArgs, GatewayLogDestinationArg, GatewayLogFormatArg, GatewayLogRotationArg,
GatewayLoggingArgs, GatewayRateLimitArgs, GatewayUsageArgs, NodeRoleArg, RuntimeBackendArg,
VideoTaskTruthSourceArg, DEFAULT_GATEWAY_HTTP2_MAX_CONCURRENT_STREAMS,
DEFAULT_GATEWAY_LISTENER_SHARDS, DEFAULT_GATEWAY_LISTEN_BACKLOG,
MAX_GATEWAY_HTTP2_MAX_CONCURRENT_STREAMS, MAX_GATEWAY_LISTENER_SHARDS,
@@ -2062,6 +2290,11 @@ mod tests {
usage: GatewayUsageArgs {
queue_terminal_events: true,
queue_lifecycle_events: true,
queue_workers: Some(4),
queue_worker_autoscale_enabled: true,
queue_worker_max_count: None,
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(),
@@ -2100,6 +2333,25 @@ mod tests {
}
}
fn test_database(driver: DatabaseDriver, max_connections: u32) -> SqlDatabaseConfig {
let url = match driver {
DatabaseDriver::Sqlite => "sqlite://./data/aether.db",
DatabaseDriver::Mysql => "mysql://root:root@localhost/aether",
DatabaseDriver::Postgres => "postgres://postgres:postgres@localhost/aether",
};
let max_connections = max_connections.max(1);
SqlDatabaseConfig::new(
driver,
url,
SqlPoolConfig {
min_connections: 1,
max_connections,
..SqlPoolConfig::default()
},
)
.expect("test database config should build")
}
#[test]
fn resolves_healthcheck_url_from_app_port() {
assert_eq!(
@@ -2260,6 +2512,183 @@ mod tests {
assert_eq!(many_cpu.max_connections, 100);
}
#[test]
fn gateway_usage_queue_workers_manual_override_wins_and_is_capped() {
let mut args = test_args();
args.usage.queue_workers = Some(72);
let database = test_database(DatabaseDriver::Postgres, 100);
let workers = args.usage.effective_queue_workers(
NodeRoleArg::All,
Some(10_000),
None,
Some(&database),
);
assert_eq!(workers, 64);
assert_eq!(args.usage.to_config(workers, 64).worker_count, 64);
}
#[test]
fn gateway_usage_queue_workers_auto_uses_cpu_default_without_concurrency_hint() {
let database = test_database(DatabaseDriver::Postgres, 100);
let workers = automatic_usage_queue_workers_for_parallelism(
4,
NodeRoleArg::All,
None,
None,
Some(&database),
);
assert_eq!(workers, 4);
}
#[test]
fn gateway_usage_queue_worker_autoscale_max_uses_database_cap() {
let mut args = test_args();
args.usage.queue_workers = None;
let database = test_database(DatabaseDriver::Postgres, 40);
let workers =
args.usage
.effective_queue_workers(args.node_role, None, None, Some(&database));
let max_workers =
args.usage
.effective_queue_worker_max_count(args.node_role, Some(&database), workers);
assert_eq!(workers, 8);
assert_eq!(max_workers, 10);
}
#[test]
fn gateway_usage_queue_worker_autoscale_max_respects_explicit_override() {
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 workers =
args.usage
.effective_queue_workers(args.node_role, None, None, Some(&database));
let max_workers =
args.usage
.effective_queue_worker_max_count(args.node_role, Some(&database), workers);
assert_eq!(workers, 8);
assert_eq!(max_workers, 32);
}
#[test]
fn gateway_usage_queue_blocking_stream_lanes_only_expand_when_worker_can_spawn() {
let database = test_database(DatabaseDriver::Postgres, 100);
let args = test_args();
assert_eq!(
args.usage
.runtime_state_blocking_stream_lanes(NodeRoleArg::All, Some(&database), 10,),
Some(10)
);
assert_eq!(
args.usage.runtime_state_blocking_stream_lanes(
NodeRoleArg::Frontdoor,
Some(&database),
10,
),
None
);
assert_eq!(
args.usage
.runtime_state_blocking_stream_lanes(NodeRoleArg::All, None, 10),
None
);
let mut disabled_queue_args = args;
disabled_queue_args.usage.queue_terminal_events = false;
disabled_queue_args.usage.queue_lifecycle_events = false;
assert_eq!(
disabled_queue_args
.usage
.runtime_state_blocking_stream_lanes(NodeRoleArg::All, Some(&database), 10,),
None
);
}
#[test]
fn gateway_usage_queue_workers_auto_scales_from_request_concurrency() {
let database = test_database(DatabaseDriver::Postgres, 100);
let workers = automatic_usage_queue_workers_for_parallelism(
8,
NodeRoleArg::All,
Some(1_536),
None,
Some(&database),
);
assert_eq!(workers, 12);
}
#[test]
fn gateway_usage_queue_workers_auto_respects_effective_request_limit() {
let database = test_database(DatabaseDriver::Postgres, 100);
let workers = automatic_usage_queue_workers_for_parallelism(
8,
NodeRoleArg::All,
Some(2_048),
Some(256),
Some(&database),
);
assert_eq!(workers, 2);
}
#[test]
fn gateway_usage_queue_workers_auto_is_capped_by_database_pool() {
let database = test_database(DatabaseDriver::Postgres, 20);
let workers = automatic_usage_queue_workers_for_parallelism(
16,
NodeRoleArg::All,
Some(5_000),
None,
Some(&database),
);
assert_eq!(workers, 5);
}
#[test]
fn gateway_usage_queue_workers_auto_gives_background_nodes_more_pool_budget() {
let database = test_database(DatabaseDriver::Postgres, 20);
let workers = automatic_usage_queue_workers_for_parallelism(
16,
NodeRoleArg::Background,
Some(5_000),
None,
Some(&database),
);
assert_eq!(workers, 10);
}
#[test]
fn gateway_usage_queue_workers_auto_uses_single_worker_for_sqlite() {
let database = test_database(DatabaseDriver::Sqlite, 1);
let workers = automatic_usage_queue_workers_for_parallelism(
16,
NodeRoleArg::All,
Some(5_000),
None,
Some(&database),
);
assert_eq!(workers, 1);
}
#[test]
fn gateway_data_pool_explicit_values_override_auto_sizing() {
let mut args = test_args();
@@ -2355,8 +2784,10 @@ mod tests {
let config = args.runtime_state_config(
RuntimeBackendArg::Redis,
args.data.effective_redis_url().as_deref(),
Some(7),
);
assert_eq!(config.blocking_stream_lanes, Some(7));
assert_eq!(
config
.redis
apps/aether-gateway/src/maintenance/runtime/pool_quota_probe.rs
+17
View File
@@ -1255,6 +1255,23 @@ async fn perform_pool_quota_probe_for_provider(
);
}
}
if aether_admin::provider::quota::provider_auto_remove_quota_exhausted_keys(
provider.config.as_ref(),
) {
let auto_removed = admin_state
.cleanup_quota_exhausted_provider_catalog_keys(provider, provider_type)
.await?;
if auto_removed > 0 {
summary.auto_removed += auto_removed;
info!(
event_name = "auto_removed_quota_exhausted",
provider_id = %provider_short_id,
provider_type,
auto_removed,
"gateway pool quota probe auto-cleaned quota-exhausted provider keys"
);
}
}
let Some(endpoint) = endpoint_for_probe_with_reconcile(
state,
apps/aether-gateway/src/state/core.rs
+39 -4
View File
@@ -1357,6 +1357,15 @@ impl AppState {
definition.trigger,
));
};
if let Some(handle) = self
.usage_runtime
.spawn_worker_supervisor(self.data.clone())
{
supervisor.supervise_handle(crate::task_runtime::TASK_KEY_USAGE_QUEUE_WORKER, handle);
record_boot(crate::task_runtime::TASK_KEY_USAGE_QUEUE_WORKER);
}
let mut supervise_worker =
|task_key: &'static str, handle: Option<tokio::task::JoinHandle<()>>| {
if let Some(handle) = handle {
@@ -1365,10 +1374,6 @@ impl AppState {
}
};
supervise_worker(
crate::task_runtime::TASK_KEY_USAGE_QUEUE_WORKER,
self.usage_runtime.spawn_worker(self.data.clone()),
);
supervise_worker(
crate::task_runtime::TASK_KEY_USAGE_COUNTER_FLUSH,
spawn_usage_counter_flush_worker(self.data.clone()),
@@ -1551,6 +1556,36 @@ fn usage_runtime_metric_samples(
MetricKind::Gauge,
u64::from(snapshot.queue_lifecycle_events),
),
MetricSample::new(
"usage_runtime_queue_worker_count",
"Minimum configured number of usage queue worker consumers.",
MetricKind::Gauge,
snapshot.worker_count as u64,
),
MetricSample::new(
"usage_runtime_queue_worker_autoscale_enabled",
"Whether usage queue worker autoscaling is enabled.",
MetricKind::Gauge,
u64::from(snapshot.worker_autoscale_enabled),
),
MetricSample::new(
"usage_runtime_queue_worker_max_count",
"Maximum configured number of elastic usage queue worker consumers.",
MetricKind::Gauge,
snapshot.worker_max_count as u64,
),
MetricSample::new(
"usage_runtime_queue_worker_active_count",
"Current active usage queue worker consumers managed by the supervisor.",
MetricKind::Gauge,
snapshot.worker_active_count as u64,
),
MetricSample::new(
"usage_runtime_queue_worker_desired_count",
"Current desired usage queue worker consumers selected by autoscaling.",
MetricKind::Gauge,
snapshot.worker_desired_count as u64,
),
MetricSample::new(
"usage_runtime_retry_deferred_lifecycle_events_enabled",
"Whether deferred lifecycle usage events are scheduled for local enqueue retry.",
apps/aether-gateway/src/tests/control/admin/endpoints/routes.rs
+65
View File
@@ -667,6 +667,71 @@ async fn gateway_updates_admin_provider_endpoint_locally_with_trusted_admin_prin
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_updates_fixed_provider_endpoint_base_url_as_template_override() {
let mut provider = sample_provider("provider-codex", "codex", 10);
provider.provider_type = "codex".to_string();
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider],
vec![sample_endpoint(
"endpoint-codex-responses",
"provider-codex",
"openai:responses",
"https://chatgpt.com/backend-api/codex",
)],
vec![],
));
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(
GatewayDataState::with_provider_catalog_repository_for_tests(
provider_catalog_repository.clone(),
),
),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.put(format!(
"{gateway_url}/api/admin/endpoints/endpoint-codex-responses"
))
.header(crate::constants::GATEWAY_HEADER, "rust-phase3b")
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
.json(&json!({
"base_url": "http://127.0.0.1:18181/v1",
"max_retries": 0
}))
.send()
.await
.expect("request should succeed");
let status = response.status();
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(status, StatusCode::OK, "payload={payload}");
assert_eq!(payload["base_url"], "http://127.0.0.1:18181/v1");
let endpoints = provider_catalog_repository
.list_endpoints_by_ids(&["endpoint-codex-responses".to_string()])
.await
.expect("endpoints should read");
assert_eq!(endpoints.len(), 1);
assert_eq!(endpoints[0].base_url, "http://127.0.0.1:18181/v1");
assert_eq!(endpoints[0].max_retries, Some(0));
assert!(endpoints[0]
.config
.as_ref()
.and_then(|value| value.get("_aether_fixed_provider_template"))
.and_then(|value| value.get("overrides"))
.and_then(serde_json::Value::as_array)
.is_some_and(|items| items.iter().any(|item| item.as_str() == Some("base_url"))));
gateway_handle.abort();
}
#[tokio::test]
async fn gateway_deletes_admin_provider_endpoint_locally_with_trusted_admin_principal() {
let upstream_hits = Arc::new(Mutex::new(0usize));