feat(scheduling): add opt-in "prefer soonest reset" account selection

Adds a use-it-or-lose-it scheduling strategy: prefer accounts whose
session window resets soonest, so near-reset accounts get drained first
instead of accounts whose reset is still far away.

Both schedulers, opt-in, default behavior unchanged:

- Anthropic (gateway_service.go): new GatewaySchedulingConfig
  .PreferSoonestReset flag. When on, the layered load-aware selection
  inserts a filterBySoonestReset stage (priority -> soonest-reset ->
  load -> LRU). Accounts with no active SessionWindowEnd are treated as
  lowest priority; ties fall through to LRU.

- OpenAI/Codex (openai_account_scheduler.go): new "reset" score weight
  in GatewayOpenAIWSSchedulerScoreWeights. Soonest-reset accounts score
  higher; weight defaults to 0 (no effect).

SessionWindowEnd (upstream 5h/quota ResetsAt) is already carried in the
scheduler snapshot, so no snapshot changes are needed.

Documented in deploy/config.example.yaml. Adds unit tests for the
Anthropic filter and the OpenAI reset factor.
This commit is contained in:
kangjwme
2026-06-18 22:50:46 +08:00
parent 4a5665da5b
commit 510adf703c
6 changed files with 261 additions and 4 deletions
+10
View File
@@ -957,6 +957,9 @@ type GatewayOpenAIWSSchedulerScoreWeights struct {
Queue float64 `mapstructure:"queue"`
ErrorRate float64 `mapstructure:"error_rate"`
TTFT float64 `mapstructure:"ttft"`
// Reset 倾向「会话窗口最早重置」的账号(use-it-or-lose-it)。
// >0 时,剩余重置时间越短的账号得分越高,从而被优先用尽。默认 0(关闭,不改变原有行为)。
Reset float64 `mapstructure:"reset"`
}
// GatewayOpenAISchedulerConfig OpenAI 高级调度器配置。
@@ -1055,6 +1058,11 @@ type GatewaySchedulingConfig struct {
// 兜底层账户选择策略: "last_used"(按最后使用时间排序,默认) 或 "random"(随机)
FallbackSelectionMode string `mapstructure:"fallback_selection_mode"`
// PreferSoonestReset 开启后,负载感知选择会优先选用「会话窗口最早重置」的账号
// use-it-or-lose-it:先用尽即将重置的账号,保留重置时间还很久的账号)。
// 默认 false,保持原有「优先级 → 负载率 → LRU」行为不变。
PreferSoonestReset bool `mapstructure:"prefer_soonest_reset"`
// 负载计算
LoadBatchEnabled bool `mapstructure:"load_batch_enabled"`
LoadBatchCacheTTLMS int `mapstructure:"load_batch_cache_ttl_ms"`
@@ -1870,6 +1878,7 @@ func setDefaults() {
viper.SetDefault("gateway.openai_ws.scheduler_score_weights.queue", 0.7)
viper.SetDefault("gateway.openai_ws.scheduler_score_weights.error_rate", 0.8)
viper.SetDefault("gateway.openai_ws.scheduler_score_weights.ttft", 0.5)
viper.SetDefault("gateway.openai_ws.scheduler_score_weights.reset", 0.0)
// OpenAI HTTP upstream protocol strategy
viper.SetDefault("gateway.openai_http2.enabled", true)
viper.SetDefault("gateway.openai_http2.allow_proxy_fallback_to_http1", true)
@@ -1906,6 +1915,7 @@ func setDefaults() {
viper.SetDefault("gateway.scheduling.fallback_wait_timeout", 30*time.Second)
viper.SetDefault("gateway.scheduling.fallback_max_waiting", 100)
viper.SetDefault("gateway.scheduling.fallback_selection_mode", "last_used")
viper.SetDefault("gateway.scheduling.prefer_soonest_reset", false)
viper.SetDefault("gateway.scheduling.load_batch_enabled", true)
viper.SetDefault("gateway.scheduling.load_batch_cache_ttl_ms", 200)
viper.SetDefault("gateway.scheduling.snapshot_mget_chunk_size", 128)
+40 -3
View File
@@ -2218,13 +2218,17 @@ func (s *GatewayService) SelectAccountWithLoadAwareness(ctx context.Context, gro
}
}
// 分层过滤选择:优先级 → 负载率 → LRU
// 分层过滤选择:优先级 →(可选)最早重置 → 负载率 → LRU
for len(available) > 0 {
// 1. 取优先级最小的集合
candidates := filterByMinPriority(available)
// 2. 取负载率最低的集合
// 2. (可选)use-it-or-lose-it:优先选用会话窗口最早重置的账号
if cfg.PreferSoonestReset {
candidates = filterBySoonestReset(candidates)
}
// 3. 取负载率最低的集合
candidates = filterByMinLoadRate(candidates)
// 3. LRU 选择最久未用的账号
// 4. LRU 选择最久未用的账号
selected := selectByLRU(candidates, preferOAuth)
if selected == nil {
break
@@ -2983,6 +2987,39 @@ func filterByMinLoadRate(accounts []accountWithLoad) []accountWithLoad {
return result
}
// filterBySoonestReset 过滤出「会话窗口最早重置」的账号集合(use-it-or-lose-it)。
// 仅保留拥有未来重置时间(SessionWindowEnd 在当前时间之后)且最早的账号;
// 窗口为空或已过期的账号视为无活跃窗口、优先级最低。
// 当所有账号都没有活跃窗口时,返回原集合(不改变后续 LRU 选择)。
func filterBySoonestReset(accounts []accountWithLoad) []accountWithLoad {
if len(accounts) <= 1 {
return accounts
}
now := time.Now()
var minEnd *time.Time
for _, acc := range accounts {
end := acc.account.SessionWindowEnd
if end == nil || !now.Before(*end) {
continue
}
if minEnd == nil || end.Before(*minEnd) {
minEnd = end
}
}
if minEnd == nil {
// 没有任何账号拥有活跃窗口,保持原集合
return accounts
}
result := make([]accountWithLoad, 0, len(accounts))
for _, acc := range accounts {
end := acc.account.SessionWindowEnd
if end != nil && now.Before(*end) && end.Equal(*minEnd) {
result = append(result, acc)
}
}
return result
}
// selectByLRU 从集合中选择最久未用的账号
// 如果有多个账号具有相同的最小 LastUsedAt,则随机选择一个
func selectByLRU(accounts []accountWithLoad, preferOAuth bool) *accountWithLoad {
@@ -0,0 +1,81 @@
//go:build unit
package service
import (
"testing"
"time"
"github.com/stretchr/testify/require"
)
func accWithWindowEnd(id int64, end *time.Time) accountWithLoad {
return accountWithLoad{
account: &Account{
ID: id,
Schedulable: true,
Status: StatusActive,
SessionWindowEnd: end,
},
loadInfo: &AccountLoadInfo{AccountID: id},
}
}
func TestFilterBySoonestReset_PicksSoonestFutureWindow(t *testing.T) {
now := time.Now()
soon := now.Add(1 * time.Hour)
later := now.Add(24 * time.Hour)
accounts := []accountWithLoad{
accWithWindowEnd(1, testTimePtr(later)),
accWithWindowEnd(2, testTimePtr(soon)),
accWithWindowEnd(3, testTimePtr(later)),
}
got := filterBySoonestReset(accounts)
require.Len(t, got, 1)
require.Equal(t, int64(2), got[0].account.ID, "重置时间最早的账号被选中")
}
func TestFilterBySoonestReset_IgnoresNilAndExpiredWindows(t *testing.T) {
now := time.Now()
expired := now.Add(-1 * time.Hour)
active := now.Add(2 * time.Hour)
accounts := []accountWithLoad{
accWithWindowEnd(1, nil), // 无活跃窗口
accWithWindowEnd(2, testTimePtr(expired)), // 已过期,视为无活跃窗口
accWithWindowEnd(3, testTimePtr(active)), // 唯一活跃窗口
}
got := filterBySoonestReset(accounts)
require.Len(t, got, 1)
require.Equal(t, int64(3), got[0].account.ID, "仅保留拥有未来重置时间的账号")
}
func TestFilterBySoonestReset_NoActiveWindowReturnsAll(t *testing.T) {
now := time.Now()
expired := now.Add(-30 * time.Minute)
accounts := []accountWithLoad{
accWithWindowEnd(1, nil),
accWithWindowEnd(2, testTimePtr(expired)),
}
got := filterBySoonestReset(accounts)
require.Len(t, got, 2, "没有任何账号拥有活跃窗口时,返回原集合不做过滤")
}
func TestFilterBySoonestReset_TiedSoonestKeepsAll(t *testing.T) {
now := time.Now()
end := now.Add(90 * time.Minute)
accounts := []accountWithLoad{
accWithWindowEnd(1, testTimePtr(end)),
accWithWindowEnd(2, testTimePtr(end)),
accWithWindowEnd(3, testTimePtr(now.Add(5*time.Hour))),
}
got := filterBySoonestReset(accounts)
require.Len(t, got, 2, "并列最早重置的账号都保留,交由后续 LRU 决定")
ids := map[int64]bool{got[0].account.ID: true, got[1].account.ID: true}
require.True(t, ids[1] && ids[2])
}
func TestFilterBySoonestReset_SingleOrEmptyUnchanged(t *testing.T) {
require.Empty(t, filterBySoonestReset(nil))
single := []accountWithLoad{accWithWindowEnd(1, nil)}
require.Len(t, filterBySoonestReset(single), 1)
}
@@ -745,6 +745,35 @@ func (s *defaultOpenAIAccountScheduler) buildOpenAIAccountLoadPlan(
plan.loadSkew = calcLoadSkewByMoments(loadRateSum, loadRateSumSquares, len(candidates))
weights := s.service.openAIWSSchedulerWeights()
// Reset 因子(use-it-or-lose-it):在拥有「未来会话窗口结束时间」的账号中,
// 剩余时间越短 → 因子越接近 1(越早重置越优先用尽)。无活跃窗口的账号因子为 0。
// 仅在 weights.Reset > 0 时计算,默认关闭不影响原有行为。
minResetRemaining, maxResetRemaining := 0.0, 0.0
hasResetSample := false
if weights.Reset > 0 {
now := time.Now()
for _, candidate := range candidates {
end := candidate.account.SessionWindowEnd
if end == nil || !now.Before(*end) {
continue
}
remaining := end.Sub(now).Seconds()
if !hasResetSample {
minResetRemaining, maxResetRemaining = remaining, remaining
hasResetSample = true
continue
}
if remaining < minResetRemaining {
minResetRemaining = remaining
}
if remaining > maxResetRemaining {
maxResetRemaining = remaining
}
}
}
now := time.Now()
for i := range candidates {
item := &candidates[i]
priorityFactor := 1.0
@@ -758,12 +787,24 @@ func (s *defaultOpenAIAccountScheduler) buildOpenAIAccountLoadPlan(
if item.hasTTFT && hasTTFTSample && maxTTFT > minTTFT {
ttftFactor = 1 - clamp01((item.ttft-minTTFT)/(maxTTFT-minTTFT))
}
resetFactor := 0.0
if weights.Reset > 0 && hasResetSample {
if end := item.account.SessionWindowEnd; end != nil && now.Before(*end) {
if maxResetRemaining > minResetRemaining {
resetFactor = 1 - clamp01((end.Sub(now).Seconds()-minResetRemaining)/(maxResetRemaining-minResetRemaining))
} else {
// 所有有窗口的账号剩余时间相同:一律给满分,让其优于无窗口账号。
resetFactor = 1
}
}
}
item.score = weights.Priority*priorityFactor +
weights.Load*loadFactor +
weights.Queue*queueFactor +
weights.ErrorRate*errorFactor +
weights.TTFT*ttftFactor
weights.TTFT*ttftFactor +
weights.Reset*resetFactor
}
plan.candidates = candidates
@@ -1415,6 +1456,7 @@ func (s *OpenAIGatewayService) openAIWSSchedulerWeights() GatewayOpenAIWSSchedul
Queue: s.cfg.Gateway.OpenAIWS.SchedulerScoreWeights.Queue,
ErrorRate: s.cfg.Gateway.OpenAIWS.SchedulerScoreWeights.ErrorRate,
TTFT: s.cfg.Gateway.OpenAIWS.SchedulerScoreWeights.TTFT,
Reset: s.cfg.Gateway.OpenAIWS.SchedulerScoreWeights.Reset,
}
}
return GatewayOpenAIWSSchedulerScoreWeightsView{
@@ -1423,6 +1465,7 @@ func (s *OpenAIGatewayService) openAIWSSchedulerWeights() GatewayOpenAIWSSchedul
Queue: 0.7,
ErrorRate: 0.8,
TTFT: 0.5,
Reset: 0.0,
}
}
@@ -1432,6 +1475,8 @@ type GatewayOpenAIWSSchedulerScoreWeightsView struct {
Queue float64
ErrorRate float64
TTFT float64
// Reset 倾向「会话窗口最早重置」的账号;0 表示关闭(默认)。
Reset float64
}
func clamp01(value float64) float64 {
@@ -0,0 +1,77 @@
package service
import (
"testing"
"time"
"github.com/Wei-Shaw/sub2api/internal/config"
"github.com/stretchr/testify/require"
)
func openAIResetTestScheduler(reset float64) *defaultOpenAIAccountScheduler {
cfg := &config.Config{}
cfg.Gateway.OpenAIWS.SchedulerScoreWeights = config.GatewayOpenAIWSSchedulerScoreWeights{
Priority: 1.0,
Load: 1.0,
Queue: 0.7,
ErrorRate: 0.8,
TTFT: 0.5,
Reset: reset,
}
return &defaultOpenAIAccountScheduler{service: &OpenAIGatewayService{cfg: cfg}}
}
func openAIPlanScores(plan openAIAccountLoadPlan) map[int64]float64 {
scores := make(map[int64]float64, len(plan.candidates))
for _, c := range plan.candidates {
scores[c.account.ID] = c.score
}
return scores
}
// Reset 权重 > 0 时,会话窗口最早重置的账号应获得更高分。
func TestBuildOpenAIAccountLoadPlan_ResetWeightPrefersSoonestReset(t *testing.T) {
now := time.Now()
soon := now.Add(1 * time.Hour)
later := now.Add(20 * time.Hour)
filtered := []*Account{
{ID: 1, Priority: 0, SessionWindowEnd: &later},
{ID: 2, Priority: 0, SessionWindowEnd: &soon},
}
sched := openAIResetTestScheduler(5.0)
plan := sched.buildOpenAIAccountLoadPlan(OpenAIAccountScheduleRequest{}, filtered, map[int64]*AccountLoadInfo{})
scores := openAIPlanScores(plan)
require.Greater(t, scores[2], scores[1], "重置时间最早的账号(ID=2)得分更高")
}
// Reset 权重为 0(默认)时,窗口重置时间不应影响打分,保持原有行为。
func TestBuildOpenAIAccountLoadPlan_ResetWeightZeroNoEffect(t *testing.T) {
now := time.Now()
soon := now.Add(1 * time.Hour)
later := now.Add(20 * time.Hour)
filtered := []*Account{
{ID: 1, Priority: 0, SessionWindowEnd: &later},
{ID: 2, Priority: 0, SessionWindowEnd: &soon},
}
sched := openAIResetTestScheduler(0.0)
plan := sched.buildOpenAIAccountLoadPlan(OpenAIAccountScheduleRequest{}, filtered, map[int64]*AccountLoadInfo{})
scores := openAIPlanScores(plan)
require.Equal(t, scores[1], scores[2], "Reset 权重为 0 时两账号得分相同")
}
// 无活跃窗口的账号 reset 因子为 0,应低于拥有未来窗口的账号。
func TestBuildOpenAIAccountLoadPlan_ResetWeightIgnoresNilWindow(t *testing.T) {
now := time.Now()
soon := now.Add(2 * time.Hour)
filtered := []*Account{
{ID: 1, Priority: 0, SessionWindowEnd: nil},
{ID: 2, Priority: 0, SessionWindowEnd: &soon},
}
sched := openAIResetTestScheduler(5.0)
plan := sched.buildOpenAIAccountLoadPlan(OpenAIAccountScheduleRequest{}, filtered, map[int64]*AccountLoadInfo{})
scores := openAIPlanScores(plan)
require.Greater(t, scores[2], scores[1], "拥有活跃窗口的账号得分高于无窗口账号")
}
+7
View File
@@ -320,6 +320,9 @@ gateway:
queue: 0.7
error_rate: 0.8
ttft: 0.5
# use-it-or-lose-it:倾向「会话窗口最早重置」的账号,剩余重置时间越短得分越高。
# 0 表示关闭(默认,不改变原有行为);调大可让即将重置的账号被优先用尽。
reset: 0.0
# OpenAI 高级调度器补充配置
openai_scheduler:
# 是否允许 session_hash sticky 在账号健康度恶化时临时逃逸;false 可一键回退旧行为
@@ -421,6 +424,10 @@ gateway:
# Fallback max waiting queue size
# 兜底最大排队长度
fallback_max_waiting: 100
# Prefer the account whose session window resets soonest (use-it-or-lose-it).
# 负载感知选择时优先用尽「会话窗口最早重置」的账号;false 保持
# 原有「优先级 → 负载率 → LRU」行为(默认)。
prefer_soonest_reset: false
# Enable batch load calculation for scheduling
# 启用调度批量负载计算
load_batch_enabled: true