i/Aether
1
0
Fork 0
mirror of https://github.com/fawney19/Aether.git synced 2026-10-10 03:09:50 +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
crates/aether-data/src/repository/provider_catalog/memory.rs
+30 -1
View File
@@ -8,7 +8,8 @@ use serde_json::{json, Map, Value};
use super::{
ProviderCatalogKeyListOrder, ProviderCatalogKeyListQuery, ProviderCatalogReadRepository,
ProviderCatalogWriteRepository, StoredProviderCatalogEndpoint, StoredProviderCatalogKey,
StoredProviderCatalogKeyPage, StoredProviderCatalogKeyStats, StoredProviderCatalogProvider,
StoredProviderCatalogKeyMaintenanceSummary, StoredProviderCatalogKeyPage,
StoredProviderCatalogKeyStats, StoredProviderCatalogProvider,
};
use crate::repository::usage::{ProviderApiKeyUsageContribution, ProviderApiKeyUsageDelta};
use crate::DataLayerError;
@@ -431,6 +432,34 @@ impl ProviderCatalogReadRepository for InMemoryProviderCatalogReadRepository {
Self::list_keys_by_provider_ids(self, provider_ids).await
}
async fn list_key_maintenance_summaries_by_provider_ids(
&self,
provider_ids: &[String],
) -> Result<Vec<StoredProviderCatalogKeyMaintenanceSummary>, DataLayerError> {
let index = self.index.read().expect("provider catalog repository lock");
let mut keys = index
.keys
.values()
.filter(|key| {
provider_ids
.iter()
.any(|provider_id| provider_id == &key.provider_id)
})
.map(|key| StoredProviderCatalogKeyMaintenanceSummary {
id: key.id.clone(),
provider_id: key.provider_id.clone(),
is_active: key.is_active,
upstream_metadata: key.upstream_metadata.clone(),
})
.collect::<Vec<_>>();
keys.sort_by(|left, right| {
left.provider_id
.cmp(&right.provider_id)
.then(left.id.cmp(&right.id))
});
Ok(keys)
}
async fn list_keys_page(
&self,
query: &ProviderCatalogKeyListQuery,
crates/aether-data/src/repository/provider_catalog/mod.rs
+2 -1
View File
@@ -7,7 +7,8 @@ mod sqlite;
pub(crate) use aether_data_contracts::repository::provider_catalog::{
ProviderCatalogKeyListOrder, ProviderCatalogKeyListQuery, ProviderCatalogReadRepository,
ProviderCatalogWriteRepository, StoredProviderCatalogEndpoint, StoredProviderCatalogKey,
StoredProviderCatalogKeyPage, StoredProviderCatalogKeyStats, StoredProviderCatalogProvider,
StoredProviderCatalogKeyMaintenanceSummary, StoredProviderCatalogKeyPage,
StoredProviderCatalogKeyStats, StoredProviderCatalogProvider,
};
pub use memory::InMemoryProviderCatalogReadRepository;
pub use mysql::MysqlProviderCatalogReadRepository;
crates/aether-data/src/repository/provider_catalog/mysql.rs
+12 -2
View File
@@ -4,8 +4,8 @@ use sqlx::{mysql::MySqlRow, Row};
use super::{
InMemoryProviderCatalogReadRepository, ProviderCatalogKeyListQuery,
ProviderCatalogReadRepository, ProviderCatalogWriteRepository, StoredProviderCatalogEndpoint,
StoredProviderCatalogKey, StoredProviderCatalogKeyPage, StoredProviderCatalogKeyStats,
StoredProviderCatalogProvider,
StoredProviderCatalogKey, StoredProviderCatalogKeyMaintenanceSummary,
StoredProviderCatalogKeyPage, StoredProviderCatalogKeyStats, StoredProviderCatalogProvider,
};
use crate::driver::mysql::MysqlPool;
use crate::error::SqlResultExt;
@@ -928,6 +928,16 @@ impl ProviderCatalogReadRepository for MysqlProviderCatalogReadRepository {
.await
}
async fn list_key_maintenance_summaries_by_provider_ids(
&self,
provider_ids: &[String],
) -> Result<Vec<StoredProviderCatalogKeyMaintenanceSummary>, DataLayerError> {
self.load_memory()
.await?
.list_key_maintenance_summaries_by_provider_ids(provider_ids)
.await
}
async fn list_keys_page(
&self,
query: &ProviderCatalogKeyListQuery,
crates/aether-data/src/repository/provider_catalog/postgres.rs
+51 -1
View File
@@ -5,7 +5,8 @@ use sqlx::{postgres::PgRow, PgPool, Postgres, QueryBuilder, Row};
use super::{
ProviderCatalogKeyListOrder, ProviderCatalogKeyListQuery, ProviderCatalogReadRepository,
ProviderCatalogWriteRepository, StoredProviderCatalogEndpoint, StoredProviderCatalogKey,
StoredProviderCatalogKeyPage, StoredProviderCatalogKeyStats, StoredProviderCatalogProvider,
StoredProviderCatalogKeyMaintenanceSummary, StoredProviderCatalogKeyPage,
StoredProviderCatalogKeyStats, StoredProviderCatalogProvider,
};
use crate::{
error::{postgres_error, SqlxResultExt},
@@ -323,6 +324,16 @@ FROM provider_api_keys
WHERE provider_id IN (
"#;
const LIST_KEY_MAINTENANCE_SUMMARIES_BY_PROVIDER_IDS_PREFIX: &str = r#"
SELECT
id,
provider_id,
is_active,
upstream_metadata
FROM provider_api_keys
WHERE provider_id IN (
"#;
#[derive(Debug, Clone)]
pub struct SqlxProviderCatalogReadRepository {
pool: PgPool,
@@ -513,6 +524,27 @@ impl SqlxProviderCatalogReadRepository {
.await
}
pub async fn list_key_maintenance_summaries_by_provider_ids(
&self,
provider_ids: &[String],
) -> Result<Vec<StoredProviderCatalogKeyMaintenanceSummary>, DataLayerError> {
if provider_ids.is_empty() {
return Ok(Vec::new());
}
collect_query_rows(
build_list_query(
LIST_KEY_MAINTENANCE_SUMMARIES_BY_PROVIDER_IDS_PREFIX,
provider_ids,
" ORDER BY provider_id ASC, id ASC",
)
.build()
.fetch(&self.pool),
map_key_maintenance_summary_row,
)
.await
}
pub async fn list_keys_page(
&self,
query: &ProviderCatalogKeyListQuery,
@@ -1917,6 +1949,13 @@ impl ProviderCatalogReadRepository for SqlxProviderCatalogReadRepository {
Self::list_key_summaries_by_provider_ids(self, provider_ids).await
}
async fn list_key_maintenance_summaries_by_provider_ids(
&self,
provider_ids: &[String],
) -> Result<Vec<StoredProviderCatalogKeyMaintenanceSummary>, DataLayerError> {
Self::list_key_maintenance_summaries_by_provider_ids(self, provider_ids).await
}
async fn list_keys_page(
&self,
query: &ProviderCatalogKeyListQuery,
@@ -2255,6 +2294,17 @@ fn map_key_stats_row(row: &PgRow) -> Result<StoredProviderCatalogKeyStats, DataL
)
}
fn map_key_maintenance_summary_row(
row: &PgRow,
) -> Result<StoredProviderCatalogKeyMaintenanceSummary, DataLayerError> {
Ok(StoredProviderCatalogKeyMaintenanceSummary {
id: row_get(row, "id")?,
provider_id: row_get(row, "provider_id")?,
is_active: row_get(row, "is_active")?,
upstream_metadata: row_get(row, "upstream_metadata")?,
})
}
fn map_key_row(row: &PgRow) -> Result<StoredProviderCatalogKey, DataLayerError> {
let rpm_limit = row_get::<Option<i32>>(row, "rpm_limit")?
.map(|value| {
crates/aether-data/src/repository/provider_catalog/sqlite.rs
+53 -1
View File
@@ -4,7 +4,8 @@ use sqlx::{sqlite::SqliteRow, QueryBuilder, Row, Sqlite};
use super::{
ProviderCatalogKeyListOrder, ProviderCatalogKeyListQuery, ProviderCatalogReadRepository,
ProviderCatalogWriteRepository, StoredProviderCatalogEndpoint, StoredProviderCatalogKey,
StoredProviderCatalogKeyPage, StoredProviderCatalogKeyStats, StoredProviderCatalogProvider,
StoredProviderCatalogKeyMaintenanceSummary, StoredProviderCatalogKeyPage,
StoredProviderCatalogKeyStats, StoredProviderCatalogProvider,
};
use crate::driver::sqlite::{sqlite_optional_real, SqlitePool};
use crate::error::SqlResultExt;
@@ -278,6 +279,16 @@ FROM provider_api_keys
WHERE provider_id IN (
"#;
const LIST_KEY_MAINTENANCE_SUMMARIES_BY_PROVIDER_IDS_PREFIX: &str = r#"
SELECT
id,
provider_id,
is_active,
upstream_metadata
FROM provider_api_keys
WHERE provider_id IN (
"#;
#[derive(Debug, Clone)]
pub struct SqliteProviderCatalogReadRepository {
pool: SqlitePool,
@@ -423,6 +434,26 @@ impl SqliteProviderCatalogReadRepository {
rows.iter().map(map_key_row).collect()
}
pub async fn list_key_maintenance_summaries_by_provider_ids(
&self,
provider_ids: &[String],
) -> Result<Vec<StoredProviderCatalogKeyMaintenanceSummary>, DataLayerError> {
if provider_ids.is_empty() {
return Ok(Vec::new());
}
let rows = build_list_query(
LIST_KEY_MAINTENANCE_SUMMARIES_BY_PROVIDER_IDS_PREFIX,
provider_ids,
" ORDER BY provider_id ASC, id ASC",
)
.build()
.fetch_all(&self.pool)
.await
.map_sql_err()?;
rows.iter().map(map_key_maintenance_summary_row).collect()
}
pub async fn list_keys_page(
&self,
query: &ProviderCatalogKeyListQuery,
@@ -1322,6 +1353,13 @@ impl ProviderCatalogReadRepository for SqliteProviderCatalogReadRepository {
Self::list_key_summaries_by_provider_ids(self, provider_ids).await
}
async fn list_key_maintenance_summaries_by_provider_ids(
&self,
provider_ids: &[String],
) -> Result<Vec<StoredProviderCatalogKeyMaintenanceSummary>, DataLayerError> {
Self::list_key_maintenance_summaries_by_provider_ids(self, provider_ids).await
}
async fn list_keys_page(
&self,
query: &ProviderCatalogKeyListQuery,
@@ -1805,6 +1843,20 @@ fn map_key_stats_row(row: &SqliteRow) -> Result<StoredProviderCatalogKeyStats, D
)
}
fn map_key_maintenance_summary_row(
row: &SqliteRow,
) -> Result<StoredProviderCatalogKeyMaintenanceSummary, DataLayerError> {
Ok(StoredProviderCatalogKeyMaintenanceSummary {
id: row.try_get("id").map_sql_err()?,
provider_id: row.try_get("provider_id").map_sql_err()?,
is_active: row.try_get("is_active").map_sql_err()?,
upstream_metadata: optional_json_from_string(
row.try_get("upstream_metadata").map_sql_err()?,
"provider_api_keys.upstream_metadata",
)?,
})
}
fn map_key_row(row: &SqliteRow) -> Result<StoredProviderCatalogKey, DataLayerError> {
let total_cost_usd = sqlite_optional_real(row, "total_cost_usd")?.unwrap_or(0.0);
if !total_cost_usd.is_finite() {