From cadc45c5b88077fa4b904fb913b61fe19b308851 Mon Sep 17 00:00:00 2001 From: stabey <36232531+stabey@users.noreply.github.com> Date: Thu, 21 May 2026 16:56:56 +0800 Subject: [PATCH] fix(data): cleanup uses failed candidate status instead of 504 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The stale-pending cleanup task previously hardcoded status_code=504 and a generic timeout message for every usage row it finalized. When a request had already been observed as failing — e.g. upstream Connection reset by peer, watchdog 504, or an authenticated 4xx — the cleanup overwrote that context with a misleading "服务器超时" outcome and 504 status, hiding the real cause from the dashboards and customer. Pull the most recent failed/cancelled candidate per stale request_id and, if present, finalize the usage row with the candidate's status_code (defaulting to 502 when none was recorded) and error_message. Requests that have no terminal candidate (truly stuck pending/streaming) keep the existing 504 + timeout-message behavior, since they really are timeouts from the cleanup's perspective. Applied to all three SQL backends with parameterized UPDATE statements. The Postgres failed-candidate lookup orders by COALESCE(finished_at, started_at, created_at) DESC, matching the MySQL and SQLite ORDER BY clauses so the three backends pick the same "most recent terminal candidate" under every NULL combination of timing columns. Co-Authored-By: Claude Opus 4.7 --- .../aether-data/src/repository/usage/mysql.rs | 76 ++++++++- .../src/repository/usage/postgres/mod.rs | 84 ++++++++- .../src/repository/usage/postgres/tests.rs | 8 + .../src/repository/usage/sqlite.rs | 161 +++++++++++++++++- 4 files changed, 315 insertions(+), 14 deletions(-) diff --git a/crates/aether-data/src/repository/usage/mysql.rs b/crates/aether-data/src/repository/usage/mysql.rs index fed0498fc..2f2e95de2 100644 --- a/crates/aether-data/src/repository/usage/mysql.rs +++ b/crates/aether-data/src/repository/usage/mysql.rs @@ -519,13 +519,20 @@ WHERE request_id = ? continue; } - let error_message = stale_pending_error_message(&row.status, timeout_minutes); + let candidate_info = + latest_failed_candidate_mysql(&mut tx, &row.request_id).await?; + let (status_code, error_message) = resolve_stale_pending_failure( + candidate_info.as_ref(), + &row.status, + timeout_minutes, + ); + let status_code_i64 = i64::from(status_code); if row.billing_status == "pending" { sqlx::query( r#" UPDATE `usage` SET status = 'failed', - status_code = 504, + status_code = ?, error_message = ?, billing_status = 'void', finalized_at = ?, @@ -534,6 +541,7 @@ SET status = 'failed', WHERE request_id = ? "#, ) + .bind(status_code_i64) .bind(&error_message) .bind(to_i64(now_unix_secs, "usage finalized_at")?) .bind(&row.request_id) @@ -551,11 +559,12 @@ WHERE request_id = ? r#" UPDATE `usage` SET status = 'failed', - status_code = 504, + status_code = ?, error_message = ? WHERE request_id = ? "#, ) + .bind(status_code_i64) .bind(&error_message) .bind(&row.request_id) .execute(&mut *tx) @@ -688,6 +697,67 @@ fn stale_pending_error_message(status: &str, timeout_minutes: u64) -> String { format!("请求超时: 状态 '{status}' 超过 {timeout_minutes} 分钟未完成") } +struct FailedCandidateCleanupInfo { + status_code: Option, + error_message: Option, +} + +fn resolve_stale_pending_failure( + candidate: Option<&FailedCandidateCleanupInfo>, + status: &str, + timeout_minutes: u64, +) -> (u16, String) { + match candidate { + Some(info) => ( + info.status_code.unwrap_or(502), + info.error_message + .clone() + .unwrap_or_else(|| stale_pending_error_message(status, timeout_minutes)), + ), + None => (504, stale_pending_error_message(status, timeout_minutes)), + } +} + +async fn latest_failed_candidate_mysql( + tx: &mut sqlx::Transaction<'_, sqlx::MySql>, + request_id: &str, +) -> Result, DataLayerError> { + let row = sqlx::query( + r#" +SELECT status_code, error_message +FROM request_candidates +WHERE request_id = ? + AND status IN ('failed', 'cancelled') +ORDER BY + COALESCE(finished_at, started_at, created_at) DESC, + retry_index DESC, + candidate_index DESC +LIMIT 1 +"#, + ) + .bind(request_id) + .fetch_optional(&mut **tx) + .await + .map_sql_err()?; + + let Some(row) = row else { + return Ok(None); + }; + let status_code = row + .try_get::, _>("status_code") + .map_sql_err()? + .and_then(|value| u16::try_from(value).ok()); + let error_message = row + .try_get::, _>("error_message") + .map_sql_err()? + .map(|value| value.trim().to_string()) + .filter(|value| !value.is_empty()); + Ok(Some(FailedCandidateCleanupInfo { + status_code, + error_message, + })) +} + fn bind_upsert<'q>( mut query: sqlx::query::Query<'q, sqlx::MySql, sqlx::mysql::MySqlArguments>, usage: &'q UpsertUsageRecord, diff --git a/crates/aether-data/src/repository/usage/postgres/mod.rs b/crates/aether-data/src/repository/usage/postgres/mod.rs index de4d77a4d..c1077d7e3 100644 --- a/crates/aether-data/src/repository/usage/postgres/mod.rs +++ b/crates/aether-data/src/repository/usage/postgres/mod.rs @@ -1690,10 +1690,25 @@ SET status = 'completed', WHERE request_id = $1 "#; +const SELECT_LATEST_FAILED_CANDIDATE_FOR_STALE_REQUESTS_SQL: &str = r#" +SELECT DISTINCT ON (request_id) + request_id, + status_code, + error_message +FROM request_candidates +WHERE request_id = ANY($1) + AND status IN ('failed', 'cancelled') +ORDER BY request_id, + COALESCE(finished_at, started_at, created_at) DESC, + retry_index DESC, + candidate_index DESC, + created_at DESC +"#; + const UPDATE_FAILED_STALE_USAGE_SQL: &str = r#" UPDATE usage SET status = 'failed', - status_code = 504, + status_code = $3, error_message = $2 WHERE request_id = $1 "#; @@ -1702,7 +1717,7 @@ const UPDATE_FAILED_VOID_STALE_USAGE_SQL: &str = r#" WITH updated_usage AS ( UPDATE usage SET status = 'failed', - status_code = 504, + status_code = $4, error_message = $2, billing_status = 'void', finalized_at = $3, @@ -8292,17 +8307,44 @@ ORDER BY "usage".user_id ASC .iter() .map(|row| row.request_id.clone()) .collect::>(); - let completed_request_ids = if request_ids.is_empty() { - Vec::new() + let (completed_request_ids, failed_candidate_info) = if request_ids.is_empty() { + (Vec::new(), std::collections::HashMap::new()) } else { - sqlx::query(SELECT_COMPLETED_PENDING_REQUEST_IDS_SQL) - .bind(request_ids) + let completed = sqlx::query(SELECT_COMPLETED_PENDING_REQUEST_IDS_SQL) + .bind(&request_ids) .fetch_all(&mut *tx) .await .map_postgres_err()? .iter() .map(|row| row.try_get("request_id").map_postgres_err()) - .collect::, DataLayerError>>()? + .collect::, DataLayerError>>()?; + let failed_rows = + sqlx::query(SELECT_LATEST_FAILED_CANDIDATE_FOR_STALE_REQUESTS_SQL) + .bind(&request_ids) + .fetch_all(&mut *tx) + .await + .map_postgres_err()?; + let mut failed_map = std::collections::HashMap::new(); + for row in failed_rows { + let request_id: String = row.try_get("request_id").map_postgres_err()?; + let status_code = row + .try_get::, _>("status_code") + .map_postgres_err()? + .and_then(|value| u16::try_from(value).ok()); + let error_message = row + .try_get::, _>("error_message") + .map_postgres_err()? + .map(|value| value.trim().to_string()) + .filter(|value| !value.is_empty()); + failed_map.insert( + request_id, + FailedCandidateCleanupInfo { + status_code, + error_message, + }, + ); + } + (completed, failed_map) }; for row in stale_rows { @@ -8322,12 +8364,16 @@ ORDER BY "usage".user_id ASC continue; } - let error_message = stale_pending_error_message(&row.status, timeout_minutes); + let candidate_info = failed_candidate_info.get(&row.request_id); + let (status_code, error_message) = + resolve_stale_pending_failure(candidate_info, &row.status, timeout_minutes); + let status_code_i32 = i32::from(status_code); if row.billing_status == "pending" { sqlx::query(UPDATE_FAILED_VOID_STALE_USAGE_SQL) .bind(&row.request_id) .bind(&error_message) .bind(now) + .bind(status_code_i32) .execute(&mut *tx) .await .map_postgres_err()?; @@ -8335,6 +8381,7 @@ ORDER BY "usage".user_id ASC sqlx::query(UPDATE_FAILED_STALE_USAGE_SQL) .bind(&row.request_id) .bind(&error_message) + .bind(status_code_i32) .execute(&mut *tx) .await .map_postgres_err()?; @@ -8785,10 +8832,31 @@ struct StalePendingUsageRow { billing_status: String, } +struct FailedCandidateCleanupInfo { + status_code: Option, + error_message: Option, +} + fn stale_pending_error_message(status: &str, timeout_minutes: u64) -> String { format!("请求超时: 状态 '{status}' 超过 {timeout_minutes} 分钟未完成") } +fn resolve_stale_pending_failure( + candidate: Option<&FailedCandidateCleanupInfo>, + status: &str, + timeout_minutes: u64, +) -> (u16, String) { + match candidate { + Some(info) => ( + info.status_code.unwrap_or(502), + info.error_message + .clone() + .unwrap_or_else(|| stale_pending_error_message(status, timeout_minutes)), + ), + None => (504, stale_pending_error_message(status, timeout_minutes)), + } +} + async fn find_usage_by_request_id_in_tx( tx: &mut sqlx::Transaction<'_, Postgres>, request_id: &str, diff --git a/crates/aether-data/src/repository/usage/postgres/tests.rs b/crates/aether-data/src/repository/usage/postgres/tests.rs index 00eafe57a..9aafa5afd 100644 --- a/crates/aether-data/src/repository/usage/postgres/tests.rs +++ b/crates/aether-data/src/repository/usage/postgres/tests.rs @@ -825,6 +825,14 @@ fn usage_sql_clears_stale_failure_fields_for_non_failed_status_updates() { )); } +#[test] +fn stale_cleanup_failed_candidate_sql_orders_by_effective_timestamp() { + let sql = super::SELECT_LATEST_FAILED_CANDIDATE_FOR_STALE_REQUESTS_SQL; + assert!(sql.contains("COALESCE(finished_at, started_at, created_at) DESC")); + assert!(!sql.contains("finished_at DESC NULLS LAST")); + assert!(!sql.contains("started_at DESC NULLS LAST")); +} + #[test] fn usage_sql_does_not_allow_streaming_to_regress_back_to_pending() { assert!(super::UPSERT_SQL.contains( diff --git a/crates/aether-data/src/repository/usage/sqlite.rs b/crates/aether-data/src/repository/usage/sqlite.rs index fc8d6d619..9e7b09871 100644 --- a/crates/aether-data/src/repository/usage/sqlite.rs +++ b/crates/aether-data/src/repository/usage/sqlite.rs @@ -3525,13 +3525,20 @@ WHERE request_id = ? continue; } - let error_message = stale_pending_error_message(&row.status, timeout_minutes); + let candidate_info = + latest_failed_candidate_sqlite(&mut tx, &row.request_id).await?; + let (status_code, error_message) = resolve_stale_pending_failure( + candidate_info.as_ref(), + &row.status, + timeout_minutes, + ); + let status_code_i64 = i64::from(status_code); if row.billing_status == "pending" { sqlx::query( r#" UPDATE "usage" SET status = 'failed', - status_code = 504, + status_code = ?, error_message = ?, billing_status = 'void', finalized_at = ?, @@ -3540,6 +3547,7 @@ SET status = 'failed', WHERE request_id = ? "#, ) + .bind(status_code_i64) .bind(&error_message) .bind(to_i64(now_unix_secs, "usage finalized_at")?) .bind(&row.request_id) @@ -3557,11 +3565,12 @@ WHERE request_id = ? r#" UPDATE "usage" SET status = 'failed', - status_code = 504, + status_code = ?, error_message = ? WHERE request_id = ? "#, ) + .bind(status_code_i64) .bind(&error_message) .bind(&row.request_id) .execute(&mut *tx) @@ -3684,6 +3693,67 @@ fn stale_pending_error_message(status: &str, timeout_minutes: u64) -> String { format!("请求超时: 状态 '{status}' 超过 {timeout_minutes} 分钟未完成") } +struct FailedCandidateCleanupInfo { + status_code: Option, + error_message: Option, +} + +fn resolve_stale_pending_failure( + candidate: Option<&FailedCandidateCleanupInfo>, + status: &str, + timeout_minutes: u64, +) -> (u16, String) { + match candidate { + Some(info) => ( + info.status_code.unwrap_or(502), + info.error_message + .clone() + .unwrap_or_else(|| stale_pending_error_message(status, timeout_minutes)), + ), + None => (504, stale_pending_error_message(status, timeout_minutes)), + } +} + +async fn latest_failed_candidate_sqlite( + tx: &mut sqlx::Transaction<'_, sqlx::Sqlite>, + request_id: &str, +) -> Result, DataLayerError> { + let row = sqlx::query( + r#" +SELECT status_code, error_message +FROM request_candidates +WHERE request_id = ? + AND status IN ('failed', 'cancelled') +ORDER BY + COALESCE(finished_at, started_at, created_at) DESC, + retry_index DESC, + candidate_index DESC +LIMIT 1 +"#, + ) + .bind(request_id) + .fetch_optional(&mut **tx) + .await + .map_sql_err()?; + + let Some(row) = row else { + return Ok(None); + }; + let status_code = row + .try_get::, _>("status_code") + .map_sql_err()? + .and_then(|value| u16::try_from(value).ok()); + let error_message = row + .try_get::, _>("error_message") + .map_sql_err()? + .map(|value| value.trim().to_string()) + .filter(|value| !value.is_empty()); + Ok(Some(FailedCandidateCleanupInfo { + status_code, + error_message, + })) +} + fn bind_upsert<'q>( mut query: sqlx::query::Query<'q, sqlx::Sqlite, sqlx::sqlite::SqliteArguments<'q>>, usage: &'q UpsertUsageRecord, @@ -4116,6 +4186,91 @@ ORDER BY request_id assert_eq!(snapshot, ("void".to_string(), Some(10))); } + #[tokio::test] + async fn sqlite_usage_write_repository_cleanup_uses_failed_candidate_status_when_present() { + let pool = sqlx::sqlite::SqlitePoolOptions::new() + .max_connections(1) + .connect("sqlite::memory:") + .await + .expect("sqlite pool should connect"); + run_sqlite_migrations(&pool) + .await + .expect("sqlite migrations should run"); + seed_stats_targets(&pool).await; + + let repository = SqliteUsageWriteRepository::new(pool.clone()); + repository + .upsert(sample_usage( + "request-upstream-reset", + "pending", + "pending", + 1, + )) + .await + .expect("pending usage should upsert"); + repository + .upsert(sample_usage("request-stuck", "pending", "pending", 1)) + .await + .expect("pending usage should upsert"); + + // request-upstream-reset has a failed candidate carrying a concrete 502 status + // and a connection-reset message — cleanup should use them instead of 504. + // request-stuck has only a still-pending candidate, so cleanup should fall back to 504. + sqlx::query( + r#" +INSERT INTO request_candidates ( + id, + request_id, + candidate_index, + retry_index, + status, + status_code, + error_message, + is_cached, + created_at, + started_at, + finished_at +) VALUES + ('candidate-reset', 'request-upstream-reset', 0, 0, 'failed', 502, 'upstream connection reset by peer', 0, 1, 2, 3), + ('candidate-stuck', 'request-stuck', 0, 0, 'pending', NULL, NULL, 0, 1, NULL, NULL) +"#, + ) + .execute(&pool) + .await + .expect("request candidates should seed"); + + let summary = repository + .cleanup_stale_pending_requests(2, 10, 5, 5) + .await + .expect("cleanup should run"); + assert_eq!(summary.recovered, 0); + assert_eq!(summary.failed, 2); + + let reset = repository + .find_by_request_id("request-upstream-reset") + .await + .expect("upstream-reset usage should load") + .expect("upstream-reset usage should exist"); + assert_eq!(reset.status, "failed"); + assert_eq!(reset.status_code, Some(502)); + assert_eq!( + reset.error_message.as_deref(), + Some("upstream connection reset by peer") + ); + + let stuck = repository + .find_by_request_id("request-stuck") + .await + .expect("stuck usage should load") + .expect("stuck usage should exist"); + assert_eq!(stuck.status, "failed"); + assert_eq!(stuck.status_code, Some(504)); + assert!(stuck + .error_message + .as_deref() + .is_some_and(|message| message.contains("超过 5 分钟未完成"))); + } + #[tokio::test] async fn sqlite_usage_read_repository_reads_usage_contract_views() { let pool = sqlx::sqlite::SqlitePoolOptions::new()