diff --git a/backend/internal/service/scheduler_snapshot_full_rebuild_lifecycle_test.go b/backend/internal/service/scheduler_snapshot_full_rebuild_lifecycle_test.go new file mode 100644 index 0000000000..1cb48260be --- /dev/null +++ b/backend/internal/service/scheduler_snapshot_full_rebuild_lifecycle_test.go @@ -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 ©Group, 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()) +} diff --git a/backend/internal/service/scheduler_snapshot_full_rebuild_test.go b/backend/internal/service/scheduler_snapshot_full_rebuild_test.go index 7ef4d9ed9f..02b7548348 100644 --- a/backend/internal/service/scheduler_snapshot_full_rebuild_test.go +++ b/backend/internal/service/scheduler_snapshot_full_rebuild_test.go @@ -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) { diff --git a/backend/internal/service/scheduler_snapshot_retirement_test.go b/backend/internal/service/scheduler_snapshot_retirement_test.go index e1ac974f81..3d24179758 100644 --- a/backend/internal/service/scheduler_snapshot_retirement_test.go +++ b/backend/internal/service/scheduler_snapshot_retirement_test.go @@ -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) -} diff --git a/backend/internal/service/scheduler_snapshot_service.go b/backend/internal/service/scheduler_snapshot_service.go index 6f87e7bbf2..edf18cb56e 100644 --- a/backend/internal/service/scheduler_snapshot_service.go +++ b/backend/internal/service/scheduler_snapshot_service.go @@ -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))