perf: reduce gateway db pressure under load

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
@@ -1,17 +1,35 @@
use std::time::Duration;
use crate::cache::{AuthApiKeyFeatureCacheKey, AuthApiKeyIdentityCacheKey};
use crate::{AppState, GatewayError};
const AUTH_API_KEY_RUNTIME_JSON_CACHE_TTL: Duration = Duration::from_secs(30);
impl AppState {
pub(crate) async fn read_auth_api_key_force_capabilities(
&self,
user_id: &str,
api_key_id: &str,
) -> Result<Option<serde_json::Value>, GatewayError> {
Ok(self
.list_auth_api_key_export_records_by_ids(&[api_key_id.to_string()])
.await?
.into_iter()
.find(|record| record.api_key_id == api_key_id && record.user_id == user_id)
.and_then(|record| record.force_capabilities))
let cache_key = AuthApiKeyIdentityCacheKey::new(user_id, api_key_id);
if cache_key.is_empty() {
return Ok(None);
}
self.auth_api_key_force_capabilities_cache
.get_or_load(
cache_key,
AUTH_API_KEY_RUNTIME_JSON_CACHE_TTL,
|| async move {
let value = self
.list_auth_api_key_export_records_by_ids(&[api_key_id.to_string()])
.await?
.into_iter()
.find(|record| record.api_key_id == api_key_id && record.user_id == user_id)
.and_then(|record| record.force_capabilities);
Ok(value)
},
)
.await
}
pub(crate) async fn read_auth_api_key_feature_settings(
@@ -20,10 +38,22 @@ impl AppState {
api_key_id: &str,
is_standalone: bool,
) -> Result<Option<serde_json::Value>, GatewayError> {
self.data
.read_auth_api_key_feature_settings(user_id, api_key_id, is_standalone)
let cache_key = AuthApiKeyFeatureCacheKey::new(user_id, api_key_id, is_standalone);
if cache_key.is_empty() {
return Ok(None);
}
self.auth_api_key_feature_settings_cache
.get_or_load(
cache_key,
AUTH_API_KEY_RUNTIME_JSON_CACHE_TTL,
|| async move {
self.data
.read_auth_api_key_feature_settings(user_id, api_key_id, is_standalone)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))
},
)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))
}
pub(crate) async fn list_auth_api_key_export_records_by_user_ids(
@@ -2,11 +2,73 @@ use std::collections::{BTreeMap, BTreeSet};
use aether_data::repository::auth::{AuthApiKeyLookupKey, ResolvedAuthApiKeySnapshotReader};
use std::time::Duration;
use crate::cache::AuthSnapshotCacheKey;
use crate::data::auth::GatewayAuthApiKeySnapshot;
use crate::{AppState, GatewayError};
const AUTH_API_KEY_SNAPSHOT_RUNTIME_CACHE_TTL: Duration = Duration::from_secs(30);
use super::super::super::{AUTH_API_KEY_LAST_USED_MAX_ENTRIES, AUTH_API_KEY_LAST_USED_TTL};
impl AppState {
pub(crate) async fn read_cached_auth_api_key_snapshot(
&self,
user_id: &str,
api_key_id: &str,
now_unix_secs: u64,
) -> Result<Option<GatewayAuthApiKeySnapshot>, GatewayError> {
let cache_key = AuthSnapshotCacheKey::user_api_key_ids(user_id, api_key_id);
if cache_key.is_empty() {
return Ok(None);
}
self.auth_snapshot_cache
.get_or_load(
cache_key,
AUTH_API_KEY_SNAPSHOT_RUNTIME_CACHE_TTL,
|| async move {
self.data
.read_auth_api_key_snapshot(user_id, api_key_id, now_unix_secs)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))
},
)
.await
}
pub(crate) async fn read_cached_auth_api_key_snapshot_by_key_hash(
&self,
key_hash: &str,
now_unix_secs: u64,
) -> Result<Option<GatewayAuthApiKeySnapshot>, GatewayError> {
let cache_key = AuthSnapshotCacheKey::key_hash(key_hash);
if cache_key.is_empty() {
return Ok(None);
}
let snapshot = self
.auth_snapshot_cache
.get_or_load(
cache_key.clone(),
AUTH_API_KEY_SNAPSHOT_RUNTIME_CACHE_TTL,
|| async move {
self.data
.read_auth_api_key_snapshot_by_key_hash(key_hash, now_unix_secs)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))
},
)
.await?;
if let Some(snapshot) = snapshot.as_ref() {
self.auth_snapshot_cache.insert(
AuthSnapshotCacheKey::user_api_key_ids(&snapshot.user_id, &snapshot.api_key_id),
Some(snapshot.clone()),
AUTH_API_KEY_SNAPSHOT_RUNTIME_CACHE_TTL,
);
}
Ok(snapshot)
}
pub(crate) async fn resolve_auth_api_key_snapshots_by_ids(
&self,
api_key_ids: &[String],
@@ -2,6 +2,9 @@ use std::collections::{BTreeMap, BTreeSet};
use crate::constants::{BUILTIN_DEFAULT_USER_GROUP_ID, DEFAULT_USER_GROUP_CONFIG_KEY};
use crate::{AppState, GatewayError};
use std::time::Duration;
const USER_GROUPS_FOR_USER_CACHE_TTL: Duration = Duration::from_secs(30);
impl AppState {
pub(crate) async fn assign_default_group_to_self_registered_user(
@@ -323,10 +326,21 @@ impl AppState {
&self,
user_id: &str,
) -> Result<Vec<aether_data::repository::users::StoredUserGroup>, GatewayError> {
self.data
.list_user_groups_for_user(user_id)
let user_id = user_id.trim();
if user_id.is_empty() {
return Ok(Vec::new());
}
let cache_key = user_id.to_string();
self.user_groups_for_user_cache
.get_or_load(cache_key, USER_GROUPS_FOR_USER_CACHE_TTL, || async move {
self.data
.list_user_groups_for_user(user_id)
.await
.map(Some)
.map_err(|err| GatewayError::Internal(err.to_string()))
})
.await
.map_err(|err| GatewayError::Internal(err.to_string()))
.map(|value| value.unwrap_or_default())
}
pub(crate) async fn list_user_group_memberships_by_user_ids(
@@ -1,10 +1,18 @@
use std::time::Duration;
use crate::{AppState, GatewayError};
const USER_RUNTIME_JSON_CACHE_TTL: Duration = Duration::from_secs(30);
impl AppState {
pub(crate) async fn read_user_model_capability_settings(
&self,
user_id: &str,
) -> Result<Option<serde_json::Value>, GatewayError> {
let user_id = user_id.trim();
if user_id.is_empty() {
return Ok(None);
}
#[cfg(test)]
if let Some(store) = self.auth_user_model_capability_store.as_ref() {
if let Some(settings) = store
@@ -17,11 +25,17 @@ impl AppState {
}
}
let users = self.list_non_admin_export_users().await?;
Ok(users
.into_iter()
.find(|user| user.id == user_id)
.and_then(|user| user.model_capability_settings))
let cache_key = user_id.to_string();
self.user_model_capability_settings_cache
.get_or_load(cache_key, USER_RUNTIME_JSON_CACHE_TTL, || async move {
Ok(self
.data
.find_export_user_by_id(user_id)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))?
.and_then(|user| user.model_capability_settings))
})
.await
}
pub(crate) async fn update_user_model_capability_settings(
@@ -56,10 +70,19 @@ impl AppState {
&self,
user_id: &str,
) -> Result<Option<serde_json::Value>, GatewayError> {
self.data
.read_user_feature_settings(user_id)
let user_id = user_id.trim();
if user_id.is_empty() {
return Ok(None);
}
let cache_key = user_id.to_string();
self.user_feature_settings_cache
.get_or_load(cache_key, USER_RUNTIME_JSON_CACHE_TTL, || async move {
self.data
.read_user_feature_settings(user_id)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))
})
.await
.map_err(|err| GatewayError::Internal(err.to_string()))
}
pub(crate) async fn update_user_feature_settings(
@@ -1,5 +1,8 @@
use crate::{AppState, GatewayError};
use aether_data_contracts::repository::{candidate_selection, candidates, quota};
use std::time::Duration;
const PROVIDER_QUOTA_RUNTIME_CACHE_TTL: Duration = Duration::from_secs(5);
impl AppState {
pub(crate) async fn list_minimal_candidate_selection_rows_for_api_format(
@@ -71,10 +74,19 @@ impl AppState {
&self,
provider_id: &str,
) -> Result<Option<quota::StoredProviderQuotaSnapshot>, GatewayError> {
self.data
.find_provider_quota_by_provider_id(provider_id)
let provider_id = provider_id.trim();
if provider_id.is_empty() {
return Ok(None);
}
let cache_key = provider_id.to_string();
self.provider_quota_snapshot_cache
.get_or_load(cache_key, PROVIDER_QUOTA_RUNTIME_CACHE_TTL, || async move {
self.data
.find_provider_quota_by_provider_id(provider_id)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))
})
.await
.map_err(|err| GatewayError::Internal(err.to_string()))
}
pub(crate) async fn read_provider_quota_snapshots(