From 01792a5dbe10b303cf20ae1ec75f5b94b7c6bc2e Mon Sep 17 00:00:00 2001 From: zyycn Date: Sun, 26 Jul 2026 20:28:52 +0800 Subject: [PATCH] perf(store): index retention deletes, parallelize dashboard, expose pool budget Add partial indexes for the three retention delete predicates and delete in bounded ctid batches with per-batch commits, replacing the hourly full-table scan and unbounded single-transaction delete. Run the six dashboard-summary queries concurrently so page latency tracks the slowest query instead of their sum. Make the PostgreSQL pool budget configurable via an optional store.pool section (max_connections, acquire_timeout_seconds) with a 5-second acquire timeout default so pool exhaustion fails fast instead of queueing for 30 seconds. Co-authored-by: Claude --- backend/crates/gateway-store/src/lib.rs | 34 ++++++- .../crates/gateway-store/src/postgres/mod.rs | 15 ++- .../src/postgres/observability.rs | 43 ++++---- .../gateway-store/src/postgres/retention.rs | 98 ++++++++++++------- .../gateway-store/tests/postgres/execution.rs | 20 +++- .../gateway-store/tests/postgres/mod.rs | 12 ++- .../0002_retention_delete_indexes.sql | 13 +++ deploy/config.example.yaml | 4 + 8 files changed, 172 insertions(+), 67 deletions(-) create mode 100644 backend/migrations/0002_retention_delete_indexes.sql diff --git a/backend/crates/gateway-store/src/lib.rs b/backend/crates/gateway-store/src/lib.rs index 06bd5e2ea..35b0a219d 100644 --- a/backend/crates/gateway-store/src/lib.rs +++ b/backend/crates/gateway-store/src/lib.rs @@ -86,6 +86,37 @@ pub type StoreResult = Result; pub struct StoreConfig { database: StoreConnectionConfig, redis: StoreConnectionConfig, + #[serde(default)] + pool: StorePoolConfig, +} + +/// PostgreSQL 连接池预算；acquire 超时决定池耗尽时快速失败而非排队积压。 +#[derive(Debug, Clone, Copy, PartialEq, Eq, Deserialize)] +#[serde(deny_unknown_fields, default)] +pub struct StorePoolConfig { + pub max_connections: u32, + pub acquire_timeout_seconds: u64, +} + +impl Default for StorePoolConfig { + fn default() -> Self { + Self { + max_connections: 10, + acquire_timeout_seconds: 5, + } + } +} + +impl StorePoolConfig { + fn validate(&self) -> StoreResult<()> { + if self.max_connections == 0 || self.acquire_timeout_seconds == 0 { + return Err(StoreError::InvalidData { + entity: "store config", + message: "pool limits must be positive".to_owned(), + }); + } + Ok(()) + } } impl StoreConfig { @@ -104,6 +135,7 @@ impl StoreConfig { } self.database.validate("database")?; self.redis.validate("redis")?; + self.pool.validate()?; Ok(()) } @@ -227,7 +259,7 @@ pub async fn initialize(mut config: StoreConfig) -> StoreResult { const REDIS_NAMESPACE: &str = "codex-proxy-rs"; config.resolve_and_validate(std::path::Path::new("."))?; - let pool = postgres::connect_and_migrate(&config.database_url()?).await?; + let pool = postgres::connect_and_migrate(&config.database_url()?, config.pool).await?; let redis_client = ::redis::Client::open(config.redis_url()?) .map_err(|_| redis_unavailable("create Redis client"))?; let redis_connection = redis_client diff --git a/backend/crates/gateway-store/src/postgres/mod.rs b/backend/crates/gateway-store/src/postgres/mod.rs index f31fc20de..905d47fb7 100644 --- a/backend/crates/gateway-store/src/postgres/mod.rs +++ b/backend/crates/gateway-store/src/postgres/mod.rs @@ -3,7 +3,9 @@ use async_trait::async_trait; use sqlx::{PgPool, postgres::PgPoolOptions}; -use crate::{Revision, StoreBackend, StoreError, StoreResult, postgres_unavailable}; +use crate::{ + Revision, StoreBackend, StoreError, StorePoolConfig, StoreResult, postgres_unavailable, +}; mod admin_security_audit; mod admission_recovery; @@ -29,12 +31,19 @@ pub use snapshot::*; static MIGRATOR: sqlx::migrate::Migrator = sqlx::migrate!("../../migrations"); -/// 建立 PostgreSQL pool 并只执行冻结的 `0001_initial.sql` migration 集。 -pub async fn connect_and_migrate(database_url: &str) -> StoreResult { +/// 建立 PostgreSQL pool 并只执行冻结的 migration 集。 +pub async fn connect_and_migrate( + database_url: &str, + pool_config: StorePoolConfig, +) -> StoreResult { if database_url.trim().is_empty() { return Err(postgres_unavailable("connect PostgreSQL")); } let pool = PgPoolOptions::new() + .max_connections(pool_config.max_connections) + .acquire_timeout(std::time::Duration::from_secs( + pool_config.acquire_timeout_seconds, + )) .connect(database_url) .await .map_err(|_| postgres_unavailable("connect PostgreSQL"))?; diff --git a/backend/crates/gateway-store/src/postgres/observability.rs b/backend/crates/gateway-store/src/postgres/observability.rs index 89d39f07e..2bc5479f5 100644 --- a/backend/crates/gateway-store/src/postgres/observability.rs +++ b/backend/crates/gateway-store/src/postgres/observability.rs @@ -759,35 +759,32 @@ impl ObservabilityRepository for PgObservabilityRepository { range: ObservabilityRange, ) -> StoreResult { let filter = UsageRecordFilter::default(); - let requests = request_metrics(&self.pool, range, &filter).await?; - let attempts = attempt_metrics(&self.pool, range, &filter).await?; - let provider_accounts = provider_account_metrics(&self.pool, range.end).await?; - let trend = request_metric_series(&self.pool, range, &filter).await?; let account_usage_range = ObservabilityRange::new( range.end - TimeDelta::hours(ACCOUNT_USAGE_TIMELINE_HOURS), range.end, )?; - let account_usage = provider_account_usage( - &self.pool, + let account_usage_query = ProviderAccountUsageQuery::recent(account_usage_range, DASHBOARD_ACCOUNT_LIMIT)? - .with_hourly_request_buckets()?, - ) - .await?; - let recent_requests = list_usage_records( - &self.pool, - UsageRecordQuery { - range, - filter: UsageRecordFilter { - outcome: Some("succeeded".to_owned()), - ..UsageRecordFilter::default() - }, - cursor: None, - page: ObservabilityPageNumber::new(1)?, - page_size: ObservabilityPageSize::new(10)?, + .with_hourly_request_buckets()?; + let recent_query = UsageRecordQuery { + range, + filter: UsageRecordFilter { + outcome: Some("succeeded".to_owned()), + ..UsageRecordFilter::default() }, - ) - .await? - .items; + cursor: None, + page: ObservabilityPageNumber::new(1)?, + page_size: ObservabilityPageSize::new(10)?, + }; + // 六路独立查询并发执行；页面延迟由最慢一路决定而非各路之和。 + let (requests, attempts, provider_accounts, trend, account_usage, recent_requests) = futures::try_join!( + request_metrics(&self.pool, range, &filter), + attempt_metrics(&self.pool, range, &filter), + provider_account_metrics(&self.pool, range.end), + request_metric_series(&self.pool, range, &filter), + provider_account_usage(&self.pool, account_usage_query), + async { Ok(list_usage_records(&self.pool, recent_query).await?.items) }, + )?; Ok(DashboardObservation { range, requests, diff --git a/backend/crates/gateway-store/src/postgres/retention.rs b/backend/crates/gateway-store/src/postgres/retention.rs index d2c3f2b21..71e3dac25 100644 --- a/backend/crates/gateway-store/src/postgres/retention.rs +++ b/backend/crates/gateway-store/src/postgres/retention.rs @@ -78,47 +78,49 @@ impl RetentionRepository for PgRetentionRepository { }); } - let mut transaction = self - .pool - .begin() - .await - .map_err(|_| postgres_unavailable("begin retention transaction"))?; - let model_requests = sqlx::query( + // 分批删除，每批独立提交：清理是幂等的，无需跨表原子性； + // 有界批量避免百万行单事务的长锁与 WAL 峰值。 + let model_requests = purge_in_batches( + &self.pool, "delete from model_requests - where outcome <> 'running' - and completed_at < $1 - ($2 * interval '1 day')", + where ctid in ( + select ctid from model_requests + where outcome <> 'running' + and completed_at < $1 - ($2 * interval '1 day') + limit $3 + )", + now, + settings.usage_retention_days, + "delete expired model requests", ) - .bind(now) - .bind(i64::from(settings.usage_retention_days)) - .execute(&mut *transaction) - .await - .map_err(|_| postgres_unavailable("delete expired model requests"))? - .rows_affected(); - let ops_events = sqlx::query( + .await?; + let ops_events = purge_in_batches( + &self.pool, "delete from ops_events - where model_request_id is null - and created_at < $1 - ($2 * interval '1 day')", + where ctid in ( + select ctid from ops_events + where model_request_id is null + and created_at < $1 - ($2 * interval '1 day') + limit $3 + )", + now, + settings.ops_event_retention_days, + "delete expired ops events", ) - .bind(now) - .bind(i64::from(settings.ops_event_retention_days)) - .execute(&mut *transaction) - .await - .map_err(|_| postgres_unavailable("delete expired ops events"))? - .rows_affected(); - let admin_audit_events = sqlx::query( + .await?; + let admin_audit_events = purge_in_batches( + &self.pool, "delete from admin_audit_events - where created_at < $1 - ($2 * interval '1 day')", + where ctid in ( + select ctid from admin_audit_events + where created_at < $1 - ($2 * interval '1 day') + limit $3 + )", + now, + settings.audit_retention_days, + "delete expired admin audit events", ) - .bind(now) - .bind(i64::from(settings.audit_retention_days)) - .execute(&mut *transaction) - .await - .map_err(|_| postgres_unavailable("delete expired admin audit events"))? - .rows_affected(); - transaction - .commit() - .await - .map_err(|_| postgres_unavailable("commit retention transaction"))?; + .await?; Ok(RetentionReport { model_requests, ops_events, @@ -127,6 +129,32 @@ impl RetentionRepository for PgRetentionRepository { } } +const RETENTION_DELETE_BATCH_ROWS: i64 = 10_000; + +async fn purge_in_batches( + pool: &PgPool, + delete_sql: &'static str, + now: DateTime, + retention_days: u32, + label: &'static str, +) -> StoreResult { + let mut total = 0u64; + loop { + let deleted = sqlx::query(delete_sql) + .bind(now) + .bind(i64::from(retention_days)) + .bind(RETENTION_DELETE_BATCH_ROWS) + .execute(pool) + .await + .map_err(|_| postgres_unavailable(label))? + .rows_affected(); + total += deleted; + if deleted < RETENTION_DELETE_BATCH_ROWS.unsigned_abs() { + return Ok(total); + } + } +} + fn to_u32(value: i64) -> StoreResult { u32::try_from(value).map_err(|_| StoreError::InvalidData { entity: "retention settings", diff --git a/backend/crates/gateway-store/tests/postgres/execution.rs b/backend/crates/gateway-store/tests/postgres/execution.rs index cad4e52f3..b8aff3472 100644 --- a/backend/crates/gateway-store/tests/postgres/execution.rs +++ b/backend/crates/gateway-store/tests/postgres/execution.rs @@ -50,6 +50,22 @@ async fn merged_first_attempt_insert_should_match_sequential_semantics() { }; let repository = PgExecutionStore::new(database.pool.clone()); let started_at = Utc::now(); + sqlx::query( + "insert into provider_accounts ( + id, provider_kind, name, email, upstream_user_id, + upstream_account_id, plan_type, authentication_kind, + provider_credentials_json, credential_revision, + has_refresh_token, access_token_expires_at, enabled, availability, + availability_observed_at, created_at, updated_at + ) values ( + 'acct_merged', 'openai', 'merged', null, 'user-merged', null, null, 'oauth', + '{}'::jsonb, 1, false, $1 + interval '1 day', true, 'ready', $1, $1, $1 + )", + ) + .bind(started_at) + .execute(&database.pool) + .await + .expect("seed provider account"); let request = NewModelRequest { id: "req_merged".to_owned(), client_api_key_id: None, @@ -111,8 +127,8 @@ async fn merged_first_attempt_insert_should_match_sequential_semantics() { model_request_id: "req_merged".to_owned(), attempt_count: 2, provider_kind: "openai".to_owned(), - provider_account_id: Some("acct_next".to_owned()), - provider_account_ref: Some("acct_next".to_owned()), + provider_account_id: Some("acct_merged".to_owned()), + provider_account_ref: Some("acct_merged".to_owned()), upstream_model_id: "gpt-test".to_owned(), upstream_transport: "http_sse".to_owned(), http_version: None, diff --git a/backend/crates/gateway-store/tests/postgres/mod.rs b/backend/crates/gateway-store/tests/postgres/mod.rs index 1b87a803d..3d8f59ff4 100644 --- a/backend/crates/gateway-store/tests/postgres/mod.rs +++ b/backend/crates/gateway-store/tests/postgres/mod.rs @@ -58,6 +58,12 @@ impl TestDatabase { .execute(&mut *migration) .await .expect("apply terminal migration"); + sqlx::raw_sql(include_str!( + "../../../../migrations/0002_retention_delete_indexes.sql" + )) + .execute(&mut *migration) + .await + .expect("apply retention index migration"); migration.commit().await.expect("commit terminal migration"); Some(Self { admin, @@ -102,7 +108,7 @@ async fn connect_and_migrate_should_apply_0001_once_and_reopen_cleanly() { .database(&database) .to_url_lossy() .to_string(); - let first = connect_and_migrate(&isolated_url) + let first = connect_and_migrate(&isolated_url, gateway_store::StorePoolConfig::default()) .await .expect("apply 0001 through production migrator"); let first_table_count = sqlx::query_scalar::<_, i64>( @@ -113,7 +119,7 @@ async fn connect_and_migrate_should_apply_0001_once_and_reopen_cleanly() { .expect("count migrated tables"); first.close().await; - let second = connect_and_migrate(&isolated_url) + let second = connect_and_migrate(&isolated_url, gateway_store::StorePoolConfig::default()) .await .expect("reopen database through production migrator"); let migration_count = @@ -131,7 +137,7 @@ async fn connect_and_migrate_should_apply_0001_once_and_reopen_cleanly() { .expect("drop migration test database"); admin.close().await; - assert_eq!((first_table_count, migration_count), (8, 1)); + assert_eq!((first_table_count, migration_count), (8, 2)); } #[test] diff --git a/backend/migrations/0002_retention_delete_indexes.sql b/backend/migrations/0002_retention_delete_indexes.sql new file mode 100644 index 000000000..b214e2cfd --- /dev/null +++ b/backend/migrations/0002_retention_delete_indexes.sql @@ -0,0 +1,13 @@ +-- 保留期清理的删除谓词索引：清理任务按终态时间批量删除， +-- 无索引时每轮都会全表扫描三张审计/用量表。 + +create index if not exists idx_model_requests_retention + on model_requests (completed_at) + where outcome <> 'running'; + +create index if not exists idx_ops_events_retention + on ops_events (created_at) + where model_request_id is null; + +create index if not exists idx_admin_audit_events_retention + on admin_audit_events (created_at); diff --git a/deploy/config.example.yaml b/deploy/config.example.yaml index 3645ecb45..e65aeea18 100644 --- a/deploy/config.example.yaml +++ b/deploy/config.example.yaml @@ -36,6 +36,10 @@ store: # Docker 内由 CPR_REDIS_URL 覆盖；密码不要写进 URL。 url: 'redis://127.0.0.1:6379/' password: &redis_password '' + # 可选：PostgreSQL 连接池预算；省略时 max_connections=10、acquire_timeout_seconds=5。 + # pool: + # max_connections: 10 + # acquire_timeout_seconds: 5 admin: session_ttl_minutes: 1440