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

Optimize provider scheduler database hotspots

Browse Source
This commit is contained in:
fawney19
2026-05-24 15:47:56 +08:00
parent 78523f122d
commit 576918daa5
25 changed files with 836 additions and 109 deletions
Download Patch File Download Diff File Expand all files Collapse all files
apps/aether-gateway/src/data/state/catalog.rs
+16 -1
View File
@@ -2,7 +2,8 @@ use super::{
ApiKeyLastUsedDelta, DataLayerError, GatewayDataState, GeminiFileMappingListQuery,
GeminiFileMappingStats, ProviderCatalogKeyListQuery, PublicHealthStatusCount,
PublicHealthTimelineBucket, StoredGeminiFileMapping, StoredGeminiFileMappingListPage,
StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogKeyPage,
StoredProviderCatalogEndpoint, StoredProviderCatalogKey,
StoredProviderCatalogKeyMaintenanceSummary, StoredProviderCatalogKeyPage,
StoredProviderCatalogKeyStats, StoredProviderCatalogProvider, StoredRequestCandidate,
UpsertGeminiFileMappingRecord, UpsertRequestCandidateRecord,
};
@@ -285,6 +286,20 @@ impl GatewayDataState {
}
}
pub(crate) async fn list_provider_catalog_key_maintenance_summaries_by_provider_ids(
&self,
provider_ids: &[String],
) -> Result<Vec<StoredProviderCatalogKeyMaintenanceSummary>, DataLayerError> {
match &self.provider_catalog_reader {
Some(repository) => {
repository
.list_key_maintenance_summaries_by_provider_ids(provider_ids)
.await
}
None => Ok(Vec::new()),
}
}
pub(crate) async fn list_provider_catalog_key_page(
&self,
query: &ProviderCatalogKeyListQuery,
apps/aether-gateway/src/data/state/mod.rs
+2 -1
View File
@@ -119,7 +119,8 @@ use aether_data_contracts::repository::pool_scores::{
};
use aether_data_contracts::repository::provider_catalog::{
ProviderCatalogKeyListQuery, ProviderCatalogReadRepository, ProviderCatalogWriteRepository,
StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogKeyPage,
StoredProviderCatalogEndpoint, StoredProviderCatalogKey,
StoredProviderCatalogKeyMaintenanceSummary, StoredProviderCatalogKeyPage,
StoredProviderCatalogKeyStats, StoredProviderCatalogProvider,
};
use aether_data_contracts::repository::quota::{
apps/aether-gateway/src/dispatch/pool_scheduler.rs
+29 -18
View File
@@ -53,6 +53,7 @@ const POOL_ACTIVE_PROBE_SEALED_SKIP_REASON: &str = "pool_active_probe_sealed";
const ROUTING_PROFILE_DISALLOWED_KEY_SKIP_REASON: &str = "routing_profile_disallowed_key";
const POOL_SCORE_SCHEDULE_INTEREST_CONCURRENCY: usize = 4;
const POOL_SCORE_SCHEDULE_INTEREST_MAX_PER_BATCH: usize = 16;
const POOL_SCORE_SCHEDULE_INTEREST_MIN_INTERVAL_SECS: u64 = 60;
type PoolCatalogKeyContext = PoolMemberSignals;
@@ -862,30 +863,18 @@ impl<'a> PoolKeyCursor<'a> {
return;
}
let Ok(permit) = POOL_SCORE_SCHEDULE_INTEREST_SEMAPHORE
.clone()
.try_acquire_owned()
else {
debug!(
event_name = "pool_group_score_interest_dropped",
log_type = "event",
provider_id = %self.group.candidate.provider_id,
endpoint_id = %self.group.candidate.endpoint_id,
model_id = %self.group.candidate.model_id,
score_count = scores.len(),
"gateway pool scheduler dropped score schedule interest because the background writer is saturated"
);
return;
};
let scheduled_at = current_unix_ms() / 1000;
let provider_id = self.group.candidate.provider_id.clone();
let endpoint_id = self.group.candidate.endpoint_id.clone();
let model_id = self.group.candidate.model_id.clone();
let score_count = scores.len().min(POOL_SCORE_SCHEDULE_INTEREST_MAX_PER_BATCH);
let app = self.state.app().clone();
let feedback = scores
.iter()
.filter(|score| {
score.last_scheduled_at.is_none_or(|last_scheduled_at| {
scheduled_at.saturating_sub(last_scheduled_at)
>= POOL_SCORE_SCHEDULE_INTEREST_MIN_INTERVAL_SECS
})
})
.take(POOL_SCORE_SCHEDULE_INTEREST_MAX_PER_BATCH)
.map(|score| PoolMemberScheduleFeedback {
identity: PoolMemberIdentity {
@@ -912,6 +901,28 @@ impl<'a> PoolKeyCursor<'a> {
})),
})
.collect::<Vec<_>>();
let score_count = feedback.len();
if feedback.is_empty() {
return;
}
let Ok(permit) = POOL_SCORE_SCHEDULE_INTEREST_SEMAPHORE
.clone()
.try_acquire_owned()
else {
debug!(
event_name = "pool_group_score_interest_dropped",
log_type = "event",
provider_id = %self.group.candidate.provider_id,
endpoint_id = %self.group.candidate.endpoint_id,
model_id = %self.group.candidate.model_id,
score_count = scores.len(),
"gateway pool scheduler dropped score schedule interest because the background writer is saturated"
);
return;
};
let app = self.state.app().clone();
tokio::spawn(async move {
let _permit = permit;
apps/aether-gateway/src/handlers/admin/provider/oauth/provisioning.rs
+1
View File
@@ -201,6 +201,7 @@ pub(crate) async fn update_existing_provider_oauth_catalog_key(
.map(|duration| duration.as_secs())
.unwrap_or(0);
let mut updated = existing_key.clone();
updated.is_active = true;
updated.encrypted_api_key = Some(encrypted_api_key);
updated.encrypted_auth_config = Some(encrypted_auth_config);
updated.api_formats = provider_oauth_catalog_key_api_formats(provider_type, api_formats);
apps/aether-gateway/src/maintenance/runtime/account_self_check.rs
+13 -6
View File
@@ -312,17 +312,22 @@ async fn select_keys_for_provider(
}
let result = async {
let keys = state
.list_provider_catalog_keys_by_provider_ids(std::slice::from_ref(&provider.id))
let summaries = state
.list_provider_catalog_key_maintenance_summaries_by_provider_ids(std::slice::from_ref(
&provider.id,
))
.await?
.into_iter()
.filter(|key| key.is_active)
.filter(|summary| summary.is_active)
.collect::<Vec<_>>();
if keys.is_empty() {
if summaries.is_empty() {
return Ok(Vec::new());
}
let key_ids = keys.iter().map(|key| key.id.clone()).collect::<Vec<_>>();
let key_ids = summaries
.iter()
.map(|summary| summary.id.clone())
.collect::<Vec<_>>();
let check_stamps = load_check_timestamps(runtime, &provider.id, &key_ids).await;
let selected_ids = select_account_self_check_key_ids(
&key_ids,
@@ -344,7 +349,9 @@ async fn select_keys_for_provider(
)
.await;
let mut keys_by_id = keys
let mut keys_by_id = state
.list_provider_catalog_keys_by_ids(&selected_ids)
.await?
.into_iter()
.map(|key| (key.id.clone(), key))
.collect::<BTreeMap<_, _>>();
apps/aether-gateway/src/maintenance/runtime/pool_quota_probe.rs
+16 -9
View File
@@ -849,13 +849,15 @@ async fn select_keys_for_provider(
}
let result = async {
let keys = state
.list_provider_catalog_keys_by_provider_ids(std::slice::from_ref(&provider.id))
let summaries = state
.list_provider_catalog_key_maintenance_summaries_by_provider_ids(std::slice::from_ref(
&provider.id,
))
.await?
.into_iter()
.filter(|key| key.is_active)
.filter(|summary| summary.is_active)
.collect::<Vec<_>>();
if keys.is_empty() {
if summaries.is_empty() {
return Ok(PoolQuotaProbeSelectionOutcome::Empty);
}
@@ -864,7 +866,7 @@ async fn select_keys_for_provider(
sample_provider_pool_demand(
runtime,
&provider.id,
keys.len(),
summaries.len(),
config.max_keys_per_provider,
)
.await
@@ -873,14 +875,14 @@ async fn select_keys_for_provider(
read_provider_pool_demand_snapshot(
runtime,
&provider.id,
keys.len(),
summaries.len(),
config.max_keys_per_provider,
)
.await
}
};
let target_active_count = pool_quota_probe_target_count_for_mode(
keys.len(),
summaries.len(),
pool_config,
demand_snapshot.desired_hot,
config,
@@ -894,7 +896,10 @@ async fn select_keys_for_provider(
return Ok(PoolQuotaProbeSelectionOutcome::Empty);
}
let key_ids = keys.iter().map(|key| key.id.clone()).collect::<Vec<_>>();
let key_ids = summaries
.iter()
.map(|summary| summary.id.clone())
.collect::<Vec<_>>();
let scores_by_key = load_provider_key_account_scores(state, &provider.id, &key_ids).await;
let mut active_member_ids =
load_pruned_active_probe_member_ids(runtime, &provider.id, &key_ids, &scores_by_key)
@@ -965,7 +970,9 @@ async fn select_keys_for_provider(
)
.await;
let mut keys_by_id = keys
let mut keys_by_id = state
.list_provider_catalog_keys_by_ids(&selected_ids)
.await?
.into_iter()
.map(|key| (key.id.clone(), key))
.collect::<BTreeMap<_, _>>();
apps/aether-gateway/src/maintenance/runtime/pool_score_rebuild.rs
+242 -21
View File
@@ -230,18 +230,21 @@ pub(crate) async fn perform_pool_score_rebuild_once_with_config(
.iter()
.map(|(provider, _)| provider.id.clone())
.collect::<Vec<_>>();
let mut keys_by_provider = BTreeMap::new();
for key in state
.list_provider_catalog_keys_by_provider_ids(&provider_ids)
let mut key_ids_by_provider = BTreeMap::new();
for summary in state
.list_provider_catalog_key_maintenance_summaries_by_provider_ids(&provider_ids)
.await?
{
keys_by_provider
.entry(key.provider_id.clone())
if !summary.is_active {
continue;
}
key_ids_by_provider
.entry(summary.provider_id.clone())
.or_insert_with(Vec::new)
.push(key);
.push(summary.id);
}
for keys in keys_by_provider.values_mut() {
keys.sort_by(|left, right| left.id.cmp(&right.id));
for key_ids in key_ids_by_provider.values_mut() {
key_ids.sort();
}
let now = now_unix_secs();
@@ -261,25 +264,39 @@ pub(crate) async fn perform_pool_score_rebuild_once_with_config(
}
last_provider_index = Some(provider_index);
let (provider, pool_config) = providers[provider_index].clone();
let keys = keys_by_provider.remove(&provider.id).unwrap_or_default();
let keys = keys
.into_iter()
.filter(|key| key.is_active)
.collect::<Vec<_>>();
if keys.is_empty() {
let key_ids = key_ids_by_provider.remove(&provider.id).unwrap_or_default();
if key_ids.is_empty() {
continue;
}
let total_keys = keys.len();
let total_keys = key_ids.len();
let provider_cursor_key = provider_offset_cursor_key(&provider.id);
let provider_cursor = load_runtime_usize(state, &provider_cursor_key).await % total_keys;
let remaining_budget = config
.max_upserts_per_tick
.saturating_sub(summary.scores_upserted);
let provider_budget = remaining_budget.min(total_keys);
let mut build_items = Vec::with_capacity(provider_budget);
let selected_ids = (0..provider_budget)
.map(|offset| {
let key_index = (provider_cursor + offset) % total_keys;
key_ids[key_index].clone()
})
.collect::<Vec<_>>();
let mut selected_keys_by_id = state
.list_provider_catalog_keys_by_ids(&selected_ids)
.await?
.into_iter()
.filter(|key| key.is_active && key.provider_id == provider.id)
.map(|key| (key.id.clone(), key))
.collect::<BTreeMap<_, _>>();
let selected_keys = selected_ids
.iter()
.filter_map(|key_id| selected_keys_by_id.remove(key_id))
.collect::<Vec<_>>();
let mut build_items = Vec::with_capacity(selected_keys.len());
for offset in 0..provider_budget {
let key_index = (provider_cursor + offset) % total_keys;
let key = &keys[key_index];
let Some(key) = selected_keys.get(offset) else {
break;
};
let draft = build_provider_key_pool_score_upsert(
key,
provider.provider_type.as_str(),
@@ -287,7 +304,7 @@ pub(crate) async fn perform_pool_score_rebuild_once_with_config(
now,
pool_config.score_rules,
);
build_items.push((key_index, draft.id));
build_items.push((offset, draft.id));
}
if build_items.is_empty() {
store_runtime_usize(
@@ -319,12 +336,12 @@ pub(crate) async fn perform_pool_score_rebuild_once_with_config(
.map(|score| (score.id.clone(), score))
.collect::<BTreeMap<_, _>>();
let mut provider_upserts = 0usize;
summary.keys_seen = summary.keys_seen.saturating_add(keys.len());
summary.keys_seen = summary.keys_seen.saturating_add(total_keys);
for (key_index, score_id) in &build_items {
if summary.scores_upserted >= config.max_upserts_per_tick {
break;
}
let key = &keys[*key_index];
let key = &selected_keys[*key_index];
let existing = existing_scores.get(score_id);
let upsert = build_provider_key_pool_score_upsert(
key,
@@ -429,3 +446,207 @@ pub(crate) fn spawn_pool_score_rebuild_worker(
}
}))
}
#[cfg(test)]
mod tests {
use super::*;
use std::sync::Arc;
use aether_data::repository::pool_scores::InMemoryPoolMemberScoreRepository;
use aether_data::repository::provider_catalog::InMemoryProviderCatalogReadRepository;
use aether_data_contracts::repository::pool_scores::PoolMemberIdentity;
use aether_data_contracts::repository::provider_catalog::{
ProviderCatalogKeyListQuery, ProviderCatalogReadRepository,
StoredProviderCatalogKeyMaintenanceSummary, StoredProviderCatalogKeyPage,
StoredProviderCatalogKeyStats,
};
use aether_data_contracts::DataLayerError;
use serde_json::json;
use crate::data::GatewayDataState;
struct NoWideKeyProviderCatalogReadRepository {
inner: InMemoryProviderCatalogReadRepository,
}
impl NoWideKeyProviderCatalogReadRepository {
fn seed(
providers: Vec<StoredProviderCatalogProvider>,
endpoints: Vec<StoredProviderCatalogEndpoint>,
keys: Vec<StoredProviderCatalogKey>,
) -> Self {
Self {
inner: InMemoryProviderCatalogReadRepository::seed(providers, endpoints, keys),
}
}
}
#[async_trait::async_trait]
impl ProviderCatalogReadRepository for NoWideKeyProviderCatalogReadRepository {
async fn list_providers(
&self,
active_only: bool,
) -> Result<Vec<StoredProviderCatalogProvider>, DataLayerError> {
self.inner.list_providers(active_only).await
}
async fn list_providers_by_ids(
&self,
provider_ids: &[String],
) -> Result<Vec<StoredProviderCatalogProvider>, DataLayerError> {
self.inner.list_providers_by_ids(provider_ids).await
}
async fn list_endpoints_by_ids(
&self,
endpoint_ids: &[String],
) -> Result<Vec<StoredProviderCatalogEndpoint>, DataLayerError> {
self.inner.list_endpoints_by_ids(endpoint_ids).await
}
async fn list_endpoints_by_provider_ids(
&self,
provider_ids: &[String],
) -> Result<Vec<StoredProviderCatalogEndpoint>, DataLayerError> {
self.inner
.list_endpoints_by_provider_ids(provider_ids)
.await
}
async fn list_keys_by_ids(
&self,
key_ids: &[String],
) -> Result<Vec<StoredProviderCatalogKey>, DataLayerError> {
self.inner.list_keys_by_ids(key_ids).await
}
async fn list_keys_by_provider_ids(
&self,
_provider_ids: &[String],
) -> Result<Vec<StoredProviderCatalogKey>, DataLayerError> {
panic!("pool score rebuild should not read full provider key lists");
}
async fn list_key_summaries_by_provider_ids(
&self,
provider_ids: &[String],
) -> Result<Vec<StoredProviderCatalogKey>, DataLayerError> {
self.inner
.list_key_summaries_by_provider_ids(provider_ids)
.await
}
async fn list_key_maintenance_summaries_by_provider_ids(
&self,
provider_ids: &[String],
) -> Result<Vec<StoredProviderCatalogKeyMaintenanceSummary>, DataLayerError> {
self.inner
.list_key_maintenance_summaries_by_provider_ids(provider_ids)
.await
}
async fn list_keys_page(
&self,
query: &ProviderCatalogKeyListQuery,
) -> Result<StoredProviderCatalogKeyPage, DataLayerError> {
self.inner.list_keys_page(query).await
}
async fn list_key_stats_by_provider_ids(
&self,
provider_ids: &[String],
) -> Result<Vec<StoredProviderCatalogKeyStats>, DataLayerError> {
self.inner
.list_key_stats_by_provider_ids(provider_ids)
.await
}
}
fn provider(id: &str) -> StoredProviderCatalogProvider {
let mut provider = StoredProviderCatalogProvider::new(
id.to_string(),
id.to_string(),
None,
"openai".to_string(),
)
.expect("provider should build");
provider.config = Some(json!({ "pool_advanced": {} }));
provider
}
fn key(id: &str, active: bool) -> StoredProviderCatalogKey {
StoredProviderCatalogKey::new(
id.to_string(),
"provider-1".to_string(),
id.to_string(),
"oauth".to_string(),
None,
active,
)
.expect("key should build")
}
fn score_id(provider_id: &str, key_id: &str) -> String {
let identity =
PoolMemberIdentity::provider_api_key(provider_id.to_string(), key_id.to_string());
crate::ai_serving::provider_key_pool_score_id(
&identity,
&crate::ai_serving::provider_key_pool_score_scope(),
)
}
#[tokio::test]
async fn pool_score_rebuild_uses_maintenance_summaries_before_full_key_load() {
let provider = provider("provider-1");
let provider_catalog_repository = Arc::new(NoWideKeyProviderCatalogReadRepository::seed(
vec![provider],
Vec::new(),
vec![
key("key-a", true),
key("key-b", true),
key("key-disabled", false),
],
));
let pool_score_repository = Arc::new(InMemoryPoolMemberScoreRepository::default());
let data =
GatewayDataState::with_provider_catalog_reader_for_tests(provider_catalog_repository)
.with_pool_score_repository_for_tests(Arc::clone(&pool_score_repository));
let state = AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(data);
let summary = perform_pool_score_rebuild_once_with_config(
&state,
PoolScoreRebuildWorkerConfig {
interval: Duration::from_secs(60),
max_upserts_per_tick: 1,
},
)
.await
.expect("rebuild should succeed");
assert_eq!(
summary,
PoolScoreRebuildRunSummary {
providers_checked: 1,
providers_scored: 1,
keys_seen: 2,
scores_upserted: 1,
}
);
let scores = state
.data
.get_pool_member_scores_by_ids(&GetPoolMemberScoresByIdsQuery {
ids: vec![
score_id("provider-1", "key-a"),
score_id("provider-1", "key-b"),
],
})
.await
.expect("scores should load");
assert_eq!(scores.len(), 1);
assert_eq!(scores[0].member_id, "key-a");
}
}
apps/aether-gateway/src/state/catalog.rs
+11
View File
@@ -429,6 +429,17 @@ impl AppState {
.map_err(|err| GatewayError::Internal(err.to_string()))
}
pub(crate) async fn list_provider_catalog_key_maintenance_summaries_by_provider_ids(
&self,
provider_ids: &[String],
) -> Result<Vec<provider_catalog::StoredProviderCatalogKeyMaintenanceSummary>, GatewayError>
{
self.data
.list_provider_catalog_key_maintenance_summaries_by_provider_ids(provider_ids)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))
}
pub(crate) async fn list_provider_catalog_keys_by_ids(
&self,
key_ids: &[String],
apps/aether-gateway/src/tests/control/admin/endpoints/keys.rs
+11 -2
View File
@@ -7,8 +7,8 @@ use aether_crypto::{
use aether_data::repository::provider_catalog::InMemoryProviderCatalogReadRepository;
use aether_data_contracts::repository::provider_catalog::{
ProviderCatalogKeyListQuery, ProviderCatalogReadRepository, StoredProviderCatalogEndpoint,
StoredProviderCatalogKey, StoredProviderCatalogKeyPage, StoredProviderCatalogKeyStats,
StoredProviderCatalogProvider,
StoredProviderCatalogKey, StoredProviderCatalogKeyMaintenanceSummary,
StoredProviderCatalogKeyPage, StoredProviderCatalogKeyStats, StoredProviderCatalogProvider,
};
use aether_data_contracts::DataLayerError;
use axum::body::Body;
@@ -107,6 +107,15 @@ impl ProviderCatalogReadRepository for SummaryNullingProviderCatalogReadReposito
Ok(keys)
}
async fn list_key_maintenance_summaries_by_provider_ids(
&self,
provider_ids: &[String],
) -> Result<Vec<StoredProviderCatalogKeyMaintenanceSummary>, DataLayerError> {
self.inner
.list_key_maintenance_summaries_by_provider_ids(provider_ids)
.await
}
async fn list_keys_page(
&self,
query: &ProviderCatalogKeyListQuery,