internal500_counter_cache.go 1.6 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455
  1. package repository
  2. import (
  3. "context"
  4. "fmt"
  5. "github.com/Wei-Shaw/sub2api/internal/service"
  6. "github.com/redis/go-redis/v9"
  7. )
  8. const (
  9. internal500CounterPrefix = "internal500_count:account:"
  10. internal500CounterTTLSeconds = 86400 // 24 小时兜底
  11. )
  12. // internal500CounterIncrScript 使用 Lua 脚本原子性地增加计数并返回当前值
  13. // 如果 key 不存在,则创建并设置过期时间
  14. var internal500CounterIncrScript = redis.NewScript(`
  15. local key = KEYS[1]
  16. local ttl = tonumber(ARGV[1])
  17. local count = redis.call('INCR', key)
  18. if count == 1 then
  19. redis.call('EXPIRE', key, ttl)
  20. end
  21. return count
  22. `)
  23. type internal500CounterCache struct {
  24. rdb *redis.Client
  25. }
  26. // NewInternal500CounterCache 创建 INTERNAL 500 连续失败计数器缓存实例
  27. func NewInternal500CounterCache(rdb *redis.Client) service.Internal500CounterCache {
  28. return &internal500CounterCache{rdb: rdb}
  29. }
  30. // IncrementInternal500Count 原子递增计数并返回当前值
  31. func (c *internal500CounterCache) IncrementInternal500Count(ctx context.Context, accountID int64) (int64, error) {
  32. key := fmt.Sprintf("%s%d", internal500CounterPrefix, accountID)
  33. result, err := internal500CounterIncrScript.Run(ctx, c.rdb, []string{key}, internal500CounterTTLSeconds).Int64()
  34. if err != nil {
  35. return 0, fmt.Errorf("increment internal500 count: %w", err)
  36. }
  37. return result, nil
  38. }
  39. // ResetInternal500Count 清零计数器(成功响应时调用)
  40. func (c *internal500CounterCache) ResetInternal500Count(ctx context.Context, accountID int64) error {
  41. key := fmt.Sprintf("%s%d", internal500CounterPrefix, accountID)
  42. return c.rdb.Del(ctx, key).Err()
  43. }