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

fix: harden concurrency limits and high-RPM runtime paths

Browse Source 
Bound request, stream, queue, and shutdown resource lifetimes. Reduce scheduler and Redis hot-path work and isolate database maintenance. Include regression coverage, load probes, and concurrency audit results.
This commit is contained in:
elky
2026-09-10 08:14:58 +08:00
parent 361952ada9
commit ecc16673eb
149 changed files with 27963 additions and 1926 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/aether-data/runtime/src/backend/maintenance/postgres.rs
+23 -4
View File
@@ -16,17 +16,36 @@ impl PostgresBackend {
table_names: &[&str],
) -> Result<DatabaseMaintenanceSummary, DataLayerError> {
let mut summary = DatabaseMaintenanceSummary::default();
if table_names.is_empty() {
return Ok(summary);
}
// VACUUM cannot run inside a transaction. Discard this connection on every
// exit path so its longer session deadlines never leak into request queries.
let mut conn = self.pool().acquire().await.map_postgres_err()?;
conn.close_on_drop();
sqlx::query("SET statement_timeout = '5min'")
.execute(&mut *conn)
.await
.map_postgres_err()?;
sqlx::query("SET lock_timeout = '30s'")
.execute(&mut *conn)
.await
.map_postgres_err()?;
for table_name in table_names {
let table_name = maintenance_identifier(table_name)?;
summary.attempted += 1;
let statement = format!("VACUUM ANALYZE \"{table_name}\"");
if sqlx::raw_sql(&statement)
.execute(self.pool())
match sqlx::query(&statement)
.execute(&mut *conn)
.await
.map_postgres_err()
.is_ok()
{
summary.succeeded += 1;
Ok(_) => summary.succeeded += 1,
Err(error) => tracing::warn!(
table_name,
error = %error,
"PostgreSQL table maintenance failed"
),
}
}
Ok(summary)
crates/aether-data/runtime/src/backend/postgres.rs
+34
View File
@@ -274,6 +274,40 @@ mod tests {
use super::PostgresBackend;
use crate::driver::postgres::{PostgresLeaseRunnerConfig, PostgresPoolConfig};
#[tokio::test]
async fn maintenance_and_aggregation_futures_are_send() {
fn assert_send(_: impl Send) {}
let backend = PostgresBackend::from_config(PostgresPoolConfig {
database_url: "postgres://localhost/aether".to_string(),
min_connections: 0,
..PostgresPoolConfig::default()
})
.unwrap();
let now = chrono::Utc::now();
let daily = crate::StatsDailyAggregationInput {
target_day_utc: now,
aggregated_at: now,
};
let hourly = crate::StatsHourlyAggregationInput {
target_hour_utc: now,
aggregated_at: now,
};
let wallet = crate::WalletDailyUsageAggregationInput {
billing_date: "2026-09-09".to_string(),
billing_timezone: "UTC".to_string(),
window_start_unix_secs: 0,
window_end_unix_secs: 86_400,
aggregated_at_unix_secs: 86_400,
};
// Drop without polling: these are compile-time checks for spawned workers.
assert_send(backend.run_table_maintenance(&["usage"]));
assert_send(backend.aggregate_stats_daily(&daily));
assert_send(backend.aggregate_stats_hourly(&hourly));
assert_send(backend.aggregate_wallet_daily_usage(&wallet));
}
#[tokio::test]
async fn backend_retains_config_and_pool() {
let config = PostgresPoolConfig {
crates/aether-data/runtime/src/backend/stats/postgres_daily/mod.rs
+6
View File
@@ -66,6 +66,12 @@ async fn perform_stats_aggregation_for_day(
) -> Result<StatsDailyAggregationSummary, sqlx::Error> {
let day_end_utc = day_start_utc + chrono::Duration::days(1);
let mut tx = pool.begin().await?;
sqlx::query("SET LOCAL statement_timeout = '5min'")
.execute(&mut *tx)
.await?;
sqlx::query("SET LOCAL lock_timeout = '30s'")
.execute(&mut *tx)
.await?;
let aggregate_row = sqlx::query(SELECT_STATS_DAILY_AGGREGATE_SQL)
.bind(day_start_utc)
.bind(day_end_utc)
crates/aether-data/runtime/src/backend/stats/postgres_hourly/mod.rs
+6
View File
@@ -65,6 +65,12 @@ async fn perform_stats_hourly_aggregation_for_hour(
) -> Result<StatsHourlyAggregationSummary, sqlx::Error> {
let hour_end = hour_utc + chrono::Duration::hours(1);
let mut tx = pool.begin().await?;
sqlx::query("SET LOCAL statement_timeout = '5min'")
.execute(&mut *tx)
.await?;
sqlx::query("SET LOCAL lock_timeout = '30s'")
.execute(&mut *tx)
.await?;
let row = sqlx::query(SELECT_STATS_HOURLY_AGGREGATE_SQL)
.bind(hour_utc)
crates/aether-data/runtime/src/backend/wallet/postgres.rs
+8
View File
@@ -104,6 +104,14 @@ impl PostgresBackend {
let window_end = unix_secs_to_utc(input.window_end_unix_secs, "window_end")?;
let aggregated_at = unix_secs_to_utc(input.aggregated_at_unix_secs, "aggregated_at")?;
let mut tx = self.pool().begin().await.map_postgres_err()?;
sqlx::query("SET LOCAL statement_timeout = '5min'")
.execute(&mut *tx)
.await
.map_postgres_err()?;
sqlx::query("SET LOCAL lock_timeout = '30s'")
.execute(&mut *tx)
.await
.map_postgres_err()?;
let aggregated_wallets = sqlx::query(UPSERT_WALLET_DAILY_USAGE_LEDGER_SQL)
.bind(window_start)
crates/aether-data/runtime/src/lifecycle/backfill/postgres.rs
+1 -1
View File
@@ -57,7 +57,7 @@ pub(super) struct AppliedBackfill {
}
pub async fn run_backfills(pool: &PgPool) -> Result<(), MigrateError> {
let mut conn = pool.acquire().await?;
let mut conn = aether_data_postgres::acquire_postgres_migration_connection(pool).await?;
if BACKFILL_MIGRATOR.locking {
conn.lock().await?;
crates/aether-data/runtime/src/repository/candidates/memory.rs
+172 -44
View File
@@ -45,7 +45,54 @@ fn merge_extra_data(
#[derive(Debug, Default)]
pub struct InMemoryRequestCandidateRepository {
by_id: RwLock<BTreeMap<String, StoredRequestCandidate>>,
rows: RwLock<CandidateRows>,
}
#[derive(Debug, Default)]
struct CandidateRows {
by_id: BTreeMap<String, StoredRequestCandidate>,
by_request: BTreeMap<String, BTreeSet<String>>,
by_created: BTreeSet<(std::cmp::Reverse<u64>, String)>,
}
impl CandidateRows {
fn remove(&mut self, id: &str) -> Option<StoredRequestCandidate> {
let row = self.by_id.remove(id)?;
self.by_created
.remove(&(std::cmp::Reverse(row.created_at_unix_ms), row.id.clone()));
if let Some(ids) = self.by_request.get_mut(&row.request_id) {
ids.remove(id);
if ids.is_empty() {
self.by_request.remove(&row.request_id);
}
}
Some(row)
}
fn insert(&mut self, row: StoredRequestCandidate) -> &StoredRequestCandidate {
// Keep all indexes behind one lock and sanitize every insertion. Reads
// can clone these records without rebuilding their diagnostic JSON.
let row = sanitize_stored_candidate(row);
self.remove(&row.id);
self.by_request
.entry(row.request_id.clone())
.or_default()
.insert(row.id.clone());
self.by_created
.insert((std::cmp::Reverse(row.created_at_unix_ms), row.id.clone()));
self.by_id
.entry(row.id.clone())
.insert_entry(row)
.into_mut()
}
fn for_request(&self, request_id: &str) -> impl Iterator<Item = &StoredRequestCandidate> {
self.by_request
.get(request_id)
.into_iter()
.flatten()
.filter_map(|id| self.by_id.get(id))
}
}
impl InMemoryRequestCandidateRepository {
@@ -53,12 +100,12 @@ impl InMemoryRequestCandidateRepository {
where
I: IntoIterator<Item = StoredRequestCandidate>,
{
let mut by_id = BTreeMap::new();
for item in items.into_iter().map(sanitize_stored_candidate) {
by_id.insert(item.id.clone(), item);
let mut rows = CandidateRows::default();
for item in items {
rows.insert(item);
}
Self {
by_id: RwLock::new(by_id),
rows: RwLock::new(rows),
}
}
}
@@ -70,13 +117,11 @@ impl RequestCandidateReadRepository for InMemoryRequestCandidateRepository {
request_id: &str,
) -> Result<Vec<StoredRequestCandidate>, DataLayerError> {
let mut rows = self
.by_id
.rows
.read()
.expect("request candidate repository lock")
.values()
.filter(|row| row.request_id == request_id)
.for_request(request_id)
.cloned()
.map(sanitize_stored_candidate)
.collect::<Vec<_>>();
rows.sort_by(|left, right| {
left.candidate_index
@@ -95,17 +140,14 @@ impl RequestCandidateReadRepository for InMemoryRequestCandidateRepository {
return Ok(Vec::new());
}
let mut rows = self
.by_id
.read()
.expect("request candidate repository lock")
.values()
let rows = self.rows.read().expect("request candidate repository lock");
Ok(rows
.by_created
.iter()
.take(limit)
.filter_map(|(_, id)| rows.by_id.get(id))
.cloned()
.map(sanitize_stored_candidate)
.collect::<Vec<_>>();
rows.sort_by_key(|entry| std::cmp::Reverse(entry.created_at_unix_ms));
rows.truncate(limit);
Ok(rows)
.collect())
}
async fn list_by_provider_id(
@@ -118,19 +160,33 @@ impl RequestCandidateReadRepository for InMemoryRequestCandidateRepository {
}
let mut rows = self
.by_id
.rows
.read()
.expect("request candidate repository lock")
.by_id
.values()
.filter(|row| row.provider_id.as_deref() == Some(provider_id))
.cloned()
.map(sanitize_stored_candidate)
.collect::<Vec<_>>();
rows.sort_by_key(|entry| std::cmp::Reverse(entry.created_at_unix_ms));
rows.truncate(limit);
Ok(rows)
}
async fn list_recent_runtime(
&self,
limit: usize,
) -> Result<Vec<StoredRequestCandidate>, DataLayerError> {
let rows = self.rows.read().expect("request candidate repository lock");
Ok(rows
.by_created
.iter()
.take(limit)
.filter_map(|(_, id)| rows.by_id.get(id))
.map(StoredRequestCandidate::runtime_snapshot)
.collect())
}
async fn list_finalized_by_endpoint_ids_since(
&self,
endpoint_ids: &[String],
@@ -143,9 +199,10 @@ impl RequestCandidateReadRepository for InMemoryRequestCandidateRepository {
let endpoint_ids = endpoint_ids.iter().cloned().collect::<BTreeSet<_>>();
let mut rows = self
.by_id
.rows
.read()
.expect("request candidate repository lock")
.by_id
.values()
.filter(|row| {
row.endpoint_id
@@ -160,7 +217,6 @@ impl RequestCandidateReadRepository for InMemoryRequestCandidateRepository {
)
})
.cloned()
.map(sanitize_stored_candidate)
.collect::<Vec<_>>();
rows.sort_by_key(|entry| std::cmp::Reverse(entry.created_at_unix_ms));
rows.truncate(limit);
@@ -179,9 +235,10 @@ impl RequestCandidateReadRepository for InMemoryRequestCandidateRepository {
let endpoint_ids = endpoint_ids.iter().cloned().collect::<BTreeSet<_>>();
let mut counts = BTreeMap::<(String, &'static str), u64>::new();
for row in self
.by_id
.rows
.read()
.expect("request candidate repository lock")
.by_id
.values()
{
let Some(endpoint_id) = row.endpoint_id.as_ref() else {
@@ -243,9 +300,10 @@ impl RequestCandidateReadRepository for InMemoryRequestCandidateRepository {
let mut buckets = BTreeMap::<(String, u32), PublicHealthTimelineBucket>::new();
for row in self
.by_id
.rows
.read()
.expect("request candidate repository lock")
.by_id
.values()
{
let Some(endpoint_id) = row.endpoint_id.as_ref() else {
@@ -319,19 +377,17 @@ impl RequestCandidateWriteRepository for InMemoryRequestCandidateRepository {
candidate.sanitize_for_persistence();
candidate.validate()?;
let mut by_id = self
.by_id
let mut rows = self
.rows
.write()
.expect("request candidate repository lock");
let existing = by_id
.values()
let existing = rows
.for_request(&candidate.request_id)
.find(|row| {
row.request_id == candidate.request_id
&& row.candidate_index == candidate.candidate_index
row.candidate_index == candidate.candidate_index
&& row.retry_index == candidate.retry_index
})
.cloned()
.map(sanitize_stored_candidate);
.cloned();
let preserve_existing_lifecycle = existing.as_ref().is_some_and(|row| {
request_candidate_lifecycle_would_regress(row.status, candidate.status)
@@ -449,10 +505,7 @@ impl RequestCandidateWriteRepository for InMemoryRequestCandidateRepository {
.or_else(|| existing.as_ref().and_then(|row| row.finished_at_unix_ms))
},
};
let stored = sanitize_stored_candidate(stored);
by_id.insert(stored.id.clone(), stored.clone());
Ok(stored)
Ok(rows.insert(stored).clone())
}
async fn delete_created_before(
@@ -464,11 +517,12 @@ impl RequestCandidateWriteRepository for InMemoryRequestCandidateRepository {
return Ok(0);
}
let mut by_id = self
.by_id
let mut rows = self
.rows
.write()
.expect("request candidate repository lock");
let mut ids = by_id
let mut ids = rows
.by_id
.values()
.filter(|row| row.created_at_unix_ms < created_before_unix_secs * 1000)
.map(|row| (row.created_at_unix_ms, row.id.clone()))
@@ -477,7 +531,7 @@ impl RequestCandidateWriteRepository for InMemoryRequestCandidateRepository {
let mut deleted = 0usize;
for (_, id) in ids.into_iter().take(limit) {
if by_id.remove(&id).is_some() {
if rows.remove(&id).is_some() {
deleted += 1;
}
}
@@ -546,6 +600,79 @@ mod tests {
assert_eq!(rows[1].request_id, "req-1");
}
#[tokio::test]
async fn request_index_tracks_replaced_ids_and_removes_empty_requests() {
let repository = InMemoryRequestCandidateRepository::seed([
sample_candidate("same-id", "old-request", 100),
sample_candidate("same-id", "new-request", 200),
sample_candidate("other-id", "new-request", 300),
]);
assert!(repository
.list_by_request_id("old-request")
.await
.unwrap()
.is_empty());
assert_eq!(
repository
.list_by_request_id("new-request")
.await
.unwrap()
.len(),
2
);
assert_eq!(repository.delete_created_before(1, 1).await.unwrap(), 1);
let remaining = repository.list_by_request_id("new-request").await.unwrap();
assert_eq!(remaining.len(), 1);
assert_eq!(remaining[0].id, "other-id");
assert_eq!(repository.delete_created_before(1, 1).await.unwrap(), 1);
let rows = repository.rows.read().unwrap();
assert!(rows.by_id.is_empty());
assert!(rows.by_request.is_empty());
assert!(rows.by_created.is_empty());
}
#[tokio::test]
async fn recent_index_preserves_equal_timestamp_order_and_replaced_dates() {
let repository = InMemoryRequestCandidateRepository::seed([
sample_candidate("b", "req-b", 400),
sample_candidate("a", "req-a", 200),
sample_candidate("c", "req-c", 200),
sample_candidate("b", "req-b", 100),
]);
let recent = repository.list_recent(2).await.unwrap();
assert_eq!(
recent.iter().map(|row| row.id.as_str()).collect::<Vec<_>>(),
["a", "c"]
);
assert_eq!(repository.list_recent(0).await.unwrap().len(), 0);
assert_eq!(repository.rows.read().unwrap().by_created.len(), 3);
}
#[tokio::test]
async fn runtime_reads_keep_metadata_without_diagnostic_payloads() {
let mut candidate = sample_candidate("candidate", "request", 100);
candidate.extra_data = Some(json!({"upstream_response": {"body": "x".repeat(32_768)}}));
candidate.error_message = Some("diagnostic detail".into());
candidate.required_capabilities = Some(json!({"vision": true}));
candidate.concurrent_requests = Some(17);
let repository = InMemoryRequestCandidateRepository::seed([candidate]);
let full = repository.list_recent(1).await.unwrap();
let runtime = repository.list_recent_runtime(1).await.unwrap();
assert_eq!(
runtime,
full.iter()
.map(StoredRequestCandidate::runtime_snapshot)
.collect::<Vec<_>>()
);
assert_eq!(runtime[0].concurrent_requests, Some(17));
assert!(runtime[0].extra_data.is_none());
assert!(runtime[0].error_message.is_none());
assert!(runtime[0].required_capabilities.is_none());
assert!(full[0].extra_data.is_some());
assert_eq!(repository.list_recent(1).await.unwrap(), full);
assert!(repository.list_recent_runtime(0).await.unwrap().is_empty());
}
#[tokio::test]
async fn lists_recent_request_candidates_in_descending_created_order() {
let repository = InMemoryRequestCandidateRepository::seed(vec![
@@ -602,10 +729,11 @@ mod tests {
{
let stored = repository
.by_id
.rows
.read()
.expect("request candidate repository lock");
let candidate = stored
.by_id
.get("cand-raw")
.expect("seeded candidate should exist");
assert_eq!(
@@ -628,10 +756,10 @@ mod tests {
bypassed_candidate.id = "cand-bypassed".to_string();
bypassed_candidate.request_id = "req-bypassed".to_string();
repository
.by_id
.rows
.write()
.expect("request candidate repository lock")
.insert(bypassed_candidate.id.clone(), bypassed_candidate);
.insert(bypassed_candidate);
let rows = repository
.list_recent(10)
crates/aether-data/runtime/src/repository/usage/memory/tests.rs
+14
View File
@@ -66,6 +66,7 @@ fn sample_usage(request_id: &str, created_at_unix_ms: i64) -> StoredRequestUsage
fn sample_upsert_usage_record(request_id: &str) -> UpsertUsageRecord {
UpsertUsageRecord {
capture_retention: Default::default(),
request_id: request_id.to_string(),
user_id: None,
api_key_id: None,
@@ -535,6 +536,7 @@ async fn stale_pending_update_does_not_regress_finalized_usage() {
let repository = InMemoryUsageReadRepository::default();
repository
.upsert(UpsertUsageRecord {
capture_retention: Default::default(),
request_id: "req-finalized-1".to_string(),
user_id: Some("user-1".to_string()),
api_key_id: Some("api-key-1".to_string()),
@@ -608,6 +610,7 @@ async fn stale_pending_update_does_not_regress_finalized_usage() {
repository
.upsert(UpsertUsageRecord {
capture_retention: Default::default(),
request_id: "req-finalized-1".to_string(),
user_id: Some("user-1".to_string()),
api_key_id: Some("api-key-1".to_string()),
@@ -695,6 +698,7 @@ async fn upsert_allows_completed_recovery_after_void_failure() {
let repository = InMemoryUsageReadRepository::default();
repository
.upsert(UpsertUsageRecord {
capture_retention: Default::default(),
request_id: "req-recover-1".to_string(),
user_id: Some("user-1".to_string()),
api_key_id: Some("api-key-1".to_string()),
@@ -768,6 +772,7 @@ async fn upsert_allows_completed_recovery_after_void_failure() {
repository
.upsert(UpsertUsageRecord {
capture_retention: Default::default(),
request_id: "req-recover-1".to_string(),
user_id: Some("user-1".to_string()),
api_key_id: Some("api-key-1".to_string()),
@@ -1009,6 +1014,7 @@ async fn stale_pending_update_does_not_regress_streaming_usage() {
let repository = InMemoryUsageReadRepository::default();
repository
.upsert(UpsertUsageRecord {
capture_retention: Default::default(),
request_id: "req-streaming-1".to_string(),
user_id: Some("user-1".to_string()),
api_key_id: Some("api-key-1".to_string()),
@@ -1084,6 +1090,7 @@ async fn stale_pending_update_does_not_regress_streaming_usage() {
repository
.upsert(UpsertUsageRecord {
capture_retention: Default::default(),
request_id: "req-streaming-1".to_string(),
user_id: Some("user-1".to_string()),
api_key_id: Some("api-key-1".to_string()),
@@ -1339,6 +1346,7 @@ async fn upsert_writes_usage_record() {
let repository = InMemoryUsageReadRepository::default();
let stored = repository
.upsert(UpsertUsageRecord {
capture_retention: Default::default(),
request_id: "req-upsert-1".to_string(),
user_id: Some("user-1".to_string()),
api_key_id: Some("key-1".to_string()),
@@ -1430,6 +1438,7 @@ async fn upsert_defaults_created_at_to_second_timestamp() {
let repository = InMemoryUsageReadRepository::default();
let stored = repository
.upsert(UpsertUsageRecord {
capture_retention: Default::default(),
request_id: "req-upsert-ms-default".to_string(),
user_id: None,
api_key_id: None,
@@ -1509,6 +1518,7 @@ async fn upsert_does_not_backfill_legacy_output_price_from_request_metadata() {
let repository = InMemoryUsageReadRepository::default();
let stored = repository
.upsert(UpsertUsageRecord {
capture_retention: Default::default(),
request_id: "req-upsert-price-metadata".to_string(),
user_id: None,
api_key_id: None,
@@ -1591,6 +1601,7 @@ async fn upsert_does_not_backfill_typed_body_refs_from_request_metadata() {
let repository = InMemoryUsageReadRepository::default();
let stored = repository
.upsert(UpsertUsageRecord {
capture_retention: Default::default(),
request_id: "req-upsert-body-ref-metadata".to_string(),
user_id: None,
api_key_id: None,
@@ -1673,6 +1684,7 @@ async fn upsert_keeps_typed_routing_fields_out_of_request_metadata() {
let repository = InMemoryUsageReadRepository::default();
let stored = repository
.upsert(UpsertUsageRecord {
capture_retention: Default::default(),
request_id: "req-upsert-routing-metadata".to_string(),
user_id: None,
api_key_id: None,
@@ -1771,6 +1783,7 @@ async fn upsert_does_not_persist_legacy_display_columns_for_new_rows() {
let repository = InMemoryUsageReadRepository::default();
let stored = repository
.upsert(UpsertUsageRecord {
capture_retention: Default::default(),
request_id: "req-upsert-display-columns".to_string(),
user_id: Some("user-1".to_string()),
api_key_id: Some("key-1".to_string()),
@@ -1921,6 +1934,7 @@ async fn upsert_preserves_existing_legacy_display_columns_when_new_write_omits_t
}]);
let stored = repository
.upsert(UpsertUsageRecord {
capture_retention: Default::default(),
request_id: "req-existing-display-columns".to_string(),
user_id: Some("user-1".to_string()),
api_key_id: Some("key-1".to_string()),
crates/aether-data/runtime/src/repository/usage/mod.rs
+1
View File
@@ -57,6 +57,7 @@ mod tests {
#[test]
fn strip_deprecated_usage_display_fields_clears_legacy_display_columns() {
let usage = strip_deprecated_usage_display_fields(UpsertUsageRecord {
capture_retention: Default::default(),
request_id: "req-1".to_string(),
user_id: Some("user-1".to_string()),
api_key_id: Some("key-1".to_string()),