billing_cache_service.go 48 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192939495969798991001011021031041051061071081091101111121131141151161171181191201211221231241251261271281291301311321331341351361371381391401411421431441451461471481491501511521531541551561571581591601611621631641651661671681691701711721731741751761771781791801811821831841851861871881891901911921931941951961971981992002012022032042052062072082092102112122132142152162172182192202212222232242252262272282292302312322332342352362372382392402412422432442452462472482492502512522532542552562572582592602612622632642652662672682692702712722732742752762772782792802812822832842852862872882892902912922932942952962972982993003013023033043053063073083093103113123133143153163173183193203213223233243253263273283293303313323333343353363373383393403413423433443453463473483493503513523533543553563573583593603613623633643653663673683693703713723733743753763773783793803813823833843853863873883893903913923933943953963973983994004014024034044054064074084094104114124134144154164174184194204214224234244254264274284294304314324334344354364374384394404414424434444454464474484494504514524534544554564574584594604614624634644654664674684694704714724734744754764774784794804814824834844854864874884894904914924934944954964974984995005015025035045055065075085095105115125135145155165175185195205215225235245255265275285295305315325335345355365375385395405415425435445455465475485495505515525535545555565575585595605615625635645655665675685695705715725735745755765775785795805815825835845855865875885895905915925935945955965975985996006016026036046056066076086096106116126136146156166176186196206216226236246256266276286296306316326336346356366376386396406416426436446456466476486496506516526536546556566576586596606616626636646656666676686696706716726736746756766776786796806816826836846856866876886896906916926936946956966976986997007017027037047057067077087097107117127137147157167177187197207217227237247257267277287297307317327337347357367377387397407417427437447457467477487497507517527537547557567577587597607617627637647657667677687697707717727737747757767777787797807817827837847857867877887897907917927937947957967977987998008018028038048058068078088098108118128138148158168178188198208218228238248258268278288298308318328338348358368378388398408418428438448458468478488498508518528538548558568578588598608618628638648658668678688698708718728738748758768778788798808818828838848858868878888898908918928938948958968978988999009019029039049059069079089099109119129139149159169179189199209219229239249259269279289299309319329339349359369379389399409419429439449459469479489499509519529539549559569579589599609619629639649659669679689699709719729739749759769779789799809819829839849859869879889899909919929939949959969979989991000100110021003100410051006100710081009101010111012101310141015101610171018101910201021102210231024102510261027102810291030103110321033103410351036103710381039104010411042104310441045104610471048104910501051105210531054105510561057105810591060106110621063106410651066106710681069107010711072107310741075107610771078107910801081108210831084108510861087108810891090109110921093109410951096109710981099110011011102110311041105110611071108110911101111111211131114111511161117111811191120112111221123112411251126112711281129113011311132113311341135113611371138113911401141114211431144114511461147114811491150115111521153115411551156115711581159116011611162116311641165116611671168116911701171117211731174117511761177117811791180118111821183118411851186118711881189119011911192119311941195119611971198119912001201120212031204120512061207120812091210121112121213121412151216121712181219122012211222122312241225122612271228122912301231123212331234123512361237123812391240124112421243124412451246124712481249125012511252125312541255125612571258125912601261126212631264126512661267126812691270127112721273127412751276127712781279128012811282128312841285128612871288128912901291129212931294129512961297129812991300130113021303130413051306130713081309131013111312131313141315131613171318131913201321132213231324132513261327132813291330133113321333133413351336133713381339134013411342134313441345134613471348134913501351135213531354135513561357135813591360136113621363136413651366136713681369137013711372
  1. package service
  2. import (
  3. "context"
  4. "fmt"
  5. "strconv"
  6. "sync"
  7. "sync/atomic"
  8. "time"
  9. "github.com/Wei-Shaw/sub2api/internal/config"
  10. infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors"
  11. "github.com/Wei-Shaw/sub2api/internal/pkg/logger"
  12. "github.com/Wei-Shaw/sub2api/internal/pkg/timezone"
  13. "golang.org/x/sync/singleflight"
  14. )
  15. // 错误定义
  16. // 注:ErrInsufficientBalance在redeem_service.go中定义
  17. // 注:ErrDailyLimitExceeded/ErrWeeklyLimitExceeded/ErrMonthlyLimitExceeded在subscription_service.go中定义
  18. // errBillingCacheUnavailable 内部哨兵:用于 quota 校验路径在 cache==nil 时
  19. // 与"Redis 故障"走同一条 fail-open + DB 一次性检查的分支。
  20. var errBillingCacheUnavailable = fmt.Errorf("billing cache unavailable")
  21. var (
  22. ErrSubscriptionInvalid = infraerrors.Forbidden("SUBSCRIPTION_INVALID", "subscription is invalid or expired")
  23. ErrBillingServiceUnavailable = infraerrors.ServiceUnavailable("BILLING_SERVICE_ERROR", "Billing service temporarily unavailable. Please retry later.")
  24. // RPM 超限错误。gateway_handler 负责映射为 HTTP 429。
  25. ErrGroupRPMExceeded = infraerrors.TooManyRequests("GROUP_RPM_EXCEEDED", "group requests-per-minute limit exceeded")
  26. ErrUserRPMExceeded = infraerrors.TooManyRequests("USER_RPM_EXCEEDED", "user requests-per-minute limit exceeded")
  27. // user × platform quota(HTTP 429 Too Many Requests + Retry-After header)。
  28. // 选用 429 而非 403:限额耗尽属于"暂时性资源用尽,重试可恢复"的场景(RFC 6585),
  29. // 大量 SDK(如 OpenAI 兼容客户端)只对 429 触发自动退避并读取 Retry-After,
  30. // 用 403 会被视为"权限不足,重试无意义"导致客户端直接报错且不退避。
  31. ErrUserPlatformDailyQuotaExhausted = infraerrors.TooManyRequests("USER_PLATFORM_DAILY_QUOTA_EXHAUSTED", "Daily usage quota exhausted for this platform.")
  32. ErrUserPlatformWeeklyQuotaExhausted = infraerrors.TooManyRequests("USER_PLATFORM_WEEKLY_QUOTA_EXHAUSTED", "Weekly usage quota exhausted for this platform.")
  33. ErrUserPlatformMonthlyQuotaExhausted = infraerrors.TooManyRequests("USER_PLATFORM_MONTHLY_QUOTA_EXHAUSTED", "Monthly usage quota exhausted for this platform.")
  34. )
  35. // subscriptionCacheData 订阅缓存数据结构(内部使用)
  36. type subscriptionCacheData struct {
  37. Status string
  38. ExpiresAt time.Time
  39. DailyUsage float64
  40. WeeklyUsage float64
  41. MonthlyUsage float64
  42. Version int64
  43. }
  44. // 缓存写入任务类型
  45. type cacheWriteKind int
  46. const (
  47. cacheWriteSetBalance cacheWriteKind = iota
  48. cacheWriteSetSubscription
  49. cacheWriteUpdateSubscriptionUsage
  50. cacheWriteDeductBalance
  51. cacheWriteUpdateRateLimitUsage
  52. )
  53. // 异步缓存写入工作池配置
  54. //
  55. // 性能优化说明:
  56. // 原实现在请求热路径中使用 goroutine 异步更新缓存,存在以下问题:
  57. // 1. 每次请求创建新 goroutine,高并发下产生大量短生命周期 goroutine
  58. // 2. 无法控制并发数量,可能导致 Redis 连接耗尽
  59. // 3. goroutine 创建/销毁带来额外开销
  60. //
  61. // 新实现使用固定大小的工作池:
  62. // 1. 预创建 10 个 worker goroutine,避免频繁创建销毁
  63. // 2. 使用带缓冲的 channel(1000)作为任务队列,平滑写入峰值
  64. // 3. 非阻塞写入,队列满时关键任务同步回退,非关键任务丢弃并告警
  65. // 4. 统一超时控制,避免慢操作阻塞工作池
  66. const (
  67. cacheWriteWorkerCount = 10 // 工作协程数量
  68. cacheWriteBufferSize = 1000 // 任务队列缓冲大小
  69. cacheWriteTimeout = 2 * time.Second // 单个写入操作超时
  70. cacheWriteDropLogInterval = 5 * time.Second // 丢弃日志节流间隔
  71. balanceLoadTimeout = 3 * time.Second
  72. )
  73. // cacheWriteTask 缓存写入任务
  74. type cacheWriteTask struct {
  75. kind cacheWriteKind
  76. userID int64
  77. groupID int64
  78. apiKeyID int64
  79. balance float64
  80. amount float64
  81. subscriptionData *subscriptionCacheData
  82. }
  83. // apiKeyRateLimitLoader defines the interface for loading rate limit data from DB.
  84. type apiKeyRateLimitLoader interface {
  85. GetRateLimitData(ctx context.Context, keyID int64) (*APIKeyRateLimitData, error)
  86. }
  87. type subscriptionCacheInvalidationPubSub interface {
  88. PublishSubscriptionCacheInvalidation(ctx context.Context, cacheKey string) error
  89. SubscribeSubscriptionCacheInvalidation(ctx context.Context, handler func(cacheKey string)) error
  90. }
  91. // BillingCacheService 计费缓存服务
  92. // 负责余额和订阅数据的缓存管理,提供高性能的计费资格检查
  93. type BillingCacheService struct {
  94. cache BillingCache
  95. userRepo UserRepository
  96. subRepo UserSubscriptionRepository
  97. apiKeyRateLimitLoader apiKeyRateLimitLoader
  98. userRPMCache UserRPMCache
  99. userGroupRateRepo UserGroupRateRepository
  100. cfg *config.Config
  101. circuitBreaker *billingCircuitBreaker
  102. userPlatformQuotaRepo UserPlatformQuotaRepository
  103. cacheWriteChan chan cacheWriteTask
  104. cacheWriteWg sync.WaitGroup
  105. cacheWriteStopOnce sync.Once
  106. cacheWriteMu sync.RWMutex
  107. stopped atomic.Bool
  108. balanceLoadSF singleflight.Group
  109. quotaLoadSF singleflight.Group
  110. // 丢弃日志节流计数器(减少高负载下日志噪音)
  111. cacheWriteDropFullCount uint64
  112. cacheWriteDropFullLastLog int64
  113. cacheWriteDropClosedCount uint64
  114. cacheWriteDropClosedLastLog int64
  115. }
  116. // NewBillingCacheService 创建计费缓存服务
  117. func NewBillingCacheService(
  118. cache BillingCache,
  119. userRepo UserRepository,
  120. subRepo UserSubscriptionRepository,
  121. apiKeyRepo APIKeyRepository,
  122. userRPMCache UserRPMCache,
  123. userGroupRateRepo UserGroupRateRepository,
  124. cfg *config.Config,
  125. userPlatformQuotaRepo UserPlatformQuotaRepository,
  126. ) *BillingCacheService {
  127. svc := &BillingCacheService{
  128. cache: cache,
  129. userRepo: userRepo,
  130. subRepo: subRepo,
  131. apiKeyRateLimitLoader: apiKeyRepo,
  132. userRPMCache: userRPMCache,
  133. userGroupRateRepo: userGroupRateRepo,
  134. cfg: cfg,
  135. userPlatformQuotaRepo: userPlatformQuotaRepo,
  136. }
  137. svc.circuitBreaker = newBillingCircuitBreaker(cfg.Billing.CircuitBreaker)
  138. svc.startCacheWriteWorkers()
  139. return svc
  140. }
  141. // Stop 关闭缓存写入工作池
  142. func (s *BillingCacheService) Stop() {
  143. s.cacheWriteStopOnce.Do(func() {
  144. s.stopped.Store(true)
  145. s.cacheWriteMu.Lock()
  146. ch := s.cacheWriteChan
  147. if ch != nil {
  148. close(ch)
  149. }
  150. s.cacheWriteMu.Unlock()
  151. if ch == nil {
  152. return
  153. }
  154. s.cacheWriteWg.Wait()
  155. s.cacheWriteMu.Lock()
  156. if s.cacheWriteChan == ch {
  157. s.cacheWriteChan = nil
  158. }
  159. s.cacheWriteMu.Unlock()
  160. })
  161. }
  162. func (s *BillingCacheService) startCacheWriteWorkers() {
  163. ch := make(chan cacheWriteTask, cacheWriteBufferSize)
  164. s.cacheWriteChan = ch
  165. for i := 0; i < cacheWriteWorkerCount; i++ {
  166. s.cacheWriteWg.Add(1)
  167. go s.cacheWriteWorker(ch)
  168. }
  169. }
  170. // enqueueCacheWrite 尝试将任务入队,队列满时返回 false(并记录告警)。
  171. func (s *BillingCacheService) enqueueCacheWrite(task cacheWriteTask) (enqueued bool) {
  172. if s.stopped.Load() {
  173. s.logCacheWriteDrop(task, "closed")
  174. return false
  175. }
  176. s.cacheWriteMu.RLock()
  177. defer s.cacheWriteMu.RUnlock()
  178. if s.cacheWriteChan == nil {
  179. s.logCacheWriteDrop(task, "closed")
  180. return false
  181. }
  182. select {
  183. case s.cacheWriteChan <- task:
  184. return true
  185. default:
  186. // 队列满时不阻塞主流程,交由调用方决定是否同步回退。
  187. s.logCacheWriteDrop(task, "full")
  188. return false
  189. }
  190. }
  191. func (s *BillingCacheService) cacheWriteWorker(ch <-chan cacheWriteTask) {
  192. defer s.cacheWriteWg.Done()
  193. for task := range ch {
  194. ctx, cancel := context.WithTimeout(context.Background(), cacheWriteTimeout)
  195. switch task.kind {
  196. case cacheWriteSetBalance:
  197. s.setBalanceCache(ctx, task.userID, task.balance)
  198. case cacheWriteSetSubscription:
  199. s.setSubscriptionCache(ctx, task.userID, task.groupID, task.subscriptionData)
  200. case cacheWriteUpdateSubscriptionUsage:
  201. if s.cache != nil {
  202. if err := s.cache.UpdateSubscriptionUsage(ctx, task.userID, task.groupID, task.amount); err != nil {
  203. logger.LegacyPrintf("service.billing_cache", "Warning: update subscription cache failed for user %d group %d: %v", task.userID, task.groupID, err)
  204. }
  205. }
  206. case cacheWriteDeductBalance:
  207. if s.cache != nil {
  208. if err := s.cache.DeductUserBalance(ctx, task.userID, task.amount); err != nil {
  209. logger.LegacyPrintf("service.billing_cache", "Warning: deduct balance cache failed for user %d: %v", task.userID, err)
  210. }
  211. }
  212. case cacheWriteUpdateRateLimitUsage:
  213. if s.cache != nil {
  214. if err := s.cache.UpdateAPIKeyRateLimitUsage(ctx, task.apiKeyID, task.amount); err != nil {
  215. logger.LegacyPrintf("service.billing_cache", "Warning: update rate limit usage cache failed for api key %d: %v", task.apiKeyID, err)
  216. }
  217. }
  218. }
  219. cancel()
  220. }
  221. }
  222. // cacheWriteKindName 用于日志中的任务类型标识,便于排查丢弃原因。
  223. func cacheWriteKindName(kind cacheWriteKind) string {
  224. switch kind {
  225. case cacheWriteSetBalance:
  226. return "set_balance"
  227. case cacheWriteSetSubscription:
  228. return "set_subscription"
  229. case cacheWriteUpdateSubscriptionUsage:
  230. return "update_subscription_usage"
  231. case cacheWriteDeductBalance:
  232. return "deduct_balance"
  233. case cacheWriteUpdateRateLimitUsage:
  234. return "update_rate_limit_usage"
  235. default:
  236. return "unknown"
  237. }
  238. }
  239. // logCacheWriteDrop 使用节流方式记录丢弃情况,并汇总丢弃数量。
  240. func (s *BillingCacheService) logCacheWriteDrop(task cacheWriteTask, reason string) {
  241. var (
  242. countPtr *uint64
  243. lastPtr *int64
  244. )
  245. switch reason {
  246. case "full":
  247. countPtr = &s.cacheWriteDropFullCount
  248. lastPtr = &s.cacheWriteDropFullLastLog
  249. case "closed":
  250. countPtr = &s.cacheWriteDropClosedCount
  251. lastPtr = &s.cacheWriteDropClosedLastLog
  252. default:
  253. return
  254. }
  255. atomic.AddUint64(countPtr, 1)
  256. now := time.Now().UnixNano()
  257. last := atomic.LoadInt64(lastPtr)
  258. if now-last < int64(cacheWriteDropLogInterval) {
  259. return
  260. }
  261. if !atomic.CompareAndSwapInt64(lastPtr, last, now) {
  262. return
  263. }
  264. dropped := atomic.SwapUint64(countPtr, 0)
  265. if dropped == 0 {
  266. return
  267. }
  268. logger.LegacyPrintf("service.billing_cache", "Warning: cache write queue %s, dropped %d tasks in last %s (latest kind=%s user %d group %d)",
  269. reason,
  270. dropped,
  271. cacheWriteDropLogInterval,
  272. cacheWriteKindName(task.kind),
  273. task.userID,
  274. task.groupID,
  275. )
  276. }
  277. // ============================================
  278. // 余额缓存方法
  279. // ============================================
  280. // GetUserBalance 获取用户余额(优先从缓存读取)
  281. func (s *BillingCacheService) GetUserBalance(ctx context.Context, userID int64) (float64, error) {
  282. if s.cache == nil {
  283. // Redis不可用,直接查询数据库
  284. return s.getUserBalanceFromDB(ctx, userID)
  285. }
  286. // 尝试从缓存读取
  287. balance, err := s.cache.GetUserBalance(ctx, userID)
  288. if err == nil {
  289. return balance, nil
  290. }
  291. // 缓存未命中:singleflight 合并同一 userID 的并发回源请求。
  292. value, err, _ := s.balanceLoadSF.Do(strconv.FormatInt(userID, 10), func() (any, error) {
  293. loadCtx, cancel := context.WithTimeout(context.Background(), balanceLoadTimeout)
  294. defer cancel()
  295. balance, err := s.getUserBalanceFromDB(loadCtx, userID)
  296. if err != nil {
  297. return nil, err
  298. }
  299. // 异步建立缓存
  300. _ = s.enqueueCacheWrite(cacheWriteTask{
  301. kind: cacheWriteSetBalance,
  302. userID: userID,
  303. balance: balance,
  304. })
  305. return balance, nil
  306. })
  307. if err != nil {
  308. return 0, err
  309. }
  310. balance, ok := value.(float64)
  311. if !ok {
  312. return 0, fmt.Errorf("unexpected balance type: %T", value)
  313. }
  314. return balance, nil
  315. }
  316. // getUserBalanceFromDB 从数据库获取用户余额
  317. func (s *BillingCacheService) getUserBalanceFromDB(ctx context.Context, userID int64) (float64, error) {
  318. user, err := s.userRepo.GetByID(ctx, userID)
  319. if err != nil {
  320. return 0, fmt.Errorf("get user balance: %w", err)
  321. }
  322. return user.Balance, nil
  323. }
  324. // setBalanceCache 设置余额缓存
  325. func (s *BillingCacheService) setBalanceCache(ctx context.Context, userID int64, balance float64) {
  326. if s.cache == nil {
  327. return
  328. }
  329. if err := s.cache.SetUserBalance(ctx, userID, balance); err != nil {
  330. logger.LegacyPrintf("service.billing_cache", "Warning: set balance cache failed for user %d: %v", userID, err)
  331. }
  332. }
  333. // DeductBalanceCache 扣减余额缓存(同步调用)
  334. func (s *BillingCacheService) DeductBalanceCache(ctx context.Context, userID int64, amount float64) error {
  335. if s.cache == nil {
  336. return nil
  337. }
  338. return s.cache.DeductUserBalance(ctx, userID, amount)
  339. }
  340. // QueueDeductBalance 异步扣减余额缓存
  341. func (s *BillingCacheService) QueueDeductBalance(userID int64, amount float64) {
  342. if s.cache == nil {
  343. return
  344. }
  345. // 队列满时同步回退,避免关键扣减被静默丢弃。
  346. if s.enqueueCacheWrite(cacheWriteTask{
  347. kind: cacheWriteDeductBalance,
  348. userID: userID,
  349. amount: amount,
  350. }) {
  351. return
  352. }
  353. ctx, cancel := context.WithTimeout(context.Background(), cacheWriteTimeout)
  354. defer cancel()
  355. if err := s.DeductBalanceCache(ctx, userID, amount); err != nil {
  356. logger.LegacyPrintf("service.billing_cache", "Warning: deduct balance cache fallback failed for user %d: %v", userID, err)
  357. }
  358. }
  359. // InvalidateUserBalance 失效用户余额缓存
  360. func (s *BillingCacheService) InvalidateUserBalance(ctx context.Context, userID int64) error {
  361. if s.cache == nil {
  362. return nil
  363. }
  364. if err := s.cache.InvalidateUserBalance(ctx, userID); err != nil {
  365. logger.LegacyPrintf("service.billing_cache", "Warning: invalidate balance cache failed for user %d: %v", userID, err)
  366. return err
  367. }
  368. return nil
  369. }
  370. // ============================================
  371. // 订阅缓存方法
  372. // ============================================
  373. // GetSubscriptionStatus 获取订阅状态(优先从缓存读取)
  374. func (s *BillingCacheService) GetSubscriptionStatus(ctx context.Context, userID, groupID int64) (*subscriptionCacheData, error) {
  375. if s.cache == nil {
  376. return s.getSubscriptionFromDB(ctx, userID, groupID)
  377. }
  378. // 尝试从缓存读取
  379. cacheData, err := s.cache.GetSubscriptionCache(ctx, userID, groupID)
  380. if err == nil && cacheData != nil {
  381. return s.convertFromPortsData(cacheData), nil
  382. }
  383. // 缓存未命中,从数据库读取
  384. data, err := s.getSubscriptionFromDB(ctx, userID, groupID)
  385. if err != nil {
  386. return nil, err
  387. }
  388. // 异步建立缓存
  389. _ = s.enqueueCacheWrite(cacheWriteTask{
  390. kind: cacheWriteSetSubscription,
  391. userID: userID,
  392. groupID: groupID,
  393. subscriptionData: data,
  394. })
  395. return data, nil
  396. }
  397. func (s *BillingCacheService) convertFromPortsData(data *SubscriptionCacheData) *subscriptionCacheData {
  398. return &subscriptionCacheData{
  399. Status: data.Status,
  400. ExpiresAt: data.ExpiresAt,
  401. DailyUsage: data.DailyUsage,
  402. WeeklyUsage: data.WeeklyUsage,
  403. MonthlyUsage: data.MonthlyUsage,
  404. Version: data.Version,
  405. }
  406. }
  407. func (s *BillingCacheService) convertToPortsData(data *subscriptionCacheData) *SubscriptionCacheData {
  408. return &SubscriptionCacheData{
  409. Status: data.Status,
  410. ExpiresAt: data.ExpiresAt,
  411. DailyUsage: data.DailyUsage,
  412. WeeklyUsage: data.WeeklyUsage,
  413. MonthlyUsage: data.MonthlyUsage,
  414. Version: data.Version,
  415. }
  416. }
  417. // getSubscriptionFromDB 从数据库获取订阅数据
  418. func (s *BillingCacheService) getSubscriptionFromDB(ctx context.Context, userID, groupID int64) (*subscriptionCacheData, error) {
  419. sub, err := s.subRepo.GetActiveByUserIDAndGroupID(ctx, userID, groupID)
  420. if err != nil {
  421. return nil, fmt.Errorf("get subscription: %w", err)
  422. }
  423. return &subscriptionCacheData{
  424. Status: sub.Status,
  425. ExpiresAt: sub.ExpiresAt,
  426. DailyUsage: sub.DailyUsageUSD,
  427. WeeklyUsage: sub.WeeklyUsageUSD,
  428. MonthlyUsage: sub.MonthlyUsageUSD,
  429. Version: sub.UpdatedAt.Unix(),
  430. }, nil
  431. }
  432. // setSubscriptionCache 设置订阅缓存
  433. func (s *BillingCacheService) setSubscriptionCache(ctx context.Context, userID, groupID int64, data *subscriptionCacheData) {
  434. if s.cache == nil || data == nil {
  435. return
  436. }
  437. if err := s.cache.SetSubscriptionCache(ctx, userID, groupID, s.convertToPortsData(data)); err != nil {
  438. logger.LegacyPrintf("service.billing_cache", "Warning: set subscription cache failed for user %d group %d: %v", userID, groupID, err)
  439. }
  440. }
  441. // UpdateSubscriptionUsage 更新订阅用量缓存(同步调用)
  442. func (s *BillingCacheService) UpdateSubscriptionUsage(ctx context.Context, userID, groupID int64, costUSD float64) error {
  443. if s.cache == nil {
  444. return nil
  445. }
  446. return s.cache.UpdateSubscriptionUsage(ctx, userID, groupID, costUSD)
  447. }
  448. // QueueUpdateSubscriptionUsage 异步更新订阅用量缓存
  449. func (s *BillingCacheService) QueueUpdateSubscriptionUsage(userID, groupID int64, costUSD float64) {
  450. if s.cache == nil {
  451. return
  452. }
  453. // 队列满时同步回退,确保订阅用量及时更新。
  454. if s.enqueueCacheWrite(cacheWriteTask{
  455. kind: cacheWriteUpdateSubscriptionUsage,
  456. userID: userID,
  457. groupID: groupID,
  458. amount: costUSD,
  459. }) {
  460. return
  461. }
  462. ctx, cancel := context.WithTimeout(context.Background(), cacheWriteTimeout)
  463. defer cancel()
  464. if err := s.UpdateSubscriptionUsage(ctx, userID, groupID, costUSD); err != nil {
  465. logger.LegacyPrintf("service.billing_cache", "Warning: update subscription cache fallback failed for user %d group %d: %v", userID, groupID, err)
  466. }
  467. }
  468. // InvalidateSubscription 失效指定订阅缓存
  469. func (s *BillingCacheService) InvalidateSubscription(ctx context.Context, userID, groupID int64) error {
  470. if s.cache == nil {
  471. return nil
  472. }
  473. if err := s.cache.InvalidateSubscriptionCache(ctx, userID, groupID); err != nil {
  474. logger.LegacyPrintf("service.billing_cache", "Warning: invalidate subscription cache failed for user %d group %d: %v", userID, groupID, err)
  475. return err
  476. }
  477. return nil
  478. }
  479. func (s *BillingCacheService) PublishSubscriptionCacheInvalidation(ctx context.Context, cacheKey string) error {
  480. if s.cache == nil {
  481. return nil
  482. }
  483. pubsub, ok := s.cache.(subscriptionCacheInvalidationPubSub)
  484. if !ok {
  485. return nil
  486. }
  487. return pubsub.PublishSubscriptionCacheInvalidation(ctx, cacheKey)
  488. }
  489. func (s *BillingCacheService) SubscribeSubscriptionCacheInvalidation(ctx context.Context, handler func(cacheKey string)) error {
  490. if s.cache == nil {
  491. return nil
  492. }
  493. pubsub, ok := s.cache.(subscriptionCacheInvalidationPubSub)
  494. if !ok {
  495. return nil
  496. }
  497. return pubsub.SubscribeSubscriptionCacheInvalidation(ctx, handler)
  498. }
  499. // InvalidateAPIKeyRateLimit invalidates the Redis rate-limit usage cache for an API key.
  500. func (s *BillingCacheService) InvalidateAPIKeyRateLimit(ctx context.Context, keyID int64) error {
  501. if s.cache == nil {
  502. return nil
  503. }
  504. if err := s.cache.InvalidateAPIKeyRateLimit(ctx, keyID); err != nil {
  505. logger.LegacyPrintf("service.billing_cache", "Warning: invalidate api key rate limit cache failed for key %d: %v", keyID, err)
  506. return err
  507. }
  508. return nil
  509. }
  510. // ============================================
  511. // API Key 限速缓存方法
  512. // ============================================
  513. // checkAPIKeyRateLimits checks rate limit windows for an API key.
  514. // It loads usage from Redis cache (falling back to DB on cache miss),
  515. // resets expired windows in-memory and triggers async DB reset,
  516. // and returns an error if any window limit is exceeded.
  517. func (s *BillingCacheService) checkAPIKeyRateLimits(ctx context.Context, apiKey *APIKey) error {
  518. if s.cache == nil {
  519. // No cache: fall back to reading from DB directly
  520. if s.apiKeyRateLimitLoader == nil {
  521. return nil
  522. }
  523. data, err := s.apiKeyRateLimitLoader.GetRateLimitData(ctx, apiKey.ID)
  524. if err != nil {
  525. return nil // Don't block requests on DB errors
  526. }
  527. return s.evaluateRateLimits(ctx, apiKey, data.Usage5h, data.Usage1d, data.Usage7d,
  528. data.Window5hStart, data.Window1dStart, data.Window7dStart)
  529. }
  530. cacheData, err := s.cache.GetAPIKeyRateLimit(ctx, apiKey.ID)
  531. if err != nil {
  532. // Cache miss: load from DB and populate cache
  533. if s.apiKeyRateLimitLoader == nil {
  534. return nil
  535. }
  536. dbData, dbErr := s.apiKeyRateLimitLoader.GetRateLimitData(ctx, apiKey.ID)
  537. if dbErr != nil {
  538. return nil // Don't block requests on DB errors
  539. }
  540. // Build cache entry from DB data
  541. cacheEntry := &APIKeyRateLimitCacheData{
  542. Usage5h: dbData.Usage5h,
  543. Usage1d: dbData.Usage1d,
  544. Usage7d: dbData.Usage7d,
  545. }
  546. if dbData.Window5hStart != nil {
  547. cacheEntry.Window5h = dbData.Window5hStart.Unix()
  548. }
  549. if dbData.Window1dStart != nil {
  550. cacheEntry.Window1d = dbData.Window1dStart.Unix()
  551. }
  552. if dbData.Window7dStart != nil {
  553. cacheEntry.Window7d = dbData.Window7dStart.Unix()
  554. }
  555. _ = s.cache.SetAPIKeyRateLimit(ctx, apiKey.ID, cacheEntry)
  556. cacheData = cacheEntry
  557. }
  558. var w5h, w1d, w7d *time.Time
  559. if cacheData.Window5h > 0 {
  560. t := time.Unix(cacheData.Window5h, 0)
  561. w5h = &t
  562. }
  563. if cacheData.Window1d > 0 {
  564. t := time.Unix(cacheData.Window1d, 0)
  565. w1d = &t
  566. }
  567. if cacheData.Window7d > 0 {
  568. t := time.Unix(cacheData.Window7d, 0)
  569. w7d = &t
  570. }
  571. return s.evaluateRateLimits(ctx, apiKey, cacheData.Usage5h, cacheData.Usage1d, cacheData.Usage7d, w5h, w1d, w7d)
  572. }
  573. // evaluateRateLimits checks usage against limits, triggering async resets for expired windows.
  574. func (s *BillingCacheService) evaluateRateLimits(ctx context.Context, apiKey *APIKey, usage5h, usage1d, usage7d float64, w5h, w1d, w7d *time.Time) error {
  575. needsReset := false
  576. // Reset expired windows in-memory for check purposes
  577. if IsWindowExpired(w5h, RateLimitWindow5h) {
  578. usage5h = 0
  579. needsReset = true
  580. }
  581. if IsWindowExpired(w1d, RateLimitWindow1d) {
  582. usage1d = 0
  583. needsReset = true
  584. }
  585. if IsWindowExpired(w7d, RateLimitWindow7d) {
  586. usage7d = 0
  587. needsReset = true
  588. }
  589. // Trigger async DB reset if any window expired
  590. if needsReset {
  591. keyID := apiKey.ID
  592. go func() {
  593. resetCtx, cancel := context.WithTimeout(context.Background(), cacheWriteTimeout)
  594. defer cancel()
  595. if s.apiKeyRateLimitLoader != nil {
  596. // Use the repo directly - reset then reload cache
  597. if loader, ok := s.apiKeyRateLimitLoader.(interface {
  598. ResetRateLimitWindows(ctx context.Context, id int64) error
  599. }); ok {
  600. if err := loader.ResetRateLimitWindows(resetCtx, keyID); err != nil {
  601. logger.LegacyPrintf("service.billing_cache", "Warning: reset rate limit windows failed for api key %d: %v", keyID, err)
  602. }
  603. }
  604. }
  605. // Invalidate cache so next request loads fresh data
  606. if s.cache != nil {
  607. if err := s.cache.InvalidateAPIKeyRateLimit(resetCtx, keyID); err != nil {
  608. logger.LegacyPrintf("service.billing_cache", "Warning: invalidate rate limit cache failed for api key %d: %v", keyID, err)
  609. }
  610. }
  611. }()
  612. }
  613. // Check limits
  614. if apiKey.RateLimit5h > 0 && usage5h >= apiKey.RateLimit5h {
  615. return ErrAPIKeyRateLimit5hExceeded
  616. }
  617. if apiKey.RateLimit1d > 0 && usage1d >= apiKey.RateLimit1d {
  618. return ErrAPIKeyRateLimit1dExceeded
  619. }
  620. if apiKey.RateLimit7d > 0 && usage7d >= apiKey.RateLimit7d {
  621. return ErrAPIKeyRateLimit7dExceeded
  622. }
  623. return nil
  624. }
  625. // QueueUpdateAPIKeyRateLimitUsage asynchronously updates rate limit usage in the cache.
  626. func (s *BillingCacheService) QueueUpdateAPIKeyRateLimitUsage(apiKeyID int64, cost float64) {
  627. if s.cache == nil {
  628. return
  629. }
  630. s.enqueueCacheWrite(cacheWriteTask{
  631. kind: cacheWriteUpdateRateLimitUsage,
  632. apiKeyID: apiKeyID,
  633. amount: cost,
  634. })
  635. }
  636. // IncrementUserPlatformQuotaUsage 同步累加 user × platform usage 到 Redis 缓存。
  637. //
  638. // 设计:同步写入而非异步入队。同步写确保下次 preflight 立即看到最新 usage,
  639. // 把 TOCTOU 超支窗口限制在并发 in-flight 请求数量内(而非随时间无限累积)。
  640. // 写延迟通常 < 1ms(本地 Redis),换取 quota 视图实时性的取舍合理。
  641. //
  642. // Redis 写失败用 ALERT 级 log;DB 持久化由 caller 单独 goroutine 兜底(gateway_service.go)。
  643. func (s *BillingCacheService) IncrementUserPlatformQuotaUsage(userID int64, platform string, cost float64) {
  644. if s.cache == nil {
  645. return
  646. }
  647. if platform == "" || cost <= 0 {
  648. return
  649. }
  650. ctx, cancel := context.WithTimeout(context.Background(), cacheWriteTimeout)
  651. defer cancel()
  652. ttl := time.Duration(s.cfg.Billing.UserPlatformQuotaCacheTTLSeconds) * time.Second
  653. markDirty := s.cfg.Database.UserPlatformQuotaFlusherEnabled
  654. if err := s.cache.IncrUserPlatformQuotaUsageCache(ctx, userID, platform, cost, ttl, markDirty); err != nil {
  655. logger.LegacyPrintf("service.billing_cache",
  656. "ALERT: incr user platform quota cache failed user=%d platform=%s cost=%f: %v",
  657. userID, platform, cost, err)
  658. }
  659. }
  660. // ============================================
  661. // 统一检查方法
  662. // ============================================
  663. // CheckBillingEligibility 检查用户是否有资格发起请求
  664. // 余额模式:检查缓存余额 > 0
  665. // 订阅模式:检查缓存用量未超过限额(Group限额从参数传入)
  666. // platform 为请求的目标平台(如 "anthropic"),传空串 "" 时跳过 user × platform quota 检查。
  667. func (s *BillingCacheService) CheckBillingEligibility(ctx context.Context, user *User, apiKey *APIKey, group *Group, subscription *UserSubscription, platform string) error {
  668. // 简易模式:跳过所有计费检查
  669. if s.cfg.RunMode == config.RunModeSimple {
  670. return nil
  671. }
  672. if s.circuitBreaker != nil && !s.circuitBreaker.Allow() {
  673. return ErrBillingServiceUnavailable
  674. }
  675. // 判断计费模式
  676. isSubscriptionMode := group != nil && group.IsSubscriptionType() && subscription != nil
  677. if isSubscriptionMode {
  678. if err := s.checkSubscriptionEligibility(ctx, user.ID, group, subscription); err != nil {
  679. return err
  680. }
  681. } else {
  682. if err := s.checkBalanceEligibility(ctx, user.ID); err != nil {
  683. return err
  684. }
  685. }
  686. // user × platform quota 仅在 standard(余额)模式生效;订阅模式豁免
  687. if !isSubscriptionMode {
  688. if err := s.checkUserPlatformQuotaEligibility(ctx, user.ID, platform); err != nil {
  689. return err
  690. }
  691. }
  692. // Check API Key rate limits (applies to both billing modes)
  693. if apiKey != nil && apiKey.HasRateLimits() {
  694. if err := s.checkAPIKeyRateLimits(ctx, apiKey); err != nil {
  695. return err
  696. }
  697. }
  698. // RPM 限流:级联回落(Override → Group → User),放在最后以避免为注定失败的请求增加计数。
  699. if err := s.checkRPM(ctx, user, group); err != nil {
  700. return err
  701. }
  702. return nil
  703. }
  704. // checkRPM 执行并行 RPM 限流,所有适用的限制同时生效,任一超限即拒绝:
  705. //
  706. // 1. (用户, 分组) rpm_override — 最细粒度:管理员为特定用户在特定分组设定的专属限额。
  707. // override=0 表示该用户在该分组免检(绿灯),但 user 级全局上限仍然生效。
  708. // 2. group.rpm_limit — 分组级:该分组的统一 RPM 容量(仅当无 override 时生效)。
  709. // 3. user.rpm_limit — 用户级全局硬上限:无论 override/group 如何配置,始终生效。
  710. //
  711. // 与旧版"级联互斥"设计不同,新版确保 user.rpm_limit 作为全局天花板不会被 group 或 override 覆盖。
  712. // Redis 故障一律 fail-open(打 warning,不阻塞业务)。
  713. func (s *BillingCacheService) checkRPM(ctx context.Context, user *User, group *Group) error {
  714. if s == nil || s.userRPMCache == nil || user == nil {
  715. return nil
  716. }
  717. // ── 第一层:分组级检查(override 或 group.rpm_limit) ──
  718. if group != nil {
  719. // 解析 override:优先从 auth cache snapshot,nil 时回退 DB。
  720. var override *int
  721. if user.UserGroupRPMOverride != nil {
  722. override = user.UserGroupRPMOverride
  723. } else if s.userGroupRateRepo != nil {
  724. dbOverride, err := s.userGroupRateRepo.GetRPMOverrideByUserAndGroup(ctx, user.ID, group.ID)
  725. if err != nil {
  726. logger.LegacyPrintf(
  727. "service.billing_cache",
  728. "Warning: rpm override lookup failed for user=%d group=%d: %v",
  729. user.ID, group.ID, err,
  730. )
  731. } else {
  732. override = dbOverride
  733. }
  734. }
  735. if override != nil {
  736. // override=0 → 该用户在该分组免检(但 user 级仍会在下面检查)。
  737. if *override > 0 {
  738. count, incErr := s.userRPMCache.IncrementUserGroupRPM(ctx, user.ID, group.ID)
  739. if incErr != nil {
  740. logger.LegacyPrintf(
  741. "service.billing_cache",
  742. "Warning: rpm increment (override) failed for user=%d group=%d: %v",
  743. user.ID, group.ID, incErr,
  744. )
  745. // fail-open
  746. } else if count > *override {
  747. return ErrGroupRPMExceeded
  748. }
  749. }
  750. // override 命中后跳过 group.rpm_limit(override 替代 group),但不 return——继续检查 user 级。
  751. } else if group.RPMLimit > 0 {
  752. // 无 override,检查 group.rpm_limit。
  753. count, err := s.userRPMCache.IncrementUserGroupRPM(ctx, user.ID, group.ID)
  754. if err != nil {
  755. logger.LegacyPrintf(
  756. "service.billing_cache",
  757. "Warning: rpm increment (group) failed for user=%d group=%d: %v",
  758. user.ID, group.ID, err,
  759. )
  760. // fail-open
  761. } else if count > group.RPMLimit {
  762. return ErrGroupRPMExceeded
  763. }
  764. }
  765. }
  766. // ── 第二层:用户级全局硬上限(始终生效) ──
  767. if user.RPMLimit > 0 {
  768. count, err := s.userRPMCache.IncrementUserRPM(ctx, user.ID)
  769. if err != nil {
  770. logger.LegacyPrintf(
  771. "service.billing_cache",
  772. "Warning: rpm increment (user) failed for user=%d: %v",
  773. user.ID, err,
  774. )
  775. return nil // fail-open
  776. }
  777. if count > user.RPMLimit {
  778. return ErrUserRPMExceeded
  779. }
  780. }
  781. return nil
  782. }
  783. func (s *BillingCacheService) minimumBalanceReserve() float64 {
  784. if s == nil || s.cfg == nil || s.cfg.Billing.MinimumBalanceReserve <= 0 {
  785. return 0
  786. }
  787. return s.cfg.Billing.MinimumBalanceReserve
  788. }
  789. func (s *BillingCacheService) balanceBelowEligibilityThreshold(balance float64) bool {
  790. if balance <= 0 {
  791. return true
  792. }
  793. minimumReserve := s.minimumBalanceReserve()
  794. return minimumReserve > 0 && balance < minimumReserve
  795. }
  796. // checkBalanceEligibility 检查余额模式资格
  797. func (s *BillingCacheService) checkBalanceEligibility(ctx context.Context, userID int64) error {
  798. balance, err := s.GetUserBalance(ctx, userID)
  799. if err != nil {
  800. if s.circuitBreaker != nil {
  801. s.circuitBreaker.OnFailure(err)
  802. }
  803. logger.LegacyPrintf("service.billing_cache", "ALERT: billing balance check failed for user %d: %v", userID, err)
  804. return ErrBillingServiceUnavailable.WithCause(err)
  805. }
  806. if s.circuitBreaker != nil {
  807. s.circuitBreaker.OnSuccess()
  808. }
  809. if s.balanceBelowEligibilityThreshold(balance) {
  810. return ErrInsufficientBalance
  811. }
  812. return nil
  813. }
  814. // checkSubscriptionEligibility 检查订阅模式资格
  815. func (s *BillingCacheService) checkSubscriptionEligibility(ctx context.Context, userID int64, group *Group, subscription *UserSubscription) error {
  816. // 获取订阅缓存数据
  817. subData, err := s.GetSubscriptionStatus(ctx, userID, group.ID)
  818. if err != nil {
  819. if s.circuitBreaker != nil {
  820. s.circuitBreaker.OnFailure(err)
  821. }
  822. logger.LegacyPrintf("service.billing_cache", "ALERT: billing subscription check failed for user %d group %d: %v", userID, group.ID, err)
  823. return ErrBillingServiceUnavailable.WithCause(err)
  824. }
  825. if s.circuitBreaker != nil {
  826. s.circuitBreaker.OnSuccess()
  827. }
  828. // 检查订阅状态
  829. if subData.Status != SubscriptionStatusActive {
  830. return ErrSubscriptionInvalid
  831. }
  832. // 检查是否过期
  833. if time.Now().After(subData.ExpiresAt) {
  834. return ErrSubscriptionInvalid
  835. }
  836. // 检查限额(使用传入的Group限额配置)
  837. if group.HasDailyLimit() && subData.DailyUsage >= *group.DailyLimitUSD {
  838. return ErrDailyLimitExceeded
  839. }
  840. if group.HasWeeklyLimit() && subData.WeeklyUsage >= *group.WeeklyLimitUSD {
  841. return ErrWeeklyLimitExceeded
  842. }
  843. if group.HasMonthlyLimit() && subData.MonthlyUsage >= *group.MonthlyLimitUSD {
  844. return ErrMonthlyLimitExceeded
  845. }
  846. return nil
  847. }
  848. type billingCircuitBreakerState int
  849. const (
  850. billingCircuitClosed billingCircuitBreakerState = iota
  851. billingCircuitOpen
  852. billingCircuitHalfOpen
  853. )
  854. type billingCircuitBreaker struct {
  855. mu sync.Mutex
  856. state billingCircuitBreakerState
  857. failures int
  858. openedAt time.Time
  859. failureThreshold int
  860. resetTimeout time.Duration
  861. halfOpenRequests int
  862. halfOpenRemaining int
  863. }
  864. func newBillingCircuitBreaker(cfg config.CircuitBreakerConfig) *billingCircuitBreaker {
  865. if !cfg.Enabled {
  866. return nil
  867. }
  868. resetTimeout := time.Duration(cfg.ResetTimeoutSeconds) * time.Second
  869. if resetTimeout <= 0 {
  870. resetTimeout = 30 * time.Second
  871. }
  872. halfOpen := cfg.HalfOpenRequests
  873. if halfOpen <= 0 {
  874. halfOpen = 1
  875. }
  876. threshold := cfg.FailureThreshold
  877. if threshold <= 0 {
  878. threshold = 5
  879. }
  880. return &billingCircuitBreaker{
  881. state: billingCircuitClosed,
  882. failureThreshold: threshold,
  883. resetTimeout: resetTimeout,
  884. halfOpenRequests: halfOpen,
  885. }
  886. }
  887. func (b *billingCircuitBreaker) Allow() bool {
  888. b.mu.Lock()
  889. defer b.mu.Unlock()
  890. switch b.state {
  891. case billingCircuitClosed:
  892. return true
  893. case billingCircuitOpen:
  894. if time.Since(b.openedAt) < b.resetTimeout {
  895. return false
  896. }
  897. b.state = billingCircuitHalfOpen
  898. b.halfOpenRemaining = b.halfOpenRequests
  899. logger.LegacyPrintf("service.billing_cache", "ALERT: billing circuit breaker entering half-open state")
  900. fallthrough
  901. case billingCircuitHalfOpen:
  902. if b.halfOpenRemaining <= 0 {
  903. return false
  904. }
  905. b.halfOpenRemaining--
  906. return true
  907. default:
  908. return false
  909. }
  910. }
  911. func (b *billingCircuitBreaker) OnFailure(err error) {
  912. if b == nil {
  913. return
  914. }
  915. b.mu.Lock()
  916. defer b.mu.Unlock()
  917. switch b.state {
  918. case billingCircuitOpen:
  919. return
  920. case billingCircuitHalfOpen:
  921. b.state = billingCircuitOpen
  922. b.openedAt = time.Now()
  923. b.halfOpenRemaining = 0
  924. logger.LegacyPrintf("service.billing_cache", "ALERT: billing circuit breaker opened after half-open failure: %v", err)
  925. return
  926. default:
  927. b.failures++
  928. if b.failures >= b.failureThreshold {
  929. b.state = billingCircuitOpen
  930. b.openedAt = time.Now()
  931. b.halfOpenRemaining = 0
  932. logger.LegacyPrintf("service.billing_cache", "ALERT: billing circuit breaker opened after %d failures: %v", b.failures, err)
  933. }
  934. }
  935. }
  936. func (b *billingCircuitBreaker) OnSuccess() {
  937. if b == nil {
  938. return
  939. }
  940. b.mu.Lock()
  941. defer b.mu.Unlock()
  942. previousState := b.state
  943. previousFailures := b.failures
  944. b.state = billingCircuitClosed
  945. b.failures = 0
  946. b.halfOpenRemaining = 0
  947. // 只有状态真正发生变化时才记录日志
  948. if previousState != billingCircuitClosed {
  949. logger.LegacyPrintf("service.billing_cache", "ALERT: billing circuit breaker closed (was %s)", circuitStateString(previousState))
  950. } else if previousFailures > 0 {
  951. logger.LegacyPrintf("service.billing_cache", "INFO: billing circuit breaker failures reset from %d", previousFailures)
  952. }
  953. }
  954. func circuitStateString(state billingCircuitBreakerState) string {
  955. switch state {
  956. case billingCircuitClosed:
  957. return "closed"
  958. case billingCircuitOpen:
  959. return "open"
  960. case billingCircuitHalfOpen:
  961. return "half-open"
  962. default:
  963. return "unknown"
  964. }
  965. }
  966. // checkUserPlatformQuotaEligibility 在 standard 模式下检查 user × platform 日/周/月 quota。
  967. // 返回 nil = 允许;返回 ErrUserPlatform{Daily/Weekly/Monthly}QuotaExhausted = 拒绝(带 window_resets_at metadata)。
  968. // checkUserPlatformQuotaEligibility 检查用户在指定平台的 USD 配额。
  969. //
  970. // 流程(Redis-first / DB-fallback):
  971. // 1. 先读 Redis cache;若命中且 SchemaVersion==1,直接用 entry 中的 limits 和 window_start 做校验,
  972. // 免除 DB 查询。
  973. // 2. cache MISS 或旧版 entry(SchemaVersion==0)→ 查 DB 回填完整 entry(含 limits/window_start)。
  974. // 3. Redis 故障(err != nil)→ fail-open,查 DB 做一次性检查,不回填。
  975. func (s *BillingCacheService) checkUserPlatformQuotaEligibility(
  976. ctx context.Context,
  977. userID int64,
  978. platform string,
  979. ) error {
  980. if platform == "" || s.userPlatformQuotaRepo == nil {
  981. return nil
  982. }
  983. // cache 未配置(如简化部署 / 单测路径)→ 直接走 DB 查询,避免 nil panic。
  984. // 其他 check* 方法(balance/subscription/rate-limit)也有类似守卫。
  985. var (
  986. entry *UserPlatformQuotaCacheEntry
  987. ok bool
  988. cacheErr error
  989. )
  990. if s.cache != nil {
  991. entry, ok, cacheErr = s.cache.GetUserPlatformQuotaCache(ctx, userID, platform)
  992. } else {
  993. // 标记为"cache 故障"分支:跳过 HIT 路径、不回填、走 DB 一次性检查
  994. cacheErr = errBillingCacheUnavailable
  995. }
  996. // --- cache HIT with current schema → 直接用 entry,不查 DB ---
  997. if cacheErr == nil && ok && entry != nil && entry.SchemaVersion == UserPlatformQuotaCacheSchemaV1 {
  998. now := time.Now()
  999. dailyUsage := entry.DailyUsageUSD
  1000. weeklyUsage := entry.WeeklyUsageUSD
  1001. monthlyUsage := entry.MonthlyUsageUSD
  1002. // 若窗口已更新(DB 已重置但 cache 尚未失效),将对应 usage 清零再做比较,
  1003. // 同时记录新窗口起点用于后续刷新 cache entry。
  1004. // 本次请求用本地清零值继续判断;DB 层 IncrementUsageWithReset 已有窗口自愈能力,
  1005. // 持久化数据始终正确。
  1006. windowExpired := false
  1007. newDailyStart := entry.DailyWindowStart
  1008. newWeeklyStart := entry.WeeklyWindowStart
  1009. newMonthlyStart := entry.MonthlyWindowStart
  1010. if quotaWindowExpired(entry.DailyWindowStart, timezone.StartOfDay(now)) {
  1011. dailyUsage = 0
  1012. windowExpired = true
  1013. dayStart := timezone.StartOfDay(now)
  1014. newDailyStart = &dayStart
  1015. }
  1016. if quotaWindowExpired(entry.WeeklyWindowStart, timezone.StartOfWeek(now)) {
  1017. weeklyUsage = 0
  1018. windowExpired = true
  1019. weekStart := timezone.StartOfWeek(now)
  1020. newWeeklyStart = &weekStart
  1021. }
  1022. if monthlyQuotaWindowExpired(entry.MonthlyWindowStart, now) {
  1023. monthlyUsage = 0
  1024. windowExpired = true
  1025. monthStart := now
  1026. newMonthlyStart = &monthStart
  1027. }
  1028. // 检测到任意窗口过期:用 reset 后的 entry 覆盖 Redis(而非 Delete)。
  1029. // 旧实现 Delete 后,期间到达的 IncrUserPlatformQuotaUsage 调用让 Lua 看到
  1030. // EXISTS=0 直接 return 0,并发请求的 cost 永久丢失,直到下次 cache MISS 回填。
  1031. // 改为 SetCache 原子覆盖:key 不断链,Lua INCR 可在新窗口 entry 上正确累加。
  1032. // 超时 50ms:覆盖正常路径与可接受抖动;Redis 异常时 hot path 不阻塞超过此值。
  1033. // 用 context.Background()+短超时,避免请求 ctx 取消导致刷新丢失。
  1034. // 显式 setCancel()(而非 defer):缩短 context 生命周期,避免 defer 延迟到函数返回。
  1035. // isSentinel 判定「该 entry 无任何 limit」,涵盖两类,跨窗口命中时都跳过 refresh:
  1036. // 1) A3 回填的 sentinel(DB 无行,短 TTL):refresh 会把短 TTL 误升级为 86400s,有害;
  1037. // 2) DB 有行但三 limit 全未配置的用户(TTL 86400s):refresh 纯属无意义(TTL 升级本身无害)。
  1038. // 两类的 enforcement(下方 limit!=nil 比较)都因 limit 全 nil 永远放行,跳过 refresh 均正确。
  1039. isSentinel := entry.DailyLimitUSD == nil && entry.WeeklyLimitUSD == nil && entry.MonthlyLimitUSD == nil
  1040. if windowExpired && s.cache != nil && !isSentinel {
  1041. refreshed := &UserPlatformQuotaCacheEntry{
  1042. DailyUsageUSD: dailyUsage,
  1043. WeeklyUsageUSD: weeklyUsage,
  1044. MonthlyUsageUSD: monthlyUsage,
  1045. SchemaVersion: UserPlatformQuotaCacheSchemaV1,
  1046. DailyLimitUSD: entry.DailyLimitUSD,
  1047. WeeklyLimitUSD: entry.WeeklyLimitUSD,
  1048. MonthlyLimitUSD: entry.MonthlyLimitUSD,
  1049. DailyWindowStart: newDailyStart,
  1050. WeeklyWindowStart: newWeeklyStart,
  1051. MonthlyWindowStart: newMonthlyStart,
  1052. }
  1053. ttl := time.Duration(s.cfg.Billing.UserPlatformQuotaCacheTTLSeconds) * time.Second
  1054. setCtx, setCancel := context.WithTimeout(context.Background(), 50*time.Millisecond)
  1055. if setErr := s.cache.SetUserPlatformQuotaCache(setCtx, userID, platform, refreshed, ttl); setErr != nil {
  1056. logger.LegacyPrintf("service.billing_cache",
  1057. "Warning: refresh expired user platform quota cache failed user=%d platform=%s: %v",
  1058. userID, platform, setErr)
  1059. }
  1060. setCancel()
  1061. }
  1062. if entry.DailyLimitUSD != nil && dailyUsage >= *entry.DailyLimitUSD {
  1063. return withWindowResetsMetadata(ErrUserPlatformDailyQuotaExhausted, nextDailyReset(now))
  1064. }
  1065. if entry.WeeklyLimitUSD != nil && weeklyUsage >= *entry.WeeklyLimitUSD {
  1066. return withWindowResetsMetadata(ErrUserPlatformWeeklyQuotaExhausted, nextWeeklyReset(now))
  1067. }
  1068. if entry.MonthlyLimitUSD != nil && monthlyUsage >= *entry.MonthlyLimitUSD {
  1069. return withWindowResetsMetadata(ErrUserPlatformMonthlyQuotaExhausted, nextMonthlyResetFrom(entry.MonthlyWindowStart, now))
  1070. }
  1071. return nil
  1072. }
  1073. // --- cache MISS、旧版 entry 或 Redis 故障 → 查 DB(singleflight 合并并发回源)---
  1074. // 使用 DoChan 而非 Do:avoid sharing the first caller's ctx among all dedupe followers.
  1075. // 若第一个 caller 的 ctx 被取消(客户端断连),后续 caller 不受影响,仍由各自 ctx 控制超时。
  1076. sfKey := strconv.FormatInt(userID, 10) + ":" + platform
  1077. ch := s.quotaLoadSF.DoChan(sfKey, func() (any, error) {
  1078. // 子查询用 detached context + 短超时,独立于任何 caller 的请求 ctx,
  1079. // 防止"第一个 caller ctx 取消"使所有 follower 一起 fail。
  1080. bgCtx, bgCancel := context.WithTimeout(context.Background(), 3*time.Second)
  1081. defer bgCancel()
  1082. return s.userPlatformQuotaRepo.GetByUserPlatform(bgCtx, userID, platform)
  1083. })
  1084. var (
  1085. v any
  1086. dbErr error
  1087. )
  1088. select {
  1089. case res := <-ch:
  1090. v, dbErr = res.Val, res.Err
  1091. case <-ctx.Done():
  1092. // 当前 caller 的 ctx 被取消:fail-open,不阻断 (此请求已无意义)。
  1093. logger.LegacyPrintf("service.billing_cache", "Warning: user platform quota check ctx cancelled user=%d platform=%s: %v (fail-open)", userID, platform, ctx.Err())
  1094. return nil
  1095. }
  1096. if dbErr != nil {
  1097. logger.LegacyPrintf("service.billing_cache", "Warning: load user platform quota failed user=%d platform=%s: %v (fail-open)", userID, platform, dbErr)
  1098. return nil
  1099. }
  1100. rec, _ := v.(*UserPlatformQuotaRecord)
  1101. if rec == nil {
  1102. // 仅在 cache 可用且本次 GET 未出错时回填 sentinel:Redis GET 故障(cacheErr!=nil)
  1103. // 时不回填,与下方 line ~1201 "Redis 故障时 fail-open:不回填" 保持一致,
  1104. // 避免在 Redis 异常期做一次注定失败的 SET。
  1105. if s.cache != nil && cacheErr == nil {
  1106. now := time.Now()
  1107. startOfDay := timezone.StartOfDay(now)
  1108. startOfWeek := timezone.StartOfWeek(now)
  1109. sentinel := &UserPlatformQuotaCacheEntry{
  1110. SchemaVersion: UserPlatformQuotaCacheSchemaV1,
  1111. DailyWindowStart: &startOfDay,
  1112. WeeklyWindowStart: &startOfWeek,
  1113. MonthlyWindowStart: &now,
  1114. // limits 全 nil, usage 全 0(零值)
  1115. }
  1116. sentinelTTL := time.Duration(s.cfg.Billing.UserPlatformQuotaSentinelTTLSeconds) * time.Second
  1117. if sentinelTTL <= 0 {
  1118. // 防御:TTL<=0 时 Redis EXPIRE 会立即删除整个 key(见 billing_cache.go 的 pipe.Expire),
  1119. // sentinel 不持久化 → 每请求击穿 DB。配置缺失/误配为 0 时 fallback 到 1h。
  1120. sentinelTTL = time.Hour
  1121. }
  1122. setCtx, setCancel := context.WithTimeout(context.Background(), 50*time.Millisecond)
  1123. if setErr := s.cache.SetUserPlatformQuotaCache(setCtx, userID, platform, sentinel, sentinelTTL); setErr != nil {
  1124. userPlatformQuotaSentinelSetCacheErrorTotal.Add(1)
  1125. logger.LegacyPrintf("service.billing_cache", "Warning: set sentinel quota cache failed user=%d platform=%s: %v", userID, platform, setErr)
  1126. }
  1127. setCancel()
  1128. }
  1129. return nil
  1130. }
  1131. now := time.Now()
  1132. dailyUsage := rec.DailyUsageUSD
  1133. weeklyUsage := rec.WeeklyUsageUSD
  1134. monthlyUsage := rec.MonthlyUsageUSD
  1135. if quotaWindowExpired(rec.DailyWindowStart, timezone.StartOfDay(now)) {
  1136. dailyUsage = 0
  1137. }
  1138. if quotaWindowExpired(rec.WeeklyWindowStart, timezone.StartOfWeek(now)) {
  1139. weeklyUsage = 0
  1140. }
  1141. if monthlyQuotaWindowExpired(rec.MonthlyWindowStart, now) {
  1142. monthlyUsage = 0
  1143. }
  1144. // Redis 故障时 fail-open:不回填,直接用 DB 数据做一次性检查
  1145. if cacheErr != nil {
  1146. if rec.DailyLimitUSD != nil && dailyUsage >= *rec.DailyLimitUSD {
  1147. return withWindowResetsMetadata(ErrUserPlatformDailyQuotaExhausted, nextDailyReset(now))
  1148. }
  1149. if rec.WeeklyLimitUSD != nil && weeklyUsage >= *rec.WeeklyLimitUSD {
  1150. return withWindowResetsMetadata(ErrUserPlatformWeeklyQuotaExhausted, nextWeeklyReset(now))
  1151. }
  1152. if rec.MonthlyLimitUSD != nil && monthlyUsage >= *rec.MonthlyLimitUSD {
  1153. return withWindowResetsMetadata(ErrUserPlatformMonthlyQuotaExhausted, nextMonthlyResetFrom(rec.MonthlyWindowStart, now))
  1154. }
  1155. return nil
  1156. }
  1157. // cache MISS 或旧版 entry → 回填完整 entry(含 limits 和 window_start)
  1158. newEntry := &UserPlatformQuotaCacheEntry{
  1159. DailyUsageUSD: dailyUsage,
  1160. WeeklyUsageUSD: weeklyUsage,
  1161. MonthlyUsageUSD: monthlyUsage,
  1162. SchemaVersion: UserPlatformQuotaCacheSchemaV1,
  1163. DailyLimitUSD: rec.DailyLimitUSD,
  1164. WeeklyLimitUSD: rec.WeeklyLimitUSD,
  1165. MonthlyLimitUSD: rec.MonthlyLimitUSD,
  1166. DailyWindowStart: rec.DailyWindowStart,
  1167. WeeklyWindowStart: rec.WeeklyWindowStart,
  1168. MonthlyWindowStart: rec.MonthlyWindowStart,
  1169. }
  1170. if s.cache != nil {
  1171. ttl := time.Duration(s.cfg.Billing.UserPlatformQuotaCacheTTLSeconds) * time.Second
  1172. // 与 HIT 过期回填路径(上文 SetCache 调用)保持一致:用 context.Background()+50ms,
  1173. // 避免请求 ctx 提前取消(客户端断连/上游超时)导致 cache 回填失败,
  1174. // 让下一次 preflight 仍然 MISS 并击穿到 DB(高并发下增大 DB 压力)。
  1175. setCtx, setCancel := context.WithTimeout(context.Background(), 50*time.Millisecond)
  1176. if setErr := s.cache.SetUserPlatformQuotaCache(setCtx, userID, platform, newEntry, ttl); setErr != nil {
  1177. logger.LegacyPrintf("service.billing_cache", "Warning: set user platform quota cache failed user=%d platform=%s: %v", userID, platform, setErr)
  1178. }
  1179. setCancel()
  1180. }
  1181. if rec.DailyLimitUSD != nil && dailyUsage >= *rec.DailyLimitUSD {
  1182. return withWindowResetsMetadata(ErrUserPlatformDailyQuotaExhausted, nextDailyReset(now))
  1183. }
  1184. if rec.WeeklyLimitUSD != nil && weeklyUsage >= *rec.WeeklyLimitUSD {
  1185. return withWindowResetsMetadata(ErrUserPlatformWeeklyQuotaExhausted, nextWeeklyReset(now))
  1186. }
  1187. if rec.MonthlyLimitUSD != nil && monthlyUsage >= *rec.MonthlyLimitUSD {
  1188. return withWindowResetsMetadata(ErrUserPlatformMonthlyQuotaExhausted, nextMonthlyResetFrom(rec.MonthlyWindowStart, now))
  1189. }
  1190. return nil
  1191. }
  1192. // withWindowResetsMetadata 给 quota error 附加 window_resets_at metadata(RFC3339)。
  1193. func withWindowResetsMetadata(err error, resetAt time.Time) error {
  1194. appErr, ok := err.(*infraerrors.ApplicationError)
  1195. if !ok || appErr == nil {
  1196. return err
  1197. }
  1198. return appErr.WithMetadata(map[string]string{
  1199. "window_resets_at": resetAt.Format(time.RFC3339),
  1200. })
  1201. }
  1202. // nextDailyReset 计算下一个日窗口起点(次日全局时区 0 点)。
  1203. // 必须与 timezone.StartOfDay 同口径,否则 Retry-After 会偏差。
  1204. func nextDailyReset(now time.Time) time.Time {
  1205. return timezone.StartOfDay(now).AddDate(0, 0, 1)
  1206. }
  1207. // nextWeeklyReset 计算下一个周窗口起点(下周一全局时区 0 点)。
  1208. // 必须与 timezone.StartOfWeek 同口径,否则 Retry-After 会偏差。
  1209. func nextWeeklyReset(now time.Time) time.Time {
  1210. return timezone.StartOfWeek(now).AddDate(0, 0, 7)
  1211. }
  1212. // nextMonthlyResetFrom 返回 30 天滚动窗口的下次重置时间(start + 30d)。
  1213. // start 为 nil(未初始化)或已过期(now-start >= 30d,与 monthlyQuotaWindowExpired 同口径)时
  1214. // 退化为 now+30d:过期窗口会在下次 increment 时重置为 now,下次重置即 now+30d;
  1215. // 否则按 start 计算会得到一个过去的时间,使 Retry-After 落回 fallback 并触发客户端紧凑重试。
  1216. func nextMonthlyResetFrom(start *time.Time, now time.Time) time.Time {
  1217. if start == nil || now.Sub(*start) >= 30*24*time.Hour {
  1218. return now.Add(30 * 24 * time.Hour)
  1219. }
  1220. return start.Add(30 * 24 * time.Hour)
  1221. }
  1222. // quotaWindowExpired 判断窗口是否已过期:start 为 nil(未初始化)或在 currWindowStart 之前视为已过期。
  1223. func quotaWindowExpired(start *time.Time, currWindowStart time.Time) bool {
  1224. if start == nil {
  1225. return true
  1226. }
  1227. return start.Before(currWindowStart)
  1228. }
  1229. // monthlyQuotaWindowExpired 判断 30 天滚动月度窗口是否已过期。
  1230. // 过期条件:now - start >= 30×24h(与订阅模式 NeedsMonthlyReset 语义一致)。
  1231. // start 为 nil 时视为已过期(未初始化窗口)。
  1232. func monthlyQuotaWindowExpired(start *time.Time, now time.Time) bool {
  1233. if start == nil {
  1234. return true
  1235. }
  1236. return now.Sub(*start) >= 30*24*time.Hour
  1237. }
  1238. // HasUserPlatformQuotaLimit 判断该 user×platform 是否设了任一非 nil limit。
  1239. // 写入点守卫:无 limit 直接跳过 Redis 写 + 脏集标记,消除无谓写入。
  1240. // fail-safe:任何不确定(simple 模式除外)都返回 true 维持写入。
  1241. func (s *BillingCacheService) HasUserPlatformQuotaLimit(ctx context.Context, userID int64, platform string) bool {
  1242. if s.cfg.RunMode == config.RunModeSimple {
  1243. return false
  1244. }
  1245. if s.cache == nil {
  1246. return true
  1247. }
  1248. entry, ok, err := s.cache.GetUserPlatformQuotaCache(ctx, userID, platform)
  1249. if err != nil || !ok || entry == nil {
  1250. return true
  1251. }
  1252. return entry.DailyLimitUSD != nil || entry.WeeklyLimitUSD != nil || entry.MonthlyLimitUSD != nil
  1253. }