| 12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091929394959697989910010110210310410510610710810911011111211311411511611711811912012112212312412512612712812913013113213313413513613713813914014114214314414514614714814915015115215315415515615715815916016116216316416516616716816917017117217317417517617717817918018118218318418518618718818919019119219319419519619719819920020120220320420520620720820921021121221321421521621721821922022122222322422522622722822923023123223323423523623723823924024124224324424524624724824925025125225325425525625725825926026126226326426526626726826927027127227327427527627727827928028128228328428528628728828929029129229329429529629729829930030130230330430530630730830931031131231331431531631731831932032132232332432532632732832933033133233333433533633733833934034134234334434534634734834935035135235335435535635735835936036136236336436536636736836937037137237337437537637737837938038138238338438538638738838939039139239339439539639739839940040140240340440540640740840941041141241341441541641741841942042142242342442542642742842943043143243343443543643743843944044144244344444544644744844945045145245345445545645745845946046146246346446546646746846947047147247347447547647747847948048148248348448548648748848949049149249349449549649749849950050150250350450550650750850951051151251351451551651751851952052152252352452552652752852953053153253353453553653753853954054154254354454554654754854955055155255355455555655755855956056156256356456556656756856957057157257357457557657757857958058158258358458558658758858959059159259359459559659759859960060160260360460560660760860961061161261361461561661761861962062162262362462562662762862963063163263363463563663763863964064164264364464564664764864965065165265365465565665765865966066166266366466566666766866967067167267367467567667767867968068168268368468568668768868969069169269369469569669769869970070170270370470570670770870971071171271371471571671771871972072172272372472572672772872973073173273373473573673773873974074174274374474574674774874975075175275375475575675775875976076176276376476576676776876977077177277377477577677777877978078178278378478578678778878979079179279379479579679779879980080180280380480580680780880981081181281381481581681781881982082182282382482582682782882983083183283383483583683783883984084184284384484584684784884985085185285385485585685785885986086186286386486586686786886987087187287387487587687787887988088188288388488588688788888989089189289389489589689789889990090190290390490590690790890991091191291391491591691791891992092192292392492592692792892993093193293393493593693793893994094194294394494594694794894995095195295395495595695795895996096196296396496596696796896997097197297397497597697797897998098198298398498598698798898999099199299399499599699799899910001001100210031004100510061007100810091010101110121013101410151016101710181019102010211022102310241025102610271028102910301031103210331034103510361037103810391040104110421043104410451046104710481049105010511052105310541055105610571058105910601061106210631064106510661067106810691070107110721073107410751076107710781079108010811082108310841085108610871088108910901091109210931094 |
- 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",
- // 透传开关必须进投影:候选过滤(ListSchedulableAccounts)读的是本投影,
- // 而 Account.IsModelSupported 靠 extra 上的这两个键短路 model_mapping 白名单。
- // 裁掉它们,透传账号在选号阶段会退回按(常为过期的)白名单判定并被误判为
- // model_not_supported —— 转发阶段却仍按透传工作,表现为"单独测账号能通、
- // 走网关报 no available accounts"。
- "openai_passthrough",
- "openai_oauth_passthrough",
- "codex_fingerprint_mode",
- "request_integrity_mode",
- "account_traffic_control",
- "codex_fingerprint_seed",
- "anti_degrade",
- "anti_degradation",
- "protection_scope",
- "proxy_mode",
- "enable_tls_fingerprint",
- "tls_fingerprint_builtin",
- "tls_fingerprint_profile_id",
- "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
- }
|