mirror of
https://github.com/fawney19/Aether.git
synced 2026-10-08 10:27:46 +08:00
perf: reduce gateway db pressure under load
This commit is contained in:
@@ -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(
|
||||
|
||||
Reference in New Issue
Block a user