mirror of
https://github.com/Wei-Shaw/sub2api.git
synced 2026-10-16 07:33:06 +08:00
Yoga Sakti
4805f069e4
`filterSchedulerExtra` 的白名单没有 `base_rpm` / `rpm_strategy` / `rpm_sticky_buffer`，写进 `sched:meta:<id>` 的精简投影因此丢掉整套 RPM 配置。 候选过滤走 `ListSchedulableAccounts`，缓存命中时读的正是这份投影： `isAccountSchedulableForRPM` 拿到的账号上 `GetBaseRPM()` 返回 0，`baseRPM <= 0` 直接放行，已配置限流的账号在选号阶段被当作未开启 RPM，单个自然分钟内被超额调度。 缓存未命中时回退数据库读到完整 extra，所以现象是间歇性的。 选中的账号随后由 `hydrateSelectedAccount` 从 `sched:acc:<id>` 重新水合， RPM 计数与用户消息队列限速读的是完整账号，不受影响；缺的只是选号阶段这道门。 修法沿用投影里既有的处理方式（见 `openai_passthrough` 一组注释）：把判定依赖的 键补进白名单，而不是在判定侧另开一条读取完整账号的路径。 回归测试覆盖序列化往返后的三个键：`tiered` 子用例证明缓冲被裁掉时红区会被误判成 黄区，`sticky_exempt` 子用例证明策略被裁掉时会退回三区判定。 Refs #7225
1185 lines
47 KiB
Go
1185 lines
47 KiB
Go
//go:build unit
|
||
|
||
package repository
|
||
|
||
import (
|
||
"context"
|
||
"encoding/hex"
|
||
"encoding/json"
|
||
"fmt"
|
||
"strconv"
|
||
"strings"
|
||
"testing"
|
||
"time"
|
||
|
||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||
"github.com/alicebob/miniredis/v2"
|
||
"github.com/redis/go-redis/v9"
|
||
"github.com/stretchr/testify/require"
|
||
)
|
||
|
||
func newSchedulerCacheUnit(t *testing.T) *schedulerCache {
|
||
cache, _ := newSchedulerCacheUnitWithRedis(t)
|
||
return cache
|
||
}
|
||
|
||
func newSchedulerCacheUnitWithRedis(t *testing.T) (*schedulerCache, *miniredis.Miniredis) {
|
||
t.Helper()
|
||
mr := miniredis.RunT(t)
|
||
rdb := redis.NewClient(&redis.Options{Addr: mr.Addr()})
|
||
t.Cleanup(func() { _ = rdb.Close() })
|
||
cache, ok := newSchedulerCacheWithChunkSizes(rdb, defaultSchedulerSnapshotMGetChunkSize, defaultSchedulerSnapshotWriteChunkSize).(*schedulerCache)
|
||
require.True(t, ok)
|
||
return cache, mr
|
||
}
|
||
|
||
func TestSchedulerCacheWriteAccountIDsSkipsUnencodableTimes(t *testing.T) {
|
||
ctx := context.Background()
|
||
cache := newSchedulerCacheUnit(t)
|
||
invalidTime := time.Date(10000, time.January, 1, 0, 0, 0, 0, time.UTC)
|
||
|
||
accountIDs, err := cache.writeAccountIDs(ctx, []service.Account{
|
||
{ID: 111, Platform: service.PlatformOpenAI, Type: service.AccountTypeAPIKey},
|
||
{ID: 112, Platform: service.PlatformOpenAI, Type: service.AccountTypeAPIKey, ExpiresAt: &invalidTime},
|
||
})
|
||
require.NoError(t, err)
|
||
require.Equal(t, []int64{111}, accountIDs)
|
||
|
||
cached, err := cache.GetAccount(ctx, 111)
|
||
require.NoError(t, err)
|
||
require.NotNil(t, cached)
|
||
|
||
invalid, err := cache.GetAccount(ctx, 112)
|
||
require.NoError(t, err)
|
||
require.Nil(t, invalid)
|
||
}
|
||
|
||
func TestSchedulerCacheSetAccountClearsUnencodablePayload(t *testing.T) {
|
||
ctx := context.Background()
|
||
cache := newSchedulerCacheUnit(t)
|
||
|
||
account := service.Account{ID: 113, Platform: service.PlatformOpenAI, Type: service.AccountTypeAPIKey}
|
||
require.NoError(t, cache.SetAccount(ctx, &account))
|
||
|
||
invalidTime := time.Date(10000, time.January, 1, 0, 0, 0, 0, time.UTC)
|
||
account.ExpiresAt = &invalidTime
|
||
require.NoError(t, cache.SetAccount(ctx, &account))
|
||
|
||
cached, err := cache.GetAccount(ctx, account.ID)
|
||
require.NoError(t, err)
|
||
require.Nil(t, cached)
|
||
}
|
||
|
||
func TestSchedulerCacheUpdateLastUsedClearsUnencodablePayload(t *testing.T) {
|
||
ctx := context.Background()
|
||
cache := newSchedulerCacheUnit(t)
|
||
account := service.Account{ID: 114, Platform: service.PlatformOpenAI, Type: service.AccountTypeAPIKey}
|
||
require.NoError(t, cache.SetAccount(ctx, &account))
|
||
|
||
invalidTime := time.Date(10000, time.January, 1, 0, 0, 0, 0, time.UTC)
|
||
require.NoError(t, cache.UpdateLastUsed(ctx, map[int64]time.Time{account.ID: invalidTime}))
|
||
|
||
cached, err := cache.GetAccount(ctx, account.ID)
|
||
require.NoError(t, err)
|
||
require.Nil(t, cached)
|
||
}
|
||
|
||
func TestSchedulerCacheSnapshotAccountIDReusePreservesPayloadAndMembers(t *testing.T) {
|
||
ctx := context.Background()
|
||
cache, _ := newSchedulerCacheUnitWithRedis(t)
|
||
invalidTime := time.Date(10000, time.January, 1, 0, 0, 0, 0, time.UTC)
|
||
validOne := service.Account{
|
||
ID: 701,
|
||
Name: "first",
|
||
Platform: service.PlatformOpenAI,
|
||
Type: service.AccountTypeOAuth,
|
||
Credentials: map[string]any{"model_mapping": map[string]any{"z": "last", "a": "first"}},
|
||
Extra: map[string]any{"mixed_scheduling": true},
|
||
GroupIDs: []int64{17},
|
||
}
|
||
validTwo := service.Account{ID: 702, Name: "second", Platform: service.PlatformOpenAI, Type: service.AccountTypeAPIKey}
|
||
invalid := service.Account{ID: 799, Platform: service.PlatformOpenAI, Type: service.AccountTypeAPIKey, ExpiresAt: &invalidTime}
|
||
accounts := []service.Account{validOne, invalid, validTwo, validOne}
|
||
|
||
single := service.SchedulerBucket{GroupID: 17, Platform: service.PlatformOpenAI, Mode: service.SchedulerModeSingle}
|
||
singleToken, err := cache.CaptureBucketWriteToken(ctx, single)
|
||
require.NoError(t, err)
|
||
accountIDs, err := cache.SetSnapshotAndReturnAccountIDs(ctx, single, singleToken, accounts)
|
||
require.NoError(t, err)
|
||
require.Equal(t, []int64{701, 702, 701}, accountIDs, "应保留可编码账号的原顺序和重复项")
|
||
|
||
wantFull, err := json.Marshal(validOne)
|
||
require.NoError(t, err)
|
||
wantMeta, err := json.Marshal(buildSchedulerMetadataAccount(validOne))
|
||
require.NoError(t, err)
|
||
fullBefore, err := cache.rdb.Get(ctx, schedulerAccountKey("701")).Bytes()
|
||
require.NoError(t, err)
|
||
metaBefore, err := cache.rdb.Get(ctx, schedulerAccountMetaKey("701")).Bytes()
|
||
require.NoError(t, err)
|
||
require.Equal(t, wantFull, fullBefore)
|
||
require.Equal(t, wantMeta, metaBefore)
|
||
|
||
forced := service.SchedulerBucket{GroupID: 17, Platform: service.PlatformOpenAI, Mode: service.SchedulerModeForced}
|
||
forcedToken, err := cache.CaptureBucketWriteToken(ctx, forced)
|
||
require.NoError(t, err)
|
||
require.NoError(t, cache.SetSnapshotByAccountIDs(ctx, forced, forcedToken, accountIDs))
|
||
|
||
fullAfter, err := cache.rdb.Get(ctx, schedulerAccountKey("701")).Bytes()
|
||
require.NoError(t, err)
|
||
metaAfter, err := cache.rdb.Get(ctx, schedulerAccountMetaKey("701")).Bytes()
|
||
require.NoError(t, err)
|
||
require.Equal(t, fullBefore, fullAfter, "ID-only 路径不得重写完整账号键")
|
||
require.Equal(t, metaBefore, metaAfter, "ID-only 路径不得重写调度元数据键")
|
||
|
||
for _, bucket := range []service.SchedulerBucket{single, forced} {
|
||
version, err := cache.rdb.Get(ctx, schedulerBucketKey(schedulerActivePrefix, bucket)).Result()
|
||
require.NoError(t, err)
|
||
members, err := cache.rdb.ZRange(ctx, schedulerSnapshotKey(bucket, version), 0, -1).Result()
|
||
require.NoError(t, err)
|
||
require.Equal(t, []string{"702", "701"}, members, bucket.String())
|
||
}
|
||
missing, err := cache.GetAccount(ctx, invalid.ID)
|
||
require.NoError(t, err)
|
||
require.Nil(t, missing)
|
||
}
|
||
|
||
func TestSchedulerCacheSetSnapshotMatchesIDPublishing(t *testing.T) {
|
||
ctx := context.Background()
|
||
cache, _ := newSchedulerCacheUnitWithRedis(t)
|
||
invalidTime := time.Date(10000, time.January, 1, 0, 0, 0, 0, time.UTC)
|
||
validOne := service.Account{
|
||
ID: 721,
|
||
Name: "first",
|
||
Platform: service.PlatformOpenAI,
|
||
Type: service.AccountTypeOAuth,
|
||
Credentials: map[string]any{"model_mapping": map[string]any{"source": "target"}},
|
||
Extra: map[string]any{"mixed_scheduling": true},
|
||
GroupIDs: []int64{21},
|
||
}
|
||
validTwo := service.Account{ID: 722, Name: "second", Platform: service.PlatformOpenAI, Type: service.AccountTypeAPIKey}
|
||
invalid := service.Account{ID: 799, Platform: service.PlatformOpenAI, Type: service.AccountTypeAPIKey, ExpiresAt: &invalidTime}
|
||
accounts := []service.Account{validOne, invalid, validTwo, validOne}
|
||
|
||
normal := service.SchedulerBucket{GroupID: 21, Platform: service.PlatformOpenAI, Mode: service.SchedulerModeSingle}
|
||
normalToken, err := cache.CaptureBucketWriteToken(ctx, normal)
|
||
require.NoError(t, err)
|
||
require.NoError(t, cache.SetSnapshot(ctx, normal, normalToken, accounts))
|
||
|
||
fullBefore, err := cache.rdb.Get(ctx, schedulerAccountKey("721")).Bytes()
|
||
require.NoError(t, err)
|
||
metaBefore, err := cache.rdb.Get(ctx, schedulerAccountMetaKey("721")).Bytes()
|
||
require.NoError(t, err)
|
||
|
||
idOnly := service.SchedulerBucket{GroupID: 21, Platform: service.PlatformOpenAI, Mode: service.SchedulerModeForced}
|
||
idOnlyToken, err := cache.CaptureBucketWriteToken(ctx, idOnly)
|
||
require.NoError(t, err)
|
||
accountIDs, err := cache.SetSnapshotAndReturnAccountIDs(ctx, idOnly, idOnlyToken, accounts)
|
||
require.NoError(t, err)
|
||
require.Equal(t, []int64{721, 722, 721}, accountIDs)
|
||
|
||
fullAfter, err := cache.rdb.Get(ctx, schedulerAccountKey("721")).Bytes()
|
||
require.NoError(t, err)
|
||
metaAfter, err := cache.rdb.Get(ctx, schedulerAccountMetaKey("721")).Bytes()
|
||
require.NoError(t, err)
|
||
require.Equal(t, fullBefore, fullAfter, "普通快照和 ID 发布必须写入相同完整账号 payload")
|
||
require.Equal(t, metaBefore, metaAfter, "普通快照和 ID 发布必须写入相同元数据 payload")
|
||
|
||
for _, bucket := range []service.SchedulerBucket{normal, idOnly} {
|
||
version, err := cache.rdb.Get(ctx, schedulerBucketKey(schedulerActivePrefix, bucket)).Result()
|
||
require.NoError(t, err)
|
||
members, err := cache.rdb.ZRange(ctx, schedulerSnapshotKey(bucket, version), 0, -1).Result()
|
||
require.NoError(t, err)
|
||
require.Equal(t, []string{"722", "721"}, members, bucket.String())
|
||
}
|
||
}
|
||
|
||
func TestSchedulerCacheSnapshotAccountIDReuseKeepsEmptySnapshotSemantics(t *testing.T) {
|
||
ctx := context.Background()
|
||
cache := newSchedulerCacheUnit(t)
|
||
invalidTime := time.Date(10000, time.January, 1, 0, 0, 0, 0, time.UTC)
|
||
accounts := []service.Account{{ID: 811, Platform: service.PlatformOpenAI, ExpiresAt: &invalidTime}}
|
||
|
||
single := service.SchedulerBucket{GroupID: 18, Platform: service.PlatformOpenAI, Mode: service.SchedulerModeSingle}
|
||
singleToken, err := cache.CaptureBucketWriteToken(ctx, single)
|
||
require.NoError(t, err)
|
||
accountIDs, err := cache.SetSnapshotAndReturnAccountIDs(ctx, single, singleToken, accounts)
|
||
require.NoError(t, err)
|
||
require.Empty(t, accountIDs)
|
||
|
||
forced := service.SchedulerBucket{GroupID: 18, Platform: service.PlatformOpenAI, Mode: service.SchedulerModeForced}
|
||
forcedToken, err := cache.CaptureBucketWriteToken(ctx, forced)
|
||
require.NoError(t, err)
|
||
require.NoError(t, cache.SetSnapshotByAccountIDs(ctx, forced, forcedToken, accountIDs))
|
||
|
||
for _, bucket := range []service.SchedulerBucket{single, forced} {
|
||
ready, err := cache.rdb.Get(ctx, schedulerBucketKey(schedulerReadyPrefix, bucket)).Result()
|
||
require.NoError(t, err)
|
||
require.Equal(t, "1", ready)
|
||
snapshot, hit, err := cache.GetSnapshot(ctx, bucket)
|
||
require.NoError(t, err)
|
||
require.False(t, hit, bucket.String())
|
||
require.Nil(t, snapshot)
|
||
}
|
||
}
|
||
|
||
func TestSchedulerCacheSetSnapshotByAccountIDsKeepsFencing(t *testing.T) {
|
||
ctx := context.Background()
|
||
cache := newSchedulerCacheUnit(t)
|
||
bucket := service.SchedulerBucket{GroupID: 19, Platform: service.PlatformOpenAI, Mode: service.SchedulerModeForced}
|
||
|
||
err := cache.SetSnapshotByAccountIDs(ctx, bucket, service.SchedulerBucketWriteToken{}, []int64{901})
|
||
require.ErrorIs(t, err, service.ErrSchedulerBucketWriteFenced)
|
||
_, err = cache.rdb.Get(ctx, schedulerBucketKey(schedulerVersionPrefix, bucket)).Result()
|
||
require.ErrorIs(t, err, redis.Nil)
|
||
|
||
token, err := cache.CaptureBucketWriteToken(ctx, bucket)
|
||
require.NoError(t, err)
|
||
require.NoError(t, cache.RetireBucket(ctx, bucket))
|
||
err = cache.SetSnapshotByAccountIDs(ctx, bucket, token, []int64{901})
|
||
require.ErrorIs(t, err, service.ErrSchedulerBucketRetired)
|
||
}
|
||
|
||
func TestSchedulerCacheSetSnapshotByAccountIDsDoesNotResurrectDeletedAccount(t *testing.T) {
|
||
ctx := context.Background()
|
||
cache := newSchedulerCacheUnit(t)
|
||
account := service.Account{ID: 902, Platform: service.PlatformOpenAI, Type: service.AccountTypeOAuth}
|
||
single := service.SchedulerBucket{GroupID: 20, Platform: service.PlatformOpenAI, Mode: service.SchedulerModeSingle}
|
||
singleToken, err := cache.CaptureBucketWriteToken(ctx, single)
|
||
require.NoError(t, err)
|
||
accountIDs, err := cache.SetSnapshotAndReturnAccountIDs(ctx, single, singleToken, []service.Account{account})
|
||
require.NoError(t, err)
|
||
require.Equal(t, []int64{account.ID}, accountIDs)
|
||
require.NoError(t, cache.DeleteAccount(ctx, account.ID))
|
||
|
||
forced := service.SchedulerBucket{GroupID: 20, Platform: service.PlatformOpenAI, Mode: service.SchedulerModeForced}
|
||
forcedToken, err := cache.CaptureBucketWriteToken(ctx, forced)
|
||
require.NoError(t, err)
|
||
require.NoError(t, cache.SetSnapshotByAccountIDs(ctx, forced, forcedToken, accountIDs))
|
||
|
||
full, err := cache.GetAccount(ctx, account.ID)
|
||
require.NoError(t, err)
|
||
require.Nil(t, full, "ID-only 发布不得复活已删除的完整账号键")
|
||
snapshot, hit, err := cache.GetSnapshot(ctx, forced)
|
||
require.NoError(t, err)
|
||
require.False(t, hit, "元数据缺失时必须安全回源，而不是返回残缺快照")
|
||
require.Nil(t, snapshot)
|
||
}
|
||
|
||
func TestMarshalSchedulerCacheAccountKeepsEncodingJSONWireFormat(t *testing.T) {
|
||
cases := []struct {
|
||
name string
|
||
account service.Account
|
||
}{
|
||
{name: "nil collections", account: service.Account{ID: 801}},
|
||
{name: "empty collections", account: service.Account{
|
||
ID: 802,
|
||
Credentials: map[string]any{},
|
||
Extra: map[string]any{},
|
||
GroupIDs: []int64{},
|
||
Groups: []*service.Group{},
|
||
}},
|
||
{name: "nested maps and escaping", account: service.Account{
|
||
ID: 803,
|
||
Credentials: map[string]any{"model_mapping": map[string]any{"z": "<last>", "a": "&first"}},
|
||
Extra: map[string]any{"mixed_scheduling": true},
|
||
}},
|
||
}
|
||
for _, tc := range cases {
|
||
t.Run(tc.name, func(t *testing.T) {
|
||
full, meta, err := marshalSchedulerCacheAccount(tc.account)
|
||
require.NoError(t, err)
|
||
wantFull, err := json.Marshal(tc.account)
|
||
require.NoError(t, err)
|
||
wantMeta, err := json.Marshal(buildSchedulerMetadataAccount(tc.account))
|
||
require.NoError(t, err)
|
||
require.Equal(t, wantFull, full)
|
||
require.Equal(t, wantMeta, meta)
|
||
})
|
||
}
|
||
}
|
||
|
||
func TestBuildSchedulerMetadataAccount_KeepsOpenAIWSFlags(t *testing.T) {
|
||
account := service.Account{
|
||
ID: 42,
|
||
Platform: service.PlatformOpenAI,
|
||
Type: service.AccountTypeOAuth,
|
||
Extra: map[string]any{
|
||
"openai_oauth_responses_websockets_v2_enabled": true,
|
||
"openai_oauth_responses_websockets_v2_mode": service.OpenAIWSIngressModePassthrough,
|
||
"openai_ws_force_http": true,
|
||
"openai_responses_mode": "force_chat_completions",
|
||
"openai_responses_supported": false,
|
||
"codex_fingerprint_mode": "session",
|
||
"codex_fingerprint_seed": "11111111-1111-4111-8111-111111111111",
|
||
"mixed_scheduling": true,
|
||
"unused_large_field": "drop-me",
|
||
},
|
||
}
|
||
|
||
got := buildSchedulerMetadataAccount(account)
|
||
|
||
require.Equal(t, true, got.Extra["openai_oauth_responses_websockets_v2_enabled"])
|
||
require.Equal(t, service.OpenAIWSIngressModePassthrough, got.Extra["openai_oauth_responses_websockets_v2_mode"])
|
||
require.Equal(t, true, got.Extra["openai_ws_force_http"])
|
||
require.Equal(t, "force_chat_completions", got.Extra["openai_responses_mode"])
|
||
require.Equal(t, false, got.Extra["openai_responses_supported"])
|
||
require.Equal(t, "session", got.Extra["codex_fingerprint_mode"])
|
||
require.Equal(t, "11111111-1111-4111-8111-111111111111", got.Extra["codex_fingerprint_seed"])
|
||
require.Equal(t, true, got.Extra["mixed_scheduling"])
|
||
require.Nil(t, got.Extra["unused_large_field"])
|
||
}
|
||
|
||
func TestBuildSchedulerMetadataAccount_KeepsGrokMediaEligibility(t *testing.T) {
|
||
t.Run("explicit override", func(t *testing.T) {
|
||
account := service.Account{
|
||
ID: 43,
|
||
Platform: service.PlatformGrok,
|
||
Type: service.AccountTypeOAuth,
|
||
Extra: map[string]any{
|
||
service.GrokMediaEligibleExtraKey: false,
|
||
"unused_large_field": "drop-me",
|
||
},
|
||
}
|
||
|
||
got := buildSchedulerMetadataAccount(account)
|
||
|
||
eligible, reason := got.GrokMediaGenerationEligibility()
|
||
require.False(t, eligible)
|
||
require.Equal(t, "override_disabled", reason)
|
||
require.Equal(t, false, got.Extra[service.GrokMediaEligibleExtraKey])
|
||
require.Nil(t, got.Extra["unused_large_field"])
|
||
})
|
||
|
||
t.Run("forbidden billing observation", func(t *testing.T) {
|
||
account := service.Account{
|
||
ID: 44,
|
||
Platform: service.PlatformGrok,
|
||
Type: service.AccountTypeOAuth,
|
||
Extra: map[string]any{
|
||
"grok_billing_snapshot": map[string]any{
|
||
"status_code": 200,
|
||
"weekly_status_code": 403,
|
||
"monthly_status_code": 200,
|
||
},
|
||
},
|
||
}
|
||
|
||
got := buildSchedulerMetadataAccount(account)
|
||
|
||
eligible, reason := got.GrokMediaGenerationEligibility()
|
||
require.False(t, eligible)
|
||
require.Equal(t, "billing_forbidden", reason)
|
||
require.NotNil(t, got.Extra["grok_billing_snapshot"])
|
||
})
|
||
}
|
||
|
||
func TestBuildSchedulerMetadataAccount_KeepsSlimGroupMembership(t *testing.T) {
|
||
account := service.Account{
|
||
ID: 42,
|
||
Platform: service.PlatformAnthropic,
|
||
GroupIDs: []int64{7, 9, 7, 0},
|
||
AccountGroups: []service.AccountGroup{
|
||
{
|
||
AccountID: 42,
|
||
GroupID: 7,
|
||
Priority: 2,
|
||
Account: &service.Account{ID: 42, Name: "drop-from-metadata"},
|
||
Group: &service.Group{ID: 7, Name: "drop-from-metadata"},
|
||
},
|
||
{
|
||
AccountID: 42,
|
||
GroupID: 11,
|
||
Priority: 3,
|
||
Group: &service.Group{ID: 11, Name: "drop-from-metadata"},
|
||
},
|
||
{
|
||
AccountID: 42,
|
||
GroupID: 0,
|
||
Priority: 4,
|
||
},
|
||
},
|
||
}
|
||
|
||
got := buildSchedulerMetadataAccount(account)
|
||
|
||
require.Equal(t, []int64{7, 9, 11}, got.GroupIDs)
|
||
require.Len(t, got.AccountGroups, 2)
|
||
require.Equal(t, int64(42), got.AccountGroups[0].AccountID)
|
||
require.Equal(t, int64(7), got.AccountGroups[0].GroupID)
|
||
require.Equal(t, 2, got.AccountGroups[0].Priority)
|
||
require.Nil(t, got.AccountGroups[0].Account)
|
||
require.Nil(t, got.AccountGroups[0].Group)
|
||
require.Equal(t, int64(11), got.AccountGroups[1].GroupID)
|
||
require.Nil(t, got.Groups)
|
||
}
|
||
|
||
func TestBuildSchedulerMetadataAccount_KeepsQuotaAutoPauseFields(t *testing.T) {
|
||
account := service.Account{
|
||
ID: 88,
|
||
Extra: map[string]any{
|
||
"codex_5h_used_percent": 12.34,
|
||
"codex_7d_used_percent": 56.78,
|
||
"codex_5h_reset_at": "2026-05-29T10:00:00Z",
|
||
"codex_7d_reset_at": "2026-06-01T10:00:00Z",
|
||
"codex_5h_reset_after_seconds": 300,
|
||
"codex_7d_reset_after_seconds": 600,
|
||
"codex_usage_updated_at": "2026-05-29T09:00:00Z",
|
||
"auto_pause_5h_threshold": 0.95,
|
||
"auto_pause_7d_threshold": 0.96,
|
||
"auto_pause_5h_disabled": true,
|
||
"auto_pause_7d_disabled": false,
|
||
},
|
||
}
|
||
|
||
got := buildSchedulerMetadataAccount(account)
|
||
|
||
require.Equal(t, 12.34, got.Extra["codex_5h_used_percent"])
|
||
require.Equal(t, 56.78, got.Extra["codex_7d_used_percent"])
|
||
require.Equal(t, "2026-05-29T10:00:00Z", got.Extra["codex_5h_reset_at"])
|
||
require.Equal(t, "2026-06-01T10:00:00Z", got.Extra["codex_7d_reset_at"])
|
||
require.Equal(t, 300, got.Extra["codex_5h_reset_after_seconds"])
|
||
require.Equal(t, 600, got.Extra["codex_7d_reset_after_seconds"])
|
||
require.Equal(t, "2026-05-29T09:00:00Z", got.Extra["codex_usage_updated_at"])
|
||
require.Equal(t, 0.95, got.Extra["auto_pause_5h_threshold"])
|
||
require.Equal(t, 0.96, got.Extra["auto_pause_7d_threshold"])
|
||
require.Equal(t, true, got.Extra["auto_pause_5h_disabled"])
|
||
require.Equal(t, false, got.Extra["auto_pause_7d_disabled"])
|
||
}
|
||
|
||
func TestBuildSchedulerMetadataAccount_KeepsQuotaStateForCachedAccounts(t *testing.T) {
|
||
now := time.Now().UTC()
|
||
activeStart := now.Add(-time.Hour).Format(time.RFC3339)
|
||
expiredDailyStart := now.Add(-25 * time.Hour).Format(time.RFC3339)
|
||
expiredWeeklyStart := now.Add(-8 * 24 * time.Hour).Format(time.RFC3339)
|
||
weeklyResetDay := float64(now.AddDate(0, 0, 1).Weekday())
|
||
|
||
cases := []struct {
|
||
name string
|
||
platform string
|
||
typ string
|
||
extra map[string]any
|
||
quotaExceeded bool
|
||
}{
|
||
{
|
||
name: "anthropic api key total quota exhausted", platform: service.PlatformAnthropic, typ: service.AccountTypeAPIKey,
|
||
extra: map[string]any{"quota_limit": 10.0, "quota_used": 10.0}, quotaExceeded: true,
|
||
},
|
||
{
|
||
name: "gemini api key rolling daily quota exhausted", platform: service.PlatformGemini, typ: service.AccountTypeAPIKey,
|
||
extra: map[string]any{
|
||
"quota_daily_limit": 20.0, "quota_daily_used": 20.0,
|
||
"quota_daily_start": activeStart, "quota_daily_reset_mode": "rolling",
|
||
}, quotaExceeded: true,
|
||
},
|
||
{
|
||
name: "gemini api key expired rolling daily window", platform: service.PlatformGemini, typ: service.AccountTypeAPIKey,
|
||
extra: map[string]any{
|
||
"quota_daily_limit": 20.0, "quota_daily_used": 20.0,
|
||
"quota_daily_start": expiredDailyStart, "quota_daily_reset_mode": "rolling",
|
||
},
|
||
},
|
||
{
|
||
name: "bedrock fixed weekly quota exhausted", platform: service.PlatformAnthropic, typ: service.AccountTypeBedrock,
|
||
extra: map[string]any{
|
||
"quota_weekly_limit": 30.0, "quota_weekly_used": 30.0, "quota_weekly_start": activeStart,
|
||
"quota_weekly_reset_mode": "fixed", "quota_weekly_reset_day": weeklyResetDay,
|
||
"quota_weekly_reset_hour": 0.0, "quota_reset_timezone": "UTC",
|
||
}, quotaExceeded: true,
|
||
},
|
||
{
|
||
name: "bedrock expired fixed weekly window", platform: service.PlatformAnthropic, typ: service.AccountTypeBedrock,
|
||
extra: map[string]any{
|
||
"quota_weekly_limit": 30.0, "quota_weekly_used": 30.0, "quota_weekly_start": expiredWeeklyStart,
|
||
"quota_weekly_reset_mode": "fixed", "quota_weekly_reset_day": weeklyResetDay,
|
||
"quota_weekly_reset_hour": 0.0, "quota_reset_timezone": "UTC",
|
||
},
|
||
},
|
||
}
|
||
for i, tc := range cases {
|
||
t.Run(tc.name, func(t *testing.T) {
|
||
extra := make(map[string]any, len(tc.extra)+1)
|
||
for key, value := range tc.extra {
|
||
extra[key] = value
|
||
}
|
||
extra["unrelated"] = "drop me"
|
||
account := service.Account{
|
||
ID: int64(46690 + i), Platform: tc.platform, Type: tc.typ, Extra: extra,
|
||
Status: service.StatusActive, Schedulable: true,
|
||
}
|
||
cache := newSchedulerCacheUnit(t)
|
||
ctx := context.Background()
|
||
bucket := service.SchedulerBucket{GroupID: int64(46690 + i), Platform: tc.platform, Mode: service.SchedulerModeSingle}
|
||
token, err := cache.CaptureBucketWriteToken(ctx, bucket)
|
||
require.NoError(t, err)
|
||
require.NoError(t, cache.SetSnapshot(ctx, bucket, token, []service.Account{account}))
|
||
|
||
snapshot, hit, err := cache.GetSnapshot(ctx, bucket)
|
||
require.NoError(t, err)
|
||
require.True(t, hit)
|
||
require.Len(t, snapshot, 1)
|
||
cached := snapshot[0]
|
||
require.Equal(t, tc.extra, cached.Extra)
|
||
require.NotContains(t, cached.Extra, "unrelated")
|
||
require.Equal(t, tc.quotaExceeded, cached.IsQuotaExceeded())
|
||
require.Equal(t, !tc.quotaExceeded, cached.IsSchedulable())
|
||
})
|
||
}
|
||
}
|
||
|
||
func TestBuildSchedulerMetadataAccount_KeepsModelRateLimits(t *testing.T) {
|
||
account := service.Account{
|
||
ID: 90,
|
||
Platform: service.PlatformAntigravity,
|
||
Extra: map[string]any{
|
||
"model_rate_limits": map[string]any{
|
||
"gemini-3-flash": map[string]any{
|
||
"rate_limit_reset_at": "2026-05-30T10:10:00Z",
|
||
},
|
||
"antigravity:gemini": map[string]any{
|
||
"rate_limit_reset_at": "2026-05-30T10:10:00Z",
|
||
},
|
||
},
|
||
"unused_large_field": "drop-me",
|
||
},
|
||
}
|
||
|
||
got := buildSchedulerMetadataAccount(account)
|
||
|
||
limits, ok := got.Extra["model_rate_limits"].(map[string]any)
|
||
require.True(t, ok)
|
||
require.Contains(t, limits, "gemini-3-flash")
|
||
require.Contains(t, limits, "antigravity:gemini")
|
||
require.Nil(t, got.Extra["unused_large_field"])
|
||
}
|
||
|
||
func TestBuildSchedulerMetadataAccount_KeepsSparkShadowRoutingIdentity(t *testing.T) {
|
||
parentID := int64(100)
|
||
account := service.Account{
|
||
ID: 200,
|
||
Platform: service.PlatformOpenAI,
|
||
Type: service.AccountTypeOAuth,
|
||
ParentAccountID: &parentID,
|
||
QuotaDimension: service.QuotaDimensionSpark,
|
||
Credentials: map[string]any{
|
||
"model_mapping": map[string]any{
|
||
"gpt-5.3-codex-spark": "gpt-5.3-codex-spark",
|
||
},
|
||
"compact_model_mapping": map[string]any{
|
||
"gpt-5.4": "gpt-5.4-openai-compact",
|
||
},
|
||
"access_token": "drop-me",
|
||
},
|
||
}
|
||
|
||
got := buildSchedulerMetadataAccount(account)
|
||
|
||
require.NotNil(t, got.ParentAccountID)
|
||
require.Equal(t, parentID, *got.ParentAccountID)
|
||
require.Equal(t, service.QuotaDimensionSpark, got.QuotaDimension)
|
||
require.Equal(t, map[string]any{"gpt-5.3-codex-spark": "gpt-5.3-codex-spark"}, got.Credentials["model_mapping"])
|
||
require.Equal(t, map[string]any{"gpt-5.4": "gpt-5.4-openai-compact"}, got.Credentials["compact_model_mapping"])
|
||
require.Nil(t, got.Credentials["access_token"])
|
||
}
|
||
|
||
func TestSchedulerCacheBucketRetirementFencesWritersAndReopen(t *testing.T) {
|
||
ctx := context.Background()
|
||
cache, mr := newSchedulerCacheUnitWithRedis(t)
|
||
bucket := service.SchedulerBucket{GroupID: 41, Platform: service.PlatformOpenAI, Mode: service.SchedulerModeSingle}
|
||
otherBucket := service.SchedulerBucket{GroupID: 42, Platform: service.PlatformOpenAI, Mode: service.SchedulerModeSingle}
|
||
account := service.Account{ID: 4101, Platform: service.PlatformOpenAI, Type: service.AccountTypeAPIKey}
|
||
|
||
token, err := cache.CaptureBucketWriteToken(ctx, bucket)
|
||
require.NoError(t, err)
|
||
require.True(t, token.ValidFor(bucket))
|
||
require.NoError(t, cache.SetSnapshot(ctx, bucket, token, []service.Account{account}))
|
||
|
||
// A token is bound to the full bucket identity, not just an epoch number.
|
||
err = cache.SetSnapshot(ctx, otherBucket, token, []service.Account{account})
|
||
require.ErrorIs(t, err, service.ErrSchedulerBucketWriteFenced)
|
||
_, err = cache.rdb.Get(ctx, schedulerBucketKey(schedulerVersionPrefix, otherBucket)).Result()
|
||
require.ErrorIs(t, err, redis.Nil)
|
||
otherAccount := service.Account{ID: 4201, Platform: service.PlatformOpenAI, Type: service.AccountTypeAPIKey}
|
||
otherToken, err := cache.CaptureBucketWriteToken(ctx, otherBucket)
|
||
require.NoError(t, err)
|
||
require.NoError(t, cache.SetSnapshot(ctx, otherBucket, otherToken, []service.Account{otherAccount}))
|
||
otherEpoch := otherToken.Epoch
|
||
|
||
activeVersion, err := cache.rdb.Get(ctx, schedulerBucketKey(schedulerActivePrefix, bucket)).Result()
|
||
require.NoError(t, err)
|
||
require.NoError(t, cache.RetireBucket(ctx, bucket))
|
||
retiredEpoch, err := cache.rdb.Get(ctx, schedulerBucketKey(schedulerEpochPrefix, bucket)).Int64()
|
||
require.NoError(t, err)
|
||
require.Greater(t, retiredEpoch, token.Epoch)
|
||
|
||
// Retirement is idempotent and does not advance the epoch again.
|
||
require.NoError(t, cache.RetireBucket(ctx, bucket))
|
||
retiredEpochAgain, err := cache.rdb.Get(ctx, schedulerBucketKey(schedulerEpochPrefix, bucket)).Int64()
|
||
require.NoError(t, err)
|
||
require.Equal(t, retiredEpoch, retiredEpochAgain)
|
||
|
||
// New readers miss because ready/active were removed atomically. A reader that
|
||
// captured activeVersion before retirement may still finish against that version.
|
||
_, hit, err := cache.GetSnapshot(ctx, bucket)
|
||
require.NoError(t, err)
|
||
require.False(t, hit)
|
||
ids, err := cache.rdb.ZRange(ctx, schedulerSnapshotKey(bucket, activeVersion), 0, -1).Result()
|
||
require.NoError(t, err)
|
||
require.Equal(t, []string{"4101"}, ids)
|
||
ttl, err := cache.rdb.TTL(ctx, schedulerSnapshotKey(bucket, activeVersion)).Result()
|
||
require.NoError(t, err)
|
||
require.Positive(t, ttl)
|
||
require.LessOrEqual(t, ttl, time.Duration(snapshotGraceTTLSeconds)*time.Second)
|
||
|
||
buckets, err := cache.ListBuckets(ctx)
|
||
require.NoError(t, err)
|
||
require.NotContains(t, buckets, bucket)
|
||
require.Contains(t, buckets, otherBucket)
|
||
otherSnapshot, otherHit, err := cache.GetSnapshot(ctx, otherBucket)
|
||
require.NoError(t, err)
|
||
require.True(t, otherHit)
|
||
require.Len(t, otherSnapshot, 1)
|
||
require.Equal(t, otherAccount.ID, otherSnapshot[0].ID)
|
||
otherEpochAfter, err := cache.rdb.Get(ctx, schedulerBucketKey(schedulerEpochPrefix, otherBucket)).Int64()
|
||
require.NoError(t, err)
|
||
require.Equal(t, otherEpoch, otherEpochAfter)
|
||
|
||
_, err = cache.CaptureBucketWriteToken(ctx, bucket)
|
||
require.ErrorIs(t, err, service.ErrSchedulerBucketRetired)
|
||
versionBeforeRejectedWrite, err := cache.rdb.Get(ctx, schedulerBucketKey(schedulerVersionPrefix, bucket)).Int64()
|
||
require.NoError(t, err)
|
||
err = cache.SetSnapshot(ctx, bucket, token, []service.Account{account})
|
||
require.ErrorIs(t, err, service.ErrSchedulerBucketRetired)
|
||
versionAfterRejectedWrite, err := cache.rdb.Get(ctx, schedulerBucketKey(schedulerVersionPrefix, bucket)).Int64()
|
||
require.NoError(t, err)
|
||
require.Equal(t, versionBeforeRejectedWrite, versionAfterRejectedWrite, "fenced writers must not allocate a new version")
|
||
retired, err := cache.rdb.Exists(ctx, schedulerBucketKey(schedulerRetiredPrefix, bucket)).Result()
|
||
require.NoError(t, err)
|
||
require.EqualValues(t, 1, retired, "ordinary writers must never clear the tombstone")
|
||
mr.FastForward(time.Duration(snapshotGraceTTLSeconds+1) * time.Second)
|
||
exists, err := cache.rdb.Exists(ctx, schedulerSnapshotKey(bucket, activeVersion)).Result()
|
||
require.NoError(t, err)
|
||
require.Zero(t, exists, "retired active snapshot must expire after the in-flight grace period")
|
||
|
||
newToken, err := cache.ReopenBucket(ctx, bucket)
|
||
require.NoError(t, err)
|
||
require.True(t, newToken.ValidFor(bucket))
|
||
require.Equal(t, retiredEpoch, newToken.Epoch)
|
||
reopenedAgain, err := cache.ReopenBucket(ctx, bucket)
|
||
require.NoError(t, err)
|
||
require.Equal(t, newToken, reopenedAgain, "reopen must be idempotent within one retirement generation")
|
||
err = cache.SetSnapshot(ctx, bucket, token, []service.Account{account})
|
||
require.ErrorIs(t, err, service.ErrSchedulerBucketWriteFenced)
|
||
require.NoError(t, cache.SetSnapshot(ctx, bucket, newToken, []service.Account{account}))
|
||
reopenedWhileOpen, err := cache.ReopenBucket(ctx, bucket)
|
||
require.NoError(t, err)
|
||
require.Equal(t, newToken, reopenedWhileOpen)
|
||
|
||
snapshot, hit, err := cache.GetSnapshot(ctx, bucket)
|
||
require.NoError(t, err)
|
||
require.True(t, hit)
|
||
require.Len(t, snapshot, 1)
|
||
require.Equal(t, account.ID, snapshot[0].ID)
|
||
}
|
||
|
||
func TestSchedulerCacheActivationIsFencedAfterRetire(t *testing.T) {
|
||
ctx := context.Background()
|
||
cache := newSchedulerCacheUnit(t)
|
||
bucket := service.SchedulerBucket{GroupID: 51, Platform: service.PlatformAnthropic, Mode: service.SchedulerModeMixed}
|
||
account := service.Account{ID: 5101, Platform: service.PlatformAnthropic, Type: service.AccountTypeAPIKey}
|
||
|
||
token, err := cache.CaptureBucketWriteToken(ctx, bucket)
|
||
require.NoError(t, err)
|
||
version, err := cache.allocateSnapshotVersion(ctx, bucket, token)
|
||
require.NoError(t, err)
|
||
_, err = cache.writeSnapshotVersionAndReturnAccountIDs(ctx, bucket, version, []service.Account{account})
|
||
require.NoError(t, err)
|
||
|
||
// Deterministic race C: retirement and authoritative reopen both happen after
|
||
// INCR/write but before the old writer activates.
|
||
require.NoError(t, cache.RetireBucket(ctx, bucket))
|
||
_, err = cache.ReopenBucket(ctx, bucket)
|
||
require.NoError(t, err)
|
||
err = cache.activateSnapshotVersion(ctx, bucket, token, version)
|
||
require.ErrorIs(t, err, service.ErrSchedulerBucketWriteFenced)
|
||
|
||
exists, err := cache.rdb.Exists(ctx, schedulerSnapshotKey(bucket, version)).Result()
|
||
require.NoError(t, err)
|
||
require.Zero(t, exists, "fenced activation must delete its unpublished snapshot")
|
||
exists, err = cache.rdb.Exists(
|
||
ctx,
|
||
schedulerBucketKey(schedulerReadyPrefix, bucket),
|
||
schedulerBucketKey(schedulerActivePrefix, bucket),
|
||
).Result()
|
||
require.NoError(t, err)
|
||
require.Zero(t, exists)
|
||
buckets, err := cache.ListBuckets(ctx)
|
||
require.NoError(t, err)
|
||
require.NotContains(t, buckets, bucket)
|
||
}
|
||
|
||
func TestSchedulerCacheConcurrentReopenReturnsSameToken(t *testing.T) {
|
||
ctx := context.Background()
|
||
cache := newSchedulerCacheUnit(t)
|
||
bucket := service.SchedulerBucket{GroupID: 53, Platform: service.PlatformOpenAI, Mode: service.SchedulerModeForced}
|
||
account := service.Account{ID: 5301, Platform: service.PlatformOpenAI, Type: service.AccountTypeAPIKey}
|
||
|
||
oldToken, err := cache.CaptureBucketWriteToken(ctx, bucket)
|
||
require.NoError(t, err)
|
||
require.NoError(t, cache.RetireBucket(ctx, bucket))
|
||
|
||
type reopenResult struct {
|
||
token service.SchedulerBucketWriteToken
|
||
err error
|
||
}
|
||
start := make(chan struct{})
|
||
results := make(chan reopenResult, 2)
|
||
for range 2 {
|
||
go func() {
|
||
<-start
|
||
token, err := cache.ReopenBucket(ctx, bucket)
|
||
results <- reopenResult{token: token, err: err}
|
||
}()
|
||
}
|
||
close(start)
|
||
first := <-results
|
||
second := <-results
|
||
require.NoError(t, first.err)
|
||
require.NoError(t, second.err)
|
||
require.Equal(t, first.token, second.token)
|
||
require.Greater(t, first.token.Epoch, oldToken.Epoch)
|
||
|
||
require.ErrorIs(t, cache.SetSnapshot(ctx, bucket, oldToken, []service.Account{account}), service.ErrSchedulerBucketWriteFenced)
|
||
require.NoError(t, cache.SetSnapshot(ctx, bucket, first.token, []service.Account{account}))
|
||
}
|
||
|
||
func TestSchedulerCacheReopenExpiresPreviousActiveSnapshot(t *testing.T) {
|
||
ctx := context.Background()
|
||
cache, mr := newSchedulerCacheUnitWithRedis(t)
|
||
bucket := service.SchedulerBucket{GroupID: 52, Platform: service.PlatformGemini, Mode: service.SchedulerModeForced}
|
||
account := service.Account{ID: 5201, Platform: service.PlatformGemini, Type: service.AccountTypeAPIKey}
|
||
|
||
oldToken, err := cache.CaptureBucketWriteToken(ctx, bucket)
|
||
require.NoError(t, err)
|
||
require.NoError(t, cache.SetSnapshot(ctx, bucket, oldToken, []service.Account{account}))
|
||
oldVersion, err := cache.rdb.Get(ctx, schedulerBucketKey(schedulerActivePrefix, bucket)).Result()
|
||
require.NoError(t, err)
|
||
retiredEpoch := oldToken.Epoch + 1
|
||
require.NoError(t, cache.rdb.Set(ctx, schedulerBucketKey(schedulerEpochPrefix, bucket), retiredEpoch, 0).Err())
|
||
require.NoError(t, cache.rdb.Set(ctx, schedulerBucketKey(schedulerRetiredPrefix, bucket), retiredEpoch, 0).Err())
|
||
|
||
newToken, err := cache.ReopenBucket(ctx, bucket)
|
||
require.NoError(t, err)
|
||
require.Equal(t, retiredEpoch, newToken.Epoch)
|
||
_, hit, err := cache.GetSnapshot(ctx, bucket)
|
||
require.NoError(t, err)
|
||
require.False(t, hit)
|
||
ttl, err := cache.rdb.TTL(ctx, schedulerSnapshotKey(bucket, oldVersion)).Result()
|
||
require.NoError(t, err)
|
||
require.Positive(t, ttl)
|
||
require.LessOrEqual(t, ttl, time.Duration(snapshotGraceTTLSeconds)*time.Second)
|
||
|
||
require.ErrorIs(t, cache.SetSnapshot(ctx, bucket, oldToken, []service.Account{account}), service.ErrSchedulerBucketWriteFenced)
|
||
mr.FastForward(time.Duration(snapshotGraceTTLSeconds+1) * time.Second)
|
||
exists, err := cache.rdb.Exists(ctx, schedulerSnapshotKey(bucket, oldVersion)).Result()
|
||
require.NoError(t, err)
|
||
require.Zero(t, exists)
|
||
require.NoError(t, cache.SetSnapshot(ctx, bucket, newToken, []service.Account{account}))
|
||
}
|
||
|
||
func TestSchedulerCacheGroupLifecycleLeaseConcurrentAcquireSingleOwner(t *testing.T) {
|
||
ctx := context.Background()
|
||
cache := newSchedulerCacheUnit(t)
|
||
const groupID int64 = 71
|
||
|
||
type result struct {
|
||
lease service.SchedulerGroupLifecycleLease
|
||
acquired bool
|
||
err error
|
||
}
|
||
start := make(chan struct{})
|
||
results := make(chan result, 32)
|
||
for range 32 {
|
||
go func() {
|
||
<-start
|
||
lease, acquired, err := cache.TryAcquireGroupLifecycleLease(ctx, groupID, time.Minute)
|
||
results <- result{lease: lease, acquired: acquired, err: err}
|
||
}()
|
||
}
|
||
close(start)
|
||
|
||
var owner service.SchedulerGroupLifecycleLease
|
||
acquiredCount := 0
|
||
for range 32 {
|
||
got := <-results
|
||
require.NoError(t, got.err)
|
||
if got.acquired {
|
||
acquiredCount++
|
||
owner = got.lease
|
||
require.True(t, got.lease.ValidFor(groupID))
|
||
} else {
|
||
require.Equal(t, service.SchedulerGroupLifecycleLease{}, got.lease)
|
||
}
|
||
}
|
||
require.Equal(t, 1, acquiredCount)
|
||
require.Len(t, owner.OwnerToken, schedulerGroupLifecycleOwnerTokenBytes*2)
|
||
require.Equal(t, strings.ToLower(owner.OwnerToken), owner.OwnerToken)
|
||
decodedOwner, err := hex.DecodeString(owner.OwnerToken)
|
||
require.NoError(t, err)
|
||
require.Len(t, decodedOwner, schedulerGroupLifecycleOwnerTokenBytes)
|
||
|
||
require.NoError(t, cache.ReleaseGroupLifecycleLease(ctx, owner))
|
||
next, acquired, err := cache.TryAcquireGroupLifecycleLease(ctx, groupID, time.Minute)
|
||
require.NoError(t, err)
|
||
require.True(t, acquired)
|
||
require.True(t, next.ValidFor(groupID))
|
||
require.NotEqual(t, owner.OwnerToken, next.OwnerToken)
|
||
require.NoError(t, cache.ReleaseGroupLifecycleLease(ctx, next))
|
||
}
|
||
|
||
func TestSchedulerCacheGroupLifecycleLeaseStaleReleaseCannotDeleteSuccessor(t *testing.T) {
|
||
ctx := context.Background()
|
||
cache, mr := newSchedulerCacheUnitWithRedis(t)
|
||
const groupID int64 = 72
|
||
const ttl = time.Minute
|
||
|
||
first, acquired, err := cache.TryAcquireGroupLifecycleLease(ctx, groupID, ttl)
|
||
require.NoError(t, err)
|
||
require.True(t, acquired)
|
||
|
||
mr.FastForward(ttl + time.Second)
|
||
second, acquired, err := cache.TryAcquireGroupLifecycleLease(ctx, groupID, ttl)
|
||
require.NoError(t, err)
|
||
require.True(t, acquired)
|
||
require.NotEqual(t, first.OwnerToken, second.OwnerToken)
|
||
|
||
require.ErrorIs(t, cache.ReleaseGroupLifecycleLease(ctx, first), service.ErrSchedulerGroupLifecycleLeaseLost)
|
||
owner, err := cache.rdb.Get(ctx, schedulerGroupLifecycleLockKey(groupID)).Result()
|
||
require.NoError(t, err)
|
||
require.Equal(t, second.OwnerToken, owner)
|
||
|
||
_, acquired, err = cache.TryAcquireGroupLifecycleLease(ctx, groupID, ttl)
|
||
require.NoError(t, err)
|
||
require.False(t, acquired)
|
||
require.NoError(t, cache.ReleaseGroupLifecycleLease(ctx, second))
|
||
require.ErrorIs(t, cache.ReleaseGroupLifecycleLease(ctx, second), service.ErrSchedulerGroupLifecycleLeaseLost)
|
||
}
|
||
|
||
func TestSchedulerCacheGroupLifecycleLeaseExpiredReleaseIsLost(t *testing.T) {
|
||
ctx := context.Background()
|
||
cache, mr := newSchedulerCacheUnitWithRedis(t)
|
||
const groupID int64 = 73
|
||
const ttl = time.Minute
|
||
|
||
lease, acquired, err := cache.TryAcquireGroupLifecycleLease(ctx, groupID, ttl)
|
||
require.NoError(t, err)
|
||
require.True(t, acquired)
|
||
mr.FastForward(ttl + time.Second)
|
||
|
||
require.ErrorIs(t, cache.ReleaseGroupLifecycleLease(ctx, lease), service.ErrSchedulerGroupLifecycleLeaseLost)
|
||
}
|
||
|
||
func TestSchedulerCacheGroupLifecycleLeaseWrongOwnerAndCrossGroupAreLost(t *testing.T) {
|
||
ctx := context.Background()
|
||
cache := newSchedulerCacheUnit(t)
|
||
const firstGroupID int64 = 74
|
||
const secondGroupID int64 = 75
|
||
|
||
first, acquired, err := cache.TryAcquireGroupLifecycleLease(ctx, firstGroupID, time.Minute)
|
||
require.NoError(t, err)
|
||
require.True(t, acquired)
|
||
second, acquired, err := cache.TryAcquireGroupLifecycleLease(ctx, secondGroupID, time.Minute)
|
||
require.NoError(t, err)
|
||
require.True(t, acquired, "different groups must acquire independently")
|
||
require.NotEqual(t, first.OwnerToken, second.OwnerToken)
|
||
|
||
wrongOwner := first
|
||
wrongOwner.OwnerToken = strings.Repeat("0", schedulerGroupLifecycleOwnerTokenBytes*2)
|
||
if wrongOwner.OwnerToken == first.OwnerToken {
|
||
wrongOwner.OwnerToken = strings.Repeat("1", schedulerGroupLifecycleOwnerTokenBytes*2)
|
||
}
|
||
require.ErrorIs(t, cache.ReleaseGroupLifecycleLease(ctx, wrongOwner), service.ErrSchedulerGroupLifecycleLeaseLost)
|
||
|
||
crossGroup := first
|
||
crossGroup.GroupID = secondGroupID
|
||
require.ErrorIs(t, cache.ReleaseGroupLifecycleLease(ctx, crossGroup), service.ErrSchedulerGroupLifecycleLeaseLost)
|
||
|
||
firstOwner, err := cache.rdb.Get(ctx, schedulerGroupLifecycleLockKey(firstGroupID)).Result()
|
||
require.NoError(t, err)
|
||
require.Equal(t, first.OwnerToken, firstOwner)
|
||
secondOwner, err := cache.rdb.Get(ctx, schedulerGroupLifecycleLockKey(secondGroupID)).Result()
|
||
require.NoError(t, err)
|
||
require.Equal(t, second.OwnerToken, secondOwner)
|
||
|
||
require.NoError(t, cache.ReleaseGroupLifecycleLease(ctx, first))
|
||
require.NoError(t, cache.ReleaseGroupLifecycleLease(ctx, second))
|
||
}
|
||
|
||
func TestSchedulerCacheGroupLifecycleLeaseCanceledContextFailsClosed(t *testing.T) {
|
||
cache := newSchedulerCacheUnit(t)
|
||
canceledCtx, cancel := context.WithCancel(context.Background())
|
||
cancel()
|
||
|
||
lease, acquired, err := cache.TryAcquireGroupLifecycleLease(canceledCtx, 76, time.Minute)
|
||
require.ErrorIs(t, err, context.Canceled)
|
||
require.False(t, acquired)
|
||
require.Equal(t, service.SchedulerGroupLifecycleLease{}, lease)
|
||
|
||
ctx := context.Background()
|
||
lease, acquired, err = cache.TryAcquireGroupLifecycleLease(ctx, 76, time.Minute)
|
||
require.NoError(t, err)
|
||
require.True(t, acquired)
|
||
require.ErrorIs(t, cache.ReleaseGroupLifecycleLease(canceledCtx, lease), context.Canceled)
|
||
owner, err := cache.rdb.Get(ctx, schedulerGroupLifecycleLockKey(lease.GroupID)).Result()
|
||
require.NoError(t, err)
|
||
require.Equal(t, lease.OwnerToken, owner)
|
||
require.NoError(t, cache.ReleaseGroupLifecycleLease(ctx, lease))
|
||
}
|
||
|
||
func TestSchedulerCacheGroupLifecycleLeaseRejectsInvalidInput(t *testing.T) {
|
||
ctx := context.Background()
|
||
cache := newSchedulerCacheUnit(t)
|
||
|
||
lease, acquired, err := cache.TryAcquireGroupLifecycleLease(ctx, 0, time.Minute)
|
||
require.ErrorIs(t, err, service.ErrSchedulerGroupLifecycleLeaseInvalid)
|
||
require.False(t, acquired)
|
||
require.Equal(t, service.SchedulerGroupLifecycleLease{}, lease)
|
||
|
||
lease, acquired, err = cache.TryAcquireGroupLifecycleLease(ctx, 73, 0)
|
||
require.ErrorIs(t, err, service.ErrSchedulerGroupLifecycleLeaseInvalid)
|
||
require.False(t, acquired)
|
||
require.Equal(t, service.SchedulerGroupLifecycleLease{}, lease)
|
||
|
||
canceledCtx, cancel := context.WithCancel(ctx)
|
||
cancel()
|
||
require.ErrorIs(t, cache.ReleaseGroupLifecycleLease(canceledCtx, service.SchedulerGroupLifecycleLease{}), service.ErrSchedulerGroupLifecycleLeaseInvalid)
|
||
keys, err := cache.rdb.DBSize(ctx).Result()
|
||
require.NoError(t, err)
|
||
require.Zero(t, keys)
|
||
}
|
||
|
||
var schedulerCachePayloadBenchmarkSink int
|
||
|
||
func BenchmarkSchedulerCacheAccountPayloadReuse(b *testing.B) {
|
||
for _, size := range []int{1, 100, 10_000} {
|
||
accounts := schedulerCacheBenchmarkAccounts(size)
|
||
b.Run(fmt.Sprintf("pair_baseline_%d_accounts", size), func(b *testing.B) {
|
||
b.ReportAllocs()
|
||
for i := 0; i < b.N; i++ {
|
||
first, err := benchmarkSchedulerLegacySnapshotPayload(accounts)
|
||
if err != nil {
|
||
b.Fatal(err)
|
||
}
|
||
second, err := benchmarkSchedulerLegacySnapshotPayload(accounts)
|
||
if err != nil {
|
||
b.Fatal(err)
|
||
}
|
||
schedulerCachePayloadBenchmarkSink = first + second
|
||
}
|
||
})
|
||
b.Run(fmt.Sprintf("pair_reuse_%d_accounts", size), func(b *testing.B) {
|
||
b.ReportAllocs()
|
||
for i := 0; i < b.N; i++ {
|
||
ids, total, err := benchmarkSchedulerReusableSnapshotPayload(accounts)
|
||
if err != nil {
|
||
b.Fatal(err)
|
||
}
|
||
// 第二个桶仍构造成员，只跳过账号 JSON 与全局账号键。
|
||
total += len(schedulerSnapshotMembers(ids))
|
||
schedulerCachePayloadBenchmarkSink = total
|
||
}
|
||
})
|
||
b.Run(fmt.Sprintf("first_baseline_%d_accounts", size), func(b *testing.B) {
|
||
b.ReportAllocs()
|
||
for i := 0; i < b.N; i++ {
|
||
total, err := benchmarkSchedulerLegacySnapshotPayload(accounts)
|
||
if err != nil {
|
||
b.Fatal(err)
|
||
}
|
||
schedulerCachePayloadBenchmarkSink = total
|
||
}
|
||
})
|
||
b.Run(fmt.Sprintf("first_reuse_%d_accounts", size), func(b *testing.B) {
|
||
b.ReportAllocs()
|
||
for i := 0; i < b.N; i++ {
|
||
ids, total, err := benchmarkSchedulerReusableSnapshotPayload(accounts)
|
||
if err != nil {
|
||
b.Fatal(err)
|
||
}
|
||
total += len(ids)
|
||
schedulerCachePayloadBenchmarkSink = total
|
||
}
|
||
})
|
||
}
|
||
}
|
||
|
||
func benchmarkSchedulerLegacySnapshotPayload(accounts []service.Account) (int, error) {
|
||
cacheable := make([]service.Account, 0, len(accounts))
|
||
total := 0
|
||
for _, account := range accounts {
|
||
full, meta, err := marshalSchedulerCacheAccount(account)
|
||
if err != nil {
|
||
continue
|
||
}
|
||
total += len(full) + len(meta)
|
||
cacheable = append(cacheable, account)
|
||
}
|
||
members := make([]redis.Z, 0, len(cacheable))
|
||
for idx, account := range cacheable {
|
||
members = append(members, redis.Z{Score: float64(idx), Member: strconv.FormatInt(account.ID, 10)})
|
||
}
|
||
return total + len(members), nil
|
||
}
|
||
|
||
func benchmarkSchedulerReusableSnapshotPayload(accounts []service.Account) ([]int64, int, error) {
|
||
accountIDs := make([]int64, 0, len(accounts))
|
||
total := 0
|
||
for _, account := range accounts {
|
||
full, meta, err := marshalSchedulerCacheAccount(account)
|
||
if err != nil {
|
||
continue
|
||
}
|
||
total += len(full) + len(meta)
|
||
accountIDs = append(accountIDs, account.ID)
|
||
}
|
||
total += len(schedulerSnapshotMembers(accountIDs))
|
||
return accountIDs, total, nil
|
||
}
|
||
|
||
func schedulerCacheBenchmarkAccounts(size int) []service.Account {
|
||
largeValue := strings.Repeat("x", 4096)
|
||
credentials := map[string]any{
|
||
"api_key": "benchmark-key",
|
||
"model_mapping": map[string]any{"z-model": "z-target", "a-model": "a-target"},
|
||
"large_value": largeValue,
|
||
}
|
||
extra := map[string]any{
|
||
"mixed_scheduling": true,
|
||
"model_rate_limits": map[string]any{
|
||
"z-model": map[string]any{"rate_limit_reset_at": "2026-07-16T00:00:00Z"},
|
||
"a-model": map[string]any{"rate_limit_reset_at": "2026-07-16T00:00:00Z"},
|
||
},
|
||
"large_value": largeValue,
|
||
}
|
||
accounts := make([]service.Account, size)
|
||
for i := range accounts {
|
||
id := int64(i + 1)
|
||
accounts[i] = service.Account{
|
||
ID: id,
|
||
Name: "benchmark-account",
|
||
Platform: service.PlatformOpenAI,
|
||
Type: service.AccountTypeOAuth,
|
||
Credentials: credentials,
|
||
Extra: extra,
|
||
GroupIDs: []int64{7, 9},
|
||
AccountGroups: []service.AccountGroup{
|
||
{AccountID: id, GroupID: 7, Priority: 1},
|
||
{AccountID: id, GroupID: 9, Priority: 2},
|
||
},
|
||
}
|
||
}
|
||
return accounts
|
||
}
|
||
|
||
// 调度投影必须保留 OpenAI 透传开关。
|
||
//
|
||
// 候选过滤走 ListSchedulableAccounts，读的是 buildSchedulerMetadataAccount 产出的精简投影；
|
||
// Account.IsModelSupported 又靠 extra 上的透传开关短路 model_mapping 白名单（#4936）。
|
||
// 一旦投影把开关裁掉、却保留了白名单，透传账号在选号阶段就会退回白名单判定并被误判成
|
||
// model_not_supported，而转发阶段（读完整账号）仍按透传工作 —— 表现为"单独测这个账号能通、
|
||
// 走网关却报 no available accounts"。#4936 修的是判定逻辑，这里守的是喂给判定的输入。
|
||
func TestBuildSchedulerMetadataAccount_KeepsOpenAIPassthroughForModelGate(t *testing.T) {
|
||
for _, key := range []string{"openai_passthrough", "openai_oauth_passthrough"} {
|
||
t.Run(key, func(t *testing.T) {
|
||
account := service.Account{
|
||
ID: 383,
|
||
Platform: service.PlatformOpenAI,
|
||
Type: service.AccountTypeOAuth,
|
||
Credentials: map[string]any{
|
||
// 账号从白名单模式切到透传后常见的残留映射，未列出请求的模型。
|
||
"model_mapping": map[string]any{"gpt-5.5": "gpt-5.5"},
|
||
"access_token": "drop-me",
|
||
},
|
||
Extra: map[string]any{key: true},
|
||
}
|
||
require.True(t, account.IsModelSupported("gpt-5.6-sol"),
|
||
"前置条件：透传账号本应放行白名单外的模型")
|
||
|
||
meta := buildSchedulerMetadataAccount(account)
|
||
|
||
// 走一遍真实的序列化/反序列化路径（写入 sched:meta 再由 decodeCachedAccount 读回）。
|
||
payload, err := json.Marshal(meta)
|
||
require.NoError(t, err)
|
||
var restored service.Account
|
||
require.NoError(t, json.Unmarshal(payload, &restored))
|
||
|
||
require.Equal(t, true, restored.Extra[key])
|
||
require.True(t, restored.IsOpenAIPassthroughEnabled())
|
||
require.True(t, restored.IsModelSupported("gpt-5.6-sol"),
|
||
"投影裁掉透传开关会让透传账号在候选过滤阶段被误判为 model_not_supported")
|
||
// 白名单本身仍需保留：非透传账号依赖它做模型门。
|
||
require.Equal(t, map[string]any{"gpt-5.5": "gpt-5.5"}, restored.Credentials["model_mapping"])
|
||
})
|
||
}
|
||
}
|
||
|
||
// 调度投影必须保留账号的 RPM 配置。
|
||
//
|
||
// 候选过滤走 ListSchedulableAccounts，读的是 buildSchedulerMetadataAccount 产出的精简投影；
|
||
// isAccountSchedulableForRPM 靠 extra 上的 base_rpm 判定，缺失时 GetBaseRPM() 返回 0，
|
||
// baseRPM <= 0 直接放行，已配置限流的账号在选号阶段被当作未开启 RPM（#7225）。
|
||
// rpm_strategy / rpm_sticky_buffer 同理：只保留 base_rpm 会让三区判定退回默认值，
|
||
// 红区（不可调度）被误判成黄区（仅粘性），超额请求照样被转发。
|
||
func TestBuildSchedulerMetadataAccount_KeepsRPMFieldsForRPMGate(t *testing.T) {
|
||
newAccount := func(extra map[string]any) service.Account {
|
||
return service.Account{
|
||
ID: 7225,
|
||
Platform: service.PlatformAnthropic,
|
||
Type: service.AccountTypeOAuth,
|
||
Extra: extra,
|
||
}
|
||
}
|
||
roundTrip := func(t *testing.T, account service.Account) service.Account {
|
||
t.Helper()
|
||
// 走一遍真实的序列化/反序列化路径（写入 sched:meta 再由 decodeCachedAccount 读回）。
|
||
payload, err := json.Marshal(buildSchedulerMetadataAccount(account))
|
||
require.NoError(t, err)
|
||
var restored service.Account
|
||
require.NoError(t, json.Unmarshal(payload, &restored))
|
||
return restored
|
||
}
|
||
|
||
t.Run("tiered", func(t *testing.T) {
|
||
account := newAccount(map[string]any{
|
||
"base_rpm": 10,
|
||
"rpm_strategy": "tiered",
|
||
"rpm_sticky_buffer": 1,
|
||
})
|
||
require.Equal(t, service.WindowCostNotSchedulable, account.CheckRPMSchedulability(11),
|
||
"前置条件：base_rpm=10 + 缓冲 1，当前 11 次已进红区")
|
||
|
||
restored := roundTrip(t, account)
|
||
|
||
require.Equal(t, 10, restored.GetBaseRPM())
|
||
require.Equal(t, "tiered", restored.GetRPMStrategy())
|
||
require.Equal(t, 1, restored.GetRPMStickyBuffer())
|
||
require.Equal(t, service.WindowCostNotSchedulable, restored.CheckRPMSchedulability(11),
|
||
"投影裁掉 RPM 配置会让已配置限流的账号在候选过滤阶段被放行")
|
||
})
|
||
|
||
t.Run("sticky_exempt", func(t *testing.T) {
|
||
account := newAccount(map[string]any{
|
||
"base_rpm": 10,
|
||
"rpm_strategy": "sticky_exempt",
|
||
})
|
||
require.Equal(t, service.WindowCostStickyOnly, account.CheckRPMSchedulability(100),
|
||
"前置条件：粘性豁免策略没有红区")
|
||
|
||
restored := roundTrip(t, account)
|
||
|
||
require.Equal(t, 10, restored.GetBaseRPM())
|
||
require.Equal(t, "sticky_exempt", restored.GetRPMStrategy())
|
||
require.Equal(t, service.WindowCostStickyOnly, restored.CheckRPMSchedulability(100),
|
||
"投影裁掉 rpm_strategy 会让粘性豁免账号退回三区判定")
|
||
})
|
||
}
|