diff --git a/backend/Cargo.lock b/backend/Cargo.lock index a0818093b..b5962f51d 100644 --- a/backend/Cargo.lock +++ b/backend/Cargo.lock @@ -384,6 +384,7 @@ dependencies = [ "thiserror 2.0.18", "tokio", "tokio-tungstenite", + "tokio-util", "tower", "tower-http", "tracing", @@ -2873,6 +2874,7 @@ dependencies = [ "bytes", "futures-core", "futures-sink", + "futures-util", "pin-project-lite", "tokio", ] diff --git a/backend/Cargo.toml b/backend/Cargo.toml index 77eeec89d..c17fc58e1 100644 --- a/backend/Cargo.toml +++ b/backend/Cargo.toml @@ -36,6 +36,7 @@ subtle = "2.6.1" tar = "0.4.44" thiserror = "2.0.18" tokio = { version = "1.52.3", features = ["macros", "rt-multi-thread", "signal", "fs", "time", "sync", "net", "test-util"] } +tokio-util = { version = "0.7.18", features = ["rt"] } tokio-tungstenite = { version = "0.28.0", features = ["proxy", "rustls-tls-native-roots"] } tower = { version = "0.5.3", features = ["util"] } tower-http = { version = "0.6.11", features = ["trace", "cors", "request-id", "timeout", "fs"] } diff --git a/backend/src/api/admin/auth_routes.rs b/backend/src/api/admin/auth_routes.rs index b202a4ee4..4c6c9b5f5 100644 --- a/backend/src/api/admin/auth_routes.rs +++ b/backend/src/api/admin/auth_routes.rs @@ -108,11 +108,14 @@ pub(crate) async fn logout( headers: HeaderMap, ) -> Result { if let Some(session_id) = admin_session_cookie(&headers) { - let _ = state + if let Err(error) = state .services .admin_sessions .delete_session(&session_id) - .await; + .await + { + tracing::warn!(error = %error, "failed to revoke admin session during logout"); + } } let mut response = AdminResponse::new( diff --git a/backend/src/bootstrap/services.rs b/backend/src/bootstrap/services.rs index a1281f553..d9cb02a6f 100644 --- a/backend/src/bootstrap/services.rs +++ b/backend/src/bootstrap/services.rs @@ -1,6 +1,6 @@ //! 运行时服务集合。 -use std::{env, error::Error, sync::Arc, time::Duration}; +use std::{env, error::Error, future::IntoFuture, sync::Arc, time::Duration}; use axum::Router; @@ -434,12 +434,38 @@ pub async fn run() -> Result<(), Box> { let listener = tokio::net::TcpListener::bind((host.as_str(), port)).await?; let (router, task_coordinator) = build_router(config).await?; - let serve_result = axum::serve( - listener, - router.into_make_service_with_connect_info::(), - ) - .with_graceful_shutdown(crate::bootstrap::shutdown::shutdown_signal()) - .await; + let serve_result = { + let (graceful_shutdown_tx, graceful_shutdown_rx) = tokio::sync::oneshot::channel(); + let serve = axum::serve( + listener, + router.into_make_service_with_connect_info::(), + ) + .with_graceful_shutdown(async move { + let _ = graceful_shutdown_rx.await; + }) + .into_future(); + tokio::pin!(serve); + tokio::select! { + result = &mut serve => result, + () = crate::bootstrap::shutdown::shutdown_signal() => { + let _ = graceful_shutdown_tx.send(()); + tokio::select! { + result = &mut serve => result, + () = tokio::time::sleep(HTTP_DRAIN_TIMEOUT) => { + tracing::warn!( + timeout_secs = HTTP_DRAIN_TIMEOUT.as_secs(), + "HTTP graceful shutdown timed out; cancelling remaining connections" + ); + Ok(()) + } + () = crate::bootstrap::shutdown::shutdown_signal() => { + tracing::warn!("received a second shutdown signal; cancelling remaining connections"); + Ok(()) + } + } + } + } + }; task_coordinator.shutdown().await; serve_result?; @@ -456,6 +482,8 @@ pub async fn run() -> Result<(), Box> { Ok(()) } +const HTTP_DRAIN_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(20); + async fn build_router( mut config: AppConfig, ) -> Result< diff --git a/backend/src/bootstrap/tasks/coordinator.rs b/backend/src/bootstrap/tasks/coordinator.rs index a1b7a51f3..754b2659e 100644 --- a/backend/src/bootstrap/tasks/coordinator.rs +++ b/backend/src/bootstrap/tasks/coordinator.rs @@ -136,9 +136,21 @@ impl TaskCoordinator { /// 关闭所有后台任务。 pub async fn shutdown(self) { tracing::info!("正在关闭后台任务"); - for (name, handle) in self.handles { + let shutdown_tasks = self.handles.into_iter().map(|(name, handle)| async move { handle.shutdown().await; tracing::info!(task = name, "后台任务已停止"); + }); + if tokio::time::timeout( + COORDINATOR_SHUTDOWN_TIMEOUT, + futures::future::join_all(shutdown_tasks), + ) + .await + .is_err() + { + tracing::warn!( + timeout_secs = COORDINATOR_SHUTDOWN_TIMEOUT.as_secs(), + "后台任务并行关闭超时" + ); } tracing::info!("所有后台任务已停止"); } @@ -200,6 +212,7 @@ impl SchedulerHandle { } const TASK_SHUTDOWN_TIMEOUT: Duration = Duration::from_secs(5); +const COORDINATOR_SHUTDOWN_TIMEOUT: Duration = Duration::from_secs(6); async fn wait_for_task_shutdown(handle: JoinHandle<()>) { let mut handle = handle; diff --git a/backend/src/bootstrap/tasks/model_refresh.rs b/backend/src/bootstrap/tasks/model_refresh.rs index ab0463294..abeacbf03 100644 --- a/backend/src/bootstrap/tasks/model_refresh.rs +++ b/backend/src/bootstrap/tasks/model_refresh.rs @@ -85,10 +85,8 @@ impl ModelRefreshTask { self.installation_id.as_deref(), ) .await; - for selected in &selected_accounts { - self.account_pool - .release_without_request_usage(&selected.account.id) - .await; + for selected in selected_accounts { + selected.release().await; } let result = result?; diff --git a/backend/src/bootstrap/tasks/token_refresh.rs b/backend/src/bootstrap/tasks/token_refresh.rs index 50eb9f404..9d686a831 100644 --- a/backend/src/bootstrap/tasks/token_refresh.rs +++ b/backend/src/bootstrap/tasks/token_refresh.rs @@ -92,7 +92,7 @@ where fn shutdown(&mut self) -> super::periodic::TaskFuture<'_, ()> { Box::pin(async move { - self.service.clear_scheduled_timers().await; + self.service.shutdown().await; }) } } diff --git a/backend/src/dispatch/recovery/cloudflare.rs b/backend/src/dispatch/recovery/cloudflare.rs index 88440ccc8..1dcb14f86 100644 --- a/backend/src/dispatch/recovery/cloudflare.rs +++ b/backend/src/dispatch/recovery/cloudflare.rs @@ -1,7 +1,5 @@ //! Cloudflare challenge and path-block recovery shared by dispatch routes. -use std::time::Duration as StdDuration; - use chrono::Utc; use crate::{ @@ -61,7 +59,17 @@ impl CloudflareRecovery { account_pool .set_cloudflare_cooldown_until(account_id, cooldown.cooldown_until) .await; - self.clear_challenge_cookies_after_cooldown(account_id, cooldown.delay_seconds); + if let Err(error) = self + .cookie_store + .expire_account_cookies_at(account_id, cooldown.cooldown_until) + .await + { + tracing::warn!( + account_id, + error = %error, + "failed to persist Cloudflare challenge cookie cleanup deadline" + ); + } } pub async fn apply_path_block(&self, account_pool: &AccountPoolService, account_id: &str) { @@ -90,17 +98,6 @@ impl CloudflareRecovery { async fn delete_account_cookies(&self, account_id: &str, reason: &str) { delete_account_cookies(&self.cookie_store, account_id, reason).await; } - - fn clear_challenge_cookies_after_cooldown(&self, account_id: &str, delay_seconds: i64) { - let cookie_store = self.cookie_store.clone(); - let account_id = account_id.to_string(); - let delay = StdDuration::from_secs(delay_seconds.max(0).cast_unsigned()); - - tokio::spawn(async move { - tokio::time::sleep(delay).await; - delete_account_cookies(&cookie_store, &account_id, "Cloudflare challenge").await; - }); - } } async fn delete_account_cookies(cookie_store: &PgCookieStore, account_id: &str, reason: &str) { diff --git a/backend/src/dispatch/service.rs b/backend/src/dispatch/service.rs index 0fb945e56..d4bce1e14 100644 --- a/backend/src/dispatch/service.rs +++ b/backend/src/dispatch/service.rs @@ -303,7 +303,7 @@ impl ResponseDispatchService { &acquired.account.id, ); self.account_pool.wait_for_request_interval(&acquired).await; - let account = acquired.account; + let account = acquired.account.clone(); let release_account_id = account.id.clone(); let attempt = trace.start_attempt(&release_account_id); let response_result = create_response_with_account_retrying_5xx( @@ -316,7 +316,7 @@ impl ResponseDispatchService { started_at, ) .await; - self.account_pool.release(&release_account_id).await; + acquired.complete().await; if let Err(error) = &response_result { self.cloudflare .capture_set_cookie_headers( diff --git a/backend/src/dispatch/stream/lifecycle.rs b/backend/src/dispatch/stream/lifecycle.rs index c3acec4d0..99ac373a9 100644 --- a/backend/src/dispatch/stream/lifecycle.rs +++ b/backend/src/dispatch/stream/lifecycle.rs @@ -197,7 +197,7 @@ impl ResponseDispatchService { ); self.account_pool.wait_for_request_interval(&acquired).await; - let account = acquired.account; + let account = acquired.account.clone(); let release_account_id = account.id.clone(); let attempt = trace.start_attempt(&release_account_id); let response_result = create_response_stream_with_account_retrying_5xx( @@ -239,7 +239,7 @@ impl ResponseDispatchService { Err(ResponseDispatchError::Upstream(error)) if is_rate_limit_upstream_error(&error) => { - self.account_pool.release(&release_account_id).await; + acquired.complete().await; exhausted_accounts.record_rate_limited( Some(&release_account_id), upstream_error_body(&error), @@ -254,7 +254,7 @@ impl ResponseDispatchService { Err(ResponseDispatchError::Upstream(error)) if is_quota_exhausted_upstream_error(&error) => { - self.account_pool.release(&release_account_id).await; + acquired.complete().await; exhausted_accounts.record_quota_exhausted( Some(&release_account_id), upstream_error_body(&error), @@ -268,7 +268,7 @@ impl ResponseDispatchService { Err(ResponseDispatchError::Upstream(error)) if is_auth_upstream_error(&error) => { - self.account_pool.release(&release_account_id).await; + acquired.complete().await; let upstream_error = upstream_error_body(&error); let account_status = auth_failure_account_status(&error); exhausted_accounts.record_auth_failure( @@ -286,7 +286,7 @@ impl ResponseDispatchService { Err(ResponseDispatchError::Upstream(error)) if is_cloudflare_challenge_upstream_error(&error) => { - self.account_pool.release(&release_account_id).await; + acquired.complete().await; exhausted_accounts.record_cloudflare_challenge( Some(&release_account_id), cloudflare_challenge_error_message(), @@ -300,7 +300,7 @@ impl ResponseDispatchService { Err(ResponseDispatchError::Upstream(error)) if is_cloudflare_path_block_upstream_error(&error) => { - self.account_pool.release(&release_account_id).await; + acquired.complete().await; exhausted_accounts.record_cloudflare_path_blocked( Some(&release_account_id), cloudflare_path_block_error_message(), @@ -314,7 +314,7 @@ impl ResponseDispatchService { Err(ResponseDispatchError::Upstream(error)) if is_model_unsupported_upstream_error(&error) => { - self.account_pool.release(&release_account_id).await; + acquired.complete().await; let upstream_error = upstream_error_body(&error); if let Some(exhausted) = exhausted_accounts .model_unsupported_retry_exhausted(upstream_error.clone()) @@ -335,7 +335,7 @@ impl ResponseDispatchService { Err(ResponseDispatchError::Upstream(error)) if is_banned_upstream_error(&error) => { - self.account_pool.release(&release_account_id).await; + acquired.complete().await; exhausted_accounts.record_auth_failure( Some(&release_account_id), AccountStatus::Banned, @@ -349,7 +349,7 @@ impl ResponseDispatchService { continue; } Err(error) => { - self.account_pool.release(&release_account_id).await; + acquired.complete().await; if let ResponseDispatchError::Upstream(upstream_error) = &error { let history_unavailable = is_history_recovery_upstream_error(upstream_error); @@ -401,7 +401,7 @@ impl ResponseDispatchService { let first_failure = match first_sse_failure(&prefetched) { Ok(failure) => failure, Err(error) => { - self.account_pool.release(&release_account_id).await; + acquired.complete().await; return_stream_dispatch_error!( ResponseDispatchError::InvalidSse(error), account_id: Some(&release_account_id), @@ -420,7 +420,7 @@ impl ResponseDispatchService { ) .await { - self.account_pool.release(&release_account_id).await; + acquired.complete().await; next_required_account_id = Some(release_account_id); continue; } @@ -429,7 +429,7 @@ impl ResponseDispatchService { if let Some(exhausted) = exhausted_accounts .model_unsupported_retry_exhausted(upstream_error.clone()) { - self.account_pool.release(&release_account_id).await; + acquired.complete().await; return_stream_dispatch_error!( ResponseDispatchError::from_exhausted_account(exhausted), account_id: Some(&release_account_id), @@ -441,7 +441,7 @@ impl ResponseDispatchService { upstream_error, ); excluded_account_ids.push(release_account_id); - self.account_pool.release(&account.id).await; + acquired.complete().await; continue; } if is_quota_exhausted_sse_failure(&failure) { @@ -453,7 +453,7 @@ impl ResponseDispatchService { .set_status(&release_account_id, AccountStatus::QuotaExhausted) .await; excluded_account_ids.push(release_account_id); - self.account_pool.release(&account.id).await; + acquired.complete().await; continue; } if is_auth_sse_failure(&failure) { @@ -469,10 +469,10 @@ impl ResponseDispatchService { .set_status(&release_account_id, account_status) .await; excluded_account_ids.push(release_account_id); - self.account_pool.release(&account.id).await; + acquired.complete().await; continue; } - self.account_pool.release(&release_account_id).await; + acquired.complete().await; record_prefetched_response_stream_failure_event( ResponseStreamFailureEventRecord { recorder: &self.recorder, @@ -505,6 +505,7 @@ impl ResponseDispatchService { let context = LiveResponseStreamContext { account_pool: Arc::clone(&self.account_pool), + account_lease: acquired, session_affinity: Arc::clone(&self.session_affinity), reasoning_replay: Arc::clone(&self.reasoning_replay), recorder: Arc::clone(&self.recorder), @@ -533,7 +534,7 @@ impl ResponseDispatchService { return Ok(spawn_live_response_stream(context, prefetched, body)); } Err(error) if is_rate_limit_upstream_error(&error) => { - self.account_pool.release(&release_account_id).await; + acquired.complete().await; exhausted_accounts.record_rate_limited( Some(&release_account_id), upstream_error_body(&error), @@ -545,7 +546,7 @@ impl ResponseDispatchService { excluded_account_ids.push(release_account_id); } Err(error) if is_quota_exhausted_upstream_error(&error) => { - self.account_pool.release(&release_account_id).await; + acquired.complete().await; exhausted_accounts.record_quota_exhausted( Some(&release_account_id), upstream_error_body(&error), @@ -556,7 +557,7 @@ impl ResponseDispatchService { excluded_account_ids.push(release_account_id); } Err(error) if is_auth_upstream_error(&error) => { - self.account_pool.release(&release_account_id).await; + acquired.complete().await; let upstream_error = upstream_error_body(&error); let account_status = auth_failure_account_status(&error); exhausted_accounts.record_auth_failure( @@ -571,7 +572,7 @@ impl ResponseDispatchService { excluded_account_ids.push(release_account_id); } Err(error) if is_cloudflare_challenge_upstream_error(&error) => { - self.account_pool.release(&release_account_id).await; + acquired.complete().await; exhausted_accounts.record_cloudflare_challenge( Some(&release_account_id), cloudflare_challenge_error_message(), @@ -582,7 +583,7 @@ impl ResponseDispatchService { excluded_account_ids.push(release_account_id); } Err(error) if is_cloudflare_path_block_upstream_error(&error) => { - self.account_pool.release(&release_account_id).await; + acquired.complete().await; exhausted_accounts.record_cloudflare_path_blocked( Some(&release_account_id), cloudflare_path_block_error_message(), @@ -593,7 +594,7 @@ impl ResponseDispatchService { excluded_account_ids.push(release_account_id); } Err(error) if is_model_unsupported_upstream_error(&error) => { - self.account_pool.release(&release_account_id).await; + acquired.complete().await; let upstream_error = upstream_error_body(&error); if let Some(exhausted) = exhausted_accounts.model_unsupported_retry_exhausted(upstream_error.clone()) @@ -611,7 +612,7 @@ impl ResponseDispatchService { excluded_account_ids.push(release_account_id); } Err(error) if is_banned_upstream_error(&error) => { - self.account_pool.release(&release_account_id).await; + acquired.complete().await; exhausted_accounts.record_auth_failure( Some(&release_account_id), AccountStatus::Banned, @@ -624,7 +625,7 @@ impl ResponseDispatchService { excluded_account_ids.push(release_account_id); } Err(error) => { - self.account_pool.release(&release_account_id).await; + acquired.complete().await; let history_unavailable = is_history_recovery_upstream_error(&error); if history_unavailable && self diff --git a/backend/src/dispatch/stream/live.rs b/backend/src/dispatch/stream/live.rs index e8cf4e0a1..91d35b108 100644 --- a/backend/src/dispatch/stream/live.rs +++ b/backend/src/dispatch/stream/live.rs @@ -24,7 +24,7 @@ use crate::{ recovery::{cloudflare::CloudflareRecovery, reasoning_replay::ReasoningReplayCache}, service::ResponseDispatchStream, }, - fleet::pool::AccountPoolService, + fleet::pool::{AccountLease, AccountPoolService}, infra::time::elapsed_millis_i64, telemetry::{recorder::Recorder, usage::types::UsageRecordLevel}, upstream::openai::{ @@ -99,7 +99,7 @@ pub(in crate::dispatch) fn spawn_live_response_stream( ) .await { - context.account_pool.release(&context.account_id).await; + context.account_lease.complete().await; return; } update_first_response_output_ms(context.started_at, &body_bytes, &mut first_token_ms); @@ -107,7 +107,7 @@ pub(in crate::dispatch) fn spawn_live_response_stream( loop { let next = tokio::select! { _ = &mut cancel_receiver => { - context.account_pool.release(&context.account_id).await; + context.account_lease.complete().await; return; } next = body.next() => next, @@ -125,7 +125,7 @@ pub(in crate::dispatch) fn spawn_live_response_stream( ) .await { - context.account_pool.release(&context.account_id).await; + context.account_lease.complete().await; return; } update_first_response_output_ms( @@ -142,7 +142,7 @@ pub(in crate::dispatch) fn spawn_live_response_stream( ) .await { - context.account_pool.release(&context.account_id).await; + context.account_lease.complete().await; return; } let detail = error.to_string(); @@ -150,7 +150,7 @@ pub(in crate::dispatch) fn spawn_live_response_stream( send_live_response_stream_tail(&sender, &mut body_bytes, Some(&detail)) .await else { - context.account_pool.release(&context.account_id).await; + context.account_lease.complete().await; return; }; finalize_live_response_stream(context, body_text, first_token_ms).await; @@ -166,12 +166,12 @@ pub(in crate::dispatch) fn spawn_live_response_stream( ) .await { - context.account_pool.release(&context.account_id).await; + context.account_lease.complete().await; return; } let Some(body_text) = send_live_response_stream_tail(&sender, &mut body_bytes, None).await else { - context.account_pool.release(&context.account_id).await; + context.account_lease.complete().await; return; }; @@ -340,6 +340,7 @@ fn missing_sse_event_separator(body: &str) -> Option<&'static str> { pub(in crate::dispatch) struct LiveResponseStreamContext { pub(in crate::dispatch) account_pool: Arc, + pub(in crate::dispatch) account_lease: AccountLease, pub(in crate::dispatch) session_affinity: Arc, pub(in crate::dispatch) reasoning_replay: Arc>, pub(in crate::dispatch) recorder: Arc, @@ -552,5 +553,5 @@ pub(in crate::dispatch) async fn finalize_live_response_stream( } } - context.account_pool.release(&context.account_id).await; + context.account_lease.complete().await; } diff --git a/backend/src/dispatch/upstream_call.rs b/backend/src/dispatch/upstream_call.rs index 77ed2a8fb..f2d202146 100644 --- a/backend/src/dispatch/upstream_call.rs +++ b/backend/src/dispatch/upstream_call.rs @@ -33,7 +33,7 @@ use crate::{ }, fleet::{ account::{Account, AccountStatus}, - pool::{AccountAcquireRequest, AccountPoolService, AcquiredAccount}, + pool::{AccountAcquireRequest, AccountLease, AccountPoolService}, quota::{ quota_from_usage, quota_snapshot_limit_reached, quota_snapshot_limit_window_seconds, quota_snapshot_reset_at, @@ -62,7 +62,7 @@ pub(crate) const QUOTA_VERIFY_LIMIT_REACHED_MESSAGE: &str = "Upstream usage quota still reports limit_reached"; pub(crate) enum QuotaVerificationDecision { - Ready(Box), + Ready(Box), RetryWithAnotherAccount, RequiredAccountUnavailable, MaxAttemptsReached, @@ -81,7 +81,7 @@ pub(crate) struct QuotaVerificationContext<'a> { pub(crate) async fn verify_acquired_quota_if_required( context: QuotaVerificationContext<'_>, - acquired: AcquiredAccount, + acquired: AccountLease, ) -> QuotaVerificationDecision { if !acquired.account.quota_verify_required { return QuotaVerificationDecision::Ready(Box::new(acquired)); @@ -133,10 +133,7 @@ pub(crate) async fn verify_acquired_quota_if_required( .apply_quota_snapshot(&account_id, "a) .await; if quota_snapshot_limit_reached("a) { - context - .account_pool - .release_without_request_usage(&account_id) - .await; + acquired.release_without_usage().await; context.excluded_account_ids.push(account_id.clone()); *context.verify_attempts += 1; let max_attempts_reached = *context.verify_attempts >= MAX_QUOTA_VERIFY_ATTEMPTS; @@ -161,7 +158,7 @@ pub(crate) async fn verify_acquired_quota_if_required( QuotaVerificationDecision::Ready(Box::new(acquired_with_verified_quota(acquired, "a))) } -fn acquired_with_verified_quota(mut acquired: AcquiredAccount, quota: &Value) -> AcquiredAccount { +fn acquired_with_verified_quota(mut acquired: AccountLease, quota: &Value) -> AccountLease { let limit_reached = quota_snapshot_limit_reached(quota); acquired.account.quota_verify_required = false; acquired.account.quota_limit_reached = limit_reached; @@ -564,7 +561,7 @@ impl ResponseDispatchService { return Err(error); } }; - let account = acquired.account; + let account = acquired.account.clone(); let release_account_id = account.id.clone(); let attempt = trace.start_attempt(&release_account_id); let response_result = create_compact_response_with_account_retrying_5xx( @@ -576,7 +573,7 @@ impl ResponseDispatchService { &account, ) .await; - self.account_pool.release(&release_account_id).await; + acquired.complete().await; if let Err(error) = &response_result { self.cloudflare .capture_set_cookie_headers( diff --git a/backend/src/fleet/cookies/store.rs b/backend/src/fleet/cookies/store.rs index 02f569fdc..0bab89310 100644 --- a/backend/src/fleet/cookies/store.rs +++ b/backend/src/fleet/cookies/store.rs @@ -118,6 +118,23 @@ impl PgCookieStore { Ok(result.rows_affected()) } + /// 将账号现有 Cookie 的过期时间收紧到指定时间点。 + pub async fn expire_account_cookies_at( + &self, + account_id: &str, + expires_at: DateTime, + ) -> PgCookieStoreResult { + let result = sqlx::query( + "update account_cookies set expires_at = least(coalesce(expires_at, $1), $1), updated_at = $2 where account_id = $3", + ) + .bind(expires_at) + .bind(Utc::now()) + .bind(account_id) + .execute(&self.pool) + .await?; + Ok(result.rows_affected()) + } + /// 删除指定时间之前过期的 Cookie。 pub async fn cleanup_expired(&self, now: DateTime) -> PgCookieStoreResult { let result = sqlx::query( diff --git a/backend/src/fleet/pool/filters.rs b/backend/src/fleet/pool/filters.rs index 1cc17433d..3d1c7ad7b 100644 --- a/backend/src/fleet/pool/filters.rs +++ b/backend/src/fleet/pool/filters.rs @@ -90,6 +90,15 @@ pub struct AcquiredAccount { pub account: Account, /// 同一账号上一个在途槽位的创建时间。 pub previous_slot_at: Option>, + pub(super) slot_id: uuid::Uuid, +} + +impl std::ops::Deref for AcquiredAccount { + type Target = Account; + + fn deref(&self) -> &Self::Target { + &self.account + } } /// 模型刷新时每个订阅计划选中的账号。 @@ -99,4 +108,5 @@ pub struct DistinctPlanAccount { pub plan_type: String, /// 被选中的账号快照。 pub account: Account, + pub(super) slot_id: uuid::Uuid, } diff --git a/backend/src/fleet/pool/mod.rs b/backend/src/fleet/pool/mod.rs index 3f9b41ccd..cbb3d5300 100644 --- a/backend/src/fleet/pool/mod.rs +++ b/backend/src/fleet/pool/mod.rs @@ -2,6 +2,8 @@ use std::{ collections::{BTreeMap, BTreeSet, VecDeque}, + fmt, + ops::{Deref, DerefMut}, sync::Arc, time::Duration as StdDuration, }; @@ -71,6 +73,108 @@ pub struct AccountPoolService { request_interval: Arc>, } +/// 一次请求对账号池 slot 的唯一所有权凭据。 +pub struct AccountLease { + acquired: AcquiredAccount, + slot: AccountSlotLease, +} + +impl AccountLease { + fn new(service: AccountPoolService, acquired: AcquiredAccount) -> Self { + let slot = AccountSlotLease::new(service, acquired.account.id.clone(), acquired.slot_id); + Self { acquired, slot } + } + + /// 完成已经发往上游的请求并记录请求用量。 + pub async fn complete(self) { + self.slot.release(true).await; + } + + /// 释放尚未用于实际上游请求的 slot。 + pub async fn release_without_usage(self) { + self.slot.release(false).await; + } +} + +impl Deref for AccountLease { + type Target = AcquiredAccount; + + fn deref(&self) -> &Self::Target { + &self.acquired + } +} + +impl DerefMut for AccountLease { + fn deref_mut(&mut self) -> &mut Self::Target { + &mut self.acquired + } +} + +impl fmt::Debug for AccountLease { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter + .debug_struct("AccountLease") + .field("acquired", &self.acquired) + .finish_non_exhaustive() + } +} + +/// 模型目录刷新使用的订阅计划账号 lease。 +pub struct DistinctPlanAccountLease { + pub plan_type: String, + pub account: Account, + slot: AccountSlotLease, +} + +impl DistinctPlanAccountLease { + /// 释放模型目录刷新占用的 slot，不计入客户端请求用量。 + pub async fn release(self) { + self.slot.release(false).await; + } +} + +struct AccountSlotLease { + service: AccountPoolService, + account_id: String, + slot_id: uuid::Uuid, + armed: bool, +} + +impl AccountSlotLease { + fn new(service: AccountPoolService, account_id: String, slot_id: uuid::Uuid) -> Self { + Self { + service, + account_id, + slot_id, + armed: true, + } + } + + async fn release(mut self, record_request_usage: bool) { + self.service + .release_slot(&self.account_id, self.slot_id, record_request_usage) + .await; + self.armed = false; + } +} + +impl Drop for AccountSlotLease { + fn drop(&mut self) { + if !self.armed { + return; + } + let Ok(runtime) = tokio::runtime::Handle::try_current() else { + return; + }; + let service = self.service.clone(); + let account_id = self.account_id.clone(); + let slot_id = self.slot_id; + runtime.spawn(async move { + service.release_slot(&account_id, slot_id, false).await; + }); + } +} + impl AccountPoolService { /// 构造运行时账号池服务。 pub fn new( @@ -176,24 +280,22 @@ impl AccountPoolService { self.capacity_summary(Utc::now()).await } - /// 释放账号在途槽位。 - pub async fn release(&self, account_id: &str) { - self.release_slot(account_id, true).await; - } - - /// 释放账号在途槽位，不累计请求用量。 - pub async fn release_without_request_usage(&self, account_id: &str) { - self.release_slot(account_id, false).await; - } - - async fn release_slot(&self, account_id: &str, record_request_usage: bool) { + async fn release_slot( + &self, + account_id: &str, + slot_id: uuid::Uuid, + record_request_usage: bool, + ) { let released = if record_request_usage { - self.pool.lock().await.release(account_id) + self.pool + .lock() + .await + .release_slot(account_id, slot_id, true) } else { self.pool .lock() .await - .release_without_request_usage(account_id) + .release_slot(account_id, slot_id, false) }; let Some(released) = released else { return; @@ -275,23 +377,42 @@ impl AccountPoolService { } /// 按请求上下文获取可用账号。 - pub async fn acquire_with(&self, request: &AccountAcquireRequest) -> Option { + pub async fn acquire_with(&self, request: &AccountAcquireRequest) -> Option { let refresh = self.pool.lock().await.acquire_with_status_refresh(request); + let lease = refresh + .acquired + .map(|acquired| AccountLease::new(self.clone(), acquired)); self.persist_runtime_account_states(refresh.refreshed_accounts) .await; - refresh.acquired + lease } /// 按订阅计划各选一个可用账号，用于刷新模型列表。 - pub async fn distinct_plan_accounts(&self, now: DateTime) -> Vec { + pub async fn distinct_plan_accounts( + &self, + now: DateTime, + ) -> Vec { let refresh = self .pool .lock() .await .distinct_plan_accounts_with_status_refresh(now); + let leases = refresh + .accounts + .into_iter() + .map(|selected| DistinctPlanAccountLease { + slot: AccountSlotLease::new( + self.clone(), + selected.account.id.clone(), + selected.slot_id, + ), + plan_type: selected.plan_type, + account: selected.account, + }) + .collect(); self.persist_runtime_account_states(refresh.refreshed_accounts) .await; - refresh.accounts + leases } async fn persist_runtime_account_states(&self, refreshed_states: Vec) { diff --git a/backend/src/fleet/pool/state.rs b/backend/src/fleet/pool/state.rs index d09ce432f..d9781894e 100644 --- a/backend/src/fleet/pool/state.rs +++ b/backend/src/fleet/pool/state.rs @@ -132,6 +132,7 @@ pub struct AccountPool { #[derive(Debug, Clone, PartialEq, Eq)] struct AccountSlot { + id: uuid::Uuid, created_at: DateTime, model: Option, } @@ -360,9 +361,9 @@ impl AccountPool { // 亲和优先：required 硬锁定该账号；preferred 若在候选内则直接选中； // 二者都不命中时，回退到调度器按当前策略选择。 let selected = if let Some(required_account_id) = &request.required_account_id { - candidates.first().cloned().map(|account| { + candidates.first().map(|account| { ( - account, + (*account).clone(), self.previous_slot_at(required_account_id), "required", ) @@ -371,10 +372,9 @@ impl AccountPool { candidates .iter() .find(|account| account.id == *preferred_account_id) - .cloned() .map(|account| { ( - account, + (*account).clone(), self.previous_slot_at(preferred_account_id), "preferred", ) @@ -411,11 +411,12 @@ impl AccountPool { ); let selected_id = selected.id.clone(); - self.push_slot(&selected_id, request.now, Some(&request.model)); + let slot_id = self.push_slot(&selected_id, request.now, Some(&request.model)); AccountAcquireWithStatusRefresh { acquired: Some(AcquiredAccount { account: selected, previous_slot_at, + slot_id, }), refreshed_accounts, } @@ -449,28 +450,32 @@ impl AccountPool { } /// 释放指定账号的一个在途槽位。 - pub fn release(&mut self, account_id: &str) -> Option { - self.release_slot(account_id, true) + pub fn release(&mut self, acquired: &AcquiredAccount) -> Option { + self.release_slot(&acquired.account.id, acquired.slot_id, true) } - /// 释放指定账号的一个在途槽位，不累计请求用量。 - pub fn release_without_request_usage( + /// 释放模型目录刷新选择的账号，不累计客户端请求用量。 + pub fn release_distinct_plan_account( &mut self, - account_id: &str, + selected: &DistinctPlanAccount, ) -> Option { - self.release_slot(account_id, false) + self.release_slot(&selected.account.id, selected.slot_id, false) } - fn release_slot( + pub(super) fn release_slot( &mut self, account_id: &str, + slot_id: uuid::Uuid, record_request_usage: bool, ) -> Option { self.accounts.get(account_id)?; let mut remove_slots = false; let slot = self.slots.get_mut(account_id).and_then(|slots| { - let slot = slots.pop_front(); + let slot = slots + .iter() + .position(|slot| slot.id == slot_id) + .and_then(|position| slots.remove(position)); remove_slots = slots.is_empty(); slot }); @@ -540,7 +545,7 @@ impl AccountPool { ) -> DistinctPlanAccountsWithStatusRefresh { self.cleanup_stale_slots(now); let refreshed_accounts = self.refresh_account_statuses(now); - let mut by_plan = IndexMap::>::new(); + let mut by_plan = IndexMap::>::new(); let max_concurrent = self.options.max_concurrent_per_account; let skip_quota_limited = self.options.skip_quota_limited; @@ -559,15 +564,12 @@ impl AccountPool { let Some(plan_type) = account.plan_type.as_ref() else { continue; }; - by_plan - .entry(plan_type.clone()) - .or_default() - .push(account.clone()); + by_plan.entry(plan_type.clone()).or_default().push(account); } } let strategy = self.options.rotation_strategy; - let mut accounts = Vec::new(); + let mut selected_accounts = Vec::new(); for (plan_type, group) in by_plan { let Some(account) = self.select_candidate(&group, now) else { continue; @@ -586,8 +588,17 @@ impl AccountPool { quota_cooldown_until = ?account.quota_cooldown_until, "account selected for model refresh" ); - self.push_slot(&account.id, now, None); - accounts.push(DistinctPlanAccount { plan_type, account }); + selected_accounts.push((plan_type, account)); + } + + let mut accounts = Vec::with_capacity(selected_accounts.len()); + for (plan_type, account) in selected_accounts { + let slot_id = self.push_slot(&account.id, now, None); + accounts.push(DistinctPlanAccount { + plan_type, + account, + slot_id, + }); } DistinctPlanAccountsWithStatusRefresh { @@ -607,14 +618,22 @@ impl AccountPool { .map(|slot| slot.created_at) } - fn push_slot(&mut self, account_id: &str, now: DateTime, model: Option<&str>) { + fn push_slot( + &mut self, + account_id: &str, + now: DateTime, + model: Option<&str>, + ) -> uuid::Uuid { + let slot_id = uuid::Uuid::new_v4(); self.slots .entry(account_id.to_string()) .or_default() .push_back(AccountSlot { + id: slot_id, created_at: now, model: model.map(str::to_string), }); + slot_id } fn mark_request_usage(&mut self, account_id: &str, now: DateTime) -> Option { @@ -667,7 +686,7 @@ impl AccountPool { } /// 过滤出本次请求的可用候选集合（委托 [`candidates::filter`]）。 - fn filter_candidates(&self, request: &AccountAcquireRequest) -> Vec { + fn filter_candidates(&self, request: &AccountAcquireRequest) -> Vec<&Account> { let slot_count = |account_id: &str| self.slots.get(account_id).map_or(0, VecDeque::len); candidates::filter( self.accounts.values(), @@ -689,7 +708,7 @@ impl AccountPool { } /// 按当前策略从候选集合中选择一个账号（委托 [`AccountScheduler::select`]）。 - fn select_candidate(&self, candidates: &[Account], now: DateTime) -> Option { + fn select_candidate(&self, candidates: &[&Account], now: DateTime) -> Option { let slot_count = |account_id: &str| self.slots.get(account_id).map_or(0, VecDeque::len); self.scheduler .select(self.options.rotation_strategy, candidates, &slot_count, now) @@ -697,16 +716,13 @@ impl AccountPool { } fn access_token_expired(account: &Account, now: DateTime) -> bool { - if account - .access_token_expires_at - .is_some_and(|expires_at| now >= expires_at) - { - return true; + if let Some(expires_at) = account.access_token_expires_at { + return now >= expires_at; } match jwt_expiry(&account.access_token, now) { JwtExpiry::Expired => true, JwtExpiry::Valid => false, - JwtExpiry::MissingOrInvalid => account.access_token_expires_at.is_none(), + JwtExpiry::MissingOrInvalid => true, } } diff --git a/backend/src/fleet/refresh/policy.rs b/backend/src/fleet/refresh/policy.rs index cc699b63b..832f9bea1 100644 --- a/backend/src/fleet/refresh/policy.rs +++ b/backend/src/fleet/refresh/policy.rs @@ -164,21 +164,6 @@ pub struct RefreshPolicy { pub refresh_concurrency: u32, } -/// 单次刷新扫描摘要。 -#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] -pub struct TokenRefreshSummary { - /// 扫描账号数。 - pub scanned: usize, - /// 成功刷新 token 的账号数。 - pub refreshed: usize, - /// 仅更新状态的账号数。 - pub status_updated: usize, - /// 无需刷新的账号数。 - pub skipped: usize, - /// 传输失败账号数。 - pub failed: usize, -} - /// 账号级刷新定时器调度摘要。 #[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] pub struct TokenTimerSummary { @@ -229,15 +214,22 @@ impl RuntimeRefreshPolicy { /// 更新当前策略。 pub fn update(&self, policy: RefreshPolicy) { - *self - .policy - .write() - .unwrap_or_else(std::sync::PoisonError::into_inner) = policy; - *self - .semaphore - .write() - .unwrap_or_else(std::sync::PoisonError::into_inner) = - Arc::new(Semaphore::new(policy.refresh_concurrency.max(1) as usize)); + let previous = { + let mut current = self + .policy + .write() + .unwrap_or_else(std::sync::PoisonError::into_inner); + let previous = *current; + *current = policy; + previous + }; + if previous.refresh_concurrency != policy.refresh_concurrency { + *self + .semaphore + .write() + .unwrap_or_else(std::sync::PoisonError::into_inner) = + Arc::new(Semaphore::new(policy.refresh_concurrency.max(1) as usize)); + } } /// 返回当前提前刷新秒数。 diff --git a/backend/src/fleet/refresh/service.rs b/backend/src/fleet/refresh/service.rs index 3e432705a..f2692aaa7 100644 --- a/backend/src/fleet/refresh/service.rs +++ b/backend/src/fleet/refresh/service.rs @@ -4,7 +4,10 @@ use std::{ collections::{HashMap, HashSet}, future::Future, pin::Pin, - sync::Arc, + sync::{ + atomic::{AtomicBool, Ordering}, + Arc, Mutex as StdMutex, + }, time::Duration, }; @@ -12,6 +15,7 @@ use chrono::{DateTime, Utc}; use secrecy::{ExposeSecret, SecretString}; use thiserror::Error; use tokio::{sync::Mutex, task::JoinHandle, time::sleep}; +use tokio_util::task::TaskTracker; use tracing::{debug, info, warn}; use uuid::Uuid; @@ -29,13 +33,14 @@ use super::{ default_refresh_retry_delays, duration_until, is_permanent_refresh_failure_status, jittered_refresh_at, jwt_expiration, same_refresh_time, scheduled_at_from_delay, stable_jittered_duration, token_refresh_status_eligible, RefreshError, RefreshPolicy, - RefreshScheduler, RuntimeRefreshPolicy, TokenRefreshSummary, TokenTimerSummary, + RefreshScheduler, RuntimeRefreshPolicy, TokenTimerSummary, PERMANENT_FAILURE_CONFIRMATION_THRESHOLD, RECOVERY_DELAY_JITTER, RECOVERY_DELAY_SECONDS, RETRY_DELAY_JITTER, }, }; const REFRESH_LEASE_TTL_SECONDS: i64 = 5 * 60; +const REFRESH_TASK_SHUTDOWN_TIMEOUT: Duration = Duration::from_secs(4); /// 运行时 token refresh 服务。 pub struct TokenRefreshService @@ -48,11 +53,14 @@ where refresh_leases: Option, lease_owner: String, retry_delays: Vec, - in_flight: Arc>>, + in_flight: Arc>>, timers: Arc>>, + tasks: TaskTracker, + shutting_down: Arc, } struct ScheduledRefreshTimer { + id: Uuid, scheduled_at: DateTime, handle: JoinHandle<()>, } @@ -62,6 +70,90 @@ enum ScheduleTimerResult { Scheduled, Replaced, Unchanged, + ShuttingDown, +} + +struct InFlightRefreshGuard { + in_flight: Arc>>, + account_id: String, +} + +impl Drop for InFlightRefreshGuard { + fn drop(&mut self) { + self.in_flight + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .remove(&self.account_id); + } +} + +struct RefreshLeaseGuard { + store: Option, + tasks: Option, + account_id: String, + owner: String, + armed: bool, +} + +impl RefreshLeaseGuard { + fn local(account_id: &str, owner: &str) -> Self { + Self { + store: None, + tasks: None, + account_id: account_id.to_string(), + owner: owner.to_string(), + armed: false, + } + } + + fn distributed( + store: RedisRefreshLeaseStore, + tasks: TaskTracker, + account_id: &str, + owner: &str, + ) -> Self { + Self { + store: Some(store), + tasks: Some(tasks), + account_id: account_id.to_string(), + owner: owner.to_string(), + armed: true, + } + } + + async fn release(mut self) -> Result<(), RedisRefreshLeaseStoreError> { + let Some(store) = self.store.as_ref() else { + return Ok(()); + }; + store.release(&self.account_id, &self.owner).await?; + self.armed = false; + Ok(()) + } +} + +impl Drop for RefreshLeaseGuard { + fn drop(&mut self) { + if !self.armed { + return; + } + let (Some(store), Some(tasks), Ok(runtime)) = ( + self.store.clone(), + self.tasks.clone(), + tokio::runtime::Handle::try_current(), + ) else { + return; + }; + let account_id = self.account_id.clone(); + let owner = self.owner.clone(); + drop(tasks.spawn_on( + async move { + if let Err(error) = store.release(&account_id, &owner).await { + warn!(account_id, error = %error, "failed to release cancelled token refresh lease"); + } + }, + &runtime, + )); + } } enum TokenRefreshOutcome { @@ -108,6 +200,8 @@ where retry_delays: self.retry_delays.clone(), in_flight: self.in_flight.clone(), timers: self.timers.clone(), + tasks: self.tasks.clone(), + shutting_down: self.shutting_down.clone(), } } } @@ -126,8 +220,10 @@ where refresh_leases: None, lease_owner: refresh_lease_owner(), retry_delays: default_refresh_retry_delays(), - in_flight: Arc::new(Mutex::new(HashSet::new())), + in_flight: Arc::new(StdMutex::new(HashSet::new())), timers: Arc::new(Mutex::new(HashMap::new())), + tasks: TaskTracker::new(), + shutting_down: Arc::new(AtomicBool::new(false)), } } @@ -168,7 +264,7 @@ where let mut refreshing = self .in_flight .lock() - .await + .unwrap_or_else(std::sync::PoisonError::into_inner) .iter() .filter(|account_id| requested_ids.contains(account_id.as_str())) .cloned() @@ -198,6 +294,9 @@ where &self, now: DateTime, ) -> TokenRefreshServiceResult { + if self.shutting_down.load(Ordering::Acquire) { + return Ok(TokenTimerSummary::default()); + } let accounts = self .store .list_pool_accounts() @@ -217,19 +316,6 @@ where summary.skipped += 1; continue; }; - if delay.is_zero() { - if self.clear_scheduled_timer(&account.id).await { - summary.replaced += 1; - } - let outcome = self.refresh_scheduled_account(&account.id, now).await?; - if matches!(outcome, TokenRefreshOutcome::Skipped) { - summary.skipped += 1; - } else { - summary.immediate += 1; - } - continue; - } - let scheduled_at = scheduled_at_from_delay(now, delay); self.persist_next_refresh_at_if_changed(&account, Some(scheduled_at)) .await?; @@ -243,6 +329,10 @@ where } ScheduleTimerResult::Replaced => summary.replaced += 1, ScheduleTimerResult::Scheduled => {} + ScheduleTimerResult::ShuttingDown => { + summary.skipped += 1; + continue; + } } if delay.is_zero() { @@ -260,77 +350,6 @@ where self.timers.lock().await.len() } - /// 在指定时间点执行一次到期刷新扫描。 - pub async fn refresh_due_accounts_once_at( - &self, - now: DateTime, - ) -> TokenRefreshServiceResult { - let accounts = self - .store - .list_pool_accounts() - .await - .map_err(TokenRefreshServiceError::AccountStore)?; - let mut summary = TokenRefreshSummary { - scanned: accounts.len(), - ..TokenRefreshSummary::default() - }; - - for account in accounts { - if self.clear_refresh_schedule_if_ineligible(&account).await? { - summary.skipped += 1; - continue; - } - - if account.next_refresh_at.is_some_and(|value| value > now) { - summary.skipped += 1; - continue; - } - - if !self.scheduler.should_refresh_account_at(&account, now) { - summary.skipped += 1; - continue; - } - - if !self.try_mark_in_flight(&account.id).await { - summary.skipped += 1; - continue; - } - - let lease_acquired = match self.try_acquire_refresh_lease(&account.id, now).await { - Ok(acquired) => acquired, - Err(error) => { - self.release_in_flight(&account.id).await; - return Err(error); - } - }; - if !lease_acquired { - self.release_in_flight(&account.id).await; - summary.skipped += 1; - continue; - } - - let outcome = self - .refresh_account_with_status_transitions(&account, now) - .await; - - self.release_in_flight(&account.id).await; - self.release_refresh_lease(&account.id).await?; - - let outcome = outcome?; - self.schedule_next_refresh_after_outcome(&outcome, Utc::now()) - .await; - - match outcome { - TokenRefreshOutcome::Refreshed(_) => summary.refreshed += 1, - TokenRefreshOutcome::StatusUpdated => summary.status_updated += 1, - TokenRefreshOutcome::Skipped => summary.skipped += 1, - TokenRefreshOutcome::Failed => summary.failed += 1, - }; - } - - Ok(summary) - } - fn timer_delay_for_account(&self, account: &Account, now: DateTime) -> Option { if let Some(next_refresh_at) = account.next_refresh_at.filter(|value| *value > now) { return Some(duration_until(now, next_refresh_at)); @@ -390,24 +409,42 @@ where scheduled_at: DateTime, ) -> Pin + Send + '_>> { Box::pin(async move { - { - let timers = self.timers.lock().await; - if timers - .get(&account_id) - .is_some_and(|timer| timer.scheduled_at == scheduled_at) - { - return ScheduleTimerResult::Unchanged; - } + let mut timers = self.timers.lock().await; + if self.shutting_down.load(Ordering::Acquire) { + return ScheduleTimerResult::ShuttingDown; } - - let replaced = self.clear_scheduled_timer(&account_id).await; + if timers + .get(&account_id) + .is_some_and(|timer| timer.scheduled_at == scheduled_at) + { + return ScheduleTimerResult::Unchanged; + } + let replaced = timers.remove(&account_id).is_some_and(|timer| { + timer.handle.abort(); + true + }); let task = self.clone(); let timer_account_id = account_id.clone(); - let handle = tokio::spawn(async move { + let timer_id = Uuid::new_v4(); + let handle = self.tasks.spawn(async move { if !delay.is_zero() { sleep(delay).await; } - task.timers.lock().await.remove(&timer_account_id); + let owns_timer = { + let mut timers = task.timers.lock().await; + if timers + .get(&timer_account_id) + .is_some_and(|timer| timer.id == timer_id) + { + timers.remove(&timer_account_id); + true + } else { + false + } + }; + if !owns_timer || task.shutting_down.load(Ordering::Acquire) { + return; + } if let Err(error) = task .refresh_scheduled_account(&timer_account_id, scheduled_at) .await @@ -420,9 +457,10 @@ where } }); - self.timers.lock().await.insert( + timers.insert( account_id, ScheduledRefreshTimer { + id: timer_id, scheduled_at, handle, }, @@ -452,6 +490,23 @@ where } } + /// 停止新调度，取消尚未触发的 timer，并等待已进入刷新流程的任务结束。 + pub async fn shutdown(&self) { + self.shutting_down.store(true, Ordering::Release); + self.clear_scheduled_timers().await; + self.tasks.close(); + if tokio::time::timeout(REFRESH_TASK_SHUTDOWN_TIMEOUT, self.tasks.wait()) + .await + .is_err() + { + warn!( + remaining_tasks = self.tasks.len(), + timeout_secs = REFRESH_TASK_SHUTDOWN_TIMEOUT.as_secs(), + "waiting for in-flight token refresh tasks timed out" + ); + } + } + async fn refresh_scheduled_account( &self, account_id: &str, @@ -469,28 +524,19 @@ where return Ok(TokenRefreshOutcome::Skipped); } - if !self.try_mark_in_flight(&account.id).await { + let Some(_in_flight) = self.try_mark_in_flight(&account.id) else { return Ok(TokenRefreshOutcome::Skipped); - } - - let lease_acquired = match self.try_acquire_refresh_lease(&account.id, now).await { - Ok(acquired) => acquired, - Err(error) => { - self.release_in_flight(&account.id).await; - return Err(error); - } }; - if !lease_acquired { - self.release_in_flight(&account.id).await; + + let Some(refresh_lease) = self.try_acquire_refresh_lease(&account.id, now).await? else { return Ok(TokenRefreshOutcome::Skipped); - } + }; let outcome = self .refresh_account_with_status_transitions(&account, now) .await; - self.release_in_flight(&account.id).await; - self.release_refresh_lease(&account.id).await?; + refresh_lease.release().await?; let outcome = outcome?; self.schedule_next_refresh_after_outcome(&outcome, Utc::now()) @@ -638,23 +684,30 @@ where true } - async fn try_mark_in_flight(&self, account_id: &str) -> bool { - self.in_flight.lock().await.insert(account_id.to_string()) - } - - async fn release_in_flight(&self, account_id: &str) { - self.in_flight.lock().await.remove(account_id); + fn try_mark_in_flight(&self, account_id: &str) -> Option { + let inserted = self + .in_flight + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .insert(account_id.to_string()); + inserted.then(|| InFlightRefreshGuard { + in_flight: self.in_flight.clone(), + account_id: account_id.to_string(), + }) } async fn try_acquire_refresh_lease( &self, account_id: &str, now: DateTime, - ) -> TokenRefreshServiceResult { + ) -> TokenRefreshServiceResult> { let Some(refresh_leases) = self.refresh_leases.as_ref() else { - return Ok(true); + return Ok(Some(RefreshLeaseGuard::local( + account_id, + &self.lease_owner, + ))); }; - refresh_leases + let acquired = refresh_leases .try_acquire( account_id, &self.lease_owner, @@ -662,18 +715,15 @@ where now, ) .await - .map_err(TokenRefreshServiceError::Lease) - } - - async fn release_refresh_lease(&self, account_id: &str) -> TokenRefreshServiceResult<()> { - let Some(refresh_leases) = self.refresh_leases.as_ref() else { - return Ok(()); - }; - refresh_leases - .release(account_id, &self.lease_owner) - .await - .map(|_| ()) - .map_err(TokenRefreshServiceError::Lease) + .map_err(TokenRefreshServiceError::Lease)?; + Ok(acquired.then(|| { + RefreshLeaseGuard::distributed( + refresh_leases.clone(), + self.tasks.clone(), + account_id, + &self.lease_owner, + ) + })) } } diff --git a/backend/src/fleet/scheduler/candidates.rs b/backend/src/fleet/scheduler/candidates.rs index b6a469e8d..a431b36e6 100644 --- a/backend/src/fleet/scheduler/candidates.rs +++ b/backend/src/fleet/scheduler/candidates.rs @@ -52,10 +52,9 @@ pub fn filter<'a>( accounts: impl Iterator, filter: &CandidateFilter<'_>, request: &CandidateRequest<'_>, -) -> Vec { +) -> Vec<&'a Account> { let mut candidates = accounts .filter(|account| is_base_available(account, filter, request)) - .cloned() .collect::>(); if let Some(required_account_id) = request.required_account_id { @@ -64,7 +63,7 @@ pub fn filter<'a>( } if let Some(best_tier) = best_available_tier(&candidates, filter.tier_priority) { - candidates.retain(|account| account.plan_type.as_deref() == Some(best_tier.as_str())); + candidates.retain(|account| account.plan_type.as_deref() == Some(best_tier)); } candidates } @@ -109,13 +108,16 @@ fn is_model_allowed( } /// 在候选集合中找出订阅层级优先级最高的可用层级。 -fn best_available_tier(candidates: &[Account], tier_priority: &[String]) -> Option { +fn best_available_tier<'a>( + candidates: &[&Account], + tier_priority: &'a [String], +) -> Option<&'a str> { for tier in tier_priority { let matched = candidates .iter() .any(|account| account.plan_type.as_deref() == Some(tier.as_str())); if matched { - return Some(tier.clone()); + return Some(tier); } } None diff --git a/backend/src/fleet/scheduler/mod.rs b/backend/src/fleet/scheduler/mod.rs index 5dbb37366..24af3fd28 100644 --- a/backend/src/fleet/scheduler/mod.rs +++ b/backend/src/fleet/scheduler/mod.rs @@ -55,7 +55,7 @@ impl RotationStrategy { /// 策略选择的运行时输入。 pub struct SelectionInput<'a> { /// 已过滤的候选账号集合。 - pub candidates: &'a [Account], + pub candidates: &'a [&'a Account], /// 按账号 ID 读取当前在途槽位数。 pub slot_count: &'a dyn Fn(&str) -> usize, /// 账号 EWMA 反馈存储。 @@ -114,7 +114,7 @@ impl AccountScheduler { pub fn select( &self, strategy: RotationStrategy, - candidates: &[Account], + candidates: &[&Account], slot_count: &dyn Fn(&str) -> usize, now: DateTime, ) -> Option { @@ -165,11 +165,11 @@ fn select_with_strategy( } pub(crate) fn select_by( - candidates: &[Account], + candidates: &[&Account], cursor: &mut usize, compare: impl Fn(&Account, &Account) -> Ordering, ) -> Option { - let mut sorted = candidates.iter().collect::>(); + let mut sorted = candidates.to_vec(); sorted.sort_by(|a, b| compare(a, b)); let best = *sorted.first()?; let tied_count = sorted diff --git a/backend/src/fleet/scheduler/strategy/round_robin.rs b/backend/src/fleet/scheduler/strategy/round_robin.rs index 41cf36129..faa403926 100644 --- a/backend/src/fleet/scheduler/strategy/round_robin.rs +++ b/backend/src/fleet/scheduler/strategy/round_robin.rs @@ -15,5 +15,5 @@ pub fn select(input: &SelectionInput<'_>, cursor: &mut usize) -> Option *cursor %= candidates.len(); let index = *cursor; *cursor = cursor.wrapping_add(1); - Some(candidates[index].clone()) + Some((*candidates[index]).clone()) } diff --git a/backend/src/fleet/scheduler/strategy/smart.rs b/backend/src/fleet/scheduler/strategy/smart.rs index f43b8f187..6be1df89b 100644 --- a/backend/src/fleet/scheduler/strategy/smart.rs +++ b/backend/src/fleet/scheduler/strategy/smart.rs @@ -78,7 +78,7 @@ pub struct ScoreBreakdown { /// Smart 策略选择：按加权得分降序排序，在并列最高分之间用轮转游标破并列。 pub fn select(input: &SelectionInput<'_>, cursor: &mut usize) -> Option { - let ranked = rank_candidates( + let ranked = rank_candidate_refs( input.candidates, input.weights, input.now.timestamp(), @@ -94,7 +94,7 @@ pub fn select(input: &SelectionInput<'_>, cursor: &mut usize) -> Option let index = *cursor % tied_count; *cursor = cursor.wrapping_add(1); let account_index = ranked[index].0; - Some(input.candidates[account_index].clone()) + Some((*input.candidates[account_index]).clone()) } /// 对候选集合按加权得分排序（降序），返回 `(账号索引, 打分明细)` 列表。 @@ -106,6 +106,17 @@ pub fn rank_candidates( now_secs: i64, slot_count: &dyn Fn(&str) -> usize, feedback: &FeedbackStats, +) -> Vec<(usize, ScoreBreakdown)> { + let candidates = candidates.iter().collect::>(); + rank_candidate_refs(&candidates, weights, now_secs, slot_count, feedback) +} + +fn rank_candidate_refs( + candidates: &[&Account], + weights: &ScoreWeights, + now_secs: i64, + slot_count: &dyn Fn(&str) -> usize, + feedback: &FeedbackStats, ) -> Vec<(usize, ScoreBreakdown)> { let samples = candidates .iter() @@ -142,7 +153,7 @@ struct NormalizationContext { impl NormalizationContext { fn build( - candidates: &[Account], + candidates: &[&Account], samples: &[FeedbackSample], now_secs: i64, slot_count: &dyn Fn(&str) -> usize, diff --git a/backend/src/fleet/scheduler/strategy/sticky.rs b/backend/src/fleet/scheduler/strategy/sticky.rs index 7ed6301ed..9e1ef3225 100644 --- a/backend/src/fleet/scheduler/strategy/sticky.rs +++ b/backend/src/fleet/scheduler/strategy/sticky.rs @@ -13,7 +13,7 @@ use crate::fleet::scheduler::{compare_last_used, SelectionInput}; /// `_cursor` 未使用：Sticky 由 `last_used_at` 决定顺序，不参与轮转破并列。 pub fn select(input: &SelectionInput<'_>, _cursor: &mut usize) -> Option { let candidates = input.candidates; - let mut selected = candidates.first()?; + let mut selected = *candidates.first()?; for candidate in &candidates[1..] { if compare_last_used( candidate.last_used_at.as_deref(), diff --git a/backend/src/infra/logging.rs b/backend/src/infra/logging.rs index d60935e31..9f364cb2c 100644 --- a/backend/src/infra/logging.rs +++ b/backend/src/infra/logging.rs @@ -508,89 +508,3 @@ fn report_dropped_logs(counters: &mut [NamedErrorCounter], writer: &mut impl Wri let _ = writeln!(writer, "{warning}"); } } - -#[cfg(test)] -mod tests { - use super::*; - - struct FirstWriteBlocking { - started_tx: mpsc::SyncSender<()>, - release_rx: Option>, - } - - impl Write for FirstWriteBlocking { - fn write(&mut self, buffer: &[u8]) -> io::Result { - if let Some(release_rx) = self.release_rx.take() { - self.started_tx - .send(()) - .map_err(|error| io::Error::other(error.to_string()))?; - release_rx - .recv() - .map_err(|error| io::Error::other(error.to_string()))?; - } - Ok(buffer.len()) - } - - fn flush(&mut self) -> io::Result<()> { - Ok(()) - } - } - - #[test] - fn tracing_config_should_reject_disabled_outputs() { - let config = TracingConfig::new("info", false, None); - - assert!(config.validate().is_err()); - } - - #[test] - fn tracing_config_should_accept_each_output_independently() { - let directory = tempfile::tempdir().unwrap(); - let stdout_only = TracingConfig::new("info", true, None); - let file_only = TracingConfig::new( - "info", - false, - Some(RotationConfig::new(directory.path(), 14, 1024, 20)), - ); - - assert_eq!( - (stdout_only.validate().is_ok(), file_only.validate().is_ok()), - (true, true) - ); - } - - #[test] - fn dropped_log_report_should_emit_emergency_json_for_full_buffer() { - let (started_tx, started_rx) = mpsc::sync_channel(0); - let (release_tx, release_rx) = mpsc::sync_channel(0); - let (mut writer, guard) = NonBlockingBuilder::default() - .lossy(true) - .buffered_lines_limit(1) - .finish(FirstWriteBlocking { - started_tx, - release_rx: Some(release_rx), - }); - let counter = writer.error_counter(); - - writer.write_all(b"first\n").unwrap(); - started_rx.recv().unwrap(); - writer.write_all(b"second\n").unwrap(); - writer.write_all(b"dropped\n").unwrap(); - - let mut output = Vec::new(); - report_dropped_logs(&mut [NamedErrorCounter::new("test", counter)], &mut output); - release_tx.send(()).unwrap(); - drop(writer); - drop(guard); - - let warning: serde_json::Value = serde_json::from_slice(&output).unwrap(); - assert_eq!( - ( - warning["level"].as_str(), - warning["fields"]["output"].as_str(), - warning["fields"]["dropped"].as_u64(), - ), - (Some("ERROR"), Some("test"), Some(1)) - ); - } -} diff --git a/backend/src/keys/service.rs b/backend/src/keys/service.rs index 2cc98e4d3..c01db5b7c 100644 --- a/backend/src/keys/service.rs +++ b/backend/src/keys/service.rs @@ -87,10 +87,13 @@ impl KeyVerifier { let updates = self.take_pending_last_used(); if let Err(error) = self.store.touch_last_used_batch(&updates).await { tracing::error!(error = %error, "failed to flush client key last_used_at batch"); - self.pending_last_used - .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner) - .extend(updates); + merge_last_used_updates( + &mut self + .pending_last_used + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner), + updates, + ); } self.flush_scheduled.store(false, Ordering::Release); if !self @@ -113,6 +116,22 @@ impl KeyVerifier { } } +fn merge_last_used_updates( + pending: &mut BTreeMap>, + updates: BTreeMap>, +) { + for (id, updated_at) in updates { + pending + .entry(id) + .and_modify(|pending_at| { + if updated_at > *pending_at { + *pending_at = updated_at; + } + }) + .or_insert(updated_at); + } +} + fn map_store_error(error: PgClientKeyStoreError) -> ClientKeyStoreError { ClientKeyStoreError::OperationFailed { message: error.to_string(), diff --git a/backend/src/update/archive.rs b/backend/src/update/archive.rs index d32952290..60ae8d9d1 100644 --- a/backend/src/update/archive.rs +++ b/backend/src/update/archive.rs @@ -127,25 +127,62 @@ pub(super) fn replace_release_files( let binary_backup = backup_path_for(exe_path); if binary_backup.exists() { if let Err(error) = fs::remove_file(&binary_backup) { + let mut rollback_errors = Vec::new(); if web_replaced { - let _ = restore_dir(web_dist_dir, &web_backup); + collect_rollback_error( + &mut rollback_errors, + "restore web assets", + restore_dir(web_dist_dir, &web_backup), + ); } - return Err(internal_error("Failed to remove old binary backup", error)); + return Err(update_error_with_rollback( + "Failed to remove old binary backup", + error, + rollback_errors, + )); } } if let Err(error) = move_file(exe_path, &binary_backup) { + let mut rollback_errors = Vec::new(); if web_replaced { - let _ = restore_dir(web_dist_dir, &web_backup); + collect_rollback_error( + &mut rollback_errors, + "restore web assets", + restore_dir(web_dist_dir, &web_backup), + ); } - return Err(internal_error("Binary backup failed", error)); + return Err(update_error_with_rollback( + "Binary backup failed", + error, + rollback_errors, + )); } if let Err(error) = move_file(&extracted.binary_path, exe_path) { - let _ = fs::remove_file(exe_path); - let _ = move_file(&binary_backup, exe_path); - if web_replaced { - let _ = restore_dir(web_dist_dir, &web_backup); + let mut rollback_errors = Vec::new(); + if exe_path.exists() { + collect_rollback_error( + &mut rollback_errors, + "remove partial replacement binary", + fs::remove_file(exe_path), + ); } - return Err(internal_error("Binary replace failed", error)); + collect_rollback_error( + &mut rollback_errors, + "restore previous binary", + move_file(&binary_backup, exe_path), + ); + if web_replaced { + collect_rollback_error( + &mut rollback_errors, + "restore web assets", + restore_dir(web_dist_dir, &web_backup), + ); + } + return Err(update_error_with_rollback( + "Binary replace failed", + error, + rollback_errors, + )); } Ok(()) } @@ -159,14 +196,43 @@ fn replace_dir(current: &Path, backup: &Path, replacement: &Path) -> Result<(), move_dir(current, backup).map_err(internal_error_with("Failed to backup web assets"))?; } if let Err(error) = move_dir(replacement, current) { + let mut rollback_errors = Vec::new(); if backup.exists() { - let _ = restore_dir(current, backup); + collect_rollback_error( + &mut rollback_errors, + "restore previous web assets", + restore_dir(current, backup), + ); } - return Err(internal_error("Failed to replace web assets", error)); + return Err(update_error_with_rollback( + "Failed to replace web assets", + error, + rollback_errors, + )); } Ok(()) } +fn collect_rollback_error(errors: &mut Vec, action: &'static str, result: io::Result<()>) { + if let Err(error) = result { + errors.push(format!("{action}: {error}")); + } +} + +fn update_error_with_rollback( + context: &'static str, + error: impl std::fmt::Display, + rollback_errors: Vec, +) -> UpdateError { + if rollback_errors.is_empty() { + return internal_error(context, error); + } + UpdateError::internal(format!( + "{context}: {error}; rollback failed: {}", + rollback_errors.join("; ") + )) +} + fn restore_dir(current: &Path, backup: &Path) -> io::Result<()> { if current.exists() { fs::remove_dir_all(current)?; diff --git a/backend/src/upstream/openai/transport/client.rs b/backend/src/upstream/openai/transport/client.rs index 8bf36e579..b1df57eb4 100644 --- a/backend/src/upstream/openai/transport/client.rs +++ b/backend/src/upstream/openai/transport/client.rs @@ -65,11 +65,11 @@ pub fn build_reqwest_client(force_http11: bool) -> Result let cache_key = (force_http11, custom_ca_env_cache_key()); static CLIENTS: OnceLock = OnceLock::new(); let cache = CLIENTS.get_or_init(|| Mutex::new(HashMap::new())); - let mut clients = cache + if let Some(client) = cache .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner); - - if let Some(client) = clients.get(&cache_key) { + .unwrap_or_else(std::sync::PoisonError::into_inner) + .get(&cache_key) + { return Ok(client.clone()); } @@ -88,8 +88,10 @@ pub fn build_reqwest_client(force_http11: bool) -> Result builder }; let client = build_reqwest_client_with_custom_ca(builder)?; - clients.insert(cache_key, client.clone()); - Ok(client) + let mut clients = cache + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + Ok(clients.entry(cache_key).or_insert(client).clone()) } // --------------------------------------------------------------------------- diff --git a/backend/src/upstream/openai/transport/websocket.rs b/backend/src/upstream/openai/transport/websocket.rs index 3122967f7..b88253393 100644 --- a/backend/src/upstream/openai/transport/websocket.rs +++ b/backend/src/upstream/openai/transport/websocket.rs @@ -41,7 +41,7 @@ use crate::upstream::openai::protocol::websocket::{ use super::websocket_pool::{ CodexWebSocketConnectionMetadata, CodexWebSocketPool, CodexWebSocketPoolKey, - PooledWebSocketConnection, WebSocketPoolAcquire, WebSocketPoolDecision, + PooledWebSocketConnection, WebSocketPoolAcquire, WebSocketPoolDecision, WebSocketPoolLease, }; use super::websocket_pump::{PumpKeepalive, PumpedWebSocket, RawWsStream}; use super::{diagnostics::CodexUpstreamDiagnostics, response_meta}; @@ -441,10 +441,15 @@ pub(crate) async fn execute_response_create_request_with_pool( }; match pool.acquire(&key).await { - WebSocketPoolAcquire::Reused(connection) => { - let result = - execute_pooled_response_create_request(request, pool, key, *connection, started_at) - .await; + WebSocketPoolAcquire::Reused { connection, lease } => { + let result = execute_pooled_response_create_request( + request, + pool, + lease, + *connection, + started_at, + ) + .await; match result { Err(CodexWebSocketExchangeError::ReusedConnectionDiedBeforeFirstOutput { .. @@ -467,21 +472,16 @@ pub(crate) async fn execute_response_create_request_with_pool( Err(error) => Err(error), } } - WebSocketPoolAcquire::FreshReserved => { - let result = execute_fresh_pooled_response_create_request( - request, - pool, - key.clone(), - started_at, - ) - .await; + WebSocketPoolAcquire::FreshReserved(lease) => { + let result = + execute_fresh_pooled_response_create_request(request, pool, lease, started_at) + .await; match result { Ok(mut exchange) => { exchange.pool_decision = Some(WebSocketPoolDecision::new()); Ok(exchange) } Err(error) if is_first_token_timeout(&error) => { - pool.discard(&key).await; let mut exchange = execute_fresh_response_create_request_with_retries( request, started_at, @@ -492,10 +492,7 @@ pub(crate) async fn execute_response_create_request_with_pool( exchange.pool_decision = Some(WebSocketPoolDecision::new()); Ok(exchange) } - Err(error) => { - pool.discard(&key).await; - Err(error) - } + Err(error) => Err(error), } } WebSocketPoolAcquire::Bypass(reason) => { @@ -527,14 +524,10 @@ pub(crate) async fn execute_response_create_request_stream_with_pool( }; match pool.acquire(&key).await { - WebSocketPoolAcquire::Reused(connection) => { - let result = execute_pooled_response_create_request_stream( - request, - pool.clone(), - key, - *connection, - ) - .await; + WebSocketPoolAcquire::Reused { connection, lease } => { + let result = + execute_pooled_response_create_request_stream(request, pool, lease, *connection) + .await; match result { Err(CodexWebSocketExchangeError::ReusedConnectionDiedBeforeFirstOutput { .. @@ -556,20 +549,15 @@ pub(crate) async fn execute_response_create_request_stream_with_pool( Err(error) => Err(error), } } - WebSocketPoolAcquire::FreshReserved => { - let result = execute_fresh_pooled_response_create_request_stream( - request, - pool.clone(), - key.clone(), - ) - .await; + WebSocketPoolAcquire::FreshReserved(lease) => { + let result = + execute_fresh_pooled_response_create_request_stream(request, pool, lease).await; match result { Ok(mut exchange) => { exchange.pool_decision = Some(WebSocketPoolDecision::new()); Ok(exchange) } Err(error) if is_first_token_timeout(&error) => { - pool.discard(&key).await; let mut exchange = execute_fresh_response_create_request_stream_with_retries( request, pool.first_token_timeout(), @@ -579,10 +567,7 @@ pub(crate) async fn execute_response_create_request_stream_with_pool( exchange.pool_decision = Some(WebSocketPoolDecision::new()); Ok(exchange) } - Err(error) => { - pool.discard(&key).await; - Err(error) - } + Err(error) => Err(error), } } WebSocketPoolAcquire::Bypass(reason) => { @@ -601,19 +586,29 @@ pub(crate) async fn execute_response_create_request_stream_with_pool( async fn execute_fresh_pooled_response_create_request( request: &CodexWebSocketRequest, pool: &CodexWebSocketPool, - key: CodexWebSocketPoolKey, + lease: WebSocketPoolLease, started_at: Instant, ) -> Result { let (websocket, response) = - connect_pumped_websocket(request.connection(), pool.keepalive()).await?; + match connect_pumped_websocket(request.connection(), pool.keepalive()).await { + Ok(connected) => connected, + Err(error) => { + lease.discard().await; + return Err(error); + } + }; let first_token_started_at = Instant::now(); - websocket + if let Err(error) = websocket .send(Message::Text(request.payload_text().to_string().into())) - .await?; + .await + { + lease.discard().await; + return Err(error.into()); + } let now = Instant::now(); let metadata = websocket_connection_metadata(&response); - let (exchange, websocket, metadata) = collect_websocket_response( + let (exchange, websocket, metadata) = match collect_websocket_response( websocket, metadata, false, @@ -621,30 +616,45 @@ async fn execute_fresh_pooled_response_create_request( first_token_started_at, pool.first_token_timeout(), ) - .await?; - pool.put( - key, - PooledWebSocketConnection { + .await + { + Ok(exchange) => exchange, + Err(error) => { + lease.discard().await; + return Err(error); + } + }; + lease + .put(PooledWebSocketConnection { websocket, metadata: reusable_websocket_metadata(metadata), created_at: now, - }, - ) - .await; + }) + .await; Ok(exchange) } async fn execute_fresh_pooled_response_create_request_stream( request: &CodexWebSocketRequest, - pool: CodexWebSocketPool, - key: CodexWebSocketPoolKey, + pool: &CodexWebSocketPool, + lease: WebSocketPoolLease, ) -> Result { let (mut websocket, response) = - connect_pumped_websocket(request.connection(), pool.keepalive()).await?; + match connect_pumped_websocket(request.connection(), pool.keepalive()).await { + Ok(connected) => connected, + Err(error) => { + lease.discard().await; + return Err(error); + } + }; let first_token_started_at = Instant::now(); - websocket + if let Err(error) = websocket .send(Message::Text(request.payload_text().to_string().into())) - .await?; + .await + { + lease.discard().await; + return Err(error.into()); + } let now = Instant::now(); let mut metadata = websocket_connection_metadata(&response); @@ -659,6 +669,7 @@ async fn execute_fresh_pooled_response_create_request_stream( Ok(prefetched_frames) => prefetched_frames, Err(error) => { websocket.close().await; + lease.discard().await; return Err(error); } }; @@ -666,8 +677,7 @@ async fn execute_fresh_pooled_response_create_request_stream( websocket, metadata, Some(WebSocketStreamPoolReturn { - pool, - key, + lease, created_at: now, }), prefetched_frames, @@ -677,7 +687,7 @@ async fn execute_fresh_pooled_response_create_request_stream( async fn execute_pooled_response_create_request( request: &CodexWebSocketRequest, pool: &CodexWebSocketPool, - key: CodexWebSocketPoolKey, + lease: WebSocketPoolLease, connection: PooledWebSocketConnection, started_at: Instant, ) -> Result { @@ -687,7 +697,7 @@ async fn execute_pooled_response_create_request( .send(Message::Text(request.payload_text().to_string().into())) .await { - pool.discard(&key).await; + lease.discard().await; return Err( CodexWebSocketExchangeError::ReusedConnectionDiedBeforeFirstOutput { message: error.to_string(), @@ -707,19 +717,17 @@ async fn execute_pooled_response_create_request( .await { Ok((exchange, websocket, metadata)) => { - pool.put( - key, - PooledWebSocketConnection { + lease + .put(PooledWebSocketConnection { websocket, metadata: reusable_websocket_metadata(metadata), created_at, - }, - ) - .await; + }) + .await; Ok(exchange) } Err(error) => { - pool.discard(&key).await; + lease.discard().await; Err(error) } } @@ -727,8 +735,8 @@ async fn execute_pooled_response_create_request( async fn execute_pooled_response_create_request_stream( request: &CodexWebSocketRequest, - pool: CodexWebSocketPool, - key: CodexWebSocketPoolKey, + pool: &CodexWebSocketPool, + lease: WebSocketPoolLease, connection: PooledWebSocketConnection, ) -> Result { let mut websocket = connection.websocket; @@ -739,7 +747,7 @@ async fn execute_pooled_response_create_request_stream( .send(Message::Text(request.payload_text().to_string().into())) .await { - pool.discard(&key).await; + lease.discard().await; return Err( CodexWebSocketExchangeError::ReusedConnectionDiedBeforeFirstOutput { message: error.to_string(), @@ -757,7 +765,7 @@ async fn execute_pooled_response_create_request_stream( { Ok(prefetched_frames) => prefetched_frames, Err(error) => { - pool.discard(&key).await; + lease.discard().await; websocket.close().await; return Err(reused_stream_prefetch_error(error)); } @@ -766,11 +774,7 @@ async fn execute_pooled_response_create_request_stream( Ok(stream_websocket_response( websocket, metadata, - Some(WebSocketStreamPoolReturn { - pool, - key, - created_at, - }), + Some(WebSocketStreamPoolReturn { lease, created_at }), prefetched_frames, )) } diff --git a/backend/src/upstream/openai/transport/websocket_frames.rs b/backend/src/upstream/openai/transport/websocket_frames.rs index 0b7db70a8..2cb95b06c 100644 --- a/backend/src/upstream/openai/transport/websocket_frames.rs +++ b/backend/src/upstream/openai/transport/websocket_frames.rs @@ -399,8 +399,7 @@ pub(super) fn reusable_websocket_metadata( // --------------------------------------------------------------------------- pub(super) struct WebSocketStreamPoolReturn { - pub(super) pool: CodexWebSocketPool, - pub(super) key: CodexWebSocketPoolKey, + pub(super) lease: WebSocketPoolLease, pub(super) created_at: Instant, } @@ -716,15 +715,12 @@ async fn finish_stream_websocket( return; }; pool_return - .pool - .put( - pool_return.key, - PooledWebSocketConnection { - websocket, - metadata: reusable_websocket_metadata(metadata), - created_at: pool_return.created_at, - }, - ) + .lease + .put(PooledWebSocketConnection { + websocket, + metadata: reusable_websocket_metadata(metadata), + created_at: pool_return.created_at, + }) .await; } @@ -739,7 +735,7 @@ async fn discard_stream_websocket( "discarding Responses WebSocket stream" ); if let Some(pool_return) = pool_return { - pool_return.pool.discard(&pool_return.key).await; + pool_return.lease.discard().await; } websocket.close().await; } diff --git a/backend/src/upstream/openai/transport/websocket_pool.rs b/backend/src/upstream/openai/transport/websocket_pool.rs index 4dfed24d3..1036da43f 100644 --- a/backend/src/upstream/openai/transport/websocket_pool.rs +++ b/backend/src/upstream/openai/transport/websocket_pool.rs @@ -9,6 +9,7 @@ use std::{ use serde_json::{json, Value}; use sha2::{Digest, Sha256}; use tokio::sync::Mutex; +use uuid::Uuid; use super::websocket_pump::{PumpKeepalive, PumpedWebSocket}; @@ -160,7 +161,7 @@ impl CodexWebSocketPool { return WebSocketPoolAcquire::Bypass(WebSocketPoolBypassReason::Disabled); } match state.slots.get(key) { - Some(WebSocketPoolSlot::Busy) => { + Some(WebSocketPoolSlot::Busy(_)) => { return WebSocketPoolAcquire::Bypass(WebSocketPoolBypassReason::Busy); } Some(WebSocketPoolSlot::Idle(_)) => { @@ -172,8 +173,11 @@ impl CodexWebSocketPool { if connection.created_at.elapsed() < self.config.max_age && !connection.websocket.is_closed() { - state.slots.insert(key.clone(), WebSocketPoolSlot::Busy); - return WebSocketPoolAcquire::Reused(connection); + let lease = WebSocketPoolLease::reserve(self.clone(), key.clone()); + state + .slots + .insert(key.clone(), WebSocketPoolSlot::Busy(lease.reservation)); + return WebSocketPoolAcquire::Reused { connection, lease }; } expired_connection = Some(*connection); } @@ -185,8 +189,11 @@ impl CodexWebSocketPool { { WebSocketPoolAcquire::Bypass(WebSocketPoolBypassReason::Cap) } else { - state.slots.insert(key.clone(), WebSocketPoolSlot::Busy); - WebSocketPoolAcquire::FreshReserved + let lease = WebSocketPoolLease::reserve(self.clone(), key.clone()); + state + .slots + .insert(key.clone(), WebSocketPoolSlot::Busy(lease.reservation)); + WebSocketPoolAcquire::FreshReserved(lease) } }; @@ -197,9 +204,10 @@ impl CodexWebSocketPool { acquire } - pub(crate) async fn put( + async fn put_reserved( &self, - key: CodexWebSocketPoolKey, + key: &CodexWebSocketPoolKey, + reservation_id: Uuid, connection: PooledWebSocketConnection, ) { let mut connection = Some(connection); @@ -208,13 +216,18 @@ impl CodexWebSocketPool { let expired = connection .as_ref() .is_some_and(|connection| connection.created_at.elapsed() >= self.config.max_age); - if expired || state.shutting_down || !self.config.enabled { - state.slots.remove(&key); - } else if matches!(state.slots.get(&key), Some(WebSocketPoolSlot::Busy)) { + let owns_reservation = matches!( + state.slots.get(key), + Some(WebSocketPoolSlot::Busy(reservation)) + if reservation.id == reservation_id + ); + if owns_reservation && (expired || state.shutting_down || !self.config.enabled) { + state.slots.remove(key); + } else if owns_reservation { if let Some(connection) = connection.take() { state .slots - .insert(key, WebSocketPoolSlot::Idle(Box::new(connection))); + .insert(key.clone(), WebSocketPoolSlot::Idle(Box::new(connection))); } } } @@ -223,16 +236,13 @@ impl CodexWebSocketPool { } } - pub(crate) async fn discard(&self, key: &CodexWebSocketPoolKey) { - let idle_connection = { - let mut state = self.inner.lock().await; - match state.slots.remove(key) { - Some(WebSocketPoolSlot::Idle(connection)) => Some(*connection), - Some(WebSocketPoolSlot::Busy) | None => None, - } - }; - if let Some(connection) = idle_connection { - close_pooled_connection(connection).await; + async fn discard_reserved(&self, key: &CodexWebSocketPoolKey, reservation_id: Uuid) { + let mut state = self.inner.lock().await; + if matches!( + state.slots.get(key), + Some(WebSocketPoolSlot::Busy(reservation)) if reservation.id == reservation_id + ) { + state.slots.remove(key); } } @@ -266,20 +276,20 @@ impl CodexWebSocketPool { .drain() .filter_map(|(_, slot)| match slot { WebSocketPoolSlot::Idle(connection) => Some(*connection), - WebSocketPoolSlot::Busy => None, + WebSocketPoolSlot::Busy(_) => None, }) .collect::>() }; close_pooled_connections(idle_connections).await; } - /// 维护 idle 连接：清扫已被后台 pump 判定死亡或已超龄的连接。 + /// 维护池 slot：清扫已死亡或超龄的 idle 连接，以及异常残留的 Busy reservation。 /// /// 保活（ping/pong）与失活检测已下沉到每条连接的 pump 任务内部，此处不再做 /// 同步 ping 探活，只负责把「后台已标记 closed」或「超过 max_age」的 idle /// 连接从池中摘除并关闭，避免它们占用 slot。 pub async fn maintain_idle_connections(&self) { - let close = self.take_dead_idle_connections().await; + let close = self.take_expired_slots().await; close_pooled_connections(close).await; } @@ -313,7 +323,7 @@ impl CodexWebSocketPool { self.inner.lock().await.shutting_down } - async fn take_dead_idle_connections(&self) -> Vec { + async fn take_expired_slots(&self) -> Vec { let mut close = Vec::new(); let now = Instant::now(); let mut state = self.inner.lock().await; @@ -329,12 +339,26 @@ impl CodexWebSocketPool { { Some(key.clone()) } - WebSocketPoolSlot::Idle(_) | WebSocketPoolSlot::Busy => None, + WebSocketPoolSlot::Busy(reservation) + if now.duration_since(reservation.reserved_at) >= self.config.max_age => + { + Some(key.clone()) + } + WebSocketPoolSlot::Idle(_) | WebSocketPoolSlot::Busy(_) => None, }) .collect::>(); for key in keys { - if let Some(WebSocketPoolSlot::Idle(connection)) = state.slots.remove(&key) { - close.push(*connection); + match state.slots.remove(&key) { + Some(WebSocketPoolSlot::Idle(connection)) => close.push(*connection), + Some(WebSocketPoolSlot::Busy(reservation)) => { + tracing::warn!( + account_id = key.account_id(), + conversation_id_hash = key.conversation_id_hash(), + reservation_id = %reservation.id, + "removed stale WebSocket pool reservation" + ); + } + None => {} } } close @@ -363,15 +387,76 @@ pub(crate) struct PooledWebSocketConnection { enum WebSocketPoolSlot { Idle(Box), - Busy, + Busy(WebSocketPoolReservation), } pub(crate) enum WebSocketPoolAcquire { - Reused(Box), - FreshReserved, + Reused { + connection: Box, + lease: WebSocketPoolLease, + }, + FreshReserved(WebSocketPoolLease), Bypass(WebSocketPoolBypassReason), } +#[derive(Clone, Copy)] +struct WebSocketPoolReservation { + id: Uuid, + reserved_at: Instant, +} + +pub(crate) struct WebSocketPoolLease { + pool: CodexWebSocketPool, + key: CodexWebSocketPoolKey, + reservation: WebSocketPoolReservation, + armed: bool, +} + +impl WebSocketPoolLease { + fn reserve(pool: CodexWebSocketPool, key: CodexWebSocketPoolKey) -> Self { + Self { + pool, + key, + reservation: WebSocketPoolReservation { + id: Uuid::new_v4(), + reserved_at: Instant::now(), + }, + armed: true, + } + } + + pub(crate) async fn put(mut self, connection: PooledWebSocketConnection) { + self.pool + .put_reserved(&self.key, self.reservation.id, connection) + .await; + self.armed = false; + } + + pub(crate) async fn discard(mut self) { + self.pool + .discard_reserved(&self.key, self.reservation.id) + .await; + self.armed = false; + } +} + +impl Drop for WebSocketPoolLease { + fn drop(&mut self) { + if !self.armed { + return; + } + let Ok(runtime) = tokio::runtime::Handle::try_current() else { + return; + }; + let pool = self.pool.clone(); + let key = self.key.clone(); + let reservation_id = self.reservation.id; + runtime.spawn(async move { + pool.discard_reserved(&key, reservation_id).await; + }); + } +} + #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum WebSocketPoolBypassReason { Disabled, diff --git a/backend/src/upstream/openai/transport/websocket_pump.rs b/backend/src/upstream/openai/transport/websocket_pump.rs index e61c2a71d..3463f0efd 100644 --- a/backend/src/upstream/openai/transport/websocket_pump.rs +++ b/backend/src/upstream/openai/transport/websocket_pump.rs @@ -58,10 +58,12 @@ enum PumpCommand { }, } +const PUMP_MESSAGE_BUFFER: usize = 64; + /// 由后台 pump 任务托管的 WebSocket 连接句柄。 pub(crate) struct PumpedWebSocket { tx_command: mpsc::Sender, - rx_message: mpsc::UnboundedReceiver>, + rx_message: mpsc::Receiver>, closed: Arc, pump: Option>, } @@ -78,7 +80,7 @@ impl PumpedWebSocket { /// 用底层流启动一个 pump 任务并返回句柄。 pub(crate) fn new(inner: RawWsStream, keepalive: PumpKeepalive) -> Self { let (tx_command, rx_command) = mpsc::channel::(PUMP_COMMAND_BUFFER); - let (tx_message, rx_message) = mpsc::unbounded_channel(); + let (tx_message, rx_message) = mpsc::channel(PUMP_MESSAGE_BUFFER); let closed = Arc::new(AtomicBool::new(false)); let closed_for_task = Arc::clone(&closed); let pump = tokio::spawn(async move { @@ -138,7 +140,7 @@ impl Drop for PumpedWebSocket { async fn pump_loop( mut inner: RawWsStream, mut rx_command: mpsc::Receiver, - tx_message: mpsc::UnboundedSender>, + tx_message: mpsc::Sender>, keepalive: PumpKeepalive, ) { let mut last_activity = Instant::now(); @@ -191,7 +193,7 @@ async fn pump_loop( Some(Ok(message)) => { last_activity = Instant::now(); let is_close = matches!(message, Message::Close(_)); - if tx_message.send(Ok(message)).is_err() { + if tx_message.try_send(Ok(message)).is_err() { break; } if is_close { @@ -199,7 +201,7 @@ async fn pump_loop( } } Some(Err(err)) => { - let _ = tx_message.send(Err(err)); + let _ = tx_message.try_send(Err(err)); break; } } diff --git a/backend/tests/dispatch/service.rs b/backend/tests/dispatch/service.rs index 0cf3928a3..ebf38c624 100644 --- a/backend/tests/dispatch/service.rs +++ b/backend/tests/dispatch/service.rs @@ -20,7 +20,11 @@ use codex_proxy_rs::{ build_conversation_identity, compute_variant_hash, prepare_variant_identity, SessionAffinityEntry, }, - fleet::{account::AccountStatus, cookies::PgCookieStore, pool::AccountAcquireRequest}, + fleet::{ + account::AccountStatus, + cookies::PgCookieStore, + pool::{AccountAcquireRequest, AccountLease}, + }, infra::identity::hash_credential, telemetry::usage::store::{PgUsageRecordStore, UsageRecordFilter}, upstream::openai::protocol::responses::CodexResponsesRequest, diff --git a/backend/tests/dispatch/service/responses_websocket.rs b/backend/tests/dispatch/service/responses_websocket.rs index d388a9cd0..0de9d52d1 100644 --- a/backend/tests/dispatch/service/responses_websocket.rs +++ b/backend/tests/dispatch/service/responses_websocket.rs @@ -1671,7 +1671,7 @@ async fn responses_websocket_should_preserve_known_history_on_fallback_account() config.auth.max_concurrent_per_account = 1; }) .await; - seed_primary_affinity_and_hold_primary_slot(&state).await; + let _held_primary_slot = seed_primary_affinity_and_hold_primary_slot(&state).await; let response = app .oneshot(responses_json_request( @@ -1709,7 +1709,7 @@ async fn responses_websocket_stream_should_preserve_known_history_on_fallback_ac config.auth.max_concurrent_per_account = 1; }) .await; - seed_primary_affinity_and_hold_primary_slot(&state).await; + let _held_primary_slot = seed_primary_affinity_and_hold_primary_slot(&state).await; let response = app .oneshot(responses_json_request( @@ -2362,7 +2362,7 @@ async fn responses_stream_with_previous_response_id_should_record_websocket_audi ); } -async fn seed_primary_affinity_and_hold_primary_slot(state: &AppState) { +async fn seed_primary_affinity_and_hold_primary_slot(state: &AppState) -> AccountLease { state .services .session_affinity @@ -2391,4 +2391,5 @@ async fn seed_primary_affinity_and_hold_primary_slot(state: &AppState) { .await .unwrap(); assert_eq!(held_primary_slot.account.id, "acct_primary"); + held_primary_slot } diff --git a/backend/tests/fleet/cookies/store.rs b/backend/tests/fleet/cookies/store.rs index cc292dbdf..cc8bf589f 100644 --- a/backend/tests/fleet/cookies/store.rs +++ b/backend/tests/fleet/cookies/store.rs @@ -59,6 +59,53 @@ async fn cookie_store_should_cleanup_expired_cookies() { assert_eq!(remaining.0, 1); } +#[tokio::test] +async fn cookie_store_should_persist_account_cleanup_deadline() { + let (pool, _dir) = cookie_store_parts("cookies-cleanup-deadline").await; + seed_account(&pool, "acct_deadline").await; + let store = PgCookieStore::new(pool.clone()); + store + .capture_set_cookie( + "acct_deadline", + "cf_clearance=keep-until-deadline; Domain=.chatgpt.com; Path=/", + ) + .await + .unwrap(); + let deadline = Utc::now() + Duration::seconds(30); + + let updated = store + .expire_account_cookies_at("acct_deadline", deadline) + .await + .unwrap(); + let stored_deadline: (Option>,) = + sqlx::query_as("select expires_at from account_cookies where account_id = $1") + .bind("acct_deadline") + .fetch_one(&pool) + .await + .unwrap(); + + assert_eq!(updated, 1); + assert_eq!( + stored_deadline.0.map(|value| value.timestamp_micros()), + Some(deadline.timestamp_micros()) + ); + assert_eq!( + store + .cookie_header("acct_deadline", "chatgpt.com") + .await + .unwrap() + .as_deref(), + Some("cf_clearance=keep-until-deadline") + ); + assert_eq!( + store + .cleanup_expired(deadline + Duration::seconds(1)) + .await + .unwrap(), + 1 + ); +} + async fn cookie_store_parts(name: &str) -> (PgPool, crate::support::storage::TestDatabaseGuard) { init_test_db(name).await } diff --git a/backend/tests/fleet/pool/mod.rs b/backend/tests/fleet/pool/mod.rs index b7639ba97..0183cd02f 100644 --- a/backend/tests/fleet/pool/mod.rs +++ b/backend/tests/fleet/pool/mod.rs @@ -2,10 +2,10 @@ use std::collections::{BTreeMap, BTreeSet}; use chrono::{Duration, TimeZone, Utc}; use codex_proxy_rs::fleet::{ - account::{Account, AccountStatus}, + account::AccountStatus, pool::{ AccountAcquireRequest, AccountPool, AccountPoolOptions, AccountWindowUsageDelta, - RotationStrategy, + AcquiredAccount, RotationStrategy, }, }; use serde_json::json; @@ -19,9 +19,8 @@ fn fixed_time() -> chrono::DateTime { Utc.with_ymd_and_hms(2026, 6, 11, 8, 0, 0).unwrap() } -fn acquire_account(pool: &mut AccountPool, model: &str) -> Option { +fn acquire_account(pool: &mut AccountPool, model: &str) -> Option { pool.acquire_with(&AccountAcquireRequest::new(model, Utc::now())) - .map(|acquired| acquired.account) } fn test_jwt(exp_offset_seconds: i64) -> String { diff --git a/backend/tests/fleet/pool/quota.rs b/backend/tests/fleet/pool/quota.rs index 1003f39ac..c5bc3a268 100644 --- a/backend/tests/fleet/pool/quota.rs +++ b/backend/tests/fleet/pool/quota.rs @@ -225,7 +225,7 @@ fn sync_rate_limit_window_should_reset_window_counters_when_drift_crosses_thresh let acquired = pool .acquire_with(&AccountAcquireRequest::new("gpt-5.5", now)) .unwrap(); - pool.release(&acquired.account.id); + pool.release(&acquired); let released = pool.get(&acquired.account.id).unwrap(); assert_eq!(released.window_request_count, 1); diff --git a/backend/tests/fleet/pool/runtime.rs b/backend/tests/fleet/pool/runtime.rs index 68fb12ff2..ececf3396 100644 --- a/backend/tests/fleet/pool/runtime.rs +++ b/backend/tests/fleet/pool/runtime.rs @@ -81,6 +81,43 @@ async fn runtime_account_pool_should_restore_accounts_by_added_at_and_id() { assert_eq!(acquired.account.id, "acct_a"); } +#[tokio::test] +async fn runtime_account_lease_drop_should_release_capacity() { + let (pool, _dir) = init_test_db("runtime-account-pool-lease-drop").await; + insert_account(&pool, "acct_drop").await; + + let store = PgAccountStore::new(pool.clone()); + let service = AccountPoolService::new( + Arc::new(store) as Arc, + Arc::new(PgAccountUsageStore::new(pool)), + AccountPoolOptions { + max_concurrent_per_account: 1, + ..AccountPoolOptions::default() + }, + 0, + ); + service + .restore_from_store() + .await + .expect("account pool should restore"); + let lease = service + .acquire_with(&AccountAcquireRequest::new("gpt-5.5", Utc::now())) + .await + .expect("active account should be acquired"); + + drop(lease); + tokio::time::timeout(std::time::Duration::from_secs(1), async { + loop { + if service.capacity_summary_now().await.available_slots == 1 { + break; + } + tokio::task::yield_now().await; + } + }) + .await + .expect("dropped lease should release its slot"); +} + async fn insert_account(pool: &sqlx::PgPool, id: &str) { sqlx::query( "insert into accounts (id, email, chatgpt_account_id, chatgpt_user_id, access_token, access_token_expires_at, status, added_at, updated_at) values ($1, $2, $3, $4, $5, $6, $7, $8, $9)", diff --git a/backend/tests/fleet/pool/selection.rs b/backend/tests/fleet/pool/selection.rs index b035f068e..5794d6747 100644 --- a/backend/tests/fleet/pool/selection.rs +++ b/backend/tests/fleet/pool/selection.rs @@ -22,7 +22,7 @@ fn account_pool_should_respect_max_concurrent_slots_per_account() { assert_ne!(first.id, second.id); assert!(third.is_none()); - pool.release(&first.id); + pool.release(&first); assert_eq!(acquire_account(&mut pool, "gpt-5.5").unwrap().id, first.id); } @@ -94,8 +94,8 @@ fn round_robin_should_keep_ts_cursor_when_candidates_recover() { let first = acquire_account(&mut pool, "gpt-5.5").unwrap(); let second = acquire_account(&mut pool, "gpt-5.5").unwrap(); - pool.release(&first.id); - pool.release(&second.id); + pool.release(&first); + pool.release(&second); let third = acquire_account(&mut pool, "gpt-5.5").unwrap(); assert_eq!(third.id, "acct_c"); @@ -155,6 +155,20 @@ fn acquire_should_skip_accounts_with_expired_token_metadata() { assert!(acquired.is_none()); } +#[test] +fn acquire_should_trust_future_token_expiry_metadata_without_parsing_jwt() { + let now = fixed_time(); + let mut active = crate::support::accounts::test_account("active", AccountStatus::Active); + active.access_token = test_jwt(-60); + active.access_token_expires_at = Some(now + Duration::hours(1)); + let mut pool = AccountPool::default(); + pool.insert(active); + + let acquired = pool.acquire_with(&AccountAcquireRequest::new("gpt-5.5", now)); + + assert_eq!(acquired.unwrap().account.id, "active"); +} + #[test] fn acquire_should_skip_accounts_with_expired_jwt_when_metadata_is_missing() { let now = Utc::now(); @@ -444,7 +458,12 @@ fn account_pool_distinct_plan_accounts_should_filter_like_model_refresh() { assert_eq!(selected.len(), 1); assert_eq!(selected[0].plan_type, "plus"); assert_eq!(selected[0].account.id, "plus-ok"); - assert_eq!(pool.release("plus-ok").unwrap().model, None); + assert_eq!( + pool.release_distinct_plan_account(&selected[0]) + .unwrap() + .model, + None + ); } #[test] @@ -609,7 +628,7 @@ fn smart_should_balance_released_requests_by_runtime_usage() { let first = pool .acquire_with(&AccountAcquireRequest::new("gpt-5.5", now)) .unwrap(); - pool.release(&first.account.id); + pool.release(&first); let second = pool .acquire_with(&AccountAcquireRequest::new( "gpt-5.5", @@ -675,7 +694,7 @@ fn smart_should_not_drain_earliest_reset_account_when_other_accounts_are_idle() )) .unwrap(); selected_ids.push(acquired.account.id.clone()); - pool.release(&acquired.account.id); + pool.release(&acquired); } assert_eq!( @@ -795,7 +814,7 @@ fn account_pool_should_rotate_tied_quota_reset_priority_accounts() { )); let first = acquire_account(&mut pool, "gpt-5.5").unwrap(); - pool.release(&first.id); + pool.release(&first); let second = acquire_account(&mut pool, "gpt-5.5").unwrap(); assert_ne!(first.id, second.id); @@ -820,14 +839,11 @@ fn quota_reset_priority_should_prefer_lru_before_tie_rotation() { pool.insert(older); pool.insert(newer); - assert_eq!( - pool.acquire_with(&AccountAcquireRequest::new("gpt-5.5", now)) - .unwrap() - .account - .id, - "seed" - ); - pool.release("seed"); + let seed = pool + .acquire_with(&AccountAcquireRequest::new("gpt-5.5", now)) + .unwrap(); + assert_eq!(seed.account.id, "seed"); + pool.release(&seed); let acquired = pool .acquire_with(&AccountAcquireRequest::new( "gpt-5.5", @@ -947,7 +963,7 @@ fn quota_reset_priority_should_not_mutate_candidate_order() { let selected = pool .acquire_with(&AccountAcquireRequest::new("gpt-5.5", now)) .unwrap(); - pool.release(&selected.account.id); + pool.release(&selected); pool.set_options(AccountPoolOptions { rotation_strategy: RotationStrategy::RoundRobin, ..AccountPoolOptions::default() diff --git a/backend/tests/fleet/pool/usage_window.rs b/backend/tests/fleet/pool/usage_window.rs index 20f0ffe38..e5f1f6bb5 100644 --- a/backend/tests/fleet/pool/usage_window.rs +++ b/backend/tests/fleet/pool/usage_window.rs @@ -91,6 +91,38 @@ fn quota_reset_priority_should_compare_request_count_when_only_one_account_has_w assert_eq!(acquired.account.id, "reset_unknown"); } +#[test] +fn release_should_remove_the_exact_concurrent_slot() { + let now = fixed_time(); + let mut pool = AccountPool::with_options(AccountPoolOptions { + max_concurrent_per_account: 2, + ..AccountPoolOptions::default() + }); + pool.insert(crate::support::accounts::test_account( + "acct_a", + AccountStatus::Active, + )); + + let first = pool + .acquire_with(&AccountAcquireRequest::new("model-first", now)) + .unwrap(); + let second = pool + .acquire_with(&AccountAcquireRequest::new( + "model-second", + now + Duration::milliseconds(1), + )) + .unwrap(); + + assert_eq!( + pool.release(&second).unwrap().model.as_deref(), + Some("model-second") + ); + assert_eq!( + pool.release(&first).unwrap().model.as_deref(), + Some("model-first") + ); +} + #[test] fn release_should_mark_runtime_usage_window() { let now = fixed_time(); @@ -106,7 +138,7 @@ fn release_should_mark_runtime_usage_window() { assert_eq!(acquired.account.request_count, 0); assert_eq!(acquired.account.window_request_count, 0); - pool.release(&acquired.account.id); + pool.release(&acquired); let released = pool.get(&acquired.account.id).unwrap(); assert_eq!(released.request_count, 1); @@ -129,7 +161,7 @@ fn release_should_start_usage_window_without_known_reset() { .acquire_with(&AccountAcquireRequest::new("gpt-5.5", now)) .unwrap(); - pool.release(&acquired.account.id); + pool.release(&acquired); let released = pool.get(&acquired.account.id).unwrap(); assert!(released.window_started_at.is_some()); @@ -153,7 +185,7 @@ fn release_should_mark_usage_after_stale_slot_cleanup() { .unwrap(); let summary = pool.capacity_summary(now + Duration::minutes(6)); - pool.release(&acquired.account.id); + pool.release(&acquired); let released = pool.get(&acquired.account.id).unwrap(); assert_eq!(summary.used_slots, 0); diff --git a/backend/tests/fleet/refresh/failures.rs b/backend/tests/fleet/refresh/failures.rs index a719c29f8..a384f7bfb 100644 --- a/backend/tests/fleet/refresh/failures.rs +++ b/backend/tests/fleet/refresh/failures.rs @@ -44,18 +44,18 @@ async fn token_refresh_task_should_keep_business_status_before_refresher_call() refresher, ); - let summary = task - .refresh_due_accounts_once_at(now) + let timer_summary = task + .schedule_account_timers_once_at(now) .await - .expect("refresh should succeed"); - let stored = store - .get("acct-status-marker") - .await - .expect("account should load") - .expect("account should exist"); + .expect("refresh timer should be scheduled"); + let stored = wait_for_account(&store, "acct-status-marker", |account| { + account.access_token.expose_secret() == new_access_token.as_str() + }) + .await; + task.shutdown().await; let observed_statuses = observed_statuses.lock().await.clone(); - assert_eq!(summary.refreshed, 1); + assert_eq!(timer_summary.immediate, 1); assert_eq!(observed_statuses, [AccountStatus::Active]); assert_eq!(stored.status, AccountStatus::Active); assert_eq!( @@ -96,8 +96,8 @@ async fn token_refresh_task_should_skip_expired_account_after_restart() { refresher.clone(), ); - let summary = task - .refresh_due_accounts_once_at(now) + let timer_summary = task + .schedule_account_timers_once_at(now) .await .expect("expired account should be skipped"); let stored = store @@ -106,7 +106,7 @@ async fn token_refresh_task_should_skip_expired_account_after_restart() { .expect("account should load") .expect("account should exist"); - assert_eq!(summary.skipped, 1); + assert_eq!(timer_summary.skipped, 1); assert_eq!(refresher.calls.load(Ordering::SeqCst), 0); assert_eq!(stored.status, AccountStatus::Expired); assert!(stored.next_refresh_at.is_none()); @@ -151,18 +151,18 @@ async fn token_refresh_task_should_retry_transport_failure_before_recovery() { ) .with_retry_delays(vec![StdDuration::ZERO; 4]); - let summary = task - .refresh_due_accounts_once_at(now) + let timer_summary = task + .schedule_account_timers_once_at(now) .await - .expect("transport failure should be summarized"); - let stored = store - .get("acct-refresh-transport") - .await - .expect("account should load") - .expect("account should exist"); + .expect("refresh timer should be scheduled"); + let stored = wait_for_account(&store, "acct-refresh-transport", |account| { + account.next_refresh_at.is_some_and(|next| next > now) + }) + .await; + task.shutdown().await; let observed_statuses = observed_statuses.lock().await.clone(); - assert_eq!(summary.failed, 1); + assert_eq!(timer_summary.immediate, 1); assert_eq!(observed_statuses, [AccountStatus::Active; 5]); assert_eq!(stored.status, AccountStatus::Active); assert!(stored.next_refresh_at.is_some_and(|next| next > now)); @@ -221,36 +221,36 @@ async fn token_refresh_task_should_delay_recovery_after_retry_exhaustion() { .with_retry_delays(vec![StdDuration::ZERO; 4]); let failed = task - .refresh_due_accounts_once_at(now) + .schedule_account_timers_once_at(now) .await - .expect("first scan should summarize retry exhaustion"); - let recovery_at = store - .get("acct-refresh-delayed-recovery") - .await - .expect("account should load") - .expect("account should exist") - .next_refresh_at - .expect("retry exhaustion should persist recovery time"); + .expect("first refresh timer should be scheduled"); + let recovery_at = wait_for_account(&store, "acct-refresh-delayed-recovery", |account| { + account.next_refresh_at.is_some_and(|next| next > now) + }) + .await + .next_refresh_at + .expect("retry exhaustion should persist recovery time"); let delayed = task - .refresh_due_accounts_once_at(now + Duration::minutes(5)) + .schedule_account_timers_once_at(now + Duration::minutes(5)) .await - .expect("recovery window should skip refresh"); + .expect("recovery timer should be scheduled"); let refreshed = task - .refresh_due_accounts_once_at(recovery_at + Duration::seconds(1)) + .schedule_account_timers_once_at(recovery_at + Duration::seconds(1)) .await - .expect("recovery window should allow refresh"); - let stored = store - .get("acct-refresh-delayed-recovery") - .await - .expect("account should load") - .expect("account should exist"); + .expect("due recovery timer should replace future timer"); + let stored = wait_for_account(&store, "acct-refresh-delayed-recovery", |account| { + account.access_token.expose_secret() == new_access_token.as_str() + }) + .await; + task.shutdown().await; let observed_statuses = observed_statuses.lock().await.clone(); - assert_eq!(failed.failed, 1); + assert_eq!(failed.immediate, 1); assert!(recovery_at >= now + Duration::minutes(8)); assert!(recovery_at <= now + Duration::minutes(12)); - assert_eq!(delayed.skipped, 1); - assert_eq!(refreshed.refreshed, 1); + assert_eq!(delayed.scheduled, 1); + assert_eq!(refreshed.immediate, 1); + assert_eq!(refreshed.replaced, 1); assert_eq!(observed_statuses, [AccountStatus::Active; 6]); assert_eq!(stored.status, AccountStatus::Active); assert_eq!( @@ -299,17 +299,20 @@ async fn token_refresh_task_should_not_reuse_stale_refresh_token_after_retryable ) .with_retry_delays(vec![StdDuration::ZERO; 4]); - let summary = task - .refresh_due_accounts_once_at(now) + let timer_summary = task + .schedule_account_timers_once_at(now) .await - .expect("stale refresh token should skip retry"); - let stored = store - .get("acct-refresh-stale-rt") - .await - .expect("account should load") - .expect("account should exist"); + .expect("refresh timer should be scheduled"); + let stored = wait_for_account(&store, "acct-refresh-stale-rt", |account| { + account + .refresh_token + .as_ref() + .is_some_and(|token| token.expose_secret() == "refresh-rotated") + }) + .await; + task.shutdown().await; - assert_eq!(summary.skipped, 1); + assert_eq!(timer_summary.immediate, 1); assert_eq!(calls.load(Ordering::SeqCst), 1); assert_eq!( stored.access_token.expose_secret(), @@ -368,18 +371,18 @@ async fn token_refresh_task_should_confirm_invalid_grant_before_expiring_account ) .with_retry_delays(vec![StdDuration::ZERO; 4]); - let summary = task - .refresh_due_accounts_once_at(now) + let timer_summary = task + .schedule_account_timers_once_at(now) .await - .expect("confirmed permanent failure should update status"); - let stored = store - .get("acct-refresh-invalid-grant") - .await - .expect("account should load") - .expect("account should exist"); + .expect("refresh timer should be scheduled"); + let stored = wait_for_account(&store, "acct-refresh-invalid-grant", |account| { + account.status == AccountStatus::Expired + }) + .await; + task.shutdown().await; let observed_statuses = observed_statuses.lock().await.clone(); - assert_eq!(summary.status_updated, 1); + assert_eq!(timer_summary.immediate, 1); assert_eq!(observed_statuses, [AccountStatus::Active; 2]); assert_eq!(stored.status, AccountStatus::Expired); assert_eq!( diff --git a/backend/tests/fleet/refresh/mod.rs b/backend/tests/fleet/refresh/mod.rs index bc3b394df..aae9a73c2 100644 --- a/backend/tests/fleet/refresh/mod.rs +++ b/backend/tests/fleet/refresh/mod.rs @@ -11,8 +11,9 @@ use async_trait::async_trait; use chrono::{Duration, TimeZone, Utc}; use codex_proxy_rs::fleet::account::AccountStatus; use codex_proxy_rs::fleet::refresh::{RedisRefreshLeaseStore, RefreshPolicy}; -use codex_proxy_rs::fleet::store::PgAccountStore; -use codex_proxy_rs::fleet::store::{AccountClaimsUpdate, NewAccount}; +use codex_proxy_rs::fleet::store::{ + AccountClaimsUpdate, NewAccount, PgAccountStore, StoredAccount, +}; use codex_proxy_rs::upstream::openai::token_client::{RefreshFailure, TokenPair, TokenRefresher}; use secrecy::{ExposeSecret, SecretString}; use serde_json::json; @@ -177,3 +178,42 @@ fn test_jwt(exp: i64) -> String { let payload = json!({ "exp": exp }); crate::support::jwt::unsigned_jwt(&payload) } + +async fn wait_for_account( + store: &PgAccountStore, + account_id: &str, + predicate: impl Fn(&StoredAccount) -> bool, +) -> StoredAccount { + timeout(StdDuration::from_secs(2), async { + loop { + let account = store + .get(account_id) + .await + .expect("account should load") + .expect("account should exist"); + if predicate(&account) { + return account; + } + sleep(StdDuration::from_millis(20)).await; + } + }) + .await + .unwrap_or_else(|_| panic!("account {account_id} did not reach expected state")) +} + +async fn wait_for_no_scheduled_timers( + service: &codex_proxy_rs::fleet::refresh::TokenRefreshService, +) where + C: TokenRefresher, +{ + timeout(StdDuration::from_secs(2), async { + loop { + if service.scheduled_timer_count().await == 0 { + return; + } + tokio::task::yield_now().await; + } + }) + .await + .expect("scheduled refresh timers should finish"); +} diff --git a/backend/tests/fleet/refresh/scheduling.rs b/backend/tests/fleet/refresh/scheduling.rs index a75ae3c45..9b55008e7 100644 --- a/backend/tests/fleet/refresh/scheduling.rs +++ b/backend/tests/fleet/refresh/scheduling.rs @@ -43,18 +43,19 @@ async fn token_refresh_task_should_skip_account_when_refresh_lease_is_held() { ) .with_refresh_lease_store(leases); - let summary = task - .refresh_due_accounts_once_at(now) + let timer_summary = task + .schedule_account_timers_once_at(now) .await - .expect("refresh scan should skip lease-held account"); + .expect("refresh timer should be scheduled"); + wait_for_no_scheduled_timers(&task).await; + task.shutdown().await; let stored = store .get("acct-lease-held") .await .expect("account should load") .expect("account should exist"); - assert_eq!(summary.skipped, 1); - assert_eq!(summary.refreshed, 0); + assert_eq!(timer_summary.immediate, 1); assert_eq!(refresher.calls.load(Ordering::SeqCst), 0); assert_eq!( stored.access_token.expose_secret(), @@ -96,6 +97,88 @@ async fn token_refresh_task_should_skip_duplicate_in_flight_refresh() { refresh_token: None, }), }; + let task = codex_proxy_rs::fleet::refresh::TokenRefreshService::new( + store.clone(), + RefreshPolicy { + refresh_margin_seconds: 300, + refresh_concurrency: 1, + }, + refresher.clone(), + ); + + let started = refresher.started.notified(); + tokio::pin!(started); + let first = task + .schedule_account_timers_once_at(now) + .await + .expect("first refresh timer should be scheduled"); + timeout(StdDuration::from_secs(2), &mut started) + .await + .expect("first refresh should start"); + + let Ok(second) = timeout( + StdDuration::from_millis(100), + task.schedule_account_timers_once_at(now), + ) + .await + else { + refresher.release.notify_waiters(); + task.shutdown().await; + panic!("second timer scan should not wait on the in-flight refresh"); + }; + let second = second.expect("second refresh timer should be scheduled"); + wait_for_no_scheduled_timers(&task).await; + + refresher.release.notify_waiters(); + task.shutdown().await; + let stored = wait_for_account(&store, "acct-refresh-in-flight", |account| { + account.access_token.expose_secret() == new_access_token.as_str() + }) + .await; + + assert_eq!(first.immediate, 1); + assert_eq!(second.immediate, 1); + assert_eq!(refresher.calls.load(Ordering::SeqCst), 1); + assert_eq!(stored.status, AccountStatus::Active); + assert_eq!( + stored.access_token.expose_secret(), + new_access_token.as_str() + ); +} + +#[tokio::test] +async fn token_refresh_shutdown_should_wait_for_in_flight_refresh() { + let (pool, _guard) = init_test_db("token-refresh-shutdown").await; + let store = PgAccountStore::new(pool); + let now = Utc::now(); + let new_access_token = test_jwt((now + Duration::hours(1)).timestamp()); + store + .insert(NewAccount { + id: "acct-refresh-shutdown".to_string(), + email: Some("shutdown@example.com".to_string()), + account_id: Some("chatgpt-shutdown".to_string()), + user_id: Some("user-shutdown".to_string()), + label: None, + plan_type: Some("plus".to_string()), + access_token: SecretString::new( + test_jwt((now + Duration::seconds(30)).timestamp()).into(), + ), + refresh_token: Some(SecretString::new("refresh-shutdown".to_string().into())), + added_at: None, + access_token_expires_at: Some(now + Duration::seconds(30)), + status: AccountStatus::Active, + }) + .await + .expect("account should be inserted"); + let refresher = BlockingTokenRefresher { + calls: Arc::new(AtomicUsize::new(0)), + started: Arc::new(Notify::new()), + release: Arc::new(Notify::new()), + response: Ok(TokenPair { + access_token: new_access_token.clone(), + refresh_token: None, + }), + }; let task = Arc::new(codex_proxy_rs::fleet::refresh::TokenRefreshService::new( store.clone(), RefreshPolicy { @@ -104,45 +187,34 @@ async fn token_refresh_task_should_skip_duplicate_in_flight_refresh() { }, refresher.clone(), )); + let started = refresher.started.notified(); + tokio::pin!(started); + task.schedule_account_timers_once_at(now) + .await + .expect("refresh timer should be scheduled"); + timeout(StdDuration::from_secs(2), &mut started) + .await + .expect("refresh should start"); - let first = { - let task = task.clone(); - tokio::spawn(async move { task.refresh_due_accounts_once_at(now).await }) - }; - refresher.started.notified().await; - - let second = { - let task = task.clone(); - tokio::spawn(async move { task.refresh_due_accounts_once_at(now).await }) - }; - let Ok(result) = timeout(StdDuration::from_millis(100), second).await else { - refresher.release.notify_waiters(); - let _ = first.await; - panic!("second scan should skip instead of waiting on duplicate refresh"); - }; - let second = result - .expect("second scan should join") - .expect("second scan should succeed"); + let shutdown_task = task.clone(); + let shutdown = tokio::spawn(async move { shutdown_task.shutdown().await }); + tokio::pin!(shutdown); + assert!(timeout(StdDuration::from_millis(100), &mut shutdown) + .await + .is_err()); refresher.release.notify_waiters(); - let first = first + timeout(StdDuration::from_secs(2), &mut shutdown) .await - .expect("first scan should join") - .expect("first scan should succeed"); - let stored = store - .get("acct-refresh-in-flight") - .await - .expect("account should load") - .expect("account should exist"); + .expect("shutdown should finish after refresh") + .expect("shutdown task should join"); + let stored = wait_for_account(&store, "acct-refresh-shutdown", |account| { + account.access_token.expose_secret() == new_access_token.as_str() + }) + .await; - assert_eq!(first.refreshed, 1); - assert_eq!(second.skipped, 1); assert_eq!(refresher.calls.load(Ordering::SeqCst), 1); - assert_eq!(stored.status, AccountStatus::Active); - assert_eq!( - stored.access_token.expose_secret(), - new_access_token.as_str() - ); + assert_eq!(stored.access_token.expose_secret(), new_access_token); } #[tokio::test] @@ -198,6 +270,7 @@ async fn token_refresh_task_should_schedule_future_per_account_timer_without_ref assert_ne!(scheduled_at, exact_margin_at); assert!(scheduled_at >= now + Duration::seconds(255)); assert!(scheduled_at <= now + Duration::seconds(345)); + task.shutdown().await; } #[tokio::test] @@ -268,6 +341,7 @@ async fn token_refresh_task_should_cancel_future_timer_without_refreshing_expire ); assert!(stored.next_refresh_at.is_none()); assert_eq!(task.scheduled_timer_count().await, 0); + task.shutdown().await; } #[tokio::test] @@ -327,9 +401,9 @@ async fn token_refresh_task_should_clear_stale_refresh_time_without_refresh_toke .expect("stale due refresh time should persist"); let due_summary = task - .refresh_due_accounts_once_at(now) + .schedule_account_timers_once_at(now) .await - .expect("due refresh scan should skip account without refresh token"); + .expect("due timer scan should skip account without refresh token"); let stored = store .get("acct-refresh-no-rt-clear-next") .await @@ -346,6 +420,7 @@ async fn token_refresh_task_should_clear_stale_refresh_time_without_refresh_toke stored.access_token.expose_secret(), old_access_token.as_str() ); + task.shutdown().await; } #[tokio::test] @@ -398,12 +473,10 @@ async fn token_refresh_task_should_fire_per_account_timer_at_refresh_time() { timeout(StdDuration::from_secs(2), &mut started) .await .expect("scheduled timer should trigger refresh"); - tokio::task::yield_now().await; - let stored = store - .get("acct-refresh-timer") - .await - .expect("account should load") - .expect("account should exist"); + let stored = wait_for_account(&store, "acct-refresh-timer", |account| { + account.access_token.expose_secret() == new_access_token.as_str() + }) + .await; assert_eq!(summary.immediate, 1); assert_eq!(refresher.calls.load(Ordering::SeqCst), 1); @@ -413,6 +486,7 @@ async fn token_refresh_task_should_fire_per_account_timer_at_refresh_time() { stored.access_token.expose_secret(), new_access_token.as_str() ); + task.shutdown().await; } #[tokio::test] @@ -468,22 +542,10 @@ async fn token_refresh_task_should_fire_quota_exhausted_timer_without_clearing_s .await .expect("quota-exhausted timer should trigger refresh"); - let mut stored = store - .get("acct-refresh-quota-exhausted-timer") - .await - .expect("account should load") - .expect("account should exist"); - for _ in 0..50 { - if stored.access_token.expose_secret() == new_access_token.as_str() { - break; - } - sleep(StdDuration::from_millis(20)).await; - stored = store - .get("acct-refresh-quota-exhausted-timer") - .await - .expect("account should load") - .expect("account should exist"); - } + let stored = wait_for_account(&store, "acct-refresh-quota-exhausted-timer", |account| { + account.access_token.expose_secret() == new_access_token.as_str() + }) + .await; assert_eq!(summary.scheduled, 1); assert_eq!(refresher.calls.load(Ordering::SeqCst), 1); @@ -493,6 +555,7 @@ async fn token_refresh_task_should_fire_quota_exhausted_timer_without_clearing_s new_access_token.as_str() ); assert!(stored.next_refresh_at.is_some()); + task.shutdown().await; } #[tokio::test] @@ -546,22 +609,10 @@ async fn token_refresh_task_should_reschedule_next_timer_after_scheduled_refresh .await .expect("scheduled timer should trigger refresh"); - let mut stored = store - .get("acct-refresh-next-timer") - .await - .expect("account should load") - .expect("account should exist"); - for _ in 0..50 { - if stored.access_token.expose_secret() == new_access_token.as_str() { - break; - } - sleep(StdDuration::from_millis(20)).await; - stored = store - .get("acct-refresh-next-timer") - .await - .expect("account should load") - .expect("account should exist"); - } + let stored = wait_for_account(&store, "acct-refresh-next-timer", |account| { + account.access_token.expose_secret() == new_access_token.as_str() + }) + .await; assert_eq!(summary.scheduled, 1); assert_eq!(refresher.calls.load(Ordering::SeqCst), 1); @@ -570,4 +621,5 @@ async fn token_refresh_task_should_reschedule_next_timer_after_scheduled_refresh new_access_token.as_str() ); assert_eq!(task.scheduled_timer_count().await, 1); + task.shutdown().await; } diff --git a/backend/tests/fleet/refresh/success.rs b/backend/tests/fleet/refresh/success.rs index b0453930f..c99b2d5e5 100644 --- a/backend/tests/fleet/refresh/success.rs +++ b/backend/tests/fleet/refresh/success.rs @@ -40,19 +40,19 @@ async fn token_refresh_task_should_persist_refreshed_access_token_and_keep_refre refresher, ); - let summary = task - .refresh_due_accounts_once_at(now) + let timer_summary = task + .schedule_account_timers_once_at(now) .await - .expect("refresh should succeed"); - let stored = store - .get("acct-refresh") - .await - .expect("account should load") - .expect("account should exist"); + .expect("refresh timer should be scheduled"); + let stored = wait_for_account(&store, "acct-refresh", |account| { + account.access_token.expose_secret() == new_access_token.as_str() + }) + .await; + task.shutdown().await; assert_eq!( ( - summary.refreshed, + timer_summary.immediate, stored.access_token.expose_secret(), stored .refresh_token @@ -111,17 +111,17 @@ async fn token_refresh_task_should_persist_rotated_refresh_token_when_returned() refresher, ); - let summary = task - .refresh_due_accounts_once_at(now) + let timer_summary = task + .schedule_account_timers_once_at(now) .await - .expect("refresh should succeed"); - let stored = store - .get("acct-refresh-rotated") - .await - .expect("account should load") - .expect("account should exist"); + .expect("refresh timer should be scheduled"); + let stored = wait_for_account(&store, "acct-refresh-rotated", |account| { + account.access_token.expose_secret() == new_access_token.as_str() + }) + .await; + task.shutdown().await; - assert_eq!(summary.refreshed, 1); + assert_eq!(timer_summary.immediate, 1); assert_eq!(stored.access_token.expose_secret(), new_access_token); assert_eq!( stored @@ -178,18 +178,18 @@ async fn token_refresh_task_should_refresh_quota_exhausted_account_without_clear refresher, ); - let summary = task - .refresh_due_accounts_once_at(now) + let timer_summary = task + .schedule_account_timers_once_at(now) .await - .expect("quota-exhausted account should refresh token"); - let stored = store - .get("acct-refresh-quota-exhausted") - .await - .expect("account should load") - .expect("account should exist"); + .expect("quota-exhausted refresh timer should be scheduled"); + let stored = wait_for_account(&store, "acct-refresh-quota-exhausted", |account| { + account.access_token.expose_secret() == new_access_token.as_str() + }) + .await; + task.shutdown().await; let observed_statuses = observed_statuses.lock().await.clone(); - assert_eq!(summary.refreshed, 1); + assert_eq!(timer_summary.immediate, 1); assert_eq!(observed_statuses, [AccountStatus::QuotaExhausted]); assert_eq!(stored.status, AccountStatus::QuotaExhausted); assert_eq!( @@ -251,15 +251,15 @@ async fn token_refresh_task_should_not_clear_quota_limit_when_exhausted_token_re refresher, ); - let summary = task - .refresh_due_accounts_once_at(now) + let timer_summary = task + .schedule_account_timers_once_at(now) .await - .expect("quota-exhausted account should refresh token"); - let stored = store - .get("acct-exhausted-quota") - .await - .expect("account should load") - .expect("account should exist"); + .expect("quota-exhausted refresh timer should be scheduled"); + let stored = wait_for_account(&store, "acct-exhausted-quota", |account| { + account.access_token.expose_secret() == new_access_token.as_str() + }) + .await; + task.shutdown().await; let quota_state: (bool, Option>) = sqlx::query_as( "select quota_limit_reached, quota_cooldown_until from accounts where id = $1", ) @@ -268,7 +268,7 @@ async fn token_refresh_task_should_not_clear_quota_limit_when_exhausted_token_re .await .unwrap(); - assert_eq!(summary.refreshed, 1); + assert_eq!(timer_summary.immediate, 1); assert_eq!(stored.status, AccountStatus::QuotaExhausted); assert_eq!( stored.access_token.expose_secret(), @@ -325,19 +325,18 @@ async fn token_refresh_task_should_persist_success_after_transient_transport_ret ) .with_retry_delays(vec![StdDuration::ZERO; 4]); - let summary = task - .refresh_due_accounts_once_at(now) + let timer_summary = task + .schedule_account_timers_once_at(now) .await - .expect("retry should eventually refresh"); - let stored = store - .get("acct-refresh-transient-success") - .await - .expect("account should load") - .expect("account should exist"); + .expect("refresh timer should be scheduled"); + let stored = wait_for_account(&store, "acct-refresh-transient-success", |account| { + account.access_token.expose_secret() == new_access_token.as_str() + }) + .await; + task.shutdown().await; let observed_statuses = observed_statuses.lock().await.clone(); - assert_eq!(summary.refreshed, 1); - assert_eq!(summary.failed, 0); + assert_eq!(timer_summary.immediate, 1); assert_eq!(observed_statuses, [AccountStatus::Active; 2]); assert_eq!(stored.status, AccountStatus::Active); assert_eq!( diff --git a/backend/tests/infra/logging.rs b/backend/tests/infra/logging.rs index b13557289..0a21b765d 100644 --- a/backend/tests/infra/logging.rs +++ b/backend/tests/infra/logging.rs @@ -2,10 +2,17 @@ use std::{fs, io::Write}; use chrono::{Duration, Utc}; use codex_proxy_rs::infra::{ - logging::{build_file_appender, RotationConfig}, + logging::{build_file_appender, init_tracing, RotationConfig, TracingConfig}, time::china_date, }; +#[test] +fn tracing_config_should_reject_disabled_outputs() { + let result = init_tracing(&TracingConfig::new("info", false, None)); + + assert!(result.is_err()); +} + #[test] fn rolling_appender_should_write_daily_codex_log_file() { let dir = tempfile::tempdir().unwrap(); diff --git a/backend/tests/upstream/openai/transport/websocket_pool.rs b/backend/tests/upstream/openai/transport/websocket_pool.rs index ce8774b9f..f0c67dabd 100644 --- a/backend/tests/upstream/openai/transport/websocket_pool.rs +++ b/backend/tests/upstream/openai/transport/websocket_pool.rs @@ -601,6 +601,68 @@ async fn websocket_pool_should_bypass_busy_key_with_one_shot_connections() { ); } +#[tokio::test] +async fn websocket_pool_should_release_fresh_reservation_when_prefetch_is_cancelled() { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + let (first_request_received_tx, first_request_received_rx) = tokio::sync::oneshot::channel(); + let server = tokio::spawn(async move { + let (first_stream, _) = listener.accept().await.unwrap(); + let mut first_websocket = accept_codex_test_websocket(first_stream).await; + let _first_message = first_websocket.next().await.unwrap().unwrap(); + first_request_received_tx.send(()).unwrap(); + + let (second_stream, _) = listener.accept().await.unwrap(); + let mut second_websocket = accept_codex_test_websocket(second_stream).await; + let _second_message = second_websocket.next().await.unwrap().unwrap(); + second_websocket + .send(Message::Text( + completed_websocket_response("resp_after_cancelled_prefetch", 2, 1).into(), + )) + .await + .unwrap(); + second_websocket.close(None).await.unwrap(); + drop(first_websocket); + }); + let pool = Arc::new(CodexWebSocketPool::default()); + let backend = CodexBackendClient::new( + reqwest::Client::builder().no_proxy().build().unwrap(), + format!("http://{addr}"), + crate::support::fingerprint::runtime_test_fingerprint(), + ) + .with_websocket_pool(pool); + let request = pooled_websocket_request("conversation-cancelled-prefetch"); + + { + let first_request = backend.create_response_stream( + &request, + request_context("req_cancelled_prefetch", Some("chatgpt-account")), + ); + tokio::pin!(first_request); + tokio::select! { + _ = &mut first_request => { + panic!("prefetch request unexpectedly completed"); + } + result = first_request_received_rx => { + result.expect("server should observe the first request"); + } + } + } + tokio::time::sleep(Duration::from_millis(100)).await; + + let second = backend + .create_response( + &request, + request_context("req_after_cancelled_prefetch", Some("chatgpt-account")), + ) + .await + .expect("request after cancelled prefetch should succeed"); + server.await.unwrap(); + + assert!(second.body.contains("resp_after_cancelled_prefetch")); + assert_eq!(second.websocket_pool_decision.unwrap().kind(), "new"); +} + #[tokio::test] async fn websocket_pool_should_release_slot_when_client_drops_stream() { let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); diff --git a/docs/concurrency-review.md b/docs/concurrency-review.md new file mode 100644 index 000000000..e46f8f19a --- /dev/null +++ b/docs/concurrency-review.md @@ -0,0 +1,137 @@ +# 并发与错误处理审查 + +> 审查分支：`feat/postgres-redis-migration` +> 最终复核：2026-07-11 +> 覆盖范围：账号池、WebSocket 连接池、token 刷新、HTTP 停机、日志回灌、更新回滚及后台清理任务。 + +## 结论 + +本轮确认的取消泄漏、停机数据窗口、timer 竞态和错误吞没均已修复。当前并发设计遵守以下不变量： + +1. 每次资源占用都有唯一 ID 和唯一 lease，释放只命中该 lease 持有的资源。 +2. 请求 future 被取消时，lease 的 `Drop` 路径负责回收资源；正常完成路径消费 lease，二者不会重复释放其他并发请求。 +3. token 刷新只有账号级 timer 一条生产执行路径，不保留旧扫描兼容入口。 +4. HTTP 排水、后台任务关闭和 token 刷新收尾均有明确时间上限。 +5. `backend/src` 不包含测试模块或测试函数；测试统一位于 `backend/tests`。 + +## 问题处理结果 + +### P0：WebSocket Busy reservation 取消泄漏 + +状态：已修复。 + +- `WebSocketPoolReservation` 使用 UUID 标识单次占用，并记录 `reserved_at`。 +- `WebSocketPoolAcquire::Reused` 和 `FreshReserved` 都返回 `WebSocketPoolLease`。 +- `put`、`discard` 只在 reservation ID 匹配时修改 slot，旧任务不能覆盖新 reservation。 +- 请求在 connect、send 或首帧 prefetch 期间被取消时，lease 的 `Drop` 会条件删除自己的 Busy slot。 +- 维护任务会清理超过 `max_age` 的异常 Busy reservation，作为进程内最终兜底。 +- 流式转发任务接管 lease 后，连接归还权随任务生命周期移动。 + +回归测试：`backend/tests/upstream/openai/transport/websocket_pool.rs`，覆盖首帧 prefetch 取消、下游主动断开、复用和池容量限制。 + +### P1：账号池 slot 取消泄漏 + +状态：已修复。 + +- 每个 `AccountSlot` 使用 UUID 标识，不再按账号 ID 模糊释放队首 slot。 +- `AccountPoolService::acquire_with` 返回 `AccountLease`；lease 在状态持久化 await 之前创建，覆盖 acquire 内部取消窗口。 +- `AccountLease::complete` 精确释放 slot 并记录真实上游请求用量。 +- `AccountLease::release_without_usage` 用于尚未发起上游请求的配额预检分支。 +- lease 的 `Drop` 精确释放自己的 slot，不会释放同账号的其他并发请求。 +- 模型刷新使用独立的 `DistinctPlanAccountLease`，不保留旧 `release(account_id)` 兼容 API。 +- 流式任务持有完整 lease；下游断开、任务取消或错误退出都会回收 slot。 + +回归测试：`backend/tests/fleet/pool/`，覆盖并发 slot 精确释放、lease Drop 容量恢复和运行时用量持久化。 + +### P1：HTTP 和后台任务停机无界 + +状态：已修复。 + +- 第一次关闭信号触发 Axum graceful shutdown。 +- HTTP 在途请求最多排水 20 秒；第二次关闭信号立即结束排水。 +- 排水超时后先丢弃 serve future，取消剩余连接，再关闭后台服务。 +- 后台任务并行关闭；单任务上限 5 秒，协调器总上限 6 秒。 +- 总预算不超过 Compose 的 `stop_grace_period: 30s`。 + +### P1：token 刷新任务不可见及数据丢失窗口 + +状态：已修复。 + +- 所有 timer 和已开始的刷新任务统一注册到 `TaskTracker`。 +- 到期账号不再在周期 tick 内同步执行，而是进入相同的 tracked timer 路径。 +- `shutdown` 停止新调度、取消仍在睡眠的 timer，并等待已进入刷新流程的任务结束。 +- Redis 刷新租约使用 owner 校验释放；取消清理任务也由同一个 `TaskTracker` 管理。 +- `in_flight` 使用同步 RAII guard，panic 和取消都能移除账号 ID。 +- 已删除无生产调用方的 `refresh_due_accounts_once_at` 和 `TokenRefreshSummary`，不存在第二套测试专用执行路径。 + +回归测试：`backend/tests/fleet/refresh/`，覆盖结构化关闭、重复 in-flight 跳过、租约竞争、重试、恢复时间和 timer 替换。 + +### P1：账号选择锁内开销随账号数放大 + +状态：已修复。 + +- `access_token_expires_at` 存在时直接比较时间；仅在该字段缺失时解析 JWT。 +- 候选过滤返回账号借用，不再克隆全部候选。 +- Smart、QuotaResetPriority、RoundRobin 和 Sticky 均在借用切片上选择，只克隆最终账号。 +- 按订阅计划选择模型刷新账号时同样先借用分组，最终选择后才克隆。 +- 保留账号池单锁，不引入会破坏跨账号调度一致性的分片锁。 + +回归测试：`backend/tests/fleet/pool/selection.rs` 和 `backend/tests/fleet/scheduler/`。 + +### P2：timer 调度竞态 + +状态：已修复。 + +- `schedule_account_timer` 的检查、旧 timer 取消、新任务创建和 map 写入在一次锁持有期间完成。 +- 每个 timer 使用 UUID；任务醒来时仅能移除与自身 UUID 匹配的 entry。 +- 旧 timer 不能删除或覆盖同账号的新 timer。 + +### P2：错误处理和边界问题 + +状态：已修复。 + +- 管理端 logout 删除 Redis 会话失败时记录 `warn`，仍清除客户端 cookie。 +- API key `last_used_at` flush 失败回灌按时间取最大值，不覆盖并发写入的新时间戳。 +- 自更新回滚失败不再静默；主错误会携带所有回滚失败上下文。 +- `RuntimeRefreshPolicy` 仅在并发数变化时替换 semaphore，普通设置更新不会瞬时放大并发。 +- WebSocket pump 入站队列改为容量 64；缓冲满时关闭异常连接，不阻塞 pump 命令处理。 +- reqwest client cache miss 的 TLS 构建和 CA 文件读取移到同步 Mutex 之外。 +- Cloudflare challenge cookie 的清理时间持久化到 `account_cookies.expires_at`；冷却期间保留，截止后读取失效并由周期任务删除，重启不会丢失清理计划。 +- Cloudflare path-block 仍立即清除 cookie，因为该信号表示当前路径状态已不可继续使用。 + +## 锁与生命周期复核 + +- std `Mutex` / `RwLock` guard 不跨 `.await`。 +- 需要跨 `.await` 的共享状态使用 Tokio 锁。 +- 账号池锁内只处理内存状态，PostgreSQL 和 Redis 写入均在锁外。 +- WebSocket pool 的连接关闭在锁外执行。 +- 账号 slot、WebSocket reservation、token in-flight 和 Redis refresh lease 都有取消兜底。 +- crate 已启用 `#![forbid(unsafe_code)]`，生产代码无 `unsafe` 路径。 + +## 验证 + +必须通过以下门禁后才可提交： + +```bash +cargo fmt --manifest-path backend/Cargo.toml -- --check +cargo clippy --manifest-path backend/Cargo.toml --all-targets --all-features --locked -- -D warnings +CPR_TEST_REDIS_URL='redis://:codex_proxy@127.0.0.1:6379' \ + cargo test --manifest-path backend/Cargo.toml --all-features --locked +pnpm --dir frontend format:check +pnpm --dir frontend build +docker build -f deploy/Dockerfile --target runtime -t codex-proxy-rs:concurrency-review . +``` + +2026-07-11 最终验证结果： + +- Rust format：通过。 +- Rust Clippy：`--all-targets --all-features -D warnings` 通过，0 warning。 +- 后端完整测试：630 passed，0 failed。 +- 前端 Prettier、`vue-tsc`、Vite build：全部通过。 +- RustSec `cargo audit`：扫描 382 个依赖，0 vulnerability。 +- 前端生产依赖审计：0 known vulnerability。 +- `backend/src` 测试属性扫描：0 处。 +- runtime 镜像：`codex-proxy-rs:concurrency-review`，镜像 ID `sha256:96a9bb2a828b16ddae55af1c3923255d4167c6c33c341b7146b62c857c30f84b`。 +- 容器健康检查：`/healthz` 返回 204，PostgreSQL 和 Redis 均为 healthy。 +- 真实 SIGTERM 重启：0.319 秒完成，全部后台任务按关闭流程退出。 +- 本次启动及重启日志：0 条 `WARN`，0 条 `ERROR`。