diff --git a/backend/cmd/server/wire_gen.go b/backend/cmd/server/wire_gen.go index b2774d7232..675552de3e 100644 --- a/backend/cmd/server/wire_gen.go +++ b/backend/cmd/server/wire_gen.go @@ -137,7 +137,8 @@ func initializeApplication(buildInfo handler.BuildInfo) (*Application, error) { tlsFingerprintProfileCache := repository.NewTLSFingerprintProfileCache(redisClient) tlsFingerprintProfileService := service.NewTLSFingerprintProfileService(tlsFingerprintProfileRepository, tlsFingerprintProfileCache) channelRepository := repository.NewChannelRepository(db) - channelService := service.NewChannelService(channelRepository, groupRepository, apiKeyAuthCacheInvalidator, pricingService) + channelCachePubSub := repository.NewChannelCache(redisClient) + channelService := service.NewChannelService(channelRepository, groupRepository, apiKeyAuthCacheInvalidator, pricingService, channelCachePubSub) modelPricingResolver := service.NewModelPricingResolver(channelService, billingService) compositeModelRouteRepository := repository.NewCompositeModelRouteRepository(client) compositeRouteResolver := service.NewCompositeRouteResolver(compositeModelRouteRepository) diff --git a/backend/go.sum b/backend/go.sum index cbb257f12f..3c46484df8 100644 --- a/backend/go.sum +++ b/backend/go.sum @@ -255,6 +255,8 @@ github.com/google/go-tpm-tools v0.3.13-0.20230620182252-4639ecce2aba/go.mod h1:E github.com/google/gofuzz v1.0.0/go.mod h1:dBl0BpW6vV/+mYPU4Po3pmUjxk6FQPldtuIdl/M65Eg= github.com/google/pprof v0.0.0-20250317173921-a4b03ec1a45e h1:ijClszYn+mADRFY17kjQEVQ1XRhq2/JR1M3sGqeJoxs= github.com/google/pprof v0.0.0-20250317173921-a4b03ec1a45e/go.mod h1:boTsfXsheKC2y+lKOCMpSfarhxDeIzfZG1jqGcPl3cA= +github.com/google/subcommands v1.2.0 h1:vWQspBTo2nEqTUFita5/KeEWlUL8kQObDFbub/EN9oE= +github.com/google/subcommands v1.2.0/go.mod h1:ZjhPrFU+Olkh9WazFPsl27BQ4UPiG37m3yTrtFlrHVk= github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0= github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo= github.com/google/wire v0.7.0 h1:JxUKI6+CVBgCO2WToKy/nQk0sS+amI9z9EjVmdaocj4= diff --git a/backend/internal/handler/openai_gateway_handler_test.go b/backend/internal/handler/openai_gateway_handler_test.go index 86c56d4901..00303e0ff2 100644 --- a/backend/internal/handler/openai_gateway_handler_test.go +++ b/backend/internal/handler/openai_gateway_handler_test.go @@ -2868,7 +2868,7 @@ func runOpenAIResponsesWebSocketUsageLogCase(t *testing.T, tc openAIResponsesWSU BillingModelSource: tc.billingModelSource, }}, groupPlatforms: map[int64]string{groupID: service.PlatformOpenAI}, - }, nil, nil, nil) + }, nil, nil, nil, nil) } billingCacheSvc := service.NewBillingCacheService(nil, nil, nil, nil, nil, nil, cfg, nil) diff --git a/backend/internal/repository/channel_cache.go b/backend/internal/repository/channel_cache.go new file mode 100644 index 0000000000..ff8e75b61e --- /dev/null +++ b/backend/internal/repository/channel_cache.go @@ -0,0 +1,53 @@ +package repository + +import ( + "context" + "log/slog" + + "github.com/Wei-Shaw/sub2api/internal/service" + "github.com/redis/go-redis/v9" +) + +const channelCachePubSubKey = "channel_cache_updated" + +type channelCache struct { + rdb *redis.Client +} + +// NewChannelCache creates the Redis-backed channel cache invalidation bus. +func NewChannelCache(rdb *redis.Client) service.ChannelCachePubSub { + return &channelCache{rdb: rdb} +} + +// NotifyUpdate notifies all instances to invalidate their local channel cache. +func (c *channelCache) NotifyUpdate(ctx context.Context) error { + return c.rdb.Publish(ctx, channelCachePubSubKey, "refresh").Err() +} + +// SubscribeUpdates subscribes to channel cache invalidation notifications. +func (c *channelCache) SubscribeUpdates(ctx context.Context, handler func()) { + go func() { + pubsub := c.rdb.Subscribe(ctx, channelCachePubSubKey) + defer func() { + if err := pubsub.Close(); err != nil { + slog.Warn("failed to close channel cache subscriber", "error", err) + } + }() + + messages := pubsub.Channel() + for { + select { + case <-ctx.Done(): + return + case message, ok := <-messages: + if !ok { + slog.Warn("channel cache subscriber stopped", "reason", "channel_closed") + return + } + if message != nil { + handler() + } + } + } + }() +} diff --git a/backend/internal/repository/channel_cache_test.go b/backend/internal/repository/channel_cache_test.go new file mode 100644 index 0000000000..89c45249ef --- /dev/null +++ b/backend/internal/repository/channel_cache_test.go @@ -0,0 +1,44 @@ +package repository + +import ( + "context" + "testing" + "time" + + "github.com/alicebob/miniredis/v2" + "github.com/redis/go-redis/v9" + "github.com/stretchr/testify/require" +) + +func TestChannelCachePubSub(t *testing.T) { + redisServer := miniredis.RunT(t) + publisherRedis := redis.NewClient(&redis.Options{Addr: redisServer.Addr()}) + subscriberRedis := redis.NewClient(&redis.Options{Addr: redisServer.Addr()}) + t.Cleanup(func() { + _ = publisherRedis.Close() + _ = subscriberRedis.Close() + }) + + publisher := NewChannelCache(publisherRedis) + subscriber := NewChannelCache(subscriberRedis) + ctx, cancel := context.WithCancel(context.Background()) + t.Cleanup(cancel) + received := make(chan struct{}, 1) + subscriber.SubscribeUpdates(ctx, func() { + select { + case received <- struct{}{}: + default: + } + }) + + require.Eventually(t, func() bool { + return redisServer.PubSubNumSub(channelCachePubSubKey)[channelCachePubSubKey] == 1 + }, time.Second, 10*time.Millisecond) + require.NoError(t, publisher.NotifyUpdate(context.Background())) + + select { + case <-received: + case <-time.After(time.Second): + t.Fatal("channel cache invalidation notification was not received") + } +} diff --git a/backend/internal/repository/wire.go b/backend/internal/repository/wire.go index e6904f833c..eae41bde41 100644 --- a/backend/internal/repository/wire.go +++ b/backend/internal/repository/wire.go @@ -137,6 +137,7 @@ var ProviderSet = wire.NewSet( NewRefreshTokenCache, NewErrorPassthroughCache, NewTLSFingerprintProfileCache, + NewChannelCache, NewContentModerationHashCache, // Encryptors diff --git a/backend/internal/service/billing_token_cost_request_test.go b/backend/internal/service/billing_token_cost_request_test.go index 8392b8cf49..76809514ab 100644 --- a/backend/internal/service/billing_token_cost_request_test.go +++ b/backend/internal/service/billing_token_cost_request_test.go @@ -23,7 +23,7 @@ func newTokenCostTestEnv(t *testing.T, groupPlatform string, pricing []ChannelMo return map[int64]string{100: groupPlatform}, nil }, } - cs := NewChannelService(repo, nil, nil, nil) + cs := NewChannelService(repo, nil, nil, nil, nil) bs := NewBillingService(&config.Config{}, catalog) return bs, NewModelPricingResolver(cs, bs) } diff --git a/backend/internal/service/channel_available_test.go b/backend/internal/service/channel_available_test.go index 2b7b23f7e5..d382c5fc2d 100644 --- a/backend/internal/service/channel_available_test.go +++ b/backend/internal/service/channel_available_test.go @@ -75,7 +75,7 @@ func newAvailableChannelService(channels []Channel, groupRepo GroupRepository) * repo := &mockChannelRepository{ listAllFn: func(ctx context.Context) ([]Channel, error) { return channels, nil }, } - return NewChannelService(repo, groupRepo, nil, nil) + return NewChannelService(repo, groupRepo, nil, nil, nil) } func TestListAvailable_EmptyActiveGroups_NoGroupsAttached(t *testing.T) { @@ -134,7 +134,7 @@ func TestListAvailable_ListAllErrorPropagates(t *testing.T) { listAllFn: func(ctx context.Context) ([]Channel, error) { return nil, sentinel }, } groupRepo := &stubGroupRepoForAvailable{} - svc := NewChannelService(repo, groupRepo, nil, nil) + svc := NewChannelService(repo, groupRepo, nil, nil, nil) out, err := svc.ListAvailable(context.Background()) require.Nil(t, out) require.ErrorIs(t, err, sentinel) diff --git a/backend/internal/service/channel_service.go b/backend/internal/service/channel_service.go index b3fe7dc63a..f32b297679 100644 --- a/backend/internal/service/channel_service.go +++ b/backend/internal/service/channel_service.go @@ -147,14 +147,22 @@ const ( channelCacheTTL = 10 * time.Minute channelErrorTTL = 5 * time.Second // DB 错误时的短缓存 channelCacheDBTimeout = 10 * time.Second + channelCacheNotifyTTL = 3 * time.Second ) +// ChannelCachePubSub broadcasts channel cache invalidations between instances. +type ChannelCachePubSub interface { + NotifyUpdate(ctx context.Context) error + SubscribeUpdates(ctx context.Context, handler func()) +} + // ChannelService 渠道管理服务 type ChannelService struct { repo ChannelRepository groupRepo GroupRepository authCacheInvalidator APIKeyAuthCacheInvalidator pricingService *PricingService // 用于「可用渠道」展示时回落到全局定价；可为 nil（测试场景） + cachePubSub ChannelCachePubSub cache atomic.Value // *channelCache cacheSF singleflight.Group @@ -163,13 +171,15 @@ type ChannelService struct { // NewChannelService 创建渠道服务实例。 // pricingService 仅供 ListAvailable 在渠道未配置定价时回落到全局 LiteLLM 数据； // 计费热路径走独立的 ModelPricingResolver，与此参数无关。可传 nil。 -func NewChannelService(repo ChannelRepository, groupRepo GroupRepository, authCacheInvalidator APIKeyAuthCacheInvalidator, pricingService *PricingService) *ChannelService { +func NewChannelService(repo ChannelRepository, groupRepo GroupRepository, authCacheInvalidator APIKeyAuthCacheInvalidator, pricingService *PricingService, cachePubSub ChannelCachePubSub) *ChannelService { s := &ChannelService{ repo: repo, groupRepo: groupRepo, authCacheInvalidator: authCacheInvalidator, pricingService: pricingService, + cachePubSub: cachePubSub, } + s.subscribeCacheUpdates(context.Background()) return s } @@ -383,13 +393,41 @@ func (s *ChannelService) InvalidateCache() { } func (s *ChannelService) invalidateCache() { - s.cache.Store((*channelCache)(nil)) - s.cacheSF.Forget("channel_cache") + s.clearCache() // 主动重建缓存，确保 CRUD 后立即生效 if _, err := s.buildCache(context.Background()); err != nil { slog.Warn("failed to rebuild channel cache after invalidation", "error", err) } + + s.notifyCacheUpdate() +} + +// clearCache clears only the in-process snapshot. Keeping this separate from +// invalidateCache prevents notifications received from Redis from being +// published again in a loop. +func (s *ChannelService) clearCache() { + s.cache.Store((*channelCache)(nil)) + s.cacheSF.Forget("channel_cache") +} + +func (s *ChannelService) notifyCacheUpdate() { + if s.cachePubSub == nil { + return + } + + ctx, cancel := context.WithTimeout(context.Background(), channelCacheNotifyTTL) + defer cancel() + if err := s.cachePubSub.NotifyUpdate(ctx); err != nil { + slog.Warn("failed to publish channel cache invalidation", "error", err) + } +} + +func (s *ChannelService) subscribeCacheUpdates(ctx context.Context) { + if s.cachePubSub == nil { + return + } + s.cachePubSub.SubscribeUpdates(ctx, s.clearCache) } // matchWildcard 在通配符定价中查找匹配项（最先匹配到优先） diff --git a/backend/internal/service/channel_service_test.go b/backend/internal/service/channel_service_test.go index e443b80b59..ec5613f009 100644 --- a/backend/internal/service/channel_service_test.go +++ b/backend/internal/service/channel_service_test.go @@ -6,6 +6,7 @@ import ( "context" "errors" "net/http" + "sync" "testing" infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors" @@ -175,6 +176,27 @@ type mockChannelAuthCacheInvalidator struct { invalidatedUserIDs []int64 } +type mockChannelCachePubSub struct { + mu sync.Mutex + handlers []func() +} + +func (m *mockChannelCachePubSub) NotifyUpdate(context.Context) error { + m.mu.Lock() + handlers := append([]func(){}, m.handlers...) + m.mu.Unlock() + for _, handler := range handlers { + handler() + } + return nil +} + +func (m *mockChannelCachePubSub) SubscribeUpdates(_ context.Context, handler func()) { + m.mu.Lock() + m.handlers = append(m.handlers, handler) + m.mu.Unlock() +} + func (m *mockChannelAuthCacheInvalidator) InvalidateAuthCacheByKey(_ context.Context, key string) { m.invalidatedKeys = append(m.invalidatedKeys, key) } @@ -192,11 +214,11 @@ func (m *mockChannelAuthCacheInvalidator) InvalidateAuthCacheByGroupID(_ context // --------------------------------------------------------------------------- func newTestChannelService(repo *mockChannelRepository) *ChannelService { - return NewChannelService(repo, nil, nil, nil) + return NewChannelService(repo, nil, nil, nil, nil) } func newTestChannelServiceWithAuth(repo *mockChannelRepository, auth *mockChannelAuthCacheInvalidator) *ChannelService { - return NewChannelService(repo, nil, auth, nil) + return NewChannelService(repo, nil, auth, nil, nil) } // makeStandardRepo returns a repo that serves one active channel with anthropic pricing @@ -1384,6 +1406,42 @@ func TestInvalidateCache(t *testing.T) { require.Equal(t, 2, callCount) // rebuilt } +func TestInvalidateCachePublishesToOtherInstances(t *testing.T) { + cachePubSub := &mockChannelCachePubSub{} + publisher := NewChannelService(&mockChannelRepository{}, nil, nil, nil, cachePubSub) + updated := false + subscriberRepo := &mockChannelRepository{ + listAllFn: func(_ context.Context) ([]Channel, error) { + model := "old-model" + if updated { + model = "new-model" + } + return []Channel{{ + ID: 1, + Status: StatusActive, + GroupIDs: []int64{10}, + ModelPricing: []ChannelModelPricing{{ + ID: 100, + Platform: PlatformAnthropic, + Models: []string{model}, + }}, + }}, nil + }, + getGroupPlatformsFn: func(_ context.Context, _ []int64) (map[int64]string, error) { + return map[int64]string{10: PlatformAnthropic}, nil + }, + } + subscriber := NewChannelService(subscriberRepo, nil, nil, nil, cachePubSub) + + require.NotNil(t, subscriber.GetChannelModelPricing(context.Background(), 10, "old-model")) + require.Nil(t, subscriber.GetChannelModelPricing(context.Background(), 10, "new-model")) + + updated = true + publisher.invalidateCache() + + require.NotNil(t, subscriber.GetChannelModelPricing(context.Background(), 10, "new-model")) +} + // =========================================================================== // 5. CRUD Methods // =========================================================================== diff --git a/backend/internal/service/model_plaza_service_test.go b/backend/internal/service/model_plaza_service_test.go index d54aca70ad..bb9d87b4b9 100644 --- a/backend/internal/service/model_plaza_service_test.go +++ b/backend/internal/service/model_plaza_service_test.go @@ -339,7 +339,7 @@ func newPlazaServiceWithBilling(channels []Channel, groups []Group, groupPlatfor return groupPlatforms, nil }, } - cs := NewChannelService(repo, nil, nil, nil) + cs := NewChannelService(repo, nil, nil, nil, nil) bs := NewBillingService(&config.Config{}, catalog) return NewModelPlazaService(repo, &stubGroupRepoForAvailable{activeGroups: groups}, catalog, bs, NewModelPricingResolver(cs, bs)) } diff --git a/backend/internal/service/model_pricing_resolver_test.go b/backend/internal/service/model_pricing_resolver_test.go index 3fb3e1c63f..a0373a1247 100644 --- a/backend/internal/service/model_pricing_resolver_test.go +++ b/backend/internal/service/model_pricing_resolver_test.go @@ -230,7 +230,7 @@ func newResolverWithChannel(t *testing.T, pricing []ChannelModelPricing) *ModelP return map[int64]string{groupID: "anthropic"}, nil }, } - cs := NewChannelService(repo, nil, nil, nil) + cs := NewChannelService(repo, nil, nil, nil, nil) bs := newTestBillingServiceForResolver() return NewModelPricingResolver(cs, bs) } @@ -566,7 +566,7 @@ func TestResolve_WithChannelOverride_CacheError(t *testing.T) { return nil, errors.New("database unavailable") }, } - cs := NewChannelService(repo, nil, nil, nil) + cs := NewChannelService(repo, nil, nil, nil, nil) bs := newTestBillingServiceForResolver() r := NewModelPricingResolver(cs, bs)