李建琦
6d655c9903
Release / update-version (push) Has been cancelled
Release / build-frontend (push) Has been cancelled
Release / release (push) Has been cancelled
Release / sync-version-file (push) Has been cancelled
CI / shell (push) Canceled after 0s
CI / test (push) Canceled after 0s
CI / frontend (push) Canceled after 0s
CI / golangci-lint (push) Canceled after 0s
Security Scan / backend-security (push) Canceled after 0s
Security Scan / frontend-security (push) Canceled after 0s
1079 lines
33 KiB
Go
1079 lines
33 KiB
Go
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
|
|
}
|
|
keys := []string{"model_mapping", "compact_model_mapping", "api_key", "project_id", "oauth_type", "plan_type"}
|
|
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{
|
|
"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",
|
|
"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",
|
|
"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",
|
|
"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
|
|
}
|