清理全量重建中的过期调度桶

This commit is contained in:
jjaw
2026-07-15 04:23:06 +08:00
parent 31a0efc747
commit 8edad3a455
4 changed files with 901 additions and 90 deletions
@@ -0,0 +1,666 @@
//go:build unit
package service
import (
"context"
"errors"
"sync"
"testing"
"time"
"github.com/Wei-Shaw/sub2api/internal/config"
"github.com/stretchr/testify/require"
)
type fullRebuildLifecycleCache struct {
*groupLifecycleTestCache
mu sync.Mutex
captureAttempts []SchedulerBucket
captureErrors map[string]error
lockBusyOnce map[string]bool
watermark int64
watermarkWrites []int64
}
func newFullRebuildLifecycleCache(buckets ...SchedulerBucket) *fullRebuildLifecycleCache {
return &fullRebuildLifecycleCache{
groupLifecycleTestCache: newGroupLifecycleTestCache(buckets...),
captureErrors: make(map[string]error),
lockBusyOnce: make(map[string]bool),
}
}
func (c *fullRebuildLifecycleCache) CaptureBucketWriteToken(ctx context.Context, bucket SchedulerBucket) (SchedulerBucketWriteToken, error) {
c.mu.Lock()
c.captureAttempts = append(c.captureAttempts, bucket)
err := c.captureErrors[bucket.String()]
c.mu.Unlock()
if err != nil {
return SchedulerBucketWriteToken{}, err
}
return c.retirementRaceCache.CaptureBucketWriteToken(ctx, bucket)
}
func (c *fullRebuildLifecycleCache) ListBuckets(ctx context.Context) ([]SchedulerBucket, error) {
buckets, err := c.groupLifecycleTestCache.ListBuckets(ctx)
if err != nil {
return nil, err
}
c.retirementRaceCache.mu.Lock()
defer c.retirementRaceCache.mu.Unlock()
registered := make([]SchedulerBucket, 0, len(buckets))
for _, bucket := range buckets {
if !c.retirementRaceCache.retired[bucket.String()] {
registered = append(registered, bucket)
}
}
return registered, nil
}
func (c *fullRebuildLifecycleCache) TryLockBucket(ctx context.Context, bucket SchedulerBucket, ttl time.Duration) (bool, error) {
c.mu.Lock()
busy := c.lockBusyOnce[bucket.String()]
if busy {
delete(c.lockBusyOnce, bucket.String())
}
c.mu.Unlock()
if busy {
return false, nil
}
return c.groupLifecycleTestCache.TryLockBucket(ctx, bucket, ttl)
}
func (c *fullRebuildLifecycleCache) GetOutboxWatermark(context.Context) (int64, error) {
c.mu.Lock()
defer c.mu.Unlock()
return c.watermark, nil
}
func (c *fullRebuildLifecycleCache) SetOutboxWatermark(_ context.Context, id int64) error {
c.mu.Lock()
defer c.mu.Unlock()
c.watermark = id
c.watermarkWrites = append(c.watermarkWrites, id)
return nil
}
func (c *fullRebuildLifecycleCache) captureAttemptCount() int {
c.mu.Lock()
defer c.mu.Unlock()
return len(c.captureAttempts)
}
func (c *fullRebuildLifecycleCache) currentWatermark() int64 {
c.mu.Lock()
defer c.mu.Unlock()
return c.watermark
}
func (c *fullRebuildLifecycleCache) totalSetAttempts() int {
c.retirementRaceCache.mu.Lock()
defer c.retirementRaceCache.mu.Unlock()
var total int
for _, attempts := range c.retirementRaceCache.setAttempts {
total += attempts
}
return total
}
type fullRebuildLifecycleGroupRepo struct {
GroupRepository
mu sync.Mutex
activeIDs []int64
activeIDsErr error
listActiveErr error
fresh map[int64]*Group
freshErr map[int64]error
activeIDCalls int
listActiveCalls int
freshCalls []int64
}
func (r *fullRebuildLifecycleGroupRepo) ListActiveIDs(context.Context) ([]int64, error) {
r.mu.Lock()
defer r.mu.Unlock()
r.activeIDCalls++
return append([]int64(nil), r.activeIDs...), r.activeIDsErr
}
func (r *fullRebuildLifecycleGroupRepo) ListActive(context.Context) ([]Group, error) {
r.mu.Lock()
defer r.mu.Unlock()
r.listActiveCalls++
return nil, r.listActiveErr
}
func (r *fullRebuildLifecycleGroupRepo) GetByIDLite(_ context.Context, id int64) (*Group, error) {
r.mu.Lock()
defer r.mu.Unlock()
r.freshCalls = append(r.freshCalls, id)
if err := r.freshErr[id]; err != nil {
return nil, err
}
group := r.fresh[id]
if group == nil {
return nil, ErrGroupNotFound
}
copyGroup := *group
return &copyGroup, nil
}
func (r *fullRebuildLifecycleGroupRepo) stats() (activeIDs, listActive int, fresh []int64) {
r.mu.Lock()
defer r.mu.Unlock()
return r.activeIDCalls, r.listActiveCalls, append([]int64(nil), r.freshCalls...)
}
type fullRebuildFallbackGroupRepo struct {
GroupRepository
mu sync.Mutex
groups []Group
err error
listCalls int
}
func (r *fullRebuildFallbackGroupRepo) ListActive(context.Context) ([]Group, error) {
r.mu.Lock()
defer r.mu.Unlock()
r.listCalls++
return append([]Group(nil), r.groups...), r.err
}
type fullRebuildAccountCall struct {
groupID int64
platform string
}
type fullRebuildAccountRepo struct {
AccountRepository
mu sync.Mutex
calls []fullRebuildAccountCall
beforeFirst func()
once sync.Once
}
func (r *fullRebuildAccountRepo) record(groupID int64, platform string) ([]Account, error) {
r.mu.Lock()
r.calls = append(r.calls, fullRebuildAccountCall{groupID: groupID, platform: platform})
beforeFirst := r.beforeFirst
r.mu.Unlock()
if beforeFirst != nil {
r.once.Do(beforeFirst)
}
return []Account{{ID: 1, Platform: platform, Status: StatusActive, Schedulable: true}}, nil
}
func (r *fullRebuildAccountRepo) ListSchedulableByGroupIDAndPlatform(_ context.Context, groupID int64, platform string) ([]Account, error) {
return r.record(groupID, platform)
}
func (r *fullRebuildAccountRepo) ListSchedulableByGroupIDAndPlatforms(_ context.Context, groupID int64, platforms []string) ([]Account, error) {
return r.record(groupID, firstPlatform(platforms))
}
func (r *fullRebuildAccountRepo) ListSchedulableUngroupedByPlatform(_ context.Context, platform string) ([]Account, error) {
return r.record(0, platform)
}
func (r *fullRebuildAccountRepo) ListSchedulableUngroupedByPlatforms(_ context.Context, platforms []string) ([]Account, error) {
return r.record(0, firstPlatform(platforms))
}
func (r *fullRebuildAccountRepo) ListSchedulableByPlatform(_ context.Context, platform string) ([]Account, error) {
return r.record(0, platform)
}
func (r *fullRebuildAccountRepo) ListSchedulableByPlatforms(_ context.Context, platforms []string) ([]Account, error) {
return r.record(0, firstPlatform(platforms))
}
func (r *fullRebuildAccountRepo) callCount() int {
r.mu.Lock()
defer r.mu.Unlock()
return len(r.calls)
}
func (r *fullRebuildAccountRepo) groupCallCount(groupID int64) int {
r.mu.Lock()
defer r.mu.Unlock()
var count int
for _, call := range r.calls {
if call.groupID == groupID {
count++
}
}
return count
}
func firstPlatform(platforms []string) string {
if len(platforms) == 0 {
return ""
}
return platforms[0]
}
func newFullRebuildLifecycleService(
cache SchedulerCache,
outbox SchedulerOutboxRepository,
accounts AccountRepository,
groups GroupRepository,
runMode string,
) *SchedulerSnapshotService {
return NewSchedulerSnapshotService(cache, outbox, accounts, groups, &config.Config{RunMode: runMode})
}
func TestSchedulerFullRebuildActiveTombstoneDoesNotBlockFollowingGroupEvent(t *testing.T) {
const groupID int64 = 101
canonical := schedulerBucketsForGroup(groupID)
cache := newFullRebuildLifecycleCache()
require.NoError(t, cache.retirementRaceCache.RetireBucket(context.Background(), canonical[0]))
groups := &fullRebuildLifecycleGroupRepo{
activeIDs: []int64{groupID},
fresh: map[int64]*Group{
groupID: {ID: groupID, Status: StatusActive, Hydrated: true},
},
freshErr: make(map[int64]error),
}
accounts := &fullRebuildAccountRepo{}
outbox := &outboxCleanupRepo{events: []SchedulerOutboxEvent{
{ID: 1, EventType: SchedulerOutboxEventFullRebuild},
{ID: 2, EventType: SchedulerOutboxEventGroupChanged, GroupID: ptrInt64(groupID)},
}}
svc := newFullRebuildLifecycleService(cache, outbox, accounts, groups, config.RunModeStandard)
svc.pollOutbox()
require.Equal(t, int64(2), cache.currentWatermark())
cache.mu.Lock()
require.Equal(t, []int64{2}, cache.watermarkWrites)
cache.mu.Unlock()
activeCalls, fallbackCalls, freshCalls := groups.stats()
require.Equal(t, 1, activeCalls)
require.Zero(t, fallbackCalls)
require.Equal(t, []int64{groupID, groupID}, freshCalls)
require.Len(t, cache.tokens(), 24, "full rebuild and the following group event must each run fresh authority")
_, reopenHeld := cache.lifecycleMutationLeaseStates()
require.Len(t, reopenHeld, 24)
for _, held := range reopenHeld {
require.True(t, held)
}
require.Equal(t, 36, accounts.callCount())
}
func TestSchedulerFullRebuildGlobalReadErrorsFailBeforeMutationOrDB(t *testing.T) {
t.Run("list buckets", func(t *testing.T) {
cache := newFullRebuildLifecycleCache()
cache.listErr = errors.New("registry failed")
groups := &fullRebuildLifecycleGroupRepo{fresh: make(map[int64]*Group), freshErr: make(map[int64]error)}
accounts := &fullRebuildAccountRepo{}
svc := newFullRebuildLifecycleService(cache, nil, accounts, groups, config.RunModeStandard)
err := svc.rebuildFullSnapshot(context.Background(), "test")
require.ErrorIs(t, err, cache.listErr)
activeCalls, fallbackCalls, freshCalls := groups.stats()
require.Zero(t, activeCalls)
require.Zero(t, fallbackCalls)
require.Empty(t, freshCalls)
requireFullRebuildNoMutationOrDB(t, cache, accounts)
})
t.Run("list active ids without fallback", func(t *testing.T) {
cache := newFullRebuildLifecycleCache()
groups := &fullRebuildLifecycleGroupRepo{
activeIDsErr: errors.New("active ids failed"),
listActiveErr: errors.New("fallback must not run"),
fresh: make(map[int64]*Group),
freshErr: make(map[int64]error),
}
accounts := &fullRebuildAccountRepo{}
svc := newFullRebuildLifecycleService(cache, nil, accounts, groups, config.RunModeStandard)
err := svc.rebuildFullSnapshot(context.Background(), "test")
require.ErrorIs(t, err, groups.activeIDsErr)
activeCalls, fallbackCalls, freshCalls := groups.stats()
require.Equal(t, 1, activeCalls)
require.Zero(t, fallbackCalls)
require.Empty(t, freshCalls)
requireFullRebuildNoMutationOrDB(t, cache, accounts)
})
t.Run("list active fallback", func(t *testing.T) {
cache := newFullRebuildLifecycleCache()
groups := &fullRebuildFallbackGroupRepo{err: errors.New("active groups failed")}
accounts := &fullRebuildAccountRepo{}
svc := newFullRebuildLifecycleService(cache, nil, accounts, groups, config.RunModeStandard)
err := svc.rebuildFullSnapshot(context.Background(), "test")
require.ErrorIs(t, err, groups.err)
groups.mu.Lock()
require.Equal(t, 1, groups.listCalls)
groups.mu.Unlock()
requireFullRebuildNoMutationOrDB(t, cache, accounts)
})
}
func TestSchedulerFullRebuildFreshActivePreparesEveryTokenBeforeFirstDB(t *testing.T) {
const groupID int64 = 102
historical := SchedulerBucket{GroupID: groupID, Platform: "legacy", Mode: "unknown"}
cache := newFullRebuildLifecycleCache(historical)
groups := &fullRebuildLifecycleGroupRepo{
fresh: map[int64]*Group{
groupID: {ID: groupID, Status: StatusActive, Hydrated: true},
},
freshErr: make(map[int64]error),
}
accounts := &fullRebuildAccountRepo{}
var capturesAtFirstDB int
accounts.beforeFirst = func() {
capturesAtFirstDB = cache.captureAttemptCount()
held, reopenCount := cache.leaseHeldAndTokenCount()
require.False(t, held)
require.Equal(t, 12, reopenCount)
require.Equal(t, 13, capturesAtFirstDB, "C(0) and the historical bucket must be captured before DB")
}
svc := newFullRebuildLifecycleService(cache, nil, accounts, groups, config.RunModeStandard)
require.NoError(t, svc.rebuildFullSnapshot(context.Background(), "test"))
require.Equal(t, capturesAtFirstDB, cache.captureAttemptCount())
require.Equal(t, 25, accounts.callCount())
require.Equal(t, 13, accounts.groupCallCount(groupID))
_, historicalPublished := cache.counts(historical)
require.Equal(t, 1, historicalPublished)
activeCalls, fallbackCalls, freshCalls := groups.stats()
require.Equal(t, 1, activeCalls)
require.Zero(t, fallbackCalls)
require.Equal(t, []int64{groupID}, freshCalls)
_, _, listCalls := cache.lifecycleCounts()
require.Equal(t, 1, listCalls, "known Rg must prevent a second registry read")
}
func TestSchedulerFullRebuildOrdinaryCaptureErrorReturnsBeforeFirstDB(t *testing.T) {
first := SchedulerBucket{GroupID: 0, Platform: "legacy-a", Mode: "unknown"}
last := SchedulerBucket{GroupID: 0, Platform: "legacy-b", Mode: "unknown"}
cache := newFullRebuildLifecycleCache(first, last)
wantErr := errors.New("capture failed")
cache.captureErrors[last.String()] = wantErr
groups := &fullRebuildLifecycleGroupRepo{fresh: make(map[int64]*Group), freshErr: make(map[int64]error)}
accounts := &fullRebuildAccountRepo{}
svc := newFullRebuildLifecycleService(cache, nil, accounts, groups, config.RunModeStandard)
err := svc.rebuildFullSnapshot(context.Background(), "test")
require.ErrorIs(t, err, wantErr)
require.Equal(t, 14, cache.captureAttemptCount(), "all canonical and ordinary captures must be attempted before returning")
require.Zero(t, accounts.callCount())
require.Zero(t, cache.totalSetAttempts())
}
func TestSchedulerFullRebuildPreservesGroupZeroActiveHistoricalAndInvalidRegistryBuckets(t *testing.T) {
const groupID int64 = 103
groupZeroHistorical := SchedulerBucket{GroupID: 0, Platform: "legacy-zero", Mode: "unknown"}
activeHistorical := SchedulerBucket{GroupID: groupID, Platform: "legacy-active", Mode: "unknown"}
activeForced := SchedulerBucket{GroupID: groupID, Platform: PlatformOpenAI, Mode: SchedulerModeForced}
activeMixed := SchedulerBucket{GroupID: groupID, Platform: PlatformAnthropic, Mode: SchedulerModeMixed}
invalidHistorical := SchedulerBucket{GroupID: -7, Platform: "legacy-invalid", Mode: "unknown"}
cache := newFullRebuildLifecycleCache(groupZeroHistorical, activeHistorical, activeForced, activeMixed, invalidHistorical)
groups := &fullRebuildFallbackGroupRepo{groups: []Group{{ID: groupID, Status: StatusActive}}}
accounts := &fullRebuildAccountRepo{}
svc := newFullRebuildLifecycleService(cache, nil, accounts, groups, config.RunModeStandard)
require.NoError(t, svc.rebuildFullSnapshot(context.Background(), "test"))
require.Equal(t, 27, cache.captureAttemptCount())
require.Equal(t, 27, accounts.callCount())
groups.mu.Lock()
require.Equal(t, 1, groups.listCalls)
groups.mu.Unlock()
require.Empty(t, cache.retiredBuckets())
require.Empty(t, cache.tokens())
for _, bucket := range []SchedulerBucket{groupZeroHistorical, activeHistorical, activeForced, activeMixed, invalidHistorical} {
_, published := cache.counts(bucket)
require.Equal(t, 1, published, bucket.String())
}
}
func TestSchedulerFullRebuildActiveTombstoneFreshInactiveOrMissingFiltersAllGroupTasks(t *testing.T) {
for _, tc := range []struct {
name string
group *Group
}{
{name: "inactive", group: &Group{ID: 104, Status: StatusDisabled, Hydrated: true}},
{name: "missing"},
} {
t.Run(tc.name, func(t *testing.T) {
const groupID int64 = 104
canonical := schedulerBucketsForGroup(groupID)
historical := SchedulerBucket{GroupID: groupID, Platform: "legacy", Mode: "unknown"}
cache := newFullRebuildLifecycleCache(historical)
require.NoError(t, cache.retirementRaceCache.RetireBucket(context.Background(), canonical[5]))
groups := &fullRebuildLifecycleGroupRepo{
activeIDs: []int64{groupID},
fresh: make(map[int64]*Group),
freshErr: make(map[int64]error),
}
if tc.group != nil {
groups.fresh[groupID] = tc.group
}
accounts := &fullRebuildAccountRepo{}
svc := newFullRebuildLifecycleService(cache, nil, accounts, groups, config.RunModeStandard)
require.NoError(t, svc.rebuildFullSnapshot(context.Background(), "test"))
require.Zero(t, accounts.groupCallCount(groupID))
require.Equal(t, 12, accounts.groupCallCount(0))
require.Empty(t, cache.tokens())
require.Equal(t, bucketStrings(append(canonical, historical)), bucketStrings(cache.retiredBuckets()))
for _, bucket := range append(canonical, historical) {
attempts, published := cache.counts(bucket)
require.Zero(t, attempts, bucket.String())
require.Zero(t, published, bucket.String())
}
_, _, listCalls := cache.lifecycleCounts()
require.Equal(t, 1, listCalls)
})
}
}
func TestSchedulerFullRebuildStaleCandidatesAreSortedAndNeverReopenedAcrossRounds(t *testing.T) {
historical := []SchedulerBucket{
{GroupID: 3, Platform: "legacy", Mode: "unknown"},
{GroupID: 1, Platform: "legacy", Mode: "unknown"},
{GroupID: 2, Platform: "legacy", Mode: "unknown"},
}
cache := newFullRebuildLifecycleCache(historical...)
groups := &fullRebuildLifecycleGroupRepo{
fresh: map[int64]*Group{
1: {ID: 1, Status: StatusDisabled, Hydrated: true},
3: {ID: 3, Status: StatusDisabled, Hydrated: true},
},
freshErr: make(map[int64]error),
}
accounts := &fullRebuildAccountRepo{}
svc := newFullRebuildLifecycleService(cache, nil, accounts, groups, config.RunModeStandard)
require.NoError(t, svc.rebuildFullSnapshot(context.Background(), "first"))
retireCallsAfterFirst := len(cache.retiredBuckets())
require.NoError(t, svc.rebuildFullSnapshot(context.Background(), "second"))
_, _, freshCalls := groups.stats()
require.Equal(t, []int64{1, 2, 3}, freshCalls)
require.Equal(t, retireCallsAfterFirst, len(cache.retiredBuckets()))
require.Empty(t, cache.tokens())
require.Zero(t, accounts.groupCallCount(1))
require.Zero(t, accounts.groupCallCount(2))
require.Zero(t, accounts.groupCallCount(3))
_, _, listCalls := cache.lifecycleCounts()
require.Equal(t, 2, listCalls, "each round must use only its one global registry snapshot")
for _, bucket := range historical {
require.Contains(t, bucketStrings(cache.retiredBuckets()), bucket.String())
}
}
func TestSchedulerFullRebuildPartialLifecycleFailureReturnsBeforeDBAndRetries(t *testing.T) {
historical := []SchedulerBucket{
{GroupID: 1, Platform: "legacy", Mode: "unknown"},
{GroupID: 2, Platform: "legacy", Mode: "unknown"},
{GroupID: 3, Platform: "legacy", Mode: "unknown"},
}
wantErr := errors.New("fresh group query failed")
cache := newFullRebuildLifecycleCache(historical...)
groups := &fullRebuildLifecycleGroupRepo{
fresh: map[int64]*Group{
1: {ID: 1, Status: StatusDisabled, Hydrated: true},
3: {ID: 3, Status: StatusDisabled, Hydrated: true},
},
freshErr: map[int64]error{2: wantErr},
}
accounts := &fullRebuildAccountRepo{}
svc := newFullRebuildLifecycleService(cache, nil, accounts, groups, config.RunModeStandard)
err := svc.triggerFullRebuild("first")
require.ErrorIs(t, err, wantErr)
_, _, freshCalls := groups.stats()
require.Equal(t, []int64{1, 2}, freshCalls)
require.Zero(t, accounts.callCount())
require.Zero(t, cache.totalSetAttempts())
require.Equal(t, 13, len(cache.retiredBuckets()))
groups.mu.Lock()
delete(groups.freshErr, 2)
groups.fresh[2] = &Group{ID: 2, Status: StatusDisabled, Hydrated: true}
groups.mu.Unlock()
require.NoError(t, svc.triggerFullRebuild("retry"))
_, _, freshCalls = groups.stats()
require.Equal(t, []int64{1, 2, 2, 3}, freshCalls)
require.Equal(t, 39, len(cache.retiredBuckets()))
require.Equal(t, 12, accounts.callCount())
require.Empty(t, cache.tokens())
}
func TestSchedulerFullRebuildActiveTombstoneLazyRecoveryDiscardsPartialCaptureTasks(t *testing.T) {
const groupID int64 = 105
canonical := schedulerBucketsForGroup(groupID)
historical := SchedulerBucket{GroupID: groupID, Platform: "legacy", Mode: "unknown"}
cache := newFullRebuildLifecycleCache(canonical[0], canonical[4], historical)
require.NoError(t, cache.retirementRaceCache.RetireBucket(context.Background(), canonical[5]))
groups := &fullRebuildLifecycleGroupRepo{
activeIDs: []int64{groupID},
fresh: map[int64]*Group{
groupID: {ID: groupID, Status: StatusActive, Hydrated: true},
},
freshErr: make(map[int64]error),
}
accounts := &fullRebuildAccountRepo{}
var capturesAtFirstDB int
accounts.beforeFirst = func() {
capturesAtFirstDB = cache.captureAttemptCount()
held, reopenCount := cache.leaseHeldAndTokenCount()
require.False(t, held)
require.Equal(t, 12, reopenCount)
require.Equal(t, 19, capturesAtFirstDB)
}
svc := newFullRebuildLifecycleService(cache, nil, accounts, groups, config.RunModeStandard)
require.NoError(t, svc.rebuildFullSnapshot(context.Background(), "test"))
require.Equal(t, capturesAtFirstDB, cache.captureAttemptCount())
require.Equal(t, 25, accounts.callCount())
for _, bucket := range canonical {
attempts, published := cache.counts(bucket)
require.Equal(t, 1, attempts, "discarded pre-recovery tokens must never publish: %s", bucket.String())
require.Equal(t, 1, published, bucket.String())
}
_, published := cache.counts(historical)
require.Equal(t, 1, published)
_, _, listCalls := cache.lifecycleCounts()
require.Equal(t, 1, listCalls)
}
func TestSchedulerFullRebuildSimpleModePreservesRegistryWithoutLifecycleAuthority(t *testing.T) {
registered := []SchedulerBucket{
{GroupID: 0, Platform: "legacy-zero", Mode: "unknown"},
{GroupID: 106, Platform: "legacy-positive", Mode: "unknown"},
{GroupID: -8, Platform: "legacy-negative", Mode: "unknown"},
}
cache := newFullRebuildLifecycleCache(registered...)
groups := &fullRebuildLifecycleGroupRepo{
activeIDsErr: errors.New("simple mode must not query groups"),
fresh: make(map[int64]*Group),
freshErr: make(map[int64]error),
}
accounts := &fullRebuildAccountRepo{}
svc := newFullRebuildLifecycleService(cache, nil, accounts, groups, config.RunModeSimple)
require.NoError(t, svc.rebuildFullSnapshot(context.Background(), "test"))
activeCalls, fallbackCalls, freshCalls := groups.stats()
require.Zero(t, activeCalls)
require.Zero(t, fallbackCalls)
require.Empty(t, freshCalls)
require.Equal(t, 15, cache.captureAttemptCount())
require.Equal(t, 15, accounts.callCount())
require.Equal(t, 15, accounts.groupCallCount(0))
require.Empty(t, cache.retiredBuckets())
require.Empty(t, cache.tokens())
for _, bucket := range registered {
_, published := cache.counts(bucket)
require.Equal(t, 1, published, bucket.String())
}
}
func TestSchedulerFullRebuildFreshReopenLockBusyRetriesWithoutBlockingOrdinaryTasks(t *testing.T) {
const groupID int64 = 107
canonical := schedulerBucketsForGroup(groupID)
cache := newFullRebuildLifecycleCache()
require.NoError(t, cache.retirementRaceCache.RetireBucket(context.Background(), canonical[0]))
cache.lockBusyOnce[canonical[0].String()] = true
groups := &fullRebuildLifecycleGroupRepo{
activeIDs: []int64{groupID},
fresh: map[int64]*Group{
groupID: {ID: groupID, Status: StatusActive, Hydrated: true},
},
freshErr: make(map[int64]error),
}
accounts := &fullRebuildAccountRepo{}
outbox := &outboxCleanupRepo{events: []SchedulerOutboxEvent{{ID: 1, EventType: SchedulerOutboxEventFullRebuild}}}
svc := newFullRebuildLifecycleService(cache, outbox, accounts, groups, config.RunModeStandard)
svc.pollOutbox()
require.Zero(t, cache.currentWatermark())
_, groupZeroPublished := cache.counts(schedulerCanonicalBuckets(0)[0])
require.Equal(t, 1, groupZeroPublished, "ordinary tasks must still run when one strict Reopen task is busy")
require.Equal(t, 23, accounts.callCount())
svc.pollOutbox()
require.Equal(t, int64(1), cache.currentWatermark())
require.Equal(t, 47, accounts.callCount())
_, busyBucketPublished := cache.counts(canonical[0])
require.Equal(t, 1, busyBucketPublished)
activeCalls, fallbackCalls, freshCalls := groups.stats()
require.Equal(t, 2, activeCalls)
require.Zero(t, fallbackCalls)
require.Equal(t, []int64{groupID}, freshCalls)
}
func TestSchedulerFullRebuildOrdinaryLockBusyKeepsExistingSkipSemantics(t *testing.T) {
busyBucket := schedulerCanonicalBuckets(0)[0]
cache := newFullRebuildLifecycleCache()
cache.lockBusyOnce[busyBucket.String()] = true
groups := &fullRebuildLifecycleGroupRepo{fresh: make(map[int64]*Group), freshErr: make(map[int64]error)}
accounts := &fullRebuildAccountRepo{}
svc := newFullRebuildLifecycleService(cache, nil, accounts, groups, config.RunModeStandard)
require.NoError(t, svc.rebuildFullSnapshot(context.Background(), "test"))
require.Equal(t, 11, accounts.callCount())
attempts, published := cache.counts(busyBucket)
require.Zero(t, attempts)
require.Zero(t, published)
}
func requireFullRebuildNoMutationOrDB(t *testing.T, cache *fullRebuildLifecycleCache, accounts *fullRebuildAccountRepo) {
t.Helper()
require.Zero(t, cache.captureAttemptCount())
require.Empty(t, cache.retiredBuckets())
require.Empty(t, cache.tokens())
require.Zero(t, accounts.callCount())
require.Zero(t, cache.totalSetAttempts())
}
@@ -17,6 +17,7 @@ type schedulerFullRebuildTestCache struct {
mu sync.Mutex
listErr error
listCalls int
captures int
lockCalls int
}
@@ -35,6 +36,9 @@ func (c *schedulerFullRebuildTestCache) TryLockBucket(context.Context, Scheduler
}
func (c *schedulerFullRebuildTestCache) CaptureBucketWriteToken(_ context.Context, bucket SchedulerBucket) (SchedulerBucketWriteToken, error) {
c.mu.Lock()
c.captures++
c.mu.Unlock()
return SchedulerBucketWriteToken{Bucket: bucket, Epoch: 1}, nil
}
@@ -129,7 +133,7 @@ func TestSchedulerSnapshotServiceFullRebuildRunsAgainForSequentialRequest(t *tes
require.Equal(t, requested, completed)
}
func TestSchedulerSnapshotServiceInitialFullRebuildFallsBackWhenListBucketsFails(t *testing.T) {
func TestSchedulerSnapshotServiceInitialFullRebuildFailsClosedWhenListBucketsFails(t *testing.T) {
cache := &schedulerFullRebuildTestCache{listErr: errors.New("list buckets failed")}
svc := NewSchedulerSnapshotService(cache, nil, nil, nil, nil)
@@ -137,13 +141,18 @@ func TestSchedulerSnapshotServiceInitialFullRebuildFallsBackWhenListBucketsFails
cache.mu.Lock()
listCalls := cache.listCalls
captures := cache.captures
lockCalls := cache.lockCalls
cache.mu.Unlock()
require.Equal(t, 1, listCalls)
require.Positive(t, lockCalls, "startup should rebuild default buckets after ListBuckets fails")
require.Zero(t, captures)
require.Zero(t, lockCalls)
requested, completed := schedulerFullRebuildState(svc)
require.EqualValues(t, 1, requested)
require.Equal(t, requested, completed)
svc.fullRebuildStateMu.Lock()
require.ErrorIs(t, svc.fullRebuildLastErr, cache.listErr)
svc.fullRebuildStateMu.Unlock()
}
func schedulerFullRebuildState(svc *SchedulerSnapshotService) (requested uint64, completed uint64) {
@@ -4,7 +4,6 @@ package service
import (
"context"
"errors"
"sync"
"testing"
"time"
@@ -161,7 +160,7 @@ func TestSchedulerFullRebuildCapturesAllRegistryTokensBeforeDBLoad(t *testing.T)
return []Account{{ID: 6101, Platform: PlatformOpenAI, Status: StatusActive, Schedulable: true}}, nil
},
}
svc := NewSchedulerSnapshotService(cache, nil, repo, nil, &config.Config{
svc := NewSchedulerSnapshotService(cache, nil, repo, &retirementGroupRepo{groups: []Group{{ID: 61, Status: StatusActive}}}, &config.Config{
RunMode: config.RunModeStandard,
Gateway: config.GatewayConfig{Scheduling: config.GatewaySchedulingConfig{
DbFallbackEnabled: true,
@@ -177,7 +176,7 @@ func TestSchedulerFullRebuildCapturesAllRegistryTokensBeforeDBLoad(t *testing.T)
}
captures, reopens := cache.captureAndReopenCounts()
require.Equal(t, 2, captures, "all registry tokens must be captured before the first DB load")
require.Equal(t, 24, captures, "group0 and active-group canonical tokens must be captured before the first DB load")
require.Zero(t, reopens)
require.NoError(t, cache.RetireBucket(context.Background(), queued))
_, err := cache.ReopenBucket(context.Background(), queued)
@@ -262,28 +261,3 @@ func TestSchedulerFallbackReturnsDBAccountsWhenBucketRetired(t *testing.T) {
require.Zero(t, setAttempts)
require.Zero(t, published)
}
func TestSchedulerDefaultBucketsUseCaptureAndListActiveFailureKeepsGroupZero(t *testing.T) {
cache := newRetirementRaceCache()
svc := NewSchedulerSnapshotService(
cache,
nil,
nil,
&retirementGroupRepo{err: errors.New("list active failed")},
testConfig(),
)
buckets, err := svc.defaultBuckets(context.Background())
require.NoError(t, err)
require.NotEmpty(t, buckets)
for _, bucket := range buckets {
require.Zero(t, bucket.GroupID)
}
tasks, err := svc.prepareBucketWriteTasks(context.Background(), buckets)
require.NoError(t, err)
require.Len(t, tasks, len(buckets))
captures, reopens := cache.captureAndReopenCounts()
require.Equal(t, len(buckets), captures)
require.Zero(t, reopens)
}
@@ -6,6 +6,7 @@ import (
"errors"
"fmt"
"log/slog"
"sort"
"strconv"
"sync"
"time"
@@ -47,6 +48,10 @@ type schedulerGroupLifecyclePlan struct {
tasks []schedulerBucketWriteTask
}
type schedulerActiveGroupIDLister interface {
ListActiveIDs(ctx context.Context) ([]int64, error)
}
type SchedulerSnapshotService struct {
cache SchedulerCache
outboxRepo SchedulerOutboxRepository
@@ -225,18 +230,7 @@ func (s *SchedulerSnapshotService) runInitialRebuild() {
_ = s.coalesceFullRebuild(func() error {
ctx, cancel := context.WithTimeout(context.Background(), 2*time.Minute)
defer cancel()
buckets, err := s.cache.ListBuckets(ctx)
if err != nil {
logger.LegacyPrintf("service.scheduler_snapshot", "[Scheduler] list buckets failed: %v", err)
}
if len(buckets) == 0 {
buckets, err = s.defaultBuckets(ctx)
if err != nil {
logger.LegacyPrintf("service.scheduler_snapshot", "[Scheduler] default buckets failed: %v", err)
return err
}
}
if err := s.rebuildBuckets(ctx, buckets, "startup"); err != nil {
if err := s.rebuildFullSnapshot(ctx, "startup"); err != nil {
logger.LegacyPrintf("service.scheduler_snapshot", "[Scheduler] rebuild startup failed: %v", err)
return err
}
@@ -665,6 +659,10 @@ func schedulerBucketsForGroup(groupID int64) []SchedulerBucket {
if groupID <= 0 {
return nil
}
return schedulerCanonicalBuckets(groupID)
}
func schedulerCanonicalBuckets(groupID int64) []SchedulerBucket {
buckets := make([]SchedulerBucket, 0, 12)
for _, platform := range schedulerSnapshotPlatforms() {
buckets = append(buckets,
@@ -718,10 +716,8 @@ func (s *SchedulerSnapshotService) bucketsForPlatform(platform string, groupIDs
func (s *SchedulerSnapshotService) rebuildBuckets(ctx context.Context, buckets []SchedulerBucket, reason string) error {
tasks, firstErr := s.prepareBucketWriteTasks(ctx, buckets)
for _, task := range tasks {
if err := s.rebuildBucketWithToken(ctx, task, reason); err != nil && firstErr == nil {
firstErr = err
}
if err := s.rebuildPreparedBucketTasks(ctx, tasks, reason, false); err != nil && firstErr == nil {
firstErr = err
}
return firstErr
}
@@ -748,8 +744,14 @@ func (s *SchedulerSnapshotService) prepareBucketWriteTasks(ctx context.Context,
return tasks, firstErr
}
func (s *SchedulerSnapshotService) rebuildBucketWithToken(ctx context.Context, task schedulerBucketWriteTask, reason string) error {
return s.rebuildBucketWithTokenPolicy(ctx, task, reason, false)
func (s *SchedulerSnapshotService) rebuildPreparedBucketTasks(ctx context.Context, tasks []schedulerBucketWriteTask, reason string, strict bool) error {
var firstErr error
for _, task := range tasks {
if err := s.rebuildBucketWithTokenPolicy(ctx, task, reason, strict); err != nil && firstErr == nil {
firstErr = err
}
}
return firstErr
}
func (s *SchedulerSnapshotService) rebuildBucketWithTokenPolicy(ctx context.Context, task schedulerBucketWriteTask, reason string, strict bool) error {
@@ -801,21 +803,213 @@ func (s *SchedulerSnapshotService) triggerFullRebuild(reason string) error {
return s.coalesceFullRebuild(func() error {
ctx, cancel := context.WithTimeout(context.Background(), 2*time.Minute)
defer cancel()
return s.rebuildFullSnapshot(ctx, reason)
})
}
buckets, err := s.cache.ListBuckets(ctx)
func (s *SchedulerSnapshotService) rebuildFullSnapshot(ctx context.Context, reason string) error {
if s.cache == nil {
return ErrSchedulerCacheNotReady
}
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 {
logger.LegacyPrintf("service.scheduler_snapshot", "[Scheduler] list buckets failed: %v", err)
return err
}
if len(buckets) == 0 {
buckets, err = s.defaultBuckets(ctx)
if err != nil {
logger.LegacyPrintf("service.scheduler_snapshot", "[Scheduler] default buckets failed: %v", err)
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...)
}
return s.rebuildBuckets(ctx, buckets, reason)
})
}
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
}
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 {
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...)
if err := s.rebuildPreparedBucketTasks(ctx, reopened, reason, true); err != nil {
firstErr = err
}
if err := s.rebuildPreparedBucketTasks(ctx, captured, reason, false); 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 {
@@ -1053,38 +1247,6 @@ func (s *SchedulerSnapshotService) fullRebuildInterval() time.Duration {
return time.Duration(sec) * time.Second
}
func (s *SchedulerSnapshotService) defaultBuckets(ctx context.Context) ([]SchedulerBucket, error) {
buckets := make([]SchedulerBucket, 0)
platforms := []string{PlatformAnthropic, PlatformGemini, PlatformOpenAI, PlatformAntigravity, PlatformGrok}
for _, platform := range platforms {
buckets = append(buckets, SchedulerBucket{GroupID: 0, Platform: platform, Mode: SchedulerModeSingle})
buckets = append(buckets, SchedulerBucket{GroupID: 0, Platform: platform, Mode: SchedulerModeForced})
if platform == PlatformAnthropic || platform == PlatformGemini {
buckets = append(buckets, SchedulerBucket{GroupID: 0, Platform: platform, Mode: SchedulerModeMixed})
}
}
if s.isRunModeSimple() || s.groupRepo == nil {
return dedupeBuckets(buckets), nil
}
groups, err := s.groupRepo.ListActive(ctx)
if err != nil {
return dedupeBuckets(buckets), nil
}
for _, group := range groups {
if group.Platform == "" {
continue
}
buckets = append(buckets, SchedulerBucket{GroupID: group.ID, Platform: group.Platform, Mode: SchedulerModeSingle})
buckets = append(buckets, SchedulerBucket{GroupID: group.ID, Platform: group.Platform, Mode: SchedulerModeForced})
if group.Platform == PlatformAnthropic || group.Platform == PlatformGemini {
buckets = append(buckets, SchedulerBucket{GroupID: group.ID, Platform: group.Platform, Mode: SchedulerModeMixed})
}
}
return dedupeBuckets(buckets), nil
}
func dedupeBuckets(in []SchedulerBucket) []SchedulerBucket {
seen := make(map[string]struct{}, len(in))
out := make([]SchedulerBucket, 0, len(in))