Files
sub2api/backend/internal/repository/scheduler_cache.go
T
Yuan LiuandClaude Opus 5.5 1ca7b2a519 fix(scheduler): keep auto reset-credit fields in the scheduler projection
OpenAI candidate filtering reads the Redis metadata projection, but
filterSchedulerExtra dropped auto_reset_credit_enabled, the 5h/7d
reset-credit thresholds and codex_auto_reset_credit_state. The scheduler
therefore always saw automatic reset credits as disabled, so accounts
were paused at the ops quota auto-pause threshold even when a fresh
state showed an available reset credit. The "keep scheduling until the
reset-credit threshold" branch never took effect, and accounts with
credits stayed idle until their window reset naturally.

- Keep the four keys in the scheduler projection; the test round-trips
  the real JSON payload and resolves the credit config from it.
- Once visible, every scheduling pass notifies the auto-reset worker for
  paused credit-enabled accounts, and each notification reloads the
  account. Add a 30s per-account cooldown for the scheduler hot path
  only; other event-driven notifications are unchanged and the
  per-minute scan remains the fallback.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
2026-09-23 20:02:50 +08:00

1107 lines
35 KiB
Go
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
package repository
import (
"context"
"crypto/rand"
"encoding/hex"
"encoding/json"
"fmt"
"log/slog"
"strconv"
"time"
"github.com/Wei-Shaw/sub2api/internal/service"
"github.com/redis/go-redis/v9"
)
const (
schedulerBucketSetKey = "sched:buckets"
schedulerOutboxWatermarkKey = "sched:outbox:watermark"
schedulerAccountPrefix = "sched:acc:"
schedulerAccountMetaPrefix = "sched:meta:"
schedulerAccountLastUsedPrefix = "sched:acc:last_used:"
schedulerActivePrefix = "sched:active:"
schedulerReadyPrefix = "sched:ready:"
schedulerVersionPrefix = "sched:ver:"
schedulerEpochPrefix = "sched:epoch:"
schedulerRetiredPrefix = "sched:retired:"
schedulerSnapshotPrefix = "sched:"
schedulerLockPrefix = "sched:lock:"
defaultSchedulerSnapshotMGetChunkSize = 128
defaultSchedulerSnapshotWriteChunkSize = 256
schedulerLastUsedUpdateChunkSize = 256
// snapshotGraceTTLSeconds 旧快照过期的宽限期(秒)。
// 替代立即 DEL,让正在读取旧版本的 reader 有足够时间完成 ZRANGE。
snapshotGraceTTLSeconds = 60
)
const (
schedulerGroupLifecycleLockPrefix = "sched:group:lifecycle-lock:"
schedulerGroupLifecycleOwnerTokenBytes = 16
)
var updateSchedulerLastUsedScript = redis.NewScript(`
local updated = 0
for index = 1, #ARGV do
local key_index = (index - 1) * 2 + 1
local candidate = tonumber(ARGV[index])
if candidate == nil then
return redis.error_reply('invalid last_used value')
end
if redis.call('EXISTS', KEYS[key_index]) == 1 then
local current = tonumber(redis.call('GET', KEYS[key_index + 1]))
if current == nil or candidate > current then
redis.call('SET', KEYS[key_index + 1], ARGV[index])
updated = updated + 1
end
end
end
return updated
`)
var (
// epoch 标识 bucket writer 的代际,retired key 是持久退休标记。
// Capture、allocate、activate 都在 Lua 内同时校验两者:-1 表示已退休,-2 表示 epoch 无效或与 token 代际不匹配;
// allocate 与 activate 的双重校验可拦截快照写入期间发生的 Retire。
// Retire 仅在首次退休时推进 epoch,Reopen 只清除标记并沿用该代际,因此重复调用保持幂等。
captureBucketWriteTokenScript = redis.NewScript(`
if redis.call('EXISTS', KEYS[2]) == 1 then
return -1
end
local currentEpoch = redis.call('GET', KEYS[1])
if currentEpoch == false then
redis.call('SET', KEYS[1], '1')
return 1
end
local parsedEpoch = tonumber(currentEpoch)
if parsedEpoch == nil or parsedEpoch < 1 then
return -2
end
return parsedEpoch
`)
allocateSnapshotVersionScript = redis.NewScript(`
if redis.call('EXISTS', KEYS[2]) == 1 then
return -1
end
local currentEpoch = tonumber(redis.call('GET', KEYS[1]))
local expectedEpoch = tonumber(ARGV[1])
if currentEpoch == nil or expectedEpoch == nil or currentEpoch ~= expectedEpoch then
return -2
end
return redis.call('INCR', KEYS[3])
`)
retireBucketScript = redis.NewScript(`
local retired = redis.call('GET', KEYS[2])
local currentEpoch = tonumber(redis.call('GET', KEYS[1])) or 0
if retired == false then
currentEpoch = currentEpoch + 1
if currentEpoch < 1 then
currentEpoch = 1
end
redis.call('SET', KEYS[1], tostring(currentEpoch))
redis.call('SET', KEYS[2], tostring(currentEpoch))
elseif currentEpoch < 1 then
currentEpoch = tonumber(retired) or 1
redis.call('SET', KEYS[1], tostring(currentEpoch))
end
redis.call('SREM', KEYS[3], ARGV[1])
local currentActive = redis.call('GET', KEYS[5])
if currentActive ~= false then
redis.call('EXPIRE', ARGV[2] .. currentActive, tonumber(ARGV[3]))
end
redis.call('DEL', KEYS[4], KEYS[5])
return currentEpoch
`)
reopenBucketScript = redis.NewScript(`
local currentEpochRaw = redis.call('GET', KEYS[1])
local currentEpoch = tonumber(currentEpochRaw)
local retiredEpochRaw = redis.call('GET', KEYS[2])
if retiredEpochRaw == false then
if currentEpochRaw == false then
redis.call('SET', KEYS[1], '1')
return 1
end
if currentEpoch == nil or currentEpoch < 1 then
return -2
end
return currentEpoch
end
local retiredEpoch = tonumber(retiredEpochRaw)
if retiredEpoch == nil or retiredEpoch < 1 then
return -2
end
if currentEpoch == nil or currentEpoch < retiredEpoch then
currentEpoch = retiredEpoch
end
redis.call('SET', KEYS[1], tostring(currentEpoch))
redis.call('DEL', KEYS[2])
redis.call('SREM', KEYS[3], ARGV[1])
local currentActive = redis.call('GET', KEYS[5])
if currentActive ~= false then
redis.call('EXPIRE', ARGV[2] .. currentActive, tonumber(ARGV[3]))
end
redis.call('DEL', KEYS[4], KEYS[5])
return currentEpoch
`)
// 释放租约必须先比较所有者令牌再删除,过期持有者的延迟释放不能误删继任租约。
releaseGroupLifecycleLeaseScript = redis.NewScript(`
if redis.call('GET', KEYS[1]) == ARGV[1] then
return redis.call('DEL', KEYS[1])
end
return 0
`)
// activateSnapshotScript 原子 CAS 切换快照版本。
// 仅当新版本号 >= 当前激活版本时才切换,防止并发写入导致版本回滚。
// 旧快照使用 EXPIRE 设置宽限期而非立即 DEL,避免与 reader 竞态。
//
// KEYS[1] = activeKey (sched:active:{bucket})
// KEYS[2] = readyKey (sched:ready:{bucket})
// KEYS[3] = bucketSetKey (sched:buckets)
// KEYS[4] = snapshotKey (新写入的快照 key)
// KEYS[5] = epochKey
// KEYS[6] = retiredKey
// ARGV[1] = 新版本号字符串
// ARGV[2] = bucket 字符串 (用于 SADD)
// ARGV[3] = 快照 key 前缀 (用于构造旧快照 key)
// ARGV[4] = 宽限期 TTL 秒数
// ARGV[5] = writer epoch
//
// 返回 1 = 已激活, 0 = 版本过旧未激活
activateSnapshotScript = redis.NewScript(`
if redis.call('EXISTS', KEYS[6]) == 1 then
redis.call('DEL', KEYS[4])
return -1
end
local currentEpoch = tonumber(redis.call('GET', KEYS[5]))
local expectedEpoch = tonumber(ARGV[5])
if currentEpoch == nil or expectedEpoch == nil or currentEpoch ~= expectedEpoch then
redis.call('DEL', KEYS[4])
return -2
end
local currentActive = redis.call('GET', KEYS[1])
local newVersion = tonumber(ARGV[1])
if currentActive ~= false then
local curVersion = tonumber(currentActive)
if curVersion and newVersion < curVersion then
redis.call('DEL', KEYS[4])
return 0
end
end
redis.call('SET', KEYS[1], ARGV[1])
redis.call('SET', KEYS[2], '1')
redis.call('SADD', KEYS[3], ARGV[2])
if currentActive ~= false and currentActive ~= ARGV[1] then
redis.call('EXPIRE', ARGV[3] .. currentActive, tonumber(ARGV[4]))
end
return 1
`)
)
type schedulerCache struct {
rdb *redis.Client
mgetChunkSize int
writeChunkSize int
}
func NewSchedulerCache(rdb *redis.Client) service.SchedulerCache {
return newSchedulerCacheWithChunkSizes(rdb, defaultSchedulerSnapshotMGetChunkSize, defaultSchedulerSnapshotWriteChunkSize)
}
func newSchedulerCacheWithChunkSizes(rdb *redis.Client, mgetChunkSize, writeChunkSize int) service.SchedulerCache {
if mgetChunkSize <= 0 {
mgetChunkSize = defaultSchedulerSnapshotMGetChunkSize
}
if writeChunkSize <= 0 {
writeChunkSize = defaultSchedulerSnapshotWriteChunkSize
}
return &schedulerCache{
rdb: rdb,
mgetChunkSize: mgetChunkSize,
writeChunkSize: writeChunkSize,
}
}
func (c *schedulerCache) GetSnapshot(ctx context.Context, bucket service.SchedulerBucket) ([]*service.Account, bool, error) {
readyKey := schedulerBucketKey(schedulerReadyPrefix, bucket)
readyVal, err := c.rdb.Get(ctx, readyKey).Result()
if err == redis.Nil {
return nil, false, nil
}
if err != nil {
return nil, false, err
}
if readyVal != "1" {
return nil, false, nil
}
activeKey := schedulerBucketKey(schedulerActivePrefix, bucket)
activeVal, err := c.rdb.Get(ctx, activeKey).Result()
if err == redis.Nil {
return nil, false, nil
}
if err != nil {
return nil, false, err
}
snapshotKey := schedulerSnapshotKey(bucket, activeVal)
ids, err := c.rdb.ZRange(ctx, snapshotKey, 0, -1).Result()
if err != nil {
return nil, false, err
}
if len(ids) == 0 {
// 空快照视为缓存未命中,触发数据库回退查询
// 这解决了新分组创建后立即绑定账号时的竞态条件问题
return nil, false, nil
}
keys := make([]string, 0, len(ids))
lastUsedKeys := make([]string, 0, len(ids))
for _, id := range ids {
keys = append(keys, schedulerAccountMetaKey(id))
lastUsedKeys = append(lastUsedKeys, schedulerLastUsedKey(id))
}
values, err := c.mgetChunked(ctx, keys)
if err != nil {
return nil, false, err
}
lastUsedValues, err := c.mgetChunked(ctx, lastUsedKeys)
if err != nil {
return nil, false, err
}
accounts := make([]*service.Account, 0, len(values))
for i, val := range values {
if val == nil {
return nil, false, nil
}
account, err := decodeCachedAccount(val)
if err != nil {
return nil, false, err
}
if err := applySchedulerLastUsed(account, lastUsedValues[i]); err != nil {
return nil, false, err
}
accounts = append(accounts, account)
}
return accounts, true, nil
}
func (c *schedulerCache) CaptureBucketWriteToken(ctx context.Context, bucket service.SchedulerBucket) (service.SchedulerBucketWriteToken, error) {
result, err := captureBucketWriteTokenScript.Run(ctx, c.rdb, []string{
schedulerBucketKey(schedulerEpochPrefix, bucket),
schedulerBucketKey(schedulerRetiredPrefix, bucket),
}).Int64()
if err != nil {
return service.SchedulerBucketWriteToken{}, err
}
if err := schedulerBucketWriteResultError(result, bucket); err != nil {
return service.SchedulerBucketWriteToken{}, err
}
return service.SchedulerBucketWriteToken{Bucket: bucket, Epoch: result}, nil
}
func (c *schedulerCache) RetireBucket(ctx context.Context, bucket service.SchedulerBucket) error {
snapshotKeyPrefix := fmt.Sprintf("%s%d:%s:%s:v", schedulerSnapshotPrefix, bucket.GroupID, bucket.Platform, bucket.Mode)
result, err := retireBucketScript.Run(ctx, c.rdb, []string{
schedulerBucketKey(schedulerEpochPrefix, bucket),
schedulerBucketKey(schedulerRetiredPrefix, bucket),
schedulerBucketSetKey,
schedulerBucketKey(schedulerReadyPrefix, bucket),
schedulerBucketKey(schedulerActivePrefix, bucket),
}, bucket.String(), snapshotKeyPrefix, snapshotGraceTTLSeconds).Int64()
if err != nil {
return err
}
if result < 1 {
return fmt.Errorf("retire scheduler bucket %s returned invalid epoch %d", bucket.String(), result)
}
return nil
}
func (c *schedulerCache) ReopenBucket(ctx context.Context, bucket service.SchedulerBucket) (service.SchedulerBucketWriteToken, error) {
snapshotKeyPrefix := fmt.Sprintf("%s%d:%s:%s:v", schedulerSnapshotPrefix, bucket.GroupID, bucket.Platform, bucket.Mode)
result, err := reopenBucketScript.Run(ctx, c.rdb, []string{
schedulerBucketKey(schedulerEpochPrefix, bucket),
schedulerBucketKey(schedulerRetiredPrefix, bucket),
schedulerBucketSetKey,
schedulerBucketKey(schedulerReadyPrefix, bucket),
schedulerBucketKey(schedulerActivePrefix, bucket),
}, bucket.String(), snapshotKeyPrefix, snapshotGraceTTLSeconds).Int64()
if err != nil {
return service.SchedulerBucketWriteToken{}, err
}
if err := schedulerBucketWriteResultError(result, bucket); err != nil {
return service.SchedulerBucketWriteToken{}, err
}
return service.SchedulerBucketWriteToken{Bucket: bucket, Epoch: result}, nil
}
func (c *schedulerCache) TryAcquireGroupLifecycleLease(ctx context.Context, groupID int64, ttl time.Duration) (service.SchedulerGroupLifecycleLease, bool, error) {
if groupID <= 0 {
return service.SchedulerGroupLifecycleLease{}, false, fmt.Errorf("%w: group id must be positive", service.ErrSchedulerGroupLifecycleLeaseInvalid)
}
if ttl <= 0 {
return service.SchedulerGroupLifecycleLease{}, false, fmt.Errorf("%w: ttl must be positive", service.ErrSchedulerGroupLifecycleLeaseInvalid)
}
ownerToken, err := newSchedulerGroupLifecycleOwnerToken()
if err != nil {
return service.SchedulerGroupLifecycleLease{}, false, err
}
acquired, err := c.rdb.SetNX(ctx, schedulerGroupLifecycleLockKey(groupID), ownerToken, ttl).Result()
if err != nil {
return service.SchedulerGroupLifecycleLease{}, false, err
}
if !acquired {
return service.SchedulerGroupLifecycleLease{}, false, nil
}
return service.SchedulerGroupLifecycleLease{GroupID: groupID, OwnerToken: ownerToken}, true, nil
}
func (c *schedulerCache) ReleaseGroupLifecycleLease(ctx context.Context, lease service.SchedulerGroupLifecycleLease) error {
if !lease.ValidFor(lease.GroupID) {
return service.ErrSchedulerGroupLifecycleLeaseInvalid
}
result, err := releaseGroupLifecycleLeaseScript.Run(
ctx,
c.rdb,
[]string{schedulerGroupLifecycleLockKey(lease.GroupID)},
lease.OwnerToken,
).Int64()
if err != nil {
return err
}
if result == 0 {
return fmt.Errorf("%w: group=%d", service.ErrSchedulerGroupLifecycleLeaseLost, lease.GroupID)
}
if result != 1 {
return fmt.Errorf("release scheduler group lifecycle lease returned %d", result)
}
return nil
}
func newSchedulerGroupLifecycleOwnerToken() (string, error) {
raw := make([]byte, schedulerGroupLifecycleOwnerTokenBytes)
if _, err := rand.Read(raw); err != nil {
return "", fmt.Errorf("generate scheduler group lifecycle owner token: %w", err)
}
return hex.EncodeToString(raw), nil
}
func (c *schedulerCache) SetSnapshot(ctx context.Context, bucket service.SchedulerBucket, token service.SchedulerBucketWriteToken, accounts []service.Account) error {
if !token.ValidFor(bucket) {
return fmt.Errorf("%w: bucket=%s", service.ErrSchedulerBucketWriteFenced, bucket.String())
}
// 分配版本与激活指针是两个 fencing 边界;中间写入的数据只有通过第二次校验才能发布。
version, err := c.allocateSnapshotVersion(ctx, bucket, token)
if err != nil {
return err
}
// 快照成员最终只依赖可编码账号的有序 ID;直接复用 ID 路径,避免为
// 随后立即丢弃的完整 Account 再分配一份临时切片。
if _, err := c.writeSnapshotVersionAndReturnAccountIDs(ctx, bucket, version, accounts); err != nil {
return err
}
return c.activateSnapshotVersion(ctx, bucket, token, version)
}
// SetSnapshotAndReturnAccountIDs 完整发布快照,并返回实际成功编码并写入的有序账号 ID。
// 该可选能力只供同一重建批次复用,返回前仍会完成版本激活与 fencing 校验。
func (c *schedulerCache) SetSnapshotAndReturnAccountIDs(ctx context.Context, bucket service.SchedulerBucket, token service.SchedulerBucketWriteToken, accounts []service.Account) ([]int64, error) {
if !token.ValidFor(bucket) {
return nil, fmt.Errorf("%w: bucket=%s", service.ErrSchedulerBucketWriteFenced, bucket.String())
}
// 分配版本与激活指针是两个 fencing 边界;中间写入的数据只有通过第二次校验才能发布。
version, err := c.allocateSnapshotVersion(ctx, bucket, token)
if err != nil {
return nil, err
}
accountIDs, err := c.writeSnapshotVersionAndReturnAccountIDs(ctx, bucket, version, accounts)
if err != nil {
return nil, err
}
if err := c.activateSnapshotVersion(ctx, bucket, token, version); err != nil {
return nil, err
}
return accountIDs, nil
}
// SetSnapshotByAccountIDs 复用同批次首次完整写入后得到的账号成员。
// 每个桶仍独立分配版本、写入有序集合并执行激活 fencing,只省略重复的账号 JSON 与全局键写入。
func (c *schedulerCache) SetSnapshotByAccountIDs(ctx context.Context, bucket service.SchedulerBucket, token service.SchedulerBucketWriteToken, accountIDs []int64) error {
if !token.ValidFor(bucket) {
return fmt.Errorf("%w: bucket=%s", service.ErrSchedulerBucketWriteFenced, bucket.String())
}
version, err := c.allocateSnapshotVersion(ctx, bucket, token)
if err != nil {
return err
}
if err := c.writeSnapshotAccountIDs(ctx, bucket, version, accountIDs); err != nil {
return err
}
return c.activateSnapshotVersion(ctx, bucket, token, version)
}
func (c *schedulerCache) allocateSnapshotVersion(ctx context.Context, bucket service.SchedulerBucket, token service.SchedulerBucketWriteToken) (string, error) {
result, err := allocateSnapshotVersionScript.Run(ctx, c.rdb, []string{
schedulerBucketKey(schedulerEpochPrefix, bucket),
schedulerBucketKey(schedulerRetiredPrefix, bucket),
schedulerBucketKey(schedulerVersionPrefix, bucket),
}, token.Epoch).Int64()
if err != nil {
return "", err
}
if err := schedulerBucketWriteResultError(result, bucket); err != nil {
return "", err
}
return strconv.FormatInt(result, 10), nil
}
func (c *schedulerCache) writeSnapshotVersionAndReturnAccountIDs(ctx context.Context, bucket service.SchedulerBucket, version string, accounts []service.Account) ([]int64, error) {
accountIDs, err := c.writeAccountIDs(ctx, accounts)
if err != nil {
return nil, err
}
if err := c.writeSnapshotAccountIDs(ctx, bucket, version, accountIDs); err != nil {
return nil, err
}
return accountIDs, nil
}
func (c *schedulerCache) writeSnapshotAccountIDs(ctx context.Context, bucket service.SchedulerBucket, version string, accountIDs []int64) error {
members := schedulerSnapshotMembers(accountIDs)
return c.writeSnapshotMembers(ctx, bucket, version, members)
}
func schedulerSnapshotMembers(accountIDs []int64) []redis.Z {
if len(accountIDs) == 0 {
return nil
}
// 使用序号作为 score,保持数据库返回的排序语义;重复 ID 继续交由 Redis ZADD
// 按最后一个 score 覆盖,与直接从账号切片构造成员时的行为一致。
members := make([]redis.Z, 0, len(accountIDs))
for idx, accountID := range accountIDs {
members = append(members, redis.Z{
Score: float64(idx),
Member: strconv.FormatInt(accountID, 10),
})
}
return members
}
func (c *schedulerCache) writeSnapshotMembers(ctx context.Context, bucket service.SchedulerBucket, version string, members []redis.Z) error {
if len(members) == 0 {
return nil
}
snapshotKey := schedulerSnapshotKey(bucket, version)
pipe := c.rdb.Pipeline()
for start := 0; start < len(members); start += c.writeChunkSize {
end := start + c.writeChunkSize
if end > len(members) {
end = len(members)
}
pipe.ZAdd(ctx, snapshotKey, members[start:end]...)
}
_, err := pipe.Exec(ctx)
return err
}
func (c *schedulerCache) activateSnapshotVersion(ctx context.Context, bucket service.SchedulerBucket, token service.SchedulerBucketWriteToken, version string) error {
snapshotKey := schedulerSnapshotKey(bucket, version)
// Phase 2: 原子 CAS 切换版本,同时再次校验退休状态与 writer epoch。
// Lua 脚本保证:仅当新版本 >= 当前激活版本时才切换 active 指针,
// 防止并发写入导致版本回滚。
// 旧快照使用 EXPIRE 宽限期而非立即 DEL,避免 reader 竞态。
activeKey := schedulerBucketKey(schedulerActivePrefix, bucket)
readyKey := schedulerBucketKey(schedulerReadyPrefix, bucket)
snapshotKeyPrefix := fmt.Sprintf("%s%d:%s:%s:v", schedulerSnapshotPrefix, bucket.GroupID, bucket.Platform, bucket.Mode)
keys := []string{
activeKey,
readyKey,
schedulerBucketSetKey,
snapshotKey,
schedulerBucketKey(schedulerEpochPrefix, bucket),
schedulerBucketKey(schedulerRetiredPrefix, bucket),
}
args := []any{version, bucket.String(), snapshotKeyPrefix, snapshotGraceTTLSeconds, token.Epoch}
result, err := activateSnapshotScript.Run(ctx, c.rdb, keys, args...).Int64()
if err != nil {
return err
}
return schedulerBucketWriteResultError(result, bucket)
}
func schedulerBucketWriteResultError(result int64, bucket service.SchedulerBucket) error {
switch result {
case -1:
return fmt.Errorf("%w: bucket=%s", service.ErrSchedulerBucketRetired, bucket.String())
case -2:
return fmt.Errorf("%w: bucket=%s", service.ErrSchedulerBucketWriteFenced, bucket.String())
default:
return nil
}
}
func (c *schedulerCache) GetAccount(ctx context.Context, accountID int64) (*service.Account, error) {
id := strconv.FormatInt(accountID, 10)
values, err := c.rdb.MGet(ctx, schedulerAccountKey(id), schedulerLastUsedKey(id)).Result()
if err != nil {
return nil, err
}
if len(values) != 2 || values[0] == nil {
return nil, nil
}
account, err := decodeCachedAccount(values[0])
if err != nil {
return nil, err
}
if err := applySchedulerLastUsed(account, values[1]); err != nil {
return nil, err
}
return account, nil
}
func (c *schedulerCache) SetAccount(ctx context.Context, account *service.Account) error {
if account == nil || account.ID <= 0 {
return nil
}
accountIDs, err := c.writeAccountIDs(ctx, []service.Account{*account})
if err != nil {
return err
}
if len(accountIDs) == 0 {
return c.DeleteAccount(ctx, account.ID)
}
return nil
}
func (c *schedulerCache) DeleteAccount(ctx context.Context, accountID int64) error {
if accountID <= 0 {
return nil
}
id := strconv.FormatInt(accountID, 10)
return c.rdb.Del(ctx, schedulerAccountKey(id), schedulerAccountMetaKey(id), schedulerLastUsedKey(id)).Err()
}
func (c *schedulerCache) UpdateLastUsed(ctx context.Context, updates map[int64]time.Time) error {
if len(updates) == 0 {
return nil
}
pipe := c.rdb.Pipeline()
queued := 0
keys := make([]string, 0, schedulerLastUsedUpdateChunkSize*2)
args := make([]any, 0, schedulerLastUsedUpdateChunkSize)
queueBatch := func() {
if len(args) == 0 {
return
}
updateSchedulerLastUsedScript.Eval(ctx, pipe, keys, args...)
queued++
keys = make([]string, 0, schedulerLastUsedUpdateChunkSize*2)
args = make([]any, 0, schedulerLastUsedUpdateChunkSize)
}
for id, usedAt := range updates {
if id <= 0 {
continue
}
millis, err := schedulerLastUsedMillis(usedAt)
if err != nil {
slog.Warn("scheduler cache removes account with unencodable payload",
"account_id", id,
"error", err,
)
idText := strconv.FormatInt(id, 10)
pipe.Del(ctx, schedulerAccountKey(idText), schedulerAccountMetaKey(idText), schedulerLastUsedKey(idText))
queued++
continue
}
idText := strconv.FormatInt(id, 10)
keys = append(keys, schedulerAccountKey(idText), schedulerLastUsedKey(idText))
args = append(args, millis)
if len(args) >= schedulerLastUsedUpdateChunkSize {
queueBatch()
}
}
queueBatch()
if queued == 0 {
return nil
}
_, err := pipe.Exec(ctx)
return err
}
func (c *schedulerCache) TryLockBucket(ctx context.Context, bucket service.SchedulerBucket, ttl time.Duration) (bool, error) {
key := schedulerBucketKey(schedulerLockPrefix, bucket)
return c.rdb.SetNX(ctx, key, time.Now().UnixNano(), ttl).Result()
}
func (c *schedulerCache) UnlockBucket(ctx context.Context, bucket service.SchedulerBucket) error {
key := schedulerBucketKey(schedulerLockPrefix, bucket)
return c.rdb.Del(ctx, key).Err()
}
func (c *schedulerCache) ListBuckets(ctx context.Context) ([]service.SchedulerBucket, error) {
raw, err := c.rdb.SMembers(ctx, schedulerBucketSetKey).Result()
if err != nil {
return nil, err
}
out := make([]service.SchedulerBucket, 0, len(raw))
for _, entry := range raw {
bucket, ok := service.ParseSchedulerBucket(entry)
if !ok {
continue
}
out = append(out, bucket)
}
return out, nil
}
func (c *schedulerCache) GetOutboxWatermark(ctx context.Context) (int64, error) {
val, err := c.rdb.Get(ctx, schedulerOutboxWatermarkKey).Result()
if err == redis.Nil {
return 0, nil
}
if err != nil {
return 0, err
}
id, err := strconv.ParseInt(val, 10, 64)
if err != nil {
return 0, err
}
return id, nil
}
func (c *schedulerCache) SetOutboxWatermark(ctx context.Context, id int64) error {
return c.rdb.Set(ctx, schedulerOutboxWatermarkKey, strconv.FormatInt(id, 10), 0).Err()
}
func schedulerBucketKey(prefix string, bucket service.SchedulerBucket) string {
return fmt.Sprintf("%s%d:%s:%s", prefix, bucket.GroupID, bucket.Platform, bucket.Mode)
}
func schedulerGroupLifecycleLockKey(groupID int64) string {
return schedulerGroupLifecycleLockPrefix + strconv.FormatInt(groupID, 10)
}
func schedulerSnapshotKey(bucket service.SchedulerBucket, version string) string {
return fmt.Sprintf("%s%d:%s:%s:v%s", schedulerSnapshotPrefix, bucket.GroupID, bucket.Platform, bucket.Mode, version)
}
func schedulerAccountKey(id string) string {
return schedulerAccountPrefix + id
}
func schedulerAccountMetaKey(id string) string {
return schedulerAccountMetaPrefix + id
}
func schedulerLastUsedKey(id string) string {
return schedulerAccountLastUsedPrefix + id
}
func ptrTime(t time.Time) *time.Time {
return &t
}
func schedulerLastUsedMillis(value time.Time) (int64, error) {
if _, err := value.MarshalJSON(); err != nil {
return 0, err
}
return value.UTC().UnixMilli(), nil
}
func applySchedulerLastUsed(account *service.Account, value any) error {
if account == nil || value == nil {
return nil
}
var raw string
switch typed := value.(type) {
case string:
raw = typed
case []byte:
raw = string(typed)
default:
return fmt.Errorf("unexpected last_used cache type: %T", value)
}
millis, err := strconv.ParseInt(raw, 10, 64)
if err != nil {
return fmt.Errorf("invalid last_used cache value %q: %w", raw, err)
}
lastUsedAt := time.UnixMilli(millis).UTC()
if account.LastUsedAt == nil || lastUsedAt.After(*account.LastUsedAt) {
account.LastUsedAt = ptrTime(lastUsedAt)
}
return nil
}
func decodeCachedAccount(val any) (*service.Account, error) {
var payload []byte
switch raw := val.(type) {
case string:
payload = []byte(raw)
case []byte:
payload = raw
default:
return nil, fmt.Errorf("unexpected account cache type: %T", val)
}
var account service.Account
if err := json.Unmarshal(payload, &account); err != nil {
return nil, err
}
return &account, nil
}
func (c *schedulerCache) writeAccountIDs(ctx context.Context, accounts []service.Account) ([]int64, error) {
if len(accounts) == 0 {
return nil, nil
}
pipe := c.rdb.Pipeline()
accountIDs := make([]int64, 0, len(accounts))
pending := 0
flush := func() error {
if pending == 0 {
return nil
}
if _, err := pipe.Exec(ctx); err != nil {
return err
}
pipe = c.rdb.Pipeline()
pending = 0
return nil
}
for _, account := range accounts {
fullPayload, metaPayload, err := marshalSchedulerCacheAccount(account)
if err != nil {
slog.Warn("scheduler cache skips account with unencodable payload",
"account_id", account.ID,
"error", err,
)
continue
}
id := strconv.FormatInt(account.ID, 10)
pipe.Set(ctx, schedulerAccountKey(id), fullPayload, 0)
pipe.Set(ctx, schedulerAccountMetaKey(id), metaPayload, 0)
// Keep the hot LastUsedAt side key untouched: a lagging snapshot rebuild
// must not overwrite a newer scheduler update.
accountIDs = append(accountIDs, account.ID)
pending++
if pending >= c.writeChunkSize {
if err := flush(); err != nil {
return nil, err
}
}
}
if err := flush(); err != nil {
return nil, err
}
return accountIDs, nil
}
func marshalSchedulerCacheAccount(account service.Account) ([]byte, []byte, error) {
fullPayload, err := json.Marshal(account)
if err != nil {
return nil, nil, fmt.Errorf("marshal account: %w", err)
}
metaPayload, err := json.Marshal(buildSchedulerMetadataAccount(account))
if err != nil {
return nil, nil, fmt.Errorf("marshal account metadata: %w", err)
}
return fullPayload, metaPayload, nil
}
func (c *schedulerCache) mgetChunked(ctx context.Context, keys []string) ([]any, error) {
if len(keys) == 0 {
return []any{}, nil
}
out := make([]any, 0, len(keys))
chunkSize := c.mgetChunkSize
if chunkSize <= 0 {
chunkSize = defaultSchedulerSnapshotMGetChunkSize
}
for start := 0; start < len(keys); start += chunkSize {
end := start + chunkSize
if end > len(keys) {
end = len(keys)
}
part, err := c.rdb.MGet(ctx, keys[start:end]...).Result()
if err != nil {
return nil, err
}
out = append(out, part...)
}
return out, nil
}
func buildSchedulerMetadataAccount(account service.Account) service.Account {
return service.Account{
ID: account.ID,
Name: account.Name,
Platform: account.Platform,
Type: account.Type,
Concurrency: account.Concurrency,
LoadFactor: account.LoadFactor,
Priority: account.Priority,
RateMultiplier: account.RateMultiplier,
Status: account.Status,
LastUsedAt: account.LastUsedAt,
ExpiresAt: account.ExpiresAt,
AutoPauseOnExpired: account.AutoPauseOnExpired,
Schedulable: account.Schedulable,
RateLimitedAt: account.RateLimitedAt,
RateLimitResetAt: account.RateLimitResetAt,
OverloadUntil: account.OverloadUntil,
TempUnschedulableUntil: account.TempUnschedulableUntil,
TempUnschedulableReason: account.TempUnschedulableReason,
SessionWindowStart: account.SessionWindowStart,
SessionWindowEnd: account.SessionWindowEnd,
SessionWindowStatus: account.SessionWindowStatus,
ParentAccountID: account.ParentAccountID,
QuotaDimension: account.QuotaDimension,
AccountGroups: filterSchedulerAccountGroups(account.AccountGroups),
GroupIDs: filterSchedulerGroupIDs(account.GroupIDs, account.AccountGroups),
Credentials: filterSchedulerCredentials(account.Credentials),
Extra: filterSchedulerExtra(account.Extra),
}
}
func filterSchedulerAccountGroups(accountGroups []service.AccountGroup) []service.AccountGroup {
if len(accountGroups) == 0 {
return nil
}
filtered := make([]service.AccountGroup, 0, len(accountGroups))
for _, ag := range accountGroups {
if ag.GroupID <= 0 {
continue
}
filtered = append(filtered, service.AccountGroup{
AccountID: ag.AccountID,
GroupID: ag.GroupID,
Priority: ag.Priority,
CreatedAt: ag.CreatedAt,
})
}
if len(filtered) == 0 {
return nil
}
return filtered
}
func filterSchedulerGroupIDs(groupIDs []int64, accountGroups []service.AccountGroup) []int64 {
if len(groupIDs) == 0 && len(accountGroups) == 0 {
return nil
}
seen := make(map[int64]struct{}, len(groupIDs)+len(accountGroups))
filtered := make([]int64, 0, len(groupIDs)+len(accountGroups))
for _, id := range groupIDs {
if id <= 0 {
continue
}
if _, ok := seen[id]; ok {
continue
}
seen[id] = struct{}{}
filtered = append(filtered, id)
}
for _, ag := range accountGroups {
if ag.GroupID <= 0 {
continue
}
if _, ok := seen[ag.GroupID]; ok {
continue
}
seen[ag.GroupID] = struct{}{}
filtered = append(filtered, ag.GroupID)
}
if len(filtered) == 0 {
return nil
}
return filtered
}
func filterSchedulerCredentials(credentials map[string]any) map[string]any {
if len(credentials) == 0 {
return nil
}
// Candidate-list admission evaluates the account override before hydrating
// the full account. Dropping it silently falls back to the platform threshold.
keys := []string{"model_mapping", "compact_model_mapping", "api_key", "project_id", "oauth_type", "plan_type", "account_scheduling_threshold"}
filtered := make(map[string]any)
for _, key := range keys {
if value, ok := credentials[key]; ok && value != nil {
filtered[key] = value
}
}
if len(filtered) == 0 {
return nil
}
return filtered
}
func filterSchedulerExtra(extra map[string]any) map[string]any {
if len(extra) == 0 {
return nil
}
keys := []string{
// Anthropic shared-window and Fable-only threshold checks run on this
// projection. UpdateExtra refreshes both payloads without a bucket rebuild.
"session_window_utilization",
"passive_usage_7d_utilization",
"passive_usage_7d_reset",
"passive_usage_7d_oi_utilization",
"passive_usage_7d_oi_reset",
"quota_limit",
"quota_used",
"quota_daily_limit",
"quota_daily_used",
"quota_daily_start",
"quota_daily_reset_mode",
"quota_daily_reset_hour",
"quota_weekly_limit",
"quota_weekly_used",
"quota_weekly_start",
"quota_weekly_reset_mode",
"quota_weekly_reset_day",
"quota_weekly_reset_hour",
"quota_reset_timezone",
"mixed_scheduling",
"window_cost_limit",
"window_cost_sticky_reserve",
// RPM 门与窗口费用门一样跑在本投影上:isAccountSchedulableForRPM 读 base_rpm,
// 缺失时 GetBaseRPM() 返回 0 并直接放行,已配置限流的账号会被超额调度。
"base_rpm",
"rpm_strategy",
"rpm_sticky_buffer",
"max_sessions",
"session_idle_timeout_minutes",
"openai_oauth_responses_websockets_v2_enabled",
"openai_oauth_responses_websockets_v2_mode",
"openai_apikey_responses_websockets_v2_enabled",
"openai_apikey_responses_websockets_v2_mode",
"responses_websockets_v2_enabled",
"openai_ws_enabled",
"openai_ws_force_http",
"openai_responses_mode",
"openai_responses_supported",
// 透传开关必须进投影:候选过滤(ListSchedulableAccounts)读的是本投影,
// 而 Account.IsModelSupported 靠 extra 上的这两个键短路 model_mapping 白名单。
// 裁掉它们,透传账号在选号阶段会退回按(常为过期的)白名单判定并被误判为
// model_not_supported —— 转发阶段却仍按透传工作,表现为"单独测账号能通、
// 走网关报 no available accounts"。
"openai_passthrough",
"openai_oauth_passthrough",
"codex_fingerprint_mode",
"codex_fingerprint_seed",
"codex_5h_used_percent",
"codex_7d_used_percent",
"codex_5h_reset_at",
"codex_7d_reset_at",
"codex_5h_reset_after_seconds",
"codex_7d_reset_after_seconds",
"codex_usage_updated_at",
"auto_pause_5h_threshold",
"auto_pause_7d_threshold",
"auto_pause_5h_disabled",
"auto_pause_7d_disabled",
// 自动用卡:卡可用的 OpenAI 号在暂停阈值与用卡阈值之间继续调度。
// 候选过滤读的是本投影,缺这几个键时放行分支永远不会生效,
// 账号会在暂停阈值处被一刀切停调,直到窗口自然重置。
service.OpenAIAutoResetCreditEnabledExtraKey,
service.OpenAIAutoResetCredit5hThresholdExtraKey,
service.OpenAIAutoResetCredit7dThresholdExtraKey,
service.OpenAIAutoResetCreditStateExtraKey,
"model_rate_limits",
service.UpstreamBillingProbeExtraKey,
service.GrokMediaEligibleExtraKey,
"grok_billing_snapshot",
}
filtered := make(map[string]any)
for _, key := range keys {
if value, ok := extra[key]; ok && value != nil {
if key == service.UpstreamBillingProbeExtraKey {
filteredProbe := filterSchedulerUpstreamBillingProbe(value)
if filteredProbe == nil {
continue
}
value = filteredProbe
}
filtered[key] = value
}
}
if len(filtered) == 0 {
return nil
}
return filtered
}
func filterSchedulerUpstreamBillingProbe(value any) map[string]any {
source, ok := value.(map[string]any)
if !ok {
return nil
}
status, ok := source["status"].(string)
if !ok || status == "" {
return nil
}
filtered := map[string]any{"status": status}
for _, key := range []string{"received_at", "fresh_until", "next_probe_at"} {
if field, exists := source[key]; exists && field != nil {
filtered[key] = field
}
}
data, ok := source["data"].(map[string]any)
if !ok {
return filtered
}
filteredData := make(map[string]any)
for _, key := range []string{
"billing_scope",
"resolved_rate_multiplier",
"peak_rate_enabled",
"peak_start",
"peak_end",
"peak_rate_multiplier",
"timezone",
} {
if field, exists := data[key]; exists && field != nil {
filteredData[key] = field
}
}
if len(filteredData) > 0 {
filtered["data"] = filteredData
}
return filtered
}