starred/codex-proxy-rs
1
0
Fork 0
mirror of https://github.com/zyycn/codex-proxy-rs.git synced 2026-11-06 16:06:08 +00:00
Code Issues Packages Projects Releases Wiki Activity

perf(store): index retention deletes, parallelize dashboard, expose pool budget

Browse Source 
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 <noreply@anthropic.com>
This commit is contained in:
zyycn
2026-07-26 20:28:52 +08:00
co-authored by Claude
parent a8b0fcf5ef
commit 01792a5dbe
8 changed files with 172 additions and 67 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/crates/gateway-store/src/lib.rs
+33 -1
View File
@@ -86,6 +86,37 @@ pub type StoreResult<T> = Result<T, StoreError>;
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<StoreBundle> {
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
backend/crates/gateway-store/src/postgres/mod.rs
+12 -3
View File
@@ -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<PgPool> {
/// 建立 PostgreSQL pool 并只执行冻结的 migration 集。
pub async fn connect_and_migrate(
database_url: &str,
pool_config: StorePoolConfig,
) -> StoreResult<PgPool> {
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"))?;
backend/crates/gateway-store/src/postgres/observability.rs
+20 -23
View File
@@ -759,35 +759,32 @@ impl ObservabilityRepository for PgObservabilityRepository {
range: ObservabilityRange,
) -> StoreResult<DashboardObservation> {
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,
backend/crates/gateway-store/src/postgres/retention.rs
+63 -35
View File
@@ -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<Utc>,
retention_days: u32,
label: &'static str,
) -> StoreResult<u64> {
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> {
u32::try_from(value).map_err(|_| StoreError::InvalidData {
entity: "retention settings",
backend/crates/gateway-store/tests/postgres/execution.rs
+18 -2
View File
@@ -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,
backend/crates/gateway-store/tests/postgres/mod.rs
+9 -3
View File
@@ -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]
backend/migrations/0002_retention_delete_indexes.sql
+13
View File
@@ -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);
deploy/config.example.yaml
+4
View File
@@ -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