refactor: extract upstream scheduler

Co-authored-by: Codex <noreply@openai.com>
This commit is contained in:
zyycn
2026-07-08 13:56:27 +08:00
co-authored by Codex
parent daad7c746c
commit 5483e37b00
16 changed files with 1224 additions and 274 deletions
+89 -262
View File
@@ -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<String, Account>,
slots: BTreeMap<String, VecDeque<AccountSlot>>,
options: AccountPoolOptions,
rotation_cursor: usize,
scheduler: AccountScheduler,
}
#[derive(Debug, Clone, PartialEq, Eq)]
@@ -278,16 +262,6 @@ struct AccountSlot {
model: Option<String>,
}
#[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<u64>) {
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::<String, Vec<Account>>::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<Account> {
self.select_by(candidates, |pool, a, b| pool.compare_smart(a, b))
}
fn select_quota_reset_priority(&mut self, candidates: &[Account]) -> Option<Account> {
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<Account> {
let mut sorted = candidates.iter().collect::<Vec<_>>();
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<Account> {
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<Account> {
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<Account> {
let mut candidates = self
.accounts
.values()
.filter(|account| self.is_base_available(account, request))
.cloned()
.collect::<Vec<_>>();
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<Utc>) -> 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<Utc>) -> Option<Account> {
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<Account> {
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<String> {
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<Utc>) -> Option<Account> {
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<Utc>) {
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<DateTime<Utc>>, b: Option<DateTime<Utc>>) -> 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)
}
+1
View File
@@ -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;
@@ -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<String, Vec<String>>,
/// 已成功拉取过模型列表的订阅计划。
pub fetched_model_plan_types: &'a std::collections::BTreeSet<String>,
/// 按账号 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<Utc>,
}
/// 从账号池筛出本次请求的可用候选集合。
pub fn filter<'a>(
accounts: impl Iterator<Item = &'a Account>,
filter: &CandidateFilter<'_>,
request: &CandidateRequest<'_>,
) -> Vec<Account> {
let mut candidates = accounts
.filter(|account| is_base_available(account, filter, request))
.cloned()
.collect::<Vec<_>>();
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<String, Vec<String>>,
fetched_model_plan_types: &std::collections::BTreeSet<String>,
) -> 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<String> {
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<Utc>,
) -> 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)
}
+171
View File
@@ -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<f64> {
load_sample(&self.error_rate_bits)
}
/// 读取当前 TTFT EWMA(毫秒);无样本时返回 `None`。
fn ttft_ms(&self) -> Option<f64> {
load_sample(&self.ttft_ms_bits)
}
/// 回灌一次请求结果。
fn report(&self, success: bool, first_token_ms: Option<u64>) {
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<f64> {
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<f64>,
/// TTFT EWMA(毫秒);无样本时为 `None`。
pub ttft_ms: Option<f64>,
}
/// 账号运行时反馈存储:按账号 ID 聚合 EWMA 统计。
///
/// 选择热路径只读(`sample`),请求结束后写(`report`)。内部按账号分片为独立的
/// 原子对象,`RwLock` 仅保护 map 结构本身(插入新账号时短暂写锁)。
#[derive(Debug, Default)]
pub struct FeedbackStats {
accounts: RwLock<HashMap<String, AccountFeedback>>,
}
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<u64>) {
// 快路径:账号已存在时只需读锁 + 原子更新。
{
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();
}
}
+112
View File
@@ -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<usize>,
/// 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<Utc>,
) -> Option<Account> {
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<u64>) {
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)
}
}
@@ -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<Utc>,
}
/// 按指定策略从候选集合中选择一个账号。
///
/// `cursor` 是账号池持有的轮转游标,用于在并列最优账号之间做公平轮转;不同策略对它
/// 的读写方式不同(详见各策略模块),本函数只负责透传可变引用。
pub fn select(
strategy: RotationStrategy,
input: &SelectionInput<'_>,
cursor: &mut usize,
) -> Option<Account> {
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<Account> {
let mut sorted = candidates.iter().collect::<Vec<_>>();
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<DateTime<Utc>>, b: Option<DateTime<Utc>>) -> 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)
}
@@ -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<Account> {
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()))
}
@@ -0,0 +1,20 @@
//! RoundRobin 策略:按候选顺序循环选择。
//!
//! 最简单的策略:忽略账号负载与用量,纯粹按候选集合的顺序用轮转游标依次选中,
//! 保证请求在候选账号间均匀轮转。
use crate::upstream::accounts::model::Account;
use super::SelectionInput;
/// 按候选顺序做纯轮转选择。
pub fn select(input: &SelectionInput<'_>, cursor: &mut usize) -> Option<Account> {
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())
}
@@ -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<Account> {
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::<Vec<_>>();
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::<Vec<_>>();
// 降序:得分高在前。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,
}
}
@@ -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<Account> {
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())
}
@@ -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");
+1
View File
@@ -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;
@@ -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());
}
+2
View File
@@ -0,0 +1,2 @@
pub(crate) mod feedback;
pub(crate) mod score;
+144
View File
@@ -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));
}
@@ -19,8 +19,8 @@ const model = defineModel<RotationStrategy>({ required: true })
<template>
<BaseCard
:padded="false"
title="账号选择"
description="决定每次请求如何使用账号池。"
title="调度策略"
description="决定每次请求如何调度账号池。"
header-class="px-5 pt-4"
body-class="px-5 py-5"
>