i/Aether
1
0
Fork 0
mirror of https://github.com/fawney19/Aether.git synced 2026-09-02 01:10:23 +08:00
Code Issues Packages Projects Releases Wiki Activity

feat: add adaptive pool metrics and self-check

Browse Source
This commit is contained in:
fawney19
2026-05-15 01:46:24 +08:00
parent bf511f9f8c
commit 54a8312e46
34 changed files with 4521 additions and 300 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
apps/aether-gateway/src/admin_api.rs
View File

@@ -2,14 +2,15 @@ pub(crate) use crate::handlers::admin::{
admin_provider_ops_local_action_response, admin_provider_pool_config,
build_internal_control_error_response, create_provider_oauth_catalog_key,
find_duplicate_provider_oauth_key, maybe_build_local_admin_pool_response,
maybe_build_local_admin_response, provider_oauth_maintenance_endpoint_for_provider,
provider_oauth_runtime_endpoint_for_provider, provider_quota_refresh_endpoint_for_provider,
provider_type_supports_quota_refresh, reconcile_admin_fixed_provider_template_endpoints,
maybe_build_local_admin_response, persist_provider_quota_refresh_state,
provider_oauth_maintenance_endpoint_for_provider, provider_oauth_runtime_endpoint_for_provider,
provider_quota_refresh_endpoint_for_provider, provider_type_supports_quota_refresh,
reconcile_admin_fixed_provider_template_endpoints,
refresh_provider_oauth_account_state_after_update, refresh_provider_pool_quota_locally,
update_existing_provider_oauth_catalog_key, AdminAppState,
AdminGatewayProviderTransportSnapshot, AdminLocalOAuthRefreshError, AdminRequestContext,
AdminRouteRequest, AdminRouteResponse, AdminRouteResult, AdminStatsTimeRange,
AdminStatsUsageFilter,
AdminStatsUsageFilter, OAUTH_ACCOUNT_BLOCK_PREFIX, OAUTH_REQUEST_FAILED_PREFIX,
};
use crate::handlers::admin::{

282
apps/aether-gateway/src/dispatch/pool_scheduler.rs
View File

@@ -35,9 +35,11 @@ use crate::handlers::shared::provider_pool::{
AdminProviderPoolRuntimeState,
};
use crate::handlers::shared::{parse_catalog_auth_config_json, provider_key_health_summary};
use crate::maintenance::spawn_pool_quota_probe_replenish_for_request;
use crate::orchestration::LocalExecutionCandidateMetadata;
static LOAD_BALANCE_SEQUENCE: AtomicU64 = AtomicU64::new(0);
const POOL_ACTIVE_PROBE_SEALED_SKIP_REASON: &str = "pool_active_probe_sealed";
type PoolCatalogKeyContext = PoolMemberSignals;
@@ -108,6 +110,7 @@ async fn schedule_pool_page_candidates(
let key_context_by_id = read_pool_catalog_key_contexts_by_id(state, &candidates).await;
let mut runtime_by_provider = BTreeMap::new();
let mut burst_provider_ids = BTreeSet::<String>::new();
for (provider_id, (pool_config, key_ids)) in provider_runtime_requirements {
let key_ids = key_ids.into_iter().collect::<Vec<_>>();
let runtime = if key_ids.is_empty() {
@@ -122,14 +125,29 @@ async fn schedule_pool_page_candidates(
)
.await
};
if should_trigger_active_probe_burst_for_request(&pool_config, &runtime) {
burst_provider_ids.insert(provider_id.clone());
}
runtime_by_provider.insert(provider_id, runtime);
}
apply_local_execution_pool_scheduler_with_runtime_map(
let (scheduled, skipped) = apply_local_execution_pool_scheduler_with_runtime_map(
candidates,
&runtime_by_provider,
&key_context_by_id,
)
);
for skipped_candidate in &skipped {
if skipped_candidate.skip_reason == POOL_ACTIVE_PROBE_SEALED_SKIP_REASON {
burst_provider_ids.insert(skipped_candidate.candidate.provider_id.clone());
}
}
for provider_id in burst_provider_ids {
let _ = spawn_pool_quota_probe_replenish_for_request(state.app().clone(), provider_id);
}
(scheduled, skipped)
}
async fn expand_pool_group_candidate(
@@ -847,48 +865,69 @@ fn apply_local_execution_pool_scheduler_with_runtime_map(
Vec<EligibleLocalExecutionCandidate>,
Vec<SkippedLocalExecutionCandidate>,
) {
let runtime_by_provider = runtime_by_provider
let scheduler_runtime_by_provider = runtime_by_provider
.iter()
.map(|(provider_id, runtime)| (provider_id.clone(), pool_runtime_state(runtime)))
.collect::<BTreeMap<_, _>>();
let inputs = candidates
.into_iter()
.map(|candidate| {
let key_context = key_context_by_id
.get(&candidate.candidate.key_id)
.cloned()
.unwrap_or_default();
PoolCandidateInput {
facts: pool_candidate_facts(&candidate),
pool_config: pool_config_for_candidate(&candidate).map(|config| {
pool_scheduling_config(
config,
candidate.transport.provider.provider_type.as_str(),
)
}),
key_context,
candidate,
let mut inputs = Vec::new();
let mut skipped_candidates = Vec::new();
for candidate in candidates {
let key_context = key_context_by_id
.get(&candidate.candidate.key_id)
.cloned()
.unwrap_or_default();
let admin_pool_config = pool_config_for_candidate(&candidate);
if let Some(config) = admin_pool_config.as_ref() {
if should_enforce_active_probe_sealed_pool(config) {
let active_member_ids = runtime_by_provider
.get(&candidate.candidate.provider_id)
.map(|runtime| &runtime.active_probe_member_ids);
if !active_member_ids
.is_some_and(|members| members.contains(&candidate.candidate.key_id))
{
skipped_candidates.push(SkippedLocalExecutionCandidate {
candidate: candidate.candidate.clone(),
skip_reason: POOL_ACTIVE_PROBE_SEALED_SKIP_REASON,
transport: Some(candidate.transport.clone()),
ranking: candidate.ranking.clone(),
extra_data: None,
});
continue;
}
}
})
.collect::<Vec<_>>();
let outcome = run_pool_scheduler(inputs, &runtime_by_provider, pool_sort_seed().as_str());
}
let pool_config = admin_pool_config.map(|config| {
pool_scheduling_config(config, candidate.transport.provider.provider_type.as_str())
});
inputs.push(PoolCandidateInput {
facts: pool_candidate_facts(&candidate),
pool_config,
key_context,
candidate,
});
}
let outcome = run_pool_scheduler(
inputs,
&scheduler_runtime_by_provider,
pool_sort_seed().as_str(),
);
let candidates = outcome
.candidates
.into_iter()
.map(|scheduled| apply_pool_orchestration(scheduled.candidate, scheduled.orchestration))
.collect::<Vec<_>>();
let skipped_candidates = outcome
.skipped_candidates
.into_iter()
.map(|skipped| SkippedLocalExecutionCandidate {
skipped_candidates.extend(outcome.skipped_candidates.into_iter().map(|skipped| {
SkippedLocalExecutionCandidate {
candidate: skipped.candidate.candidate,
skip_reason: skipped.skip_reason,
transport: Some(skipped.candidate.transport),
ranking: skipped.candidate.ranking,
extra_data: None,
})
.collect::<Vec<_>>();
}
}));
(candidates, skipped_candidates)
}
@@ -899,6 +938,24 @@ fn pool_config_for_candidate(
admin_provider_pool_config_from_config_value(candidate.transport.provider.config.as_ref())
}
fn should_enforce_active_probe_sealed_pool(pool_config: &AdminProviderPoolConfig) -> bool {
pool_config.probing_enabled
}
fn should_trigger_active_probe_burst_for_request(
pool_config: &AdminProviderPoolConfig,
runtime: &AdminProviderPoolRuntimeState,
) -> bool {
if !should_enforce_active_probe_sealed_pool(pool_config) {
return false;
}
if runtime.provider_burst_pending {
return false;
}
let active_count = runtime.active_probe_member_ids.len();
runtime.provider_desired_hot > 0 && active_count < runtime.provider_desired_hot
}
fn pool_key_candidate_order_for_group(
group: &EligibleLocalExecutionCandidate,
) -> StoredPoolKeyCandidateOrder {
@@ -1012,7 +1069,8 @@ fn apply_pool_orchestration(
mod tests {
use super::{
apply_local_execution_pool_scheduler_with_runtime_map, build_pool_catalog_key_context,
pool_config_for_candidate, PoolCatalogKeyContext, PoolKeyCursor,
pool_config_for_candidate, should_trigger_active_probe_burst_for_request,
PoolCatalogKeyContext, PoolKeyCursor, POOL_ACTIVE_PROBE_SEALED_SKIP_REASON,
};
use crate::ai_serving::{
apply_local_runtime_candidate_terminal_reason, EligibleLocalExecutionCandidate,
@@ -1040,7 +1098,7 @@ mod tests {
};
use aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate;
use serde_json::json;
use std::collections::{BTreeMap, VecDeque};
use std::collections::{BTreeMap, BTreeSet, VecDeque};
use std::sync::Arc;
#[test]
@@ -1387,6 +1445,168 @@ mod tests {
);
}
#[test]
fn pool_scheduler_uses_only_active_probe_members_when_active_probe_enabled() {
let provider_config = Some(json!({
"pool_advanced": {
"probing_enabled": true
}
}));
let key_active = sample_eligible_candidate(
"provider-pool",
"endpoint-1",
"key-active",
10,
provider_config.clone(),
);
let key_sealed = sample_eligible_candidate(
"provider-pool",
"endpoint-1",
"key-sealed",
10,
provider_config,
);
let runtime_by_provider = BTreeMap::from([(
"provider-pool".to_string(),
AdminProviderPoolRuntimeState {
active_probe_member_ids: BTreeSet::from(["key-active".to_string()]),
..AdminProviderPoolRuntimeState::default()
},
)]);
let (scheduled, skipped) = apply_local_execution_pool_scheduler_with_runtime_map(
vec![key_sealed, key_active],
&runtime_by_provider,
&BTreeMap::new(),
);
assert_eq!(
scheduled
.iter()
.map(|item| item.candidate.key_id.as_str())
.collect::<Vec<_>>(),
vec!["key-active"]
);
assert_eq!(
skipped
.iter()
.map(|item| (item.candidate.key_id.as_str(), item.skip_reason))
.collect::<Vec<_>>(),
vec![("key-sealed", POOL_ACTIVE_PROBE_SEALED_SKIP_REASON)]
);
}
#[test]
fn pool_scheduler_keeps_pool_out_keys_sealed_when_active_probe_pool_is_empty() {
let provider_config = Some(json!({
"pool_advanced": {
"probing_enabled": true
}
}));
let key_a = sample_eligible_candidate(
"provider-pool",
"endpoint-1",
"key-a",
10,
provider_config.clone(),
);
let key_b =
sample_eligible_candidate("provider-pool", "endpoint-1", "key-b", 10, provider_config);
let runtime_by_provider = BTreeMap::from([(
"provider-pool".to_string(),
AdminProviderPoolRuntimeState::default(),
)]);
let (scheduled, skipped) = apply_local_execution_pool_scheduler_with_runtime_map(
vec![key_a, key_b],
&runtime_by_provider,
&BTreeMap::new(),
);
assert!(scheduled.is_empty());
assert_eq!(
skipped
.iter()
.map(|item| (item.candidate.key_id.as_str(), item.skip_reason))
.collect::<Vec<_>>(),
vec![
("key-a", POOL_ACTIVE_PROBE_SEALED_SKIP_REASON),
("key-b", POOL_ACTIVE_PROBE_SEALED_SKIP_REASON),
]
);
}
#[test]
fn pool_scheduler_triggers_burst_when_auto_hot_target_has_gap() {
let provider_config = Some(json!({
"pool_advanced": {
"probing_enabled": true
}
}));
let candidate =
sample_eligible_candidate("provider-pool", "endpoint-1", "key-a", 10, provider_config);
let pool_config = pool_config_for_candidate(&candidate).expect("pool config should parse");
let runtime = AdminProviderPoolRuntimeState {
active_probe_member_ids: BTreeSet::from(["key-a".to_string()]),
provider_desired_hot: 3,
..AdminProviderPoolRuntimeState::default()
};
assert!(should_trigger_active_probe_burst_for_request(
&pool_config,
&runtime
));
}
#[test]
fn pool_scheduler_ignores_legacy_threshold_fields_for_burst_target() {
let provider_config = Some(json!({
"pool_advanced": {
"probing_enabled": true,
"probing_target_percent": 60,
"probing_target_count": 10
}
}));
let candidate =
sample_eligible_candidate("provider-pool", "endpoint-1", "key-a", 10, provider_config);
let pool_config = pool_config_for_candidate(&candidate).expect("pool config should parse");
let runtime = AdminProviderPoolRuntimeState {
active_probe_member_ids: BTreeSet::from(["key-a".to_string(), "key-b".to_string()]),
provider_desired_hot: 2,
..AdminProviderPoolRuntimeState::default()
};
assert!(!should_trigger_active_probe_burst_for_request(
&pool_config,
&runtime
));
}
#[test]
fn pool_scheduler_skips_burst_when_active_probe_target_is_met() {
let provider_config = Some(json!({
"pool_advanced": {
"probing_enabled": true,
"probing_target_count": 1
}
}));
let candidate =
sample_eligible_candidate("provider-pool", "endpoint-1", "key-a", 10, provider_config);
let pool_config = pool_config_for_candidate(&candidate).expect("pool config should parse");
let runtime = AdminProviderPoolRuntimeState {
active_probe_member_ids: BTreeSet::from(["key-a".to_string()]),
provider_desired_hot: 1,
..AdminProviderPoolRuntimeState::default()
};
assert!(!should_trigger_active_probe_burst_for_request(
&pool_config,
&runtime
));
}
#[test]
fn pool_scheduler_applies_distribution_mode_before_strategy_presets() {
let key_a = sample_eligible_candidate(

20
apps/aether-gateway/src/execution_runtime/stream/execution.rs
View File

@@ -89,6 +89,9 @@ use crate::orchestration::{
LocalExecutionEffect, LocalExecutionEffectContext, LocalHealthFailureEffect,
LocalHealthSuccessEffect, LocalOAuthInvalidationEffect, LocalPoolErrorEffect,
};
use crate::provider_pool_demand::{
acquire_provider_pool_in_flight_guard, ProviderPoolInFlightGuard,
};
use crate::request_candidate_runtime::{
ensure_execution_request_candidate_slot, record_local_request_candidate_status,
record_local_request_candidate_status_snapshot, snapshot_local_request_candidate_status,
@@ -514,6 +517,14 @@ pub(crate) async fn execute_execution_runtime_stream(
.and_then(|context| context.candidate_index)
.map(|value| value.to_string())
.unwrap_or_else(|| "-".to_string());
let mut provider_pool_in_flight_guard = acquire_provider_pool_in_flight_guard(
state.runtime_state.clone(),
&plan.provider_id,
plan.request_id.as_str(),
plan.candidate_id.as_deref(),
key_id.as_str(),
)
.await;
match maybe_execute_kiro_web_search_stream(state, &plan, report_context.as_ref()).await {
Ok(Some(kiro_web_search)) => {
return execute_stream_from_frame_stream(
@@ -527,6 +538,7 @@ pub(crate) async fn execute_execution_runtime_stream(
candidate_started_unix_secs,
stream_started_at,
kiro_web_search.frame_stream,
provider_pool_in_flight_guard.take(),
)
.await;
}
@@ -578,6 +590,7 @@ pub(crate) async fn execute_execution_runtime_stream(
candidate_started_unix_secs,
stream_started_at,
chatgpt_web_image.frame_stream,
provider_pool_in_flight_guard.take(),
)
.await;
}
@@ -673,6 +686,7 @@ pub(crate) async fn execute_execution_runtime_stream(
candidate_started_unix_secs,
stream_started_at,
frame_stream,
provider_pool_in_flight_guard.take(),
)
.await;
}
@@ -737,6 +751,7 @@ pub(crate) async fn execute_execution_runtime_stream(
candidate_started_unix_secs,
stream_started_at,
frame_stream,
provider_pool_in_flight_guard.take(),
)
.await;
}
@@ -822,6 +837,7 @@ pub(crate) async fn execute_execution_runtime_stream(
candidate_started_unix_secs,
stream_started_at,
frame_stream,
provider_pool_in_flight_guard.take(),
)
.await;
}
@@ -1089,6 +1105,7 @@ async fn execute_stream_from_frame_stream(
candidate_started_unix_secs: u64,
stream_started_at: Instant,
frame_stream: BoxStream<'static, Result<Bytes, IoError>>,
in_flight_guard: Option<ProviderPoolInFlightGuard>,
) -> Result<Option<Response<Body>>, GatewayError> {
let request_id = plan.request_id.as_str();
let request_id_for_log = short_request_id(request_id);
@@ -1923,7 +1940,9 @@ async fn execute_stream_from_frame_stream(
};
let plan_kind_for_report = plan_kind.to_string();
let stream_started_at_for_report = stream_started_at;
let provider_pool_in_flight_guard_for_report = in_flight_guard;
tokio::spawn(async move {
let _provider_pool_in_flight_guard = provider_pool_in_flight_guard_for_report;
let mut provider_buffered_body = Vec::new();
let mut buffered_body = Vec::new();
let mut provider_body_truncated = false;
@@ -3136,6 +3155,7 @@ mod tests {
crate::clock::current_unix_ms(),
Instant::now(),
frame_stream,
None,
)
.await
.expect("execution should succeed")

9
apps/aether-gateway/src/execution_runtime/sync/execution.rs
View File

@@ -59,6 +59,7 @@ use crate::orchestration::{
LocalExecutionEffect, LocalExecutionEffectContext, LocalHealthFailureEffect,
LocalHealthSuccessEffect, LocalOAuthInvalidationEffect, LocalPoolErrorEffect,
};
use crate::provider_pool_demand::acquire_provider_pool_in_flight_guard;
use crate::request_candidate_runtime::{
ensure_execution_request_candidate_slot, record_local_request_candidate_extra_data,
record_local_request_candidate_status,
@@ -1066,6 +1067,14 @@ async fn execute_execution_runtime_sync_impl(
},
)
.await;
let _provider_pool_in_flight_guard = acquire_provider_pool_in_flight_guard(
state.runtime_state.clone(),
&plan.provider_id,
plan_request_id.as_str(),
plan_candidate_id.as_deref(),
key_id.as_str(),
)
.await;
#[cfg(not(test))]
let mut result = {
match maybe_execute_chatgpt_web_image_sync(state, &plan, report_context.as_ref()).await {

9
apps/aether-gateway/src/handlers/admin/mod.rs
View File

@@ -26,8 +26,10 @@ pub(crate) use self::provider::oauth::provisioning::{
create_provider_oauth_catalog_key, update_existing_provider_oauth_catalog_key,
};
pub(crate) use self::provider::oauth::quota::dispatch::refresh_provider_pool_quota_locally;
pub(crate) use self::provider::oauth::quota::shared::provider_quota_refresh_endpoint_for_provider;
pub(crate) use self::provider::oauth::quota::shared::provider_type_supports_quota_refresh;
pub(crate) use self::provider::oauth::quota::shared::{
persist_provider_quota_refresh_state, provider_quota_refresh_endpoint_for_provider,
provider_type_supports_quota_refresh,
};
pub(crate) use self::provider::oauth::runtime::{
provider_oauth_maintenance_endpoint_for_provider, provider_oauth_runtime_endpoint_for_provider,
refresh_provider_oauth_account_state_after_update,
@@ -35,6 +37,9 @@ pub(crate) use self::provider::oauth::runtime::{
pub(crate) use self::provider::ops::providers::actions::admin_provider_ops_local_action_response;
pub(crate) use self::provider::pool::config::admin_provider_pool_config;
pub(crate) use self::provider::pool_admin::maybe_build_local_admin_pool_response;
pub(crate) use self::provider::shared::payloads::{
OAUTH_ACCOUNT_BLOCK_PREFIX, OAUTH_REQUEST_FAILED_PREFIX,
};
pub(crate) use self::provider::write::provider::reconcile_admin_fixed_provider_template_endpoints;
pub(crate) use self::provider::{
maybe_build_local_admin_provider_oauth_response, maybe_build_local_admin_providers_response,

91
apps/aether-gateway/src/handlers/admin/provider/pool/config.rs
View File

@@ -37,6 +37,38 @@ fn json_f64(value: &Value) -> Option<f64> {
})
}
fn parse_pool_probe_target_percent(pool_advanced: &Map<String, Value>) -> Option<f64> {
pool_advanced
.get("probing_target_percent")
.or_else(|| pool_advanced.get("probing_active_target_percent"))
.or_else(|| pool_advanced.get("active_probe_target_percent"))
.and_then(json_f64)
.filter(|value| value.is_finite() && *value > 0.0)
.map(|value| value.clamp(0.0, 100.0))
}
fn parse_pool_probe_target_count(pool_advanced: &Map<String, Value>) -> Option<u64> {
pool_advanced
.get("probing_target_count")
.or_else(|| pool_advanced.get("probing_active_target_count"))
.or_else(|| pool_advanced.get("active_probe_target_count"))
.and_then(json_u64)
.filter(|value| *value > 0)
.map(|value| value.min(100_000))
}
fn parse_pool_account_self_check_method(pool_advanced: &Map<String, Value>) -> String {
pool_advanced
.get("account_self_check_method")
.or_else(|| pool_advanced.get("self_check_method"))
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(str::to_ascii_lowercase)
.filter(|value| matches!(value.as_str(), "quota_refresh" | "custom_request"))
.unwrap_or_else(|| "quota_refresh".to_string())
}
fn pool_score_weight(object: &Map<String, Value>, names: &[&str], current: f64) -> f64 {
names
.iter()
@@ -388,7 +420,14 @@ pub(crate) fn admin_provider_pool_config_from_config_value(
health_policy_enabled: true,
probing_enabled: false,
probing_interval_minutes: 10,
probing_target_percent: None,
probing_target_count: None,
probe_concurrency: 4,
account_self_check_enabled: false,
account_self_check_interval_minutes: 60,
account_self_check_concurrency: 4,
account_self_check_method: "quota_refresh".to_string(),
account_self_check_request: None,
score_top_n: 128,
score_fallback_scan_limit: 1024,
score_rules: PoolMemberScoreRules::default(),
@@ -456,12 +495,39 @@ pub(crate) fn admin_provider_pool_config_from_config_value(
.filter(|value| *value > 0)
.map(|value| value.min(1440))
.unwrap_or(10),
probing_target_percent: parse_pool_probe_target_percent(pool_advanced),
probing_target_count: parse_pool_probe_target_count(pool_advanced),
probe_concurrency: pool_advanced
.get("probe_concurrency")
.and_then(json_u64)
.filter(|value| *value > 0)
.map(|value| value.min(64))
.unwrap_or(4),
account_self_check_enabled: pool_advanced
.get("account_self_check_enabled")
.or_else(|| pool_advanced.get("self_check_enabled"))
.and_then(Value::as_bool)
.unwrap_or(false),
account_self_check_interval_minutes: pool_advanced
.get("account_self_check_interval_minutes")
.or_else(|| pool_advanced.get("self_check_interval_minutes"))
.and_then(json_u64)
.filter(|value| *value > 0)
.map(|value| value.min(1440))
.unwrap_or(60),
account_self_check_concurrency: pool_advanced
.get("account_self_check_concurrency")
.or_else(|| pool_advanced.get("self_check_concurrency"))
.and_then(json_u64)
.filter(|value| *value > 0)
.map(|value| value.min(64))
.unwrap_or(4),
account_self_check_method: parse_pool_account_self_check_method(pool_advanced),
account_self_check_request: pool_advanced
.get("account_self_check_request")
.or_else(|| pool_advanced.get("self_check_request"))
.filter(|value| value.is_object())
.cloned(),
score_top_n: pool_advanced
.get("score_top_n")
.and_then(json_u64)
@@ -547,7 +613,18 @@ mod tests {
"health_policy_enabled": false,
"probing_enabled": true,
"probing_interval_minutes": 20,
"probing_target_percent": 25,
"probing_target_count": 3,
"probe_concurrency": 6,
"account_self_check_enabled": true,
"account_self_check_interval_minutes": 90,
"account_self_check_concurrency": 5,
"account_self_check_method": "custom_request",
"account_self_check_request": {
"method": "GET",
"path": "/v1/me",
"success_status_codes": [200]
},
"score_top_n": 256,
"score_fallback_scan_limit": 2048,
"score_rules": {
@@ -584,7 +661,21 @@ mod tests {
assert!(!config.health_policy_enabled);
assert!(config.probing_enabled);
assert_eq!(config.probing_interval_minutes, 20);
assert_eq!(config.probing_target_percent, Some(25.0));
assert_eq!(config.probing_target_count, Some(3));
assert_eq!(config.probe_concurrency, 6);
assert!(config.account_self_check_enabled);
assert_eq!(config.account_self_check_interval_minutes, 90);
assert_eq!(config.account_self_check_concurrency, 5);
assert_eq!(config.account_self_check_method, "custom_request");
assert_eq!(
config
.account_self_check_request
.as_ref()
.and_then(|value| value.get("path"))
.and_then(serde_json::Value::as_str),
Some("/v1/me")
);
assert_eq!(config.score_top_n, 256);
assert_eq!(config.score_fallback_scan_limit, 2048);
assert_eq!(config.score_rules.weights.manual_priority, 0.4);

47
apps/aether-gateway/src/handlers/admin/provider/pool/runtime/reads.rs
View File

@@ -5,10 +5,15 @@ use super::keys::{
};
use crate::handlers::admin::provider::pool::config::admin_provider_pool_cache_affinity_enabled;
use crate::handlers::admin::provider::shared::support::{
AdminProviderPoolConfig, AdminProviderPoolRuntimeState,
admin_provider_pool_quota_probe_active_members_key, AdminProviderPoolConfig,
AdminProviderPoolRuntimeState,
};
use crate::maintenance::PoolQuotaProbeWorkerConfig;
use crate::provider_pool_demand::{
provider_pool_burst_pending, read_provider_pool_demand_snapshot,
};
use aether_runtime_state::{DataLayerError, RuntimeState};
use std::collections::BTreeMap;
use std::collections::{BTreeMap, BTreeSet};
use std::time::{SystemTime, UNIX_EPOCH};
use tracing::warn;
@@ -19,6 +24,10 @@ fn current_unix_secs() -> u64 {
.as_secs()
}
fn should_load_active_probe_members(pool_config: &AdminProviderPoolConfig) -> bool {
pool_config.probing_enabled
}
pub(crate) async fn read_admin_provider_pool_cooldown_counts(
runtime: &RuntimeState,
provider_ids: &[String],
@@ -102,6 +111,40 @@ pub(crate) async fn read_admin_provider_pool_runtime_state(
}
}
if should_load_active_probe_members(pool_config) {
state.active_probe_member_ids = runtime
.set_members(&admin_provider_pool_quota_probe_active_members_key(
provider_id,
))
.await
.map(|values| {
values
.into_iter()
.map(|value| value.trim().to_string())
.filter(|value| !value.is_empty())
.collect::<BTreeSet<_>>()
})
.unwrap_or_default();
}
let probe_config = PoolQuotaProbeWorkerConfig::from_env();
let demand_snapshot = read_provider_pool_demand_snapshot(
runtime,
provider_id,
key_ids.len(),
probe_config.max_keys_per_provider,
)
.await;
state.provider_in_flight = demand_snapshot.in_flight;
state.provider_ema_in_flight = demand_snapshot.ema_in_flight;
state.provider_desired_hot = if pool_config.probing_enabled {
demand_snapshot.desired_hot
} else {
0
};
state.provider_burst_pending =
pool_config.probing_enabled && provider_pool_burst_pending(runtime, provider_id).await;
if !cooldown_keys.is_empty() {
let cooldown_reasons = runtime
.kv_get_many(&cooldown_keys)

10
apps/aether-gateway/src/handlers/admin/provider/pool/runtime/status.rs
View File

@@ -23,6 +23,11 @@ pub(crate) async fn build_admin_provider_pool_status_payload(
"pool_enabled": false,
"total_keys": 0,
"total_sticky_sessions": 0,
"provider_hot_count": 0,
"provider_desired_hot": 0,
"provider_in_flight": 0,
"provider_ema_in_flight": 0.0,
"provider_burst_pending": false,
"keys": [],
}));
};
@@ -68,6 +73,11 @@ pub(crate) async fn build_admin_provider_pool_status_payload(
"pool_enabled": true,
"total_keys": key_payloads.len(),
"total_sticky_sessions": runtime.total_sticky_sessions,
"provider_hot_count": runtime.active_probe_member_ids.len(),
"provider_desired_hot": runtime.provider_desired_hot,
"provider_in_flight": runtime.provider_in_flight,
"provider_ema_in_flight": runtime.provider_ema_in_flight,
"provider_burst_pending": runtime.provider_burst_pending,
"keys": key_payloads,
}))
}

7
apps/aether-gateway/src/handlers/admin/provider/pool/runtime/writes.rs
View File

@@ -615,7 +615,14 @@ mod tests {
health_policy_enabled: true,
probing_enabled: false,
probing_interval_minutes: 10,
probing_target_percent: None,
probing_target_count: None,
probe_concurrency: 4,
account_self_check_enabled: false,
account_self_check_interval_minutes: 60,
account_self_check_concurrency: 4,
account_self_check_method: "quota_refresh".to_string(),
account_self_check_request: None,
score_top_n: 128,
score_fallback_scan_limit: 1024,
score_rules: aether_pool_core::PoolMemberScoreRules::default(),

82
apps/aether-gateway/src/handlers/admin/provider/pool_admin/read_routes/overview.rs
View File

@@ -5,7 +5,12 @@ use super::{
read_admin_provider_pool_cooldown_counts,
ADMIN_POOL_PROVIDER_CATALOG_READER_UNAVAILABLE_DETAIL,
};
use crate::handlers::admin::provider::shared::support::admin_provider_pool_quota_probe_active_members_key;
use crate::handlers::admin::request::AdminAppState;
use crate::maintenance::PoolQuotaProbeWorkerConfig;
use crate::provider_pool_demand::{
provider_pool_burst_pending, read_provider_pool_demand_snapshot,
};
use crate::GatewayError;
use aether_admin::provider::pool as admin_provider_pool_pure;
use axum::{
@@ -14,6 +19,7 @@ use axum::{
response::{IntoResponse, Response},
Json,
};
use serde_json::{json, Value};
pub(super) async fn build_admin_pool_overview_response(
state: &AdminAppState<'_>,
@@ -60,17 +66,77 @@ pub(super) async fn build_admin_pool_overview_response(
.map(|item| (item.provider_id.clone(), item))
.collect::<BTreeMap<_, _>>();
let probe_config = PoolQuotaProbeWorkerConfig::from_env();
let mut runtime_metrics_by_provider = BTreeMap::new();
for (provider, pool_config) in &pool_enabled_providers {
let active_keys = key_stats_by_provider
.get(&provider.id)
.map(|item| item.active_keys as usize)
.unwrap_or(0);
let hot_count = if pool_config.probing_enabled {
state
.runtime_state()
.set_len(&admin_provider_pool_quota_probe_active_members_key(
&provider.id,
))
.await
.unwrap_or(0)
} else {
0
};
let demand_snapshot = read_provider_pool_demand_snapshot(
state.runtime_state(),
&provider.id,
active_keys,
probe_config.max_keys_per_provider,
)
.await;
let burst_pending = pool_config.probing_enabled
&& provider_pool_burst_pending(state.runtime_state(), &provider.id).await;
runtime_metrics_by_provider.insert(
provider.id.clone(),
json!({
"provider_hot_count": hot_count,
"provider_desired_hot": if pool_config.probing_enabled {
demand_snapshot.desired_hot
} else {
0
},
"provider_in_flight": demand_snapshot.in_flight,
"provider_ema_in_flight": demand_snapshot.ema_in_flight,
"provider_burst_pending": burst_pending,
}),
);
}
let providers = pool_enabled_providers
.into_iter()
.map(|(provider, _)| provider)
.collect::<Vec<_>>();
Ok(
Json(admin_provider_pool_pure::build_admin_pool_overview_payload(
&providers,
&key_stats_by_provider,
&cooldown_counts_by_provider,
))
.into_response(),
)
let mut payload = admin_provider_pool_pure::build_admin_pool_overview_payload(
&providers,
&key_stats_by_provider,
&cooldown_counts_by_provider,
);
if let Some(items) = payload.get_mut("items").and_then(Value::as_array_mut) {
for item in items {
let Some(provider_id) = item.get("provider_id").and_then(Value::as_str) else {
continue;
};
let Some(metrics) = runtime_metrics_by_provider.get(provider_id) else {
continue;
};
let Some(item_object) = item.as_object_mut() else {
continue;
};
if let Some(metrics_object) = metrics.as_object() {
for (key, value) in metrics_object {
item_object.insert(key.clone(), value.clone());
}
}
}
}
Ok(Json(payload).into_response())
}

20
apps/aether-gateway/src/handlers/admin/provider/shared/support.rs
View File

@@ -2,15 +2,21 @@ use crate::handlers::admin::request::AdminAppState;
use crate::LocalProviderDeleteTaskState;
use aether_pool_core::PoolMemberScoreRules;
use serde_json::json;
use std::collections::BTreeMap;
use std::collections::{BTreeMap, BTreeSet};
pub(crate) const ADMIN_PROVIDER_MAPPING_PREVIEW_MAX_KEYS: usize = 200;
pub(crate) const ADMIN_PROVIDER_MAPPING_PREVIEW_MAX_MODELS: usize = 500;
pub(crate) const ADMIN_PROVIDER_MAPPING_PREVIEW_FETCH_LIMIT: usize = 10_000;
pub(crate) const ADMIN_PROVIDER_POOL_SCAN_BATCH: u64 = 200;
pub(crate) const ADMIN_PROVIDER_POOL_QUOTA_PROBE_ACTIVE_SET_PREFIX: &str =
"ap:quota_probe:active_members";
pub(crate) const ADMIN_PROVIDER_OAUTH_DATA_UNAVAILABLE_DETAIL: &str =
"Admin provider OAuth data unavailable";
pub(crate) fn admin_provider_pool_quota_probe_active_members_key(provider_id: &str) -> String {
format!("{ADMIN_PROVIDER_POOL_QUOTA_PROBE_ACTIVE_SET_PREFIX}:{provider_id}")
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub(crate) struct AdminProviderPoolSchedulingPreset {
pub(crate) preset: String,
@@ -40,7 +46,14 @@ pub(crate) struct AdminProviderPoolConfig {
pub(crate) health_policy_enabled: bool,
pub(crate) probing_enabled: bool,
pub(crate) probing_interval_minutes: u64,
pub(crate) probing_target_percent: Option<f64>,
pub(crate) probing_target_count: Option<u64>,
pub(crate) probe_concurrency: u64,
pub(crate) account_self_check_enabled: bool,
pub(crate) account_self_check_interval_minutes: u64,
pub(crate) account_self_check_concurrency: u64,
pub(crate) account_self_check_method: String,
pub(crate) account_self_check_request: Option<serde_json::Value>,
pub(crate) score_top_n: u64,
pub(crate) score_fallback_scan_limit: u64,
pub(crate) score_rules: PoolMemberScoreRules,
@@ -54,6 +67,11 @@ pub(crate) struct AdminProviderPoolRuntimeState {
pub(crate) total_sticky_sessions: usize,
pub(crate) sticky_sessions_by_key: BTreeMap<String, usize>,
pub(crate) sticky_bound_key_id: Option<String>,
pub(crate) active_probe_member_ids: BTreeSet<String>,
pub(crate) provider_in_flight: usize,
pub(crate) provider_ema_in_flight: f64,
pub(crate) provider_desired_hot: usize,
pub(crate) provider_burst_pending: bool,
pub(crate) cooldown_reason_by_key: BTreeMap<String, String>,
pub(crate) cooldown_ttl_by_key: BTreeMap<String, u64>,
pub(crate) cost_window_usage_by_key: BTreeMap<String, u64>,

3
apps/aether-gateway/src/handlers/shared/provider_pool.rs
View File

@@ -8,6 +8,7 @@ pub(crate) use super::super::admin::provider::pool::runtime::{
release_admin_provider_pool_key_lease,
};
pub(crate) use super::super::admin::provider::shared::support::{
AdminProviderPoolConfig, AdminProviderPoolRuntimeState, AdminProviderPoolSchedulingPreset,
admin_provider_pool_quota_probe_active_members_key, AdminProviderPoolConfig,
AdminProviderPoolRuntimeState, AdminProviderPoolSchedulingPreset,
AdminProviderPoolUnschedulableRule, ADMIN_PROVIDER_POOL_SCAN_BATCH,
};

1
apps/aether-gateway/src/lib.rs
View File

@@ -53,6 +53,7 @@ mod oauth;
mod orchestration;
mod privacy;
mod provider_key_auth;
mod provider_pool_demand;
pub(crate) use aether_provider_transport as provider_transport;
mod rate_limit;
mod request_candidate_runtime;

26
apps/aether-gateway/src/maintenance/mod.rs
View File

@@ -5,25 +5,27 @@ mod tests;
pub(crate) use runtime::{
cancel_proxy_upgrade_rollout, clear_proxy_upgrade_rollout_conflicts,
ensure_provider_key_pool_scores_for_keys, inspect_proxy_upgrade_rollout,
list_admin_cleanup_run_records, perform_oauth_token_refresh_once,
perform_pool_quota_probe_once, perform_provider_checkin_once, preview_manual_usage_cleanup,
rebuild_admin_stats_once, record_completed_cleanup_run, record_proxy_upgrade_traffic_success,
list_admin_cleanup_run_records, perform_account_self_check_once,
perform_oauth_token_refresh_once, perform_pool_quota_probe_once, perform_provider_checkin_once,
pool_quota_probe_target_count, preview_manual_usage_cleanup, rebuild_admin_stats_once,
record_completed_cleanup_run, record_proxy_upgrade_traffic_success,
restore_proxy_upgrade_rollout_skipped_nodes, retry_proxy_upgrade_rollout_node,
run_admin_system_cleanup_once, run_manual_usage_cleanup_once, skip_proxy_upgrade_rollout_node,
spawn_audit_cleanup_worker, spawn_db_maintenance_worker,
spawn_account_self_check_worker, spawn_audit_cleanup_worker, spawn_db_maintenance_worker,
spawn_gemini_file_mapping_cleanup_worker, spawn_oauth_token_refresh_worker,
spawn_pending_cleanup_worker, spawn_pool_monitor_worker, spawn_pool_quota_probe_worker,
spawn_pending_cleanup_worker, spawn_pool_monitor_worker,
spawn_pool_quota_probe_replenish_for_request, spawn_pool_quota_probe_worker,
spawn_pool_score_rebuild_worker, spawn_provider_checkin_worker,
spawn_proxy_node_metrics_cleanup_worker, spawn_proxy_node_stale_cleanup_worker,
spawn_proxy_upgrade_rollout_worker, spawn_request_candidate_cleanup_worker,
spawn_stats_aggregation_worker, spawn_stats_hourly_aggregation_worker,
spawn_usage_cleanup_worker, spawn_wallet_daily_usage_aggregation_worker,
start_admin_request_body_cleanup_task, start_admin_system_purge_task,
start_proxy_upgrade_rollout, AdminCleanupRunRecord, AdminCleanupTaskKind,
AdminStatsRebuildSummary, AdminSystemCleanupSummary, ManualUsageCleanupError,
OAuthTokenRefreshRunSummary, PoolQuotaProbeRunSummary, ProviderCheckinRunSummary,
ProxyUpgradeRolloutCancelSummary, ProxyUpgradeRolloutConflictClearSummary,
ProxyUpgradeRolloutNodeActionSummary, ProxyUpgradeRolloutProbeConfig,
ProxyUpgradeRolloutSkippedRestoreSummary, ProxyUpgradeRolloutStatus,
ProxyUpgradeRolloutTrackedNodeState,
start_proxy_upgrade_rollout, AccountSelfCheckRunSummary, AdminCleanupRunRecord,
AdminCleanupTaskKind, AdminStatsRebuildSummary, AdminSystemCleanupSummary,
ManualUsageCleanupError, OAuthTokenRefreshRunSummary, PoolQuotaProbeRunSummary,
PoolQuotaProbeWorkerConfig, ProviderCheckinRunSummary, ProxyUpgradeRolloutCancelSummary,
ProxyUpgradeRolloutConflictClearSummary, ProxyUpgradeRolloutNodeActionSummary,
ProxyUpgradeRolloutProbeConfig, ProxyUpgradeRolloutSkippedRestoreSummary,
ProxyUpgradeRolloutStatus, ProxyUpgradeRolloutTrackedNodeState,
};

14
apps/aether-gateway/src/maintenance/runtime.rs
View File

@@ -6,6 +6,8 @@ use crate::admin_api::admin_provider_ops_local_action_response;
use crate::data::GatewayDataState;
use crate::{AppState, GatewayError};
#[path = "runtime/account_self_check.rs"]
mod account_self_check;
#[path = "runtime/audit_cleanup.rs"]
mod audit_cleanup;
#[path = "runtime/cleanup_runs.rs"]
@@ -49,6 +51,11 @@ mod usage_cleanup;
mod wallet_daily_usage;
#[path = "runtime/workers.rs"]
mod workers;
pub(crate) use account_self_check::{
perform_account_self_check_once, perform_account_self_check_once_with_config,
select_account_self_check_key_ids, spawn_account_self_check_worker, AccountSelfCheckRunSummary,
AccountSelfCheckWorkerConfig,
};
pub(crate) use aether_data_contracts::repository::usage::{
UsageCleanupSummary, UsageCleanupWindow,
};
@@ -65,9 +72,10 @@ pub(crate) use oauth_token_refresh::{
};
use pending_cleanup::*;
pub(crate) use pool_quota_probe::{
perform_pool_quota_probe_once, perform_pool_quota_probe_once_with_config,
select_pool_quota_probe_key_ids, spawn_pool_quota_probe_worker, PoolQuotaProbeRunSummary,
PoolQuotaProbeWorkerConfig,
perform_pool_quota_probe_once, perform_pool_quota_probe_once_for_provider_with_config,
perform_pool_quota_probe_once_with_config, pool_quota_probe_target_count,
select_pool_quota_probe_key_ids, spawn_pool_quota_probe_replenish_for_request,
spawn_pool_quota_probe_worker, PoolQuotaProbeRunSummary, PoolQuotaProbeWorkerConfig,
};
pub(crate) use pool_score_rebuild::{
ensure_provider_key_pool_scores_for_keys, perform_pool_score_rebuild_once,

1278
apps/aether-gateway/src/maintenance/runtime/account_self_check.rs Normal file
View File

File diff suppressed because it is too large Load Diff

1597
apps/aether-gateway/src/maintenance/runtime/pool_quota_probe.rs
View File

File diff suppressed because it is too large Load Diff

27
apps/aether-gateway/src/maintenance/runtime/tests.rs
View File

@@ -20,15 +20,15 @@ use super::{
next_stats_hourly_aggregation_run_after, pending_cleanup_batch_size,
pending_cleanup_timeout_minutes, plan_pending_cleanup_batch, provider_checkin_schedule,
proxy_node_metrics_cleanup_settings, record_proxy_upgrade_traffic_success,
run_db_maintenance_with, run_proxy_upgrade_rollout_once, spawn_audit_cleanup_worker,
spawn_db_maintenance_worker, spawn_oauth_token_refresh_worker, spawn_pending_cleanup_worker,
spawn_pool_monitor_worker, spawn_pool_quota_probe_worker, spawn_provider_checkin_worker,
spawn_proxy_node_stale_cleanup_worker, spawn_proxy_upgrade_rollout_worker,
spawn_stats_aggregation_worker, spawn_stats_hourly_aggregation_worker,
spawn_usage_cleanup_worker, spawn_wallet_daily_usage_aggregation_worker,
start_proxy_upgrade_rollout, stats_aggregation_target_day,
stats_hourly_aggregation_target_hour, summarize_database_pool, usage_cleanup_settings,
usage_cleanup_window, usage_cleanup_window_with_override,
run_db_maintenance_with, run_proxy_upgrade_rollout_once, spawn_account_self_check_worker,
spawn_audit_cleanup_worker, spawn_db_maintenance_worker, spawn_oauth_token_refresh_worker,
spawn_pending_cleanup_worker, spawn_pool_monitor_worker, spawn_pool_quota_probe_worker,
spawn_provider_checkin_worker, spawn_proxy_node_stale_cleanup_worker,
spawn_proxy_upgrade_rollout_worker, spawn_stats_aggregation_worker,
spawn_stats_hourly_aggregation_worker, spawn_usage_cleanup_worker,
spawn_wallet_daily_usage_aggregation_worker, start_proxy_upgrade_rollout,
stats_aggregation_target_day, stats_hourly_aggregation_target_hour, summarize_database_pool,
usage_cleanup_settings, usage_cleanup_window, usage_cleanup_window_with_override,
wallet_daily_usage_aggregation_target, AppState, DbMaintenanceRunSummary,
FailedPendingUsageRow, GatewayDataState, ProxyNodeMetricsCleanupSettings,
ProxyUpgradeRolloutProbeConfig, StalePendingUsageRow, UsageCleanupSettings, USAGE_CLEANUP_HOUR,
@@ -98,6 +98,15 @@ async fn spawn_pool_quota_probe_worker_skips_when_provider_catalog_unavailable()
assert!(spawn_pool_quota_probe_worker(state).is_none());
}
#[tokio::test]
async fn spawn_account_self_check_worker_skips_when_provider_catalog_unavailable() {
let state = AppState::new()
.expect("gateway state should build")
.with_data_state_for_tests(GatewayDataState::disabled());
assert!(spawn_account_self_check_worker(state).is_none());
}
fn sample_connected_proxy_node(
node_id: &str,
heartbeat_interval: i32,

416
apps/aether-gateway/src/provider_pool_demand.rs Normal file
View File

@@ -0,0 +1,416 @@
use std::sync::{
atomic::{AtomicBool, Ordering},
Arc,
};
use std::time::{Duration, SystemTime, UNIX_EPOCH};
use aether_runtime_state::RuntimeState;
use serde::{Deserialize, Serialize};
use tokio::task::JoinHandle;
use tracing::debug;
use uuid::Uuid;
const PROVIDER_POOL_IN_FLIGHT_TOKENS_PREFIX: &str = "ap:provider_pool:in_flight";
const PROVIDER_POOL_DEMAND_SNAPSHOT_PREFIX: &str = "ap:provider_pool:demand";
const PROVIDER_POOL_BURST_PENDING_PREFIX: &str = "ap:quota_probe:burst_pending";
const PROVIDER_POOL_IN_FLIGHT_TOKEN_TTL_MS: u64 = 120_000;
const PROVIDER_POOL_IN_FLIGHT_RENEW_MS: u64 = 30_000;
const PROVIDER_POOL_DEMAND_SNAPSHOT_TTL_SECONDS: u64 = 6 * 60 * 60;
const PROVIDER_POOL_DEMAND_ALPHA: f64 = 0.2;
const PROVIDER_POOL_DEMAND_HEADROOM: f64 = 1.2;
const PROVIDER_POOL_DEMAND_FLOOR: usize = 2;
#[derive(Debug, Clone, Copy, PartialEq, Serialize, Deserialize)]
pub(crate) struct ProviderPoolDemandSnapshot {
pub(crate) in_flight: usize,
pub(crate) ema_in_flight: f64,
pub(crate) desired_hot: usize,
pub(crate) sampled_at_unix_ms: u64,
}
pub(crate) struct ProviderPoolInFlightGuard {
runtime: Arc<RuntimeState>,
tokens_key: String,
token: String,
stop_renewal: Arc<AtomicBool>,
renew_handle: Option<JoinHandle<()>>,
released: bool,
}
impl ProviderPoolInFlightGuard {
pub(crate) async fn release(mut self) {
self.release_inner().await;
}
async fn release_inner(&mut self) {
if self.released {
return;
}
self.released = true;
self.stop_renewal.store(true, Ordering::Release);
if let Some(handle) = self.renew_handle.take() {
handle.abort();
}
if let Err(err) = self
.runtime
.score_remove(&self.tokens_key, &self.token)
.await
{
debug!(
error = ?err,
"gateway provider pool demand: failed to release in-flight token"
);
}
}
}
impl Drop for ProviderPoolInFlightGuard {
fn drop(&mut self) {
if self.released {
return;
}
self.released = true;
self.stop_renewal.store(true, Ordering::Release);
if let Some(handle) = self.renew_handle.take() {
handle.abort();
}
let runtime = self.runtime.clone();
let tokens_key = self.tokens_key.clone();
let token = self.token.clone();
if let Ok(handle) = tokio::runtime::Handle::try_current() {
handle.spawn(async move {
if let Err(err) = runtime.score_remove(&tokens_key, &token).await {
debug!(
error = ?err,
"gateway provider pool demand: failed to release dropped in-flight token"
);
}
});
}
}
}
fn current_unix_ms() -> u64 {
let millis = SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap_or_default()
.as_millis();
u64::try_from(millis).unwrap_or(u64::MAX)
}
fn in_flight_tokens_key(provider_id: &str) -> String {
format!("{PROVIDER_POOL_IN_FLIGHT_TOKENS_PREFIX}:{provider_id}")
}
fn demand_snapshot_key(provider_id: &str) -> String {
format!("{PROVIDER_POOL_DEMAND_SNAPSHOT_PREFIX}:{provider_id}")
}
pub(crate) fn provider_pool_burst_pending_key(provider_id: &str) -> String {
format!("{PROVIDER_POOL_BURST_PENDING_PREFIX}:{provider_id}")
}
fn build_in_flight_token(request_id: &str, candidate_id: Option<&str>, key_id: &str) -> String {
let request_id = request_id.trim();
let candidate_id = candidate_id
.map(str::trim)
.filter(|value| !value.is_empty())
.unwrap_or("-");
let key_id = key_id.trim();
format!(
"{}:{}:{}:{}",
current_unix_ms(),
request_id,
candidate_id,
if key_id.is_empty() { "-" } else { key_id },
) + &format!(":{}", Uuid::new_v4())
}
fn token_expiry_score(now_ms: u64) -> f64 {
now_ms.saturating_add(PROVIDER_POOL_IN_FLIGHT_TOKEN_TTL_MS) as f64
}
fn spawn_in_flight_renewal(
runtime: Arc<RuntimeState>,
tokens_key: String,
token: String,
stop: Arc<AtomicBool>,
) -> JoinHandle<()> {
tokio::spawn(async move {
let mut interval =
tokio::time::interval(Duration::from_millis(PROVIDER_POOL_IN_FLIGHT_RENEW_MS));
interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
loop {
interval.tick().await;
if stop.load(Ordering::Acquire) {
break;
}
if let Err(err) = runtime
.score_set(&tokens_key, &token, token_expiry_score(current_unix_ms()))
.await
{
debug!(
error = ?err,
"gateway provider pool demand: failed to renew in-flight token"
);
}
}
})
}
pub(crate) async fn acquire_provider_pool_in_flight_guard(
runtime: Arc<RuntimeState>,
provider_id: &str,
request_id: &str,
candidate_id: Option<&str>,
key_id: &str,
) -> Option<ProviderPoolInFlightGuard> {
let provider_id = provider_id.trim();
if provider_id.is_empty() {
return None;
}
let tokens_key = in_flight_tokens_key(provider_id);
let token = build_in_flight_token(request_id, candidate_id, key_id);
if let Err(err) = runtime
.score_set(&tokens_key, &token, token_expiry_score(current_unix_ms()))
.await
{
debug!(
provider_id,
error = ?err,
"gateway provider pool demand: failed to acquire in-flight token"
);
return None;
}
let stop_renewal = Arc::new(AtomicBool::new(false));
let renew_handle = spawn_in_flight_renewal(
runtime.clone(),
tokens_key.clone(),
token.clone(),
stop_renewal.clone(),
);
Some(ProviderPoolInFlightGuard {
runtime,
tokens_key,
token,
stop_renewal,
renew_handle: Some(renew_handle),
released: false,
})
}
pub(crate) async fn provider_pool_live_in_flight_count(
runtime: &RuntimeState,
provider_id: &str,
) -> usize {
let provider_id = provider_id.trim();
if provider_id.is_empty() {
return 0;
}
let key = in_flight_tokens_key(provider_id);
let now_ms = current_unix_ms() as f64;
if let Err(err) = runtime.score_remove_by_score(&key, now_ms).await {
debug!(
provider_id,
error = ?err,
"gateway provider pool demand: failed to prune expired in-flight tokens"
);
}
runtime.score_len(&key).await.unwrap_or(0)
}
pub(crate) async fn provider_pool_burst_pending(runtime: &RuntimeState, provider_id: &str) -> bool {
let provider_id = provider_id.trim();
if provider_id.is_empty() {
return false;
}
runtime
.kv_exists(&provider_pool_burst_pending_key(provider_id))
.await
.unwrap_or(false)
}
pub(crate) fn provider_pool_desired_hot(
in_flight: usize,
ema_in_flight: f64,
total_active_keys: usize,
max_keys_per_provider: usize,
) -> usize {
let cap = total_active_keys.min(max_keys_per_provider);
if cap == 0 {
return 0;
}
let signal = ema_in_flight
.max(in_flight as f64)
.max(0.0)
.min(usize::MAX as f64);
let desired = (signal * PROVIDER_POOL_DEMAND_HEADROOM).ceil() as usize;
desired.max(PROVIDER_POOL_DEMAND_FLOOR.min(cap)).min(cap)
}
fn parse_stored_demand_snapshot(raw: Option<String>) -> Option<ProviderPoolDemandSnapshot> {
let mut snapshot: ProviderPoolDemandSnapshot = serde_json::from_str(&raw?).ok()?;
if !snapshot.ema_in_flight.is_finite() || snapshot.ema_in_flight < 0.0 {
snapshot.ema_in_flight = 0.0;
}
Some(snapshot)
}
async fn stored_demand_snapshot(
runtime: &RuntimeState,
provider_id: &str,
) -> Option<ProviderPoolDemandSnapshot> {
runtime
.kv_get(&demand_snapshot_key(provider_id))
.await
.ok()
.and_then(parse_stored_demand_snapshot)
}
pub(crate) async fn read_provider_pool_demand_snapshot(
runtime: &RuntimeState,
provider_id: &str,
total_active_keys: usize,
max_keys_per_provider: usize,
) -> ProviderPoolDemandSnapshot {
let in_flight = provider_pool_live_in_flight_count(runtime, provider_id).await;
let stored = stored_demand_snapshot(runtime, provider_id).await;
let ema_in_flight = stored
.as_ref()
.map(|snapshot| snapshot.ema_in_flight)
.unwrap_or(in_flight as f64);
ProviderPoolDemandSnapshot {
in_flight,
ema_in_flight,
desired_hot: provider_pool_desired_hot(
in_flight,
ema_in_flight,
total_active_keys,
max_keys_per_provider,
),
sampled_at_unix_ms: stored
.map(|snapshot| snapshot.sampled_at_unix_ms)
.unwrap_or(0),
}
}
pub(crate) async fn sample_provider_pool_demand(
runtime: &RuntimeState,
provider_id: &str,
total_active_keys: usize,
max_keys_per_provider: usize,
) -> ProviderPoolDemandSnapshot {
let in_flight = provider_pool_live_in_flight_count(runtime, provider_id).await;
let previous = stored_demand_snapshot(runtime, provider_id).await;
let previous_ema = previous
.as_ref()
.map(|snapshot| snapshot.ema_in_flight)
.unwrap_or(in_flight as f64);
let ema_in_flight = if previous.is_some() {
previous_ema.mul_add(
1.0 - PROVIDER_POOL_DEMAND_ALPHA,
in_flight as f64 * PROVIDER_POOL_DEMAND_ALPHA,
)
} else {
in_flight as f64
}
.max(0.0);
let sampled_at_unix_ms = current_unix_ms();
let snapshot = ProviderPoolDemandSnapshot {
in_flight,
ema_in_flight,
desired_hot: provider_pool_desired_hot(
in_flight,
ema_in_flight,
total_active_keys,
max_keys_per_provider,
),
sampled_at_unix_ms,
};
if let Ok(serialized) = serde_json::to_string(&snapshot) {
if let Err(err) = runtime
.kv_set(
&demand_snapshot_key(provider_id),
serialized,
Some(Duration::from_secs(
PROVIDER_POOL_DEMAND_SNAPSHOT_TTL_SECONDS,
)),
)
.await
{
debug!(
provider_id,
error = ?err,
"gateway provider pool demand: failed to persist demand snapshot"
);
}
}
snapshot
}
#[cfg(test)]
mod tests {
use super::*;
use aether_runtime_state::{MemoryRuntimeStateConfig, RuntimeState};
#[tokio::test]
async fn in_flight_guard_tracks_and_releases_provider_tokens() {
let runtime = Arc::new(RuntimeState::memory(MemoryRuntimeStateConfig::default()));
let guard = acquire_provider_pool_in_flight_guard(
runtime.clone(),
"provider-1",
"request-1",
Some("candidate-1"),
"key-1",
)
.await
.expect("guard should be acquired");
assert_eq!(
provider_pool_live_in_flight_count(runtime.as_ref(), "provider-1").await,
1
);
guard.release().await;
assert_eq!(
provider_pool_live_in_flight_count(runtime.as_ref(), "provider-1").await,
0
);
}
#[tokio::test]
async fn demand_snapshot_uses_instant_in_flight_for_fast_rise_and_ema_for_fall() {
let runtime = RuntimeState::memory(MemoryRuntimeStateConfig::default());
for idx in 0..10 {
let guard = acquire_provider_pool_in_flight_guard(
Arc::new(runtime.clone()),
"provider-1",
"request-1",
Some(&format!("candidate-{idx}")),
"key-1",
)
.await
.expect("guard");
std::mem::forget(guard);
}
let high = sample_provider_pool_demand(&runtime, "provider-1", 100, 50).await;
assert_eq!(high.in_flight, 10);
assert_eq!(high.desired_hot, 12);
let key = in_flight_tokens_key("provider-1");
let _ = runtime.score_remove_by_score(&key, f64::INFINITY).await;
let low = sample_provider_pool_demand(&runtime, "provider-1", 100, 50).await;
assert_eq!(low.in_flight, 0);
assert!(low.ema_in_flight > 0.0);
assert!(low.desired_hot >= PROVIDER_POOL_DEMAND_FLOOR);
assert!(low.desired_hot < high.desired_hot);
}
}

5
apps/aether-gateway/src/state/core.rs
View File

@@ -40,6 +40,7 @@ use super::super::router::RequestAdmissionError;
use super::super::{control::GatewayControlDecision, error::GatewayError};
use super::super::{provider_transport, usage};
use crate::maintenance::spawn_account_self_check_worker;
use crate::maintenance::spawn_audit_cleanup_worker;
use crate::maintenance::spawn_db_maintenance_worker;
use crate::maintenance::spawn_gemini_file_mapping_cleanup_worker;
@@ -1172,6 +1173,10 @@ impl AppState {
crate::task_runtime::TASK_KEY_POOL_QUOTA_PROBE,
spawn_pool_quota_probe_worker(self.clone()),
);
supervise_worker(
crate::task_runtime::TASK_KEY_ACCOUNT_SELF_CHECK,
spawn_account_self_check_worker(self.clone()),
);
supervise_worker(
crate::task_runtime::TASK_KEY_POOL_SCORE_REBUILD,
spawn_pool_score_rebuild_worker(self.clone()),

9
apps/aether-gateway/src/task_runtime/mod.rs
View File

@@ -22,6 +22,7 @@ pub(crate) const TASK_KEY_VIDEO_TASK_POLLER: &str = "video.task.poller";
pub(crate) const TASK_KEY_MODEL_FETCH_WORKER: &str = "model.fetch.worker";
pub(crate) const TASK_KEY_PROVIDER_QUOTA_RESET: &str = "provider.quota.reset.worker";
pub(crate) const TASK_KEY_POOL_QUOTA_PROBE: &str = "pool.quota.probe.worker";
pub(crate) const TASK_KEY_ACCOUNT_SELF_CHECK: &str = "account.self_check.worker";
pub(crate) const TASK_KEY_POOL_SCORE_REBUILD: &str = "pool.score.rebuild.worker";
pub(crate) const TASK_KEY_POOL_MONITOR: &str = "pool.monitor.worker";
pub(crate) const TASK_KEY_AUDIT_CLEANUP: &str = "maintenance.audit.cleanup";
@@ -102,6 +103,14 @@ const TASK_DEFINITIONS: &[TaskDefinition] = &[
true,
RETRY_ONCE,
),
TaskDefinition::new(
TASK_KEY_ACCOUNT_SELF_CHECK,
TaskKind::Scheduled,
"interval",
true,
true,
RETRY_ONCE,
),
TaskDefinition::new(
TASK_KEY_POOL_SCORE_REBUILD,
TaskKind::Scheduled,

5
apps/aether-gateway/src/tests/control/admin/pool.rs
View File

@@ -236,6 +236,11 @@ async fn gateway_handles_admin_pool_overview_locally_with_trusted_admin_principa
assert_eq!(items[0]["active_keys"], 1);
assert_eq!(items[0]["cooldown_count"], 0);
assert_eq!(items[0]["pool_enabled"], true);
assert_eq!(items[0]["provider_hot_count"], 0);
assert_eq!(items[0]["provider_desired_hot"], 0);
assert_eq!(items[0]["provider_in_flight"], 0);
assert_eq!(items[0]["provider_ema_in_flight"], 0.0);
assert_eq!(items[0]["provider_burst_pending"], false);
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();

5
apps/aether-gateway/src/tests/control/admin/providers.rs
View File

@@ -1839,6 +1839,11 @@ async fn gateway_handles_admin_provider_pool_status_locally_with_trusted_admin_p
assert_eq!(payload["pool_enabled"], true);
assert_eq!(payload["total_keys"], 1);
assert_eq!(payload["total_sticky_sessions"], 0);
assert_eq!(payload["provider_hot_count"], 0);
assert_eq!(payload["provider_desired_hot"], 0);
assert_eq!(payload["provider_in_flight"], 0);
assert_eq!(payload["provider_ema_in_flight"], 0.0);
assert_eq!(payload["provider_burst_pending"], false);
let keys = payload["keys"].as_array().expect("keys should be an array");
assert_eq!(keys.len(), 1);
assert_eq!(keys[0]["key_id"], "key-openai-pool");