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

perf: reduce gateway db pressure under load

Browse Source
This commit is contained in:
elky
2026-06-22 00:08:48 +08:00
parent 541cc197c4
commit f75894acbb
52 changed files with 4440 additions and 252 deletions
Download Patch File Download Diff File Expand all files Collapse all files
apps/aether-gateway/src/control/auth/resolution.rs
+195 -30
View File
@@ -1,4 +1,4 @@
use std::time::Duration;
use std::{sync::OnceLock, time::Duration};
use aether_data_contracts::repository::provider_catalog::{
StoredProviderCatalogEndpoint, StoredProviderCatalogProvider,
@@ -26,7 +26,14 @@ use super::types::{
use crate::headers::header_value_str;
const AUTH_CONTEXT_CACHE_TTL: Duration = Duration::from_secs(60);
const AUTH_CONTEXT_CACHE_MAX_ENTRIES: usize = 256;
const AUTH_CONTEXT_NEGATIVE_CACHE_TTL: Duration = Duration::from_secs(10);
const AUTH_CONTEXT_CACHE_MAX_ENTRIES: usize = 10_000;
const AUTH_CONTEXT_CACHE_MAX_ENTRIES_ENV: &str = "AETHER_GATEWAY_AUTH_CONTEXT_CACHE_MAX_ENTRIES";
const AUTH_CONTEXT_CACHE_REFRESH_ON_HIT_ENV: &str =
"AETHER_GATEWAY_AUTH_CONTEXT_CACHE_REFRESH_ON_HIT";
const AUTH_CONTEXT_NEGATIVE_CACHE_TTL_SECS_ENV: &str =
"AETHER_GATEWAY_AUTH_CONTEXT_NEGATIVE_CACHE_TTL_SECS";
const AUTH_CONTEXT_NEGATIVE_CACHE_KEY_PREFIX: &str = "negative:";
#[derive(Debug, Clone, Deserialize)]
struct AntigravityBearerBridgeConfig {
@@ -121,28 +128,49 @@ pub(in super::super) async fn resolve_control_decision_auth(
decision.admin_principal = Some(admin_principal);
}
if let Some(auth_context) = resolve_data_backed_auth_context(
state,
headers,
uri,
decision.auth_endpoint_signature.as_deref(),
)
.await?
{
log_auth_context_resolution(trace_id, &decision, &auth_context);
decision.local_auth_rejection = auth_context.local_rejection.clone();
if !auth_context.user_id.is_empty() && !auth_context.api_key_id.is_empty() {
if let Some(cache_key) = decision
.auth_endpoint_signature
.as_deref()
.and_then(|signature| build_auth_context_cache_key(headers, uri, signature))
{
put_cached_auth_context(state, cache_key, auth_context.clone());
}
decision.auth_context = Some(auth_context);
let auth_context_cache_key = decision
.auth_endpoint_signature
.as_deref()
.and_then(|signature| build_auth_context_cache_key(headers, uri, signature));
let mut resolved_auth_context = None;
if let Some(cache_key) = auth_context_cache_key.as_deref() {
if let Some(auth_context) = get_cached_auth_context(state, cache_key) {
resolved_auth_context = if auth_context_cache_refresh_on_hit() {
let refreshed = refresh_execution_runtime_auth_context(
state,
auth_context,
decision.auth_endpoint_signature.as_deref(),
)
.await?;
put_cached_auth_context(state, cache_key.to_string(), refreshed.clone());
Some(refreshed)
} else {
Some(auth_context)
};
}
}
if resolved_auth_context.is_none() {
resolved_auth_context = resolve_data_backed_auth_context(
state,
headers,
uri,
decision.auth_endpoint_signature.as_deref(),
)
.await?;
if let (Some(cache_key), Some(auth_context)) = (
auth_context_cache_key.as_ref(),
resolved_auth_context.as_ref(),
) {
put_cached_auth_context(state, cache_key.clone(), auth_context.clone());
}
}
if let Some(auth_context) = resolved_auth_context {
apply_resolved_auth_context_to_decision(trace_id, &mut decision, auth_context);
}
if decision.local_auth_rejection.is_some() {
log_local_auth_rejection(trace_id, &decision);
return Ok(ControlDecisionAuthResolution::Resolved(decision));
@@ -470,6 +498,9 @@ pub(crate) async fn resolve_execution_runtime_auth_context(
let _ = trace_id;
if let Some(auth_context) = decision.auth_context.clone() {
if !auth_context_cache_refresh_on_hit() {
return Ok(Some(auth_context));
}
return Ok(Some(
refresh_execution_runtime_auth_context(
state,
@@ -489,6 +520,10 @@ pub(crate) async fn resolve_execution_runtime_auth_context(
};
if let Some(auth_context) = get_cached_auth_context(state, &cache_key) {
if !auth_context_cache_refresh_on_hit() {
return Ok(Some(auth_context));
}
let refreshed = refresh_execution_runtime_auth_context(
state,
auth_context,
@@ -567,14 +602,88 @@ fn put_cached_auth_context(
cache_key: String,
auth_context: GatewayControlAuthContext,
) {
let (cache_key, ttl) = if is_negative_auth_context(&auth_context) {
let ttl = auth_context_negative_cache_ttl();
if ttl.is_zero() {
return;
}
(
negative_auth_context_cache_key(&cache_key),
AUTH_CONTEXT_CACHE_TTL.max(ttl),
)
} else {
(cache_key, AUTH_CONTEXT_CACHE_TTL)
};
state.auth_context_cache.insert(
cache_key,
auth_context,
AUTH_CONTEXT_CACHE_TTL,
AUTH_CONTEXT_CACHE_MAX_ENTRIES,
ttl,
auth_context_cache_max_entries(),
);
}
fn auth_context_cache_max_entries() -> usize {
static MAX_ENTRIES: OnceLock<usize> = OnceLock::new();
*MAX_ENTRIES.get_or_init(|| {
std::env::var(AUTH_CONTEXT_CACHE_MAX_ENTRIES_ENV)
.ok()
.and_then(|value| value.trim().parse::<usize>().ok())
.filter(|value| *value > 0)
.unwrap_or(AUTH_CONTEXT_CACHE_MAX_ENTRIES)
})
}
fn auth_context_cache_refresh_on_hit() -> bool {
static REFRESH_ON_HIT: OnceLock<bool> = OnceLock::new();
*REFRESH_ON_HIT.get_or_init(|| {
std::env::var(AUTH_CONTEXT_CACHE_REFRESH_ON_HIT_ENV)
.ok()
.map(|value| {
matches!(
value.trim().to_ascii_lowercase().as_str(),
"1" | "true" | "yes" | "on"
)
})
.unwrap_or(true)
})
}
fn auth_context_negative_cache_ttl() -> Duration {
static NEGATIVE_TTL: OnceLock<Duration> = OnceLock::new();
*NEGATIVE_TTL.get_or_init(|| {
std::env::var(AUTH_CONTEXT_NEGATIVE_CACHE_TTL_SECS_ENV)
.ok()
.and_then(|value| value.trim().parse::<u64>().ok())
.map(Duration::from_secs)
.unwrap_or(AUTH_CONTEXT_NEGATIVE_CACHE_TTL)
})
}
fn negative_auth_context_cache_key(cache_key: &str) -> String {
format!("{AUTH_CONTEXT_NEGATIVE_CACHE_KEY_PREFIX}{cache_key}")
}
fn is_negative_auth_context(auth_context: &GatewayControlAuthContext) -> bool {
auth_context.user_id.is_empty()
|| auth_context.api_key_id.is_empty()
|| matches!(
auth_context.local_rejection,
Some(GatewayLocalAuthRejection::InvalidApiKey)
)
}
fn apply_resolved_auth_context_to_decision(
trace_id: &str,
decision: &mut GatewayControlDecision,
auth_context: GatewayControlAuthContext,
) {
log_auth_context_resolution(trace_id, decision, &auth_context);
decision.local_auth_rejection = auth_context.local_rejection.clone();
if !auth_context.user_id.is_empty() && !auth_context.api_key_id.is_empty() {
decision.auth_context = Some(auth_context);
}
}
pub(super) async fn resolve_data_backed_auth_context(
state: &AppState,
headers: &http::HeaderMap,
@@ -600,10 +709,8 @@ pub(super) async fn resolve_data_backed_auth_context(
}
Some(GatewayPrincipalCandidate::ApiKeyHash { key_hash, .. }) => {
let snapshot = state
.data
.read_auth_api_key_snapshot_by_key_hash(&key_hash, now_unix_secs)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))?;
.read_cached_auth_api_key_snapshot_by_key_hash(&key_hash, now_unix_secs)
.await?;
let Some(snapshot) = snapshot else {
return Ok(Some(GatewayControlAuthContext {
user_id: String::new(),
@@ -1052,6 +1159,15 @@ fn endpoint_matches_requested_provider(
}
fn get_cached_auth_context(state: &AppState, cache_key: &str) -> Option<GatewayControlAuthContext> {
let negative_ttl = auth_context_negative_cache_ttl();
if !negative_ttl.is_zero() {
if let Some(auth_context) = state
.auth_context_cache
.get_fresh(&negative_auth_context_cache_key(cache_key), negative_ttl)
{
return Some(auth_context);
}
}
state
.auth_context_cache
.get_fresh(cache_key, AUTH_CONTEXT_CACHE_TTL)
@@ -1074,10 +1190,11 @@ mod tests {
use axum::http::{HeaderMap, Uri};
use super::{
resolve_data_backed_auth_context, resolve_execution_runtime_auth_context,
get_cached_auth_context, resolve_control_decision_auth, resolve_data_backed_auth_context,
resolve_execution_runtime_auth_context, ControlDecisionAuthResolution,
GatewayLocalAuthRejection,
};
use crate::control::auth::credentials::hash_api_key;
use crate::control::auth::credentials::{build_auth_context_cache_key, hash_api_key};
use crate::control::GatewayControlDecision;
use crate::data::GatewayDataState;
use crate::AppState;
@@ -1139,6 +1256,53 @@ mod tests {
.expect("endpoint should build")
}
#[tokio::test]
async fn control_auth_caches_invalid_api_key_rejections() {
let repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed([]));
let data = GatewayDataState::with_auth_api_key_repository_for_tests(repository);
let state = AppState::new()
.expect("state should build")
.with_data_state_for_tests(data);
let mut headers = HeaderMap::new();
headers.insert(
http::header::AUTHORIZATION,
"Bearer sk-missing-for-negative-cache".parse().unwrap(),
);
let request_uri = uri("/v1/chat/completions");
let decision = GatewayControlDecision::synthetic(
"/v1/chat/completions",
Some("ai_public".to_string()),
Some("openai".to_string()),
Some("chat".to_string()),
Some("openai:chat".to_string()),
);
let ControlDecisionAuthResolution::Resolved(first) = resolve_control_decision_auth(
&state,
&headers,
&request_uri,
"trace-invalid-auth-cache",
decision,
)
.await
.expect("auth resolution should succeed");
assert_eq!(
first.local_auth_rejection,
Some(GatewayLocalAuthRejection::InvalidApiKey)
);
let cache_key = build_auth_context_cache_key(&headers, &request_uri, "openai:chat")
.expect("cache key should exist");
let cached = get_cached_auth_context(&state, &cache_key)
.expect("invalid API key rejection should be cached");
assert_eq!(
cached.local_rejection,
Some(GatewayLocalAuthRejection::InvalidApiKey)
);
assert!(cached.user_id.is_empty());
assert!(cached.api_key_id.is_empty());
}
#[tokio::test]
async fn data_backed_api_key_auth_touches_last_used_once_per_throttle_window() {
let api_key = "sk-test-touch";
@@ -1326,7 +1490,8 @@ mod tests {
second.local_rejection,
Some(GatewayLocalAuthRejection::BalanceDenied {
remaining: Some(0.0),
})
}),
"cached auth context should revalidate wallet state before execution"
);
assert!(!second.access_allowed);
}