diff --git a/backend/internal/repository/scheduler_cache.go b/backend/internal/repository/scheduler_cache.go index b3cccefff9..fde3c06645 100644 --- a/backend/internal/repository/scheduler_cache.go +++ b/backend/internal/repository/scheduler_cache.go @@ -395,6 +395,43 @@ func (c *schedulerCache) SetSnapshot(ctx context.Context, bucket service.Schedul return c.activateSnapshotVersion(ctx, bucket, token, version) } +// SetSnapshotAndReturnAccountIDs 完整发布快照，并返回 writeAccounts 实际接受的有序账号 ID。 +// 该可选能力只供同一重建批次复用，返回前仍会完成版本激活与 fencing 校验。 +func (c *schedulerCache) SetSnapshotAndReturnAccountIDs(ctx context.Context, bucket service.SchedulerBucket, token service.SchedulerBucketWriteToken, accounts []service.Account) ([]int64, error) { + if !token.ValidFor(bucket) { + return nil, fmt.Errorf("%w: bucket=%s", service.ErrSchedulerBucketWriteFenced, bucket.String()) + } + // 分配版本与激活指针是两个 fencing 边界；中间写入的数据只有通过第二次校验才能发布。 + version, err := c.allocateSnapshotVersion(ctx, bucket, token) + if err != nil { + return nil, err + } + accountIDs, err := c.writeSnapshotVersionAndReturnAccountIDs(ctx, bucket, version, accounts) + if err != nil { + return nil, err + } + if err := c.activateSnapshotVersion(ctx, bucket, token, version); err != nil { + return nil, err + } + return accountIDs, nil +} + +// SetSnapshotByAccountIDs 复用同批次首次完整写入后得到的账号成员。 +// 每个桶仍独立分配版本、写入有序集合并执行激活 fencing，只省略重复的账号 JSON 与全局键写入。 +func (c *schedulerCache) SetSnapshotByAccountIDs(ctx context.Context, bucket service.SchedulerBucket, token service.SchedulerBucketWriteToken, accountIDs []int64) error { + if !token.ValidFor(bucket) { + return fmt.Errorf("%w: bucket=%s", service.ErrSchedulerBucketWriteFenced, bucket.String()) + } + version, err := c.allocateSnapshotVersion(ctx, bucket, token) + if err != nil { + return err + } + if err := c.writeSnapshotAccountIDs(ctx, bucket, version, accountIDs); err != nil { + return err + } + return c.activateSnapshotVersion(ctx, bucket, token, version) +} + func (c *schedulerCache) allocateSnapshotVersion(ctx context.Context, bucket service.SchedulerBucket, token service.SchedulerBucketWriteToken) (string, error) { result, err := allocateSnapshotVersionScript.Run(ctx, c.rdb, []string{ schedulerBucketKey(schedulerEpochPrefix, bucket), @@ -411,35 +448,74 @@ func (c *schedulerCache) allocateSnapshotVersion(ctx context.Context, bucket ser } func (c *schedulerCache) writeSnapshotVersion(ctx context.Context, bucket service.SchedulerBucket, version string, accounts []service.Account) error { - snapshotKey := schedulerSnapshotKey(bucket, version) cacheableAccounts, err := c.writeAccounts(ctx, accounts) if err != nil { return err } + return c.writeSnapshotAccounts(ctx, bucket, version, cacheableAccounts) +} - if len(cacheableAccounts) > 0 { - // 使用序号作为 score，保持数据库返回的排序语义。 - members := make([]redis.Z, 0, len(cacheableAccounts)) - for idx, account := range cacheableAccounts { - members = append(members, redis.Z{ - Score: float64(idx), - Member: strconv.FormatInt(account.ID, 10), - }) - } - pipe := c.rdb.Pipeline() - for start := 0; start < len(members); start += c.writeChunkSize { - end := start + c.writeChunkSize - if end > len(members) { - end = len(members) - } - pipe.ZAdd(ctx, snapshotKey, members[start:end]...) - } - if _, err := pipe.Exec(ctx); err != nil { - return err - } +func (c *schedulerCache) writeSnapshotVersionAndReturnAccountIDs(ctx context.Context, bucket service.SchedulerBucket, version string, accounts []service.Account) ([]int64, error) { + accountIDs, err := c.writeAccountIDs(ctx, accounts) + if err != nil { + return nil, err } + if err := c.writeSnapshotAccountIDs(ctx, bucket, version, accountIDs); err != nil { + return nil, err + } + return accountIDs, nil +} - return nil +func (c *schedulerCache) writeSnapshotAccounts(ctx context.Context, bucket service.SchedulerBucket, version string, accounts []service.Account) error { + if len(accounts) == 0 { + return nil + } + members := make([]redis.Z, 0, len(accounts)) + for idx, account := range accounts { + members = append(members, redis.Z{ + Score: float64(idx), + Member: strconv.FormatInt(account.ID, 10), + }) + } + return c.writeSnapshotMembers(ctx, bucket, version, members) +} + +func (c *schedulerCache) writeSnapshotAccountIDs(ctx context.Context, bucket service.SchedulerBucket, version string, accountIDs []int64) error { + members := schedulerSnapshotMembers(accountIDs) + return c.writeSnapshotMembers(ctx, bucket, version, members) +} + +func schedulerSnapshotMembers(accountIDs []int64) []redis.Z { + if len(accountIDs) == 0 { + return nil + } + // 使用序号作为 score，保持数据库返回的排序语义；重复 ID 继续交由 Redis ZADD + // 按最后一个 score 覆盖，与直接从账号切片构造成员时的行为一致。 + members := make([]redis.Z, 0, len(accountIDs)) + for idx, accountID := range accountIDs { + members = append(members, redis.Z{ + Score: float64(idx), + Member: strconv.FormatInt(accountID, 10), + }) + } + return members +} + +func (c *schedulerCache) writeSnapshotMembers(ctx context.Context, bucket service.SchedulerBucket, version string, members []redis.Z) error { + if len(members) == 0 { + return nil + } + snapshotKey := schedulerSnapshotKey(bucket, version) + pipe := c.rdb.Pipeline() + for start := 0; start < len(members); start += c.writeChunkSize { + end := start + c.writeChunkSize + if end > len(members) { + end = len(members) + } + pipe.ZAdd(ctx, snapshotKey, members[start:end]...) + } + _, err := pipe.Exec(ctx) + return err } func (c *schedulerCache) activateSnapshotVersion(ctx context.Context, bucket service.SchedulerBucket, token service.SchedulerBucketWriteToken, version string) error { @@ -644,12 +720,28 @@ func decodeCachedAccount(val any) (*service.Account, error) { } func (c *schedulerCache) writeAccounts(ctx context.Context, accounts []service.Account) ([]service.Account, error) { + cacheableAccounts, _, err := c.writeAccountPayloads(ctx, accounts, false) + return cacheableAccounts, err +} + +func (c *schedulerCache) writeAccountIDs(ctx context.Context, accounts []service.Account) ([]int64, error) { + _, accountIDs, err := c.writeAccountPayloads(ctx, accounts, true) + return accountIDs, err +} + +func (c *schedulerCache) writeAccountPayloads(ctx context.Context, accounts []service.Account, collectIDs bool) ([]service.Account, []int64, error) { if len(accounts) == 0 { - return nil, nil + return nil, nil, nil } pipe := c.rdb.Pipeline() - cacheableAccounts := make([]service.Account, 0, len(accounts)) + var cacheableAccounts []service.Account + var accountIDs []int64 + if collectIDs { + accountIDs = make([]int64, 0, len(accounts)) + } else { + cacheableAccounts = make([]service.Account, 0, len(accounts)) + } pending := 0 flush := func() error { if pending == 0 { @@ -676,19 +768,24 @@ func (c *schedulerCache) writeAccounts(ctx context.Context, accounts []service.A id := strconv.FormatInt(account.ID, 10) pipe.Set(ctx, schedulerAccountKey(id), fullPayload, 0) pipe.Set(ctx, schedulerAccountMetaKey(id), metaPayload, 0) - cacheableAccounts = append(cacheableAccounts, account) + // 复用路径只保留有序 ID，避免先物化完整账号切片再做第二次扫描。 + if collectIDs { + accountIDs = append(accountIDs, account.ID) + } else { + cacheableAccounts = append(cacheableAccounts, account) + } pending++ if pending >= c.writeChunkSize { if err := flush(); err != nil { - return nil, err + return nil, nil, err } } } if err := flush(); err != nil { - return nil, err + return nil, nil, err } - return cacheableAccounts, nil + return cacheableAccounts, accountIDs, nil } func marshalSchedulerCacheAccount(account service.Account) ([]byte, []byte, error) { diff --git a/backend/internal/repository/scheduler_cache_unit_test.go b/backend/internal/repository/scheduler_cache_unit_test.go index f933603c47..230dc59903 100644 --- a/backend/internal/repository/scheduler_cache_unit_test.go +++ b/backend/internal/repository/scheduler_cache_unit_test.go @@ -5,6 +5,9 @@ package repository import ( "context" "encoding/hex" + "encoding/json" + "fmt" + "strconv" "strings" "testing" "time" @@ -82,6 +85,170 @@ func TestSchedulerCacheUpdateLastUsedClearsUnencodablePayload(t *testing.T) { require.Nil(t, cached) } +func TestSchedulerCacheSnapshotAccountIDReusePreservesPayloadAndMembers(t *testing.T) { + ctx := context.Background() + cache, _ := newSchedulerCacheUnitWithRedis(t) + invalidTime := time.Date(10000, time.January, 1, 0, 0, 0, 0, time.UTC) + validOne := service.Account{ + ID: 701, + Name: "first", + Platform: service.PlatformOpenAI, + Type: service.AccountTypeOAuth, + Credentials: map[string]any{"model_mapping": map[string]any{"z": "last", "a": "first"}}, + Extra: map[string]any{"mixed_scheduling": true}, + GroupIDs: []int64{17}, + } + validTwo := service.Account{ID: 702, Name: "second", Platform: service.PlatformOpenAI, Type: service.AccountTypeAPIKey} + invalid := service.Account{ID: 799, Platform: service.PlatformOpenAI, Type: service.AccountTypeAPIKey, ExpiresAt: &invalidTime} + accounts := []service.Account{validOne, invalid, validTwo, validOne} + + single := service.SchedulerBucket{GroupID: 17, Platform: service.PlatformOpenAI, Mode: service.SchedulerModeSingle} + singleToken, err := cache.CaptureBucketWriteToken(ctx, single) + require.NoError(t, err) + accountIDs, err := cache.SetSnapshotAndReturnAccountIDs(ctx, single, singleToken, accounts) + require.NoError(t, err) + require.Equal(t, []int64{701, 702, 701}, accountIDs, "应保留可编码账号的原顺序和重复项") + + wantFull, err := json.Marshal(validOne) + require.NoError(t, err) + wantMeta, err := json.Marshal(buildSchedulerMetadataAccount(validOne)) + require.NoError(t, err) + fullBefore, err := cache.rdb.Get(ctx, schedulerAccountKey("701")).Bytes() + require.NoError(t, err) + metaBefore, err := cache.rdb.Get(ctx, schedulerAccountMetaKey("701")).Bytes() + require.NoError(t, err) + require.Equal(t, wantFull, fullBefore) + require.Equal(t, wantMeta, metaBefore) + + forced := service.SchedulerBucket{GroupID: 17, Platform: service.PlatformOpenAI, Mode: service.SchedulerModeForced} + forcedToken, err := cache.CaptureBucketWriteToken(ctx, forced) + require.NoError(t, err) + require.NoError(t, cache.SetSnapshotByAccountIDs(ctx, forced, forcedToken, accountIDs)) + + fullAfter, err := cache.rdb.Get(ctx, schedulerAccountKey("701")).Bytes() + require.NoError(t, err) + metaAfter, err := cache.rdb.Get(ctx, schedulerAccountMetaKey("701")).Bytes() + require.NoError(t, err) + require.Equal(t, fullBefore, fullAfter, "ID-only 路径不得重写完整账号键") + require.Equal(t, metaBefore, metaAfter, "ID-only 路径不得重写调度元数据键") + + for _, bucket := range []service.SchedulerBucket{single, forced} { + version, err := cache.rdb.Get(ctx, schedulerBucketKey(schedulerActivePrefix, bucket)).Result() + require.NoError(t, err) + members, err := cache.rdb.ZRange(ctx, schedulerSnapshotKey(bucket, version), 0, -1).Result() + require.NoError(t, err) + require.Equal(t, []string{"702", "701"}, members, bucket.String()) + } + missing, err := cache.GetAccount(ctx, invalid.ID) + require.NoError(t, err) + require.Nil(t, missing) +} + +func TestSchedulerCacheSnapshotAccountIDReuseKeepsEmptySnapshotSemantics(t *testing.T) { + ctx := context.Background() + cache := newSchedulerCacheUnit(t) + invalidTime := time.Date(10000, time.January, 1, 0, 0, 0, 0, time.UTC) + accounts := []service.Account{{ID: 811, Platform: service.PlatformOpenAI, ExpiresAt: &invalidTime}} + + single := service.SchedulerBucket{GroupID: 18, Platform: service.PlatformOpenAI, Mode: service.SchedulerModeSingle} + singleToken, err := cache.CaptureBucketWriteToken(ctx, single) + require.NoError(t, err) + accountIDs, err := cache.SetSnapshotAndReturnAccountIDs(ctx, single, singleToken, accounts) + require.NoError(t, err) + require.Empty(t, accountIDs) + + forced := service.SchedulerBucket{GroupID: 18, Platform: service.PlatformOpenAI, Mode: service.SchedulerModeForced} + forcedToken, err := cache.CaptureBucketWriteToken(ctx, forced) + require.NoError(t, err) + require.NoError(t, cache.SetSnapshotByAccountIDs(ctx, forced, forcedToken, accountIDs)) + + for _, bucket := range []service.SchedulerBucket{single, forced} { + ready, err := cache.rdb.Get(ctx, schedulerBucketKey(schedulerReadyPrefix, bucket)).Result() + require.NoError(t, err) + require.Equal(t, "1", ready) + snapshot, hit, err := cache.GetSnapshot(ctx, bucket) + require.NoError(t, err) + require.False(t, hit, bucket.String()) + require.Nil(t, snapshot) + } +} + +func TestSchedulerCacheSetSnapshotByAccountIDsKeepsFencing(t *testing.T) { + ctx := context.Background() + cache := newSchedulerCacheUnit(t) + bucket := service.SchedulerBucket{GroupID: 19, Platform: service.PlatformOpenAI, Mode: service.SchedulerModeForced} + + err := cache.SetSnapshotByAccountIDs(ctx, bucket, service.SchedulerBucketWriteToken{}, []int64{901}) + require.ErrorIs(t, err, service.ErrSchedulerBucketWriteFenced) + _, err = cache.rdb.Get(ctx, schedulerBucketKey(schedulerVersionPrefix, bucket)).Result() + require.ErrorIs(t, err, redis.Nil) + + token, err := cache.CaptureBucketWriteToken(ctx, bucket) + require.NoError(t, err) + require.NoError(t, cache.RetireBucket(ctx, bucket)) + err = cache.SetSnapshotByAccountIDs(ctx, bucket, token, []int64{901}) + require.ErrorIs(t, err, service.ErrSchedulerBucketRetired) +} + +func TestSchedulerCacheSetSnapshotByAccountIDsDoesNotResurrectDeletedAccount(t *testing.T) { + ctx := context.Background() + cache := newSchedulerCacheUnit(t) + account := service.Account{ID: 902, Platform: service.PlatformOpenAI, Type: service.AccountTypeOAuth} + single := service.SchedulerBucket{GroupID: 20, Platform: service.PlatformOpenAI, Mode: service.SchedulerModeSingle} + singleToken, err := cache.CaptureBucketWriteToken(ctx, single) + require.NoError(t, err) + accountIDs, err := cache.SetSnapshotAndReturnAccountIDs(ctx, single, singleToken, []service.Account{account}) + require.NoError(t, err) + require.Equal(t, []int64{account.ID}, accountIDs) + require.NoError(t, cache.DeleteAccount(ctx, account.ID)) + + forced := service.SchedulerBucket{GroupID: 20, Platform: service.PlatformOpenAI, Mode: service.SchedulerModeForced} + forcedToken, err := cache.CaptureBucketWriteToken(ctx, forced) + require.NoError(t, err) + require.NoError(t, cache.SetSnapshotByAccountIDs(ctx, forced, forcedToken, accountIDs)) + + full, err := cache.GetAccount(ctx, account.ID) + require.NoError(t, err) + require.Nil(t, full, "ID-only 发布不得复活已删除的完整账号键") + snapshot, hit, err := cache.GetSnapshot(ctx, forced) + require.NoError(t, err) + require.False(t, hit, "元数据缺失时必须安全回源，而不是返回残缺快照") + require.Nil(t, snapshot) +} + +func TestMarshalSchedulerCacheAccountKeepsEncodingJSONWireFormat(t *testing.T) { + cases := []struct { + name string + account service.Account + }{ + {name: "nil collections", account: service.Account{ID: 801}}, + {name: "empty collections", account: service.Account{ + ID: 802, + Credentials: map[string]any{}, + Extra: map[string]any{}, + GroupIDs: []int64{}, + Groups: []*service.Group{}, + }}, + {name: "nested maps and escaping", account: service.Account{ + ID: 803, + Credentials: map[string]any{"model_mapping": map[string]any{"z": "", "a": "&first"}}, + Extra: map[string]any{"mixed_scheduling": true}, + }}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + full, meta, err := marshalSchedulerCacheAccount(tc.account) + require.NoError(t, err) + wantFull, err := json.Marshal(tc.account) + require.NoError(t, err) + wantMeta, err := json.Marshal(buildSchedulerMetadataAccount(tc.account)) + require.NoError(t, err) + require.Equal(t, wantFull, full) + require.Equal(t, wantMeta, meta) + }) + } +} + func TestBuildSchedulerMetadataAccount_KeepsOpenAIWSFlags(t *testing.T) { account := service.Account{ ID: 42, @@ -611,3 +778,126 @@ func TestSchedulerCacheGroupLifecycleLeaseRejectsInvalidInput(t *testing.T) { require.NoError(t, err) require.Zero(t, keys) } + +var schedulerCachePayloadBenchmarkSink int + +func BenchmarkSchedulerCacheAccountPayloadReuse(b *testing.B) { + for _, size := range []int{1, 100, 10_000} { + accounts := schedulerCacheBenchmarkAccounts(size) + b.Run(fmt.Sprintf("pair_baseline_%d_accounts", size), func(b *testing.B) { + b.ReportAllocs() + for i := 0; i < b.N; i++ { + first, err := benchmarkSchedulerLegacySnapshotPayload(accounts) + if err != nil { + b.Fatal(err) + } + second, err := benchmarkSchedulerLegacySnapshotPayload(accounts) + if err != nil { + b.Fatal(err) + } + schedulerCachePayloadBenchmarkSink = first + second + } + }) + b.Run(fmt.Sprintf("pair_reuse_%d_accounts", size), func(b *testing.B) { + b.ReportAllocs() + for i := 0; i < b.N; i++ { + ids, total, err := benchmarkSchedulerReusableSnapshotPayload(accounts) + if err != nil { + b.Fatal(err) + } + // 第二个桶仍构造成员，只跳过账号 JSON 与全局账号键。 + total += len(schedulerSnapshotMembers(ids)) + schedulerCachePayloadBenchmarkSink = total + } + }) + b.Run(fmt.Sprintf("first_baseline_%d_accounts", size), func(b *testing.B) { + b.ReportAllocs() + for i := 0; i < b.N; i++ { + total, err := benchmarkSchedulerLegacySnapshotPayload(accounts) + if err != nil { + b.Fatal(err) + } + schedulerCachePayloadBenchmarkSink = total + } + }) + b.Run(fmt.Sprintf("first_reuse_%d_accounts", size), func(b *testing.B) { + b.ReportAllocs() + for i := 0; i < b.N; i++ { + ids, total, err := benchmarkSchedulerReusableSnapshotPayload(accounts) + if err != nil { + b.Fatal(err) + } + total += len(ids) + schedulerCachePayloadBenchmarkSink = total + } + }) + } +} + +func benchmarkSchedulerLegacySnapshotPayload(accounts []service.Account) (int, error) { + cacheable := make([]service.Account, 0, len(accounts)) + total := 0 + for _, account := range accounts { + full, meta, err := marshalSchedulerCacheAccount(account) + if err != nil { + continue + } + total += len(full) + len(meta) + cacheable = append(cacheable, account) + } + members := make([]redis.Z, 0, len(cacheable)) + for idx, account := range cacheable { + members = append(members, redis.Z{Score: float64(idx), Member: strconv.FormatInt(account.ID, 10)}) + } + return total + len(members), nil +} + +func benchmarkSchedulerReusableSnapshotPayload(accounts []service.Account) ([]int64, int, error) { + accountIDs := make([]int64, 0, len(accounts)) + total := 0 + for _, account := range accounts { + full, meta, err := marshalSchedulerCacheAccount(account) + if err != nil { + continue + } + total += len(full) + len(meta) + accountIDs = append(accountIDs, account.ID) + } + total += len(schedulerSnapshotMembers(accountIDs)) + return accountIDs, total, nil +} + +func schedulerCacheBenchmarkAccounts(size int) []service.Account { + largeValue := strings.Repeat("x", 4096) + credentials := map[string]any{ + "api_key": "benchmark-key", + "model_mapping": map[string]any{"z-model": "z-target", "a-model": "a-target"}, + "large_value": largeValue, + } + extra := map[string]any{ + "mixed_scheduling": true, + "model_rate_limits": map[string]any{ + "z-model": map[string]any{"rate_limit_reset_at": "2026-07-16T00:00:00Z"}, + "a-model": map[string]any{"rate_limit_reset_at": "2026-07-16T00:00:00Z"}, + }, + "large_value": largeValue, + } + accounts := make([]service.Account, size) + for i := range accounts { + id := int64(i + 1) + accounts[i] = service.Account{ + ID: id, + Name: "benchmark-account", + Platform: service.PlatformOpenAI, + Type: service.AccountTypeOAuth, + Credentials: credentials, + Extra: extra, + GroupIDs: []int64{7, 9}, + AccountGroups: []service.AccountGroup{ + {AccountID: id, GroupID: 7, Priority: 1}, + {AccountID: id, GroupID: 9, Priority: 2}, + }, + } + } + return accounts +} diff --git a/backend/internal/service/scheduler_snapshot_batch_query_test.go b/backend/internal/service/scheduler_snapshot_batch_query_test.go index a2e7ae8e85..221610bdbf 100644 --- a/backend/internal/service/scheduler_snapshot_batch_query_test.go +++ b/backend/internal/service/scheduler_snapshot_batch_query_test.go @@ -116,6 +116,72 @@ type batchSnapshotCache struct { beforeSet func() } +type batchSnapshotAccountIDCache struct { + *batchSnapshotCache + + reuseMu sync.Mutex + fullCalls map[SchedulerBucket]int + idOnlyCalls map[SchedulerBucket]int + idOnlyError map[SchedulerBucket]error + fullLateErr map[SchedulerBucket]error + returnEmpty bool +} + +func newBatchSnapshotAccountIDCache() *batchSnapshotAccountIDCache { + return &batchSnapshotAccountIDCache{ + batchSnapshotCache: newBatchSnapshotCache(), + fullCalls: make(map[SchedulerBucket]int), + idOnlyCalls: make(map[SchedulerBucket]int), + idOnlyError: make(map[SchedulerBucket]error), + fullLateErr: make(map[SchedulerBucket]error), + } +} + +func (c *batchSnapshotAccountIDCache) SetSnapshotAndReturnAccountIDs(ctx context.Context, bucket SchedulerBucket, token SchedulerBucketWriteToken, accounts []Account) ([]int64, error) { + c.reuseMu.Lock() + c.fullCalls[bucket]++ + c.reuseMu.Unlock() + if err := c.batchSnapshotCache.SetSnapshot(ctx, bucket, token, accounts); err != nil { + return nil, err + } + c.reuseMu.Lock() + lateErr := c.fullLateErr[bucket] + returnEmpty := c.returnEmpty + c.reuseMu.Unlock() + if lateErr != nil { + return nil, lateErr + } + if returnEmpty { + return []int64{}, nil + } + ids := make([]int64, 0, len(accounts)) + for _, account := range accounts { + ids = append(ids, account.ID) + } + return ids, nil +} + +func (c *batchSnapshotAccountIDCache) SetSnapshotByAccountIDs(ctx context.Context, bucket SchedulerBucket, token SchedulerBucketWriteToken, accountIDs []int64) error { + c.reuseMu.Lock() + c.idOnlyCalls[bucket]++ + err := c.idOnlyError[bucket] + c.reuseMu.Unlock() + if err != nil { + return err + } + accounts := make([]Account, 0, len(accountIDs)) + for _, id := range accountIDs { + accounts = append(accounts, Account{ID: id}) + } + return c.batchSnapshotCache.SetSnapshot(ctx, bucket, token, accounts) +} + +func (c *batchSnapshotAccountIDCache) reuseCounts(bucket SchedulerBucket) (full, idOnly int) { + c.reuseMu.Lock() + defer c.reuseMu.Unlock() + return c.fullCalls[bucket], c.idOnlyCalls[bucket] +} + func newBatchSnapshotCache() *batchSnapshotCache { return &batchSnapshotCache{ captured: make(map[SchedulerBucket]SchedulerBucketWriteToken), @@ -231,6 +297,186 @@ func TestSchedulerRebuildBatchReusesSingleForcedQueryAndKeepsSnapshotsIndependen } } +func TestSchedulerRebuildBatchReusesAccountPayloadForSingleForced(t *testing.T) { + const groupID int64 = 211 + single := SchedulerBucket{GroupID: groupID, Platform: PlatformOpenAI, Mode: SchedulerModeSingle} + forced := SchedulerBucket{GroupID: groupID, Platform: PlatformOpenAI, Mode: SchedulerModeForced} + cache := newBatchSnapshotAccountIDCache() + repo := newBatchAccountQueryRepo() + svc := newBatchQueryTestService(cache, repo, config.RunModeStandard) + + for run := 1; run <= 2; run++ { + require.NoError(t, svc.rebuildBuckets(context.Background(), []SchedulerBucket{single, forced}, "reuse")) + full, idOnly := cache.reuseCounts(single) + require.Equal(t, run, full) + require.Zero(t, idOnly) + full, idOnly = cache.reuseCounts(forced) + require.Zero(t, full) + require.Equal(t, run, idOnly) + } + require.Equal(t, 2, repo.callCount(batchAccountQueryKey{groupID: groupID, platform: PlatformOpenAI}), "账号载荷不得跨重建批次复用") +} + +func TestSchedulerRebuildBatchDoesNotReuseAccountPayloadAfterFirstWriterFailure(t *testing.T) { + const groupID int64 = 212 + single := SchedulerBucket{GroupID: groupID, Platform: PlatformOpenAI, Mode: SchedulerModeSingle} + forced := SchedulerBucket{GroupID: groupID, Platform: PlatformOpenAI, Mode: SchedulerModeForced} + wantErr := errors.New("snapshot write failed") + cache := newBatchSnapshotAccountIDCache() + cache.setErrors[single] = wantErr + svc := newBatchQueryTestService(cache, newBatchAccountQueryRepo(), config.RunModeStandard) + + err := svc.rebuildBuckets(context.Background(), []SchedulerBucket{single, forced}, "failure") + require.ErrorIs(t, err, wantErr) + full, idOnly := cache.reuseCounts(single) + require.Equal(t, 1, full) + require.Zero(t, idOnly) + full, idOnly = cache.reuseCounts(forced) + require.Zero(t, full) + require.Zero(t, idOnly) + _, attempts, _, writes := cache.bucketState(forced) + require.Equal(t, 1, attempts, "首次完整写失败后，后续桶必须走原 SetSnapshot") + require.Len(t, writes, 1) +} + +func TestSchedulerRebuildBatchDoesNotReuseAccountPayloadAfterLateFirstWriterFailure(t *testing.T) { + const groupID int64 = 216 + single := SchedulerBucket{GroupID: groupID, Platform: PlatformOpenAI, Mode: SchedulerModeSingle} + forced := SchedulerBucket{GroupID: groupID, Platform: PlatformOpenAI, Mode: SchedulerModeForced} + wantErr := errors.New("snapshot activation failed") + cache := newBatchSnapshotAccountIDCache() + cache.fullLateErr[single] = wantErr + svc := newBatchQueryTestService(cache, newBatchAccountQueryRepo(), config.RunModeStandard) + + err := svc.rebuildBuckets(context.Background(), []SchedulerBucket{single, forced}, "late-failure") + require.ErrorIs(t, err, wantErr) + full, idOnly := cache.reuseCounts(single) + require.Equal(t, 1, full) + require.Zero(t, idOnly) + full, idOnly = cache.reuseCounts(forced) + require.Zero(t, full) + require.Zero(t, idOnly) + _, attempts, _, writes := cache.bucketState(forced) + require.Equal(t, 1, attempts, "首次激活失败后不得登记可复用 ID") + require.Len(t, writes, 1) +} + +func TestSchedulerRebuildBatchDoesNotReuseAccountPayloadAfterLockBusy(t *testing.T) { + const groupID int64 = 213 + single := SchedulerBucket{GroupID: groupID, Platform: PlatformOpenAI, Mode: SchedulerModeSingle} + forced := SchedulerBucket{GroupID: groupID, Platform: PlatformOpenAI, Mode: SchedulerModeForced} + cache := newBatchSnapshotAccountIDCache() + cache.lockBusy[single] = true + svc := newBatchQueryTestService(cache, newBatchAccountQueryRepo(), config.RunModeStandard) + + require.NoError(t, svc.rebuildBuckets(context.Background(), []SchedulerBucket{single, forced}, "busy")) + full, idOnly := cache.reuseCounts(single) + require.Zero(t, full) + require.Zero(t, idOnly) + full, idOnly = cache.reuseCounts(forced) + require.Zero(t, full) + require.Zero(t, idOnly) + _, attempts, _, writes := cache.bucketState(forced) + require.Equal(t, 1, attempts) + require.Len(t, writes, 1) +} + +func TestSchedulerRebuildBatchKeepsMixedAndDifferentQueriesOnFullWrites(t *testing.T) { + const groupID int64 = 214 + openAISingle := SchedulerBucket{GroupID: groupID, Platform: PlatformOpenAI, Mode: SchedulerModeSingle} + openAIForced := SchedulerBucket{GroupID: groupID, Platform: PlatformOpenAI, Mode: SchedulerModeForced} + anthropicSingle := SchedulerBucket{GroupID: groupID, Platform: PlatformAnthropic, Mode: SchedulerModeSingle} + anthropicMixed := SchedulerBucket{GroupID: groupID, Platform: PlatformAnthropic, Mode: SchedulerModeMixed} + cache := newBatchSnapshotAccountIDCache() + svc := newBatchQueryTestService(cache, newBatchAccountQueryRepo(), config.RunModeStandard) + + require.NoError(t, svc.rebuildBuckets(context.Background(), []SchedulerBucket{openAISingle, openAIForced, anthropicSingle, anthropicMixed}, "scope")) + full, idOnly := cache.reuseCounts(openAISingle) + require.Equal(t, 1, full) + require.Zero(t, idOnly) + full, idOnly = cache.reuseCounts(openAIForced) + require.Zero(t, full) + require.Equal(t, 1, idOnly) + full, idOnly = cache.reuseCounts(anthropicSingle) + require.Zero(t, full) + require.Zero(t, idOnly) + _, attempts, _, writes := cache.bucketState(anthropicSingle) + require.Equal(t, 1, attempts) + require.Len(t, writes, 1) + full, idOnly = cache.reuseCounts(anthropicMixed) + require.Zero(t, full) + require.Zero(t, idOnly) + _, attempts, _, _ = cache.bucketState(anthropicMixed) + require.Equal(t, 1, attempts, "mixed 桶必须继续走原 SetSnapshot") +} + +func TestSchedulerRebuildBatchPropagatesAccountIDOnlyWriteFailure(t *testing.T) { + const groupID int64 = 215 + single := SchedulerBucket{GroupID: groupID, Platform: PlatformOpenAI, Mode: SchedulerModeSingle} + forced := SchedulerBucket{GroupID: groupID, Platform: PlatformOpenAI, Mode: SchedulerModeForced} + wantErr := errors.New("id-only write failed") + cache := newBatchSnapshotAccountIDCache() + cache.idOnlyError[forced] = wantErr + svc := newBatchQueryTestService(cache, newBatchAccountQueryRepo(), config.RunModeStandard) + + err := svc.rebuildBuckets(context.Background(), []SchedulerBucket{single, forced}, "id-error") + require.ErrorIs(t, err, wantErr) + full, idOnly := cache.reuseCounts(forced) + require.Zero(t, full, "ID-only 失败不得静默回退为完整写") + require.Equal(t, 1, idOnly) +} + +func TestSchedulerRebuildBatchReusesSuccessfulEmptyAccountIDs(t *testing.T) { + const groupID int64 = 217 + single := SchedulerBucket{GroupID: groupID, Platform: PlatformOpenAI, Mode: SchedulerModeSingle} + forced := SchedulerBucket{GroupID: groupID, Platform: PlatformOpenAI, Mode: SchedulerModeForced} + cache := newBatchSnapshotAccountIDCache() + cache.returnEmpty = true + svc := newBatchQueryTestService(cache, newBatchAccountQueryRepo(), config.RunModeStandard) + + require.NoError(t, svc.rebuildBuckets(context.Background(), []SchedulerBucket{single, forced}, "empty")) + full, idOnly := cache.reuseCounts(single) + require.Equal(t, 1, full) + require.Zero(t, idOnly) + full, idOnly = cache.reuseCounts(forced) + require.Zero(t, full) + require.Equal(t, 1, idOnly, "已成功缓存的空 ID 集也必须通过 map presence 复用") + _, _, _, writes := cache.bucketState(forced) + require.Len(t, writes, 1) + require.Empty(t, writes[0].accounts) +} + +func TestSchedulerRebuildBatchReusesAccountPayloadForSimpleGroupZero(t *testing.T) { + single := SchedulerBucket{GroupID: 0, Platform: PlatformOpenAI, Mode: SchedulerModeSingle} + forced := SchedulerBucket{GroupID: 0, Platform: PlatformOpenAI, Mode: SchedulerModeForced} + cache := newBatchSnapshotAccountIDCache() + svc := newBatchQueryTestService(cache, newBatchAccountQueryRepo(), config.RunModeSimple) + + require.NoError(t, svc.rebuildBuckets(context.Background(), []SchedulerBucket{single, forced}, "simple")) + full, idOnly := cache.reuseCounts(single) + require.Equal(t, 1, full) + require.Zero(t, idOnly) + full, idOnly = cache.reuseCounts(forced) + require.Zero(t, full) + require.Equal(t, 1, idOnly) +} + +func TestSchedulerAccountQueryCacheReleasesSnapshotAccountIDs(t *testing.T) { + single := schedulerBucketWriteTask{bucket: SchedulerBucket{GroupID: 218, Platform: PlatformOpenAI, Mode: SchedulerModeSingle}} + forced := schedulerBucketWriteTask{bucket: SchedulerBucket{GroupID: 218, Platform: PlatformOpenAI, Mode: SchedulerModeForced}} + queries := newSchedulerAccountQueryCache([]schedulerBucketWriteTask{single, forced}) + key, ok := schedulerAccountQueryKeyForBucket(single.bucket) + require.True(t, ok) + queries.snapshotAccountIDs[key] = []int64{1, 2} + + queries.release(single.bucket) + require.Contains(t, queries.snapshotAccountIDs, key) + queries.release(forced.bucket) + require.NotContains(t, queries.snapshotAccountIDs, key) + require.Empty(t, queries.remaining) + require.Empty(t, queries.accounts) +} + func TestSchedulerRebuildBatchKeepsMixedAndDifferentKeysIndependent(t *testing.T) { const groupID int64 = 202 buckets := []SchedulerBucket{ diff --git a/backend/internal/service/scheduler_snapshot_service.go b/backend/internal/service/scheduler_snapshot_service.go index a5bccdb265..4add904399 100644 --- a/backend/internal/service/scheduler_snapshot_service.go +++ b/backend/internal/service/scheduler_snapshot_service.go @@ -55,14 +55,24 @@ type schedulerAccountQueryKey struct { // mixed 与历史模式保持独立。每个 task 都用 defer 消费 remaining，最后一个消费者会立即释放结果， // 避免把账号切片的生命周期扩大到整轮 full rebuild。 type schedulerAccountQueryCache struct { - remaining map[schedulerAccountQueryKey]int - accounts map[schedulerAccountQueryKey][]Account + remaining map[schedulerAccountQueryKey]int + accounts map[schedulerAccountQueryKey][]Account + snapshotAccountIDs map[schedulerAccountQueryKey][]int64 +} + +// schedulerSnapshotAccountIDWriter 是 SchedulerCache 的可选批次优化能力。 +// 首次完整发布成功后返回实际可编码账号 ID；同一查询结果的后续桶只需发布这些 ID， +// 避免重复序列化并覆盖全局账号缓存。未实现该接口的缓存继续走原 SetSnapshot 路径。 +type schedulerSnapshotAccountIDWriter interface { + SetSnapshotAndReturnAccountIDs(ctx context.Context, bucket SchedulerBucket, token SchedulerBucketWriteToken, accounts []Account) ([]int64, error) + SetSnapshotByAccountIDs(ctx context.Context, bucket SchedulerBucket, token SchedulerBucketWriteToken, accountIDs []int64) error } func newSchedulerAccountQueryCache(taskSets ...[]schedulerBucketWriteTask) *schedulerAccountQueryCache { queries := &schedulerAccountQueryCache{ - remaining: make(map[schedulerAccountQueryKey]int), - accounts: make(map[schedulerAccountQueryKey][]Account), + remaining: make(map[schedulerAccountQueryKey]int), + accounts: make(map[schedulerAccountQueryKey][]Account), + snapshotAccountIDs: make(map[schedulerAccountQueryKey][]int64), } for _, tasks := range taskSets { for _, task := range tasks { @@ -93,6 +103,7 @@ func (c *schedulerAccountQueryCache) release(bucket SchedulerBucket) { if remaining <= 0 { delete(c.remaining, key) delete(c.accounts, key) + delete(c.snapshotAccountIDs, key) return } c.remaining[key] = remaining @@ -869,7 +880,7 @@ func (s *SchedulerSnapshotService) rebuildBucketWithTokenPolicyAndQueryCache( logger.LegacyPrintf("service.scheduler_snapshot", "[Scheduler] rebuild failed: bucket=%s reason=%s err=%v", bucket.String(), reason, err) return err } - if err := s.cache.SetSnapshot(rebuildCtx, bucket, task.token, accounts); err != nil { + if err := s.setRebuildSnapshot(rebuildCtx, task, accounts, queries); err != nil { if errors.Is(err, ErrSchedulerBucketRetired) || errors.Is(err, ErrSchedulerBucketWriteFenced) { slog.Debug("[Scheduler] rebuild fenced", "bucket", bucket.String(), "reason", reason) if strict { @@ -884,6 +895,38 @@ func (s *SchedulerSnapshotService) rebuildBucketWithTokenPolicyAndQueryCache( return nil } +func (s *SchedulerSnapshotService) setRebuildSnapshot( + ctx context.Context, + task schedulerBucketWriteTask, + accounts []Account, + queries *schedulerAccountQueryCache, +) error { + writer, ok := s.cache.(schedulerSnapshotAccountIDWriter) + key, reusable := schedulerAccountQueryKeyForBucket(task.bucket) + if !ok || queries == nil || !reusable { + return s.cache.SetSnapshot(ctx, task.bucket, task.token, accounts) + } + + if accountIDs, exists := queries.snapshotAccountIDs[key]; exists { + return writer.SetSnapshotByAccountIDs(ctx, task.bucket, task.token, accountIDs) + } + if queries.remaining[key] <= 1 { + return s.cache.SetSnapshot(ctx, task.bucket, task.token, accounts) + } + + accountIDs, err := writer.SetSnapshotAndReturnAccountIDs(ctx, task.bucket, task.token, accounts) + if err != nil { + return err + } + if queries.remaining[key] > 1 { + // 必须保存 writeAccounts 实际接受的有序 ID，不能从原账号切片重新推导； + // 否则不可编码账号会只出现在后续桶中，破坏两个快照的成员一致性。 + // 返回切片由当前批次独占，直接接管可避免 10k 账号场景再次复制。 + queries.snapshotAccountIDs[key] = accountIDs + } + return nil +} + func (s *SchedulerSnapshotService) triggerFullRebuild(reason string) error { if s.cache == nil { return ErrSchedulerCacheNotReady