diff --git a/backend/internal/service/scheduler_snapshot_full_rebuild_lifecycle_test.go b/backend/internal/service/scheduler_snapshot_full_rebuild_lifecycle_test.go index e52adbb4c8..e4489b81c7 100644 --- a/backend/internal/service/scheduler_snapshot_full_rebuild_lifecycle_test.go +++ b/backend/internal/service/scheduler_snapshot_full_rebuild_lifecycle_test.go @@ -290,13 +290,14 @@ func TestSchedulerFullRebuildActiveTombstoneDoesNotBlockFollowingGroupEvent(t *t require.Equal(t, 1, activeCalls) require.Zero(t, fallbackCalls) require.Equal(t, []int64{groupID, groupID}, freshCalls) - require.Len(t, cache.tokens(), 36, "full rebuild and the following group event must each run fresh authority") + canonicalCount := len(canonical) + require.Len(t, cache.tokens(), 2*canonicalCount, "full rebuild and the following group event must each run fresh authority") _, reopenHeld := cache.lifecycleMutationLeaseStates() - require.Len(t, reopenHeld, 36) + require.Len(t, reopenHeld, 2*canonicalCount) for _, held := range reopenHeld { require.True(t, held) } - require.Equal(t, 30, accounts.callCount()) + require.Equal(t, 3*schedulerCanonicalAccountQueryCount(), accounts.callCount()) } func TestSchedulerFullRebuildGlobalReadErrorsFailBeforeMutationOrDB(t *testing.T) { @@ -367,15 +368,15 @@ func TestSchedulerFullRebuildFreshActivePreparesEveryTokenBeforeFirstDB(t *testi capturesAtFirstDB = cache.captureAttemptCount() held, reopenCount := cache.leaseHeldAndTokenCount() require.False(t, held) - require.Equal(t, 18, reopenCount) - require.Equal(t, 19, capturesAtFirstDB, "C(0) and the historical bucket must be captured before DB") + require.Equal(t, schedulerCanonicalBucketCount(), reopenCount) + require.Equal(t, schedulerCanonicalBucketCount()+1, 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, 21, accounts.callCount()) - require.Equal(t, 11, accounts.groupCallCount(groupID)) + require.Equal(t, 2*schedulerCanonicalAccountQueryCount()+1, accounts.callCount()) + require.Equal(t, schedulerCanonicalAccountQueryCount()+1, accounts.groupCallCount(groupID)) _, historicalPublished := cache.counts(historical) require.Equal(t, 1, historicalPublished) activeCalls, fallbackCalls, freshCalls := groups.stats() @@ -398,7 +399,7 @@ func TestSchedulerFullRebuildOrdinaryCaptureErrorReturnsBeforeFirstDB(t *testing err := svc.rebuildFullSnapshot(context.Background(), "test") require.ErrorIs(t, err, wantErr) - require.Equal(t, 20, cache.captureAttemptCount(), "all canonical and ordinary captures must be attempted before returning") + require.Equal(t, schedulerCanonicalBucketCount()+2, cache.captureAttemptCount(), "all canonical and ordinary captures must be attempted before returning") require.Zero(t, accounts.callCount()) require.Zero(t, cache.totalSetAttempts()) } @@ -416,8 +417,8 @@ func TestSchedulerFullRebuildPreservesGroupZeroActiveHistoricalAndInvalidRegistr svc := newFullRebuildLifecycleService(cache, nil, accounts, groups, config.RunModeStandard) require.NoError(t, svc.rebuildFullSnapshot(context.Background(), "test")) - require.Equal(t, 39, cache.captureAttemptCount()) - require.Equal(t, 23, accounts.callCount()) + require.Equal(t, 2*schedulerCanonicalBucketCount()+3, cache.captureAttemptCount()) + require.Equal(t, 2*schedulerCanonicalAccountQueryCount()+3, accounts.callCount()) groups.mu.Lock() require.Equal(t, 1, groups.listCalls) groups.mu.Unlock() @@ -456,7 +457,7 @@ func TestSchedulerFullRebuildActiveTombstoneFreshInactiveOrMissingFiltersAllGrou require.NoError(t, svc.rebuildFullSnapshot(context.Background(), "test")) require.Zero(t, accounts.groupCallCount(groupID)) - require.Equal(t, 10, accounts.groupCallCount(0)) + require.Equal(t, schedulerCanonicalAccountQueryCount(), accounts.groupCallCount(0)) require.Empty(t, cache.tokens()) require.Equal(t, bucketStrings(append(canonical, historical)), bucketStrings(cache.retiredBuckets())) for _, bucket := range append(canonical, historical) { @@ -528,7 +529,7 @@ func TestSchedulerFullRebuildPartialLifecycleFailureReturnsBeforeDBAndRetries(t require.Equal(t, []int64{1, 2}, freshCalls) require.Zero(t, accounts.callCount()) require.Zero(t, cache.totalSetAttempts()) - require.Equal(t, 19, len(cache.retiredBuckets())) + require.Equal(t, schedulerCanonicalBucketCount()+1, len(cache.retiredBuckets())) groups.mu.Lock() delete(groups.freshErr, 2) @@ -537,8 +538,8 @@ func TestSchedulerFullRebuildPartialLifecycleFailureReturnsBeforeDBAndRetries(t require.NoError(t, svc.triggerFullRebuild("retry")) _, _, freshCalls = groups.stats() require.Equal(t, []int64{1, 2, 2, 3}, freshCalls) - require.Equal(t, 57, len(cache.retiredBuckets())) - require.Equal(t, 10, accounts.callCount()) + require.Equal(t, 3*(schedulerCanonicalBucketCount()+1), len(cache.retiredBuckets())) + require.Equal(t, schedulerCanonicalAccountQueryCount(), accounts.callCount()) require.Empty(t, cache.tokens()) } @@ -561,14 +562,15 @@ func TestSchedulerFullRebuildActiveTombstoneLazyRecoveryDiscardsPartialCaptureTa capturesAtFirstDB = cache.captureAttemptCount() held, reopenCount := cache.leaseHeldAndTokenCount() require.False(t, held) - require.Equal(t, 18, reopenCount) - require.Equal(t, 25, capturesAtFirstDB) + require.Equal(t, len(canonical), reopenCount) + // C(0) + captures through the retired canonical index (inclusive) + historical ordinary bucket. + require.Equal(t, schedulerCanonicalBucketCount()+6+1, 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, 21, accounts.callCount()) + require.Equal(t, 2*schedulerCanonicalAccountQueryCount()+1, 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()) @@ -600,9 +602,9 @@ func TestSchedulerFullRebuildSimpleModePreservesRegistryWithoutLifecycleAuthorit require.Zero(t, activeCalls) require.Zero(t, fallbackCalls) require.Empty(t, freshCalls) - require.Equal(t, 21, cache.captureAttemptCount()) - require.Equal(t, 13, accounts.callCount()) - require.Equal(t, 13, accounts.groupCallCount(0)) + require.Equal(t, schedulerCanonicalBucketCount()+len(registered), cache.captureAttemptCount()) + require.Equal(t, schedulerCanonicalAccountQueryCount()+len(registered), accounts.callCount()) + require.Equal(t, schedulerCanonicalAccountQueryCount()+len(registered), accounts.groupCallCount(0)) require.Empty(t, cache.retiredBuckets()) require.Empty(t, cache.tokens()) for _, bucket := range registered { @@ -632,11 +634,11 @@ func TestSchedulerFullRebuildFreshReopenLockBusyRetriesWithoutBlockingOrdinaryTa 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, 20, accounts.callCount()) + require.Equal(t, 2*schedulerCanonicalAccountQueryCount(), accounts.callCount()) svc.pollOutbox() require.Equal(t, int64(1), cache.currentWatermark()) - require.Equal(t, 40, accounts.callCount()) + require.Equal(t, 4*schedulerCanonicalAccountQueryCount(), accounts.callCount()) _, busyBucketPublished := cache.counts(canonical[0]) require.Equal(t, 1, busyBucketPublished) activeCalls, fallbackCalls, freshCalls := groups.stats() @@ -654,7 +656,7 @@ func TestSchedulerFullRebuildOrdinaryLockBusyKeepsExistingSkipSemantics(t *testi svc := newFullRebuildLifecycleService(cache, nil, accounts, groups, config.RunModeStandard) require.NoError(t, svc.rebuildFullSnapshot(context.Background(), "test")) - require.Equal(t, 10, accounts.callCount()) + require.Equal(t, schedulerCanonicalAccountQueryCount(), accounts.callCount()) attempts, published := cache.counts(busyBucket) require.Zero(t, attempts) require.Zero(t, published) diff --git a/backend/internal/service/scheduler_snapshot_group_lifecycle_test.go b/backend/internal/service/scheduler_snapshot_group_lifecycle_test.go index f9e75ed158..7eda02f737 100644 --- a/backend/internal/service/scheduler_snapshot_group_lifecycle_test.go +++ b/backend/internal/service/scheduler_snapshot_group_lifecycle_test.go @@ -325,18 +325,18 @@ func newGroupLifecycleTestService(cache SchedulerCache, accounts AccountReposito } func expectedGroupLifecycleBuckets(groupID int64) []SchedulerBucket { - platforms := schedulerSnapshotPlatforms() - buckets := make([]SchedulerBucket, 0, 18) - for _, platform := range platforms { - buckets = append(buckets, - SchedulerBucket{GroupID: groupID, Platform: platform, Mode: SchedulerModeSingle}, - SchedulerBucket{GroupID: groupID, Platform: platform, Mode: SchedulerModeForced}, - ) + return schedulerCanonicalBuckets(groupID) +} + +func schedulerCanonicalAccountQueryCount() int { + count := 0 + for _, platform := range schedulerSnapshotPlatforms() { + count++ if platform == PlatformAnthropic || platform == PlatformGemini { - buckets = append(buckets, SchedulerBucket{GroupID: groupID, Platform: platform, Mode: SchedulerModeMixed}) + count++ } } - return buckets + return count } func bucketStrings(buckets []SchedulerBucket) map[string]struct{} { @@ -441,7 +441,7 @@ func TestSchedulerGroupLifecycleActiveReopensAndRebuildsAllCurrentBuckets(t *tes accounts.beforeLoad = func() { held, tokenCount := cache.leaseHeldAndTokenCount() require.False(t, held, "the group lifecycle lease must be released before the first account query") - require.Equal(t, 18, tokenCount, "all reopen tokens must be prepared before the first account query") + require.Equal(t, len(current), tokenCount, "all reopen tokens must be prepared before the first account query") } svc := newGroupLifecycleTestService(cache, accounts, groups, config.RunModeStandard) seen := make(map[batchSeenKey]struct{}) @@ -453,8 +453,8 @@ func TestSchedulerGroupLifecycleActiveReopensAndRebuildsAllCurrentBuckets(t *tes registered, err := cache.retirementRaceCache.ListBuckets(context.Background()) require.NoError(t, err) require.Contains(t, bucketStrings(registered), historical.String()) - require.Len(t, cache.tokens(), 18) - require.Equal(t, 10, accounts.callCount()) + require.Len(t, cache.tokens(), len(current)) + require.Equal(t, schedulerCanonicalAccountQueryCount(), accounts.callCount()) require.Equal(t, 1, accounts.platformCallCount(PlatformOpenAI)) for _, bucket := range current { _, published := cache.counts(bucket) @@ -472,16 +472,16 @@ func TestSchedulerGroupLifecycleActiveReopensAndRebuildsAllCurrentBuckets(t *tes require.True(t, cache.releaseDeadline) require.NoError(t, cache.releaseCtxErr) _, reopenHeld := cache.lifecycleMutationLeaseStates() - require.Len(t, reopenHeld, 18) + require.Len(t, reopenHeld, len(current)) for _, held := range reopenHeld { require.True(t, held) } lockTTLs, unlockCalls := cache.lockStats() - require.Len(t, lockTTLs, 18) + require.Len(t, lockTTLs, len(current)) for _, ttl := range lockTTLs { require.Equal(t, 30*time.Second, ttl) } - require.Equal(t, 18, unlockCalls) + require.Equal(t, len(current), unlockCalls) requireLifecycleSeen(t, seen, groupID) } @@ -497,8 +497,8 @@ func TestSchedulerGroupLifecycleInactiveThenActiveAuthoritativelyReopens(t *test groups.set(&Group{ID: groupID, Status: StatusActive, Hydrated: true}, nil) require.NoError(t, svc.handleGroupEvent(context.Background(), ptrInt64(groupID), make(map[batchSeenKey]struct{}))) - require.Len(t, cache.tokens(), 18) - require.Equal(t, 10, accounts.callCount()) + require.Len(t, cache.tokens(), schedulerCanonicalBucketCount()) + require.Equal(t, schedulerCanonicalAccountQueryCount(), accounts.callCount()) for _, bucket := range expectedGroupLifecycleBuckets(groupID) { _, published := cache.counts(bucket) require.Equal(t, 1, published, bucket.String()) @@ -546,15 +546,16 @@ func TestSchedulerGroupLifecycleEpochPreventsABA(t *testing.T) { groups.set(&Group{ID: groupID, Status: StatusActive, Hydrated: true}, nil) require.NoError(t, svc.handleGroupEvent(context.Background(), ptrInt64(groupID), make(map[batchSeenKey]struct{}))) firstActiveTokens := cache.tokens() - require.Len(t, firstActiveTokens, 18) + canonicalCount := schedulerCanonicalBucketCount() + require.Len(t, firstActiveTokens, canonicalCount) groups.set(&Group{ID: groupID, Status: StatusDisabled, Hydrated: true}, nil) require.NoError(t, svc.handleGroupEvent(context.Background(), ptrInt64(groupID), make(map[batchSeenKey]struct{}))) groups.set(&Group{ID: groupID, Status: StatusActive, Hydrated: true}, nil) require.NoError(t, svc.handleGroupEvent(context.Background(), ptrInt64(groupID), make(map[batchSeenKey]struct{}))) allTokens := cache.tokens() - require.Len(t, allTokens, 36) - require.Greater(t, allTokens[18].Epoch, firstActiveTokens[0].Epoch) + require.Len(t, allTokens, 2*canonicalCount) + require.Greater(t, allTokens[canonicalCount].Epoch, firstActiveTokens[0].Epoch) require.ErrorIs(t, cache.SetSnapshot(context.Background(), firstActiveTokens[0].Bucket, firstActiveTokens[0], nil), ErrSchedulerBucketWriteFenced) } @@ -571,11 +572,11 @@ func TestSchedulerGroupLifecycleSeenIsIndependentAndDeduplicatesGroupEvents(t *t require.NoError(t, svc.handleGroupEvent(context.Background(), ptrInt64(groupID), seen)) require.Equal(t, 1, groups.callCount()) - require.Equal(t, 10, accounts.callCount()) + require.Equal(t, schedulerCanonicalAccountQueryCount(), accounts.callCount()) requireLifecycleSeen(t, seen, groupID) require.NoError(t, svc.handleGroupEvent(context.Background(), ptrInt64(groupID), seen)) require.Equal(t, 1, groups.callCount()) - require.Equal(t, 10, accounts.callCount()) + require.Equal(t, schedulerCanonicalAccountQueryCount(), accounts.callCount()) } func TestSchedulerGroupLifecycleFailuresDoNotMarkSeen(t *testing.T) { diff --git a/backend/internal/service/scheduler_snapshot_retirement_test.go b/backend/internal/service/scheduler_snapshot_retirement_test.go index 6fa2b7ddc5..7b5a8368d3 100644 --- a/backend/internal/service/scheduler_snapshot_retirement_test.go +++ b/backend/internal/service/scheduler_snapshot_retirement_test.go @@ -176,7 +176,7 @@ func TestSchedulerFullRebuildCapturesAllRegistryTokensBeforeDBLoad(t *testing.T) } captures, reopens := cache.captureAndReopenCounts() - require.Equal(t, 36, captures, "group0 and active-group canonical tokens must be captured before the first DB load") + require.Equal(t, 2*schedulerCanonicalBucketCount(), 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) diff --git a/backend/internal/service/scheduler_snapshot_service.go b/backend/internal/service/scheduler_snapshot_service.go index c307c302ea..6ee67979e1 100644 --- a/backend/internal/service/scheduler_snapshot_service.go +++ b/backend/internal/service/scheduler_snapshot_service.go @@ -836,9 +836,21 @@ func schedulerBucketsForGroup(groupID int64) []SchedulerBucket { return schedulerCanonicalBuckets(groupID) } -func schedulerCanonicalBuckets(groupID int64) []SchedulerBucket { - buckets := make([]SchedulerBucket, 0, 18) +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}, @@ -855,7 +867,7 @@ func (s *SchedulerSnapshotService) rebuildByGroupIDs(ctx context.Context, groupI if len(groupIDs) == 0 { return nil } - buckets := make([]SchedulerBucket, 0, len(groupIDs)*18) + buckets := make([]SchedulerBucket, 0, len(groupIDs)*schedulerCanonicalBucketCount()) for _, platform := range schedulerSnapshotPlatforms() { buckets = append(buckets, s.bucketsForPlatform(platform, groupIDs, seen)...) }