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

Improve gateway transport and usage runtime

Browse Source
This commit is contained in:
elky
2026-06-25 22:36:27 +08:00
parent d336d1a7fa
commit 6f00e9fc67
112 changed files with 12456 additions and 1387 deletions
Download Patch File Download Diff File Expand all files Collapse all files
apps/aether-gateway/src/state/app.rs
+248 -23
View File
@@ -2,13 +2,15 @@ use std::collections::HashMap;
use std::sync::atomic::AtomicU64;
use std::sync::Arc;
use std::sync::Mutex as StdMutex;
use std::sync::RwLock as StdRwLock;
use std::time::Duration;
use aether_data::repository::users::StoredUserGroup;
use aether_data_contracts::repository::billing::UserDailyQuotaAvailabilityRecord;
use aether_data_contracts::repository::quota::StoredProviderQuotaSnapshot;
use aether_runtime::ConcurrencyGate;
use aether_runtime_state::{RuntimeSemaphore, RuntimeState};
use dashmap::DashMap;
use tokio::sync::Mutex as TokioMutex;
use super::super::async_task::{VideoTaskPollerConfig, VideoTaskService};
use super::super::cache::{
@@ -39,14 +41,25 @@ 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_CANDIDATE_PLANNING_GATE_LIMIT: usize = 1024;
const DEFAULT_UPSTREAM_EXECUTION_GATE_LIMIT: usize = 2000;
const DEFAULT_UPSTREAM_TARGET_GATE_LIMIT: usize = 2000;
const DEFAULT_UPSTREAM_EXECUTION_GATE_LIMIT: usize = 10_000;
const DEFAULT_UPSTREAM_TARGET_GATE_LIMIT: usize = 10_000;
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 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 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";
const INTERNAL_GATE_QUEUE_BUDGET_MS_ENV: &str = "AETHER_GATEWAY_INTERNAL_GATE_QUEUE_BUDGET_MS";
const DEFAULT_AUTH_CAPACITY_CACHE_TTL_MS: u64 = 500;
const MIN_AUTH_CAPACITY_CACHE_TTL_MS: u64 = 10;
const MAX_AUTH_CAPACITY_CACHE_TTL_MS: u64 = 10_000;
const AUTH_CAPACITY_CACHE_TTL_MS_ENV: &str = "AETHER_GATEWAY_AUTH_CAPACITY_CACHE_TTL_MS";
#[cfg(test)]
type TestExecutionRuntimeSyncOverrideFn = dyn Fn(
@@ -73,6 +86,7 @@ pub(crate) struct FrontdoorRuntimeGuardConfig {
pub(crate) request_body_read_timeout: Duration,
pub(crate) local_execution_planning_timeout: Duration,
pub(crate) internal_gate_queue_budget: Duration,
pub(crate) auth_capacity_cache_ttl: Duration,
pub(crate) candidate_planning_gate_limit: Option<usize>,
pub(crate) upstream_execution_gate_limit: Option<usize>,
pub(crate) upstream_target_gate_limit: Option<usize>,
@@ -99,18 +113,15 @@ impl FrontdoorRuntimeGuardConfig {
1,
MAX_INTERNAL_GATE_QUEUE_BUDGET_MS,
),
candidate_planning_gate_limit: env_optional_usize(
CANDIDATE_PLANNING_GATE_LIMIT_ENV,
DEFAULT_CANDIDATE_PLANNING_GATE_LIMIT,
),
upstream_execution_gate_limit: env_optional_usize(
UPSTREAM_EXECUTION_GATE_LIMIT_ENV,
DEFAULT_UPSTREAM_EXECUTION_GATE_LIMIT,
),
upstream_target_gate_limit: env_optional_usize(
UPSTREAM_TARGET_GATE_LIMIT_ENV,
DEFAULT_UPSTREAM_TARGET_GATE_LIMIT,
auth_capacity_cache_ttl: env_cache_duration_ms(
AUTH_CAPACITY_CACHE_TTL_MS_ENV,
DEFAULT_AUTH_CAPACITY_CACHE_TTL_MS,
MIN_AUTH_CAPACITY_CACHE_TTL_MS,
MAX_AUTH_CAPACITY_CACHE_TTL_MS,
),
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(),
}
}
@@ -125,6 +136,7 @@ impl FrontdoorRuntimeGuardConfig {
internal_gate_queue_budget: Duration::from_millis(
DEFAULT_INTERNAL_GATE_QUEUE_BUDGET_MS,
),
auth_capacity_cache_ttl: Duration::from_millis(DEFAULT_AUTH_CAPACITY_CACHE_TTL_MS),
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),
@@ -142,15 +154,152 @@ fn env_duration_ms(key: &str, default_ms: u64, min_ms: u64, max_ms: u64) -> Dura
Duration::from_millis(ms)
}
fn env_optional_usize(key: &str, default_value: usize) -> Option<usize> {
match std::env::var(key)
fn env_cache_duration_ms(key: &str, default_ms: u64, min_ms: u64, max_ms: u64) -> Duration {
let Some(raw) = std::env::var(key)
.ok()
.and_then(|value| value.trim().parse::<usize>().ok())
{
Some(0) => None,
Some(value) => Some(value.max(1)),
None => Some(default_value.max(1)),
.map(|value| value.trim().to_string())
.filter(|value| !value.is_empty())
else {
return Duration::from_millis(default_ms);
};
let Some(parsed) = raw.parse::<u64>().ok() else {
return Duration::from_millis(default_ms);
};
if parsed == 0 {
return Duration::ZERO;
}
Duration::from_millis(parsed.clamp(min_ms, max_ms))
}
#[derive(Debug, Clone, Copy)]
struct GateAutoProfile {
floor: usize,
cap: usize,
per_cpu: usize,
fd_divisor: Option<usize>,
}
#[derive(Debug, Clone, Copy)]
struct GateAutoCapacity {
cpu_parallelism: usize,
fd_soft_limit: usize,
}
const CANDIDATE_PLANNING_GATE_AUTO_PROFILE: GateAutoProfile = GateAutoProfile {
floor: DEFAULT_CANDIDATE_PLANNING_GATE_LIMIT,
cap: MAX_CANDIDATE_PLANNING_GATE_LIMIT,
per_cpu: CANDIDATE_PLANNING_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,
per_cpu: UPSTREAM_EXECUTION_GATE_LIMIT_PER_CPU,
fd_divisor: Some(2),
};
const UPSTREAM_TARGET_GATE_AUTO_PROFILE: GateAutoProfile = GateAutoProfile {
floor: DEFAULT_UPSTREAM_TARGET_GATE_LIMIT,
cap: MAX_UPSTREAM_TARGET_GATE_LIMIT,
per_cpu: UPSTREAM_TARGET_GATE_LIMIT_PER_CPU,
fd_divisor: Some(4),
};
fn candidate_planning_gate_limit_from_env() -> Option<usize> {
env_gate_limit(
CANDIDATE_PLANNING_GATE_LIMIT_ENV,
CANDIDATE_PLANNING_GATE_AUTO_PROFILE,
)
}
fn upstream_execution_gate_limit_from_env() -> Option<usize> {
env_gate_limit(
UPSTREAM_EXECUTION_GATE_LIMIT_ENV,
UPSTREAM_EXECUTION_GATE_AUTO_PROFILE,
)
}
pub(crate) fn upstream_target_gate_limit_from_env() -> Option<usize> {
env_gate_limit(
UPSTREAM_TARGET_GATE_LIMIT_ENV,
UPSTREAM_TARGET_GATE_AUTO_PROFILE,
)
}
pub(crate) fn upstream_target_gate_auto_limit() -> usize {
auto_gate_limit(
UPSTREAM_TARGET_GATE_AUTO_PROFILE,
current_gate_auto_capacity(),
)
}
fn env_gate_limit(key: &str, profile: GateAutoProfile) -> Option<usize> {
let raw = std::env::var(key).ok();
parse_gate_limit_value(raw.as_deref(), profile, current_gate_auto_capacity())
}
fn parse_gate_limit_value(
raw: Option<&str>,
profile: GateAutoProfile,
capacity: GateAutoCapacity,
) -> Option<usize> {
let Some(value) = raw.map(str::trim).filter(|value| !value.is_empty()) else {
return Some(auto_gate_limit(profile, capacity));
};
let normalized = value.to_ascii_lowercase();
match normalized.as_str() {
"auto" => Some(auto_gate_limit(profile, capacity)),
"off" | "none" | "disabled" | "disable" => None,
_ => match value.parse::<usize>() {
Ok(0) => None,
Ok(limit) => Some(limit.max(1)),
Err(_) => Some(auto_gate_limit(profile, capacity)),
},
}
}
fn auto_gate_limit(profile: GateAutoProfile, capacity: GateAutoCapacity) -> usize {
let cpu_limit = capacity
.cpu_parallelism
.max(1)
.saturating_mul(profile.per_cpu.max(1));
let mut limit = cpu_limit.max(profile.floor).min(profile.cap);
if let Some(fd_divisor) = profile.fd_divisor.filter(|value| *value > 0) {
let fd_budget = capacity
.fd_soft_limit
.saturating_sub(GATE_LIMIT_FD_RESERVE)
.checked_div(fd_divisor)
.unwrap_or(1)
.max(1);
limit = limit.min(fd_budget);
}
limit.max(1)
}
fn current_gate_auto_capacity() -> GateAutoCapacity {
GateAutoCapacity {
cpu_parallelism: std::thread::available_parallelism()
.map(|value| value.get())
.unwrap_or(1)
.max(1),
fd_soft_limit: soft_fd_limit().unwrap_or(1024).max(1),
}
}
fn soft_fd_limit() -> Option<usize> {
#[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 usize::try_from(limit.rlim_cur).ok();
}
}
None
}
#[derive(Debug, Clone)]
@@ -178,8 +327,15 @@ pub struct AppState {
pub(crate) auth_api_key_force_capabilities_cache:
Arc<JsonValueCache<AuthApiKeyIdentityCacheKey>>,
pub(crate) auth_api_key_feature_settings_cache: Arc<JsonValueCache<AuthApiKeyFeatureCacheKey>>,
pub(crate) auth_daily_quota_availability_cache:
Arc<ValueCache<String, UserDailyQuotaAvailabilityRecord>>,
pub(crate) auth_wallet_snapshot_cache:
Arc<ValueCache<String, aether_data::repository::wallet::StoredWalletSnapshot>>,
pub(crate) auth_request_cost_upper_bound_cache: Arc<ValueCache<String, f64>>,
pub(crate) provider_quota_snapshot_cache: Arc<ValueCache<String, StoredProviderQuotaSnapshot>>,
pub(crate) user_groups_for_user_cache: Arc<ValueCache<String, Vec<StoredUserGroup>>>,
pub(crate) routing_group_selection_cache:
Arc<ValueCache<String, crate::routing::GatewayRoutingGroupSelection>>,
pub(crate) auth_api_key_last_used_cache: Arc<AuthApiKeyLastUsedCache>,
pub(crate) oauth_refresh: Arc<provider_transport::LocalOAuthRefreshCoordinator>,
pub(crate) direct_plan_bypass_cache: Arc<DirectPlanBypassCache>,
@@ -187,6 +343,7 @@ pub struct AppState {
pub(crate) scheduler_affinity_epoch: Arc<AtomicU64>,
pub(crate) dashboard_response_cache: Arc<DashboardResponseCache>,
pub(crate) system_config_cache: Arc<SystemConfigCache>,
pub(crate) candidate_row_page_cache: Arc<super::super::cache::CandidateRowPageCache>,
pub(crate) candidate_page_cache: Arc<super::super::cache::CandidatePageCache>,
pub(crate) candidate_resolved_page_cache: Arc<super::super::cache::CandidateResolvedPageCache>,
pub(crate) chat_pii_redaction_runtime_config_cache:
@@ -197,10 +354,12 @@ pub struct AppState {
pub(crate) frontdoor_user_rpm: Arc<FrontdoorUserRpmLimiter>,
pub(crate) tunnel: crate::tunnel::EmbeddedTunnelState,
pub(crate) provider_transport_snapshot_cache:
Arc<StdRwLock<HashMap<ProviderTransportSnapshotCacheKey, CachedProviderTransportSnapshot>>>,
Arc<DashMap<ProviderTransportSnapshotCacheKey, CachedProviderTransportSnapshot>>,
pub(crate) provider_transport_snapshot_inflight:
Arc<DashMap<ProviderTransportSnapshotCacheKey, Arc<TokioMutex<()>>>>,
pub(crate) provider_key_rpm_resets: Arc<StdMutex<HashMap<String, u64>>>,
pub(crate) local_execution_runtime_miss_diagnostics:
Arc<StdMutex<HashMap<String, LocalExecutionRuntimeMissDiagnostic>>>,
Arc<DashMap<String, LocalExecutionRuntimeMissDiagnostic>>,
pub(crate) admin_monitoring_error_stats_reset_at: Arc<StdMutex<Option<u64>>>,
pub(crate) provider_delete_tasks: Arc<StdMutex<HashMap<String, LocalProviderDeleteTaskState>>>,
#[cfg(test)]
@@ -262,3 +421,69 @@ pub struct AppState {
#[cfg(test)]
pub(crate) provider_oauth_token_url_overrides: Arc<StdMutex<HashMap<String, String>>>,
}
#[cfg(test)]
mod tests {
use super::*;
const TEST_PROFILE: GateAutoProfile = GateAutoProfile {
floor: 10_000,
cap: 16_384,
per_cpu: 1024,
fd_divisor: Some(2),
};
const TEST_CAPACITY: GateAutoCapacity = GateAutoCapacity {
cpu_parallelism: 12,
fd_soft_limit: 1_048_576,
};
#[test]
fn gate_limit_parser_defaults_to_auto() {
assert_eq!(
parse_gate_limit_value(None, TEST_PROFILE, TEST_CAPACITY),
Some(12_288)
);
assert_eq!(
parse_gate_limit_value(Some("auto"), TEST_PROFILE, TEST_CAPACITY),
Some(12_288)
);
assert_eq!(
parse_gate_limit_value(Some(" AUTO "), TEST_PROFILE, TEST_CAPACITY),
Some(12_288)
);
}
#[test]
fn gate_limit_parser_accepts_fixed_numbers() {
assert_eq!(
parse_gate_limit_value(Some("4096"), TEST_PROFILE, TEST_CAPACITY),
Some(4096)
);
}
#[test]
fn gate_limit_parser_accepts_off_and_legacy_zero() {
for value in ["off", "none", "disabled", "disable", "0"] {
assert_eq!(
parse_gate_limit_value(Some(value), TEST_PROFILE, TEST_CAPACITY),
None
);
}
}
#[test]
fn auto_gate_limit_respects_fd_budget_when_fd_limit_is_low() {
assert_eq!(
parse_gate_limit_value(
Some("auto"),
TEST_PROFILE,
GateAutoCapacity {
cpu_parallelism: 32,
fd_soft_limit: 1024,
},
),
Some(448)
);
}
}
apps/aether-gateway/src/state/catalog.rs
+104 -2
View File
@@ -672,7 +672,30 @@ impl AppState {
.await
.map_err(|err| GatewayError::Internal(err.to_string()))?;
if updated.is_some() {
self.invalidate_provider_health_routing_caches();
self.invalidate_provider_runtime_state_caches();
}
Ok(updated)
}
pub(crate) async fn update_provider_catalog_key_success_health_state(
&self,
key_id: &str,
is_active: bool,
health_by_format: Option<&serde_json::Value>,
circuit_breaker_by_format: Option<&serde_json::Value>,
) -> Result<bool, GatewayError> {
let updated = self
.data
.update_provider_catalog_key_health_state(
key_id,
is_active,
health_by_format,
circuit_breaker_by_format,
)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))?;
if updated {
self.invalidate_provider_runtime_state_caches();
}
Ok(updated)
}
@@ -920,7 +943,8 @@ mod tests {
};
use async_trait::async_trait;
use crate::cache::SchedulerAffinityTarget;
use crate::cache::{CandidatePageCacheKey, SchedulerAffinityTarget};
use crate::data::auth::GatewayAuthApiKeySnapshot;
use crate::data::GatewayDataState;
use crate::AppState;
@@ -969,6 +993,35 @@ mod tests {
.expect("key should build")
}
fn sample_auth_snapshot() -> GatewayAuthApiKeySnapshot {
GatewayAuthApiKeySnapshot {
user_id: "user-1".to_string(),
username: "alice".to_string(),
email: None,
user_role: "user".to_string(),
user_auth_source: "local".to_string(),
user_is_active: true,
user_is_deleted: false,
user_rate_limit: None,
user_allowed_providers: None,
user_allowed_api_formats: None,
user_allowed_models: None,
api_key_id: "api-key-1".to_string(),
api_key_name: Some("default".to_string()),
api_key_is_active: true,
api_key_is_locked: false,
api_key_is_standalone: false,
api_key_rate_limit: None,
api_key_concurrent_limit: None,
api_key_expires_at_unix_secs: None,
api_key_allowed_providers: None,
api_key_allowed_api_formats: None,
api_key_allowed_models: None,
api_key_ip_rules: None,
currently_usable: true,
}
}
fn sample_admin_global_model() -> StoredAdminGlobalModel {
StoredAdminGlobalModel::new(
"global-1".to_string(),
@@ -1260,4 +1313,53 @@ mod tests {
Some(target)
);
}
#[tokio::test]
async fn provider_catalog_runtime_state_update_keeps_candidate_page_cache() {
let repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![sample_provider()],
vec![sample_endpoint()],
vec![sample_key()],
));
let state = AppState::new()
.expect("app state should build")
.with_data_state_for_tests(
GatewayDataState::with_provider_catalog_repository_for_tests(repository)
.with_encryption_key_for_tests("test-encryption-key"),
);
let ttl = Duration::from_secs(300);
let cache_key = CandidatePageCacheKey::new(
"gpt-5",
"openai:chat",
true,
&sample_auth_snapshot(),
None,
None,
None,
state.scheduler_affinity_epoch(),
"fixed_order",
true,
None,
);
state.candidate_page_cache.insert(
cache_key.clone(),
Some(Arc::new(crate::cache::CandidatePageSnapshot {
candidates: Vec::new(),
skipped_candidates: Vec::new(),
})),
ttl,
);
assert!(state.candidate_page_cache.get(&cache_key, ttl).is_some());
let mut updated_key = sample_key();
updated_key.status_snapshot = Some(serde_json::json!({"source": "runtime"}));
let updated = state
.update_provider_catalog_key_runtime_state(&updated_key)
.await
.expect("runtime state update should succeed");
assert!(updated.is_some());
assert!(state.candidate_page_cache.get(&cache_key, ttl).is_some());
}
}
apps/aether-gateway/src/state/core.rs
+145 -20
View File
@@ -20,6 +20,7 @@ use aether_runtime_state::{
RuntimeSemaphoreSnapshot, RuntimeState,
};
use aether_scheduler_core::PROVIDER_KEY_RPM_WINDOW_SECS;
use dashmap::DashMap;
use tracing::warn;
use super::{
@@ -100,6 +101,10 @@ fn system_config_key_affects_chat_pii_redaction(key: &str) -> bool {
.starts_with(CHAT_PII_REDACTION_SYSTEM_CONFIG_PREFIX)
}
fn system_config_key_affects_provider_transport_snapshot(key: &str) -> bool {
key.trim() == "enable_format_conversion"
}
impl AppState {
pub async fn prewarm_chat_pii_redaction_runtime_config(&self) -> Result<bool, String> {
crate::privacy::read_chat_pii_redaction_runtime_config(self)
@@ -185,6 +190,7 @@ impl AppState {
self.clear_provider_transport_snapshot_cache();
self.invalidate_scheduler_affinity_cache();
self.invalidate_auth_context_cache();
self.candidate_row_page_cache.clear();
self.candidate_resolved_page_cache.clear();
self.system_config_cache.clear();
self.frontdoor_user_rpm.clear_system_default_cache();
@@ -193,6 +199,7 @@ impl AppState {
.clone()
.with_usage_worker_queue(Self::usage_worker_queue_for(&self.runtime_state)),
);
self.candidate_row_page_cache.clear();
self.candidate_page_cache.clear();
self.candidate_resolved_page_cache.clear();
self.tunnel = crate::tunnel::EmbeddedTunnelState::with_data_and_runtime_state(
@@ -275,8 +282,12 @@ impl AppState {
user_feature_settings_cache: Arc::new(JsonValueCache::default()),
auth_api_key_force_capabilities_cache: Arc::new(JsonValueCache::default()),
auth_api_key_feature_settings_cache: Arc::new(JsonValueCache::default()),
auth_daily_quota_availability_cache: Arc::new(ValueCache::default()),
auth_wallet_snapshot_cache: Arc::new(ValueCache::default()),
auth_request_cost_upper_bound_cache: Arc::new(ValueCache::default()),
provider_quota_snapshot_cache: Arc::new(ValueCache::default()),
user_groups_for_user_cache: Arc::new(ValueCache::default()),
routing_group_selection_cache: Arc::new(ValueCache::default()),
auth_api_key_last_used_cache: Arc::new(AuthApiKeyLastUsedCache::default()),
oauth_refresh: Arc::new(provider_transport::LocalOAuthRefreshCoordinator::new()),
direct_plan_bypass_cache: Arc::new(DirectPlanBypassCache::default()),
@@ -284,6 +295,7 @@ impl AppState {
scheduler_affinity_epoch: Arc::new(AtomicU64::new(0)),
dashboard_response_cache: Arc::new(DashboardResponseCache::default()),
system_config_cache: Arc::new(SystemConfigCache::default()),
candidate_row_page_cache: Arc::new(crate::cache::CandidateRowPageCache::default()),
candidate_page_cache: Arc::new(crate::cache::CandidatePageCache::default()),
candidate_resolved_page_cache: Arc::new(
crate::cache::CandidateResolvedPageCache::default(),
@@ -300,9 +312,10 @@ impl AppState {
data,
runtime_state.clone(),
),
provider_transport_snapshot_cache: Arc::new(std::sync::RwLock::new(HashMap::new())),
provider_transport_snapshot_cache: Arc::new(DashMap::new()),
provider_transport_snapshot_inflight: Arc::new(DashMap::new()),
provider_key_rpm_resets: Arc::new(StdMutex::new(HashMap::new())),
local_execution_runtime_miss_diagnostics: Arc::new(StdMutex::new(HashMap::new())),
local_execution_runtime_miss_diagnostics: Arc::new(DashMap::new()),
admin_monitoring_error_stats_reset_at: Arc::new(StdMutex::new(None)),
provider_delete_tasks: Arc::new(StdMutex::new(HashMap::new())),
#[cfg(test)]
@@ -671,12 +684,18 @@ impl AppState {
&self.chat_pii_redaction_runtime_config_cache,
);
}
if deleted && system_config_key_affects_provider_transport_snapshot(key) {
self.clear_provider_transport_snapshot_cache();
}
Ok(deleted)
}
pub(crate) fn invalidate_provider_routing_caches(&self) {
self.data.clear_minimal_candidate_selection_cache();
self.data.clear_routing_group_cache();
self.data.clear_provider_catalog_cache();
self.routing_group_selection_cache.clear();
self.candidate_row_page_cache.clear();
self.candidate_page_cache.clear();
self.candidate_resolved_page_cache.clear();
self.clear_provider_transport_snapshot_cache();
@@ -686,11 +705,18 @@ impl AppState {
pub(crate) fn invalidate_provider_health_routing_caches(&self) {
self.data.clear_minimal_candidate_selection_cache();
self.data.clear_provider_catalog_cache();
self.candidate_row_page_cache.clear();
self.candidate_page_cache.clear();
self.candidate_resolved_page_cache.clear();
self.clear_provider_transport_snapshot_cache();
}
pub(crate) fn invalidate_provider_runtime_state_caches(&self) {
self.data.clear_minimal_candidate_selection_cache();
self.data.clear_provider_catalog_cache();
self.clear_provider_transport_snapshot_cache();
}
pub(crate) fn invalidate_auth_context_cache(&self) {
self.auth_context_cache.clear();
self.auth_snapshot_cache.clear();
@@ -698,8 +724,13 @@ impl AppState {
self.user_feature_settings_cache.clear();
self.auth_api_key_force_capabilities_cache.clear();
self.auth_api_key_feature_settings_cache.clear();
self.auth_daily_quota_availability_cache.clear();
self.auth_wallet_snapshot_cache.clear();
self.auth_request_cost_upper_bound_cache.clear();
self.provider_quota_snapshot_cache.clear();
self.user_groups_for_user_cache.clear();
self.routing_group_selection_cache.clear();
self.candidate_row_page_cache.clear();
self.candidate_page_cache.clear();
self.candidate_resolved_page_cache.clear();
}
@@ -721,6 +752,9 @@ impl AppState {
&self.chat_pii_redaction_runtime_config_cache,
);
}
if system_config_key_affects_provider_transport_snapshot(key) {
self.clear_provider_transport_snapshot_cache();
}
}
pub(crate) async fn read_admin_system_stats(
@@ -1041,6 +1075,9 @@ impl AppState {
if let Some(queue) = self.request_candidate_queue.as_ref() {
samples.extend(queue.metric_samples());
}
samples.extend(usage_runtime_metric_samples(
&self.usage_runtime.metrics_snapshot(),
));
samples.extend(
crate::execution_runtime::transport::direct_reqwest_client_cache_metric_samples(),
);
@@ -1066,8 +1103,6 @@ impl AppState {
pub(crate) fn clear_local_execution_runtime_miss_diagnostic(&self, trace_id: &str) {
self.local_execution_runtime_miss_diagnostics
.lock()
.expect("local execution runtime miss diagnostics should lock")
.remove(trace_id);
}
@@ -1076,17 +1111,17 @@ impl AppState {
trace_id: &str,
diagnostic: LocalExecutionRuntimeMissDiagnostic,
) {
let mut diagnostics = self
if self
.local_execution_runtime_miss_diagnostics
.lock()
.expect("local execution runtime miss diagnostics should lock");
if diagnostics
.get(trace_id)
.is_some_and(|existing| should_preserve_runtime_miss_diagnostic(existing, &diagnostic))
.is_some_and(|existing| {
should_preserve_runtime_miss_diagnostic(existing.value(), &diagnostic)
})
{
return;
}
diagnostics.insert(trace_id.to_string(), diagnostic);
self.local_execution_runtime_miss_diagnostics
.insert(trace_id.to_string(), diagnostic);
}
pub(crate) fn mutate_local_execution_runtime_miss_diagnostic<F>(
@@ -1096,12 +1131,11 @@ impl AppState {
) where
F: FnOnce(&mut LocalExecutionRuntimeMissDiagnostic),
{
let mut diagnostics = self
if let Some(mut diagnostic) = self
.local_execution_runtime_miss_diagnostics
.lock()
.expect("local execution runtime miss diagnostics should lock");
if let Some(diagnostic) = diagnostics.get_mut(trace_id) {
mutate(diagnostic);
.get_mut(trace_id)
{
mutate(&mut diagnostic);
}
}
@@ -1110,10 +1144,10 @@ impl AppState {
trace_id: &str,
) -> bool {
self.local_execution_runtime_miss_diagnostics
.lock()
.expect("local execution runtime miss diagnostics should lock")
.get(trace_id)
.is_some_and(runtime_miss_diagnostic_has_candidate_signal)
.is_some_and(|diagnostic| {
runtime_miss_diagnostic_has_candidate_signal(diagnostic.value())
})
}
pub(crate) fn take_local_execution_runtime_miss_diagnostic(
@@ -1121,9 +1155,8 @@ impl AppState {
trace_id: &str,
) -> Option<LocalExecutionRuntimeMissDiagnostic> {
self.local_execution_runtime_miss_diagnostics
.lock()
.expect("local execution runtime miss diagnostics should lock")
.remove(trace_id)
.map(|(_, diagnostic)| diagnostic)
}
pub(crate) async fn try_acquire_request_permit(
@@ -1231,6 +1264,7 @@ impl AppState {
.fetch_add(1, Ordering::AcqRel)
.saturating_add(1);
self.scheduler_affinity_cache.clear();
self.candidate_row_page_cache.clear();
self.candidate_page_cache.clear();
self.candidate_resolved_page_cache.clear();
next_epoch
@@ -1495,6 +1529,97 @@ fn database_pool_metric_samples(summary: &aether_data::DatabasePoolSummary) -> V
]
}
fn usage_runtime_metric_samples(
snapshot: &usage::UsageRuntimeMetricsSnapshot,
) -> Vec<MetricSample> {
vec![
MetricSample::new(
"usage_runtime_enabled",
"Whether the gateway usage runtime is enabled.",
MetricKind::Gauge,
u64::from(snapshot.enabled),
),
MetricSample::new(
"usage_runtime_queue_terminal_events_enabled",
"Whether terminal usage events are queued before settlement.",
MetricKind::Gauge,
u64::from(snapshot.queue_terminal_events),
),
MetricSample::new(
"usage_runtime_queue_lifecycle_events_enabled",
"Whether lifecycle usage events are queued.",
MetricKind::Gauge,
u64::from(snapshot.queue_lifecycle_events),
),
MetricSample::new(
"usage_runtime_retry_deferred_lifecycle_events_enabled",
"Whether deferred lifecycle usage events are scheduled for local enqueue retry.",
MetricKind::Gauge,
u64::from(snapshot.retry_deferred_lifecycle_events),
),
MetricSample::new(
"usage_runtime_terminal_enqueue_in_flight",
"Current terminal usage enqueue operations in flight.",
MetricKind::Gauge,
snapshot.terminal_enqueue_in_flight,
),
MetricSample::new(
"usage_runtime_terminal_enqueue_deferred_total",
"Total terminal usage enqueue operations deferred by circuit or in-flight limits.",
MetricKind::Counter,
snapshot.terminal_enqueue_deferred_total,
),
MetricSample::new(
"usage_runtime_terminal_enqueue_deferred_retry_total",
"Total deferred terminal usage events scheduled for local retry.",
MetricKind::Counter,
snapshot.terminal_enqueue_deferred_retry_total,
),
MetricSample::new(
"usage_runtime_terminal_enqueue_failed_total",
"Total terminal usage enqueue failures that opened the terminal enqueue circuit.",
MetricKind::Counter,
snapshot.terminal_enqueue_failed_total,
),
MetricSample::new(
"usage_runtime_lifecycle_enqueue_in_flight",
"Current lifecycle usage enqueue operations in flight.",
MetricKind::Gauge,
snapshot.lifecycle_enqueue_in_flight,
),
MetricSample::new(
"usage_runtime_lifecycle_enqueue_deferred_total",
"Total lifecycle usage enqueue operations deferred by circuit or in-flight limits.",
MetricKind::Counter,
snapshot.lifecycle_enqueue_deferred_total,
),
MetricSample::new(
"usage_runtime_lifecycle_enqueue_deferred_dropped_total",
"Total deferred lifecycle usage events dropped instead of retrying.",
MetricKind::Counter,
snapshot.lifecycle_enqueue_deferred_dropped_total,
),
MetricSample::new(
"usage_runtime_lifecycle_enqueue_deferred_retry_total",
"Total deferred lifecycle usage events scheduled for local retry.",
MetricKind::Counter,
snapshot.lifecycle_enqueue_deferred_retry_total,
),
MetricSample::new(
"usage_runtime_lifecycle_enqueue_failed_total",
"Total lifecycle usage enqueue failures that opened the lifecycle enqueue circuit.",
MetricKind::Counter,
snapshot.lifecycle_enqueue_failed_total,
),
MetricSample::new(
"usage_runtime_enqueue_retry_scheduled_total",
"Total usage events scheduled into the local enqueue retry dispatcher.",
MetricKind::Counter,
snapshot.enqueue_retry_scheduled_total,
),
]
}
fn should_preserve_runtime_miss_diagnostic(
existing: &LocalExecutionRuntimeMissDiagnostic,
next: &LocalExecutionRuntimeMissDiagnostic,
apps/aether-gateway/src/state/mod.rs
+4 -1
View File
@@ -27,7 +27,10 @@ pub(crate) use self::admin_types::{
UserDailyQuotaAvailabilityRecord, UserPlanEntitlementRecord,
};
pub use self::app::AppState;
pub(crate) use self::app::FrontdoorRuntimeGuardConfig;
pub(crate) use self::app::{
upstream_target_gate_auto_limit, upstream_target_gate_limit_from_env,
FrontdoorRuntimeGuardConfig,
};
pub(crate) use self::cache::{
CachedProviderTransportSnapshot, AUTH_API_KEY_LAST_USED_MAX_ENTRIES,
AUTH_API_KEY_LAST_USED_TTL, PROVIDER_TRANSPORT_SNAPSHOT_CACHE_MAX_ENTRIES,
apps/aether-gateway/src/state/oauth.rs
+64 -37
View File
@@ -479,34 +479,29 @@ impl<'a> provider_transport::LocalOAuthHttpExecutor for GatewayLocalOAuthHttpExe
impl AppState {
pub(crate) fn clear_provider_transport_snapshot_cache(&self) {
self.provider_transport_snapshot_cache
.write()
.expect("provider transport snapshot cache should lock")
.clear();
self.provider_transport_snapshot_cache.clear();
self.provider_transport_snapshot_inflight.clear();
}
fn get_cached_provider_transport_snapshot_arc(
&self,
cache_key: &ProviderTransportSnapshotCacheKey,
) -> Option<Arc<provider_transport::GatewayProviderTransportSnapshot>> {
let cached = {
let cache = self
.provider_transport_snapshot_cache
.read()
.expect("provider transport snapshot cache should lock");
cache.get(cache_key).cloned()
}?;
let cached = self
.provider_transport_snapshot_cache
.get(cache_key)
.map(|entry| entry.clone())?;
if cached.loaded_at.elapsed() <= PROVIDER_TRANSPORT_SNAPSHOT_CACHE_STALE_TTL {
return Some(cached.snapshot);
}
let mut cache = self
if self
.provider_transport_snapshot_cache
.write()
.expect("provider transport snapshot cache should lock");
if cache.get(cache_key).is_some_and(|entry| {
entry.loaded_at.elapsed() > PROVIDER_TRANSPORT_SNAPSHOT_CACHE_STALE_TTL
}) {
cache.remove(cache_key);
.get(cache_key)
.is_some_and(|entry| {
entry.loaded_at.elapsed() > PROVIDER_TRANSPORT_SNAPSHOT_CACHE_STALE_TTL
})
{
self.provider_transport_snapshot_cache.remove(cache_key);
}
None
}
@@ -516,19 +511,19 @@ impl AppState {
cache_key: ProviderTransportSnapshotCacheKey,
snapshot: Arc<provider_transport::GatewayProviderTransportSnapshot>,
) {
let mut cache = self
.provider_transport_snapshot_cache
.write()
.expect("provider transport snapshot cache should lock");
if cache.len() >= PROVIDER_TRANSPORT_SNAPSHOT_CACHE_MAX_ENTRIES {
cache.retain(|_, entry| {
if self.provider_transport_snapshot_cache.len()
>= PROVIDER_TRANSPORT_SNAPSHOT_CACHE_MAX_ENTRIES
{
self.provider_transport_snapshot_cache.retain(|_, entry| {
entry.loaded_at.elapsed() <= PROVIDER_TRANSPORT_SNAPSHOT_CACHE_STALE_TTL
});
if cache.len() >= PROVIDER_TRANSPORT_SNAPSHOT_CACHE_MAX_ENTRIES {
cache.clear();
if self.provider_transport_snapshot_cache.len()
>= PROVIDER_TRANSPORT_SNAPSHOT_CACHE_MAX_ENTRIES
{
self.provider_transport_snapshot_cache.clear();
}
}
cache.insert(
self.provider_transport_snapshot_cache.insert(
cache_key,
CachedProviderTransportSnapshot {
loaded_at: std::time::Instant::now(),
@@ -828,18 +823,37 @@ impl AppState {
return Ok(Some(snapshot));
}
let snapshot = self
.read_provider_transport_snapshot_uncached(provider_id, endpoint_id, key_id)
.await?;
match snapshot {
Some(snapshot) => {
let snapshot = self.apply_global_format_conversion_override(snapshot).await;
let snapshot = Arc::new(snapshot);
self.put_cached_provider_transport_snapshot(cache_key, Arc::clone(&snapshot));
let inflight = self
.provider_transport_snapshot_inflight
.entry(cache_key.clone())
.or_insert_with(|| Arc::new(tokio::sync::Mutex::new(())))
.clone();
let result = {
let _inflight_guard = inflight.lock().await;
if let Some(snapshot) = self.get_cached_provider_transport_snapshot_arc(&cache_key) {
Ok(Some(snapshot))
} else {
match self
.read_provider_transport_snapshot_uncached(provider_id, endpoint_id, key_id)
.await
{
Ok(Some(snapshot)) => {
let snapshot = self.apply_global_format_conversion_override(snapshot).await;
let snapshot = Arc::new(snapshot);
self.put_cached_provider_transport_snapshot(
cache_key.clone(),
Arc::clone(&snapshot),
);
Ok(Some(snapshot))
}
Ok(None) => Ok(None),
Err(err) => Err(err),
}
}
None => Ok(None),
}
};
self.provider_transport_snapshot_inflight
.remove_if(&cache_key, |_, current| Arc::ptr_eq(current, &inflight));
result
}
pub(crate) async fn read_provider_transport_snapshot(
@@ -1768,6 +1782,19 @@ mod tests {
assert!(!snapshot.provider.enable_format_conversion);
}
#[tokio::test]
async fn provider_transport_snapshot_inflight_entry_is_removed_after_read() {
let state = state_with_global_format_conversion(false);
let snapshot = state
.read_provider_transport_snapshot_arc("provider-1", "endpoint-1", "key-1")
.await
.expect("snapshot read should succeed");
assert!(snapshot.is_some());
assert!(state.provider_transport_snapshot_inflight.is_empty());
}
#[test]
fn normalizes_local_openai_refresh_token_expired_response() {
let body = r#"{"error":{"message":"Could not validate your refresh token. Please try signing in again.","type":"invalid_request_error","param":null,"code":"refresh_token_expired"}}"#;
apps/aether-gateway/src/state/runtime/billing/admin.rs
+19
View File
@@ -514,4 +514,23 @@ impl AppState {
.await
.map_err(data_error)
}
pub(crate) async fn find_user_daily_quota_availability_for_auth(
&self,
user_id: &str,
) -> Result<Option<UserDailyQuotaAvailabilityRecord>, GatewayError> {
let user_id = user_id.trim();
if user_id.is_empty() {
return Ok(None);
}
let ttl = self.frontdoor_runtime_guards.auth_capacity_cache_ttl;
if ttl.is_zero() {
return self.find_user_daily_quota_availability(user_id).await;
}
self.auth_daily_quota_availability_cache
.get_or_load(user_id.to_string(), ttl, || async move {
self.find_user_daily_quota_availability(user_id).await
})
.await
}
}
apps/aether-gateway/src/state/runtime/wallet/reads.rs
+54 -3
View File
@@ -42,19 +42,70 @@ impl AppState {
api_key_id: &str,
api_key_is_standalone: bool,
) -> Result<Option<aether_data::repository::wallet::StoredWalletSnapshot>, GatewayError> {
let user_id = user_id.trim();
let api_key_id = api_key_id.trim();
let lookup = if api_key_is_standalone {
if api_key_id.trim().is_empty() {
if api_key_id.is_empty() {
None
} else {
Some((
format!("api_key:{api_key_id}"),
aether_data::repository::wallet::WalletLookupKey::ApiKeyId(api_key_id),
))
}
} else if !user_id.is_empty() {
Some((
format!("user:{user_id}"),
aether_data::repository::wallet::WalletLookupKey::UserId(user_id),
))
} else if !api_key_id.is_empty() {
Some((
format!("api_key:{api_key_id}"),
aether_data::repository::wallet::WalletLookupKey::ApiKeyId(api_key_id),
))
} else {
None
};
let Some((cache_key, lookup)) = lookup else {
return Ok(None);
};
let ttl = self.frontdoor_runtime_guards.auth_capacity_cache_ttl;
if ttl.is_zero() {
return self.find_wallet(lookup).await;
}
self.auth_wallet_snapshot_cache
.get_or_load(
cache_key,
ttl,
|| async move { self.find_wallet(lookup).await },
)
.await
}
pub(crate) async fn read_wallet_snapshot_for_auth_uncached(
&self,
user_id: &str,
api_key_id: &str,
api_key_is_standalone: bool,
) -> Result<Option<aether_data::repository::wallet::StoredWalletSnapshot>, GatewayError> {
let user_id = user_id.trim();
let api_key_id = api_key_id.trim();
let lookup = if api_key_is_standalone {
if api_key_id.is_empty() {
None
} else {
Some(aether_data::repository::wallet::WalletLookupKey::ApiKeyId(
api_key_id,
))
}
} else if !user_id.trim().is_empty() {
} else if !user_id.is_empty() {
Some(aether_data::repository::wallet::WalletLookupKey::UserId(
user_id,
))
} else if !api_key_id.trim().is_empty() {
} else if !api_key_id.is_empty() {
Some(aether_data::repository::wallet::WalletLookupKey::ApiKeyId(
api_key_id,
))
apps/aether-gateway/src/state/testing.rs
+1
View File
@@ -18,6 +18,7 @@ use crate::{provider_transport, usage};
impl AppState {
pub(crate) fn with_data_state_for_tests(mut self, data_state: GatewayDataState) -> Self {
self.replace_data_state(Arc::new(data_state));
self.request_candidate_queue = None;
self
}