rate_limit.lua 1.2 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344
  1. -- 令牌桶限流器
  2. -- KEYS[1]: 限流器唯一标识
  3. -- ARGV[1]: 请求令牌数 (通常为1)
  4. -- ARGV[2]: 令牌生成速率 (每秒)
  5. -- ARGV[3]: 桶容量
  6. local key = KEYS[1]
  7. local requested = tonumber(ARGV[1])
  8. local rate = tonumber(ARGV[2])
  9. local capacity = tonumber(ARGV[3])
  10. -- 获取当前时间(Redis服务器时间)
  11. local now = redis.call('TIME')
  12. local nowInSeconds = tonumber(now[1])
  13. -- 获取桶状态
  14. local bucket = redis.call('HMGET', key, 'tokens', 'last_time')
  15. local tokens = tonumber(bucket[1])
  16. local last_time = tonumber(bucket[2])
  17. -- 初始化桶(首次请求或过期)
  18. if not tokens or not last_time then
  19. tokens = capacity
  20. last_time = nowInSeconds
  21. else
  22. -- 计算新增令牌
  23. local elapsed = nowInSeconds - last_time
  24. local add_tokens = elapsed * rate
  25. tokens = math.min(capacity, tokens + add_tokens)
  26. last_time = nowInSeconds
  27. end
  28. -- 判断是否允许请求
  29. local allowed = false
  30. if tokens >= requested then
  31. tokens = tokens - requested
  32. allowed = true
  33. end
  34. ---- 更新桶状态并设置过期时间
  35. redis.call('HMSET', key, 'tokens', tokens, 'last_time', last_time)
  36. --redis.call('EXPIRE', key, math.ceil(capacity / rate) + 60) -- 适当延长过期时间
  37. return allowed and 1 or 0