billing_cache.go 21 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643
  1. package repository
  2. import (
  3. "context"
  4. "errors"
  5. "fmt"
  6. "log"
  7. "math/rand/v2"
  8. "strconv"
  9. "strings"
  10. "time"
  11. "github.com/Wei-Shaw/sub2api/internal/service"
  12. "github.com/redis/go-redis/v9"
  13. )
  14. const (
  15. billingBalanceKeyPrefix = "billing:balance:"
  16. billingSubKeyPrefix = "billing:sub:"
  17. billingRateLimitKeyPrefix = "apikey:rate:"
  18. subCacheInvalidateChannel = "subscription:cache:invalidate"
  19. billingCacheTTL = 5 * time.Minute
  20. billingCacheJitter = 30 * time.Second
  21. rateLimitCacheTTL = 7 * 24 * time.Hour // 7 days matches the longest window
  22. // Rate limit window durations — must match service.RateLimitWindow* constants.
  23. rateLimitWindow5h = 5 * time.Hour
  24. rateLimitWindow1d = 24 * time.Hour
  25. rateLimitWindow7d = 7 * 24 * time.Hour
  26. )
  27. // jitteredTTL 返回带随机抖动的 TTL,防止缓存雪崩
  28. func jitteredTTL() time.Duration {
  29. // 只做“减法抖动”,确保实际 TTL 不会超过 billingCacheTTL(避免上界预期被打破)。
  30. if billingCacheJitter <= 0 {
  31. return billingCacheTTL
  32. }
  33. jitter := time.Duration(rand.IntN(int(billingCacheJitter)))
  34. return billingCacheTTL - jitter
  35. }
  36. // billingBalanceKey generates the Redis key for user balance cache.
  37. func billingBalanceKey(userID int64) string {
  38. return fmt.Sprintf("%s%d", billingBalanceKeyPrefix, userID)
  39. }
  40. // billingSubKey generates the Redis key for subscription cache.
  41. func billingSubKey(userID, groupID int64) string {
  42. return fmt.Sprintf("%s%d:%d", billingSubKeyPrefix, userID, groupID)
  43. }
  44. const (
  45. subFieldStatus = "status"
  46. subFieldExpiresAt = "expires_at"
  47. subFieldDailyUsage = "daily_usage"
  48. subFieldWeeklyUsage = "weekly_usage"
  49. subFieldMonthlyUsage = "monthly_usage"
  50. subFieldVersion = "version"
  51. )
  52. // billingRateLimitKey generates the Redis key for API key rate limit cache.
  53. func billingRateLimitKey(keyID int64) string {
  54. return fmt.Sprintf("%s%d", billingRateLimitKeyPrefix, keyID)
  55. }
  56. const (
  57. rateLimitFieldUsage5h = "usage_5h"
  58. rateLimitFieldUsage1d = "usage_1d"
  59. rateLimitFieldUsage7d = "usage_7d"
  60. rateLimitFieldWindow5h = "window_5h"
  61. rateLimitFieldWindow1d = "window_1d"
  62. rateLimitFieldWindow7d = "window_7d"
  63. )
  64. var (
  65. deductBalanceScript = redis.NewScript(`
  66. local current = redis.call('GET', KEYS[1])
  67. if current == false then
  68. return 0
  69. end
  70. local newVal = tonumber(current) - tonumber(ARGV[1])
  71. redis.call('SET', KEYS[1], newVal)
  72. redis.call('EXPIRE', KEYS[1], ARGV[2])
  73. return 1
  74. `)
  75. updateSubUsageScript = redis.NewScript(`
  76. local exists = redis.call('EXISTS', KEYS[1])
  77. if exists == 0 then
  78. return 0
  79. end
  80. local cost = tonumber(ARGV[1])
  81. redis.call('HINCRBYFLOAT', KEYS[1], 'daily_usage', cost)
  82. redis.call('HINCRBYFLOAT', KEYS[1], 'weekly_usage', cost)
  83. redis.call('HINCRBYFLOAT', KEYS[1], 'monthly_usage', cost)
  84. redis.call('EXPIRE', KEYS[1], ARGV[2])
  85. return 1
  86. `)
  87. // updateRateLimitUsageScript atomically increments all three rate limit usage counters
  88. // with window expiration checking. If a window has expired, its usage is reset to cost
  89. // (instead of accumulated) and the window timestamp is updated, matching the DB-side
  90. // IncrementRateLimitUsage semantics.
  91. //
  92. // ARGV: [1]=cost, [2]=ttl_seconds, [3]=now_unix, [4]=window_5h_seconds, [5]=window_1d_seconds, [6]=window_7d_seconds
  93. updateRateLimitUsageScript = redis.NewScript(`
  94. local exists = redis.call('EXISTS', KEYS[1])
  95. if exists == 0 then
  96. return 0
  97. end
  98. local cost = tonumber(ARGV[1])
  99. local now = tonumber(ARGV[3])
  100. local win5h = tonumber(ARGV[4])
  101. local win1d = tonumber(ARGV[5])
  102. local win7d = tonumber(ARGV[6])
  103. -- Helper: check if window is expired and update usage + window accordingly
  104. -- Returns nothing, modifies the hash in-place.
  105. local function update_window(usage_field, window_field, window_duration)
  106. local w = tonumber(redis.call('HGET', KEYS[1], window_field) or 0)
  107. if w == 0 or (now - w) >= window_duration then
  108. -- Window expired or never started: reset usage to cost, start new window
  109. redis.call('HSET', KEYS[1], usage_field, tostring(cost))
  110. redis.call('HSET', KEYS[1], window_field, tostring(now))
  111. else
  112. -- Window still valid: accumulate
  113. redis.call('HINCRBYFLOAT', KEYS[1], usage_field, cost)
  114. end
  115. end
  116. update_window('usage_5h', 'window_5h', win5h)
  117. update_window('usage_1d', 'window_1d', win1d)
  118. update_window('usage_7d', 'window_7d', win7d)
  119. redis.call('EXPIRE', KEYS[1], ARGV[2])
  120. return 1
  121. `)
  122. )
  123. type billingCache struct {
  124. rdb *redis.Client
  125. }
  126. func NewBillingCache(rdb *redis.Client) service.BillingCache {
  127. return &billingCache{rdb: rdb}
  128. }
  129. func (c *billingCache) GetUserBalance(ctx context.Context, userID int64) (float64, error) {
  130. key := billingBalanceKey(userID)
  131. val, err := c.rdb.Get(ctx, key).Result()
  132. if err != nil {
  133. return 0, err
  134. }
  135. return strconv.ParseFloat(val, 64)
  136. }
  137. func (c *billingCache) SetUserBalance(ctx context.Context, userID int64, balance float64) error {
  138. key := billingBalanceKey(userID)
  139. return c.rdb.Set(ctx, key, balance, jitteredTTL()).Err()
  140. }
  141. func (c *billingCache) DeductUserBalance(ctx context.Context, userID int64, amount float64) error {
  142. key := billingBalanceKey(userID)
  143. _, err := deductBalanceScript.Run(ctx, c.rdb, []string{key}, amount, int(jitteredTTL().Seconds())).Result()
  144. if err != nil && !errors.Is(err, redis.Nil) {
  145. log.Printf("Warning: deduct balance cache failed for user %d: %v", userID, err)
  146. return err
  147. }
  148. return nil
  149. }
  150. func (c *billingCache) InvalidateUserBalance(ctx context.Context, userID int64) error {
  151. key := billingBalanceKey(userID)
  152. return c.rdb.Del(ctx, key).Err()
  153. }
  154. func (c *billingCache) GetSubscriptionCache(ctx context.Context, userID, groupID int64) (*service.SubscriptionCacheData, error) {
  155. key := billingSubKey(userID, groupID)
  156. result, err := c.rdb.HGetAll(ctx, key).Result()
  157. if err != nil {
  158. return nil, err
  159. }
  160. if len(result) == 0 {
  161. return nil, redis.Nil
  162. }
  163. return c.parseSubscriptionCache(result)
  164. }
  165. func (c *billingCache) parseSubscriptionCache(data map[string]string) (*service.SubscriptionCacheData, error) {
  166. result := &service.SubscriptionCacheData{}
  167. result.Status = data[subFieldStatus]
  168. if result.Status == "" {
  169. return nil, errors.New("invalid cache: missing status")
  170. }
  171. if expiresStr, ok := data[subFieldExpiresAt]; ok {
  172. expiresAt, err := strconv.ParseInt(expiresStr, 10, 64)
  173. if err == nil {
  174. result.ExpiresAt = time.Unix(expiresAt, 0)
  175. }
  176. }
  177. if dailyStr, ok := data[subFieldDailyUsage]; ok {
  178. result.DailyUsage, _ = strconv.ParseFloat(dailyStr, 64)
  179. }
  180. if weeklyStr, ok := data[subFieldWeeklyUsage]; ok {
  181. result.WeeklyUsage, _ = strconv.ParseFloat(weeklyStr, 64)
  182. }
  183. if monthlyStr, ok := data[subFieldMonthlyUsage]; ok {
  184. result.MonthlyUsage, _ = strconv.ParseFloat(monthlyStr, 64)
  185. }
  186. if versionStr, ok := data[subFieldVersion]; ok {
  187. result.Version, _ = strconv.ParseInt(versionStr, 10, 64)
  188. }
  189. return result, nil
  190. }
  191. func (c *billingCache) SetSubscriptionCache(ctx context.Context, userID, groupID int64, data *service.SubscriptionCacheData) error {
  192. if data == nil {
  193. return nil
  194. }
  195. key := billingSubKey(userID, groupID)
  196. fields := map[string]any{
  197. subFieldStatus: data.Status,
  198. subFieldExpiresAt: data.ExpiresAt.Unix(),
  199. subFieldDailyUsage: data.DailyUsage,
  200. subFieldWeeklyUsage: data.WeeklyUsage,
  201. subFieldMonthlyUsage: data.MonthlyUsage,
  202. subFieldVersion: data.Version,
  203. }
  204. pipe := c.rdb.Pipeline()
  205. pipe.HSet(ctx, key, fields)
  206. pipe.Expire(ctx, key, jitteredTTL())
  207. _, err := pipe.Exec(ctx)
  208. return err
  209. }
  210. func (c *billingCache) UpdateSubscriptionUsage(ctx context.Context, userID, groupID int64, cost float64) error {
  211. key := billingSubKey(userID, groupID)
  212. _, err := updateSubUsageScript.Run(ctx, c.rdb, []string{key}, cost, int(jitteredTTL().Seconds())).Result()
  213. if err != nil && !errors.Is(err, redis.Nil) {
  214. log.Printf("Warning: update subscription usage cache failed for user %d group %d: %v", userID, groupID, err)
  215. return err
  216. }
  217. return nil
  218. }
  219. func (c *billingCache) InvalidateSubscriptionCache(ctx context.Context, userID, groupID int64) error {
  220. key := billingSubKey(userID, groupID)
  221. return c.rdb.Del(ctx, key).Err()
  222. }
  223. func (c *billingCache) PublishSubscriptionCacheInvalidation(ctx context.Context, cacheKey string) error {
  224. return c.rdb.Publish(ctx, subCacheInvalidateChannel, cacheKey).Err()
  225. }
  226. func (c *billingCache) SubscribeSubscriptionCacheInvalidation(ctx context.Context, handler func(cacheKey string)) error {
  227. pubsub := c.rdb.Subscribe(ctx, subCacheInvalidateChannel)
  228. _, err := pubsub.Receive(ctx)
  229. if err != nil {
  230. _ = pubsub.Close()
  231. return fmt.Errorf("subscribe to subscription cache invalidation: %w", err)
  232. }
  233. go func() {
  234. defer func() {
  235. if err := pubsub.Close(); err != nil {
  236. log.Printf("Warning: failed to close subscription cache invalidation pubsub: %v", err)
  237. }
  238. }()
  239. ch := pubsub.Channel()
  240. for {
  241. select {
  242. case <-ctx.Done():
  243. return
  244. case msg, ok := <-ch:
  245. if !ok {
  246. return
  247. }
  248. if msg != nil {
  249. handler(msg.Payload)
  250. }
  251. }
  252. }
  253. }()
  254. return nil
  255. }
  256. func (c *billingCache) GetAPIKeyRateLimit(ctx context.Context, keyID int64) (*service.APIKeyRateLimitCacheData, error) {
  257. key := billingRateLimitKey(keyID)
  258. result, err := c.rdb.HGetAll(ctx, key).Result()
  259. if err != nil {
  260. return nil, err
  261. }
  262. if len(result) == 0 {
  263. return nil, redis.Nil
  264. }
  265. data := &service.APIKeyRateLimitCacheData{}
  266. if v, ok := result[rateLimitFieldUsage5h]; ok {
  267. data.Usage5h, _ = strconv.ParseFloat(v, 64)
  268. }
  269. if v, ok := result[rateLimitFieldUsage1d]; ok {
  270. data.Usage1d, _ = strconv.ParseFloat(v, 64)
  271. }
  272. if v, ok := result[rateLimitFieldUsage7d]; ok {
  273. data.Usage7d, _ = strconv.ParseFloat(v, 64)
  274. }
  275. if v, ok := result[rateLimitFieldWindow5h]; ok {
  276. data.Window5h, _ = strconv.ParseInt(v, 10, 64)
  277. }
  278. if v, ok := result[rateLimitFieldWindow1d]; ok {
  279. data.Window1d, _ = strconv.ParseInt(v, 10, 64)
  280. }
  281. if v, ok := result[rateLimitFieldWindow7d]; ok {
  282. data.Window7d, _ = strconv.ParseInt(v, 10, 64)
  283. }
  284. return data, nil
  285. }
  286. func (c *billingCache) SetAPIKeyRateLimit(ctx context.Context, keyID int64, data *service.APIKeyRateLimitCacheData) error {
  287. if data == nil {
  288. return nil
  289. }
  290. key := billingRateLimitKey(keyID)
  291. fields := map[string]any{
  292. rateLimitFieldUsage5h: data.Usage5h,
  293. rateLimitFieldUsage1d: data.Usage1d,
  294. rateLimitFieldUsage7d: data.Usage7d,
  295. rateLimitFieldWindow5h: data.Window5h,
  296. rateLimitFieldWindow1d: data.Window1d,
  297. rateLimitFieldWindow7d: data.Window7d,
  298. }
  299. pipe := c.rdb.Pipeline()
  300. pipe.HSet(ctx, key, fields)
  301. pipe.Expire(ctx, key, rateLimitCacheTTL)
  302. _, err := pipe.Exec(ctx)
  303. return err
  304. }
  305. func (c *billingCache) UpdateAPIKeyRateLimitUsage(ctx context.Context, keyID int64, cost float64) error {
  306. key := billingRateLimitKey(keyID)
  307. now := time.Now().Unix()
  308. _, err := updateRateLimitUsageScript.Run(ctx, c.rdb, []string{key},
  309. cost,
  310. int(rateLimitCacheTTL.Seconds()),
  311. now,
  312. int(rateLimitWindow5h.Seconds()),
  313. int(rateLimitWindow1d.Seconds()),
  314. int(rateLimitWindow7d.Seconds()),
  315. ).Result()
  316. if err != nil && !errors.Is(err, redis.Nil) {
  317. log.Printf("Warning: update rate limit usage cache failed for api key %d: %v", keyID, err)
  318. return err
  319. }
  320. return nil
  321. }
  322. func (c *billingCache) InvalidateAPIKeyRateLimit(ctx context.Context, keyID int64) error {
  323. key := billingRateLimitKey(keyID)
  324. return c.rdb.Del(ctx, key).Err()
  325. }
  326. // ============================================
  327. // user × platform quota 缓存
  328. // ============================================
  329. // userPlatformQuotaCacheKey 构造 Redis key
  330. func userPlatformQuotaCacheKey(userID int64, platform string) string {
  331. return fmt.Sprintf("billing:user_platform_quota:%d:%s", userID, platform)
  332. }
  333. // parseUserPlatformQuotaHash 将 Redis HGETALL 返回的 map[string]string 反序列化为
  334. // *service.UserPlatformQuotaCacheEntry。空 map(key 不存在)返回 nil。
  335. // GetUserPlatformQuotaCache 和 BatchGetUserPlatformQuotaCache 共用此函数,确保解析逻辑一致。
  336. func parseUserPlatformQuotaHash(m map[string]string) *service.UserPlatformQuotaCacheEntry {
  337. if len(m) == 0 {
  338. return nil
  339. }
  340. parseFloat := func(s string) float64 {
  341. if s == "" {
  342. return 0
  343. }
  344. f, err := strconv.ParseFloat(s, 64)
  345. if err != nil {
  346. log.Printf("billing_cache: corrupt quota usage field %q (using 0): %v", s, err)
  347. return 0
  348. }
  349. return f
  350. }
  351. parseFloatPtr := func(s string) *float64 {
  352. if s == "" {
  353. return nil
  354. }
  355. f, err := strconv.ParseFloat(s, 64)
  356. if err != nil {
  357. return nil
  358. }
  359. return &f
  360. }
  361. parseTimePtr := func(s string) *time.Time {
  362. if s == "" {
  363. return nil
  364. }
  365. n, err := strconv.ParseInt(s, 10, 64)
  366. if err != nil {
  367. return nil
  368. }
  369. t := time.Unix(n, 0).UTC()
  370. return &t
  371. }
  372. parseInt64 := func(s string) int64 {
  373. n, _ := strconv.ParseInt(s, 10, 64)
  374. return n
  375. }
  376. return &service.UserPlatformQuotaCacheEntry{
  377. DailyUsageUSD: parseFloat(m["daily_usage"]),
  378. WeeklyUsageUSD: parseFloat(m["weekly_usage"]),
  379. MonthlyUsageUSD: parseFloat(m["monthly_usage"]),
  380. Version: parseInt64(m["version"]),
  381. SchemaVersion: parseInt64(m["schema_version"]),
  382. DailyLimitUSD: parseFloatPtr(m["daily_limit"]),
  383. WeeklyLimitUSD: parseFloatPtr(m["weekly_limit"]),
  384. MonthlyLimitUSD: parseFloatPtr(m["monthly_limit"]),
  385. DailyWindowStart: parseTimePtr(m["daily_window_start"]),
  386. WeeklyWindowStart: parseTimePtr(m["weekly_window_start"]),
  387. MonthlyWindowStart: parseTimePtr(m["monthly_window_start"]),
  388. }
  389. }
  390. func (c *billingCache) GetUserPlatformQuotaCache(ctx context.Context, userID int64, platform string) (*service.UserPlatformQuotaCacheEntry, bool, error) {
  391. key := userPlatformQuotaCacheKey(userID, platform)
  392. m, err := c.rdb.HGetAll(ctx, key).Result()
  393. if err != nil {
  394. return nil, false, err
  395. }
  396. entry := parseUserPlatformQuotaHash(m)
  397. if entry == nil {
  398. // 空 map → key 不存在 → MISS
  399. return nil, false, nil
  400. }
  401. return entry, true, nil
  402. }
  403. func (c *billingCache) SetUserPlatformQuotaCache(ctx context.Context, userID int64, platform string, entry *service.UserPlatformQuotaCacheEntry, ttl time.Duration) error {
  404. if entry == nil {
  405. return nil
  406. }
  407. key := userPlatformQuotaCacheKey(userID, platform)
  408. pipe := c.rdb.TxPipeline()
  409. // 浮点可空字段:nil → 空字符串(读取时 parseFloatPtr 返回 nil,表示无限额)
  410. fmtFloatPtr := func(p *float64) string {
  411. if p == nil {
  412. return ""
  413. }
  414. return strconv.FormatFloat(*p, 'f', -1, 64)
  415. }
  416. // time.Time 可空字段:nil → 空字符串;有值 → unix 秒
  417. fmtTimePtr := func(p *time.Time) string {
  418. if p == nil {
  419. return ""
  420. }
  421. return strconv.FormatInt(p.Unix(), 10)
  422. }
  423. pipe.HSet(ctx, key,
  424. "daily_usage", entry.DailyUsageUSD,
  425. "weekly_usage", entry.WeeklyUsageUSD,
  426. "monthly_usage", entry.MonthlyUsageUSD,
  427. "version", entry.Version,
  428. "schema_version", entry.SchemaVersion,
  429. "daily_limit", fmtFloatPtr(entry.DailyLimitUSD),
  430. "weekly_limit", fmtFloatPtr(entry.WeeklyLimitUSD),
  431. "monthly_limit", fmtFloatPtr(entry.MonthlyLimitUSD),
  432. "daily_window_start", fmtTimePtr(entry.DailyWindowStart),
  433. "weekly_window_start", fmtTimePtr(entry.WeeklyWindowStart),
  434. "monthly_window_start", fmtTimePtr(entry.MonthlyWindowStart),
  435. )
  436. pipe.Expire(ctx, key, ttl)
  437. _, err := pipe.Exec(ctx)
  438. return err
  439. }
  440. func (c *billingCache) DeleteUserPlatformQuotaCache(ctx context.Context, userID int64, platform string) error {
  441. return c.rdb.Del(ctx, userPlatformQuotaCacheKey(userID, platform)).Err()
  442. }
  443. // updateUserPlatformQuotaUsageScript 缓存累加:EXISTS + schema_version 双重守卫。
  444. // 旧版 entry(schema_version != ARGV[3],包括缺字段的 0 值)不参与累加,由上层走 DB fallback 后
  445. // SetCache 重建为新版 entry —— 若此处仍累加,上层覆盖时会丢失这部分增量,导致 Redis usage 比真实偏小。
  446. // key 不存在同样跳过(由下次 SetCache 重建)。
  447. // KEYS[1] = hash key
  448. // KEYS[2] = 脏集 key(dirty set)
  449. // ARGV[1] = cost (string float)
  450. // ARGV[2] = ttl seconds
  451. // ARGV[3] = expected schema_version (Go 侧 UserPlatformQuotaCacheSchemaV1)
  452. // ARGV[4] = dirty set member(空串则不 SADD)
  453. // ARGV[5] = 脏集兜底 TTL 秒
  454. const updateUserPlatformQuotaUsageScript = `
  455. if redis.call("EXISTS", KEYS[1]) == 0 then
  456. return 0
  457. end
  458. local ver = redis.call("HGET", KEYS[1], "schema_version")
  459. if ver == false or tonumber(ver) ~= tonumber(ARGV[3]) then
  460. return 0
  461. end
  462. redis.call("HINCRBYFLOAT", KEYS[1], "daily_usage", ARGV[1])
  463. redis.call("HINCRBYFLOAT", KEYS[1], "weekly_usage", ARGV[1])
  464. redis.call("HINCRBYFLOAT", KEYS[1], "monthly_usage", ARGV[1])
  465. redis.call("HINCRBY", KEYS[1], "version", 1)
  466. redis.call("EXPIRE", KEYS[1], ARGV[2])
  467. if ARGV[4] ~= "" then
  468. redis.call("SADD", KEYS[2], ARGV[4])
  469. redis.call("EXPIRE", KEYS[2], ARGV[5])
  470. end
  471. return 1
  472. `
  473. // userPlatformQuotaDirtySetKey 返回脏集(dirty set)的 Redis key。
  474. // 使用与 userPlatformQuotaCacheKey 相同的前缀 "billing:"。
  475. func userPlatformQuotaDirtySetKey() string { return "billing:" + "upq:dirty" }
  476. // userPlatformQuotaDirtyTTLSeconds 脏集兜底 TTL(秒):初始 SADD(Lua)与 Readd 共用,
  477. // 确保 flusher 长期停摆时脏集最终过期;正常运行因持续 SADD 不断续期。
  478. const userPlatformQuotaDirtyTTLSeconds = 86400
  479. // userPlatformQuotaDirtyMember 构造脏集成员字符串 "userID:platform"。
  480. func userPlatformQuotaDirtyMember(userID int64, platform string) string {
  481. return strconv.FormatInt(userID, 10) + ":" + platform
  482. }
  483. func (c *billingCache) IncrUserPlatformQuotaUsageCache(ctx context.Context, userID int64, platform string, cost float64, ttl time.Duration, markDirty bool) error {
  484. member := ""
  485. if markDirty {
  486. member = userPlatformQuotaDirtyMember(userID, platform)
  487. }
  488. _, err := c.rdb.Eval(ctx, updateUserPlatformQuotaUsageScript,
  489. []string{userPlatformQuotaCacheKey(userID, platform), userPlatformQuotaDirtySetKey()},
  490. strconv.FormatFloat(cost, 'f', -1, 64),
  491. int(ttl.Seconds()),
  492. service.UserPlatformQuotaCacheSchemaV1,
  493. member,
  494. userPlatformQuotaDirtyTTLSeconds,
  495. ).Result()
  496. if err != nil && !errors.Is(err, redis.Nil) {
  497. return err
  498. }
  499. return nil
  500. }
  501. // parseUserPlatformQuotaDirtyMember 将脏集成员字符串 "userID:platform" 解析为
  502. // service.UserPlatformQuotaKey。解析失败返回 ok=false。
  503. func parseUserPlatformQuotaDirtyMember(m string) (service.UserPlatformQuotaKey, bool) {
  504. parts := strings.SplitN(m, ":", 2)
  505. if len(parts) != 2 {
  506. return service.UserPlatformQuotaKey{}, false
  507. }
  508. uid, err := strconv.ParseInt(parts[0], 10, 64)
  509. if err != nil {
  510. return service.UserPlatformQuotaKey{}, false
  511. }
  512. return service.UserPlatformQuotaKey{UserID: uid, Platform: parts[1]}, true
  513. }
  514. // PopDirtyUserPlatformQuotaKeys 从脏集随机弹出最多 n 个 key。
  515. // 脏集为空时返回 (nil, nil)。
  516. func (c *billingCache) PopDirtyUserPlatformQuotaKeys(ctx context.Context, n int) ([]service.UserPlatformQuotaKey, error) {
  517. members, err := c.rdb.SPopN(ctx, userPlatformQuotaDirtySetKey(), int64(n)).Result()
  518. if err != nil {
  519. if errors.Is(err, redis.Nil) {
  520. return nil, nil
  521. }
  522. return nil, err
  523. }
  524. keys := make([]service.UserPlatformQuotaKey, 0, len(members))
  525. for _, m := range members {
  526. k, ok := parseUserPlatformQuotaDirtyMember(m)
  527. if !ok {
  528. log.Printf("billing_cache: skipping invalid dirty member %q", m)
  529. continue
  530. }
  531. keys = append(keys, k)
  532. }
  533. return keys, nil
  534. }
  535. // ReaddDirtyUserPlatformQuotaKeys 将 keys 重新加入脏集(flush 失败时回填)。
  536. // 通过 pipeline 同时执行 SAdd + Expire,确保 Readd 后脏集具有兜底 TTL。
  537. // 空切片时直接返回 nil。
  538. func (c *billingCache) ReaddDirtyUserPlatformQuotaKeys(ctx context.Context, keys []service.UserPlatformQuotaKey) error {
  539. if len(keys) == 0 {
  540. return nil
  541. }
  542. dirtyKey := userPlatformQuotaDirtySetKey()
  543. members := make([]any, len(keys))
  544. for i, k := range keys {
  545. members[i] = userPlatformQuotaDirtyMember(k.UserID, k.Platform)
  546. }
  547. pipe := c.rdb.Pipeline()
  548. pipe.SAdd(ctx, dirtyKey, members...)
  549. pipe.Expire(ctx, dirtyKey, userPlatformQuotaDirtyTTLSeconds*time.Second)
  550. _, err := pipe.Exec(ctx)
  551. return err
  552. }
  553. // BatchGetUserPlatformQuotaCache 通过 Pipeline 批量 HGETALL 获取多个 user×platform 的
  554. // quota cache。返回切片与 keys 顺序、长度对齐;MISS 或解析失败位置返回 nil。
  555. func (c *billingCache) BatchGetUserPlatformQuotaCache(ctx context.Context, keys []service.UserPlatformQuotaKey) ([]*service.UserPlatformQuotaCacheEntry, error) {
  556. if len(keys) == 0 {
  557. return nil, nil
  558. }
  559. pipe := c.rdb.Pipeline()
  560. cmds := make([]*redis.MapStringStringCmd, len(keys))
  561. for i, k := range keys {
  562. cmds[i] = pipe.HGetAll(ctx, userPlatformQuotaCacheKey(k.UserID, k.Platform))
  563. }
  564. if _, err := pipe.Exec(ctx); err != nil && !errors.Is(err, redis.Nil) {
  565. return nil, err
  566. }
  567. results := make([]*service.UserPlatformQuotaCacheEntry, len(keys))
  568. for i, cmd := range cmds {
  569. m, err := cmd.Result()
  570. if err != nil {
  571. if !errors.Is(err, redis.Nil) {
  572. log.Printf("billing_cache: BatchGet HGETALL cmd[%d] failed: %v (skip, self-heal)", i, err)
  573. }
  574. // 单个命令失败 → 对应位置 nil,继续
  575. continue
  576. }
  577. results[i] = parseUserPlatformQuotaHash(m)
  578. }
  579. return results, nil
  580. }