复用重建批次内的账号查询

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)
+				}
+			}
+		})
+	}
+}

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)

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())
 			}
 		})
 	}

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),