fix: speed up admin pool loading

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

View File

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

View File

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

View File

@@ -291,7 +291,8 @@ CREATE TABLE IF NOT EXISTS provider_api_keys (
metadata TEXT, metadata TEXT,
created_at BIGINT NOT NULL, created_at BIGINT NOT NULL,
updated_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 ( CREATE TABLE IF NOT EXISTS pool_member_scores (

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;

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: - -- Name: idx_provider_api_keys_provider_created_at_desc; Type: INDEX; Schema: public; Owner: -
-- --

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);

View File

@@ -300,6 +300,7 @@ CREATE TABLE IF NOT EXISTS provider_api_keys (
updated_at INTEGER NOT NULL 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 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 ( CREATE TABLE IF NOT EXISTS pool_member_scores (
id TEXT PRIMARY KEY, id TEXT PRIMARY KEY,

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);

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: - -- Name: idx_provider_api_keys_provider_created_at_desc; Type: INDEX; Schema: public; Owner: -
-- --

View File

@@ -124,7 +124,8 @@ CREATE TABLE IF NOT EXISTS provider_api_keys (
metadata TEXT, metadata TEXT,
created_at BIGINT NOT NULL, created_at BIGINT NOT NULL,
updated_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 ( CREATE TABLE IF NOT EXISTS pool_member_scores (

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: - -- Name: idx_provider_api_keys_provider_created_at_desc; Type: INDEX; Schema: public; Owner: -
-- --

View File

@@ -133,6 +133,7 @@ CREATE TABLE IF NOT EXISTS provider_api_keys (
updated_at INTEGER NOT NULL 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 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 ( CREATE TABLE IF NOT EXISTS pool_member_scores (
id TEXT PRIMARY KEY, id TEXT PRIMARY KEY,

View File

@@ -127,7 +127,8 @@ CREATE TABLE IF NOT EXISTS provider_api_keys (
`created_at` BIGINT NOT NULL, `created_at` BIGINT NOT NULL,
`updated_at` BIGINT NOT NULL, `updated_at` BIGINT NOT NULL,
PRIMARY KEY (`id`), 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 ( CREATE TABLE IF NOT EXISTS pool_member_scores (

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); 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 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 ( CREATE TABLE IF NOT EXISTS public.pool_member_scores (
id character varying(192) NOT NULL, id character varying(192) NOT NULL,

View File

@@ -125,6 +125,7 @@ CREATE TABLE IF NOT EXISTS provider_api_keys (
updated_at INTEGER NOT NULL 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 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 ( CREATE TABLE IF NOT EXISTS pool_member_scores (
id TEXT PRIMARY KEY NOT NULL, id TEXT PRIMARY KEY NOT NULL,

View File

@@ -579,6 +579,10 @@ type = "unix_seconds"
name = "provider_api_keys_provider_id_idx" name = "provider_api_keys_provider_id_idx"
columns = ["provider_id"] 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] [table.pool_member_scores]
domain = "provider_catalog" domain = "provider_catalog"
order = 44 order = 44

View File

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

View File

@@ -309,6 +309,7 @@ fn empty_database_snapshot_covers_current_cutoff_versions() {
20260518000000, 20260518000000,
20260519000000, 20260519000000,
20260519120000, 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")); .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("ix_usage_counter_deltas_unprocessed"));
assert!(EMPTY_DATABASE_SNAPSHOT_SQL.contains("idx_entitlement_usage_entitlement_date")); 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("idx_video_tasks_due_poll"));
assert!(EMPTY_DATABASE_SNAPSHOT_SQL.contains("request_count bigint DEFAULT 0")); assert!(EMPTY_DATABASE_SNAPSHOT_SQL.contains("request_count bigint DEFAULT 0"));
assert!(EMPTY_DATABASE_SNAPSHOT_SQL.contains("usage_count bigint DEFAULT 0 NOT NULL")); 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, 20260518000000,
20260519000000, 20260519000000,
20260519120000, 20260519120000,
20260519130000,
] ]
); );
assert_eq!( assert_eq!(
@@ -622,6 +625,7 @@ fn mysql_and_sqlite_migrations_include_enabled_incrementals() {
20260518000000, 20260518000000,
20260519000000, 20260519000000,
20260519120000, 20260519120000,
20260519130000,
] ]
); );
} }
@@ -1142,6 +1146,7 @@ fn pending_migrations_from_applied_skips_versions_already_applied() {
20260518000000, 20260518000000,
20260519000000, 20260519000000,
20260519120000, 20260519120000,
20260519130000,
] ]
); );
} }