| 12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091929394959697989910010110210310410510610710810911011111211311411511611711811912012112212312412512612712812913013113213313413513613713813914014114214314414514614714814915015115215315415515615715815916016116216316416516616716816917017117217317417517617717817918018118218318418518618718818919019119219319419519619719819920020120220320420520620720820921021121221321421521621721821922022122222322422522622722822923023123223323423523623723823924024124224324424524624724824925025125225325425525625725825926026126226326426526626726826927027127227327427527627727827928028128228328428528628728828929029129229329429529629729829930030130230330430530630730830931031131231331431531631731831932032132232332432532632732832933033133233333433533633733833934034134234334434534634734834935035135235335435535635735835936036136236336436536636736836937037137237337437537637737837938038138238338438538638738838939039139239339439539639739839940040140240340440540640740840941041141241341441541641741841942042142242342442542642742842943043143243343443543643743843944044144244344444544644744844945045145245345445545645745845946046146246346446546646746846947047147247347447547647747847948048148248348448548648748848949049149249349449549649749849950050150250350450550650750850951051151251351451551651751851952052152252352452552652752852953053153253353453553653753853954054154254354454554654754854955055155255355455555655755855956056156256356456556656756856957057157257357457557657757857958058158258358458558658758858959059159259359459559659759859960060160260360460560660760860961061161261361461561661761861962062162262362462562662762862963063163263363463563663763863964064164264364464564664764864965065165265365465565665765865966066166266366466566666766866967067167267367467567667767867968068168268368468568668768868969069169269369469569669769869970070170270370470570670770870971071171271371471571671771871972072172272372472572672772872973073173273373473573673773873974074174274374474574674774874975075175275375475575675775875976076176276376476576676776876977077177277377477577677777877978078178278378478578678778878979079179279379479579679779879980080180280380480580680780880981081181281381481581681781881982082182282382482582682782882983083183283383483583683783883984084184284384484584684784884985085185285385485585685785885986086186286386486586686786886987087187287387487587687787887988088188288388488588688788888989089189289389489589689789889990090190290390490590690790890991091191291391491591691791891992092192292392492592692792892993093193293393493593693793893994094194294394494594694794894995095195295395495595695795895996096196296396496596696796896997097197297397497597697797897998098198298398498598698798898999099199299399499599699799899910001001100210031004100510061007100810091010101110121013101410151016101710181019102010211022102310241025102610271028102910301031103210331034103510361037103810391040104110421043104410451046104710481049105010511052105310541055105610571058105910601061106210631064106510661067106810691070107110721073107410751076107710781079108010811082108310841085108610871088108910901091109210931094109510961097109810991100110111021103110411051106110711081109111011111112111311141115111611171118111911201121112211231124112511261127112811291130113111321133113411351136113711381139114011411142114311441145114611471148114911501151115211531154115511561157115811591160116111621163116411651166116711681169117011711172117311741175117611771178117911801181118211831184118511861187118811891190119111921193119411951196119711981199120012011202120312041205120612071208120912101211121212131214121512161217121812191220122112221223122412251226122712281229123012311232123312341235123612371238123912401241124212431244124512461247124812491250125112521253125412551256125712581259126012611262126312641265126612671268126912701271127212731274127512761277127812791280128112821283128412851286128712881289129012911292129312941295129612971298129913001301130213031304130513061307130813091310131113121313131413151316131713181319132013211322132313241325132613271328132913301331133213331334133513361337133813391340134113421343134413451346134713481349135013511352135313541355135613571358135913601361136213631364136513661367136813691370137113721373137413751376137713781379138013811382138313841385138613871388138913901391139213931394139513961397139813991400140114021403140414051406140714081409141014111412141314141415141614171418141914201421142214231424142514261427142814291430143114321433143414351436143714381439144014411442144314441445144614471448144914501451145214531454145514561457145814591460146114621463146414651466146714681469147014711472147314741475147614771478147914801481148214831484148514861487148814891490149114921493149414951496149714981499150015011502150315041505150615071508150915101511151215131514151515161517151815191520152115221523152415251526152715281529153015311532153315341535153615371538153915401541154215431544154515461547154815491550155115521553155415551556155715581559156015611562156315641565156615671568156915701571157215731574157515761577157815791580158115821583158415851586158715881589159015911592159315941595159615971598159916001601160216031604160516061607160816091610161116121613161416151616161716181619162016211622162316241625162616271628162916301631163216331634163516361637163816391640164116421643164416451646164716481649165016511652165316541655165616571658165916601661166216631664166516661667166816691670167116721673167416751676167716781679168016811682168316841685168616871688168916901691169216931694169516961697169816991700170117021703170417051706170717081709171017111712171317141715171617171718171917201721172217231724172517261727172817291730173117321733 |
- package service
- import (
- "context"
- "encoding/json"
- "errors"
- "fmt"
- "log/slog"
- "sort"
- "strconv"
- "sync"
- "time"
- "github.com/Wei-Shaw/sub2api/internal/config"
- "github.com/Wei-Shaw/sub2api/internal/pkg/logger"
- )
- var (
- ErrSchedulerCacheNotReady = errors.New("scheduler cache not ready")
- ErrSchedulerFallbackLimited = errors.New("scheduler db fallback limited")
- ErrSchedulerGroupLifecycleLeaseBusy = errors.New("scheduler group lifecycle lease busy")
- ErrSchedulerBucketRebuildBusy = errors.New("scheduler bucket rebuild busy")
- )
- const (
- outboxEventTimeout = 2 * time.Minute
- schedulerOutboxCleanupBatch = 5000
- schedulerGroupLifecycleTimeout = 30 * time.Second
- schedulerGroupLifecycleLeaseTTL = 60 * time.Second
- schedulerGroupLifecycleReleaseTimeout = 2 * time.Second
- outboxRebuildRetryBaseDelay = 5 * time.Second
- outboxRebuildRetryMaxDelay = 5 * time.Minute
- outboxMaxIDErrorLogSampleInterval = time.Minute
- )
- // batchSeenKey tracks completed per-platform rebuilds and group lifecycle work
- // within one pollOutbox call.
- type batchSeenKey struct {
- groupID int64
- platform string
- lifecycle bool
- }
- type schedulerBucketWriteTask struct {
- bucket SchedulerBucket
- token SchedulerBucketWriteToken
- }
- type schedulerAccountQueryKey struct {
- groupID int64
- platform string
- }
- // 查询结果只在一次 rebuild batch 内,按原始 groupID+platform 复用成功的 single/forced 查询;
- // mixed 与历史模式保持独立。每个 task 都用 defer 消费 remaining,最后一个消费者会立即释放结果,
- // 避免把账号切片的生命周期扩大到整轮 full rebuild。
- type schedulerAccountQueryCache struct {
- remaining map[schedulerAccountQueryKey]int
- accounts map[schedulerAccountQueryKey][]Account
- snapshotAccountIDs map[schedulerAccountQueryKey][]int64
- }
- // schedulerSnapshotAccountIDWriter 是 SchedulerCache 的可选批次优化能力。
- // 首次完整发布成功后返回实际可编码账号 ID;同一查询结果的后续桶只需发布这些 ID,
- // 避免重复序列化并覆盖全局账号缓存。未实现该接口的缓存继续走原 SetSnapshot 路径。
- type schedulerSnapshotAccountIDWriter interface {
- SetSnapshotAndReturnAccountIDs(ctx context.Context, bucket SchedulerBucket, token SchedulerBucketWriteToken, accounts []Account) ([]int64, error)
- SetSnapshotByAccountIDs(ctx context.Context, bucket SchedulerBucket, token SchedulerBucketWriteToken, accountIDs []int64) error
- }
- func newSchedulerAccountQueryCache(taskSets ...[]schedulerBucketWriteTask) *schedulerAccountQueryCache {
- queries := &schedulerAccountQueryCache{
- remaining: make(map[schedulerAccountQueryKey]int),
- accounts: make(map[schedulerAccountQueryKey][]Account),
- snapshotAccountIDs: make(map[schedulerAccountQueryKey][]int64),
- }
- for _, tasks := range taskSets {
- for _, task := range tasks {
- if key, ok := schedulerAccountQueryKeyForBucket(task.bucket); ok {
- queries.remaining[key]++
- }
- }
- }
- return queries
- }
- func schedulerAccountQueryKeyForBucket(bucket SchedulerBucket) (schedulerAccountQueryKey, bool) {
- if bucket.Mode != SchedulerModeSingle && bucket.Mode != SchedulerModeForced {
- return schedulerAccountQueryKey{}, false
- }
- return schedulerAccountQueryKey{groupID: bucket.GroupID, platform: bucket.Platform}, true
- }
- func (c *schedulerAccountQueryCache) release(bucket SchedulerBucket) {
- if c == nil {
- return
- }
- key, ok := schedulerAccountQueryKeyForBucket(bucket)
- if !ok {
- return
- }
- remaining := c.remaining[key] - 1
- if remaining <= 0 {
- delete(c.remaining, key)
- delete(c.accounts, key)
- delete(c.snapshotAccountIDs, key)
- return
- }
- c.remaining[key] = remaining
- }
- type schedulerGroupLifecyclePlan struct {
- active bool
- tasks []schedulerBucketWriteTask
- }
- type schedulerActiveGroupIDLister interface {
- ListActiveIDs(ctx context.Context) ([]int64, error)
- }
- type SchedulerSnapshotService struct {
- cache SchedulerCache
- outboxRepo SchedulerOutboxRepository
- accountRepo AccountRepository
- groupRepo GroupRepository
- cfg *config.Config
- stopCh chan struct{}
- stopOnce sync.Once
- wg sync.WaitGroup
- fallbackLimit *fallbackLimiter
- lagMu sync.Mutex
- lagFailures int
- outboxRebuildLatched bool
- outboxRebuildRunning bool
- outboxRebuildFailures int
- outboxRebuildRetryAt time.Time
- outboxRebuildRetryReason string
- outboxLagWarningActive bool
- outboxMaxIDErrorLastLoggedAt time.Time
- fullRebuildRunMu sync.Mutex
- fullRebuildStateMu sync.Mutex
- fullRebuildRequested uint64
- fullRebuildCompleted uint64
- fullRebuildLastErr error
- }
- func NewSchedulerSnapshotService(
- cache SchedulerCache,
- outboxRepo SchedulerOutboxRepository,
- accountRepo AccountRepository,
- groupRepo GroupRepository,
- cfg *config.Config,
- ) *SchedulerSnapshotService {
- maxQPS := 0
- if cfg != nil {
- maxQPS = cfg.Gateway.Scheduling.DbFallbackMaxQPS
- }
- return &SchedulerSnapshotService{
- cache: cache,
- outboxRepo: outboxRepo,
- accountRepo: accountRepo,
- groupRepo: groupRepo,
- cfg: cfg,
- stopCh: make(chan struct{}),
- fallbackLimit: newFallbackLimiter(maxQPS),
- }
- }
- func (s *SchedulerSnapshotService) Start() {
- if s == nil || s.cache == nil {
- return
- }
- s.wg.Add(1)
- go func() {
- defer s.wg.Done()
- s.runInitialRebuild()
- }()
- interval := s.outboxPollInterval()
- if s.outboxRepo != nil && interval > 0 {
- s.wg.Add(1)
- go func() {
- defer s.wg.Done()
- s.runOutboxWorker(interval)
- }()
- }
- fullInterval := s.fullRebuildInterval()
- if fullInterval > 0 {
- s.wg.Add(1)
- go func() {
- defer s.wg.Done()
- s.runFullRebuildWorker(fullInterval)
- }()
- }
- }
- func (s *SchedulerSnapshotService) Stop() {
- if s == nil {
- return
- }
- s.stopOnce.Do(func() {
- close(s.stopCh)
- })
- s.wg.Wait()
- }
- func (s *SchedulerSnapshotService) ListSchedulableAccounts(ctx context.Context, groupID *int64, platform string, hasForcePlatform bool) ([]Account, bool, error) {
- useMixed := (platform == PlatformAnthropic || platform == PlatformGemini) && !hasForcePlatform
- mode := s.resolveMode(platform, hasForcePlatform)
- bucket := s.bucketFor(groupID, platform, mode)
- var writeToken SchedulerBucketWriteToken
- canPublish := false
- if err := ctx.Err(); err != nil {
- return nil, useMixed, err
- }
- if s.cache != nil {
- cached, hit, err := s.cache.GetSnapshot(ctx, bucket)
- if ctxErr := ctx.Err(); ctxErr != nil {
- return nil, useMixed, ctxErr
- }
- if err != nil {
- logger.LegacyPrintf("service.scheduler_snapshot", "[Scheduler] cache read failed: bucket=%s err=%v", bucket.String(), err)
- } else if hit {
- return derefAccounts(cached), useMixed, nil
- }
- token, err := s.cache.CaptureBucketWriteToken(ctx, bucket)
- if ctxErr := ctx.Err(); ctxErr != nil {
- return nil, useMixed, ctxErr
- }
- if err != nil {
- if errors.Is(err, ErrSchedulerBucketRetired) || errors.Is(err, ErrSchedulerBucketWriteFenced) {
- slog.Debug("[Scheduler] cache publish fenced", "bucket", bucket.String())
- } else {
- logger.LegacyPrintf("service.scheduler_snapshot", "[Scheduler] cache publish token failed: bucket=%s err=%v", bucket.String(), err)
- }
- } else {
- writeToken = token
- canPublish = true
- }
- }
- if err := s.guardFallback(ctx); err != nil {
- return nil, useMixed, err
- }
- fallbackCtx, cancel := s.withFallbackTimeout(ctx)
- defer cancel()
- accounts, err := s.loadAccountsFromDB(fallbackCtx, bucket, useMixed)
- if err != nil {
- return nil, useMixed, err
- }
- if ctxErr := ctx.Err(); ctxErr != nil {
- return nil, useMixed, ctxErr
- }
- if s.cache != nil && canPublish {
- if err := s.cache.SetSnapshot(fallbackCtx, bucket, writeToken, accounts); err != nil {
- if errors.Is(err, ErrSchedulerBucketRetired) || errors.Is(err, ErrSchedulerBucketWriteFenced) {
- slog.Debug("[Scheduler] cache publish fenced", "bucket", bucket.String())
- } else {
- logger.LegacyPrintf("service.scheduler_snapshot", "[Scheduler] cache write failed: bucket=%s err=%v", bucket.String(), err)
- }
- }
- }
- return accounts, useMixed, nil
- }
- func (s *SchedulerSnapshotService) GetAccount(ctx context.Context, accountID int64) (*Account, error) {
- if accountID <= 0 {
- return nil, nil
- }
- if err := ctx.Err(); err != nil {
- return nil, err
- }
- if s.cache != nil {
- account, err := s.cache.GetAccount(ctx, accountID)
- if ctxErr := ctx.Err(); ctxErr != nil {
- return nil, ctxErr
- }
- if err != nil {
- logger.LegacyPrintf("service.scheduler_snapshot", "[Scheduler] account cache read failed: id=%d err=%v", accountID, err)
- } else if account != nil {
- return account, nil
- }
- }
- if err := s.guardFallback(ctx); err != nil {
- return nil, err
- }
- fallbackCtx, cancel := s.withFallbackTimeout(ctx)
- defer cancel()
- return s.accountRepo.GetByID(fallbackCtx, accountID)
- }
- // GetGroupByID 获取分组信息(供调度器使用)
- func (s *SchedulerSnapshotService) GetGroupByID(ctx context.Context, groupID int64) (*Group, error) {
- if s.groupRepo == nil {
- return nil, nil
- }
- return s.groupRepo.GetByID(ctx, groupID)
- }
- // GetGroupByIDLite 获取分组配置但不加载账号计数聚合。
- // 利润门只需要平台、倍率、利润与高峰字段,GetByID 附带的那条账号计数聚合
- // 查询纯属浪费——composite / 模型路由 / fallback 每次装门都要付一次,WS 更是
- // 每个 turn 一次,且发生在「是否启用利润控制」判定之前。
- func (s *SchedulerSnapshotService) GetGroupByIDLite(ctx context.Context, groupID int64) (*Group, error) {
- if s.groupRepo == nil {
- return nil, nil
- }
- return s.groupRepo.GetByIDLite(ctx, groupID)
- }
- // UpdateAccountInCache 立即更新 Redis 中单个账号的数据(用于模型限流后立即生效)
- func (s *SchedulerSnapshotService) UpdateAccountInCache(ctx context.Context, account *Account) error {
- if s.cache == nil || account == nil {
- return nil
- }
- return s.cache.SetAccount(ctx, account)
- }
- func (s *SchedulerSnapshotService) runInitialRebuild() {
- if s.cache == nil {
- return
- }
- _ = s.coalesceFullRebuild(func() error {
- ctx, cancel := context.WithTimeout(context.Background(), 2*time.Minute)
- defer cancel()
- if err := s.rebuildFullSnapshot(ctx, "startup"); err != nil {
- logger.LegacyPrintf("service.scheduler_snapshot", "[Scheduler] rebuild startup failed: %v", err)
- return err
- }
- return nil
- })
- }
- func (s *SchedulerSnapshotService) runOutboxWorker(interval time.Duration) {
- ticker := time.NewTicker(interval)
- defer ticker.Stop()
- s.pollOutbox()
- for {
- select {
- case <-ticker.C:
- s.pollOutbox()
- case <-s.stopCh:
- return
- }
- }
- }
- func (s *SchedulerSnapshotService) runFullRebuildWorker(interval time.Duration) {
- ticker := time.NewTicker(interval)
- defer ticker.Stop()
- for {
- select {
- case <-ticker.C:
- if err := s.triggerFullRebuild("interval"); err != nil {
- logger.LegacyPrintf("service.scheduler_snapshot", "[Scheduler] full rebuild failed: %v", err)
- }
- case <-s.stopCh:
- return
- }
- }
- }
- func (s *SchedulerSnapshotService) pollOutbox() {
- if s.outboxRepo == nil || s.cache == nil {
- return
- }
- ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
- defer cancel()
- watermark, err := s.cache.GetOutboxWatermark(ctx)
- if err != nil {
- logger.LegacyPrintf("service.scheduler_snapshot", "[Scheduler] outbox watermark read failed: %v", err)
- return
- }
- events, err := s.outboxRepo.ListAfterAndReleaseDedup(ctx, watermark, 200)
- if err != nil {
- logger.LegacyPrintf("service.scheduler_snapshot", "[Scheduler] outbox poll failed: %v", err)
- return
- }
- if len(events) == 0 {
- // The outbox query itself proves there is no event after the watermark.
- // Clear degraded/retry state without adding two more repository queries to
- // the healthy one-second poll path.
- s.clearOutboxDegradedEpisode()
- return
- }
- seen := make(map[batchSeenKey]struct{})
- for _, event := range events {
- eventCtx, cancel := context.WithTimeout(context.Background(), outboxEventTimeout)
- err := s.handleOutboxEvent(eventCtx, event, seen)
- cancel()
- if err != nil {
- logger.LegacyPrintf("service.scheduler_snapshot", "[Scheduler] outbox handle failed: id=%d type=%s err=%v", event.ID, event.EventType, err)
- return
- }
- }
- lastID := events[len(events)-1].ID
- var wmErr error
- for i := range 3 {
- wmCtx, wmCancel := context.WithTimeout(context.Background(), 5*time.Second)
- wmErr = s.cache.SetOutboxWatermark(wmCtx, lastID)
- wmCancel()
- if wmErr == nil {
- break
- }
- if i < 2 {
- time.Sleep(200 * time.Millisecond)
- }
- }
- if wmErr != nil {
- logger.LegacyPrintf("service.scheduler_snapshot", "[Scheduler] outbox watermark write failed: %v", wmErr)
- return
- }
- s.cleanupConsumedOutbox(lastID)
- // 只有 watermark 成功推进后,当前批次才算已消费。延迟必须按下一条待消费事件计算,
- // 否则本批次处理越慢,越容易误触发一次更慢的全量重建,形成正反馈。
- lagCtx, lagCancel := context.WithTimeout(context.Background(), 5*time.Second)
- s.checkOutboxLag(lagCtx, lastID)
- lagCancel()
- }
- func (s *SchedulerSnapshotService) cleanupConsumedOutbox(watermark int64) {
- if s == nil || s.outboxRepo == nil || watermark <= 0 {
- return
- }
- ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
- defer cancel()
- lease, acquired, err := s.outboxRepo.TryAcquireCleanupLock(ctx)
- if err != nil {
- logger.LegacyPrintf("service.scheduler_snapshot", "[Scheduler] outbox cleanup lock failed: %v", err)
- return
- }
- if !acquired {
- return
- }
- defer lease.Release()
- for {
- deleted, err := s.outboxRepo.DeleteConsumedUpTo(ctx, watermark, schedulerOutboxCleanupBatch)
- if err != nil {
- logger.LegacyPrintf("service.scheduler_snapshot", "[Scheduler] outbox cleanup failed: watermark=%d err=%v", watermark, err)
- return
- }
- if deleted == 0 || deleted < schedulerOutboxCleanupBatch {
- return
- }
- }
- }
- func (s *SchedulerSnapshotService) handleOutboxEvent(ctx context.Context, event SchedulerOutboxEvent, seen map[batchSeenKey]struct{}) error {
- switch event.EventType {
- case SchedulerOutboxEventAccountLastUsed:
- return s.handleLastUsedEvent(ctx, event.Payload)
- case SchedulerOutboxEventAccountBulkChanged:
- return s.handleBulkAccountEvent(ctx, event.Payload, seen)
- case SchedulerOutboxEventAccountGroupsChanged:
- return s.handleAccountEvent(ctx, event.AccountID, event.Payload, seen)
- case SchedulerOutboxEventAccountChanged:
- return s.handleAccountEvent(ctx, event.AccountID, event.Payload, seen)
- case SchedulerOutboxEventGroupChanged:
- return s.handleGroupEvent(ctx, event.GroupID, seen)
- case SchedulerOutboxEventFullRebuild:
- return s.triggerFullRebuild("outbox")
- default:
- return nil
- }
- }
- func (s *SchedulerSnapshotService) handleLastUsedEvent(ctx context.Context, payload map[string]any) error {
- if s.cache == nil || payload == nil {
- return nil
- }
- raw, ok := payload["last_used"].(map[string]any)
- if !ok || len(raw) == 0 {
- return nil
- }
- updates := make(map[int64]time.Time, len(raw))
- for key, value := range raw {
- id, err := strconv.ParseInt(key, 10, 64)
- if err != nil || id <= 0 {
- continue
- }
- sec, ok := toInt64(value)
- if !ok || sec <= 0 {
- continue
- }
- updates[id] = time.Unix(sec, 0)
- }
- if len(updates) == 0 {
- return nil
- }
- return s.cache.UpdateLastUsed(ctx, updates)
- }
- func (s *SchedulerSnapshotService) handleBulkAccountEvent(ctx context.Context, payload map[string]any, seen map[batchSeenKey]struct{}) error {
- if payload == nil {
- return nil
- }
- if s.accountRepo == nil {
- return nil
- }
- rawIDs := parseInt64Slice(payload["account_ids"])
- if len(rawIDs) == 0 {
- return nil
- }
- ids := make([]int64, 0, len(rawIDs))
- seenIDs := make(map[int64]struct{}, len(rawIDs))
- for _, id := range rawIDs {
- if id <= 0 {
- continue
- }
- if _, exists := seenIDs[id]; exists {
- continue
- }
- seenIDs[id] = struct{}{}
- ids = append(ids, id)
- }
- if len(ids) == 0 {
- return nil
- }
- preloadGroupIDs := parseInt64Slice(payload["group_ids"])
- accounts, err := s.accountRepo.GetByIDs(ctx, ids)
- if err != nil {
- return err
- }
- found := make(map[int64]struct{}, len(accounts))
- rebuildGroupSet := make(map[int64]struct{}, len(preloadGroupIDs))
- for _, gid := range preloadGroupIDs {
- if gid > 0 {
- rebuildGroupSet[gid] = struct{}{}
- }
- }
- for _, account := range accounts {
- if account == nil || account.ID <= 0 {
- continue
- }
- found[account.ID] = struct{}{}
- if s.cache != nil {
- if err := s.cache.SetAccount(ctx, account); err != nil {
- return err
- }
- }
- for _, gid := range account.GroupIDs {
- if gid > 0 {
- rebuildGroupSet[gid] = struct{}{}
- }
- }
- }
- allAccountsFound := true
- for _, id := range ids {
- if _, ok := found[id]; ok {
- continue
- }
- allAccountsFound = false
- if s.cache != nil {
- if err := s.cache.DeleteAccount(ctx, id); err != nil {
- return err
- }
- }
- }
- rebuildGroupIDs := make([]int64, 0, len(rebuildGroupSet))
- for gid := range rebuildGroupSet {
- rebuildGroupIDs = append(rebuildGroupIDs, gid)
- }
- // 缺失账户无法确定原平台,保留全平台重建以避免遗留旧快照。
- if !allAccountsFound {
- return s.rebuildByGroupIDs(ctx, rebuildGroupIDs, "account_bulk_change", seen)
- }
- platformGroupSets := make(map[string]map[int64]struct{}, len(accounts))
- addPlatformGroups := func(platform string, groupIDs []int64) {
- groupSet := platformGroupSets[platform]
- if groupSet == nil {
- groupSet = make(map[int64]struct{}, len(groupIDs))
- platformGroupSets[platform] = groupSet
- }
- for _, groupID := range groupIDs {
- groupSet[groupID] = struct{}{}
- }
- }
- for _, account := range accounts {
- if account == nil || account.ID <= 0 {
- continue
- }
- accountGroupIDs := s.normalizeGroupIDs(account.GroupIDs)
- switch account.Platform {
- case PlatformAnthropic, PlatformGemini, PlatformOpenAI, PlatformGrok, PlatformKimi, PlatformZhipu, PlatformDeepseek, PlatformMiniMax:
- addPlatformGroups(account.Platform, accountGroupIDs)
- case PlatformAntigravity:
- // 批量更新可能刚关闭 mixed_scheduling,仍需清理两个兼容平台的旧快照。
- addPlatformGroups(PlatformAntigravity, accountGroupIDs)
- addPlatformGroups(PlatformAnthropic, accountGroupIDs)
- addPlatformGroups(PlatformGemini, accountGroupIDs)
- default:
- return s.rebuildByGroupIDs(ctx, rebuildGroupIDs, "account_bulk_change", seen)
- }
- }
- // payload 携带更新前的组;只扩散到本事件实际涉及的平台,避免平台间交叉重建。
- if len(preloadGroupIDs) > 0 {
- preloadGroupIDs = s.normalizeGroupIDs(preloadGroupIDs)
- for platform := range platformGroupSets {
- addPlatformGroups(platform, preloadGroupIDs)
- }
- }
- bucketCapacity := 0
- for _, groupSet := range platformGroupSets {
- bucketCapacity += len(groupSet) * 3
- }
- buckets := make([]SchedulerBucket, 0, bucketCapacity)
- for _, platform := range schedulerSnapshotPlatforms() {
- groupSet, ok := platformGroupSets[platform]
- if !ok {
- continue
- }
- platformGroupIDs := make([]int64, 0, len(groupSet))
- for groupID := range groupSet {
- platformGroupIDs = append(platformGroupIDs, groupID)
- }
- sort.Slice(platformGroupIDs, func(i, j int) bool { return platformGroupIDs[i] < platformGroupIDs[j] })
- buckets = append(buckets, s.bucketsForPlatform(platform, platformGroupIDs, seen)...)
- }
- return s.rebuildBuckets(ctx, buckets, "account_bulk_change")
- }
- func (s *SchedulerSnapshotService) handleAccountEvent(ctx context.Context, accountID *int64, payload map[string]any, seen map[batchSeenKey]struct{}) error {
- if accountID == nil || *accountID <= 0 {
- return nil
- }
- if s.accountRepo == nil {
- return nil
- }
- var groupIDs []int64
- if payload != nil {
- groupIDs = parseInt64Slice(payload["group_ids"])
- }
- account, err := s.accountRepo.GetByID(ctx, *accountID)
- if err != nil {
- if errors.Is(err, ErrAccountNotFound) {
- if s.cache != nil {
- if err := s.cache.DeleteAccount(ctx, *accountID); err != nil {
- return err
- }
- }
- return s.rebuildByGroupIDs(ctx, groupIDs, "account_miss", seen)
- }
- return err
- }
- if s.cache != nil {
- if err := s.cache.SetAccount(ctx, account); err != nil {
- return err
- }
- }
- if len(groupIDs) == 0 {
- groupIDs = account.GroupIDs
- }
- return s.rebuildByAccount(ctx, account, groupIDs, "account_change", seen)
- }
- func (s *SchedulerSnapshotService) handleGroupEvent(ctx context.Context, groupID *int64, seen map[batchSeenKey]struct{}) error {
- if groupID == nil || *groupID <= 0 || s.isRunModeSimple() {
- return nil
- }
- if seen != nil {
- if _, ok := seen[batchSeenKey{groupID: *groupID, lifecycle: true}]; ok {
- return nil
- }
- }
- return s.reconcileGroupLifecycle(ctx, *groupID, seen)
- }
- func (s *SchedulerSnapshotService) reconcileGroupLifecycle(ctx context.Context, groupID int64, seen map[batchSeenKey]struct{}) error {
- plan, err := s.prepareGroupLifecycle(ctx, groupID, nil)
- if err != nil {
- return err
- }
- if plan.active {
- queries := newSchedulerAccountQueryCache(plan.tasks)
- for _, task := range plan.tasks {
- if err := s.rebuildBucketWithTokenPolicyAndQueryCache(ctx, task, "group_change", true, queries); err != nil {
- return err
- }
- }
- }
- markGroupLifecycleSeen(seen, groupID)
- return nil
- }
- // 生命周期决策必须在所有者安全的租约内读取 fresh 且完整的分组权威状态。
- // active 仅 Reopen canonical bucket;missing/inactive 同时 Retire canonical 与已登记历史 bucket;
- // group event 路径只有在权威决策和后续重建全部成功后才会标记 seen。
- func (s *SchedulerSnapshotService) prepareGroupLifecycle(ctx context.Context, groupID int64, knownHistorical []SchedulerBucket) (plan schedulerGroupLifecyclePlan, retErr error) {
- if groupID <= 0 || s.isRunModeSimple() {
- return schedulerGroupLifecyclePlan{}, nil
- }
- if s.cache == nil || s.groupRepo == nil {
- return schedulerGroupLifecyclePlan{}, ErrSchedulerCacheNotReady
- }
- lifecycleCtx, cancel := context.WithTimeout(ctx, schedulerGroupLifecycleTimeout)
- defer cancel()
- lease, acquired, err := s.cache.TryAcquireGroupLifecycleLease(lifecycleCtx, groupID, schedulerGroupLifecycleLeaseTTL)
- if err != nil {
- return schedulerGroupLifecyclePlan{}, err
- }
- if !acquired {
- return schedulerGroupLifecyclePlan{}, fmt.Errorf("%w: group=%d", ErrSchedulerGroupLifecycleLeaseBusy, groupID)
- }
- leaseHeld := true
- defer func() {
- if leaseHeld {
- retErr = errors.Join(retErr, s.releaseGroupLifecycleLease(lease))
- }
- }()
- group, err := s.groupRepo.GetByIDLite(lifecycleCtx, groupID)
- missing := errors.Is(err, ErrGroupNotFound)
- if err != nil && !missing {
- return schedulerGroupLifecyclePlan{}, err
- }
- if err == nil && (group == nil || group.ID != groupID || !group.Hydrated) {
- return schedulerGroupLifecyclePlan{}, fmt.Errorf("untrusted scheduler group lifecycle state: group=%d", groupID)
- }
- plan = schedulerGroupLifecyclePlan{active: !missing && group.IsActive()}
- if plan.active {
- buckets := schedulerBucketsForGroup(groupID)
- plan.tasks = make([]schedulerBucketWriteTask, 0, len(buckets))
- for _, bucket := range buckets {
- token, err := s.cache.ReopenBucket(lifecycleCtx, bucket)
- if err != nil {
- return schedulerGroupLifecyclePlan{}, err
- }
- plan.tasks = append(plan.tasks, schedulerBucketWriteTask{bucket: bucket, token: token})
- }
- } else {
- registered := knownHistorical
- if registered == nil {
- registered, err = s.cache.ListBuckets(lifecycleCtx)
- if err != nil {
- return schedulerGroupLifecyclePlan{}, err
- }
- }
- buckets := schedulerBucketsForGroup(groupID)
- for _, bucket := range registered {
- if bucket.GroupID == groupID {
- buckets = append(buckets, bucket)
- }
- }
- for _, bucket := range dedupeBuckets(buckets) {
- if err := s.cache.RetireBucket(lifecycleCtx, bucket); err != nil {
- return schedulerGroupLifecyclePlan{}, err
- }
- }
- }
- releaseErr := s.releaseGroupLifecycleLease(lease)
- leaseHeld = false
- if releaseErr != nil {
- return schedulerGroupLifecyclePlan{}, releaseErr
- }
- return plan, nil
- }
- func (s *SchedulerSnapshotService) releaseGroupLifecycleLease(lease SchedulerGroupLifecycleLease) error {
- // 请求取消后仍需尝试释放自己的租约,因此使用独立且有界的后台上下文。
- releaseCtx, cancel := context.WithTimeout(context.Background(), schedulerGroupLifecycleReleaseTimeout)
- defer cancel()
- return s.cache.ReleaseGroupLifecycleLease(releaseCtx, lease)
- }
- func markGroupLifecycleSeen(seen map[batchSeenKey]struct{}, groupID int64) {
- if seen == nil {
- return
- }
- seen[batchSeenKey{groupID: groupID, lifecycle: true}] = struct{}{}
- for _, platform := range schedulerSnapshotPlatforms() {
- seen[batchSeenKey{groupID: groupID, platform: platform}] = struct{}{}
- }
- }
- func (s *SchedulerSnapshotService) rebuildByAccount(ctx context.Context, account *Account, groupIDs []int64, reason string, seen map[batchSeenKey]struct{}) error {
- if account == nil {
- return nil
- }
- groupIDs = s.normalizeGroupIDs(groupIDs)
- if len(groupIDs) == 0 {
- return nil
- }
- buckets := s.bucketsForPlatform(account.Platform, groupIDs, seen)
- if account.Platform == PlatformAntigravity && account.IsMixedSchedulingEnabled() {
- buckets = append(buckets, s.bucketsForPlatform(PlatformAnthropic, groupIDs, seen)...)
- buckets = append(buckets, s.bucketsForPlatform(PlatformGemini, groupIDs, seen)...)
- }
- return s.rebuildBuckets(ctx, buckets, reason)
- }
- func schedulerSnapshotPlatforms() [9]string {
- return [9]string{PlatformAnthropic, PlatformGemini, PlatformOpenAI, PlatformAntigravity, PlatformGrok, PlatformKimi, PlatformZhipu, PlatformDeepseek, PlatformMiniMax}
- }
- // 生命周期辅助函数有意排除 group0;full rebuild 构造 group0 canonical 集时必须显式调用 canonical helper。
- func schedulerBucketsForGroup(groupID int64) []SchedulerBucket {
- if groupID <= 0 {
- return nil
- }
- return schedulerCanonicalBuckets(groupID)
- }
- func schedulerCanonicalBucketCount() int {
- count := 0
- for _, platform := range schedulerSnapshotPlatforms() {
- count += 2
- if platform == PlatformAnthropic || platform == PlatformGemini {
- count++
- }
- }
- return count
- }
- func schedulerCanonicalBuckets(groupID int64) []SchedulerBucket {
- platforms := schedulerSnapshotPlatforms()
- buckets := make([]SchedulerBucket, 0, len(platforms)*2+2)
- for _, platform := range platforms {
- buckets = append(buckets,
- SchedulerBucket{GroupID: groupID, Platform: platform, Mode: SchedulerModeSingle},
- SchedulerBucket{GroupID: groupID, Platform: platform, Mode: SchedulerModeForced},
- )
- if platform == PlatformAnthropic || platform == PlatformGemini {
- buckets = append(buckets, SchedulerBucket{GroupID: groupID, Platform: platform, Mode: SchedulerModeMixed})
- }
- }
- return buckets
- }
- func (s *SchedulerSnapshotService) rebuildByGroupIDs(ctx context.Context, groupIDs []int64, reason string, seen map[batchSeenKey]struct{}) error {
- groupIDs = s.normalizeGroupIDs(groupIDs)
- if len(groupIDs) == 0 {
- return nil
- }
- buckets := make([]SchedulerBucket, 0, len(groupIDs)*schedulerCanonicalBucketCount())
- for _, platform := range schedulerSnapshotPlatforms() {
- buckets = append(buckets, s.bucketsForPlatform(platform, groupIDs, seen)...)
- }
- return s.rebuildBuckets(ctx, buckets, reason)
- }
- func (s *SchedulerSnapshotService) bucketsForPlatform(platform string, groupIDs []int64, seen map[batchSeenKey]struct{}) []SchedulerBucket {
- if platform == "" {
- return nil
- }
- buckets := make([]SchedulerBucket, 0, len(groupIDs)*3)
- for _, gid := range groupIDs {
- // Within a single poll batch, skip (groupID, platform) pairs that were
- // already rebuilt. The first rebuild loads fresh DB data for all accounts
- // in the group, so subsequent rebuilds for the same group+platform within
- // the same batch are redundant.
- if seen != nil {
- key := batchSeenKey{groupID: gid, platform: platform}
- if _, exists := seen[key]; exists {
- continue
- }
- seen[key] = struct{}{}
- }
- buckets = append(buckets, SchedulerBucket{GroupID: gid, Platform: platform, Mode: SchedulerModeSingle})
- buckets = append(buckets, SchedulerBucket{GroupID: gid, Platform: platform, Mode: SchedulerModeForced})
- if platform == PlatformAnthropic || platform == PlatformGemini {
- buckets = append(buckets, SchedulerBucket{GroupID: gid, Platform: platform, Mode: SchedulerModeMixed})
- }
- }
- return buckets
- }
- func (s *SchedulerSnapshotService) rebuildBuckets(ctx context.Context, buckets []SchedulerBucket, reason string) error {
- tasks, firstErr := s.prepareBucketWriteTasks(ctx, buckets)
- queries := newSchedulerAccountQueryCache(tasks)
- if err := s.rebuildPreparedBucketTasks(ctx, tasks, reason, false, queries); err != nil && firstErr == nil {
- firstErr = err
- }
- return firstErr
- }
- func (s *SchedulerSnapshotService) prepareBucketWriteTasks(ctx context.Context, buckets []SchedulerBucket) ([]schedulerBucketWriteTask, error) {
- if s.cache == nil {
- return nil, ErrSchedulerCacheNotReady
- }
- tasks := make([]schedulerBucketWriteTask, 0, len(buckets))
- var firstErr error
- for _, bucket := range buckets {
- token, err := s.cache.CaptureBucketWriteToken(ctx, bucket)
- if err != nil {
- if errors.Is(err, ErrSchedulerBucketRetired) || errors.Is(err, ErrSchedulerBucketWriteFenced) {
- continue
- }
- if firstErr == nil {
- firstErr = err
- }
- continue
- }
- tasks = append(tasks, schedulerBucketWriteTask{bucket: bucket, token: token})
- }
- return tasks, firstErr
- }
- func (s *SchedulerSnapshotService) rebuildPreparedBucketTasks(
- ctx context.Context,
- tasks []schedulerBucketWriteTask,
- reason string,
- strict bool,
- queries *schedulerAccountQueryCache,
- ) error {
- var firstErr error
- for _, task := range tasks {
- if err := s.rebuildBucketWithTokenPolicyAndQueryCache(ctx, task, reason, strict, queries); err != nil && firstErr == nil {
- firstErr = err
- }
- }
- return firstErr
- }
- func (s *SchedulerSnapshotService) rebuildBucketWithTokenPolicyAndQueryCache(
- ctx context.Context,
- task schedulerBucketWriteTask,
- reason string,
- strict bool,
- queries *schedulerAccountQueryCache,
- ) error {
- if queries != nil {
- defer queries.release(task.bucket)
- }
- if s.cache == nil {
- return ErrSchedulerCacheNotReady
- }
- bucket := task.bucket
- ok, err := s.cache.TryLockBucket(ctx, bucket, 30*time.Second)
- if err != nil {
- return err
- }
- if !ok {
- if strict {
- return fmt.Errorf("%w: bucket=%s", ErrSchedulerBucketRebuildBusy, bucket.String())
- }
- return nil
- }
- defer func() {
- _ = s.cache.UnlockBucket(ctx, bucket)
- }()
- rebuildCtx, cancel := context.WithTimeout(ctx, 30*time.Second)
- defer cancel()
- accounts, err := s.loadAccountsForRebuild(rebuildCtx, bucket, queries)
- if err != nil {
- logger.LegacyPrintf("service.scheduler_snapshot", "[Scheduler] rebuild failed: bucket=%s reason=%s err=%v", bucket.String(), reason, err)
- return err
- }
- if err := s.setRebuildSnapshot(rebuildCtx, task, accounts, queries); err != nil {
- if errors.Is(err, ErrSchedulerBucketRetired) || errors.Is(err, ErrSchedulerBucketWriteFenced) {
- slog.Debug("[Scheduler] rebuild fenced", "bucket", bucket.String(), "reason", reason)
- if strict {
- return err
- }
- return nil
- }
- logger.LegacyPrintf("service.scheduler_snapshot", "[Scheduler] rebuild cache failed: bucket=%s reason=%s err=%v", bucket.String(), reason, err)
- return err
- }
- slog.Debug("[Scheduler] rebuild ok", "bucket", bucket.String(), "reason", reason, "size", len(accounts))
- return nil
- }
- func (s *SchedulerSnapshotService) setRebuildSnapshot(
- ctx context.Context,
- task schedulerBucketWriteTask,
- accounts []Account,
- queries *schedulerAccountQueryCache,
- ) error {
- writer, ok := s.cache.(schedulerSnapshotAccountIDWriter)
- key, reusable := schedulerAccountQueryKeyForBucket(task.bucket)
- if !ok || queries == nil || !reusable {
- return s.cache.SetSnapshot(ctx, task.bucket, task.token, accounts)
- }
- if accountIDs, exists := queries.snapshotAccountIDs[key]; exists {
- return writer.SetSnapshotByAccountIDs(ctx, task.bucket, task.token, accountIDs)
- }
- if queries.remaining[key] <= 1 {
- return s.cache.SetSnapshot(ctx, task.bucket, task.token, accounts)
- }
- accountIDs, err := writer.SetSnapshotAndReturnAccountIDs(ctx, task.bucket, task.token, accounts)
- if err != nil {
- return err
- }
- if queries.remaining[key] > 1 {
- // 必须保存实际成功编码并写入的有序 ID,不能从原账号切片重新推导;
- // 否则不可编码账号会只出现在后续桶中,破坏两个快照的成员一致性。
- // 返回切片由当前批次独占,直接接管可避免 10k 账号场景再次复制。
- queries.snapshotAccountIDs[key] = accountIDs
- }
- return nil
- }
- func (s *SchedulerSnapshotService) triggerFullRebuild(reason string) error {
- if s.cache == nil {
- return ErrSchedulerCacheNotReady
- }
- return s.coalesceFullRebuild(func() error {
- ctx, cancel := context.WithTimeout(context.Background(), 2*time.Minute)
- defer cancel()
- return s.rebuildFullSnapshot(ctx, reason)
- })
- }
- func (s *SchedulerSnapshotService) rebuildFullSnapshot(ctx context.Context, reason string) error {
- if s.cache == nil {
- return ErrSchedulerCacheNotReady
- }
- // 当前模式所需的全局读取必须先成功:桶注册表始终必需,standard 还需活跃分组 ID;
- // 失败时不执行 Capture/Retire/Reopen 或 DB 查询。
- // simple 模式不获取分组生命周期权威;standard 的 stale candidate 仍须在租约内 fresh 确认后才能退休。
- registered, err := s.cache.ListBuckets(ctx)
- if err != nil {
- return err
- }
- registered = dedupeBuckets(registered)
- if s.isRunModeSimple() {
- canonical := schedulerCanonicalBuckets(0)
- captured, err := s.captureFullRebuildCanonicalTasks(ctx, canonical)
- if err != nil {
- return err
- }
- ordinary := appendBucketsExcept(nil, registered, canonical)
- return s.prepareAndRebuildFullSnapshot(ctx, captured, nil, ordinary, reason)
- }
- activeGroupIDs, err := s.listActiveSchedulerGroupIDs(ctx)
- if err != nil {
- return err
- }
- activeGroups := make(map[int64]struct{}, len(activeGroupIDs))
- for _, groupID := range activeGroupIDs {
- activeGroups[groupID] = struct{}{}
- }
- registeredByGroup := make(map[int64][]SchedulerBucket)
- for _, bucket := range registered {
- registeredByGroup[bucket.GroupID] = append(registeredByGroup[bucket.GroupID], bucket)
- }
- groupZeroCanonical := schedulerCanonicalBuckets(0)
- capturedTasks, err := s.captureFullRebuildCanonicalTasks(ctx, groupZeroCanonical)
- if err != nil {
- return err
- }
- ordinaryBuckets := appendBucketsExcept(nil, registeredByGroup[0], groupZeroCanonical)
- for groupID, buckets := range registeredByGroup {
- if groupID < 0 {
- ordinaryBuckets = append(ordinaryBuckets, buckets...)
- }
- }
- reopenedTasks := make([]schedulerBucketWriteTask, 0)
- for _, groupID := range activeGroupIDs {
- canonical := schedulerBucketsForGroup(groupID)
- canonicalTasks, captureErr := s.captureFullRebuildCanonicalTasks(ctx, canonical)
- if captureErr == nil {
- capturedTasks = append(capturedTasks, canonicalTasks...)
- ordinaryBuckets = appendBucketsExcept(ordinaryBuckets, registeredByGroup[groupID], canonical)
- continue
- }
- if !errors.Is(captureErr, ErrSchedulerBucketRetired) && !errors.Is(captureErr, ErrSchedulerBucketWriteFenced) {
- return captureErr
- }
- // A prior full_rebuild event can observe the active state committed for a
- // later group_changed event. Recover here under fresh authority so the
- // earlier event cannot block the outbox watermark before that event runs.
- knownHistorical := registeredByGroup[groupID]
- if knownHistorical == nil {
- knownHistorical = []SchedulerBucket{}
- }
- plan, err := s.prepareGroupLifecycle(ctx, groupID, knownHistorical)
- if err != nil {
- return err
- }
- if plan.active {
- reopenedTasks = append(reopenedTasks, plan.tasks...)
- ordinaryBuckets = appendBucketsExcept(ordinaryBuckets, registeredByGroup[groupID], canonical)
- }
- }
- staleGroupIDs := make([]int64, 0)
- for groupID := range registeredByGroup {
- if groupID <= 0 {
- continue
- }
- if _, active := activeGroups[groupID]; !active {
- staleGroupIDs = append(staleGroupIDs, groupID)
- }
- }
- sort.Slice(staleGroupIDs, func(i, j int) bool { return staleGroupIDs[i] < staleGroupIDs[j] })
- for _, groupID := range staleGroupIDs {
- plan, err := s.prepareGroupLifecycle(ctx, groupID, registeredByGroup[groupID])
- if err != nil {
- return err
- }
- if plan.active {
- reopenedTasks = append(reopenedTasks, plan.tasks...)
- ordinaryBuckets = appendBucketsExcept(ordinaryBuckets, registeredByGroup[groupID], schedulerBucketsForGroup(groupID))
- }
- }
- return s.prepareAndRebuildFullSnapshot(ctx, capturedTasks, reopenedTasks, ordinaryBuckets, reason)
- }
- func (s *SchedulerSnapshotService) listActiveSchedulerGroupIDs(ctx context.Context) ([]int64, error) {
- if s.groupRepo == nil {
- return nil, ErrSchedulerCacheNotReady
- }
- // 轻量接口一旦实现,其错误直接失败;只有仓储不支持该接口时才回退完整 ListActive。
- var groupIDs []int64
- if lister, ok := s.groupRepo.(schedulerActiveGroupIDLister); ok {
- ids, err := lister.ListActiveIDs(ctx)
- if err != nil {
- return nil, err
- }
- groupIDs = ids
- } else {
- groups, err := s.groupRepo.ListActive(ctx)
- if err != nil {
- return nil, err
- }
- groupIDs = make([]int64, 0, len(groups))
- for _, group := range groups {
- groupIDs = append(groupIDs, group.ID)
- }
- }
- seen := make(map[int64]struct{}, len(groupIDs))
- normalized := make([]int64, 0, len(groupIDs))
- for _, groupID := range groupIDs {
- if groupID <= 0 {
- continue
- }
- if _, ok := seen[groupID]; ok {
- continue
- }
- seen[groupID] = struct{}{}
- normalized = append(normalized, groupID)
- }
- sort.Slice(normalized, func(i, j int) bool { return normalized[i] < normalized[j] })
- return normalized, nil
- }
- func (s *SchedulerSnapshotService) prepareAndRebuildFullSnapshot(
- ctx context.Context,
- captured []schedulerBucketWriteTask,
- reopened []schedulerBucketWriteTask,
- ordinaryBuckets []SchedulerBucket,
- reason string,
- ) error {
- // 首个 DB 查询前必须完成全部普通 bucket 的 token 预备;任何预备错误都不会留下部分发布。
- // fresh Reopen task 保持严格锁与 fencing 语义,普通 captured task 继续沿用 lock busy/fence 跳过语义。
- preparedBuckets := make(map[SchedulerBucket]struct{}, len(captured)+len(reopened))
- for _, task := range captured {
- preparedBuckets[task.bucket] = struct{}{}
- }
- for _, task := range reopened {
- preparedBuckets[task.bucket] = struct{}{}
- }
- ordinaryBuckets = dedupeBuckets(ordinaryBuckets)
- toCapture := make([]SchedulerBucket, 0, len(ordinaryBuckets))
- for _, bucket := range ordinaryBuckets {
- if _, ok := preparedBuckets[bucket]; !ok {
- toCapture = append(toCapture, bucket)
- }
- }
- ordinary, firstErr := s.prepareBucketWriteTasks(ctx, toCapture)
- if firstErr != nil {
- return firstErr
- }
- captured = append(captured, ordinary...)
- queries := newSchedulerAccountQueryCache(reopened, captured)
- if err := s.rebuildPreparedBucketTasks(ctx, reopened, reason, true, queries); err != nil {
- firstErr = err
- }
- if err := s.rebuildPreparedBucketTasks(ctx, captured, reason, false, queries); err != nil && firstErr == nil {
- firstErr = err
- }
- return firstErr
- }
- func (s *SchedulerSnapshotService) captureFullRebuildCanonicalTasks(ctx context.Context, buckets []SchedulerBucket) ([]schedulerBucketWriteTask, error) {
- if s.cache == nil {
- return nil, ErrSchedulerCacheNotReady
- }
- tasks := make([]schedulerBucketWriteTask, 0, len(buckets))
- for _, bucket := range buckets {
- token, err := s.cache.CaptureBucketWriteToken(ctx, bucket)
- if err != nil {
- return nil, err
- }
- tasks = append(tasks, schedulerBucketWriteTask{bucket: bucket, token: token})
- }
- return tasks, nil
- }
- func appendBucketsExcept(dst, buckets, excluded []SchedulerBucket) []SchedulerBucket {
- excludedKeys := make(map[SchedulerBucket]struct{}, len(excluded))
- for _, bucket := range excluded {
- excludedKeys[bucket] = struct{}{}
- }
- for _, bucket := range buckets {
- if _, ok := excludedKeys[bucket]; !ok {
- dst = append(dst, bucket)
- }
- }
- return dst
- }
- func (s *SchedulerSnapshotService) coalesceFullRebuild(run func() error) error {
- s.fullRebuildStateMu.Lock()
- s.fullRebuildRequested++
- requestID := s.fullRebuildRequested
- s.fullRebuildStateMu.Unlock()
- s.fullRebuildRunMu.Lock()
- defer s.fullRebuildRunMu.Unlock()
- s.fullRebuildStateMu.Lock()
- if s.fullRebuildCompleted >= requestID {
- err := s.fullRebuildLastErr
- s.fullRebuildStateMu.Unlock()
- return err
- }
- // 当前轮重建可能早于新 outbox 事件对应事务的提交,不能让后到请求直接复用当前轮。
- // 每轮开始前记录可覆盖的请求代次,执行期间登记的请求统一合并到下一轮。
- coveredThrough := s.fullRebuildRequested
- s.fullRebuildStateMu.Unlock()
- err := run()
- s.fullRebuildStateMu.Lock()
- s.fullRebuildCompleted = coveredThrough
- s.fullRebuildLastErr = err
- s.fullRebuildStateMu.Unlock()
- return err
- }
- func (s *SchedulerSnapshotService) checkOutboxLag(ctx context.Context, watermark int64) {
- if s.cfg == nil || s.outboxRepo == nil {
- return
- }
- now := time.Now()
- oldestCreatedAt, ok, err := s.outboxRepo.FirstCreatedAtAfter(ctx, watermark)
- if err != nil {
- logger.LegacyPrintf("service.scheduler_snapshot", "[Scheduler] outbox pending event read failed: %v", err)
- return
- }
- var lag time.Duration
- if ok && !oldestCreatedAt.IsZero() {
- lag = now.Sub(oldestCreatedAt)
- }
- lagSeconds := int(lag.Seconds())
- lagWarning := ok && !oldestCreatedAt.IsZero() &&
- s.cfg.Gateway.Scheduling.OutboxLagWarnSeconds > 0 &&
- lagSeconds >= s.cfg.Gateway.Scheduling.OutboxLagWarnSeconds
- lagDegraded := ok && !oldestCreatedAt.IsZero() &&
- s.cfg.Gateway.Scheduling.OutboxLagRebuildSeconds > 0 &&
- lagSeconds >= s.cfg.Gateway.Scheduling.OutboxLagRebuildSeconds
- backlogThreshold := s.cfg.Gateway.Scheduling.OutboxBacklogRebuildRows
- backlogKnown := true
- var backlog int64
- if backlogThreshold > 0 {
- maxID, maxErr := s.outboxRepo.MaxID(ctx)
- if maxErr != nil {
- backlogKnown = false
- if s.shouldLogOutboxMaxIDError(now) {
- logger.LegacyPrintf("service.scheduler_snapshot", "[Scheduler] outbox max id read failed: %v", maxErr)
- }
- } else {
- backlog = maxID - watermark
- }
- }
- backlogDegraded := backlogKnown && backlogThreshold > 0 && backlog >= int64(backlogThreshold)
- // A successful rebuild latches the degraded episode until recovery. A failed
- // rebuild remains retryable, but only after an exponentially backed-off
- // cooldown so a one-second poll cannot create a rebuild storm.
- logLagWarning := s.shouldLogOutboxLagWarning(lagWarning)
- s.lagMu.Lock()
- fullyRecovered := !lagDegraded && backlogKnown && !backlogDegraded
- if fullyRecovered {
- s.lagFailures = 0
- s.outboxRebuildLatched = false
- s.outboxRebuildFailures = 0
- s.outboxRebuildRetryAt = time.Time{}
- s.outboxRebuildRetryReason = ""
- }
- if s.outboxRebuildRetryReason != "" {
- retryReasonActive := (s.outboxRebuildRetryReason == "outbox_lag" && lagDegraded) ||
- (s.outboxRebuildRetryReason == "outbox_backlog" && (!backlogKnown || backlogDegraded))
- if !retryReasonActive {
- s.outboxRebuildFailures = 0
- s.outboxRebuildRetryAt = time.Time{}
- s.outboxRebuildRetryReason = ""
- }
- }
- lagRetryPending := s.outboxRebuildRetryReason == "outbox_lag" && !s.outboxRebuildRetryAt.IsZero()
- if lagDegraded {
- if !s.outboxRebuildLatched && !s.outboxRebuildRunning && !lagRetryPending {
- s.lagFailures++
- }
- } else {
- s.lagFailures = 0
- }
- failures := s.lagFailures
- lagReady := lagDegraded && failures >= s.cfg.Gateway.Scheduling.OutboxLagRebuildFailures
- retryDue := s.outboxRebuildRetryReason != "" &&
- !s.outboxRebuildRetryAt.IsZero() && !now.Before(s.outboxRebuildRetryAt)
- reason := ""
- lagCanPreemptRetry := lagReady && s.outboxRebuildRetryReason != "outbox_lag"
- if !s.outboxRebuildLatched && !s.outboxRebuildRunning &&
- (s.outboxRebuildRetryAt.IsZero() || retryDue || lagCanPreemptRetry) {
- switch {
- case lagReady || (retryDue && s.outboxRebuildRetryReason == "outbox_lag" && lagDegraded):
- if s.outboxRebuildRetryReason != "" && s.outboxRebuildRetryReason != "outbox_lag" {
- s.outboxRebuildFailures = 0
- s.outboxRebuildRetryAt = time.Time{}
- s.outboxRebuildRetryReason = ""
- }
- reason = "outbox_lag"
- s.lagFailures = 0
- case backlogDegraded && (s.outboxRebuildRetryReason == "" ||
- (retryDue && s.outboxRebuildRetryReason == "outbox_backlog")):
- reason = "outbox_backlog"
- }
- if reason != "" {
- s.outboxRebuildRunning = true
- }
- }
- s.lagMu.Unlock()
- if logLagWarning {
- logger.LegacyPrintf("service.scheduler_snapshot", "[Scheduler] outbox lag warning: %ds", lagSeconds)
- }
- if reason == "" {
- return
- }
- var rebuildErr error
- switch reason {
- case "outbox_lag":
- logger.LegacyPrintf("service.scheduler_snapshot", "[Scheduler] outbox lag rebuild triggered: lag=%s failures=%d", lag, failures)
- rebuildErr = s.triggerFullRebuild(reason)
- case "outbox_backlog":
- logger.LegacyPrintf("service.scheduler_snapshot", "[Scheduler] outbox backlog rebuild triggered: backlog=%d", backlog)
- rebuildErr = s.triggerFullRebuild(reason)
- }
- s.lagMu.Lock()
- s.outboxRebuildRunning = false
- if rebuildErr == nil {
- s.outboxRebuildLatched = true
- s.outboxRebuildFailures = 0
- s.outboxRebuildRetryAt = time.Time{}
- s.outboxRebuildRetryReason = ""
- } else {
- s.outboxRebuildLatched = false
- s.outboxRebuildFailures++
- s.outboxRebuildRetryAt = time.Now().Add(outboxRebuildRetryDelay(s.outboxRebuildFailures))
- s.outboxRebuildRetryReason = reason
- }
- s.lagMu.Unlock()
- if rebuildErr != nil {
- logger.LegacyPrintf("service.scheduler_snapshot", "[Scheduler] %s rebuild failed: %v", reason, rebuildErr)
- }
- }
- func outboxRebuildRetryDelay(failures int) time.Duration {
- delay := outboxRebuildRetryBaseDelay
- for i := 1; i < failures && delay < outboxRebuildRetryMaxDelay; i++ {
- delay *= 2
- if delay >= outboxRebuildRetryMaxDelay {
- return outboxRebuildRetryMaxDelay
- }
- }
- return delay
- }
- func (s *SchedulerSnapshotService) clearOutboxDegradedEpisode() {
- if s == nil {
- return
- }
- s.lagMu.Lock()
- if s.lagFailures != 0 || s.outboxRebuildLatched || s.outboxRebuildRunning ||
- s.outboxRebuildFailures != 0 || !s.outboxRebuildRetryAt.IsZero() ||
- s.outboxRebuildRetryReason != "" || s.outboxLagWarningActive {
- s.lagFailures = 0
- s.outboxRebuildLatched = false
- s.outboxRebuildFailures = 0
- s.outboxRebuildRetryAt = time.Time{}
- s.outboxRebuildRetryReason = ""
- s.outboxLagWarningActive = false
- }
- s.lagMu.Unlock()
- }
- func (s *SchedulerSnapshotService) shouldLogOutboxMaxIDError(now time.Time) bool {
- s.lagMu.Lock()
- defer s.lagMu.Unlock()
- if !s.outboxMaxIDErrorLastLoggedAt.IsZero() && now.Sub(s.outboxMaxIDErrorLastLoggedAt) < outboxMaxIDErrorLogSampleInterval {
- return false
- }
- s.outboxMaxIDErrorLastLoggedAt = now
- return true
- }
- func (s *SchedulerSnapshotService) shouldLogOutboxLagWarning(active bool) bool {
- s.lagMu.Lock()
- defer s.lagMu.Unlock()
- shouldLog := active && !s.outboxLagWarningActive
- s.outboxLagWarningActive = active
- return shouldLog
- }
- func (s *SchedulerSnapshotService) loadAccountsFromDB(ctx context.Context, bucket SchedulerBucket, useMixed bool) ([]Account, error) {
- if s.accountRepo == nil {
- return nil, ErrSchedulerCacheNotReady
- }
- groupID := bucket.GroupID
- if s.isRunModeSimple() {
- groupID = 0
- }
- if useMixed {
- platforms := []string{bucket.Platform, PlatformAntigravity}
- var accounts []Account
- var err error
- if groupID > 0 {
- accounts, err = s.accountRepo.ListSchedulableByGroupIDAndPlatforms(ctx, groupID, platforms)
- } else if s.isRunModeSimple() {
- accounts, err = s.accountRepo.ListSchedulableByPlatforms(ctx, platforms)
- } else {
- accounts, err = s.accountRepo.ListSchedulableUngroupedByPlatforms(ctx, platforms)
- }
- if err != nil {
- return nil, err
- }
- filtered := make([]Account, 0, len(accounts))
- for _, acc := range accounts {
- if acc.Platform == PlatformAntigravity && !acc.IsMixedSchedulingEnabled() {
- continue
- }
- filtered = append(filtered, acc)
- }
- return filtered, nil
- }
- if groupID > 0 {
- return s.accountRepo.ListSchedulableByGroupIDAndPlatform(ctx, groupID, bucket.Platform)
- }
- if s.isRunModeSimple() {
- return s.accountRepo.ListSchedulableByPlatform(ctx, bucket.Platform)
- }
- return s.accountRepo.ListSchedulableUngroupedByPlatform(ctx, bucket.Platform)
- }
- func (s *SchedulerSnapshotService) loadAccountsForRebuild(
- ctx context.Context,
- bucket SchedulerBucket,
- queries *schedulerAccountQueryCache,
- ) ([]Account, error) {
- key, cacheable := schedulerAccountQueryKeyForBucket(bucket)
- if queries == nil || !cacheable {
- return s.loadAccountsFromDB(ctx, bucket, bucket.Mode == SchedulerModeMixed)
- }
- if accounts, ok := queries.accounts[key]; ok {
- return accounts, nil
- }
- if queries.remaining[key] <= 1 {
- return s.loadAccountsFromDB(ctx, bucket, false)
- }
- accounts, err := s.loadAccountsFromDB(ctx, bucket, false)
- if err != nil {
- return nil, err
- }
- queries.accounts[key] = accounts
- return accounts, nil
- }
- func (s *SchedulerSnapshotService) bucketFor(groupID *int64, platform string, mode string) SchedulerBucket {
- return SchedulerBucket{
- GroupID: s.normalizeGroupID(groupID),
- Platform: platform,
- Mode: mode,
- }
- }
- func (s *SchedulerSnapshotService) normalizeGroupID(groupID *int64) int64 {
- if s.isRunModeSimple() {
- return 0
- }
- if groupID == nil || *groupID <= 0 {
- return 0
- }
- return *groupID
- }
- func (s *SchedulerSnapshotService) normalizeGroupIDs(groupIDs []int64) []int64 {
- if s.isRunModeSimple() {
- return []int64{0}
- }
- if len(groupIDs) == 0 {
- return []int64{0}
- }
- seen := make(map[int64]struct{}, len(groupIDs))
- out := make([]int64, 0, len(groupIDs))
- for _, id := range groupIDs {
- if id <= 0 {
- continue
- }
- if _, ok := seen[id]; ok {
- continue
- }
- seen[id] = struct{}{}
- out = append(out, id)
- }
- if len(out) == 0 {
- return []int64{0}
- }
- return out
- }
- func (s *SchedulerSnapshotService) resolveMode(platform string, hasForcePlatform bool) string {
- if hasForcePlatform {
- return SchedulerModeForced
- }
- if platform == PlatformAnthropic || platform == PlatformGemini {
- return SchedulerModeMixed
- }
- return SchedulerModeSingle
- }
- func (s *SchedulerSnapshotService) guardFallback(ctx context.Context) error {
- if s.cfg == nil || s.cfg.Gateway.Scheduling.DbFallbackEnabled {
- if s.fallbackLimit == nil || s.fallbackLimit.Allow() {
- return nil
- }
- return ErrSchedulerFallbackLimited
- }
- return ErrSchedulerCacheNotReady
- }
- func (s *SchedulerSnapshotService) withFallbackTimeout(ctx context.Context) (context.Context, context.CancelFunc) {
- if s.cfg == nil || s.cfg.Gateway.Scheduling.DbFallbackTimeoutSeconds <= 0 {
- return context.WithCancel(ctx)
- }
- timeout := time.Duration(s.cfg.Gateway.Scheduling.DbFallbackTimeoutSeconds) * time.Second
- if deadline, ok := ctx.Deadline(); ok {
- remaining := time.Until(deadline)
- if remaining <= 0 {
- return context.WithCancel(ctx)
- }
- if remaining < timeout {
- timeout = remaining
- }
- }
- return context.WithTimeout(ctx, timeout)
- }
- func (s *SchedulerSnapshotService) isRunModeSimple() bool {
- return s.cfg != nil && s.cfg.RunMode == config.RunModeSimple
- }
- func (s *SchedulerSnapshotService) outboxPollInterval() time.Duration {
- if s.cfg == nil {
- return time.Second
- }
- sec := s.cfg.Gateway.Scheduling.OutboxPollIntervalSeconds
- if sec <= 0 {
- return time.Second
- }
- return time.Duration(sec) * time.Second
- }
- func (s *SchedulerSnapshotService) fullRebuildInterval() time.Duration {
- if s.cfg == nil {
- return 0
- }
- sec := s.cfg.Gateway.Scheduling.FullRebuildIntervalSeconds
- if sec <= 0 {
- return 0
- }
- return time.Duration(sec) * time.Second
- }
- func dedupeBuckets(in []SchedulerBucket) []SchedulerBucket {
- seen := make(map[string]struct{}, len(in))
- out := make([]SchedulerBucket, 0, len(in))
- for _, bucket := range in {
- key := bucket.String()
- if _, ok := seen[key]; ok {
- continue
- }
- seen[key] = struct{}{}
- out = append(out, bucket)
- }
- return out
- }
- func derefAccounts(accounts []*Account) []Account {
- if len(accounts) == 0 {
- return []Account{}
- }
- out := make([]Account, 0, len(accounts))
- for _, account := range accounts {
- if account == nil {
- continue
- }
- out = append(out, *account)
- }
- return out
- }
- func parseInt64Slice(value any) []int64 {
- raw, ok := value.([]any)
- if !ok {
- return nil
- }
- out := make([]int64, 0, len(raw))
- for _, item := range raw {
- if v, ok := toInt64(item); ok && v > 0 {
- out = append(out, v)
- }
- }
- return out
- }
- func toInt64(value any) (int64, bool) {
- switch v := value.(type) {
- case float64:
- return int64(v), true
- case int64:
- return v, true
- case int:
- return int64(v), true
- case json.Number:
- parsed, err := strconv.ParseInt(v.String(), 10, 64)
- return parsed, err == nil
- default:
- return 0, false
- }
- }
- type fallbackLimiter struct {
- maxQPS int
- mu sync.Mutex
- window time.Time
- count int
- }
- func newFallbackLimiter(maxQPS int) *fallbackLimiter {
- if maxQPS <= 0 {
- return nil
- }
- return &fallbackLimiter{
- maxQPS: maxQPS,
- window: time.Now(),
- }
- }
- func (l *fallbackLimiter) Allow() bool {
- if l == nil || l.maxQPS <= 0 {
- return true
- }
- l.mu.Lock()
- defer l.mu.Unlock()
- now := time.Now()
- if now.Sub(l.window) >= time.Second {
- l.window = now
- l.count = 0
- }
- if l.count >= l.maxQPS {
- return false
- }
- l.count++
- return true
- }
|