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

fix: speed up admin pool loading

Browse Source
This commit is contained in:
fawney19
2026-05-19 12:05:36 +08:00
parent c6a408d4e8
commit 6a1da6a5ff
18 changed files with 149 additions and 60 deletions
Download Patch File Download Diff File Expand all files Collapse all files

40
apps/aether-gateway/src/handlers/admin/provider/pool/runtime/reads.rs
View File

@@ -13,6 +13,7 @@ use crate::provider_pool_demand::{
provider_pool_burst_pending, read_provider_pool_demand_snapshot,
};
use aether_runtime_state::{DataLayerError, RuntimeState};
use futures_util::future::join_all;
use std::collections::{BTreeMap, BTreeSet};
use std::time::{SystemTime, UNIX_EPOCH};
use tracing::warn;
@@ -32,15 +33,16 @@ pub(crate) async fn read_admin_provider_pool_cooldown_counts(
runtime: &RuntimeState,
provider_ids: &[String],
) -> BTreeMap<String, usize> {
let mut counts = BTreeMap::new();
for provider_id in provider_ids {
join_all(provider_ids.iter().map(|provider_id| async move {
let count = runtime
.set_len(&pool_cooldown_index_key(provider_id))
.await
.unwrap_or(0);
counts.insert(provider_id.clone(), count);
}
counts
(provider_id.clone(), count)
}))
.await
.into_iter()
.collect()
}
pub(crate) async fn read_admin_provider_pool_runtime_state(
@@ -170,11 +172,15 @@ pub(crate) async fn read_admin_provider_pool_runtime_state(
}
let now = current_unix_secs();
for (key_id, cost_key) in key_ids.iter().zip(cost_keys) {
let window_start = now.saturating_sub(pool_config.cost_window_seconds) as f64;
let total = runtime
.score_range_by_min(&cost_key, window_start)
.await
let cost_window_start = now.saturating_sub(pool_config.cost_window_seconds) as f64;
let cost_results = join_all(
cost_keys
.iter()
.map(|cost_key| runtime.score_range_by_min(cost_key, cost_window_start)),
)
.await;
for (key_id, members) in key_ids.iter().zip(cost_results) {
let total = members
.unwrap_or_default()
.iter()
.map(|member| parse_pool_cost_member(member))
@@ -184,11 +190,15 @@ pub(crate) async fn read_admin_provider_pool_runtime_state(
}
}
for (key_id, latency_key) in key_ids.iter().zip(latency_keys) {
let window_start = now.saturating_sub(pool_config.latency_window_seconds) as f64;
let samples = runtime
.score_range_by_min(&latency_key, window_start)
.await
let latency_window_start = now.saturating_sub(pool_config.latency_window_seconds) as f64;
let latency_results = join_all(
latency_keys
.iter()
.map(|latency_key| runtime.score_range_by_min(latency_key, latency_window_start)),
)
.await;
for (key_id, members) in key_ids.iter().zip(latency_results) {
let samples = members
.unwrap_or_default()
.iter()
.map(|member| parse_pool_latency_member(member))

102
apps/aether-gateway/src/handlers/admin/provider/pool_admin/read_routes/overview.rs
View File

@@ -19,6 +19,7 @@ use axum::{
response::{IntoResponse, Response},
Json,
};
use futures_util::future::join_all;
use serde_json::{json, Value};
pub(super) async fn build_admin_pool_overview_response(
@@ -67,47 +68,66 @@ pub(super) async fn build_admin_pool_overview_response(
.collect::<BTreeMap<_, _>>();
let probe_config = PoolQuotaProbeWorkerConfig::from_env();
let mut runtime_metrics_by_provider = BTreeMap::new();
for (provider, pool_config) in &pool_enabled_providers {
let active_keys = key_stats_by_provider
.get(&provider.id)
.map(|item| item.active_keys as usize)
.unwrap_or(0);
let hot_count = if pool_config.probing_enabled {
state
.runtime_state()
.set_len(&admin_provider_pool_quota_probe_active_members_key(
&provider.id,
))
.await
.unwrap_or(0)
} else {
0
};
let demand_snapshot = read_provider_pool_demand_snapshot(
state.runtime_state(),
&provider.id,
active_keys,
probe_config.max_keys_per_provider,
)
.await;
let burst_pending = pool_config.probing_enabled
&& provider_pool_burst_pending(state.runtime_state(), &provider.id).await;
runtime_metrics_by_provider.insert(
provider.id.clone(),
json!({
"provider_hot_count": hot_count,
"provider_desired_hot": if pool_config.probing_enabled {
demand_snapshot.desired_hot
} else {
0
},
"provider_in_flight": demand_snapshot.in_flight,
"provider_ema_in_flight": demand_snapshot.ema_in_flight,
"provider_burst_pending": burst_pending,
}),
);
}
let runtime_metrics_by_provider = join_all(pool_enabled_providers.iter().map(
|(provider, pool_config)| {
let provider_id = provider.id.clone();
let probing_enabled = pool_config.probing_enabled;
let active_keys = key_stats_by_provider
.get(&provider.id)
.map(|item| item.active_keys as usize)
.unwrap_or(0);
let max_keys_per_provider = probe_config.max_keys_per_provider;
async move {
let hot_count_future = async {
if probing_enabled {
state
.runtime_state()
.set_len(&admin_provider_pool_quota_probe_active_members_key(
&provider_id,
))
.await
.unwrap_or(0)
} else {
0
}
};
let demand_snapshot_future = read_provider_pool_demand_snapshot(
state.runtime_state(),
&provider_id,
active_keys,
max_keys_per_provider,
);
let burst_pending_future = async {
probing_enabled
&& provider_pool_burst_pending(state.runtime_state(), &provider_id).await
};
let (hot_count, demand_snapshot, burst_pending) = tokio::join!(
hot_count_future,
demand_snapshot_future,
burst_pending_future
);
(
provider_id,
json!({
"provider_hot_count": hot_count,
"provider_desired_hot": if probing_enabled {
demand_snapshot.desired_hot
} else {
0
},
"provider_in_flight": demand_snapshot.in_flight,
"provider_ema_in_flight": demand_snapshot.ema_in_flight,
"provider_burst_pending": burst_pending,
}),
)
}
},
))
.await
.into_iter()
.collect::<BTreeMap<_, _>>();
let providers = pool_enabled_providers
.into_iter()

3
crates/aether-data/migrations/mysql/20260403000000_baseline.sql
View File

@@ -291,7 +291,8 @@ CREATE TABLE IF NOT EXISTS provider_api_keys (
metadata TEXT,
created_at BIGINT NOT NULL,
updated_at BIGINT NOT NULL,
KEY provider_api_keys_provider_id_idx (provider_id)
KEY provider_api_keys_provider_id_idx (provider_id),
KEY idx_provider_api_keys_provider_default_sort (provider_id, internal_priority, name, id)
);
CREATE TABLE IF NOT EXISTS pool_member_scores (

15
crates/aether-data/migrations/mysql/20260519130000_add_provider_api_key_default_pool_sort_index.sql Normal file
View File

@@ -0,0 +1,15 @@
SET @aether_provider_key_default_sort_index_sql := IF(
(
SELECT COUNT(*)
FROM information_schema.statistics
WHERE table_schema = DATABASE()
AND table_name = 'provider_api_keys'
AND index_name = 'idx_provider_api_keys_provider_default_sort'
) = 0,
'CREATE INDEX idx_provider_api_keys_provider_default_sort ON provider_api_keys (provider_id, internal_priority, name, id)',
'DO 0'
);
PREPARE aether_provider_key_default_sort_index_stmt FROM @aether_provider_key_default_sort_index_sql;
EXECUTE aether_provider_key_default_sort_index_stmt;
DEALLOCATE PREPARE aether_provider_key_default_sort_index_stmt;

8
crates/aether-data/migrations/postgres/20260403000000_baseline.sql
View File

@@ -2704,6 +2704,14 @@ CREATE INDEX IF NOT EXISTS idx_provider_api_keys_provider_active ON public.provi
--
-- Name: idx_provider_api_keys_provider_default_sort; Type: INDEX; Schema: public; Owner: -
--
CREATE INDEX IF NOT EXISTS idx_provider_api_keys_provider_default_sort ON public.provider_api_keys USING btree (provider_id, internal_priority, name, id);
--
-- Name: idx_provider_api_keys_provider_created_at_desc; Type: INDEX; Schema: public; Owner: -
--

2
crates/aether-data/migrations/postgres/20260519130000_add_provider_api_key_default_pool_sort_index.sql Normal file
View File

@@ -0,0 +1,2 @@
CREATE INDEX IF NOT EXISTS idx_provider_api_keys_provider_default_sort
ON public.provider_api_keys USING btree (provider_id, internal_priority, name, id);

1
crates/aether-data/migrations/sqlite/20260403000000_baseline.sql
View File

@@ -300,6 +300,7 @@ CREATE TABLE IF NOT EXISTS provider_api_keys (
updated_at INTEGER NOT NULL
);
CREATE INDEX IF NOT EXISTS provider_api_keys_provider_id_idx ON provider_api_keys (provider_id);
CREATE INDEX IF NOT EXISTS idx_provider_api_keys_provider_default_sort ON provider_api_keys (provider_id, internal_priority, name, id);
CREATE TABLE IF NOT EXISTS pool_member_scores (
id TEXT PRIMARY KEY,

2
crates/aether-data/migrations/sqlite/20260519130000_add_provider_api_key_default_pool_sort_index.sql Normal file
View File

@@ -0,0 +1,2 @@
CREATE INDEX IF NOT EXISTS idx_provider_api_keys_provider_default_sort
ON provider_api_keys (provider_id, internal_priority, name, id);

8
crates/aether-data/schema/bootstrap/postgres/004_indexes.sql
View File

@@ -189,6 +189,14 @@ CREATE INDEX IF NOT EXISTS idx_provider_api_keys_provider_active ON public.provi
--
-- Name: idx_provider_api_keys_provider_default_sort; Type: INDEX; Schema: public; Owner: -
--
CREATE INDEX IF NOT EXISTS idx_provider_api_keys_provider_default_sort ON public.provider_api_keys USING btree (provider_id, internal_priority, name, id);
--
-- Name: idx_provider_api_keys_provider_created_at_desc; Type: INDEX; Schema: public; Owner: -
--

3
crates/aether-data/schema/drivers/mysql/baseline/002_provider_catalog.sql
View File

@@ -124,7 +124,8 @@ CREATE TABLE IF NOT EXISTS provider_api_keys (
metadata TEXT,
created_at BIGINT NOT NULL,
updated_at BIGINT NOT NULL,
KEY provider_api_keys_provider_id_idx (provider_id)
KEY provider_api_keys_provider_id_idx (provider_id),
KEY idx_provider_api_keys_provider_default_sort (provider_id, internal_priority, name, id)
);
CREATE TABLE IF NOT EXISTS pool_member_scores (

8
crates/aether-data/schema/drivers/postgres/baseline/004_indexes.sql
View File

@@ -125,6 +125,14 @@ CREATE INDEX IF NOT EXISTS idx_provider_api_keys_provider_active ON public.provi
--
-- Name: idx_provider_api_keys_provider_default_sort; Type: INDEX; Schema: public; Owner: -
--
CREATE INDEX IF NOT EXISTS idx_provider_api_keys_provider_default_sort ON public.provider_api_keys USING btree (provider_id, internal_priority, name, id);
--
-- Name: idx_provider_api_keys_provider_created_at_desc; Type: INDEX; Schema: public; Owner: -
--

1
crates/aether-data/schema/drivers/sqlite/baseline/002_provider_catalog.sql
View File

@@ -133,6 +133,7 @@ CREATE TABLE IF NOT EXISTS provider_api_keys (
updated_at INTEGER NOT NULL
);
CREATE INDEX IF NOT EXISTS provider_api_keys_provider_id_idx ON provider_api_keys (provider_id);
CREATE INDEX IF NOT EXISTS idx_provider_api_keys_provider_default_sort ON provider_api_keys (provider_id, internal_priority, name, id);
CREATE TABLE IF NOT EXISTS pool_member_scores (
id TEXT PRIMARY KEY,

3
crates/aether-data/schema/generated/mysql/baseline/002_provider_catalog.sql
View File

@@ -127,7 +127,8 @@ CREATE TABLE IF NOT EXISTS provider_api_keys (
`created_at` BIGINT NOT NULL,
`updated_at` BIGINT NOT NULL,
PRIMARY KEY (`id`),
KEY provider_api_keys_provider_id_idx (`provider_id`)
KEY provider_api_keys_provider_id_idx (`provider_id`),
KEY idx_provider_api_keys_provider_default_sort (`provider_id`, `internal_priority`, `name`, `id`)
);
CREATE TABLE IF NOT EXISTS pool_member_scores (

1
crates/aether-data/schema/generated/postgres/baseline/002_provider_catalog.sql
View File

@@ -133,6 +133,7 @@ CREATE TABLE IF NOT EXISTS public.provider_api_keys (
ALTER TABLE ONLY public.provider_api_keys ADD CONSTRAINT provider_api_keys_pkey PRIMARY KEY (id);
CREATE INDEX IF NOT EXISTS provider_api_keys_provider_id_idx ON public.provider_api_keys USING btree (provider_id);
CREATE INDEX IF NOT EXISTS idx_provider_api_keys_provider_default_sort ON public.provider_api_keys USING btree (provider_id, internal_priority, name, id);
CREATE TABLE IF NOT EXISTS public.pool_member_scores (
id character varying(192) NOT NULL,

1
crates/aether-data/schema/generated/sqlite/baseline/002_provider_catalog.sql
View File

@@ -125,6 +125,7 @@ CREATE TABLE IF NOT EXISTS provider_api_keys (
updated_at INTEGER NOT NULL
);
CREATE INDEX IF NOT EXISTS provider_api_keys_provider_id_idx ON provider_api_keys (provider_id);
CREATE INDEX IF NOT EXISTS idx_provider_api_keys_provider_default_sort ON provider_api_keys (provider_id, internal_priority, name, id);
CREATE TABLE IF NOT EXISTS pool_member_scores (
id TEXT PRIMARY KEY NOT NULL,

4
crates/aether-data/schema/logical/002_provider_catalog.toml
View File

@@ -579,6 +579,10 @@ type = "unix_seconds"
name = "provider_api_keys_provider_id_idx"
columns = ["provider_id"]
[[table.provider_api_keys.indexes]]
name = "idx_provider_api_keys_provider_default_sort"
columns = ["provider_id", "internal_priority", "name", "id"]
[table.pool_member_scores]
domain = "provider_catalog"
order = 44

2
crates/aether-data/src/lifecycle/bootstrap/postgres.rs
View File

@@ -7,7 +7,7 @@ use tracing::info;
// Generated by build.rs from schema/bootstrap/postgres.
pub(crate) static EMPTY_DATABASE_SNAPSHOT_SQL: &str =
include_str!(concat!(env!("OUT_DIR"), "/empty_database_snapshot.sql"));
pub(crate) const EMPTY_DATABASE_SNAPSHOT_CUTOFF_VERSION: i64 = 20260519120000;
pub(crate) const EMPTY_DATABASE_SNAPSHOT_CUTOFF_VERSION: i64 = 20260519130000;
const PUBLIC_BASE_TABLE_COUNT_SQL: &str = r#"
SELECT COUNT(*)::BIGINT

5
crates/aether-data/src/lifecycle/migrate/tests.rs
View File

@@ -309,6 +309,7 @@ fn empty_database_snapshot_covers_current_cutoff_versions() {
20260518000000,
20260519000000,
20260519120000,
20260519130000,
]
);
}
@@ -384,6 +385,7 @@ fn empty_database_snapshot_sql_includes_usage_body_blobs_and_audit_admin_role()
.contains("CREATE TABLE IF NOT EXISTS public.usage_counter_deltas"));
assert!(EMPTY_DATABASE_SNAPSHOT_SQL.contains("ix_usage_counter_deltas_unprocessed"));
assert!(EMPTY_DATABASE_SNAPSHOT_SQL.contains("idx_entitlement_usage_entitlement_date"));
assert!(EMPTY_DATABASE_SNAPSHOT_SQL.contains("idx_provider_api_keys_provider_default_sort"));
assert!(EMPTY_DATABASE_SNAPSHOT_SQL.contains("idx_video_tasks_due_poll"));
assert!(EMPTY_DATABASE_SNAPSHOT_SQL.contains("request_count bigint DEFAULT 0"));
assert!(EMPTY_DATABASE_SNAPSHOT_SQL.contains("usage_count bigint DEFAULT 0 NOT NULL"));
@@ -602,6 +604,7 @@ fn mysql_and_sqlite_migrations_include_enabled_incrementals() {
20260518000000,
20260519000000,
20260519120000,
20260519130000,
]
);
assert_eq!(
@@ -622,6 +625,7 @@ fn mysql_and_sqlite_migrations_include_enabled_incrementals() {
20260518000000,
20260519000000,
20260519120000,
20260519130000,
]
);
}
@@ -1142,6 +1146,7 @@ fn pending_migrations_from_applied_skips_versions_already_applied() {
20260518000000,
20260519000000,
20260519120000,
20260519130000,
]
);
}