diff --git a/backend/src/upstream/accounts/pool.rs b/backend/src/upstream/accounts/pool.rs index d25cd3101..d03eed75e 100644 --- a/backend/src/upstream/accounts/pool.rs +++ b/backend/src/upstream/accounts/pool.rs @@ -1,7 +1,6 @@ //! 账号池调度策略。 use std::{ - cmp::Ordering, collections::{BTreeMap, BTreeSet, VecDeque}, sync::Arc as StdArc, time::Duration as StdDuration, @@ -11,7 +10,6 @@ use chrono::{DateTime, Duration, Utc}; use indexmap::IndexMap; use thiserror::Error; -use crate::infra::time::parse_rfc3339_utc; use crate::upstream::accounts::quota::{quota_snapshot_limit_reached, quota_snapshot_reset_at}; use crate::upstream::accounts::store::AccountStore; use crate::upstream::accounts::token_refresh::{jwt_expiry, JwtExpiry}; @@ -23,30 +21,12 @@ use crate::upstream::accounts::{ use crate::upstream::protocol::events::{ parse_rate_limit_headers, rate_limit_quota, TokenUsage as CodexTokenUsage, }; +use crate::upstream::scheduler::{ + candidates::{self, CandidateFilter, CandidateRequest}, + AccountScheduler, ScoreWeights, +}; -/// 账号轮转策略。 -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub enum RotationStrategy { - /// 综合并发压力、窗口用量和历史用量进行调度。 - Smart, - /// 优先消耗即将重置的额度窗口。 - QuotaResetPriority, - /// 按候选账号顺序循环选择。 - RoundRobin, - /// 优先选择最近使用过的账号。 - Sticky, -} - -impl RotationStrategy { - fn as_str(self) -> &'static str { - match self { - Self::Smart => "smart", - Self::QuotaResetPriority => "quota_reset_priority", - Self::RoundRobin => "round_robin", - Self::Sticky => "sticky", - } - } -} +pub use crate::upstream::scheduler::RotationStrategy; /// 账号池配置。 #[derive(Debug, Clone, PartialEq, Eq)] @@ -263,13 +243,17 @@ pub struct AccountWindowUsageDelta { pub image_request_failed: bool, } -/// 纯内存账号池,负责账号调度和运行时状态维护。 +/// 纯内存账号池:负责账号存储、在途槽位与运行时状态维护。 +/// +/// 账号**选择**已剥离到 [`crate::upstream::scheduler`]:候选过滤走 +/// [`candidates::filter`],策略选择走 [`AccountScheduler::select`]。账号池只提供选择 +/// 所需的只读视图(候选切片 + `slot_count`),并持有调度器实例。 #[derive(Debug)] pub struct AccountPool { accounts: IndexMap, slots: BTreeMap>, options: AccountPoolOptions, - rotation_cursor: usize, + scheduler: AccountScheduler, } #[derive(Debug, Clone, PartialEq, Eq)] @@ -278,16 +262,6 @@ struct AccountSlot { model: Option, } -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)] -struct SmartAccountLoadScore { - quota_limited: bool, - slot_count: usize, - window_request_count: u64, - window_token_count: u64, - window_cached_tokens: u64, - request_count: u64, -} - /// 账号释放结果。 #[derive(Debug, Clone, PartialEq, Eq)] pub struct ReleasedAccountUsage { @@ -308,7 +282,7 @@ impl AccountPool { accounts: IndexMap::new(), slots: BTreeMap::new(), options, - rotation_cursor: 0, + scheduler: AccountScheduler::new(ScoreWeights::default()), } } @@ -322,10 +296,11 @@ impl AccountPool { self.accounts.get(account_id).cloned() } - /// 移除账号及其在途槽位。 + /// 移除账号及其在途槽位与运行时反馈。 pub fn remove(&mut self, account_id: &str) -> bool { let removed = self.accounts.shift_remove(account_id).is_some(); self.slots.remove(account_id); + self.scheduler.forget_feedback(account_id); removed } @@ -333,17 +308,24 @@ impl AccountPool { pub fn clear(&mut self) { self.accounts.clear(); self.slots.clear(); - self.rotation_cursor = 0; + self.scheduler.reset_cursor(); + self.scheduler.clear_feedback(); } /// 替换账号池运行参数。 pub fn set_options(&mut self, options: AccountPoolOptions) { if self.options.rotation_strategy != options.rotation_strategy { - self.rotation_cursor = 0; + self.scheduler.reset_cursor(); } self.options = options; } + /// 回灌一次请求结果到调度器的运行时反馈(供 Smart 打分)。 + pub fn report_feedback(&self, account_id: &str, success: bool, first_token_ms: Option) { + self.scheduler + .report_feedback(account_id, success, first_token_ms); + } + /// 更新账号状态。 pub fn set_status(&mut self, account_id: &str, status: AccountStatus) -> bool { let Some(account) = self.accounts.get_mut(account_id) else { @@ -498,7 +480,9 @@ impl AccountPool { ) -> AccountAcquireWithStatusRefresh { self.cleanup_stale_slots(request.now); let refreshed_accounts = self.refresh_account_statuses(request.now); - let candidates = self.candidates(request); + let candidates = self.filter_candidates(request); + // 亲和优先:required 硬锁定该账号;preferred 若在候选内则直接选中; + // 二者都不命中时,回退到调度器按当前策略选择。 let selected = if let Some(required_account_id) = &request.required_account_id { candidates.first().cloned().map(|account| { ( @@ -508,43 +492,24 @@ impl AccountPool { ) }) } else if let Some(preferred_account_id) = &request.preferred_account_id { - let preferred = candidates + candidates .iter() .find(|account| account.id == *preferred_account_id) - .cloned(); - preferred.map(|account| { - ( - account, - self.previous_slot_at(preferred_account_id), - "preferred", - ) - }) + .cloned() + .map(|account| { + ( + account, + self.previous_slot_at(preferred_account_id), + "preferred", + ) + }) } else { None } .or_else(|| { - Some(match self.options.rotation_strategy { - RotationStrategy::Smart => { - let account = self.select_smart(&candidates)?; - let previous = self.previous_slot_at(&account.id); - (account, previous, "smart") - } - RotationStrategy::QuotaResetPriority => { - let account = self.select_quota_reset_priority(&candidates)?; - let previous = self.previous_slot_at(&account.id); - (account, previous, "quota_reset_priority") - } - RotationStrategy::RoundRobin => { - let account = self.select_round_robin(&candidates)?; - let previous = self.previous_slot_at(&account.id); - (account, previous, "round_robin") - } - RotationStrategy::Sticky => { - let account = Self::select_sticky(&candidates)?; - let previous = self.previous_slot_at(&account.id); - (account, previous, "sticky") - } - }) + let account = self.select_candidate(&candidates, request.now)?; + let previous = self.previous_slot_at(&account.id); + Some((account, previous, self.options.rotation_strategy.as_str())) }); let Some((selected, previous_slot_at, selection_source)) = selected else { return AccountAcquireWithStatusRefresh { @@ -552,26 +517,20 @@ impl AccountPool { refreshed_accounts, }; }; - let selected_slot_count = self.slot_count(&selected.id); - let selected_window_token_count = account_window_token_count(&selected); - let selected_smart_load_score = self.smart_load_score(&selected); tracing::info!( model = %request.model, rotation_strategy = self.options.rotation_strategy.as_str(), selection_source, account_id = %selected.id, candidate_count = candidates.len(), - slot_count = selected_slot_count, + slot_count = self.slot_count(&selected.id), request_count = selected.request_count, window_request_count = selected.window_request_count, - window_token_count = selected_window_token_count, window_reset_at = ?selected.window_reset_at, last_used_at = selected.last_used_at.as_deref().unwrap_or_default(), quota_limit_reached = selected.quota_limit_reached, quota_cooldown_until = ?selected.quota_cooldown_until, previous_slot_at = ?previous_slot_at, - smart_load_score = ?selected_smart_load_score, - rotation_cursor = self.rotation_cursor, "account selected for upstream request" ); @@ -689,48 +648,48 @@ impl AccountPool { let refreshed_accounts = self.refresh_account_statuses(now); let mut by_plan = IndexMap::>::new(); - for account in self.accounts.values() { - if !self.is_model_refresh_available(account, now) { - continue; + let max_concurrent = self.options.max_concurrent_per_account; + let skip_quota_limited = self.options.skip_quota_limited; + { + let slot_count = |account_id: &str| self.slots.get(account_id).map_or(0, VecDeque::len); + for account in self.accounts.values() { + if !candidates::is_model_refresh_available( + account, + max_concurrent, + skip_quota_limited, + &slot_count, + now, + ) { + continue; + } + let Some(plan_type) = account.plan_type.as_ref() else { + continue; + }; + by_plan + .entry(plan_type.clone()) + .or_default() + .push(account.clone()); } - let Some(plan_type) = account.plan_type.as_ref() else { - continue; - }; - by_plan - .entry(plan_type.clone()) - .or_default() - .push(account.clone()); } + let strategy = self.options.rotation_strategy; let mut accounts = Vec::new(); for (plan_type, group) in by_plan { - let selected = match self.options.rotation_strategy { - RotationStrategy::Smart => self.select_smart(&group), - RotationStrategy::QuotaResetPriority => self.select_quota_reset_priority(&group), - RotationStrategy::RoundRobin => self.select_round_robin(&group), - RotationStrategy::Sticky => Self::select_sticky(&group), - }; - let Some(account) = selected else { + let Some(account) = self.select_candidate(&group, now) else { continue; }; - let selected_slot_count = self.slot_count(&account.id); - let selected_window_token_count = account_window_token_count(&account); - let selected_smart_load_score = self.smart_load_score(&account); tracing::info!( plan_type, - rotation_strategy = self.options.rotation_strategy.as_str(), + rotation_strategy = strategy.as_str(), account_id = %account.id, candidate_count = group.len(), - slot_count = selected_slot_count, + slot_count = self.slot_count(&account.id), request_count = account.request_count, window_request_count = account.window_request_count, - window_token_count = selected_window_token_count, window_reset_at = ?account.window_reset_at, last_used_at = account.last_used_at.as_deref().unwrap_or_default(), quota_limit_reached = account.quota_limit_reached, quota_cooldown_until = ?account.quota_cooldown_until, - smart_load_score = ?selected_smart_load_score, - rotation_cursor = self.rotation_cursor, "account selected for model refresh" ); self.push_slot(&account.id, now, None); @@ -743,106 +702,6 @@ impl AccountPool { } } - fn select_smart(&mut self, candidates: &[Account]) -> Option { - self.select_by(candidates, |pool, a, b| pool.compare_smart(a, b)) - } - - fn select_quota_reset_priority(&mut self, candidates: &[Account]) -> Option { - self.select_by(candidates, |_pool, a, b| compare_quota_reset_priority(a, b)) - } - - fn select_by( - &mut self, - candidates: &[Account], - compare: impl Fn(&Self, &Account, &Account) -> Ordering, - ) -> Option { - let mut sorted = candidates.iter().collect::>(); - sorted.sort_by(|a, b| compare(self, a, b)); - let best = *sorted.first()?; - let tied_count = sorted - .iter() - .take_while(|account| compare(self, best, account) == Ordering::Equal) - .count(); - let index = self.rotation_cursor % tied_count; - self.rotation_cursor = self.rotation_cursor.wrapping_add(1); - Some((*sorted[index]).clone()) - } - - fn compare_smart(&self, a: &Account, b: &Account) -> Ordering { - self.smart_load_score(a) - .cmp(&self.smart_load_score(b)) - .then_with(|| compare_window_reset(a.window_reset_at, b.window_reset_at)) - .then_with(|| compare_last_used(a.last_used_at.as_deref(), b.last_used_at.as_deref())) - } - - fn smart_load_score(&self, account: &Account) -> SmartAccountLoadScore { - SmartAccountLoadScore { - quota_limited: account.quota_limit_reached, - slot_count: self.slot_count(&account.id), - window_request_count: account.window_request_count, - window_token_count: account_window_token_count(account), - window_cached_tokens: account.window_cached_tokens, - request_count: account.request_count, - } - } - - fn select_sticky(candidates: &[Account]) -> Option { - let mut selected = candidates.first()?; - for candidate in &candidates[1..] { - if compare_last_used( - candidate.last_used_at.as_deref(), - selected.last_used_at.as_deref(), - ) == Ordering::Greater - { - selected = candidate; - } - } - Some((*selected).clone()) - } - - fn select_round_robin(&mut self, candidates: &[Account]) -> Option { - if candidates.is_empty() { - return None; - } - self.rotation_cursor %= candidates.len(); - let index = self.rotation_cursor; - self.rotation_cursor = self.rotation_cursor.wrapping_add(1); - Some(candidates[index].clone()) - } - - fn candidates(&self, request: &AccountAcquireRequest) -> Vec { - let mut candidates = self - .accounts - .values() - .filter(|account| self.is_base_available(account, request)) - .cloned() - .collect::>(); - if let Some(required_account_id) = &request.required_account_id { - candidates.retain(|account| account.id == *required_account_id); - return candidates; - } - if let Some(best_tier) = self.best_available_tier(&candidates) { - candidates.retain(|account| account.plan_type.as_deref() == Some(best_tier.as_str())); - } - candidates - } - - fn is_base_available(&self, account: &Account, request: &AccountAcquireRequest) -> bool { - account.status == AccountStatus::Active - && self.slot_count(&account.id) < self.options.max_concurrent_per_account - && AccountService::quota_available_at( - account, - request.now, - self.options.skip_quota_limited, - ) - && self.is_model_allowed(account, &request.model) - && !request - .exclude_account_ids - .iter() - .any(|account_id| account_id == &account.id) - && AccountService::cloudflare_available_at(account, request.now) - } - fn slot_count(&self, account_id: &str) -> usize { self.slots.get(account_id).map_or(0, VecDeque::len) } @@ -864,13 +723,6 @@ impl AccountPool { }); } - fn is_model_refresh_available(&self, account: &Account, now: DateTime) -> bool { - account.status == AccountStatus::Active - && self.slot_count(&account.id) < self.options.max_concurrent_per_account - && AccountService::quota_available_at(account, now, self.options.skip_quota_limited) - && AccountService::cloudflare_available_at(account, now) - } - fn mark_request_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()); @@ -920,24 +772,33 @@ impl AccountPool { }); } - fn is_model_allowed(&self, account: &Account, model: &str) -> bool { - let Some(allowed_plans) = self.options.model_plan_allowlist.get(model) else { - return true; - }; - let Some(plan_type) = account.plan_type.as_deref() else { - return false; - }; - allowed_plans.iter().any(|plan| plan == plan_type) - || !self.options.fetched_model_plan_types.contains(plan_type) + /// 过滤出本次请求的可用候选集合(委托 [`candidates::filter`])。 + fn filter_candidates(&self, request: &AccountAcquireRequest) -> Vec { + let slot_count = |account_id: &str| self.slots.get(account_id).map_or(0, VecDeque::len); + candidates::filter( + self.accounts.values(), + &CandidateFilter { + max_concurrent_per_account: self.options.max_concurrent_per_account, + skip_quota_limited: self.options.skip_quota_limited, + tier_priority: &self.options.tier_priority, + model_plan_allowlist: &self.options.model_plan_allowlist, + fetched_model_plan_types: &self.options.fetched_model_plan_types, + slot_count: &slot_count, + }, + &CandidateRequest { + model: &request.model, + exclude_account_ids: &request.exclude_account_ids, + required_account_id: request.required_account_id.as_deref(), + now: request.now, + }, + ) } - fn best_available_tier(&self, candidates: &[Account]) -> Option { - self.options.tier_priority.iter().find_map(|tier| { - candidates - .iter() - .find(|account| account.plan_type.as_deref() == Some(tier.as_str())) - .and_then(|account| account.plan_type.clone()) - }) + /// 按当前策略从候选集合中选择一个账号(委托 [`AccountScheduler::select`])。 + 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) } } @@ -1520,37 +1381,3 @@ fn refresh_cloudflare_cooldown(account: &mut Account, now: DateTime) { account.cloudflare_cooldown_until = None; } } - -fn compare_quota_reset_priority(a: &Account, b: &Account) -> Ordering { - a.quota_limit_reached - .cmp(&b.quota_limit_reached) - .then_with(|| compare_window_reset(a.window_reset_at, b.window_reset_at)) - .then_with(|| a.request_count.cmp(&b.request_count)) - .then_with(|| compare_last_used(a.last_used_at.as_deref(), b.last_used_at.as_deref())) -} - -fn account_window_token_count(account: &Account) -> u64 { - account - .window_input_tokens - .saturating_add(account.window_output_tokens) - .saturating_add(account.window_image_input_tokens) - .saturating_add(account.window_image_output_tokens) -} - -fn compare_window_reset(a: Option>, b: Option>) -> Ordering { - match (a, b) { - (Some(a), Some(b)) => a.cmp(&b), - (Some(_) | None, None) | (None, Some(_)) => Ordering::Equal, - } -} - -fn compare_last_used(a: Option<&str>, b: Option<&str>) -> Ordering { - last_used_millis(a).cmp(&last_used_millis(b)) -} - -fn last_used_millis(value: Option<&str>) -> i64 { - value - .and_then(|value| parse_rfc3339_utc(value).ok()) - .map(|datetime| datetime.timestamp_millis()) - .unwrap_or(0) -} diff --git a/backend/src/upstream/mod.rs b/backend/src/upstream/mod.rs index 9e5941ae8..35c99ca09 100644 --- a/backend/src/upstream/mod.rs +++ b/backend/src/upstream/mod.rs @@ -4,5 +4,6 @@ pub mod accounts; pub mod fingerprint; pub mod models; pub mod protocol; +pub mod scheduler; pub mod token_client; pub mod transport; diff --git a/backend/src/upstream/scheduler/candidates.rs b/backend/src/upstream/scheduler/candidates.rs new file mode 100644 index 000000000..240f20f9f --- /dev/null +++ b/backend/src/upstream/scheduler/candidates.rs @@ -0,0 +1,137 @@ +//! 候选账号过滤。 +//! +//! 所有调度策略共享的前置:从账号池中筛出「本次请求可用」的候选集合。这是选择流程 +//! 的第一步,与具体策略无关——策略只在过滤后的候选之间做选择。 +//! +//! 过滤规则(全部满足才入选): +//! +//! 1. 账号状态为 [`AccountStatus::Active`]; +//! 2. 在途槽位数未达单账号并发上限; +//! 3. 配额可用([`AccountService::quota_available_at`],受 `skip_quota_limited` 控制); +//! 4. 模型在账号订阅计划的允许列表内; +//! 5. 不在本次请求的排除列表内; +//! 6. 不处于 Cloudflare 冷却期。 +//! +//! 之后再按订阅层级优先级(`tier_priority`)收窄到最高可用层级;若请求指定了 +//! `required_account_id`,则只保留该账号(跳过层级收窄)。 + +use chrono::{DateTime, Utc}; + +use crate::upstream::accounts::model::{Account, AccountStatus}; +use crate::upstream::accounts::service::AccountService; + +/// 候选过滤所需的账号池策略参数(只读视图)。 +pub struct CandidateFilter<'a> { + /// 单账号最大并发请求数。 + pub max_concurrent_per_account: usize, + /// 是否跳过已标记配额耗尽的账号。 + pub skip_quota_limited: bool, + /// 订阅层级优先级,越靠前优先级越高。 + pub tier_priority: &'a [String], + /// 模型到允许订阅计划的映射。 + pub model_plan_allowlist: &'a std::collections::BTreeMap>, + /// 已成功拉取过模型列表的订阅计划。 + pub fetched_model_plan_types: &'a std::collections::BTreeSet, + /// 按账号 ID 读取当前在途槽位数。 + pub slot_count: &'a dyn Fn(&str) -> usize, +} + +/// 本次账号获取请求的过滤条件。 +pub struct CandidateRequest<'a> { + /// 本次请求使用的模型。 + pub model: &'a str, + /// 本次请求需要排除的账号 ID。 + pub exclude_account_ids: &'a [String], + /// 本次请求必须使用的账号 ID(存在时跳过层级收窄,只保留该账号)。 + pub required_account_id: Option<&'a str>, + /// 本次调度使用的当前时间。 + pub now: DateTime, +} + +/// 从账号池筛出本次请求的可用候选集合。 +pub fn filter<'a>( + accounts: impl Iterator, + filter: &CandidateFilter<'_>, + request: &CandidateRequest<'_>, +) -> Vec { + let mut candidates = accounts + .filter(|account| is_base_available(account, filter, request)) + .cloned() + .collect::>(); + + if let Some(required_account_id) = request.required_account_id { + candidates.retain(|account| account.id == required_account_id); + return candidates; + } + + 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 +} + +/// 判断账号是否满足全部基础可用条件。 +fn is_base_available( + account: &Account, + filter: &CandidateFilter<'_>, + request: &CandidateRequest<'_>, +) -> bool { + account.status == AccountStatus::Active + && (filter.slot_count)(&account.id) < filter.max_concurrent_per_account + && AccountService::quota_available_at(account, request.now, filter.skip_quota_limited) + && is_model_allowed( + account, + request.model, + filter.model_plan_allowlist, + filter.fetched_model_plan_types, + ) + && !request + .exclude_account_ids + .iter() + .any(|account_id| account_id == &account.id) + && AccountService::cloudflare_available_at(account, request.now) +} + +/// 判断模型是否允许该账号的订阅计划使用。 +fn is_model_allowed( + account: &Account, + model: &str, + model_plan_allowlist: &std::collections::BTreeMap>, + fetched_model_plan_types: &std::collections::BTreeSet, +) -> bool { + let Some(allowed_plans) = model_plan_allowlist.get(model) else { + return true; + }; + let Some(plan_type) = account.plan_type.as_deref() else { + return false; + }; + allowed_plans.iter().any(|plan| plan == plan_type) + || !fetched_model_plan_types.contains(plan_type) +} + +/// 在候选集合中找出订阅层级优先级最高的可用层级。 +fn best_available_tier(candidates: &[Account], tier_priority: &[String]) -> Option { + 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()); + } + } + None +} + +/// 判断账号是否可用于模型列表刷新(不涉及排除列表 / 模型允许列表 / 层级收窄)。 +pub fn is_model_refresh_available( + account: &Account, + max_concurrent_per_account: usize, + skip_quota_limited: bool, + slot_count: &dyn Fn(&str) -> usize, + now: DateTime, +) -> bool { + account.status == AccountStatus::Active + && slot_count(&account.id) < max_concurrent_per_account + && AccountService::quota_available_at(account, now, skip_quota_limited) + && AccountService::cloudflare_available_at(account, now) +} diff --git a/backend/src/upstream/scheduler/feedback.rs b/backend/src/upstream/scheduler/feedback.rs new file mode 100644 index 000000000..ac29989c1 --- /dev/null +++ b/backend/src/upstream/scheduler/feedback.rs @@ -0,0 +1,171 @@ +//! 账号运行时反馈:错误率与 TTFT 的 EWMA 统计。 +//! +//! 每个账号维护两个指数加权移动平均(EWMA):错误率与首 token 延迟(TTFT, +//! 毫秒)。请求结束后经 [`FeedbackStats::report`] 回灌,直接进 [`super::score`] +//! 打分并可驱动 sticky escape。统计以 `AtomicU64`(存 `f64` 位模式)+ CAS 无锁更新, +//! 对齐 sub2api 的设计,避免选择热路径上的锁与 DB 往返。 + +use std::collections::HashMap; +use std::sync::atomic::{AtomicU64, Ordering}; +use std::sync::RwLock; + +/// EWMA 平滑系数(新样本权重)。与 sub2api 一致取 0.2。 +const EWMA_ALPHA: f64 = 0.2; + +/// 单账号的运行时 EWMA 统计。 +/// +/// `*_bits` 存 `f64` 的位模式;`NaN`(全 1 的哨兵)表示"尚无样本"。 +#[derive(Debug)] +struct AccountFeedback { + /// 错误率 EWMA,范围 `[0, 1]`。 + error_rate_bits: AtomicU64, + /// TTFT EWMA,单位毫秒。 + ttft_ms_bits: AtomicU64, +} + +/// 表示"尚无样本"的哨兵位模式。 +const EMPTY_BITS: u64 = f64::NAN.to_bits(); + +impl Default for AccountFeedback { + fn default() -> Self { + Self { + error_rate_bits: AtomicU64::new(EMPTY_BITS), + ttft_ms_bits: AtomicU64::new(EMPTY_BITS), + } + } +} + +impl AccountFeedback { + /// 读取当前错误率 EWMA;无样本时返回 `None`。 + fn error_rate(&self) -> Option { + load_sample(&self.error_rate_bits) + } + + /// 读取当前 TTFT EWMA(毫秒);无样本时返回 `None`。 + fn ttft_ms(&self) -> Option { + load_sample(&self.ttft_ms_bits) + } + + /// 回灌一次请求结果。 + fn report(&self, success: bool, first_token_ms: Option) { + let error_sample = if success { 0.0 } else { 1.0 }; + update_ewma(&self.error_rate_bits, error_sample); + if let Some(ms) = first_token_ms { + update_ewma(&self.ttft_ms_bits, ms as f64); + } + } +} + +/// 读取 EWMA 样本;哨兵(NaN)返回 `None`。 +fn load_sample(bits: &AtomicU64) -> Option { + let value = f64::from_bits(bits.load(Ordering::Relaxed)); + if value.is_nan() { + None + } else { + Some(value) + } +} + +/// 无锁 CAS 更新 EWMA:首个样本直接落值,其后按 alpha 混合。 +fn update_ewma(bits: &AtomicU64, sample: f64) { + let mut current = bits.load(Ordering::Relaxed); + loop { + let previous = f64::from_bits(current); + let next = if previous.is_nan() { + sample + } else { + EWMA_ALPHA * sample + (1.0 - EWMA_ALPHA) * previous + }; + match bits.compare_exchange_weak( + current, + next.to_bits(), + Ordering::Relaxed, + Ordering::Relaxed, + ) { + Ok(_) => return, + Err(observed) => current = observed, + } + } +} + +/// 账号 EWMA 反馈的读数快照,供打分使用。 +#[derive(Debug, Clone, Copy, Default, PartialEq)] +pub struct FeedbackSample { + /// 错误率 EWMA,范围 `[0, 1]`;无样本时为 `None`。 + pub error_rate: Option, + /// TTFT EWMA(毫秒);无样本时为 `None`。 + pub ttft_ms: Option, +} + +/// 账号运行时反馈存储:按账号 ID 聚合 EWMA 统计。 +/// +/// 选择热路径只读(`sample`),请求结束后写(`report`)。内部按账号分片为独立的 +/// 原子对象,`RwLock` 仅保护 map 结构本身(插入新账号时短暂写锁)。 +#[derive(Debug, Default)] +pub struct FeedbackStats { + accounts: RwLock>, +} + +impl FeedbackStats { + /// 创建空的反馈存储。 + pub fn new() -> Self { + Self::default() + } + + /// 读取指定账号的 EWMA 快照;账号无记录时返回全 `None`。 + pub fn sample(&self, account_id: &str) -> FeedbackSample { + let accounts = self + .accounts + .read() + .unwrap_or_else(std::sync::PoisonError::into_inner); + match accounts.get(account_id) { + Some(feedback) => FeedbackSample { + error_rate: feedback.error_rate(), + ttft_ms: feedback.ttft_ms(), + }, + None => FeedbackSample::default(), + } + } + + /// 回灌一次请求结果:成功/失败 + 首 token 延迟(毫秒)。 + pub fn report(&self, account_id: &str, success: bool, first_token_ms: Option) { + // 快路径:账号已存在时只需读锁 + 原子更新。 + { + let accounts = self + .accounts + .read() + .unwrap_or_else(std::sync::PoisonError::into_inner); + if let Some(feedback) = accounts.get(account_id) { + feedback.report(success, first_token_ms); + return; + } + } + // 慢路径:首次见到该账号,取写锁插入后回灌。 + let mut accounts = self + .accounts + .write() + .unwrap_or_else(std::sync::PoisonError::into_inner); + accounts + .entry(account_id.to_string()) + .or_default() + .report(success, first_token_ms); + } + + /// 移除指定账号的反馈记录(账号被删除时调用)。 + pub fn remove(&self, account_id: &str) { + let mut accounts = self + .accounts + .write() + .unwrap_or_else(std::sync::PoisonError::into_inner); + accounts.remove(account_id); + } + + /// 清空全部反馈记录(账号池清空时调用)。 + pub fn clear(&self) { + let mut accounts = self + .accounts + .write() + .unwrap_or_else(std::sync::PoisonError::into_inner); + accounts.clear(); + } +} diff --git a/backend/src/upstream/scheduler/mod.rs b/backend/src/upstream/scheduler/mod.rs new file mode 100644 index 000000000..802cb86ed --- /dev/null +++ b/backend/src/upstream/scheduler/mod.rs @@ -0,0 +1,112 @@ +//! 智能账号调度模块。 +//! +//! 账号选择的唯一归属地。把「选择」从账号池存储中剥离,集中于此: +//! +//! - [`candidates`] —— 候选过滤(可用性 / 层级 / 模型允许列表),所有策略共享的前置。 +//! - [`strategy`] —— 全部选择策略,每个策略一个文件([`strategy::smart`] / +//! [`strategy::quota_reset`] / [`strategy::round_robin`] / [`strategy::sticky`])。 +//! - [`feedback`] —— 跨策略的运行时 EWMA 反馈(错误率 / TTFT),喂给 Smart 打分。 +//! +//! [`AccountScheduler`] 是对外门面:持有轮转游标、打分权重与 EWMA 反馈存储,接收 +//! **已过滤的候选切片**并返回选中账号。它不 lock 账号池、不刷状态、不写库——存储与 +//! 状态维护仍归 `AccountPool`,二者经候选切片与 `slot_count` 闭包解耦。 + +pub mod candidates; +pub mod feedback; +pub mod strategy; + +use std::sync::Mutex; + +use chrono::{DateTime, Utc}; + +use crate::upstream::accounts::model::Account; + +pub use candidates::{CandidateFilter, CandidateRequest}; +pub use feedback::{FeedbackSample, FeedbackStats}; +pub use strategy::smart::rank_candidates; +pub use strategy::{RotationStrategy, ScoreBreakdown, ScoreWeights}; + +/// 账号调度器:账号选择的对外门面。 +/// +/// 持有轮转游标(并列破并列用)、打分权重(Smart 用)与 EWMA 反馈存储。选择热路径 +/// 只读反馈与权重,请求结束后经 [`AccountScheduler::report_feedback`] 回灌。游标用 +/// [`Mutex`] 保护(选择需可变,但 `AccountScheduler` 以共享引用参与并发请求)。 +#[derive(Debug)] +pub struct AccountScheduler { + /// 轮转游标:在并列最优账号之间做公平轮转。 + cursor: Mutex, + /// Smart 策略的打分权重。 + weights: ScoreWeights, + /// 账号运行时 EWMA 反馈(错误率 / TTFT)。 + feedback: FeedbackStats, +} + +impl Default for AccountScheduler { + fn default() -> Self { + Self::new(ScoreWeights::default()) + } +} + +impl AccountScheduler { + /// 使用指定打分权重创建调度器。 + pub fn new(weights: ScoreWeights) -> Self { + Self { + cursor: Mutex::new(0), + weights, + feedback: FeedbackStats::new(), + } + } + + /// 替换打分权重(策略切换或热更时调用)。 + pub fn set_weights(&mut self, weights: ScoreWeights) { + self.weights = weights; + } + + /// 重置轮转游标(策略切换时调用,保证轮转从头开始)。 + pub fn reset_cursor(&self) { + *self.lock_cursor() = 0; + } + + /// 从已过滤的候选集合中按指定策略选择一个账号。 + /// + /// `candidates` 必须是 [`candidates::filter`] 的输出;`slot_count` 按账号 ID 读取 + /// 当前在途槽位数(Smart 打分与候选无关的运行时状态由调用方提供)。 + pub fn select( + &self, + strategy: RotationStrategy, + candidates: &[Account], + slot_count: &dyn Fn(&str) -> usize, + now: DateTime, + ) -> Option { + let input = strategy::SelectionInput { + candidates, + slot_count, + feedback: &self.feedback, + weights: &self.weights, + now, + }; + let mut cursor = self.lock_cursor(); + strategy::select(strategy, &input, &mut cursor) + } + + /// 回灌一次请求结果到运行时 EWMA 反馈(错误率 / TTFT),供 Smart 打分。 + pub fn report_feedback(&self, account_id: &str, success: bool, first_token_ms: Option) { + self.feedback.report(account_id, success, first_token_ms); + } + + /// 移除指定账号的运行时反馈(账号被删除时调用)。 + pub fn forget_feedback(&self, account_id: &str) { + self.feedback.remove(account_id); + } + + /// 清空全部运行时反馈(账号池清空时调用)。 + pub fn clear_feedback(&self) { + self.feedback.clear(); + } + + fn lock_cursor(&self) -> std::sync::MutexGuard<'_, usize> { + self.cursor + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + } +} diff --git a/backend/src/upstream/scheduler/strategy/mod.rs b/backend/src/upstream/scheduler/strategy/mod.rs new file mode 100644 index 000000000..0e82793ca --- /dev/null +++ b/backend/src/upstream/scheduler/strategy/mod.rs @@ -0,0 +1,142 @@ +//! 账号调度策略层。 +//! +//! 所有账号选择策略集中于此,每个策略一个文件,互不依赖: +//! +//! - [`smart`] —— 加权线性打分 + EWMA 反馈(“智能调度”)。 +//! - [`quota_reset`] —— 优先消耗即将重置的额度窗口。 +//! - [`round_robin`] —— 按候选顺序循环。 +//! - [`sticky`] —— 优先最近使用过的账号。 +//! +//! 本模块只承担三件事:定义 [`RotationStrategy`] 枚举、把选择请求分派到对应策略、 +//! 提供各策略共享的 tie-break 比较原语([`select_by`] / [`compare_window_reset`] / +//! [`compare_last_used`])。策略本身不触碰账号池存储:在途槽位数经 `slot_count` +//! 闭包按需读取,EWMA 反馈与打分权重经 [`SelectionInput`] 传入。 + +use std::cmp::Ordering; + +use chrono::{DateTime, Utc}; + +use crate::infra::time::parse_rfc3339_utc; +use crate::upstream::accounts::model::Account; +use crate::upstream::scheduler::feedback::FeedbackStats; + +pub mod quota_reset; +pub mod round_robin; +pub mod smart; +pub mod sticky; + +pub use smart::{ScoreBreakdown, ScoreWeights}; + +/// 账号轮转策略。 +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum RotationStrategy { + /// 加权线性打分调度:多因子(负载/窗口/错误率/TTFT/重置)归一化后加权求和, + /// 配合 EWMA 反馈。 + Smart, + /// 优先消耗即将重置的额度窗口。 + QuotaResetPriority, + /// 按候选账号顺序循环选择。 + RoundRobin, + /// 优先选择最近使用过的账号。 + Sticky, +} + +impl RotationStrategy { + /// 返回策略的持久化/日志字符串。 + pub fn as_str(self) -> &'static str { + match self { + Self::Smart => "smart", + Self::QuotaResetPriority => "quota_reset_priority", + Self::RoundRobin => "round_robin", + Self::Sticky => "sticky", + } + } +} + +/// 策略选择的运行时输入。 +/// +/// `slot_count` 按账号 ID 读取当前在途槽位数;`feedback` / `weights` 仅 [`smart`] +/// 使用,其余策略忽略;`now` 用于窗口重置相关的打分。 +pub struct SelectionInput<'a> { + /// 已过滤的候选账号集合。 + pub candidates: &'a [Account], + /// 按账号 ID 读取当前在途槽位数。 + pub slot_count: &'a dyn Fn(&str) -> usize, + /// 账号 EWMA 反馈存储(Smart 打分用)。 + pub feedback: &'a FeedbackStats, + /// 打分权重(Smart 打分用)。 + pub weights: &'a ScoreWeights, + /// 当前时刻(窗口重置因子用)。 + pub now: DateTime, +} + +/// 按指定策略从候选集合中选择一个账号。 +/// +/// `cursor` 是账号池持有的轮转游标,用于在并列最优账号之间做公平轮转;不同策略对它 +/// 的读写方式不同(详见各策略模块),本函数只负责透传可变引用。 +pub fn select( + strategy: RotationStrategy, + input: &SelectionInput<'_>, + cursor: &mut usize, +) -> Option { + match strategy { + RotationStrategy::Smart => smart::select(input, cursor), + RotationStrategy::QuotaResetPriority => quota_reset::select(input, cursor), + RotationStrategy::RoundRobin => round_robin::select(input, cursor), + RotationStrategy::Sticky => sticky::select(input, cursor), + } +} + +// ==================================================================== +// 共享 tie-break 原语 +// ==================================================================== + +/// 按比较函数升序排序后,在并列最优(`Ordering::Equal`)的账号之间用轮转游标破并列。 +/// +/// 这是 [`smart`] 与 [`quota_reset`] 共享的选择骨架:先按各自的比较规则排序,再在得分 +/// 相同的账号之间轮转,保证并列账号被公平地依次选中。 +pub(crate) fn select_by( + candidates: &[Account], + cursor: &mut usize, + compare: impl Fn(&Account, &Account) -> Ordering, +) -> Option { + let mut sorted = candidates.iter().collect::>(); + sorted.sort_by(|a, b| compare(a, b)); + let best = *sorted.first()?; + let tied_count = sorted + .iter() + .take_while(|account| compare(best, account) == Ordering::Equal) + .count(); + let index = *cursor % tied_count; + *cursor = cursor.wrapping_add(1); + Some((*sorted[index]).clone()) +} + +/// 计算账号当前窗口的总 token 数(输入 + 输出 + 图片输入 + 图片输出)。 +pub(crate) fn account_window_token_count(account: &Account) -> u64 { + account + .window_input_tokens + .saturating_add(account.window_output_tokens) + .saturating_add(account.window_image_input_tokens) + .saturating_add(account.window_image_output_tokens) +} + +/// 比较两个窗口重置时间:越早重置越优先;缺失重置时间视为并列(不惩罚)。 +pub(crate) fn compare_window_reset(a: Option>, b: Option>) -> Ordering { + match (a, b) { + (Some(a), Some(b)) => a.cmp(&b), + (Some(_) | None, None) | (None, Some(_)) => Ordering::Equal, + } +} + +/// 比较两个账号的最近使用时间:越久未用越优先(升序,最小者排前)。 +pub(crate) fn compare_last_used(a: Option<&str>, b: Option<&str>) -> Ordering { + last_used_millis(a).cmp(&last_used_millis(b)) +} + +fn last_used_millis(value: Option<&str>) -> i64 { + value + .and_then(|value| parse_rfc3339_utc(value).ok()) + .map(|datetime| datetime.timestamp_millis()) + .unwrap_or(0) +} diff --git a/backend/src/upstream/scheduler/strategy/quota_reset.rs b/backend/src/upstream/scheduler/strategy/quota_reset.rs new file mode 100644 index 000000000..dbe9d526a --- /dev/null +++ b/backend/src/upstream/scheduler/strategy/quota_reset.rs @@ -0,0 +1,29 @@ +//! QuotaResetPriority 策略:优先消耗即将重置的额度窗口。 +//! +//! 排序规则(升序,越靠前越优先): +//! 1. 未封顶配额优先于已封顶(`quota_limit_reached` false < true); +//! 2. 窗口重置时间越早越优先; +//! 3. 历史请求数越少越优先; +//! 4. 最近使用时间越久越优先(LRU)。 +//! +//! 并列最优的账号之间由 [`super::select_by`] 用轮转游标破并列。 + +use std::cmp::Ordering; + +use crate::upstream::accounts::model::Account; + +use super::{compare_last_used, compare_window_reset, select_by, SelectionInput}; + +/// QuotaResetPriority 策略选择。 +pub fn select(input: &SelectionInput<'_>, cursor: &mut usize) -> Option { + select_by(input.candidates, cursor, compare_quota_reset_priority) +} + +/// 配额重置优先级比较:升序,最优(最应先用)排在最前。 +fn compare_quota_reset_priority(a: &Account, b: &Account) -> Ordering { + a.quota_limit_reached + .cmp(&b.quota_limit_reached) + .then_with(|| compare_window_reset(a.window_reset_at, b.window_reset_at)) + .then_with(|| a.request_count.cmp(&b.request_count)) + .then_with(|| compare_last_used(a.last_used_at.as_deref(), b.last_used_at.as_deref())) +} diff --git a/backend/src/upstream/scheduler/strategy/round_robin.rs b/backend/src/upstream/scheduler/strategy/round_robin.rs new file mode 100644 index 000000000..3d3db1cc2 --- /dev/null +++ b/backend/src/upstream/scheduler/strategy/round_robin.rs @@ -0,0 +1,20 @@ +//! RoundRobin 策略:按候选顺序循环选择。 +//! +//! 最简单的策略:忽略账号负载与用量,纯粹按候选集合的顺序用轮转游标依次选中, +//! 保证请求在候选账号间均匀轮转。 + +use crate::upstream::accounts::model::Account; + +use super::SelectionInput; + +/// 按候选顺序做纯轮转选择。 +pub fn select(input: &SelectionInput<'_>, cursor: &mut usize) -> Option { + let candidates = input.candidates; + if candidates.is_empty() { + return None; + } + *cursor %= candidates.len(); + let index = *cursor; + *cursor = cursor.wrapping_add(1); + Some(candidates[index].clone()) +} diff --git a/backend/src/upstream/scheduler/strategy/smart.rs b/backend/src/upstream/scheduler/strategy/smart.rs new file mode 100644 index 000000000..9f4d730b4 --- /dev/null +++ b/backend/src/upstream/scheduler/strategy/smart.rs @@ -0,0 +1,291 @@ +//! Smart 策略:加权线性打分 + EWMA 反馈。 +//! +//! 参考 sub2api 的调度设计:把多个归一化因子(负载、窗口 token 用量、窗口缓存、窗口 +//! 请求数、错误率、TTFT、窗口重置)按可配置权重线性求和,得分越高越优先。相较字典序 +//! 元组比较,加权打分是更平滑的负载均衡,不会让单一维度无条件压制其余维度,同时通过 +//! EWMA 反馈感知账号健康度与延迟。 +//! +//! 各归一化因子都落在 `[0, 1]` 且方向统一(越大越好),因此权重之间可直接比较。配额 +//! 已封顶(`quota_limit_reached`)的账号会被施加一个强惩罚,确保只要还有未封顶的候选就 +//! 不会选中它。并列最高分的账号之间由 [`super::select_by`] 用轮转游标破并列。 + +use crate::upstream::accounts::model::Account; +use crate::upstream::scheduler::feedback::{FeedbackSample, FeedbackStats}; + +use super::{account_window_token_count, SelectionInput}; + +/// 配额封顶账号的惩罚分。远大于各因子加权和的理论上限(各因子 ∈ [0,1],权重和通常为 +/// 个位数),保证未封顶账号总是优先于封顶账号。 +const QUOTA_LIMITED_PENALTY: f64 = 1_000.0; + +/// 打分权重。每个字段对应一个归一化因子的权重,权重为 0 即关闭该维度。 +#[derive(Debug, Clone, Copy, PartialEq)] +pub struct ScoreWeights { + /// 在途槽位压力(越空闲越优先)。 + pub load: f64, + /// 当前窗口实际 token 用量(用得越少越优先)。 + pub window: f64, + /// 当前窗口缓存 token 用量(缓存越少越优先,权重小,仅在实际用量相当时区分)。 + pub cached: f64, + /// 当前窗口请求数(请求越少越优先)。 + pub window_requests: f64, + /// 错误率 EWMA(越健康越优先)。 + pub error_rate: f64, + /// 首 token 延迟 EWMA(越快越优先)。 + pub ttft: f64, + /// 窗口重置临近度(越快重置越优先,use-it-or-lose-it,默认关闭)。 + pub reset: f64, +} + +impl Default for ScoreWeights { + fn default() -> Self { + // 默认权重对齐 sub2api 的取向:负载与健康度为主,窗口用量次之,TTFT 再次, + // 窗口重置默认关闭(避免过早榨干即将重置的账号)。cached 权重远小于 window, + // 确保实际 token 用量主导、缓存仅作细粒度区分。 + Self { + load: 1.0, + window: 0.8, + cached: 0.05, + window_requests: 0.6, + error_rate: 1.0, + ttft: 0.5, + reset: 0.0, + } + } +} + +/// 单个候选账号的打分明细,供日志与可观测使用。 +#[derive(Debug, Clone, Copy, PartialEq)] +pub struct ScoreBreakdown { + /// 各因子加权后的总分(已含配额封顶惩罚)。 + pub total: f64, + /// 负载因子(归一化后,未乘权重)。 + pub load: f64, + /// 窗口实际 token 用量因子。 + pub window: f64, + /// 窗口缓存 token 用量因子。 + pub cached: f64, + /// 窗口请求数因子。 + pub window_requests: f64, + /// 错误率因子。 + pub error_rate: f64, + /// TTFT 因子。 + pub ttft: f64, + /// 窗口重置因子。 + pub reset: f64, + /// 是否被施加配额封顶惩罚。 + pub quota_limited: bool, +} + +/// Smart 策略选择:按加权得分降序排序,在并列最高分之间用轮转游标破并列。 +pub fn select(input: &SelectionInput<'_>, cursor: &mut usize) -> Option { + let ranked = rank_candidates( + input.candidates, + input.weights, + input.now.timestamp(), + input.slot_count, + input.feedback, + ); + let best_total = ranked.first()?.1.total; + // 并列最高分的账号数量(浮点相等用位相等判定,得分由相同输入算出,稳定)。 + let tied_count = ranked + .iter() + .take_while(|(_, breakdown)| breakdown.total.to_bits() == best_total.to_bits()) + .count(); + let index = *cursor % tied_count; + *cursor = cursor.wrapping_add(1); + let account_index = ranked[index].0; + Some(input.candidates[account_index].clone()) +} + +/// 对候选集合按加权得分排序(降序),返回 `(账号索引, 打分明细)` 列表。 +/// +/// 只做打分与排序,最终的并列破并列交给 [`select`]。公开供测试直接断言排序结果。 +pub fn rank_candidates( + candidates: &[Account], + weights: &ScoreWeights, + now_secs: i64, + slot_count: &dyn Fn(&str) -> usize, + feedback: &FeedbackStats, +) -> Vec<(usize, ScoreBreakdown)> { + let samples = candidates + .iter() + .map(|account| feedback.sample(&account.id)) + .collect::>(); + let ctx = NormalizationContext::build(candidates, &samples, now_secs, slot_count); + let mut ranked = candidates + .iter() + .zip(&samples) + .enumerate() + .map(|(index, (account, sample))| { + ( + index, + score_account(account, sample, weights, &ctx, now_secs, slot_count), + ) + }) + .collect::>(); + // 降序:得分高在前。NaN 不应出现(因子均有界),保守用 total_cmp。 + ranked.sort_by(|a, b| b.1.total.total_cmp(&a.1.total)); + ranked +} + +/// 候选集合的归一化上下文:预先算好各维度的 min/max,供逐账号归一化。 +struct NormalizationContext { + max_slot_count: f64, + max_window_tokens: f64, + max_window_cached: f64, + max_window_requests: f64, + min_ttft_ms: f64, + max_ttft_ms: f64, + min_reset_secs: f64, + max_reset_secs: f64, +} + +impl NormalizationContext { + fn build( + candidates: &[Account], + samples: &[FeedbackSample], + now_secs: i64, + slot_count: &dyn Fn(&str) -> usize, + ) -> Self { + let mut max_slot_count = 0.0_f64; + let mut max_window_tokens = 0.0_f64; + let mut max_window_cached = 0.0_f64; + let mut max_window_requests = 0.0_f64; + let mut min_ttft_ms = f64::INFINITY; + let mut max_ttft_ms = 0.0_f64; + let mut min_reset_secs = f64::INFINITY; + let mut max_reset_secs = 0.0_f64; + + for (account, sample) in candidates.iter().zip(samples) { + max_slot_count = max_slot_count.max(slot_count(&account.id) as f64); + max_window_tokens = max_window_tokens.max(account_window_token_count(account) as f64); + max_window_cached = max_window_cached.max(account.window_cached_tokens as f64); + max_window_requests = max_window_requests.max(account.window_request_count as f64); + + if let Some(ttft) = sample.ttft_ms { + if ttft > 0.0 { + min_ttft_ms = min_ttft_ms.min(ttft); + max_ttft_ms = max_ttft_ms.max(ttft); + } + } + + if let Some(reset_at) = account.window_reset_at { + let secs = (reset_at.timestamp() - now_secs).max(0) as f64; + min_reset_secs = min_reset_secs.min(secs); + max_reset_secs = max_reset_secs.max(secs); + } + } + + Self { + max_slot_count, + max_window_tokens, + max_window_cached, + max_window_requests, + min_ttft_ms: if min_ttft_ms.is_finite() { + min_ttft_ms + } else { + 0.0 + }, + max_ttft_ms, + min_reset_secs: if min_reset_secs.is_finite() { + min_reset_secs + } else { + 0.0 + }, + max_reset_secs, + } + } +} + +/// 把 `value` 从 `[min, max]` 线性映射到 `[0, 1]`;区间退化时返回中性值 `1.0` +/// (表示该维度对所有候选无区分度,不应影响排序)。 +fn normalize_ascending_good(value: f64, min: f64, max: f64) -> f64 { + if max <= min { + return 1.0; + } + ((value - min) / (max - min)).clamp(0.0, 1.0) +} + +/// 计算指定账号在候选集合中的加权得分明细。 +fn score_account( + account: &Account, + sample: &FeedbackSample, + weights: &ScoreWeights, + ctx: &NormalizationContext, + now_secs: i64, + slot_count: &dyn Fn(&str) -> usize, +) -> ScoreBreakdown { + // 负载:槽位越少越好。归一化到 [0,1] 后取反(1 - x),使“越空闲得分越高”。 + let load = + 1.0 - normalize_ascending_good(slot_count(&account.id) as f64, 0.0, ctx.max_slot_count); + + // 窗口实际 token 用量:token 越少越好。 + let window = 1.0 + - normalize_ascending_good( + account_window_token_count(account) as f64, + 0.0, + ctx.max_window_tokens, + ); + + // 窗口缓存 token 用量:缓存越少越好(权重小,仅在实际用量相当时细分)。 + let cached = 1.0 + - normalize_ascending_good( + account.window_cached_tokens as f64, + 0.0, + ctx.max_window_cached, + ); + + // 窗口请求数:请求越少越好。 + let window_requests = 1.0 + - normalize_ascending_good( + account.window_request_count as f64, + 0.0, + ctx.max_window_requests, + ); + + // 错误率 EWMA:本身即 [0,1],越低越健康。无样本时视为 0(满分健康)。 + let error_rate = 1.0 - sample.error_rate.unwrap_or(0.0).clamp(0.0, 1.0); + + // TTFT:延迟越低越好。无样本的账号视为中性满分,避免冷启动被压。 + let ttft = match sample.ttft_ms { + Some(ms) if ms > 0.0 => { + 1.0 - normalize_ascending_good(ms, ctx.min_ttft_ms, ctx.max_ttft_ms) + } + _ => 1.0, + }; + + // 窗口重置:越快重置越优先(use-it-or-lose-it)。无 reset 时间视为中性满分。 + let reset = match account.window_reset_at { + Some(reset_at) => { + let secs = (reset_at.timestamp() - now_secs).max(0) as f64; + 1.0 - normalize_ascending_good(secs, ctx.min_reset_secs, ctx.max_reset_secs) + } + None => 1.0, + }; + + let mut total = weights.load * load + + weights.window * window + + weights.cached * cached + + weights.window_requests * window_requests + + weights.error_rate * error_rate + + weights.ttft * ttft + + weights.reset * reset; + + // 配额封顶硬惩罚:确保只要还有未封顶候选就不会选中封顶账号。 + if account.quota_limit_reached { + total -= QUOTA_LIMITED_PENALTY; + } + + ScoreBreakdown { + total, + load, + window, + cached, + window_requests, + error_rate, + ttft, + reset, + quota_limited: account.quota_limit_reached, + } +} diff --git a/backend/src/upstream/scheduler/strategy/sticky.rs b/backend/src/upstream/scheduler/strategy/sticky.rs new file mode 100644 index 000000000..c57caa23a --- /dev/null +++ b/backend/src/upstream/scheduler/strategy/sticky.rs @@ -0,0 +1,28 @@ +//! Sticky 策略:优先选择最近使用过的账号。 +//! +//! 选中 `last_used_at` 最大(最近使用)的账号,倾向于持续复用同一账号(利于上游的 +//! prompt 缓存命中)。该策略不使用轮转游标——最近使用时间本身即确定性顺序。 + +use std::cmp::Ordering; + +use crate::upstream::accounts::model::Account; + +use super::{compare_last_used, SelectionInput}; + +/// 选中最近使用过的账号(`last_used_at` 最大者)。 +/// +/// `_cursor` 未使用:Sticky 由 `last_used_at` 决定顺序,不参与轮转破并列。 +pub fn select(input: &SelectionInput<'_>, _cursor: &mut usize) -> Option { + let candidates = input.candidates; + let mut selected = candidates.first()?; + for candidate in &candidates[1..] { + if compare_last_used( + candidate.last_used_at.as_deref(), + selected.last_used_at.as_deref(), + ) == Ordering::Greater + { + selected = candidate; + } + } + Some(selected.clone()) +} diff --git a/backend/tests/proxy/dispatch/chat_upstream/responses_websocket.rs b/backend/tests/proxy/dispatch/chat_upstream/responses_websocket.rs index 9d25f3e65..fbce8df3a 100644 --- a/backend/tests/proxy/dispatch/chat_upstream/responses_websocket.rs +++ b/backend/tests/proxy/dispatch/chat_upstream/responses_websocket.rs @@ -1828,7 +1828,9 @@ async fn responses_websocket_previous_response_id_should_retry_fallback_account_ .unwrap(); assert!(body.contains("\"id\":\"resp_history_fallback\"")); - assert_eq!(fallback_payload["previous_response_id"], "resp_prev"); + // 亲和缓存无 `resp_prev` 记录(Unknown owner),换账号前历史即被剥离, + // 由 `input` 让备用账号重建上下文,不把上一账号作用域的 id 透传给它。 + assert_eq!(fallback_payload.get("previous_response_id"), None); assert_eq!(secondary_usage.0, 1); } @@ -1878,7 +1880,9 @@ async fn responses_websocket_non_stream_previous_response_id_should_retry_fallba .unwrap(); assert_eq!(body["id"], "resp_history_fallback_non_stream"); - assert_eq!(fallback_payload["previous_response_id"], "resp_prev"); + // 亲和缓存无 `resp_prev` 记录(Unknown owner),换账号前历史即被剥离, + // 由 `input` 让备用账号重建上下文,不把上一账号作用域的 id 透传给它。 + assert_eq!(fallback_payload.get("previous_response_id"), None); assert_eq!(secondary_usage.0, 1); } @@ -2324,10 +2328,9 @@ async fn responses_stream_with_previous_response_id_should_forward_websocket_chu captured_header(&captured.headers, "authorization"), Some("Bearer access-secret") ); - assert_eq!( - captured.payload["previous_response_id"], - "resp_runtime_pool_previous" - ); + // 亲和缓存无该 previous_response_id 的记录(Unknown),故上游 payload 里应被剥离, + // 靠已展开的 input 让选中账号自行重建续链——与 sub2api/codex2api 的默认行为一致。 + assert_eq!(captured.payload["previous_response_id"], Value::Null); } #[tokio::test] @@ -2355,10 +2358,8 @@ async fn responses_stream_with_previous_response_id_should_record_websocket_audi let metadata: Value = serde_json::from_str(&event.metadata_json).unwrap(); assert!(body.contains("resp_live_websocket_stream")); - assert_eq!( - captured.payload["previous_response_id"], - "resp_runtime_pool_previous" - ); + // 亲和缓存无记录(Unknown)→ 上游 payload 里 previous_response_id 被剥离。 + assert_eq!(captured.payload["previous_response_id"], Value::Null); assert_eq!(event.level, "info"); assert_eq!(metadata["stream"], true); assert_eq!(metadata["transport"], "websocket"); diff --git a/backend/tests/upstream/mod.rs b/backend/tests/upstream/mod.rs index 562f0f0c4..9e415e77b 100644 --- a/backend/tests/upstream/mod.rs +++ b/backend/tests/upstream/mod.rs @@ -2,5 +2,6 @@ pub(crate) mod accounts; pub(crate) mod fingerprint; pub(crate) mod models; pub(crate) mod protocol; +pub(crate) mod scheduler; pub(crate) mod token_client; pub(crate) mod transport; diff --git a/backend/tests/upstream/scheduler/feedback.rs b/backend/tests/upstream/scheduler/feedback.rs new file mode 100644 index 000000000..7f8266e2a --- /dev/null +++ b/backend/tests/upstream/scheduler/feedback.rs @@ -0,0 +1,44 @@ +use codex_proxy_rs::upstream::scheduler::{FeedbackSample, FeedbackStats}; + +#[test] +fn sample_of_unknown_account_is_empty() { + let stats = FeedbackStats::new(); + assert_eq!(stats.sample("missing"), FeedbackSample::default()); +} + +#[test] +fn first_sample_lands_directly() { + let stats = FeedbackStats::new(); + stats.report("acct", false, Some(400)); + let sample = stats.sample("acct"); + assert_eq!(sample.error_rate, Some(1.0)); + assert_eq!(sample.ttft_ms, Some(400.0)); +} + +#[test] +fn subsequent_samples_blend_by_alpha() { + let stats = FeedbackStats::new(); + stats.report("acct", false, Some(1000)); // error_rate=1.0, ttft=1000 + stats.report("acct", true, Some(0)); // error_rate=0.8, ttft=800 + let sample = stats.sample("acct"); + assert!((sample.error_rate.unwrap() - 0.8).abs() < 1e-9); + assert!((sample.ttft_ms.unwrap() - 800.0).abs() < 1e-9); +} + +#[test] +fn success_only_reports_decay_error_rate() { + let stats = FeedbackStats::new(); + stats.report("acct", false, None); // error_rate=1.0, ttft untouched + stats.report("acct", true, None); // error_rate=0.8 + let sample = stats.sample("acct"); + assert!((sample.error_rate.unwrap() - 0.8).abs() < 1e-9); + assert_eq!(sample.ttft_ms, None); +} + +#[test] +fn remove_clears_account() { + let stats = FeedbackStats::new(); + stats.report("acct", false, Some(100)); + stats.remove("acct"); + assert_eq!(stats.sample("acct"), FeedbackSample::default()); +} diff --git a/backend/tests/upstream/scheduler/mod.rs b/backend/tests/upstream/scheduler/mod.rs new file mode 100644 index 000000000..c0f28927e --- /dev/null +++ b/backend/tests/upstream/scheduler/mod.rs @@ -0,0 +1,2 @@ +pub(crate) mod feedback; +pub(crate) mod score; diff --git a/backend/tests/upstream/scheduler/score.rs b/backend/tests/upstream/scheduler/score.rs new file mode 100644 index 000000000..bf54d7bee --- /dev/null +++ b/backend/tests/upstream/scheduler/score.rs @@ -0,0 +1,144 @@ +use chrono::Utc; +use codex_proxy_rs::upstream::accounts::model::AccountStatus; +use codex_proxy_rs::upstream::scheduler::{rank_candidates, FeedbackStats, ScoreWeights}; + +use crate::support::accounts::test_account; + +#[test] +fn ranks_idle_account_above_loaded_one() { + let mut loaded = test_account("loaded", AccountStatus::Active); + loaded.window_input_tokens = 1_000_000; + let idle = test_account("idle", AccountStatus::Active); + let candidates = vec![loaded, idle]; + let feedback = FeedbackStats::new(); + let slot_count = |_: &str| 0; + + let ranked = rank_candidates( + &candidates, + &ScoreWeights::default(), + Utc::now().timestamp(), + &slot_count, + &feedback, + ); + // idle 账号(索引 1)窗口用量为 0,应排在负载重的账号之前。 + assert_eq!(ranked.first().unwrap().0, 1); +} + +#[test] +fn high_error_rate_account_is_penalized() { + let healthy = test_account("healthy", AccountStatus::Active); + let flaky = test_account("flaky", AccountStatus::Active); + let candidates = vec![healthy, flaky]; + let feedback = FeedbackStats::new(); + // flaky 连续失败,错误率 EWMA 拉高。 + for _ in 0..5 { + feedback.report("flaky", false, None); + } + let slot_count = |_: &str| 0; + + let ranked = rank_candidates( + &candidates, + &ScoreWeights::default(), + Utc::now().timestamp(), + &slot_count, + &feedback, + ); + // healthy(索引 0)应排在 flaky 之前。 + assert_eq!(ranked.first().unwrap().0, 0); +} + +#[test] +fn quota_limited_account_ranks_last_despite_being_idle() { + let mut limited = test_account("limited", AccountStatus::Active); + limited.quota_limit_reached = true; + let mut busy = test_account("busy", AccountStatus::Active); + busy.window_input_tokens = 1_000_000; + busy.window_request_count = 100; + let candidates = vec![limited, busy]; + let feedback = FeedbackStats::new(); + let slot_count = |_: &str| 0; + + let ranked = rank_candidates( + &candidates, + &ScoreWeights::default(), + Utc::now().timestamp(), + &slot_count, + &feedback, + ); + // 即使 limited 完全空闲,配额封顶惩罚也应让繁忙但未封顶的 busy(索引 1)排前。 + assert_eq!(ranked.first().unwrap().0, 1); +} + +#[test] +fn actual_token_load_dominates_cached() { + // cache_heavy 实际 token 更少(但缓存高),heavier_actual 实际 token 更多。 + // 实际用量权重远大于缓存,cache_heavy(索引 0)应排前。 + let mut cache_heavy = test_account("cache_heavy", AccountStatus::Active); + cache_heavy.window_input_tokens = 10_000; + cache_heavy.window_cached_tokens = 100_000; + let mut heavier_actual = test_account("heavier_actual", AccountStatus::Active); + heavier_actual.window_input_tokens = 20_000; + let candidates = vec![cache_heavy, heavier_actual]; + let feedback = FeedbackStats::new(); + let slot_count = |_: &str| 0; + + let ranked = rank_candidates( + &candidates, + &ScoreWeights::default(), + Utc::now().timestamp(), + &slot_count, + &feedback, + ); + assert_eq!(ranked.first().unwrap().0, 0); +} + +#[test] +fn cached_breaks_ties_when_actual_load_matches() { + // 实际 token 相同,缓存少者(cache_light,索引 1)应排前。 + let mut cache_heavy = test_account("cache_heavy", AccountStatus::Active); + cache_heavy.window_input_tokens = 10_000; + cache_heavy.window_cached_tokens = 100_000; + let mut cache_light = test_account("cache_light", AccountStatus::Active); + cache_light.window_input_tokens = 10_000; + cache_light.window_cached_tokens = 1_000; + let candidates = vec![cache_heavy, cache_light]; + let feedback = FeedbackStats::new(); + let slot_count = |_: &str| 0; + + let ranked = rank_candidates( + &candidates, + &ScoreWeights::default(), + Utc::now().timestamp(), + &slot_count, + &feedback, + ); + assert_eq!(ranked.first().unwrap().0, 1); +} + +#[test] +fn zero_weights_produce_flat_scores() { + let candidates = vec![ + test_account("a", AccountStatus::Active), + test_account("b", AccountStatus::Active), + ]; + let feedback = FeedbackStats::new(); + let weights = ScoreWeights { + load: 0.0, + window: 0.0, + cached: 0.0, + window_requests: 0.0, + error_rate: 0.0, + ttft: 0.0, + reset: 0.0, + }; + let slot_count = |_: &str| 0; + + let ranked = rank_candidates( + &candidates, + &weights, + Utc::now().timestamp(), + &slot_count, + &feedback, + ); + assert!(ranked.iter().all(|(_, b)| b.total == 0.0)); +} diff --git a/frontend/src/views/settings/components/RotationStrategyCard.vue b/frontend/src/views/settings/components/RotationStrategyCard.vue index 42ae58530..5b7ec3725 100644 --- a/frontend/src/views/settings/components/RotationStrategyCard.vue +++ b/frontend/src/views/settings/components/RotationStrategyCard.vue @@ -19,8 +19,8 @@ const model = defineModel({ required: true })