wire.go 7.6 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237
  1. package repository
  2. import (
  3. "context"
  4. "database/sql"
  5. "errors"
  6. entsql "entgo.io/ent/dialect/sql"
  7. "github.com/Wei-Shaw/sub2api/ent"
  8. "github.com/Wei-Shaw/sub2api/internal/config"
  9. "github.com/Wei-Shaw/sub2api/internal/service"
  10. "github.com/google/wire"
  11. "github.com/redis/go-redis/v9"
  12. )
  13. // ProvideConcurrencyCache 创建并发控制缓存,从配置读取 TTL 参数
  14. // 性能优化:TTL 可配置,支持长时间运行的 LLM 请求场景
  15. func ProvideConcurrencyCache(rdb *redis.Client, cfg *config.Config) service.ConcurrencyCache {
  16. waitTTLSeconds := int(cfg.Gateway.Scheduling.StickySessionWaitTimeout.Seconds())
  17. if cfg.Gateway.Scheduling.FallbackWaitTimeout > cfg.Gateway.Scheduling.StickySessionWaitTimeout {
  18. waitTTLSeconds = int(cfg.Gateway.Scheduling.FallbackWaitTimeout.Seconds())
  19. }
  20. if waitTTLSeconds <= 0 {
  21. waitTTLSeconds = cfg.Gateway.ConcurrencySlotTTLMinutes * 60
  22. }
  23. return NewConcurrencyCache(rdb, cfg.Gateway.ConcurrencySlotTTLMinutes, waitTTLSeconds)
  24. }
  25. // ProvideGitHubReleaseClient 创建 GitHub Release 客户端
  26. // 从配置中读取代理设置,支持国内服务器通过代理访问 GitHub
  27. func ProvideGitHubReleaseClient(cfg *config.Config) service.GitHubReleaseClient {
  28. return NewGitHubReleaseClient(cfg.Update.ProxyURL, cfg.Security.ProxyFallback.AllowDirectOnError)
  29. }
  30. // ProvidePricingRemoteClient 创建定价数据远程客户端
  31. // 从配置中读取代理设置,支持国内服务器通过代理访问 GitHub 上的定价数据
  32. func ProvidePricingRemoteClient(cfg *config.Config) service.PricingRemoteClient {
  33. return NewPricingRemoteClient(cfg.Update.ProxyURL, cfg.Security.ProxyFallback.AllowDirectOnError)
  34. }
  35. // ProvideSessionLimitCache 创建会话限制缓存
  36. // 用于 Anthropic OAuth/SetupToken 账号的并发会话数量控制
  37. func ProvideSessionLimitCache(rdb *redis.Client, cfg *config.Config) service.SessionLimitCache {
  38. defaultIdleTimeoutMinutes := 5 // 默认 5 分钟空闲超时
  39. if cfg != nil && cfg.Gateway.SessionIdleTimeoutMinutes > 0 {
  40. defaultIdleTimeoutMinutes = cfg.Gateway.SessionIdleTimeoutMinutes
  41. }
  42. return NewSessionLimitCache(rdb, defaultIdleTimeoutMinutes)
  43. }
  44. // ProvideSchedulerCache 创建调度快照缓存,并注入快照分块参数。
  45. func ProvideSchedulerCache(rdb *redis.Client, cfg *config.Config) service.SchedulerCache {
  46. mgetChunkSize := defaultSchedulerSnapshotMGetChunkSize
  47. writeChunkSize := defaultSchedulerSnapshotWriteChunkSize
  48. if cfg != nil {
  49. if cfg.Gateway.Scheduling.SnapshotMGetChunkSize > 0 {
  50. mgetChunkSize = cfg.Gateway.Scheduling.SnapshotMGetChunkSize
  51. }
  52. if cfg.Gateway.Scheduling.SnapshotWriteChunkSize > 0 {
  53. writeChunkSize = cfg.Gateway.Scheduling.SnapshotWriteChunkSize
  54. }
  55. }
  56. return newSchedulerCacheWithChunkSizes(rdb, mgetChunkSize, writeChunkSize)
  57. }
  58. // ProviderSet is the Wire provider set for all repositories
  59. var ProviderSet = wire.NewSet(
  60. NewIntelligentTestRepository,
  61. NewUserRepository,
  62. NewAPIKeyRepository,
  63. NewGroupRepository,
  64. NewAdminGroupRepository,
  65. NewCompositeModelRouteRepository,
  66. NewAccountRepository,
  67. NewAdminAccountRepository,
  68. NewScheduledTestPlanRepository, // 定时测试计划仓储
  69. NewScheduledTestResultRepository, // 定时测试结果仓储
  70. NewProxyRepository,
  71. NewRedeemCodeRepository,
  72. NewPromoCodeRepository,
  73. NewAnnouncementRepository,
  74. NewAnnouncementReadRepository,
  75. NewUsageLogRepository,
  76. NewUsageBillingRepository,
  77. NewBatchImageRepository,
  78. NewIdempotencyRepository,
  79. NewUsageCleanupRepository,
  80. NewDashboardAggregationRepository,
  81. NewSettingRepository,
  82. NewGlobalModelPricingRepository,
  83. NewSecurityPolicyRepository,
  84. NewSecurityPolicySessionStore,
  85. NewOpsRepository,
  86. NewAuditLogRepository,
  87. NewPasskeyRepository,
  88. NewPasskeySessionStore,
  89. NewUserSubscriptionRepository,
  90. NewUserAttributeDefinitionRepository,
  91. NewUserAttributeValueRepository,
  92. NewUserGroupRateRepository,
  93. NewErrorPassthroughRepository,
  94. NewTLSFingerprintProfileRepository,
  95. NewPluginRepository,
  96. NewChannelRepository,
  97. NewChannelMonitorRepository,
  98. NewChannelMonitorV2Repository,
  99. NewChannelMonitorRequestTemplateRepository,
  100. NewContentModerationRepository,
  101. NewAffiliateRepository,
  102. NewUserPlatformQuotaRepository, // T14: user × platform quota
  103. NewUserPlatformQuotaServiceAdapter, // T14: adapter → service.UserPlatformQuotaRepository
  104. // Cache implementations
  105. NewGatewayCache,
  106. NewBillingCache,
  107. NewAPIKeyCache,
  108. NewTempUnschedCache,
  109. NewTimeoutCounterCache,
  110. NewOpenAI403CounterCache,
  111. NewInternal500CounterCache,
  112. ProvideConcurrencyCache,
  113. ProvideSessionLimitCache,
  114. NewRPMCache,
  115. NewUserRPMCache,
  116. NewUserMsgQueueCache,
  117. NewDashboardCache,
  118. NewEmailCache,
  119. NewIdentityCache,
  120. NewRedeemCache,
  121. NewUpdateCache,
  122. NewGeminiTokenCache,
  123. NewImageTaskStore,
  124. NewBatchImageQueue,
  125. NewBatchImageDownloadLimiter,
  126. NewLeaderLockCache,
  127. ProvideSchedulerCache,
  128. NewSchedulerOutboxRepository,
  129. NewAuthCacheInvalidationOutboxRepository,
  130. NewProxyLatencyCache,
  131. NewTotpCache,
  132. NewRefreshTokenCache,
  133. NewErrorPassthroughCache,
  134. NewTLSFingerprintProfileCache,
  135. NewChannelCache,
  136. NewContentModerationHashCache,
  137. // Encryptors
  138. NewAESEncryptor,
  139. // Backup infrastructure
  140. NewPgDumper,
  141. NewS3BackupStoreFactory,
  142. // Image storage (async image task result offload)
  143. ProvideImageStorageFactory,
  144. // HTTP service ports (DI Strategy A: return interface directly)
  145. NewTurnstileVerifier,
  146. NewTencentCaptchaVerifier,
  147. NewAliyunCaptchaVerifier,
  148. ProvidePricingRemoteClient,
  149. ProvideGitHubReleaseClient,
  150. NewProxyExitInfoProber,
  151. NewClaudeUsageFetcher,
  152. NewClaudeOAuthClient,
  153. NewControlledHTTPUpstream,
  154. NewAccountTrafficCache,
  155. NewOpenAIOAuthClient,
  156. NewGrokOAuthClient,
  157. NewGeminiOAuthClient,
  158. NewGeminiCliCodeAssistClient,
  159. NewGeminiDriveClient,
  160. ProvideEnt,
  161. ProvideSQLDB,
  162. ProvideRedis,
  163. )
  164. // ProvideEnt 为依赖注入提供 Ent 客户端。
  165. //
  166. // 该函数是 InitEnt 的包装器,符合 Wire 的依赖提供函数签名要求。
  167. // Wire 会在编译时分析依赖关系,自动生成初始化代码。
  168. //
  169. // 依赖:config.Config
  170. // 提供:*ent.Client
  171. func ProvideEnt(cfg *config.Config) (*ent.Client, error) {
  172. client, _, err := InitEnt(cfg)
  173. return client, err
  174. }
  175. // ProvideImageStorageFactory 提供按需构造对象存储客户端的工厂。
  176. //
  177. // 这里返回工厂而不是实例:异步生图的开关与凭证可以在后台随时改动,客户端必须能在
  178. // 设置保存后重建,而不是在启动时定死一份。
  179. func ProvideImageStorageFactory() service.ImageStorageFactory {
  180. return func(ctx context.Context, cfg *config.ImageStorageConfig) (service.ImageStorage, error) {
  181. return NewS3ImageStorage(ctx, cfg)
  182. }
  183. }
  184. // ProvideSQLDB 从 Ent 客户端提取底层的 *sql.DB 连接。
  185. //
  186. // 某些 Repository 需要直接执行原生 SQL(如复杂的批量更新、聚合查询),
  187. // 此时需要访问底层的 sql.DB 而不是通过 Ent ORM。
  188. //
  189. // 设计说明:
  190. // - Ent 底层使用 sql.DB,通过 Driver 接口可以访问
  191. // - 这种设计允许在同一事务中混用 Ent 和原生 SQL
  192. //
  193. // 依赖:*ent.Client
  194. // 提供:*sql.DB
  195. func ProvideSQLDB(client *ent.Client) (*sql.DB, error) {
  196. if client == nil {
  197. return nil, errors.New("nil ent client")
  198. }
  199. // 从 Ent 客户端获取底层驱动
  200. drv, ok := client.Driver().(*entsql.Driver)
  201. if !ok {
  202. return nil, errors.New("ent driver does not expose *sql.DB")
  203. }
  204. // 返回驱动持有的 sql.DB 实例
  205. return drv.DB(), nil
  206. }
  207. // ProvideRedis 为依赖注入提供 Redis 客户端。
  208. //
  209. // Redis 用于:
  210. // - 分布式锁(如并发控制)
  211. // - 缓存(如用户会话、API 响应缓存)
  212. // - 速率限制
  213. // - 实时统计数据
  214. //
  215. // 依赖:config.Config
  216. // 提供:*redis.Client
  217. func ProvideRedis(cfg *config.Config) *redis.Client {
  218. return InitRedis(cfg)
  219. }