diff --git a/docs/scheduler-usage.md b/docs/scheduler-usage.md index fa9fe072e..c039e590b 100644 --- a/docs/scheduler-usage.md +++ b/docs/scheduler-usage.md @@ -169,6 +169,7 @@ admin: 2. 账户是否有 `refresh_token` 3. 账户状态是否为 `active` 或 `quota_exhausted` 4. 查看日志中的 JWT 解析错误 +5. 如果日志显示 `refresh lease held`，说明同一账户已有其他进程持有刷新租约，当前进程会跳过本次刷新 ### 问题：会话没有被清理 @@ -191,7 +192,9 @@ admin: - **定时器开销**: 每个账户一个定时器，对于大量账户建议分批调度 - **数据库查询**: `list_all_for_refresh()` 在启动时调用一次，不影响运行时性能 -- **并发控制**: `refresh_concurrency` 限制同时刷新的账户数，避免突发请求 +- **并发控制**: `refresh_concurrency` 已接入后台刷新生产路径，限制本进程同时刷新的账户数，避免突发请求 +- **刷新抖动**: 自动刷新和恢复刷新会加入 80%-120% jitter，降低大量账户同一时间命中 OAuth 上游的概率 +- **跨进程租约**: `account_refresh_leases` 表按账户持有 5 分钟刷新租约，防止多进程同时消费同一个 refresh_token ## 最佳实践 diff --git a/src/admin/http/accounts/mutate.rs b/src/admin/http/accounts/mutate.rs index ae6ee2527..beac8bf13 100644 --- a/src/admin/http/accounts/mutate.rs +++ b/src/admin/http/accounts/mutate.rs @@ -255,6 +255,12 @@ fn refresh_account_error(error: RefreshAccountError, request_id: &str) -> AdminE "Account not found", request_id, ), + RefreshAccountError::LeaseAcquire => AdminError::new( + StatusCode::INTERNAL_SERVER_ERROR, + 50001, + "Failed to acquire refresh lease", + request_id, + ), RefreshAccountError::TokenRefresherUnavailable => AdminError::new( StatusCode::INTERNAL_SERVER_ERROR, 50001, diff --git a/src/codex/accounts/model.rs b/src/codex/accounts/model.rs index 813650168..abb07da8e 100644 --- a/src/codex/accounts/model.rs +++ b/src/codex/accounts/model.rs @@ -26,6 +26,11 @@ pub struct Account { pub quota_limit_reached: bool, pub quota_cooldown_until: Option>, pub cloudflare_cooldown_until: Option>, + pub request_count: u64, + pub window_request_count: u64, + pub window_started_at: Option>, + pub window_reset_at: Option>, + pub limit_window_seconds: Option, pub added_at: String, pub last_used_at: Option, } @@ -46,6 +51,11 @@ impl Account { quota_limit_reached: false, quota_cooldown_until: None, cloudflare_cooldown_until: None, + request_count: 0, + window_request_count: 0, + window_started_at: None, + window_reset_at: None, + limit_window_seconds: None, added_at: Utc::now().to_rfc3339(), last_used_at: None, } diff --git a/src/codex/accounts/pool.rs b/src/codex/accounts/pool.rs index 7aa4c0d28..3e426243e 100644 --- a/src/codex/accounts/pool.rs +++ b/src/codex/accounts/pool.rs @@ -1,4 +1,7 @@ -use std::collections::{BTreeMap, VecDeque}; +use std::{ + cmp::Ordering, + collections::{BTreeMap, VecDeque}, +}; use chrono::{DateTime, Duration, Utc}; @@ -86,6 +89,7 @@ pub struct AccountPool { slots: BTreeMap>>, options: AccountPoolOptions, round_robin_cursor: usize, + least_used_cursor: usize, } impl Default for AccountPool { @@ -101,6 +105,7 @@ impl AccountPool { slots: BTreeMap::new(), options, round_robin_cursor: 0, + least_used_cursor: 0, } } @@ -118,6 +123,7 @@ impl AccountPool { self.accounts.clear(); self.slots.clear(); self.round_robin_cursor = 0; + self.least_used_cursor = 0; } pub fn set_model_plan_allowlist(&mut self, allowlist: BTreeMap>) { @@ -148,6 +154,10 @@ impl AccountPool { return false; }; account.last_used_at = None; + account.request_count = 0; + account.window_request_count = 0; + account.window_started_at = None; + account.window_reset_at = None; true } @@ -161,6 +171,13 @@ impl AccountPool { }; account.quota_limit_reached = true; account.quota_cooldown_until = Some(cooldown_until); + account.window_reset_at = Some(cooldown_until); + if let Ok(seconds) = (cooldown_until - Utc::now()) + .to_std() + .map(|duration| duration.as_secs()) + { + account.limit_window_seconds = Some(seconds); + } self.slots.remove(account_id); true } @@ -199,8 +216,12 @@ impl AccountPool { RotationStrategy::RoundRobin => self.select_round_robin(&candidates), RotationStrategy::Sticky => self.select_sticky(&candidates), })?; - let previous_slot_at = self.previous_slot_at(&selected.id); - self.push_slot(&selected.id, request.now); + let selected_id = selected.id.clone(); + let previous_slot_at = self.previous_slot_at(&selected_id); + self.push_slot(&selected_id, request.now); + let selected = self + .mark_usage(&selected_id, request.now) + .unwrap_or(selected); Some(AcquiredAccount { account: selected, previous_slot_at, @@ -258,16 +279,17 @@ impl AccountPool { } fn select_least_used(&mut self, candidates: &[Account]) -> Option { - let best_last_used = candidates + let best_key = candidates .iter() - .map(|account| account.last_used_at.as_deref()) - .min()?; - let tied = candidates + .map(LeastUsedGroupKey::from_account) + .min_by(compare_least_used_group)?; + let mut tied = candidates .iter() - .filter(|account| account.last_used_at.as_deref() == best_last_used) + .filter(|account| LeastUsedGroupKey::from_account(account) == best_key) .collect::>(); - let index = self.round_robin_cursor % tied.len(); - self.round_robin_cursor = (self.round_robin_cursor + 1) % tied.len(); + tied.sort_by(compare_lru_then_id); + let index = self.least_used_cursor % tied.len(); + self.least_used_cursor = self.least_used_cursor.wrapping_add(1); Some((*tied[index]).clone()) } @@ -340,6 +362,21 @@ impl AccountPool { .push_back(now); } + fn mark_usage(&mut self, account_id: &str, now: DateTime) -> Option { + let account = self.accounts.get_mut(account_id)?; + account.last_used_at = Some(now.to_rfc3339()); + account.request_count = account.request_count.saturating_add(1); + account.window_request_count = account.window_request_count.saturating_add(1); + if account.window_started_at.is_none() { + if let Some(seconds) = account.limit_window_seconds { + account.window_started_at = Some(now); + account.window_reset_at = + Some(now + Duration::seconds(seconds.min(i64::MAX as u64) as i64)); + } + } + Some(account.clone()) + } + fn cleanup_stale_slots(&mut self, now: DateTime) { let ttl = self.options.stale_slot_ttl; self.slots.retain(|_, slots| { @@ -367,3 +404,42 @@ impl AccountPool { }) } } + +#[derive(Debug, Clone, PartialEq, Eq)] +struct LeastUsedGroupKey { + quota_limited: bool, + window_reset_at: Option>, + request_count: u64, +} + +impl LeastUsedGroupKey { + fn from_account(account: &Account) -> Self { + Self { + quota_limited: account.quota_limit_reached, + window_reset_at: account.window_reset_at, + request_count: account.request_count, + } + } +} + +fn compare_least_used_group(a: &LeastUsedGroupKey, b: &LeastUsedGroupKey) -> Ordering { + a.quota_limited + .cmp(&b.quota_limited) + .then_with(|| compare_window_reset(a.window_reset_at, b.window_reset_at)) + .then_with(|| a.request_count.cmp(&b.request_count)) +} + +fn compare_window_reset(a: Option>, b: Option>) -> Ordering { + match (a, b) { + (Some(a), Some(b)) => a.cmp(&b), + (Some(_), None) => Ordering::Less, + (None, Some(_)) => Ordering::Greater, + (None, None) => Ordering::Equal, + } +} + +fn compare_lru_then_id(a: &&Account, b: &&Account) -> Ordering { + a.last_used_at + .cmp(&b.last_used_at) + .then_with(|| a.id.cmp(&b.id)) +} diff --git a/src/codex/accounts/repository.rs b/src/codex/accounts/repository.rs index 8653d3393..e98589b74 100644 --- a/src/codex/accounts/repository.rs +++ b/src/codex/accounts/repository.rs @@ -322,7 +322,7 @@ impl AccountRepository { pub async fn list_pool_accounts(&self) -> AccountRepositoryResult> { let rows = sqlx::query( - "select accounts.id, email, accounts.account_id, user_id, label, plan_type, access_token_cipher, refresh_token_cipher, access_token_expires_at, status, added_at, updated_at, quota_limit_reached, quota_cooldown_until, cloudflare_cooldown_until, account_usage.last_used_at as usage_last_used_at from accounts left join account_usage on account_usage.account_id = accounts.id order by added_at desc, accounts.id desc", + "select accounts.id, email, accounts.account_id, user_id, label, plan_type, access_token_cipher, refresh_token_cipher, access_token_expires_at, status, added_at, updated_at, quota_limit_reached, quota_cooldown_until, cloudflare_cooldown_until, coalesce(account_usage.request_count, 0) as usage_request_count, account_usage.last_used_at as usage_last_used_at from accounts left join account_usage on account_usage.account_id = accounts.id order by added_at desc, accounts.id desc", ) .fetch_all(&self.pool) .await?; @@ -352,6 +352,11 @@ impl AccountRepository { cloudflare_cooldown_until: parse_optional_rfc3339( row.get::, _>("cloudflare_cooldown_until"), )?, + request_count: row.get::("usage_request_count").max(0) as u64, + window_request_count: 0, + window_started_at: None, + window_reset_at: None, + limit_window_seconds: None, added_at: stored.added_at.to_rfc3339(), last_used_at: row.get("usage_last_used_at"), }); @@ -581,6 +586,40 @@ impl AccountRepository { Ok(result.rows_affected() > 0) } + pub async fn try_acquire_refresh_lease( + &self, + account_id: &str, + owner: &str, + lease_until: DateTime, + ) -> AccountRepositoryResult { + let now = Utc::now().to_rfc3339(); + let result = sqlx::query( + "insert into account_refresh_leases (account_id, owner, expires_at, updated_at) values (?, ?, ?, ?) on conflict(account_id) do update set owner = excluded.owner, expires_at = excluded.expires_at, updated_at = excluded.updated_at where account_refresh_leases.expires_at <= ? or account_refresh_leases.owner = excluded.owner", + ) + .bind(account_id) + .bind(owner) + .bind(lease_until.to_rfc3339()) + .bind(&now) + .bind(now) + .execute(&self.pool) + .await?; + Ok(result.rows_affected() > 0) + } + + pub async fn release_refresh_lease( + &self, + account_id: &str, + owner: &str, + ) -> AccountRepositoryResult { + let result = + sqlx::query("delete from account_refresh_leases where account_id = ? and owner = ?") + .bind(account_id) + .bind(owner) + .execute(&self.pool) + .await?; + Ok(result.rows_affected() > 0) + } + fn account_from_row( &self, row: &sqlx::sqlite::SqliteRow, diff --git a/src/codex/accounts/service/import.rs b/src/codex/accounts/service/import.rs index a0b2c0415..aab2af7a6 100644 --- a/src/codex/accounts/service/import.rs +++ b/src/codex/accounts/service/import.rs @@ -192,6 +192,11 @@ impl AccountService { quota_limit_reached: false, quota_cooldown_until: None, cloudflare_cooldown_until: None, + request_count: 0, + window_request_count: 0, + window_started_at: None, + window_reset_at: None, + limit_window_seconds: None, added_at: now, last_used_at: None, }; diff --git a/src/codex/accounts/service/mod.rs b/src/codex/accounts/service/mod.rs index e801edfce..b66d719bc 100644 --- a/src/codex/accounts/service/mod.rs +++ b/src/codex/accounts/service/mod.rs @@ -197,6 +197,7 @@ pub enum RefreshAccountError { RepositoryUnavailable, Load, NotFound, + LeaseAcquire, TokenRefresherUnavailable, StoreRefreshed, } @@ -207,6 +208,7 @@ impl std::fmt::Display for RefreshAccountError { Self::RepositoryUnavailable => write!(f, "repository unavailable"), Self::Load => write!(f, "failed to load account"), Self::NotFound => write!(f, "account not found"), + Self::LeaseAcquire => write!(f, "failed to acquire refresh lease"), Self::TokenRefresherUnavailable => write!(f, "token refresher unavailable"), Self::StoreRefreshed => write!(f, "failed to store refreshed tokens"), } diff --git a/src/codex/accounts/service/refresh.rs b/src/codex/accounts/service/refresh.rs index bf0483b58..898879217 100644 --- a/src/codex/accounts/service/refresh.rs +++ b/src/codex/accounts/service/refresh.rs @@ -1,11 +1,15 @@ +use std::sync::Arc; + +use chrono::{Duration, Utc}; use secrecy::{ExposeSecret, SecretString}; +use uuid::Uuid; use crate::{ codex::accounts::{ model::{Account, AccountStatus}, repository::{AccountRepository, StoredAccount, TokenUpdate}, }, - codex::gateway::oauth::RefreshFailure, + codex::gateway::oauth::{RefreshFailure, TokenPair, TokenRefresher}, }; use super::{ @@ -13,6 +17,8 @@ use super::{ AccountProbeResult, AccountService, RefreshAccountError, }; +const REFRESH_LEASE_TTL_SECONDS: i64 = 5 * 60; + impl AccountService { pub async fn refresh_account( &self, @@ -30,16 +36,59 @@ impl AccountService { if account.status == AccountStatus::Disabled { return Ok(skipped_probe_result(&account, "manually disabled")); } + if account.refresh_token.is_none() { + return Ok(skipped_probe_result(&account, "no refresh token")); + } + let Some(refresher) = self.token_refresher.as_ref().cloned() else { + return Err(RefreshAccountError::TokenRefresherUnavailable); + }; + let lease_owner = refresh_lease_owner(); + let lease_until = Utc::now() + Duration::seconds(REFRESH_LEASE_TTL_SECONDS); + match repo + .try_acquire_refresh_lease(account_id, &lease_owner, lease_until) + .await + { + Ok(true) => {} + Ok(false) => { + return Ok(skipped_probe_result(&account, "refresh lease held")); + } + Err(_) => return Err(RefreshAccountError::LeaseAcquire), + } + + let result = self + .refresh_account_with_lease(repo, account_id, refresher) + .await; + if let Err(error) = repo.release_refresh_lease(account_id, &lease_owner).await { + tracing::warn!( + error = %error, + account_id = %account_id, + "释放 refresh lease 失败" + ); + } + result + } + + async fn refresh_account_with_lease( + &self, + repo: &AccountRepository, + account_id: &str, + refresher: Arc, + ) -> Result { + let account = match repo.get(account_id).await { + Ok(Some(account)) => account, + Ok(None) => return Err(RefreshAccountError::NotFound), + Err(_) => return Err(RefreshAccountError::Load), + }; + if account.status == AccountStatus::Disabled { + return Ok(skipped_probe_result(&account, "manually disabled")); + } let Some(refresh_token) = account.refresh_token.as_ref() else { return Ok(skipped_probe_result(&account, "no refresh token")); }; - let Some(refresher) = self.token_refresher.as_ref() else { - return Err(RefreshAccountError::TokenRefresherUnavailable); - }; - let started_at = std::time::Instant::now(); let previous_status = account.status; - match refresher.refresh(refresh_token.expose_secret()).await { + match refresh_with_latest_disk_token_retry(repo, &account, refresh_token, &refresher).await + { Ok(tokens) => { let updated = self .persist_refreshed_account( @@ -116,6 +165,56 @@ impl AccountService { } } +async fn refresh_with_latest_disk_token_retry( + repo: &AccountRepository, + account: &StoredAccount, + refresh_token: &SecretString, + refresher: &Arc, +) -> Result { + let used_token = refresh_token.expose_secret().to_string(); + match refresher.refresh(&used_token).await { + Err(RefreshFailure::InvalidGrant) => { + if let Some(latest_token) = latest_disk_refresh_token(repo, account, &used_token).await + { + tracing::info!( + account_id = %account.id, + "检测到磁盘 refresh_token 已更新，使用最新值重试刷新" + ); + return refresher.refresh(&latest_token).await; + } + Err(RefreshFailure::InvalidGrant) + } + result => result, + } +} + +async fn latest_disk_refresh_token( + repo: &AccountRepository, + account: &StoredAccount, + used_token: &str, +) -> Option { + match repo.get(&account.id).await { + Ok(Some(latest)) => latest + .refresh_token + .as_ref() + .map(|token| token.expose_secret().to_string()) + .filter(|token| token != used_token), + Ok(None) => None, + Err(error) => { + tracing::warn!( + error = %error, + account_id = %account.id, + "刷新失败后重新读取磁盘 refresh_token 失败" + ); + None + } + } +} + +fn refresh_lease_owner() -> String { + format!("{}:{}", std::process::id(), Uuid::new_v4()) +} + fn status_for_refresh_failure(failure: RefreshFailure) -> Option { match failure { RefreshFailure::InvalidGrant => Some(AccountStatus::Expired), diff --git a/src/codex/accounts/service/runtime_pool.rs b/src/codex/accounts/service/runtime_pool.rs index afd0cae27..d7e631fd8 100644 --- a/src/codex/accounts/service/runtime_pool.rs +++ b/src/codex/accounts/service/runtime_pool.rs @@ -54,6 +54,11 @@ pub(super) fn pool_account_from_stored(account: StoredAccount) -> Account { quota_limit_reached: false, quota_cooldown_until: None, cloudflare_cooldown_until: None, + request_count: 0, + window_request_count: 0, + window_started_at: None, + window_reset_at: None, + limit_window_seconds: None, added_at: account.added_at.to_rfc3339(), last_used_at: None, } diff --git a/src/codex/serving/dispatch/affinity.rs b/src/codex/serving/dispatch/affinity.rs index c0d863422..90c661246 100644 --- a/src/codex/serving/dispatch/affinity.rs +++ b/src/codex/serving/dispatch/affinity.rs @@ -1,3 +1,4 @@ +use serde::Serialize; use std::collections::HashMap; use std::sync::Arc; use std::time::{Duration, Instant}; @@ -41,6 +42,12 @@ pub struct SessionAffinityMap { const DEFAULT_TTL: Duration = Duration::from_secs(4 * 60 * 60); // 4小时 const CLEANUP_INTERVAL: Duration = Duration::from_secs(10 * 60); // 10分钟 +pub(super) fn compute_variant_hash(value: &impl Serialize) -> Option { + serde_json::to_string(value) + .ok() + .filter(|serialized| !serialized.is_empty()) +} + impl SessionAffinityMap { pub fn new(ttl: Duration) -> Self { Self { @@ -246,6 +253,8 @@ impl Drop for SessionAffinityMap { #[cfg(test)] mod tests { use super::*; + use crate::codex::gateway::transport::types::CodexResponsesRequest; + use serde_json::json; #[tokio::test] async fn session_affinity_records_and_retrieves_account() { @@ -358,4 +367,18 @@ mod tests { assert!(map.lookup_account("resp_123").await.is_none()); } + + #[test] + fn compute_variant_hash_should_follow_serialized_request_body() { + let mut request = CodexResponsesRequest::new_http_sse( + "gpt-5.5", + "You are Codex.", + vec![json!({"role": "user", "content": "hello"})], + ); + let first = compute_variant_hash(&request).unwrap(); + request.instructions = "You are Codex with different instructions.".to_string(); + let second = compute_variant_hash(&request).unwrap(); + + assert_ne!(first, second); + } } diff --git a/src/codex/serving/dispatch/mod.rs b/src/codex/serving/dispatch/mod.rs index 74b0e4e1a..84677d347 100644 --- a/src/codex/serving/dispatch/mod.rs +++ b/src/codex/serving/dispatch/mod.rs @@ -60,7 +60,7 @@ pub(crate) use self::{ stream::{completed_response_json, CollectedResponse}, }; -use self::affinity::SessionAffinityMap; +use self::affinity::{compute_variant_hash, SessionAffinityMap}; use self::{ refresh::refresh_account_after_unauthorized, routing::request_domain, @@ -680,6 +680,7 @@ async fn record_response_affinity_with_deps( } }; + let variant_hash = compute_variant_hash(request); deps.session_affinity .record( metadata.response_id, @@ -692,7 +693,7 @@ async fn record_response_affinity_with_deps( Some(&request.instructions), usage.map(|usage| usage.input_tokens), Some(metadata.function_call_ids), - None, + variant_hash, ) .await; } diff --git a/src/codex/tasks/refresh.rs b/src/codex/tasks/refresh.rs index bf9217b50..de804491c 100644 --- a/src/codex/tasks/refresh.rs +++ b/src/codex/tasks/refresh.rs @@ -1,5 +1,10 @@ +use rand::RngExt; use std::{collections::HashMap, sync::Arc, time::Duration}; -use tokio::{sync::RwLock, task::JoinHandle, time::Instant}; +use tokio::{ + sync::{RwLock, Semaphore}, + task::JoinHandle, + time::Instant, +}; use tracing::{debug, error, info, warn}; use crate::{ @@ -18,11 +23,13 @@ pub type SchedulerResult = Result; /// - 永久失败检测（invalid_grant / invalid_token） /// - 临时失败的恢复调度（10分钟） /// - 崩溃恢复：refreshing → 立即重试，expired + refreshToken → 延迟重试 +#[derive(Clone)] pub struct RefreshScheduler { account_service: Arc, config: AppConfig, timers: Arc>>>, in_flight: Arc>>, + refresh_permits: Arc, destroyed: Arc>, } @@ -44,11 +51,15 @@ const EXPIRED_ERRORS: &[&str] = &[ impl RefreshScheduler { pub fn new(account_service: Arc, config: AppConfig) -> Self { + let refresh_permits = Arc::new(Semaphore::new(refresh_concurrency_limit( + config.auth.refresh_concurrency, + ))); Self { account_service, config, timers: Arc::new(RwLock::new(HashMap::new())), in_flight: Arc::new(RwLock::new(HashMap::new())), + refresh_permits, destroyed: Arc::new(RwLock::new(false)), } } @@ -156,36 +167,14 @@ impl RefreshScheduler { return Ok(()); } - let delay = Duration::from_secs(refresh_at - now); + let delay = jitter_delay(Duration::from_secs(refresh_at - now)); info!( account_id = %account_id, delay_secs = delay.as_secs(), "已调度 token 刷新" ); - // 启动定时器 - 直接刷新，不再调用 do_refresh - let account_id_clone = account_id.clone(); - let account_service = self.account_service.clone(); - let in_flight = self.in_flight.clone(); - - let handle = tokio::spawn(async move { - tokio::time::sleep(delay).await; - - // 标记为正在处理 - { - let mut in_flight_lock = in_flight.write().await; - if in_flight_lock.contains_key(&account_id_clone) { - return; // 已经在刷新中 - } - in_flight_lock.insert(account_id_clone.clone(), Instant::now()); - } - - // 执行刷新 - let _ = account_service.refresh_account(&account_id_clone).await; - - // 清理标记 - in_flight.write().await.remove(&account_id_clone); - }); + let handle = self.spawn_delayed_refresh(account_id.clone(), delay); self.timers.write().await.insert(account_id, handle); Ok(()) @@ -219,32 +208,14 @@ impl RefreshScheduler { return Ok(()); } - let delay = Duration::from_secs(refresh_at - now); + let delay = jitter_delay(Duration::from_secs(refresh_at - now)); info!( account_id = %account_id, delay_secs = delay.as_secs(), "已调度下一次 token 刷新" ); - // 启动定时器 - let account_id_clone = account_id.clone(); - let account_service = self.account_service.clone(); - let in_flight = self.in_flight.clone(); - - let handle = tokio::spawn(async move { - tokio::time::sleep(delay).await; - - { - let mut in_flight_lock = in_flight.write().await; - if in_flight_lock.contains_key(&account_id_clone) { - return; - } - in_flight_lock.insert(account_id_clone.clone(), Instant::now()); - } - - let _ = account_service.refresh_account(&account_id_clone).await; - in_flight.write().await.remove(&account_id_clone); - }); + let handle = self.spawn_delayed_refresh(account_id.clone(), delay); self.timers.write().await.insert(account_id, handle); Ok(()) @@ -252,16 +223,21 @@ impl RefreshScheduler { /// 调度恢复尝试（用于临时失败） async fn schedule_recovery(&self, account_id: String, delay: Duration) { - let account_id_clone = account_id.clone(); - let account_service = self.account_service.clone(); - let handle = tokio::spawn(async move { - tokio::time::sleep(delay).await; - let _ = account_service.refresh_account(&account_id_clone).await; - }); + let handle = self.spawn_delayed_refresh(account_id.clone(), jitter_delay(delay)); self.timers.write().await.insert(account_id, handle); } + fn spawn_delayed_refresh(&self, account_id: String, delay: Duration) -> JoinHandle<()> { + let scheduler = self.clone(); + tokio::spawn(async move { + tokio::time::sleep(delay).await; + // 定时器触发后先移除自身，避免刷新成功后调度下一次时 abort 当前任务。 + scheduler.timers.write().await.remove(&account_id); + scheduler.do_refresh(account_id).await; + }) + } + /// 取消单个账户的定时器 async fn clear_one(&self, account_id: &str) { if let Some(handle) = self.timers.write().await.remove(account_id) { @@ -286,6 +262,19 @@ impl RefreshScheduler { in_flight.insert(account_id.clone(), Instant::now()); } + let _permit = match self.refresh_permits.clone().acquire_owned().await { + Ok(permit) => permit, + Err(error) => { + self.in_flight.write().await.remove(&account_id); + error!( + account_id = %account_id, + error = %error, + "refresh 并发控制已关闭" + ); + return; + } + }; + // 执行刷新逻辑 let result = self.do_refresh_inner(&account_id).await; @@ -378,13 +367,11 @@ impl RefreshScheduler { error = %e, "token 刷新重试次数已耗尽" ); - // 调度恢复尝试（不在循环内调用 await，避免递归） - let account_id_recovery = account_id.to_string(); - let account_service = self.account_service.clone(); - tokio::spawn(async move { - tokio::time::sleep(Duration::from_millis(RECOVERY_DELAY_MS)).await; - let _ = account_service.refresh_account(&account_id_recovery).await; - }); + self.schedule_recovery( + account_id.to_string(), + Duration::from_millis(RECOVERY_DELAY_MS), + ) + .await; return Err(SchedulerError::AccountNotFound(format!( "Max attempts reached: {}", e @@ -455,3 +442,42 @@ impl RefreshScheduler { self.in_flight.write().await.clear(); } } + +fn refresh_concurrency_limit(configured: u32) -> usize { + configured.max(1) as usize +} + +fn jitter_delay(delay: Duration) -> Duration { + let mut rng = rand::rng(); + jitter_delay_with_factor(delay, rng.random_range(0.8..=1.2)) +} + +fn jitter_delay_with_factor(delay: Duration, factor: f64) -> Duration { + let factor = factor.clamp(0.8, 1.2); + let millis = (delay.as_millis() as f64 * factor) + .round() + .min(u64::MAX as f64); + Duration::from_millis(millis as u64) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn refresh_concurrency_limit_should_never_be_zero() { + assert_eq!(refresh_concurrency_limit(0), 1); + } + + #[test] + fn jitter_delay_with_factor_should_keep_delay_inside_original_range() { + assert_eq!( + jitter_delay_with_factor(Duration::from_secs(100), 0.8), + Duration::from_secs(80) + ); + assert_eq!( + jitter_delay_with_factor(Duration::from_secs(100), 1.2), + Duration::from_secs(120) + ); + } +} diff --git a/src/platform/storage/schema.sql b/src/platform/storage/schema.sql index 4d0d6dfe3..8a33fa1e2 100644 --- a/src/platform/storage/schema.sql +++ b/src/platform/storage/schema.sql @@ -64,6 +64,18 @@ create table if not exists accounts ( create index if not exists idx_accounts_status on accounts(status); create index if not exists idx_accounts_identity on accounts(account_id, user_id) where account_id is not null; +-- ============================================ +-- Account Refresh Leases +-- ============================================ +create table if not exists account_refresh_leases ( + account_id text primary key references accounts(id) on delete cascade, + owner text not null, + expires_at text not null, + updated_at text not null +); + +create index if not exists idx_account_refresh_leases_expires on account_refresh_leases(expires_at); + -- ============================================ -- Account Usage Statistics -- ============================================ diff --git a/tests/admin/accounts/cookies_quota.rs b/tests/admin/accounts/cookies_quota.rs index 5d33d93f8..113c3f650 100644 --- a/tests/admin/accounts/cookies_quota.rs +++ b/tests/admin/accounts/cookies_quota.rs @@ -295,7 +295,11 @@ async fn admin_account_reset_usage_should_clear_local_counters_and_pool_last_use .acquire_runtime_account("gpt-5.5") .await .unwrap(); - assert_eq!(acquired.last_used_at, None); + assert_eq!(acquired.request_count, 1); + assert_ne!( + acquired.last_used_at.as_deref(), + Some("2026-06-12T12:00:00Z") + ); } #[tokio::test] diff --git a/tests/codex_accounts.rs b/tests/codex_accounts.rs index fe04ff251..47424b2ee 100644 --- a/tests/codex_accounts.rs +++ b/tests/codex_accounts.rs @@ -7,4 +7,5 @@ mod codex_accounts { mod refresh; mod refresh_scheduler; mod repository; + mod service_refresh; } diff --git a/tests/codex_accounts/pool_scheduling.rs b/tests/codex_accounts/pool_scheduling.rs index 5863267b5..b0dfe35f0 100644 --- a/tests/codex_accounts/pool_scheduling.rs +++ b/tests/codex_accounts/pool_scheduling.rs @@ -238,6 +238,92 @@ fn account_pool_should_rotate_tied_least_used_accounts() { assert_ne!(first.id, second.id); } +#[test] +fn least_used_should_deprioritize_quota_limited_accounts_when_skip_is_disabled() { + let mut pool = AccountPool::with_options(AccountPoolOptions { + skip_quota_limited: false, + rotation_strategy: RotationStrategy::LeastUsed, + ..AccountPoolOptions::default() + }); + let mut limited = Account::test("limited", AccountStatus::Active); + limited.quota_limit_reached = true; + let mut usable = Account::test("usable", AccountStatus::Active); + usable.request_count = 100; + pool.insert(limited); + pool.insert(usable); + + let acquired = pool.acquire("gpt-5.5").unwrap(); + + assert_eq!(acquired.id, "usable"); +} + +#[test] +fn least_used_should_prefer_earlier_rate_limit_window_reset() { + let now = fixed_time(); + let mut pool = AccountPool::with_options(AccountPoolOptions { + rotation_strategy: RotationStrategy::LeastUsed, + ..AccountPoolOptions::default() + }); + let mut soon = Account::test("soon", AccountStatus::Active); + soon.window_reset_at = Some(now + Duration::seconds(30)); + let mut later = Account::test("later", AccountStatus::Active); + later.window_reset_at = Some(now + Duration::seconds(300)); + pool.insert(later); + pool.insert(soon); + + let acquired = pool + .acquire_with(AccountAcquireRequest::new("gpt-5.5", now)) + .unwrap(); + + assert_eq!(acquired.account.id, "soon"); +} + +#[test] +fn least_used_should_prefer_lower_runtime_request_count() { + let now = fixed_time(); + let mut pool = AccountPool::with_options(AccountPoolOptions { + rotation_strategy: RotationStrategy::LeastUsed, + ..AccountPoolOptions::default() + }); + let mut busy = Account::test("busy", AccountStatus::Active); + busy.request_count = 10; + let mut quiet = Account::test("quiet", AccountStatus::Active); + quiet.request_count = 2; + pool.insert(busy); + pool.insert(quiet); + + let acquired = pool + .acquire_with(AccountAcquireRequest::new("gpt-5.5", now)) + .unwrap(); + + assert_eq!(acquired.account.id, "quiet"); +} + +#[test] +fn acquire_should_mark_runtime_usage_window() { + let now = fixed_time(); + let mut pool = AccountPool::default(); + let mut account = Account::test("acct_a", AccountStatus::Active); + account.limit_window_seconds = Some(60); + pool.insert(account); + + let acquired = pool + .acquire_with(AccountAcquireRequest::new("gpt-5.5", now)) + .unwrap(); + + assert_eq!(acquired.account.request_count, 1); + assert_eq!(acquired.account.window_request_count, 1); + assert_eq!(acquired.account.window_started_at, Some(now)); + assert_eq!( + acquired.account.window_reset_at, + Some(now + Duration::seconds(60)) + ); + assert_eq!( + acquired.account.last_used_at.as_deref(), + Some(now.to_rfc3339().as_str()) + ); +} + fn fixed_time() -> chrono::DateTime { Utc.with_ymd_and_hms(2026, 6, 11, 8, 0, 0).unwrap() } diff --git a/tests/codex_accounts/repository.rs b/tests/codex_accounts/repository.rs index 692781464..d5daa0b8d 100644 --- a/tests/codex_accounts/repository.rs +++ b/tests/codex_accounts/repository.rs @@ -190,6 +190,87 @@ async fn account_repository_should_accumulate_usage_counters() { assert!(usage.last_used_at.is_some()); } +#[tokio::test] +async fn account_repository_should_load_usage_count_into_runtime_pool_accounts() { + let dir = tempfile::tempdir().unwrap(); + let db = dir.path().join("accounts.sqlite"); + let url = format!("sqlite://{}", db.display()); + let pool = connect_sqlite(&url).await.unwrap(); + let repo = AccountRepository::new(pool.clone(), SecretBox::new([17u8; 32])); + + repo.insert(NewAccount { + id: "acct_a".to_string(), + email: None, + account_id: None, + user_id: None, + label: None, + plan_type: None, + access_token: SecretString::new("access-secret".to_string().into()), + refresh_token: None, + access_token_expires_at: None, + status: AccountStatus::Active, + }) + .await + .unwrap(); + repo.record_usage( + "acct_a", + UsageDelta { + input_tokens: 1, + output_tokens: 2, + cached_tokens: 3, + }, + ) + .await + .unwrap(); + + let account = repo.list_pool_accounts().await.unwrap().remove(0); + + assert_eq!(account.request_count, 1); +} + +#[tokio::test] +async fn account_repository_should_exclude_refresh_lease_until_expiry() { + let dir = tempfile::tempdir().unwrap(); + let db = dir.path().join("accounts.sqlite"); + let url = format!("sqlite://{}", db.display()); + let pool = connect_sqlite(&url).await.unwrap(); + let repo = AccountRepository::new(pool.clone(), SecretBox::new([18u8; 32])); + let now = Utc::now(); + repo.insert(NewAccount { + id: "acct_a".to_string(), + email: None, + account_id: None, + user_id: None, + label: None, + plan_type: None, + access_token: SecretString::new("access-secret".to_string().into()), + refresh_token: None, + access_token_expires_at: None, + status: AccountStatus::Active, + }) + .await + .unwrap(); + + assert!(repo + .try_acquire_refresh_lease("acct_a", "owner-a", now + Duration::minutes(5)) + .await + .unwrap()); + assert!(!repo + .try_acquire_refresh_lease("acct_a", "owner-b", now + Duration::minutes(5)) + .await + .unwrap()); + sqlx::query("update account_refresh_leases set expires_at = ? where account_id = ?") + .bind((now - Duration::seconds(1)).to_rfc3339()) + .bind("acct_a") + .execute(&pool) + .await + .unwrap(); + assert!(repo + .try_acquire_refresh_lease("acct_a", "owner-b", now + Duration::minutes(5)) + .await + .unwrap()); +} + #[tokio::test] async fn account_repository_should_preserve_refresh_token_when_update_omits_one() { let dir = tempfile::tempdir().unwrap(); diff --git a/tests/codex_accounts/service_refresh.rs b/tests/codex_accounts/service_refresh.rs new file mode 100644 index 000000000..e00666cdd --- /dev/null +++ b/tests/codex_accounts/service_refresh.rs @@ -0,0 +1,108 @@ +use std::{ + sync::{ + atomic::{AtomicUsize, Ordering}, + Arc, + }, + time::Duration as StdDuration, +}; + +use async_trait::async_trait; +use secrecy::SecretString; +use tokio::sync::Mutex; + +use codex_proxy_rs::{ + codex::{ + accounts::{ + model::AccountStatus, + pool::AccountPool, + repository::{AccountRepository, AccountUsageRepository, NewAccount, TokenUpdate}, + service::AccountService, + }, + gateway::oauth::{RefreshFailure, TokenPair, TokenRefresher}, + }, + platform::{crypto::SecretBox, storage::db::connect_sqlite}, +}; + +use crate::support::admin_accounts::test_config; + +#[tokio::test] +async fn account_service_refresh_should_retry_with_latest_disk_refresh_token_after_invalid_grant() { + let dir = tempfile::tempdir().unwrap(); + let db = dir.path().join("accounts.sqlite"); + let url = format!("sqlite://{}", db.display()); + let pool = connect_sqlite(&url).await.unwrap(); + let repo = AccountRepository::new(pool.clone(), SecretBox::new([19u8; 32])); + repo.insert(NewAccount { + id: "acct_stale".to_string(), + email: None, + account_id: None, + user_id: None, + label: None, + plan_type: None, + access_token: SecretString::new("access-old".to_string().into()), + refresh_token: Some(SecretString::new("refresh-old".to_string().into())), + access_token_expires_at: None, + status: AccountStatus::Active, + }) + .await + .unwrap(); + let calls = Arc::new(Mutex::new(Vec::new())); + let refresher = DiskRotatingRefresher { + repo: repo.clone(), + calls: calls.clone(), + first_call_seen: Arc::new(AtomicUsize::new(0)), + }; + let service = AccountService::new( + Arc::new(test_config(url)), + Some(repo), + Some(AccountUsageRepository::new(pool)), + None, + Some(Arc::new(refresher)), + Arc::new(Mutex::new(AccountPool::default())), + ); + + let result = service.refresh_account("acct_stale").await.unwrap(); + + assert_eq!(result.outcome.as_str(), "alive"); + assert_eq!( + calls.lock().await.as_slice(), + ["refresh-old", "refresh-latest"] + ); +} + +#[derive(Clone)] +struct DiskRotatingRefresher { + repo: AccountRepository, + calls: Arc>>, + first_call_seen: Arc, +} + +#[async_trait] +impl TokenRefresher for DiskRotatingRefresher { + async fn refresh(&self, refresh_token: &str) -> Result { + self.calls.lock().await.push(refresh_token.to_string()); + if refresh_token == "refresh-old" { + self.first_call_seen.fetch_add(1, Ordering::SeqCst); + self.repo + .update_tokens( + "acct_stale", + TokenUpdate { + access_token: SecretString::new("external-access".to_string().into()), + refresh_token: Some(SecretString::new("refresh-latest".to_string().into())), + access_token_expires_at: None, + }, + ) + .await + .unwrap(); + tokio::time::sleep(StdDuration::from_millis(5)).await; + return Err(RefreshFailure::InvalidGrant); + } + if refresh_token == "refresh-latest" { + return Ok(TokenPair { + access_token: "access-after-retry".to_string(), + refresh_token: Some("refresh-after-retry".to_string()), + }); + } + Err(RefreshFailure::InvalidGrant) + } +}