user_rpm_cache.go 3.1 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108
  1. package repository
  2. import (
  3. "context"
  4. "fmt"
  5. "time"
  6. "github.com/Wei-Shaw/sub2api/internal/service"
  7. "github.com/redis/go-redis/v9"
  8. )
  9. // 用户/分组级 RPM 计数器 Redis 实现。
  10. //
  11. // 设计说明:
  12. // - key 形式:rpm:ug:{uid}:{gid}:{minute}、rpm:u:{uid}:{minute}
  13. // - 时间来源:rdb.Time()(Redis 服务端时间),避免多实例时钟漂移。
  14. // - 原子操作:TxPipeline (MULTI/EXEC) 执行 INCR+EXPIRE,兼容 Redis Cluster。
  15. // - TTL:120s,覆盖当前分钟窗口 + 少量冗余。
  16. // - 返回值语义:超限判断由调用方(billing_cache_service.checkRPM)与 RPMLimit 比较完成。
  17. const (
  18. userGroupRPMKeyPrefix = "rpm:ug:"
  19. userRPMKeyPrefix = "rpm:u:"
  20. userRPMKeyTTL = 120 * time.Second
  21. )
  22. type userRPMCacheImpl struct {
  23. rdb *redis.Client
  24. }
  25. // NewUserRPMCache 创建用户/分组级 RPM 计数器。
  26. func NewUserRPMCache(rdb *redis.Client) service.UserRPMCache {
  27. return &userRPMCacheImpl{rdb: rdb}
  28. }
  29. // minuteTS 获取当前 Redis 服务端分钟时间戳。
  30. func (c *userRPMCacheImpl) minuteTS(ctx context.Context) (int64, error) {
  31. t, err := c.rdb.Time(ctx).Result()
  32. if err != nil {
  33. return 0, fmt.Errorf("redis TIME: %w", err)
  34. }
  35. return t.Unix() / 60, nil
  36. }
  37. // atomicIncr 原子 INCR+EXPIRE。
  38. func (c *userRPMCacheImpl) atomicIncr(ctx context.Context, key string) (int, error) {
  39. pipe := c.rdb.TxPipeline()
  40. incr := pipe.Incr(ctx, key)
  41. pipe.Expire(ctx, key, userRPMKeyTTL)
  42. if _, err := pipe.Exec(ctx); err != nil {
  43. return 0, fmt.Errorf("user rpm increment: %w", err)
  44. }
  45. return int(incr.Val()), nil
  46. }
  47. // IncrementUserGroupRPM 递增 (user, group) 分钟计数。
  48. func (c *userRPMCacheImpl) IncrementUserGroupRPM(ctx context.Context, userID, groupID int64) (int, error) {
  49. minute, err := c.minuteTS(ctx)
  50. if err != nil {
  51. return 0, err
  52. }
  53. key := fmt.Sprintf("%s%d:%d:%d", userGroupRPMKeyPrefix, userID, groupID, minute)
  54. return c.atomicIncr(ctx, key)
  55. }
  56. // IncrementUserRPM 递增用户分钟计数。
  57. func (c *userRPMCacheImpl) IncrementUserRPM(ctx context.Context, userID int64) (int, error) {
  58. minute, err := c.minuteTS(ctx)
  59. if err != nil {
  60. return 0, err
  61. }
  62. key := fmt.Sprintf("%s%d:%d", userRPMKeyPrefix, userID, minute)
  63. return c.atomicIncr(ctx, key)
  64. }
  65. // GetUserGroupRPM 获取 (user, group) 当前分钟已用 RPM(只读)。
  66. func (c *userRPMCacheImpl) GetUserGroupRPM(ctx context.Context, userID, groupID int64) (int, error) {
  67. minute, err := c.minuteTS(ctx)
  68. if err != nil {
  69. return 0, err
  70. }
  71. key := fmt.Sprintf("%s%d:%d:%d", userGroupRPMKeyPrefix, userID, groupID, minute)
  72. val, err := c.rdb.Get(ctx, key).Int()
  73. if err == redis.Nil {
  74. return 0, nil
  75. }
  76. if err != nil {
  77. return 0, fmt.Errorf("user group rpm get: %w", err)
  78. }
  79. return val, nil
  80. }
  81. // GetUserRPM 获取用户当前分钟已用 RPM(只读)。
  82. func (c *userRPMCacheImpl) GetUserRPM(ctx context.Context, userID int64) (int, error) {
  83. minute, err := c.minuteTS(ctx)
  84. if err != nil {
  85. return 0, err
  86. }
  87. key := fmt.Sprintf("%s%d:%d", userRPMKeyPrefix, userID, minute)
  88. val, err := c.rdb.Get(ctx, key).Int()
  89. if err == redis.Nil {
  90. return 0, nil
  91. }
  92. if err != nil {
  93. return 0, fmt.Errorf("user rpm get: %w", err)
  94. }
  95. return val, nil
  96. }