diff --git a/backend/internal/config/config.go b/backend/internal/config/config.go index 9ed78dd3e..1d9c87a17 100644 --- a/backend/internal/config/config.go +++ b/backend/internal/config/config.go @@ -1071,6 +1071,10 @@ type GatewaySchedulingConfig struct { SnapshotMGetChunkSize int `mapstructure:"snapshot_mget_chunk_size"` // 快照重建时的缓存写入分块大小 SnapshotWriteChunkSize int `mapstructure:"snapshot_write_chunk_size"` + // 显式启用 Redis 候选索引的调度桶，格式为 "groupID:platform:mode"。 + IndexedBuckets []string `mapstructure:"indexed_buckets"` + // 候选索引一次最多返回的账号数。 + IndexedCandidateLimit int `mapstructure:"indexed_candidate_limit"` // 过期槽位清理周期（0 表示禁用） SlotCleanupInterval time.Duration `mapstructure:"slot_cleanup_interval"` @@ -1982,6 +1986,8 @@ func setDefaults() { viper.SetDefault("gateway.scheduling.load_batch_enabled", true) viper.SetDefault("gateway.scheduling.snapshot_mget_chunk_size", 128) viper.SetDefault("gateway.scheduling.snapshot_write_chunk_size", 256) + viper.SetDefault("gateway.scheduling.indexed_buckets", []string{}) + viper.SetDefault("gateway.scheduling.indexed_candidate_limit", 256) viper.SetDefault("gateway.scheduling.slot_cleanup_interval", 30*time.Second) viper.SetDefault("gateway.scheduling.db_fallback_enabled", true) viper.SetDefault("gateway.scheduling.db_fallback_timeout_seconds", 0) @@ -2849,6 +2855,14 @@ func (c *Config) Validate() error { if c.Gateway.Scheduling.SnapshotWriteChunkSize <= 0 { return fmt.Errorf("gateway.scheduling.snapshot_write_chunk_size must be positive") } + for _, rawBucket := range c.Gateway.Scheduling.IndexedBuckets { + if !isValidSchedulerBucketString(rawBucket) { + return fmt.Errorf("gateway.scheduling.indexed_buckets contains invalid bucket %q", rawBucket) + } + } + if c.Gateway.Scheduling.IndexedCandidateLimit <= 0 { + return fmt.Errorf("gateway.scheduling.indexed_candidate_limit must be positive") + } if c.Gateway.Scheduling.SlotCleanupInterval < 0 { return fmt.Errorf("gateway.scheduling.slot_cleanup_interval must be non-negative") } @@ -2902,6 +2916,17 @@ func (c *Config) Validate() error { return nil } +func isValidSchedulerBucketString(raw string) bool { + parts := strings.Split(raw, ":") + if len(parts) != 3 { + return false + } + if _, err := strconv.ParseInt(parts[0], 10, 64); err != nil { + return false + } + return strings.TrimSpace(parts[1]) != "" && strings.TrimSpace(parts[2]) != "" +} + func normalizeStringSlice(values []string) []string { if len(values) == 0 { return values diff --git a/backend/internal/repository/scheduler_cache.go b/backend/internal/repository/scheduler_cache.go index 213e5641e..5d8bebb00 100644 --- a/backend/internal/repository/scheduler_cache.go +++ b/backend/internal/repository/scheduler_cache.go @@ -4,6 +4,7 @@ import ( "context" "encoding/json" "fmt" + "math/rand" "strconv" "sync" "time" @@ -13,19 +14,27 @@ import ( ) const ( - schedulerBucketSetKey = "sched:buckets" - schedulerOutboxWatermarkKey = "sched:outbox:watermark" - schedulerAccountPrefix = "sched:acc:" - schedulerAccountMetaPrefix = "sched:meta:" - schedulerActivePrefix = "sched:active:" - schedulerReadyPrefix = "sched:ready:" - schedulerVersionPrefix = "sched:ver:" - schedulerSnapshotPrefix = "sched:" - schedulerLockPrefix = "sched:lock:" + schedulerBucketSetKey = "sched:buckets" + schedulerOutboxWatermarkKey = "sched:outbox:watermark" + schedulerAccountPrefix = "sched:acc:" + schedulerAccountMetaPrefix = "sched:meta:" + schedulerActivePrefix = "sched:active:" + schedulerReadyPrefix = "sched:ready:" + schedulerVersionPrefix = "sched:ver:" + schedulerSnapshotPrefix = "sched:" + schedulerLockPrefix = "sched:lock:" + schedulerCandidateIndexPrefix = "schedidx:v1:" + schedulerCandidateActivePrefix = "schedidx:active:" + schedulerCandidateReadyPrefix = "schedidx:ready:" + schedulerCandidateMetaPrefix = "schedidx:meta:" defaultSchedulerSnapshotMGetChunkSize = 128 defaultSchedulerSnapshotWriteChunkSize = 256 defaultSchedulerSnapshotLocalCacheTTL = 5 * time.Second + defaultSchedulerCandidateLimit = 256 + defaultSchedulerCandidateShardTarget = 512 + maxSchedulerCandidateShards = 128 + minSchedulerCandidateShardSize = 5000 // snapshotGraceTTLSeconds 旧快照过期的宽限期（秒）。 // 替代立即 DEL，让正在读取旧版本的 reader 有足够时间完成 ZRANGE。 @@ -86,6 +95,44 @@ redis.call('DEL', KEYS[2]) redis.call('SREM', KEYS[3], ARGV[2]) return 1 +`) + activateCandidateIndexScript = redis.NewScript(` +local currentActive = redis.call('GET', KEYS[1]) +local newVersion = tonumber(ARGV[1]) + +if currentActive ~= false then + local curVersion = tonumber(currentActive) + if curVersion and newVersion < curVersion then + return {0, currentActive} + end +end + +redis.call('SET', KEYS[1], ARGV[1]) +redis.call('SET', KEYS[2], '1') + +if currentActive == false then + return {1, ''} +end +return {1, currentActive} +`) + clearCandidateIndexScript = redis.NewScript(` +local currentActive = redis.call('GET', KEYS[1]) +local clearVersion = tonumber(ARGV[1]) + +if currentActive ~= false then + local curVersion = tonumber(currentActive) + if curVersion and clearVersion and curVersion > clearVersion then + return {0, currentActive} + end +end + +redis.call('DEL', KEYS[1]) +redis.call('DEL', KEYS[2]) + +if currentActive == false then + return {1, ''} +end +return {1, currentActive} `) ) @@ -93,6 +140,7 @@ type schedulerCache struct { rdb *redis.Client mgetChunkSize int writeChunkSize int + indexedBuckets map[string]struct{} localMu sync.RWMutex localSnapshots map[string]schedulerLocalSnapshot localBuckets map[int64]map[string]struct{} @@ -110,16 +158,27 @@ func NewSchedulerCache(rdb *redis.Client) service.SchedulerCache { } func newSchedulerCacheWithChunkSizes(rdb *redis.Client, mgetChunkSize, writeChunkSize int) service.SchedulerCache { + return newSchedulerCacheWithOptions(rdb, mgetChunkSize, writeChunkSize, nil) +} + +func newSchedulerCacheWithOptions(rdb *redis.Client, mgetChunkSize, writeChunkSize int, indexedBuckets []string) service.SchedulerCache { if mgetChunkSize <= 0 { mgetChunkSize = defaultSchedulerSnapshotMGetChunkSize } if writeChunkSize <= 0 { writeChunkSize = defaultSchedulerSnapshotWriteChunkSize } + indexed := make(map[string]struct{}, len(indexedBuckets)) + for _, raw := range indexedBuckets { + if bucket, ok := service.ParseSchedulerBucket(raw); ok { + indexed[bucket.String()] = struct{}{} + } + } return &schedulerCache{ rdb: rdb, mgetChunkSize: mgetChunkSize, writeChunkSize: writeChunkSize, + indexedBuckets: indexed, localSnapshots: make(map[string]schedulerLocalSnapshot), localBuckets: make(map[int64]map[string]struct{}), localTTL: defaultSchedulerSnapshotLocalCacheTTL, @@ -241,11 +300,17 @@ func (c *schedulerCache) SetSnapshot(ctx context.Context, bucket service.Schedul keys := []string{activeKey, readyKey, schedulerBucketSetKey, snapshotKey} args := []any{versionStr, bucket.String(), snapshotKeyPrefix, snapshotGraceTTLSeconds} - _, err = activateSnapshotScript.Run(ctx, c.rdb, keys, args...).Result() + activated, err := activateSnapshotScript.Run(ctx, c.rdb, keys, args...).Int() if err != nil { return err } + if activated == 0 { + return nil + } + if err := c.setCandidateIndex(ctx, bucket, versionStr, accounts); err != nil { + return err + } c.invalidateLocalSnapshot(bucket.String()) return nil } @@ -258,13 +323,102 @@ func (c *schedulerCache) clearEmptySnapshot(ctx context.Context, bucket service. keys := []string{activeKey, readyKey, schedulerBucketSetKey} args := []any{versionStr, bucket.String(), snapshotKeyPrefix, snapshotGraceTTLSeconds} - _, err := clearEmptySnapshotScript.Run(ctx, c.rdb, keys, args...).Result() - if err == nil { + activated, err := clearEmptySnapshotScript.Run(ctx, c.rdb, keys, args...).Int() + if err == nil && activated == 1 { + _ = c.clearCandidateIndex(ctx, bucket, versionStr) c.invalidateLocalSnapshot(bucket.String()) } return err } +func (c *schedulerCache) GetCandidateSnapshot(ctx context.Context, bucket service.SchedulerBucket, limit int) ([]*service.Account, bool, error) { + if c == nil || !c.isCandidateIndexEnabled(bucket) { + return nil, false, nil + } + if limit <= 0 { + limit = defaultSchedulerCandidateLimit + } + + readyVal, err := c.rdb.Get(ctx, schedulerBucketKey(schedulerCandidateReadyPrefix, bucket)).Result() + if err == redis.Nil { + return nil, false, nil + } + if err != nil { + return nil, false, err + } + if readyVal != "1" { + return nil, false, nil + } + + version, err := c.rdb.Get(ctx, schedulerBucketKey(schedulerCandidateActivePrefix, bucket)).Result() + if err == redis.Nil { + return nil, false, nil + } + if err != nil { + return nil, false, err + } + activeVersion, err := c.rdb.Get(ctx, schedulerBucketKey(schedulerActivePrefix, bucket)).Result() + if err == redis.Nil { + return nil, false, nil + } + if err != nil { + return nil, false, err + } + if activeVersion != version { + return nil, false, nil + } + + shards := 1 + size := 0 + if raw, err := c.rdb.HGet(ctx, schedulerCandidateMetaKey(bucket, version), "shards").Result(); err == nil { + if parsed, parseErr := strconv.Atoi(raw); parseErr == nil && parsed > 0 { + shards = parsed + } + } else if err != redis.Nil { + return nil, false, err + } + if raw, err := c.rdb.HGet(ctx, schedulerCandidateMetaKey(bucket, version), "size").Result(); err == nil { + if parsed, parseErr := strconv.Atoi(raw); parseErr == nil && parsed > 0 { + size = parsed + } + } else if err != redis.Nil { + return nil, false, err + } + + ids, err := c.readCandidateIDs(ctx, bucket, version, shards, size, limit) + if err != nil { + return nil, false, err + } + if len(ids) == 0 { + return nil, false, nil + } + + keys := make([]string, 0, len(ids)) + for _, id := range ids { + keys = append(keys, schedulerAccountMetaKey(id)) + } + values, err := c.mgetChunked(ctx, keys) + if err != nil { + return nil, false, err + } + + accounts := make([]*service.Account, 0, len(values)) + for _, val := range values { + if val == nil { + continue + } + account, err := decodeCachedAccount(val) + if err != nil { + return nil, false, err + } + accounts = append(accounts, account) + } + if len(accounts) == 0 { + return nil, false, nil + } + return accounts, true, nil +} + func (c *schedulerCache) GetAccount(ctx context.Context, accountID int64) (*service.Account, error) { key := schedulerAccountKey(strconv.FormatInt(accountID, 10)) val, err := c.rdb.Get(ctx, key).Result() @@ -593,6 +747,229 @@ func (c *schedulerCache) SetOutboxWatermark(ctx context.Context, id int64) error return c.rdb.Set(ctx, schedulerOutboxWatermarkKey, strconv.FormatInt(id, 10), 0).Err() } +func (c *schedulerCache) isCandidateIndexEnabled(bucket service.SchedulerBucket) bool { + if c == nil || len(c.indexedBuckets) == 0 { + return false + } + _, ok := c.indexedBuckets[bucket.String()] + return ok +} + +func (c *schedulerCache) setCandidateIndex(ctx context.Context, bucket service.SchedulerBucket, version string, accounts []service.Account) error { + if c == nil || !c.isCandidateIndexEnabled(bucket) { + return nil + } + if len(accounts) == 0 { + return c.clearCandidateIndex(ctx, bucket, version) + } + if len(accounts) <= minSchedulerCandidateShardSize { + return c.clearCandidateIndex(ctx, bucket, version) + } + + shards := schedulerCandidateShardCount(len(accounts)) + membersByKey := make(map[string][]redis.Z, shards) + for idx, account := range accounts { + if account.ID <= 0 { + continue + } + key := schedulerCandidateIndexKey(bucket, version, shards, idx) + membersByKey[key] = append(membersByKey[key], redis.Z{ + Score: float64(idx), + Member: strconv.FormatInt(account.ID, 10), + }) + } + + pipe := c.rdb.Pipeline() + for key, members := range membersByKey { + for start := 0; start < len(members); start += c.writeChunkSize { + end := start + c.writeChunkSize + if end > len(members) { + end = len(members) + } + pipe.ZAdd(ctx, key, members[start:end]...) + } + } + metaKey := schedulerCandidateMetaKey(bucket, version) + pipe.HSet(ctx, metaKey, map[string]any{ + "size": len(accounts), + "shards": shards, + }) + if _, err := pipe.Exec(ctx); err != nil { + return err + } + + activeKey := schedulerBucketKey(schedulerCandidateActivePrefix, bucket) + readyKey := schedulerBucketKey(schedulerCandidateReadyPrefix, bucket) + result, err := activateCandidateIndexScript.Run(ctx, c.rdb, []string{activeKey, readyKey}, version).Slice() + if err != nil { + return err + } + activated := false + if len(result) > 0 { + switch val := result[0].(type) { + case int64: + activated = val == 1 + case int: + activated = val == 1 + case string: + activated = val == "1" + } + } + if !activated { + _ = c.expireCandidateIndex(ctx, bucket, version) + return nil + } + if len(result) > 1 { + if oldVersion, _ := result[1].(string); oldVersion != "" && oldVersion != version { + _ = c.expireCandidateIndex(ctx, bucket, oldVersion) + } + } + return nil +} + +func (c *schedulerCache) clearCandidateIndex(ctx context.Context, bucket service.SchedulerBucket, version string) error { + if c == nil || !c.isCandidateIndexEnabled(bucket) { + return nil + } + activeKey := schedulerBucketKey(schedulerCandidateActivePrefix, bucket) + readyKey := schedulerBucketKey(schedulerCandidateReadyPrefix, bucket) + result, err := clearCandidateIndexScript.Run(ctx, c.rdb, []string{activeKey, readyKey}, version).Slice() + if err != nil { + return err + } + cleared := false + if len(result) > 0 { + switch val := result[0].(type) { + case int64: + cleared = val == 1 + case int: + cleared = val == 1 + case string: + cleared = val == "1" + } + } + if len(result) > 1 { + if oldVersion, _ := result[1].(string); oldVersion != "" { + _ = c.expireCandidateIndex(ctx, bucket, oldVersion) + } + } + if !cleared && version != "" { + _ = c.expireCandidateIndex(ctx, bucket, version) + } + return nil +} + +func (c *schedulerCache) expireCandidateIndex(ctx context.Context, bucket service.SchedulerBucket, version string) error { + shards := 1 + if raw, err := c.rdb.HGet(ctx, schedulerCandidateMetaKey(bucket, version), "shards").Result(); err == nil { + if parsed, parseErr := strconv.Atoi(raw); parseErr == nil && parsed > 0 { + shards = parsed + } + } + pipe := c.rdb.Pipeline() + pipe.Expire(ctx, schedulerCandidateMetaKey(bucket, version), snapshotGraceTTLSeconds*time.Second) + if shards <= 1 { + pipe.Expire(ctx, schedulerCandidateIndexBaseKey(bucket, version), snapshotGraceTTLSeconds*time.Second) + } else { + for shard := 0; shard < shards; shard++ { + pipe.Expire(ctx, schedulerCandidateShardKey(bucket, version, shard), snapshotGraceTTLSeconds*time.Second) + } + } + _, err := pipe.Exec(ctx) + return err +} + +func (c *schedulerCache) readCandidateIDs(ctx context.Context, bucket service.SchedulerBucket, version string, shards int, size int, limit int) ([]string, error) { + if shards <= 1 { + return c.readCandidateIDsFromZSet(ctx, schedulerCandidateIndexBaseKey(bucket, version), size, limit) + } + seen := make(map[string]struct{}, limit) + ids := make([]string, 0, limit) + perShard := limit / 4 + if perShard < 16 { + perShard = 16 + } + if perShard > limit { + perShard = limit + } + startShard := int(time.Now().UnixNano() % int64(shards)) + for offset := 0; offset < shards && len(ids) < limit; offset++ { + shard := (startShard + offset) % shards + part, err := c.readCandidateIDsFromZSet(ctx, schedulerCandidateShardKey(bucket, version, shard), schedulerCandidateShardSize(size, shards, shard), perShard) + if err != nil { + return nil, err + } + for _, id := range part { + if _, ok := seen[id]; ok { + continue + } + seen[id] = struct{}{} + ids = append(ids, id) + if len(ids) >= limit { + break + } + } + } + return ids, nil +} + +func (c *schedulerCache) readCandidateIDsFromZSet(ctx context.Context, key string, size int, limit int) ([]string, error) { + if limit <= 0 { + return nil, nil + } + if size <= limit { + return c.rdb.ZRange(ctx, key, 0, int64(limit-1)).Result() + } + start := rand.Intn(size - limit + 1) + stop := start + limit - 1 + return c.rdb.ZRange(ctx, key, int64(start), int64(stop)).Result() +} + +func schedulerCandidateShardCount(size int) int { + if size <= minSchedulerCandidateShardSize { + return 1 + } + shards := 1 + for shards < maxSchedulerCandidateShards && size/shards > defaultSchedulerCandidateShardTarget { + shards *= 2 + } + return shards +} + +func schedulerCandidateShardSize(size int, shards int, shard int) int { + if size <= 0 || shards <= 0 || shard < 0 || shard >= shards { + return 0 + } + base := size / shards + if shard < size%shards { + return base + 1 + } + return base +} + +func schedulerCandidateIndexKey(bucket service.SchedulerBucket, version string, shards int, index int) string { + if shards <= 1 { + return schedulerCandidateIndexBaseKey(bucket, version) + } + shard := index % shards + if shard < 0 { + shard = -shard + } + return schedulerCandidateShardKey(bucket, version, shard) +} + +func schedulerCandidateIndexBaseKey(bucket service.SchedulerBucket, version string) string { + return fmt.Sprintf("%s%d:%s:%s:v%s", schedulerCandidateIndexPrefix, bucket.GroupID, bucket.Platform, bucket.Mode, version) +} + +func schedulerCandidateShardKey(bucket service.SchedulerBucket, version string, shard int) string { + return fmt.Sprintf("%s:s%d", schedulerCandidateIndexBaseKey(bucket, version), shard) +} + +func schedulerCandidateMetaKey(bucket service.SchedulerBucket, version string) string { + return fmt.Sprintf("%s%d:%s:%s:v%s", schedulerCandidateMetaPrefix, bucket.GroupID, bucket.Platform, bucket.Mode, version) +} + func schedulerBucketKey(prefix string, bucket service.SchedulerBucket) string { return fmt.Sprintf("%s%d:%s:%s", prefix, bucket.GroupID, bucket.Platform, bucket.Mode) } diff --git a/backend/internal/repository/scheduler_cache_integration_test.go b/backend/internal/repository/scheduler_cache_integration_test.go index dfa4b0f0d..db59cc17b 100644 --- a/backend/internal/repository/scheduler_cache_integration_test.go +++ b/backend/internal/repository/scheduler_cache_integration_test.go @@ -149,6 +149,99 @@ func TestSchedulerCacheEmptySnapshotEvictsBucket(t *testing.T) { require.False(t, isMember) } +func TestSchedulerCacheCandidateIndexManualBucket(t *testing.T) { + ctx := context.Background() + rdb := testRedis(t) + bucket := service.SchedulerBucket{GroupID: 18, Platform: service.PlatformOpenAI, Mode: service.SchedulerModeSingle} + cache := newSchedulerCacheWithOptions(rdb, 128, 256, []string{bucket.String()}) + + accounts := make([]service.Account, 0, minSchedulerCandidateShardSize+1) + for i := 0; i < minSchedulerCandidateShardSize+1; i++ { + accounts = append(accounts, service.Account{ + ID: int64(100000 + i), + Name: "candidate-index", + Platform: service.PlatformOpenAI, + Type: service.AccountTypeOAuth, + Status: service.StatusActive, + Schedulable: true, + Concurrency: 3, + Priority: i % 10, + GroupIDs: []int64{bucket.GroupID}, + }) + } + require.NoError(t, cache.SetSnapshot(ctx, bucket, accounts)) + + candidateCache, ok := cache.(service.SchedulerCandidateCache) + require.True(t, ok) + candidates, hit, err := candidateCache.GetCandidateSnapshot(ctx, bucket, 64) + require.NoError(t, err) + require.True(t, hit) + require.Len(t, candidates, 64) + + version, err := rdb.Get(ctx, schedulerBucketKey(schedulerCandidateActivePrefix, bucket)).Result() + require.NoError(t, err) + shards, err := rdb.HGet(ctx, schedulerCandidateMetaKey(bucket, version), "shards").Int() + require.NoError(t, err) + require.Greater(t, shards, 1) + + require.NoError(t, cache.SetSnapshot(ctx, bucket, accounts[:1])) + candidates, hit, err = candidateCache.GetCandidateSnapshot(ctx, bucket, 64) + require.NoError(t, err) + require.False(t, hit) + require.Nil(t, candidates) +} + +func TestSchedulerCacheCandidateIndexManualSmallBucketMisses(t *testing.T) { + ctx := context.Background() + rdb := testRedis(t) + bucket := service.SchedulerBucket{GroupID: 18, Platform: service.PlatformOpenAI, Mode: service.SchedulerModeSingle} + cache := newSchedulerCacheWithOptions(rdb, 128, 256, []string{bucket.String()}) + + require.NoError(t, cache.SetSnapshot(ctx, bucket, []service.Account{{ + ID: 42, + Name: "small-index", + Platform: service.PlatformOpenAI, + Type: service.AccountTypeOAuth, + Status: service.StatusActive, + Schedulable: true, + Concurrency: 3, + GroupIDs: []int64{bucket.GroupID}, + }})) + + candidateCache, ok := cache.(service.SchedulerCandidateCache) + require.True(t, ok) + candidates, hit, err := candidateCache.GetCandidateSnapshot(ctx, bucket, 64) + require.NoError(t, err) + require.False(t, hit) + require.Nil(t, candidates) +} + +func TestSchedulerCacheCandidateIndexDisabledBucketMisses(t *testing.T) { + ctx := context.Background() + rdb := testRedis(t) + enabledBucket := service.SchedulerBucket{GroupID: 18, Platform: service.PlatformOpenAI, Mode: service.SchedulerModeSingle} + disabledBucket := service.SchedulerBucket{GroupID: 19, Platform: service.PlatformOpenAI, Mode: service.SchedulerModeSingle} + cache := newSchedulerCacheWithOptions(rdb, 128, 256, []string{enabledBucket.String()}) + + require.NoError(t, cache.SetSnapshot(ctx, disabledBucket, []service.Account{{ + ID: 42, + Name: "disabled-index", + Platform: service.PlatformOpenAI, + Type: service.AccountTypeOAuth, + Status: service.StatusActive, + Schedulable: true, + Concurrency: 3, + GroupIDs: []int64{disabledBucket.GroupID}, + }})) + + candidateCache, ok := cache.(service.SchedulerCandidateCache) + require.True(t, ok) + candidates, hit, err := candidateCache.GetCandidateSnapshot(ctx, disabledBucket, 64) + require.NoError(t, err) + require.False(t, hit) + require.Nil(t, candidates) +} + func schedulerBucketStrings(buckets []service.SchedulerBucket) []string { out := make([]string, 0, len(buckets)) for _, bucket := range buckets { diff --git a/backend/internal/repository/wire.go b/backend/internal/repository/wire.go index ed69e32d8..447d400e9 100644 --- a/backend/internal/repository/wire.go +++ b/backend/internal/repository/wire.go @@ -51,6 +51,7 @@ func ProvideSessionLimitCache(rdb *redis.Client, cfg *config.Config) service.Ses func ProvideSchedulerCache(rdb *redis.Client, cfg *config.Config) service.SchedulerCache { mgetChunkSize := defaultSchedulerSnapshotMGetChunkSize writeChunkSize := defaultSchedulerSnapshotWriteChunkSize + var indexedBuckets []string if cfg != nil { if cfg.Gateway.Scheduling.SnapshotMGetChunkSize > 0 { mgetChunkSize = cfg.Gateway.Scheduling.SnapshotMGetChunkSize @@ -58,8 +59,9 @@ func ProvideSchedulerCache(rdb *redis.Client, cfg *config.Config) service.Schedu if cfg.Gateway.Scheduling.SnapshotWriteChunkSize > 0 { writeChunkSize = cfg.Gateway.Scheduling.SnapshotWriteChunkSize } + indexedBuckets = cfg.Gateway.Scheduling.IndexedBuckets } - return newSchedulerCacheWithChunkSizes(rdb, mgetChunkSize, writeChunkSize) + return newSchedulerCacheWithOptions(rdb, mgetChunkSize, writeChunkSize, indexedBuckets) } // ProviderSet is the Wire provider set for all repositories diff --git a/backend/internal/service/openai_account_scheduler.go b/backend/internal/service/openai_account_scheduler.go index a83f4f7a3..1156e6108 100644 --- a/backend/internal/service/openai_account_scheduler.go +++ b/backend/internal/service/openai_account_scheduler.go @@ -850,35 +850,16 @@ func (s *defaultOpenAIAccountScheduler) selectByLoadBalance( schedGroup, _ = s.service.schedulerSnapshot.GetGroupByID(ctx, *req.GroupID) } - filtered := make([]*Account, 0, len(accounts)) - loadReq := make([]AccountWithConcurrency, 0, len(accounts)) - for i := range accounts { - account := &accounts[i] - if req.ExcludedIDs != nil { - if _, excluded := req.ExcludedIDs[account.ID]; excluded { - continue + filtered, loadReq := s.filterOpenAIAccountsForLoadBalance(ctx, accounts, req, schedGroup) + if len(filtered) == 0 { + if s.service.shouldRetryOpenAISchedulerWithoutCandidateIndex(ctx, req.GroupID) { + retryCtx := WithSchedulerCandidateIndexBypass(ctx) + accounts, err = s.service.listSchedulableAccounts(retryCtx, req.GroupID) + if err != nil { + return nil, 0, 0, 0, err } + filtered, loadReq = s.filterOpenAIAccountsForLoadBalance(retryCtx, accounts, req, schedGroup) } - if !account.IsSchedulable() || !account.IsOpenAI() { - continue - } - // require_privacy_set: 跳过 privacy 未设置的账号并标记异常 - if schedGroup != nil && schedGroup.RequirePrivacySet && !account.IsPrivacySet() { - _ = s.service.accountRepo.SetError(ctx, account.ID, - fmt.Sprintf("Privacy not set, required by group [%s]", schedGroup.Name)) - continue - } - if !s.isAccountRequestCompatible(account, req) { - continue - } - if !s.isAccountTransportCompatible(account, req.RequiredTransport) { - continue - } - filtered = append(filtered, account) - loadReq = append(loadReq, AccountWithConcurrency{ - ID: account.ID, - MaxConcurrency: account.EffectiveLoadFactor(), - }) } if len(filtered) == 0 { return nil, 0, 0, 0, noAvailableOpenAISelectionError(req.RequestedModel, false) @@ -1126,6 +1107,45 @@ func (s *defaultOpenAIAccountScheduler) selectByLoadBalance( return nil, candidateCount, topK, loadSkew, noAvailableOpenAISelectionError(req.RequestedModel, compactBlocked) } +func (s *defaultOpenAIAccountScheduler) filterOpenAIAccountsForLoadBalance( + ctx context.Context, + accounts []Account, + req OpenAIAccountScheduleRequest, + schedGroup *Group, +) ([]*Account, []AccountWithConcurrency) { + filtered := make([]*Account, 0, len(accounts)) + loadReq := make([]AccountWithConcurrency, 0, len(accounts)) + for i := range accounts { + account := &accounts[i] + if req.ExcludedIDs != nil { + if _, excluded := req.ExcludedIDs[account.ID]; excluded { + continue + } + } + if !account.IsSchedulable() || !account.IsOpenAI() { + continue + } + // require_privacy_set: 跳过 privacy 未设置的账号并标记异常 + if schedGroup != nil && schedGroup.RequirePrivacySet && !account.IsPrivacySet() { + _ = s.service.accountRepo.SetError(ctx, account.ID, + fmt.Sprintf("Privacy not set, required by group [%s]", schedGroup.Name)) + continue + } + if !s.isAccountRequestCompatible(account, req) { + continue + } + if !s.isAccountTransportCompatible(account, req.RequiredTransport) { + continue + } + filtered = append(filtered, account) + loadReq = append(loadReq, AccountWithConcurrency{ + ID: account.ID, + MaxConcurrency: account.EffectiveLoadFactor(), + }) + } + return filtered, loadReq +} + func (s *defaultOpenAIAccountScheduler) isAccountTransportCompatible(account *Account, requiredTransport OpenAIUpstreamTransport) bool { if requiredTransport == OpenAIUpstreamTransportAny || requiredTransport == OpenAIUpstreamTransportHTTPSSE { return true @@ -1433,6 +1453,26 @@ func (s *OpenAIGatewayService) openAIWSLBTopK() int { return 7 } +func (s *OpenAIGatewayService) shouldRetryOpenAISchedulerWithoutCandidateIndex(ctx context.Context, groupID *int64) bool { + if s == nil || s.schedulerSnapshot == nil || IsSchedulerCandidateIndexBypassed(ctx) { + return false + } + cfg := s.schedulingConfig() + if len(cfg.IndexedBuckets) == 0 { + return false + } + bucket := SchedulerBucket{GroupID: 0, Platform: PlatformOpenAI, Mode: SchedulerModeSingle} + if groupID != nil && *groupID > 0 { + bucket.GroupID = *groupID + } + for _, raw := range cfg.IndexedBuckets { + if raw == bucket.String() { + return true + } + } + return false +} + func (s *OpenAIGatewayService) openAIWSSchedulerWeights() GatewayOpenAIWSSchedulerScoreWeightsView { if s != nil && s.cfg != nil { return GatewayOpenAIWSSchedulerScoreWeightsView{ diff --git a/backend/internal/service/openai_account_scheduler_test.go b/backend/internal/service/openai_account_scheduler_test.go index e477a7af1..f9c7a539a 100644 --- a/backend/internal/service/openai_account_scheduler_test.go +++ b/backend/internal/service/openai_account_scheduler_test.go @@ -19,6 +19,14 @@ type openAISnapshotCacheStub struct { accountsByID map[int64]*Account } +type openAICandidateSnapshotCacheStub struct { + openAISnapshotCacheStub + candidateAccounts []*Account + candidateHits int + fullHits int + bypassHits int +} + type schedulerTestOpenAIAccountRepo struct { AccountRepository accounts []Account @@ -274,6 +282,34 @@ func (s *openAISnapshotCacheStub) GetAccount(ctx context.Context, accountID int6 return &cloned, nil } +func (s *openAICandidateSnapshotCacheStub) GetCandidateSnapshot(ctx context.Context, bucket SchedulerBucket, limit int) ([]*Account, bool, error) { + s.candidateHits++ + if len(s.candidateAccounts) == 0 { + return nil, false, nil + } + out := make([]*Account, 0, len(s.candidateAccounts)) + for _, account := range s.candidateAccounts { + if account == nil { + continue + } + cloned := *account + out = append(out, &cloned) + if len(out) >= limit { + break + } + } + return out, true, nil +} + +func (s *openAICandidateSnapshotCacheStub) GetSnapshot(ctx context.Context, bucket SchedulerBucket) ([]*Account, bool, error) { + if IsSchedulerCandidateIndexBypassed(ctx) { + s.bypassHits++ + } else { + s.fullHits++ + } + return s.openAISnapshotCacheStub.GetSnapshot(ctx, bucket) +} + func TestOpenAIGatewayService_SelectAccountWithScheduler_DefaultDisabledUsesLegacyLoadAwareness(t *testing.T) { resetOpenAIAdvancedSchedulerSettingCacheForTest() @@ -1176,6 +1212,88 @@ func TestOpenAIGatewayService_SelectAccountWithScheduler_LoadBalanceTopKFallback } } +func TestOpenAIGatewayService_SelectAccountWithScheduler_CandidateIndexFallbacksToFullSnapshot(t *testing.T) { + ctx := context.Background() + groupID := int64(18) + indexedBucket := SchedulerBucket{GroupID: groupID, Platform: PlatformOpenAI, Mode: SchedulerModeSingle} + + candidateOnly := Account{ + ID: 3101, + Platform: PlatformOpenAI, + Type: AccountTypeAPIKey, + Status: StatusActive, + Schedulable: true, + Concurrency: 1, + Priority: 0, + Credentials: map[string]any{ + "model_mapping": map[string]any{"gpt-4o": "gpt-4o"}, + }, + } + fullOnly := Account{ + ID: 3102, + Platform: PlatformOpenAI, + Type: AccountTypeAPIKey, + Status: StatusActive, + Schedulable: true, + Concurrency: 1, + Priority: 0, + Credentials: map[string]any{ + "model_mapping": map[string]any{"gpt-5.1": "gpt-5.1"}, + }, + } + candidateOnly = openAITestAccountWithGroupIfUnset(candidateOnly, groupID) + fullOnly = openAITestAccountWithGroupIfUnset(fullOnly, groupID) + + cache := &openAICandidateSnapshotCacheStub{ + openAISnapshotCacheStub: openAISnapshotCacheStub{ + snapshotAccounts: []*Account{&candidateOnly, &fullOnly}, + accountsByID: map[int64]*Account{ + candidateOnly.ID: &candidateOnly, + fullOnly.ID: &fullOnly, + }, + }, + candidateAccounts: []*Account{&candidateOnly}, + } + cfg := &config.Config{} + cfg.Gateway.Scheduling.IndexedBuckets = []string{indexedBucket.String()} + cfg.Gateway.Scheduling.IndexedCandidateLimit = 256 + cfg.Gateway.OpenAIWS.SchedulerScoreWeights.Priority = 1 + cfg.Gateway.OpenAIWS.SchedulerScoreWeights.Load = 1 + cfg.Gateway.OpenAIWS.SchedulerScoreWeights.Queue = 1 + cfg.Gateway.OpenAIWS.SchedulerScoreWeights.ErrorRate = 1 + cfg.Gateway.OpenAIWS.SchedulerScoreWeights.TTFT = 1 + + snapshot := NewSchedulerSnapshotService(cache, nil, schedulerTestOpenAIAccountRepo{accounts: []Account{candidateOnly, fullOnly}}, nil, cfg) + svc := &OpenAIGatewayService{ + accountRepo: schedulerTestOpenAIAccountRepo{accounts: []Account{candidateOnly, fullOnly}}, + cache: &schedulerTestGatewayCache{}, + cfg: cfg, + rateLimitService: newOpenAIAdvancedSchedulerRateLimitService("true"), + concurrencyService: NewConcurrencyService(schedulerTestConcurrencyCache{}), + schedulerSnapshot: snapshot, + } + + selection, decision, err := svc.SelectAccountWithScheduler( + ctx, + &groupID, + "", + "", + "gpt-5.1", + nil, + OpenAIUpstreamTransportAny, + false, + ) + require.NoError(t, err) + require.NotNil(t, selection) + require.NotNil(t, selection.Account) + require.Equal(t, fullOnly.ID, selection.Account.ID) + require.Equal(t, openAIAccountScheduleLayerLoadBalance, decision.Layer) + require.Equal(t, 1, decision.CandidateCount) + require.Equal(t, 1, cache.candidateHits) + require.Equal(t, 1, cache.bypassHits) + require.Zero(t, cache.fullHits) +} + func TestOpenAIGatewayService_OpenAIAccountSchedulerMetrics(t *testing.T) { ctx := context.Background() groupID := int64(12) diff --git a/backend/internal/service/scheduler_cache.go b/backend/internal/service/scheduler_cache.go index f9794c821..8dd3a25b0 100644 --- a/backend/internal/service/scheduler_cache.go +++ b/backend/internal/service/scheduler_cache.go @@ -20,6 +20,8 @@ type SchedulerBucket struct { Mode string } +type schedulerCandidateIndexBypassKey struct{} + func (b SchedulerBucket) String() string { return fmt.Sprintf("%d:%s:%s", b.GroupID, b.Platform, b.Mode) } @@ -43,6 +45,15 @@ func ParseSchedulerBucket(raw string) (SchedulerBucket, bool) { }, true } +func WithSchedulerCandidateIndexBypass(ctx context.Context) context.Context { + return context.WithValue(ctx, schedulerCandidateIndexBypassKey{}, true) +} + +func IsSchedulerCandidateIndexBypassed(ctx context.Context) bool { + bypass, _ := ctx.Value(schedulerCandidateIndexBypassKey{}).(bool) + return bypass +} + // SchedulerCache 负责调度快照与账号快照的缓存读写。 type SchedulerCache interface { // GetSnapshot 读取快照并返回命中与否（ready + active + 数据完整）。 @@ -68,3 +79,11 @@ type SchedulerCache interface { // SetOutboxWatermark 保存 outbox 水位。 SetOutboxWatermark(ctx context.Context, id int64) error } + +// SchedulerCandidateCache is an optional extension for caches that can return a +// small indexed candidate set instead of materializing a whole scheduler bucket. +type SchedulerCandidateCache interface { + // GetCandidateSnapshot reads a manually enabled candidate index for bucket. + // hit=false means callers should fall back to the full scheduler snapshot. + GetCandidateSnapshot(ctx context.Context, bucket SchedulerBucket, limit int) ([]*Account, bool, error) +} diff --git a/backend/internal/service/scheduler_snapshot_service.go b/backend/internal/service/scheduler_snapshot_service.go index e699b1565..542259f34 100644 --- a/backend/internal/service/scheduler_snapshot_service.go +++ b/backend/internal/service/scheduler_snapshot_service.go @@ -110,6 +110,17 @@ func (s *SchedulerSnapshotService) ListSchedulableAccounts(ctx context.Context, bucket := s.bucketFor(groupID, platform, mode) if s.cache != nil { + if candidateCache, ok := s.cache.(SchedulerCandidateCache); ok { + if !IsSchedulerCandidateIndexBypassed(ctx) { + candidateLimit := s.candidateIndexLimit() + cached, hit, err := candidateCache.GetCandidateSnapshot(ctx, bucket, candidateLimit) + if err != nil { + logger.LegacyPrintf("service.scheduler_snapshot", "[Scheduler] candidate cache read failed: bucket=%s err=%v", bucket.String(), err) + } else if hit { + return filterSchedulableAccounts(derefAccounts(cached)), useMixed, nil + } + } + } cached, hit, err := s.cache.GetSnapshot(ctx, bucket) if err != nil { logger.LegacyPrintf("service.scheduler_snapshot", "[Scheduler] cache read failed: bucket=%s err=%v", bucket.String(), err) @@ -810,6 +821,13 @@ func (s *SchedulerSnapshotService) fullRebuildInterval() time.Duration { return time.Duration(sec) * time.Second } +func (s *SchedulerSnapshotService) candidateIndexLimit() int { + if s == nil || s.cfg == nil || s.cfg.Gateway.Scheduling.IndexedCandidateLimit <= 0 { + return 256 + } + return s.cfg.Gateway.Scheduling.IndexedCandidateLimit +} + func (s *SchedulerSnapshotService) defaultBuckets(ctx context.Context) ([]SchedulerBucket, error) { buckets := make([]SchedulerBucket, 0) platforms := []string{PlatformAnthropic, PlatformGemini, PlatformOpenAI, PlatformAntigravity} diff --git a/deploy/config.example.yaml b/deploy/config.example.yaml index 12256f9d7..58b1f21ad 100644 --- a/deploy/config.example.yaml +++ b/deploy/config.example.yaml @@ -373,6 +373,13 @@ gateway: # Enable batch load calculation for scheduling # 启用调度批量负载计算 load_batch_enabled: true + # Explicit scheduler buckets that use Redis candidate index when bucket size > 5000. + # Format: "groupID:platform:mode", e.g. "18:openai:single". + # 显式启用 Redis 候选索引的调度桶；仅当桶内账号数 > 5000 时生效。 + indexed_buckets: [] + # Max accounts returned by candidate index per scheduling lookup + # 候选索引每次调度最多返回的账号数 + indexed_candidate_limit: 256 # Slot cleanup interval (duration) # 并发槽位清理周期（时间段） slot_cleanup_interval: 30s