Files
sub2api/backend/internal/repository/scheduler_cache.go
T
feeeei 05d7946af8 fix(gemini): 显式缓存错误透传收窄到上游网关账号,延长沿用创建时计费口径
- 绑定缓存的请求与缓存管理接口收到 429 / 5xx 时,只有上游网关账号的响应原样回写
  (UpstreamFailoverError.BoundUpstreamPassthrough / GeminiCachedContentUpstreamError.UpstreamGateway),
  官方账号仍走错误透传规则与通用映射,不向客户端暴露配额明细
- 记录新增 channel_usage(迁移 244)保存创建时的渠道计费口径,PATCH 延长按同一口径预检、报价与记账
- 创建 / 延长费用行的 request_id 改为固定前缀 gcache:{create|patch}:{id},不被客户端或网关请求 id 覆盖,
  按缓存 ID 即可检索一个缓存的全部费用行
- 调度快照凭据投影补齐 base_url / auth_mode / 上游显式缓存声明,显式缓存候选过滤一次算全,
  不再逐个选号后剔除
2026-10-11 01:37:04 +08:00

1111 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",
// Gemini 显式缓存的候选过滤(GeminiExplicitCacheExclusions)读本投影判定账号是否支持显式缓存:
// 官方地址、Vertex 认证方式与上游网关的显式缓存声明。缺这几个键时只能按类型粗筛,
// 不支持的账号要逐个选中后才被剔除。
"base_url", "auth_mode", service.GeminiExplicitCacheUpstreamCredentialKey, service.GeminiExplicitCacheUpstreamMaxTTLCredentialKey}
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
}