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/state/app.rs
+13 -2
View File
@@ -4,13 +4,16 @@ use std::sync::Arc;
use std::sync::Mutex as StdMutex;
use std::time::Duration;
use aether_data::repository::users::StoredUserGroup;
use aether_data_contracts::repository::quota::StoredProviderQuotaSnapshot;
use aether_runtime::ConcurrencyGate;
use aether_runtime_state::{RuntimeSemaphore, RuntimeState};
use super::super::async_task::{VideoTaskPollerConfig, VideoTaskService};
use super::super::cache::{
AuthApiKeyLastUsedCache, AuthContextCache, DashboardResponseCache, DirectPlanBypassCache,
SchedulerAffinityCache, SystemConfigCache,
AuthApiKeyFeatureCacheKey, AuthApiKeyIdentityCacheKey, AuthApiKeyLastUsedCache,
AuthContextCache, AuthSnapshotCache, DashboardResponseCache, DirectPlanBypassCache,
JsonValueCache, SchedulerAffinityCache, SystemConfigCache, ValueCache,
};
use super::super::data::GatewayDataState;
use super::super::fallback_metrics;
@@ -116,6 +119,14 @@ pub struct AppState {
pub(crate) distributed_request_gate: Option<Arc<RuntimeSemaphore>>,
pub(crate) client: reqwest::Client,
pub(crate) auth_context_cache: Arc<AuthContextCache>,
pub(crate) auth_snapshot_cache: Arc<AuthSnapshotCache>,
pub(crate) user_model_capability_settings_cache: Arc<JsonValueCache<String>>,
pub(crate) user_feature_settings_cache: Arc<JsonValueCache<String>>,
pub(crate) auth_api_key_force_capabilities_cache:
Arc<JsonValueCache<AuthApiKeyIdentityCacheKey>>,
pub(crate) auth_api_key_feature_settings_cache: Arc<JsonValueCache<AuthApiKeyFeatureCacheKey>>,
pub(crate) provider_quota_snapshot_cache: Arc<ValueCache<String, StoredProviderQuotaSnapshot>>,
pub(crate) user_groups_for_user_cache: Arc<ValueCache<String, Vec<StoredUserGroup>>>,
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>,
apps/aether-gateway/src/state/core.rs
+89 -4
View File
@@ -29,9 +29,9 @@ use super::super::async_task::{
spawn_video_task_poller, VideoTaskPollerConfig, VideoTaskService, VideoTaskTruthSourceMode,
};
use super::super::cache::{
AuthApiKeyLastUsedCache, AuthContextCache, DashboardResponseCache, DirectPlanBypassCache,
SchedulerAffinityCache, SchedulerAffinitySnapshotEntry, SchedulerAffinityTarget,
SystemConfigCache,
AuthApiKeyLastUsedCache, AuthContextCache, AuthSnapshotCache, DashboardResponseCache,
DirectPlanBypassCache, JsonValueCache, SchedulerAffinityCache, SchedulerAffinitySnapshotEntry,
SchedulerAffinityTarget, SystemConfigCache, ValueCache,
};
use super::super::data::{GatewayDataConfig, GatewayDataState};
use super::super::fallback_metrics;
@@ -62,7 +62,7 @@ use crate::maintenance::spawn_usage_cleanup_worker;
use crate::maintenance::spawn_usage_counter_flush_worker;
use crate::maintenance::spawn_wallet_daily_usage_aggregation_worker;
const SYSTEM_CONFIG_CACHE_TTL: Duration = Duration::from_secs(3);
const SYSTEM_CONFIG_CACHE_TTL: Duration = Duration::from_secs(30);
const SCHEDULER_AFFECTING_SYSTEM_CONFIG_KEYS: &[&str] = &[
"enable_format_conversion",
"keep_priority_on_conversion",
@@ -236,6 +236,13 @@ impl AppState {
distributed_request_gate: None,
client,
auth_context_cache: Arc::new(AuthContextCache::default()),
auth_snapshot_cache: Arc::new(AuthSnapshotCache::default()),
user_model_capability_settings_cache: Arc::new(JsonValueCache::default()),
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()),
provider_quota_snapshot_cache: Arc::new(ValueCache::default()),
user_groups_for_user_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()),
@@ -502,6 +509,11 @@ impl AppState {
return Ok(value);
}
let _guard = self.system_config_cache.load_guard().await;
if let Some(value) = self.system_config_cache.get(key, SYSTEM_CONFIG_CACHE_TTL) {
return Ok(value);
}
let value = self
.data
.find_system_config_value(key)
@@ -586,6 +598,13 @@ impl AppState {
pub(crate) fn invalidate_auth_context_cache(&self) {
self.auth_context_cache.clear();
self.auth_snapshot_cache.clear();
self.user_model_capability_settings_cache.clear();
self.user_feature_settings_cache.clear();
self.auth_api_key_force_capabilities_cache.clear();
self.auth_api_key_feature_settings_cache.clear();
self.provider_quota_snapshot_cache.clear();
self.user_groups_for_user_cache.clear();
}
fn remember_system_config_write(&self, key: &str, value: Option<serde_json::Value>) {
@@ -896,6 +915,9 @@ impl AppState {
),
}
}
if let Some(summary) = self.data.database_pool_summary() {
samples.extend(database_pool_metric_samples(&summary));
}
samples.extend(self.tunnel.metric_samples());
samples.extend(self.fallback_metrics.metric_samples());
samples
@@ -1279,6 +1301,69 @@ impl AppState {
}
}
fn database_pool_metric_samples(summary: &aether_data::DatabasePoolSummary) -> Vec<MetricSample> {
let labels = vec![MetricLabel::new("driver", summary.driver.to_string())];
let usage_basis_points = if summary.usage_rate.is_finite() && summary.usage_rate > 0.0 {
(summary.usage_rate * 100.0).round() as u64
} else {
0
};
let under_maintenance_pressure =
GatewayDataState::database_pool_summary_under_maintenance_pressure(summary);
vec![
MetricSample::new(
"database_pool_checked_out_connections",
"Number of database connections currently checked out from the gateway pool.",
MetricKind::Gauge,
summary.checked_out as u64,
)
.with_labels(labels.clone()),
MetricSample::new(
"database_pool_idle_connections",
"Number of idle database connections currently available in the gateway pool.",
MetricKind::Gauge,
summary.idle as u64,
)
.with_labels(labels.clone()),
MetricSample::new(
"database_pool_size_connections",
"Current number of database connections opened by the gateway pool.",
MetricKind::Gauge,
summary.pool_size as u64,
)
.with_labels(labels.clone()),
MetricSample::new(
"database_pool_max_connections",
"Configured maximum number of database connections for the gateway pool.",
MetricKind::Gauge,
summary.max_connections as u64,
)
.with_labels(labels.clone()),
MetricSample::new(
"database_pool_usage_basis_points",
"Database pool usage rate in basis points, where 10000 means 100 percent.",
MetricKind::Gauge,
usage_basis_points,
)
.with_labels(labels.clone()),
MetricSample::new(
"database_pool_idle_reserve_connections",
"Idle database connections reserved for foreground traffic before maintenance defers.",
MetricKind::Gauge,
GatewayDataState::maintenance_pool_idle_reserve(summary) as u64,
)
.with_labels(labels.clone()),
MetricSample::new(
"database_pool_under_maintenance_pressure",
"Whether maintenance workers should currently defer for foreground database pool capacity.",
MetricKind::Gauge,
u64::from(under_maintenance_pressure),
)
.with_labels(labels),
]
}
fn should_preserve_runtime_miss_diagnostic(
existing: &LocalExecutionRuntimeMissDiagnostic,
next: &LocalExecutionRuntimeMissDiagnostic,
apps/aether-gateway/src/state/runtime/api_key_exports.rs
+39 -9
View File
@@ -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(
apps/aether-gateway/src/state/runtime/auth/api_keys.rs
+62
View File
@@ -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],
apps/aether-gateway/src/state/runtime/auth/user_lifecycle.rs
+17 -3
View File
@@ -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(
apps/aether-gateway/src/state/runtime/auth/user_provisioning.rs
+31 -8
View File
@@ -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(
apps/aether-gateway/src/state/runtime/candidate_queries.rs
+15 -3
View File
@@ -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(