rpm_cache.go 4.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141
  1. package repository
  2. import (
  3. "context"
  4. "errors"
  5. "fmt"
  6. "strconv"
  7. "time"
  8. "github.com/Wei-Shaw/sub2api/internal/service"
  9. "github.com/redis/go-redis/v9"
  10. )
  11. // RPM 计数器缓存常量定义
  12. //
  13. // 设计说明:
  14. // 使用 Redis 简单计数器跟踪每个账号每分钟的请求数:
  15. // - Key: rpm:{accountID}:{minuteTimestamp}
  16. // - Value: 当前分钟内的请求计数
  17. // - TTL: 120 秒(覆盖当前分钟 + 一定冗余)
  18. //
  19. // 使用 TxPipeline(MULTI/EXEC)执行 INCR + EXPIRE,保证原子性且兼容 Redis Cluster。
  20. // 通过 rdb.Time() 获取服务端时间,避免多实例时钟不同步。
  21. //
  22. // 设计决策:
  23. // - TxPipeline vs Pipeline:Pipeline 仅合并发送但不保证原子,TxPipeline 使用 MULTI/EXEC 事务保证原子执行。
  24. // - rdb.Time() 单独调用:Pipeline/TxPipeline 中无法引用前一命令的结果,因此 TIME 必须单独调用(2 RTT)。
  25. // Lua 脚本可以做到 1 RTT,但在 Redis Cluster 中动态拼接 key 存在 CROSSSLOT 风险,选择安全性优先。
  26. const (
  27. // RPM 计数器键前缀
  28. // 格式: rpm:{accountID}:{minuteTimestamp}
  29. rpmKeyPrefix = "rpm:"
  30. // RPM 计数器 TTL(120 秒,覆盖当前分钟窗口 + 冗余)
  31. rpmKeyTTL = 120 * time.Second
  32. )
  33. // RPMCacheImpl RPM 计数器缓存 Redis 实现
  34. type RPMCacheImpl struct {
  35. rdb *redis.Client
  36. }
  37. // NewRPMCache 创建 RPM 计数器缓存
  38. func NewRPMCache(rdb *redis.Client) service.RPMCache {
  39. return &RPMCacheImpl{rdb: rdb}
  40. }
  41. // currentMinuteKey 获取当前分钟的完整 Redis key
  42. // 使用 rdb.Time() 获取 Redis 服务端时间,避免多实例时钟偏差
  43. func (c *RPMCacheImpl) currentMinuteKey(ctx context.Context, accountID int64) (string, error) {
  44. serverTime, err := c.rdb.Time(ctx).Result()
  45. if err != nil {
  46. return "", fmt.Errorf("redis TIME: %w", err)
  47. }
  48. minuteTS := serverTime.Unix() / 60
  49. return fmt.Sprintf("%s%d:%d", rpmKeyPrefix, accountID, minuteTS), nil
  50. }
  51. // currentMinuteSuffix 获取当前分钟时间戳后缀(供批量操作使用)
  52. // 使用 rdb.Time() 获取 Redis 服务端时间
  53. func (c *RPMCacheImpl) currentMinuteSuffix(ctx context.Context) (string, error) {
  54. serverTime, err := c.rdb.Time(ctx).Result()
  55. if err != nil {
  56. return "", fmt.Errorf("redis TIME: %w", err)
  57. }
  58. minuteTS := serverTime.Unix() / 60
  59. return strconv.FormatInt(minuteTS, 10), nil
  60. }
  61. // IncrementRPM 原子递增并返回当前分钟的计数
  62. // 使用 TxPipeline (MULTI/EXEC) 执行 INCR + EXPIRE,保证原子性且兼容 Redis Cluster
  63. func (c *RPMCacheImpl) IncrementRPM(ctx context.Context, accountID int64) (int, error) {
  64. key, err := c.currentMinuteKey(ctx, accountID)
  65. if err != nil {
  66. return 0, fmt.Errorf("rpm increment: %w", err)
  67. }
  68. // 使用 TxPipeline (MULTI/EXEC) 保证 INCR + EXPIRE 原子执行
  69. // EXPIRE 幂等,每次都设置不影响正确性
  70. pipe := c.rdb.TxPipeline()
  71. incrCmd := pipe.Incr(ctx, key)
  72. pipe.Expire(ctx, key, rpmKeyTTL)
  73. if _, err := pipe.Exec(ctx); err != nil {
  74. return 0, fmt.Errorf("rpm increment: %w", err)
  75. }
  76. return int(incrCmd.Val()), nil
  77. }
  78. // GetRPM 获取当前分钟的 RPM 计数
  79. func (c *RPMCacheImpl) GetRPM(ctx context.Context, accountID int64) (int, error) {
  80. key, err := c.currentMinuteKey(ctx, accountID)
  81. if err != nil {
  82. return 0, fmt.Errorf("rpm get: %w", err)
  83. }
  84. val, err := c.rdb.Get(ctx, key).Int()
  85. if errors.Is(err, redis.Nil) {
  86. return 0, nil // 当前分钟无记录
  87. }
  88. if err != nil {
  89. return 0, fmt.Errorf("rpm get: %w", err)
  90. }
  91. return val, nil
  92. }
  93. // GetRPMBatch 批量获取多个账号的 RPM 计数(使用 Pipeline)
  94. func (c *RPMCacheImpl) GetRPMBatch(ctx context.Context, accountIDs []int64) (map[int64]int, error) {
  95. if len(accountIDs) == 0 {
  96. return map[int64]int{}, nil
  97. }
  98. // 获取当前分钟后缀
  99. minuteSuffix, err := c.currentMinuteSuffix(ctx)
  100. if err != nil {
  101. return nil, fmt.Errorf("rpm batch get: %w", err)
  102. }
  103. // 使用 Pipeline 批量 GET
  104. pipe := c.rdb.Pipeline()
  105. cmds := make(map[int64]*redis.StringCmd, len(accountIDs))
  106. for _, id := range accountIDs {
  107. key := fmt.Sprintf("%s%d:%s", rpmKeyPrefix, id, minuteSuffix)
  108. cmds[id] = pipe.Get(ctx, key)
  109. }
  110. if _, err := pipe.Exec(ctx); err != nil && !errors.Is(err, redis.Nil) {
  111. return nil, fmt.Errorf("rpm batch get: %w", err)
  112. }
  113. result := make(map[int64]int, len(accountIDs))
  114. for id, cmd := range cmds {
  115. if val, err := cmd.Int(); err == nil {
  116. result[id] = val
  117. } else {
  118. result[id] = 0
  119. }
  120. }
  121. return result, nil
  122. }