From 143cab7bf5a4be75088a2d81231904f3bb06f9d3 Mon Sep 17 00:00:00 2001 From: jjaw Date: Wed, 15 Jul 2026 05:31:50 +0800 Subject: [PATCH] =?UTF-8?q?=E5=A4=8D=E7=94=A8=E9=87=8D=E5=BB=BA=E6=89=B9?= =?UTF-8?q?=E6=AC=A1=E5=86=85=E7=9A=84=E8=B4=A6=E5=8F=B7=E6=9F=A5=E8=AF=A2?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../scheduler_snapshot_batch_query_test.go | 534 ++++++++++++++++++ ...er_snapshot_full_rebuild_lifecycle_test.go | 24 +- ...scheduler_snapshot_group_lifecycle_test.go | 40 +- .../service/scheduler_snapshot_service.go | 107 +++- 4 files changed, 672 insertions(+), 33 deletions(-) create mode 100644 backend/internal/service/scheduler_snapshot_batch_query_test.go diff --git a/backend/internal/service/scheduler_snapshot_batch_query_test.go b/backend/internal/service/scheduler_snapshot_batch_query_test.go new file mode 100644 index 0000000000..a2e7ae8e85 --- /dev/null +++ b/backend/internal/service/scheduler_snapshot_batch_query_test.go @@ -0,0 +1,534 @@ +//go:build unit + +package service + +import ( + "context" + "errors" + "sync" + "testing" + "time" + + "github.com/Wei-Shaw/sub2api/internal/config" + "github.com/stretchr/testify/require" +) + +type batchAccountQueryKey struct { + groupID int64 + platform string + mixed bool +} + +type batchAccountQueryResult struct { + accounts []Account + err error +} + +type batchAccountQueryRepo struct { + AccountRepository + + mu sync.Mutex + calls map[batchAccountQueryKey]int + results map[batchAccountQueryKey][]batchAccountQueryResult + beforeRun func(batchAccountQueryKey) +} + +func newBatchAccountQueryRepo() *batchAccountQueryRepo { + return &batchAccountQueryRepo{ + calls: make(map[batchAccountQueryKey]int), + results: make(map[batchAccountQueryKey][]batchAccountQueryResult), + } +} + +func (r *batchAccountQueryRepo) ListSchedulableByGroupIDAndPlatform(_ context.Context, groupID int64, platform string) ([]Account, error) { + return r.run(batchAccountQueryKey{groupID: groupID, platform: platform}) +} + +func (r *batchAccountQueryRepo) ListSchedulableByGroupIDAndPlatforms(_ context.Context, groupID int64, platforms []string) ([]Account, error) { + return r.run(batchAccountQueryKey{groupID: groupID, platform: platforms[0], mixed: true}) +} + +func (r *batchAccountQueryRepo) ListSchedulableUngroupedByPlatform(_ context.Context, platform string) ([]Account, error) { + return r.run(batchAccountQueryKey{platform: platform}) +} + +func (r *batchAccountQueryRepo) ListSchedulableUngroupedByPlatforms(_ context.Context, platforms []string) ([]Account, error) { + return r.run(batchAccountQueryKey{platform: platforms[0], mixed: true}) +} + +func (r *batchAccountQueryRepo) ListSchedulableByPlatform(_ context.Context, platform string) ([]Account, error) { + return r.run(batchAccountQueryKey{platform: platform}) +} + +func (r *batchAccountQueryRepo) ListSchedulableByPlatforms(_ context.Context, platforms []string) ([]Account, error) { + return r.run(batchAccountQueryKey{platform: platforms[0], mixed: true}) +} + +func (r *batchAccountQueryRepo) run(key batchAccountQueryKey) ([]Account, error) { + r.mu.Lock() + r.calls[key]++ + call := r.calls[key] + results := r.results[key] + beforeRun := r.beforeRun + r.mu.Unlock() + + if beforeRun != nil { + beforeRun(key) + } + if call <= len(results) { + result := results[call-1] + return append([]Account(nil), result.accounts...), result.err + } + return []Account{{ + ID: int64(call), + Name: "source", + Platform: key.platform, + Status: StatusActive, + Schedulable: true, + }}, nil +} + +func (r *batchAccountQueryRepo) callCount(key batchAccountQueryKey) int { + r.mu.Lock() + defer r.mu.Unlock() + return r.calls[key] +} + +type batchSnapshotWrite struct { + token SchedulerBucketWriteToken + accounts []Account +} + +type batchSnapshotCache struct { + SchedulerCache + + mu sync.Mutex + nextEpoch int64 + captures []SchedulerBucket + captured map[SchedulerBucket]SchedulerBucketWriteToken + locks map[SchedulerBucket]int + lockBusy map[SchedulerBucket]bool + lockErrors map[SchedulerBucket]error + setErrors map[SchedulerBucket]error + setAttempts map[SchedulerBucket]int + writes map[SchedulerBucket][]batchSnapshotWrite + versions map[SchedulerBucket]int + beforeSet func() +} + +func newBatchSnapshotCache() *batchSnapshotCache { + return &batchSnapshotCache{ + captured: make(map[SchedulerBucket]SchedulerBucketWriteToken), + locks: make(map[SchedulerBucket]int), + lockBusy: make(map[SchedulerBucket]bool), + lockErrors: make(map[SchedulerBucket]error), + setErrors: make(map[SchedulerBucket]error), + setAttempts: make(map[SchedulerBucket]int), + writes: make(map[SchedulerBucket][]batchSnapshotWrite), + versions: make(map[SchedulerBucket]int), + } +} + +func (c *batchSnapshotCache) CaptureBucketWriteToken(_ context.Context, bucket SchedulerBucket) (SchedulerBucketWriteToken, error) { + c.mu.Lock() + defer c.mu.Unlock() + c.nextEpoch++ + token := SchedulerBucketWriteToken{Bucket: bucket, Epoch: c.nextEpoch} + c.captures = append(c.captures, bucket) + c.captured[bucket] = token + return token, nil +} + +func (c *batchSnapshotCache) TryLockBucket(_ context.Context, bucket SchedulerBucket, _ time.Duration) (bool, error) { + c.mu.Lock() + defer c.mu.Unlock() + c.locks[bucket]++ + if err := c.lockErrors[bucket]; err != nil { + return false, err + } + return !c.lockBusy[bucket], nil +} + +func (c *batchSnapshotCache) UnlockBucket(context.Context, SchedulerBucket) error { + return nil +} + +func (c *batchSnapshotCache) SetSnapshot(_ context.Context, bucket SchedulerBucket, token SchedulerBucketWriteToken, accounts []Account) error { + if c.beforeSet != nil { + c.beforeSet() + } + c.mu.Lock() + defer c.mu.Unlock() + c.setAttempts[bucket]++ + if token != c.captured[bucket] || !token.ValidFor(bucket) { + return ErrSchedulerBucketWriteFenced + } + if err := c.setErrors[bucket]; err != nil { + return err + } + c.versions[bucket]++ + c.writes[bucket] = append(c.writes[bucket], batchSnapshotWrite{ + token: token, + accounts: append([]Account(nil), accounts...), + }) + return nil +} + +func (c *batchSnapshotCache) captureCount() int { + c.mu.Lock() + defer c.mu.Unlock() + return len(c.captures) +} + +func (c *batchSnapshotCache) bucketState(bucket SchedulerBucket) (locks, attempts, version int, writes []batchSnapshotWrite) { + c.mu.Lock() + defer c.mu.Unlock() + return c.locks[bucket], c.setAttempts[bucket], c.versions[bucket], append([]batchSnapshotWrite(nil), c.writes[bucket]...) +} + +func newBatchQueryTestService(cache SchedulerCache, accounts AccountRepository, runMode string) *SchedulerSnapshotService { + return NewSchedulerSnapshotService(cache, nil, accounts, nil, &config.Config{RunMode: runMode}) +} + +func TestSchedulerRebuildBatchReusesSingleForcedQueryAndKeepsSnapshotsIndependent(t *testing.T) { + const groupID int64 = 201 + single := SchedulerBucket{GroupID: groupID, Platform: PlatformOpenAI, Mode: SchedulerModeSingle} + forced := SchedulerBucket{GroupID: groupID, Platform: PlatformOpenAI, Mode: SchedulerModeForced} + cache := newBatchSnapshotCache() + repo := newBatchAccountQueryRepo() + wantCaptures := 2 + repo.beforeRun = func(batchAccountQueryKey) { + require.Equal(t, wantCaptures, cache.captureCount(), "all tokens must be prepared before the first DB query") + wantCaptures += 2 + } + svc := newBatchQueryTestService(cache, repo, config.RunModeStandard) + + require.NoError(t, svc.rebuildBuckets(context.Background(), []SchedulerBucket{single, forced}, "first")) + queryKey := batchAccountQueryKey{groupID: groupID, platform: PlatformOpenAI} + require.Equal(t, 1, repo.callCount(queryKey)) + for _, bucket := range []SchedulerBucket{single, forced} { + locks, attempts, version, writes := cache.bucketState(bucket) + require.Equal(t, 1, locks, bucket.String()) + require.Equal(t, 1, attempts, bucket.String()) + require.Equal(t, 1, version, bucket.String()) + require.Len(t, writes, 1, bucket.String()) + require.Equal(t, "source", writes[0].accounts[0].Name, bucket.String()) + require.Equal(t, bucket, writes[0].token.Bucket) + } + _, _, _, singleWrites := cache.bucketState(single) + _, _, _, forcedWrites := cache.bucketState(forced) + require.NotEqual(t, singleWrites[0].token.Epoch, forcedWrites[0].token.Epoch) + + require.NoError(t, svc.rebuildBuckets(context.Background(), []SchedulerBucket{single, forced}, "second")) + require.Equal(t, 2, repo.callCount(queryKey), "successful results must not be cached across rebuild batches") + for _, bucket := range []SchedulerBucket{single, forced} { + locks, attempts, version, writes := cache.bucketState(bucket) + require.Equal(t, 2, locks, bucket.String()) + require.Equal(t, 2, attempts, bucket.String()) + require.Equal(t, 2, version, bucket.String()) + require.Len(t, writes, 2, bucket.String()) + require.Equal(t, "source", writes[1].accounts[0].Name, bucket.String()) + } +} + +func TestSchedulerRebuildBatchKeepsMixedAndDifferentKeysIndependent(t *testing.T) { + const groupID int64 = 202 + buckets := []SchedulerBucket{ + {GroupID: groupID, Platform: PlatformAnthropic, Mode: SchedulerModeSingle}, + {GroupID: groupID, Platform: PlatformAnthropic, Mode: SchedulerModeForced}, + {GroupID: groupID, Platform: PlatformAnthropic, Mode: SchedulerModeMixed}, + {GroupID: groupID + 1, Platform: PlatformAnthropic, Mode: SchedulerModeSingle}, + {GroupID: groupID, Platform: PlatformGemini, Mode: SchedulerModeForced}, + {GroupID: 0, Platform: PlatformOpenAI, Mode: SchedulerModeSingle}, + {GroupID: -1, Platform: PlatformOpenAI, Mode: SchedulerModeForced}, + } + cache := newBatchSnapshotCache() + repo := newBatchAccountQueryRepo() + svc := newBatchQueryTestService(cache, repo, config.RunModeStandard) + + require.NoError(t, svc.rebuildBuckets(context.Background(), buckets, "test")) + require.Equal(t, 1, repo.callCount(batchAccountQueryKey{groupID: groupID, platform: PlatformAnthropic})) + require.Equal(t, 1, repo.callCount(batchAccountQueryKey{groupID: groupID, platform: PlatformAnthropic, mixed: true})) + require.Equal(t, 1, repo.callCount(batchAccountQueryKey{groupID: groupID + 1, platform: PlatformAnthropic})) + require.Equal(t, 1, repo.callCount(batchAccountQueryKey{groupID: groupID, platform: PlatformGemini})) + require.Equal(t, 2, repo.callCount(batchAccountQueryKey{platform: PlatformOpenAI}), "group0 and a negative historical group must not share") + for _, bucket := range buckets { + locks, attempts, version, _ := cache.bucketState(bucket) + require.Equal(t, 1, locks, bucket.String()) + require.Equal(t, 1, attempts, bucket.String()) + require.Equal(t, 1, version, bucket.String()) + } +} + +func TestSchedulerRebuildBatchKeepsSimpleModeBucketGroupsIndependent(t *testing.T) { + single := SchedulerBucket{GroupID: 204, Platform: PlatformOpenAI, Mode: SchedulerModeSingle} + forced := SchedulerBucket{GroupID: 0, Platform: PlatformOpenAI, Mode: SchedulerModeForced} + cache := newBatchSnapshotCache() + repo := newBatchAccountQueryRepo() + svc := newBatchQueryTestService(cache, repo, config.RunModeSimple) + + require.NoError(t, svc.rebuildBuckets(context.Background(), []SchedulerBucket{single, forced}, "test")) + require.Equal(t, 2, repo.callCount(batchAccountQueryKey{platform: PlatformOpenAI})) +} + +func TestSchedulerRebuildBatchDoesNotCacheMixedOrHistoricalQueries(t *testing.T) { + for _, tc := range []struct { + name string + bucket SchedulerBucket + key batchAccountQueryKey + }{ + { + name: "mixed", + bucket: SchedulerBucket{GroupID: 204, Platform: PlatformAnthropic, Mode: SchedulerModeMixed}, + key: batchAccountQueryKey{groupID: 204, platform: PlatformAnthropic, mixed: true}, + }, + { + name: "historical", + bucket: SchedulerBucket{GroupID: 204, Platform: PlatformOpenAI, Mode: "unknown"}, + key: batchAccountQueryKey{groupID: 204, platform: PlatformOpenAI}, + }, + } { + t.Run(tc.name, func(t *testing.T) { + cache := newBatchSnapshotCache() + token, err := cache.CaptureBucketWriteToken(context.Background(), tc.bucket) + require.NoError(t, err) + repo := newBatchAccountQueryRepo() + svc := newBatchQueryTestService(cache, repo, config.RunModeStandard) + tasks := []schedulerBucketWriteTask{ + {bucket: tc.bucket, token: token}, + {bucket: tc.bucket, token: token}, + } + queries := newSchedulerAccountQueryCache(tasks) + + require.NoError(t, svc.rebuildPreparedBucketTasks(context.Background(), tasks, "test", false, queries)) + require.Equal(t, 2, repo.callCount(tc.key)) + require.Empty(t, queries.accounts) + locks, attempts, version, _ := cache.bucketState(tc.bucket) + require.Equal(t, 2, locks) + require.Equal(t, 2, attempts) + require.Equal(t, 2, version) + }) + } +} + +func TestSchedulerRebuildBatchRetriesQueryFailureForFollowingBucket(t *testing.T) { + const groupID int64 = 205 + single := SchedulerBucket{GroupID: groupID, Platform: PlatformOpenAI, Mode: SchedulerModeSingle} + forced := SchedulerBucket{GroupID: groupID, Platform: PlatformOpenAI, Mode: SchedulerModeForced} + wantErr := errors.New("first query failed") + key := batchAccountQueryKey{groupID: groupID, platform: PlatformOpenAI} + repo := newBatchAccountQueryRepo() + repo.results[key] = []batchAccountQueryResult{ + {err: wantErr}, + {accounts: []Account{{ID: 2051, Name: "retry", Platform: PlatformOpenAI}}}, + } + cache := newBatchSnapshotCache() + svc := newBatchQueryTestService(cache, repo, config.RunModeStandard) + + err := svc.rebuildBuckets(context.Background(), []SchedulerBucket{single, forced}, "test") + require.ErrorIs(t, err, wantErr) + require.Equal(t, 2, repo.callCount(key), "failed queries must not enter the batch cache") + _, singleAttempts, singleVersion, _ := cache.bucketState(single) + _, forcedAttempts, forcedVersion, forcedWrites := cache.bucketState(forced) + require.Zero(t, singleAttempts) + require.Zero(t, singleVersion) + require.Equal(t, 1, forcedAttempts) + require.Equal(t, 1, forcedVersion) + require.Equal(t, "retry", forcedWrites[0].accounts[0].Name) +} + +func TestSchedulerFullRebuildSharesSuccessfulQueryAcrossStrictAndOrdinarySegments(t *testing.T) { + const groupID int64 = 206 + single := SchedulerBucket{GroupID: groupID, Platform: PlatformOpenAI, Mode: SchedulerModeSingle} + forced := SchedulerBucket{GroupID: groupID, Platform: PlatformOpenAI, Mode: SchedulerModeForced} + cache := newBatchSnapshotCache() + cache.setErrors[single] = ErrSchedulerBucketWriteFenced + singleToken, err := cache.CaptureBucketWriteToken(context.Background(), single) + require.NoError(t, err) + forcedToken, err := cache.CaptureBucketWriteToken(context.Background(), forced) + require.NoError(t, err) + repo := newBatchAccountQueryRepo() + svc := newBatchQueryTestService(cache, repo, config.RunModeStandard) + + err = svc.prepareAndRebuildFullSnapshot( + context.Background(), + []schedulerBucketWriteTask{{bucket: forced, token: forcedToken}}, + []schedulerBucketWriteTask{{bucket: single, token: singleToken}}, + nil, + "test", + ) + require.ErrorIs(t, err, ErrSchedulerBucketWriteFenced) + require.Equal(t, 1, repo.callCount(batchAccountQueryKey{groupID: groupID, platform: PlatformOpenAI}), "SetSnapshot failure must not discard a successful query") + _, singleAttempts, singleVersion, _ := cache.bucketState(single) + _, forcedAttempts, forcedVersion, _ := cache.bucketState(forced) + require.Equal(t, 1, singleAttempts) + require.Zero(t, singleVersion) + require.Equal(t, 1, forcedAttempts) + require.Equal(t, 1, forcedVersion) +} + +func TestSchedulerRebuildBatchPreservesLockBusyAndFencingPolicy(t *testing.T) { + const groupID int64 = 207 + single := SchedulerBucket{GroupID: groupID, Platform: PlatformOpenAI, Mode: SchedulerModeSingle} + forced := SchedulerBucket{GroupID: groupID, Platform: PlatformOpenAI, Mode: SchedulerModeForced} + + t.Run("ordinary lock busy skips only that bucket", func(t *testing.T) { + cache := newBatchSnapshotCache() + cache.lockBusy[single] = true + repo := newBatchAccountQueryRepo() + svc := newBatchQueryTestService(cache, repo, config.RunModeStandard) + + require.NoError(t, svc.rebuildBuckets(context.Background(), []SchedulerBucket{single, forced}, "test")) + require.Equal(t, 1, repo.callCount(batchAccountQueryKey{groupID: groupID, platform: PlatformOpenAI})) + _, singleAttempts, _, _ := cache.bucketState(single) + _, forcedAttempts, forcedVersion, _ := cache.bucketState(forced) + require.Zero(t, singleAttempts) + require.Equal(t, 1, forcedAttempts) + require.Equal(t, 1, forcedVersion) + }) + + t.Run("strict lock busy is returned while ordinary work continues", func(t *testing.T) { + cache := newBatchSnapshotCache() + cache.lockBusy[single] = true + singleToken, err := cache.CaptureBucketWriteToken(context.Background(), single) + require.NoError(t, err) + forcedToken, err := cache.CaptureBucketWriteToken(context.Background(), forced) + require.NoError(t, err) + repo := newBatchAccountQueryRepo() + svc := newBatchQueryTestService(cache, repo, config.RunModeStandard) + + err = svc.prepareAndRebuildFullSnapshot( + context.Background(), + []schedulerBucketWriteTask{{bucket: forced, token: forcedToken}}, + []schedulerBucketWriteTask{{bucket: single, token: singleToken}}, + nil, + "test", + ) + require.ErrorIs(t, err, ErrSchedulerBucketRebuildBusy) + require.Equal(t, 1, repo.callCount(batchAccountQueryKey{groupID: groupID, platform: PlatformOpenAI})) + _, forcedAttempts, forcedVersion, _ := cache.bucketState(forced) + require.Equal(t, 1, forcedAttempts) + require.Equal(t, 1, forcedVersion) + }) + + t.Run("ordinary fencing stays non-fatal", func(t *testing.T) { + cache := newBatchSnapshotCache() + cache.setErrors[single] = ErrSchedulerBucketWriteFenced + repo := newBatchAccountQueryRepo() + svc := newBatchQueryTestService(cache, repo, config.RunModeStandard) + + require.NoError(t, svc.rebuildBuckets(context.Background(), []SchedulerBucket{single, forced}, "test")) + require.Equal(t, 1, repo.callCount(batchAccountQueryKey{groupID: groupID, platform: PlatformOpenAI})) + _, singleAttempts, singleVersion, _ := cache.bucketState(single) + _, forcedAttempts, forcedVersion, _ := cache.bucketState(forced) + require.Equal(t, 1, singleAttempts) + require.Zero(t, singleVersion) + require.Equal(t, 1, forcedAttempts) + require.Equal(t, 1, forcedVersion) + }) +} + +func TestSchedulerRebuildBatchReleasesResultsAfterLastConsumer(t *testing.T) { + const groups = 128 + cache := newBatchSnapshotCache() + repo := newBatchAccountQueryRepo() + tasks := make([]schedulerBucketWriteTask, 0, groups*2) + wantLockErr := errors.New("lock failed") + for i := 1; i <= groups; i++ { + groupID := int64(300 + i) + single := SchedulerBucket{GroupID: groupID, Platform: PlatformOpenAI, Mode: SchedulerModeSingle} + forced := SchedulerBucket{GroupID: groupID, Platform: PlatformOpenAI, Mode: SchedulerModeForced} + if i == 1 { + cache.lockBusy[single] = true + } + if i == 2 { + cache.lockErrors[single] = wantLockErr + } + for _, bucket := range []SchedulerBucket{single, forced} { + token, err := cache.CaptureBucketWriteToken(context.Background(), bucket) + require.NoError(t, err) + tasks = append(tasks, schedulerBucketWriteTask{bucket: bucket, token: token}) + } + } + queries := newSchedulerAccountQueryCache(tasks) + maxResident := 0 + cache.beforeSet = func() { + if resident := len(queries.accounts); resident > maxResident { + maxResident = resident + } + } + svc := newBatchQueryTestService(cache, repo, config.RunModeStandard) + + err := svc.rebuildPreparedBucketTasks(context.Background(), tasks, "test", false, queries) + require.ErrorIs(t, err, wantLockErr) + require.LessOrEqual(t, maxResident, 1, "adjacent single/forced pairs must not accumulate full-batch results") + require.Empty(t, queries.accounts) + require.Empty(t, queries.remaining) + for i := 1; i <= groups; i++ { + key := batchAccountQueryKey{groupID: int64(300 + i), platform: PlatformOpenAI} + require.Equal(t, 1, repo.callCount(key), key) + } +} + +type batchQueryBenchmarkRepo struct { + AccountRepository + accounts []Account +} + +func (r *batchQueryBenchmarkRepo) ListSchedulableByGroupIDAndPlatform(context.Context, int64, string) ([]Account, error) { + return r.accounts, nil +} + +type batchQueryBenchmarkCache struct { + SchedulerCache +} + +func (c *batchQueryBenchmarkCache) CaptureBucketWriteToken(_ context.Context, bucket SchedulerBucket) (SchedulerBucketWriteToken, error) { + return SchedulerBucketWriteToken{Bucket: bucket, Epoch: 1}, nil +} + +func (c *batchQueryBenchmarkCache) TryLockBucket(context.Context, SchedulerBucket, time.Duration) (bool, error) { + return true, nil +} + +func (c *batchQueryBenchmarkCache) UnlockBucket(context.Context, SchedulerBucket) error { + return nil +} + +var batchQueryBenchmarkAccountCount int + +func (c *batchQueryBenchmarkCache) SetSnapshot(_ context.Context, _ SchedulerBucket, _ SchedulerBucketWriteToken, accounts []Account) error { + batchQueryBenchmarkAccountCount = len(accounts) + return nil +} + +func BenchmarkSchedulerRebuildBatchQueryReuse(b *testing.B) { + const groupID int64 = 208 + buckets := []SchedulerBucket{ + {GroupID: groupID, Platform: PlatformOpenAI, Mode: SchedulerModeSingle}, + {GroupID: groupID, Platform: PlatformOpenAI, Mode: SchedulerModeForced}, + } + for _, tc := range []struct { + name string + size int + }{ + {name: "1_account", size: 1}, + {name: "10000_accounts", size: 10_000}, + } { + b.Run(tc.name, func(b *testing.B) { + accounts := make([]Account, tc.size) + svc := newBatchQueryTestService( + &batchQueryBenchmarkCache{}, + &batchQueryBenchmarkRepo{accounts: accounts}, + config.RunModeStandard, + ) + b.ReportAllocs() + b.ResetTimer() + for i := 0; i < b.N; i++ { + if err := svc.rebuildBuckets(context.Background(), buckets, "benchmark"); err != nil { + b.Fatal(err) + } + } + }) + } +} 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 1cb48260be..35ce99a8d1 100644 --- a/backend/internal/service/scheduler_snapshot_full_rebuild_lifecycle_test.go +++ b/backend/internal/service/scheduler_snapshot_full_rebuild_lifecycle_test.go @@ -292,7 +292,7 @@ func TestSchedulerFullRebuildActiveTombstoneDoesNotBlockFollowingGroupEvent(t *t for _, held := range reopenHeld { require.True(t, held) } - require.Equal(t, 36, accounts.callCount()) + require.Equal(t, 21, accounts.callCount()) } func TestSchedulerFullRebuildGlobalReadErrorsFailBeforeMutationOrDB(t *testing.T) { @@ -370,8 +370,8 @@ func TestSchedulerFullRebuildFreshActivePreparesEveryTokenBeforeFirstDB(t *testi 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)) + require.Equal(t, 15, accounts.callCount()) + require.Equal(t, 8, accounts.groupCallCount(groupID)) _, historicalPublished := cache.counts(historical) require.Equal(t, 1, historicalPublished) activeCalls, fallbackCalls, freshCalls := groups.stats() @@ -413,7 +413,7 @@ func TestSchedulerFullRebuildPreservesGroupZeroActiveHistoricalAndInvalidRegistr require.NoError(t, svc.rebuildFullSnapshot(context.Background(), "test")) require.Equal(t, 27, cache.captureAttemptCount()) - require.Equal(t, 27, accounts.callCount()) + require.Equal(t, 17, accounts.callCount()) groups.mu.Lock() require.Equal(t, 1, groups.listCalls) groups.mu.Unlock() @@ -452,7 +452,7 @@ func TestSchedulerFullRebuildActiveTombstoneFreshInactiveOrMissingFiltersAllGrou require.NoError(t, svc.rebuildFullSnapshot(context.Background(), "test")) require.Zero(t, accounts.groupCallCount(groupID)) - require.Equal(t, 12, accounts.groupCallCount(0)) + require.Equal(t, 7, accounts.groupCallCount(0)) require.Empty(t, cache.tokens()) require.Equal(t, bucketStrings(append(canonical, historical)), bucketStrings(cache.retiredBuckets())) for _, bucket := range append(canonical, historical) { @@ -534,7 +534,7 @@ func TestSchedulerFullRebuildPartialLifecycleFailureReturnsBeforeDBAndRetries(t _, _, 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.Equal(t, 7, accounts.callCount()) require.Empty(t, cache.tokens()) } @@ -564,7 +564,7 @@ func TestSchedulerFullRebuildActiveTombstoneLazyRecoveryDiscardsPartialCaptureTa require.NoError(t, svc.rebuildFullSnapshot(context.Background(), "test")) require.Equal(t, capturesAtFirstDB, cache.captureAttemptCount()) - require.Equal(t, 25, accounts.callCount()) + require.Equal(t, 15, 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()) @@ -597,8 +597,8 @@ func TestSchedulerFullRebuildSimpleModePreservesRegistryWithoutLifecycleAuthorit 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.Equal(t, 10, accounts.callCount()) + require.Equal(t, 10, accounts.groupCallCount(0)) require.Empty(t, cache.retiredBuckets()) require.Empty(t, cache.tokens()) for _, bucket := range registered { @@ -628,11 +628,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, 23, accounts.callCount()) + require.Equal(t, 14, accounts.callCount()) svc.pollOutbox() require.Equal(t, int64(1), cache.currentWatermark()) - require.Equal(t, 47, accounts.callCount()) + require.Equal(t, 28, accounts.callCount()) _, busyBucketPublished := cache.counts(canonical[0]) require.Equal(t, 1, busyBucketPublished) activeCalls, fallbackCalls, freshCalls := groups.stats() @@ -650,7 +650,7 @@ func TestSchedulerFullRebuildOrdinaryLockBusyKeepsExistingSkipSemantics(t *testi svc := newFullRebuildLifecycleService(cache, nil, accounts, groups, config.RunModeStandard) require.NoError(t, svc.rebuildFullSnapshot(context.Background(), "test")) - require.Equal(t, 11, accounts.callCount()) + require.Equal(t, 7, 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 2bd97c4663..ffc697eafc 100644 --- a/backend/internal/service/scheduler_snapshot_group_lifecycle_test.go +++ b/backend/internal/service/scheduler_snapshot_group_lifecycle_test.go @@ -255,19 +255,24 @@ func (r *groupLifecycleTestGroupRepo) callCount() int { type groupLifecycleTestAccountRepo struct { AccountRepository - mu sync.Mutex - calls int - err error - started chan struct{} - release chan struct{} - once sync.Once - beforeLoad func() - beforeLoadOnce sync.Once + mu sync.Mutex + calls int + callsByPlatform map[string]int + err error + started chan struct{} + release chan struct{} + once sync.Once + beforeLoad func() + beforeLoadOnce sync.Once } func (r *groupLifecycleTestAccountRepo) load(ctx context.Context, platform string) ([]Account, error) { r.mu.Lock() r.calls++ + if r.callsByPlatform == nil { + r.callsByPlatform = make(map[string]int) + } + r.callsByPlatform[platform]++ err := r.err started := r.started release := r.release @@ -309,6 +314,12 @@ func (r *groupLifecycleTestAccountRepo) callCount() int { return r.calls } +func (r *groupLifecycleTestAccountRepo) platformCallCount(platform string) int { + r.mu.Lock() + defer r.mu.Unlock() + return r.callsByPlatform[platform] +} + func newGroupLifecycleTestService(cache SchedulerCache, accounts AccountRepository, groups GroupRepository, runMode string) *SchedulerSnapshotService { return NewSchedulerSnapshotService(cache, nil, accounts, groups, &config.Config{RunMode: runMode}) } @@ -443,7 +454,8 @@ func TestSchedulerGroupLifecycleActiveReopensAndRebuildsAllCurrentBuckets(t *tes require.NoError(t, err) require.Contains(t, bucketStrings(registered), historical.String()) require.Len(t, cache.tokens(), 12) - require.Equal(t, 12, accounts.callCount()) + require.Equal(t, 7, accounts.callCount()) + require.Equal(t, 1, accounts.platformCallCount(PlatformOpenAI)) for _, bucket := range current { _, published := cache.counts(bucket) require.Equal(t, 1, published, bucket.String()) @@ -486,7 +498,7 @@ func TestSchedulerGroupLifecycleInactiveThenActiveAuthoritativelyReopens(t *test require.NoError(t, svc.handleGroupEvent(context.Background(), ptrInt64(groupID), make(map[batchSeenKey]struct{}))) require.Len(t, cache.tokens(), 12) - require.Equal(t, 12, accounts.callCount()) + require.Equal(t, 7, accounts.callCount()) for _, bucket := range expectedGroupLifecycleBuckets(groupID) { _, published := cache.counts(bucket) require.Equal(t, 1, published, bucket.String()) @@ -559,11 +571,11 @@ func TestSchedulerGroupLifecycleSeenIsIndependentAndDeduplicatesGroupEvents(t *t require.NoError(t, svc.handleGroupEvent(context.Background(), ptrInt64(groupID), seen)) require.Equal(t, 1, groups.callCount()) - require.Equal(t, 12, accounts.callCount()) + require.Equal(t, 7, 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, 12, accounts.callCount()) + require.Equal(t, 7, accounts.callCount()) } func TestSchedulerGroupLifecycleFailuresDoNotMarkSeen(t *testing.T) { @@ -684,8 +696,10 @@ func TestSchedulerGroupLifecycleFailuresDoNotMarkSeen(t *testing.T) { require.Zero(t, accounts.callCount()) } if tc.name == "account rebuild error" || tc.name == "set snapshot error" { - _, unlockCalls := cache.lockStats() + lockTTLs, unlockCalls := cache.lockStats() + require.Len(t, lockTTLs, 1) require.Equal(t, 1, unlockCalls) + require.Equal(t, 1, accounts.callCount()) } }) } diff --git a/backend/internal/service/scheduler_snapshot_service.go b/backend/internal/service/scheduler_snapshot_service.go index edf18cb56e..9dde73358d 100644 --- a/backend/internal/service/scheduler_snapshot_service.go +++ b/backend/internal/service/scheduler_snapshot_service.go @@ -43,6 +43,55 @@ type schedulerBucketWriteTask struct { token SchedulerBucketWriteToken } +type schedulerAccountQueryKey struct { + groupID int64 + platform string +} + +type schedulerAccountQueryCache struct { + remaining map[schedulerAccountQueryKey]int + accounts map[schedulerAccountQueryKey][]Account +} + +func newSchedulerAccountQueryCache(taskSets ...[]schedulerBucketWriteTask) *schedulerAccountQueryCache { + queries := &schedulerAccountQueryCache{ + remaining: make(map[schedulerAccountQueryKey]int), + accounts: make(map[schedulerAccountQueryKey][]Account), + } + for _, tasks := range taskSets { + for _, task := range tasks { + if key, ok := schedulerAccountQueryKeyForBucket(task.bucket); ok { + queries.remaining[key]++ + } + } + } + return queries +} + +func schedulerAccountQueryKeyForBucket(bucket SchedulerBucket) (schedulerAccountQueryKey, bool) { + if bucket.Mode != SchedulerModeSingle && bucket.Mode != SchedulerModeForced { + return schedulerAccountQueryKey{}, false + } + return schedulerAccountQueryKey{groupID: bucket.GroupID, platform: bucket.Platform}, true +} + +func (c *schedulerAccountQueryCache) release(bucket SchedulerBucket) { + if c == nil { + return + } + key, ok := schedulerAccountQueryKeyForBucket(bucket) + if !ok { + return + } + remaining := c.remaining[key] - 1 + if remaining <= 0 { + delete(c.remaining, key) + delete(c.accounts, key) + return + } + c.remaining[key] = remaining +} + type schedulerGroupLifecyclePlan struct { active bool tasks []schedulerBucketWriteTask @@ -535,8 +584,9 @@ func (s *SchedulerSnapshotService) reconcileGroupLifecycle(ctx context.Context, return err } if plan.active { + queries := newSchedulerAccountQueryCache(plan.tasks) for _, task := range plan.tasks { - if err := s.rebuildBucketWithTokenPolicy(ctx, task, "group_change", true); err != nil { + if err := s.rebuildBucketWithTokenPolicyAndQueryCache(ctx, task, "group_change", true, queries); err != nil { return err } } @@ -716,7 +766,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) - if err := s.rebuildPreparedBucketTasks(ctx, tasks, reason, false); err != nil && firstErr == nil { + queries := newSchedulerAccountQueryCache(tasks) + if err := s.rebuildPreparedBucketTasks(ctx, tasks, reason, false, queries); err != nil && firstErr == nil { firstErr = err } return firstErr @@ -744,17 +795,32 @@ func (s *SchedulerSnapshotService) prepareBucketWriteTasks(ctx context.Context, return tasks, firstErr } -func (s *SchedulerSnapshotService) rebuildPreparedBucketTasks(ctx context.Context, tasks []schedulerBucketWriteTask, reason string, strict bool) error { +func (s *SchedulerSnapshotService) rebuildPreparedBucketTasks( + ctx context.Context, + tasks []schedulerBucketWriteTask, + reason string, + strict bool, + queries *schedulerAccountQueryCache, +) error { var firstErr error for _, task := range tasks { - if err := s.rebuildBucketWithTokenPolicy(ctx, task, reason, strict); err != nil && firstErr == nil { + if err := s.rebuildBucketWithTokenPolicyAndQueryCache(ctx, task, reason, strict, queries); err != nil && firstErr == nil { firstErr = err } } return firstErr } -func (s *SchedulerSnapshotService) rebuildBucketWithTokenPolicy(ctx context.Context, task schedulerBucketWriteTask, reason string, strict bool) error { +func (s *SchedulerSnapshotService) rebuildBucketWithTokenPolicyAndQueryCache( + ctx context.Context, + task schedulerBucketWriteTask, + reason string, + strict bool, + queries *schedulerAccountQueryCache, +) error { + if queries != nil { + defer queries.release(task.bucket) + } if s.cache == nil { return ErrSchedulerCacheNotReady } @@ -776,7 +842,7 @@ func (s *SchedulerSnapshotService) rebuildBucketWithTokenPolicy(ctx context.Cont rebuildCtx, cancel := context.WithTimeout(ctx, 30*time.Second) defer cancel() - accounts, err := s.loadAccountsFromDB(rebuildCtx, bucket, bucket.Mode == SchedulerModeMixed) + accounts, err := s.loadAccountsForRebuild(rebuildCtx, bucket, queries) if err != nil { logger.LegacyPrintf("service.scheduler_snapshot", "[Scheduler] rebuild failed: bucket=%s reason=%s err=%v", bucket.String(), reason, err) return err @@ -975,10 +1041,11 @@ func (s *SchedulerSnapshotService) prepareAndRebuildFullSnapshot( return firstErr } captured = append(captured, ordinary...) - if err := s.rebuildPreparedBucketTasks(ctx, reopened, reason, true); err != nil { + queries := newSchedulerAccountQueryCache(reopened, captured) + if err := s.rebuildPreparedBucketTasks(ctx, reopened, reason, true, queries); err != nil { firstErr = err } - if err := s.rebuildPreparedBucketTasks(ctx, captured, reason, false); err != nil && firstErr == nil { + if err := s.rebuildPreparedBucketTasks(ctx, captured, reason, false, queries); err != nil && firstErr == nil { firstErr = err } return firstErr @@ -1141,6 +1208,30 @@ func (s *SchedulerSnapshotService) loadAccountsFromDB(ctx context.Context, bucke return s.accountRepo.ListSchedulableUngroupedByPlatform(ctx, bucket.Platform) } +func (s *SchedulerSnapshotService) loadAccountsForRebuild( + ctx context.Context, + bucket SchedulerBucket, + queries *schedulerAccountQueryCache, +) ([]Account, error) { + key, cacheable := schedulerAccountQueryKeyForBucket(bucket) + if queries == nil || !cacheable { + return s.loadAccountsFromDB(ctx, bucket, bucket.Mode == SchedulerModeMixed) + } + + if accounts, ok := queries.accounts[key]; ok { + return accounts, nil + } + if queries.remaining[key] <= 1 { + return s.loadAccountsFromDB(ctx, bucket, false) + } + accounts, err := s.loadAccountsFromDB(ctx, bucket, false) + if err != nil { + return nil, err + } + queries.accounts[key] = accounts + return accounts, nil +} + func (s *SchedulerSnapshotService) bucketFor(groupID *int64, platform string, mode string) SchedulerBucket { return SchedulerBucket{ GroupID: s.normalizeGroupID(groupID),