使用令牌桶算法通过redis实现限流

令牌桶算法是一种常用的限流算法，它可以平滑地控制请求的处理速率。在 Java 中结合 Redis 实现令牌桶算法，可以利用 Redis 的原子操作来保证多节点环境下的限流效果。

一 实现思路

1. 初始化令牌桶：在 Redis 中存储令牌桶的相关信息，如令牌数量、令牌桶容量、上一次更新时间等。
2. 添加令牌：根据时间间隔，计算需要添加到令牌桶中的令牌数量，并更新令牌桶的令牌数量。
3. 获取令牌：尝试从令牌桶中获取指定数量的令牌，如果令牌数量足够，则返回成功；否则返回失败。

二 Redis Lua 脚本

编写一个 Lua 脚本来处理令牌桶的逻辑。lua脚本会在 Redis 执行时保证原子性。

-- 最大令牌个数(qps)
local limit = tonumber(ARGV[1]) 
-- 令牌生成速率
local rate = tonumber(ARGV[2]) 
-- 当前时间秒*1000=毫秒
local now = redis.call('TIME')[1] * 1000
-- 剩余可用的令牌个数
local tokens = tonumber(redis.call('GET', KEYS[1]) or 0)
-- 上次访问时间戳(单位毫秒)
local lastTimestamp = tonumber(redis.call('GET', KEYS[2]) or 0)
-- 本次访问时间戳-上次访问时间戳
local elapsedTime = now - lastTimestamp
-- 计算本次访问与上次访问之间需要放入的令牌个数=上次剩余的令牌数+本次访问与上次访问之间差的令牌数
local newTokens = tokens + elapsedTime * rate / 1000
-- 如果需要放入的令牌个数超过限制的最大令牌个数，则直接使用最大令牌个数
if newTokens > limit then
	newTokens = limit 
end 
-- 如果需要放入的令牌个数小于1，则意味着无法获取到锁
if newTokens < 1 then
	return 0 
else
-- 获取到锁，则更新剩余可获取的锁个数-1，更新本次访问时间戳
	redis.call('SET', KEYS[1], newTokens - 1)
	redis.call('SET', KEYS[2], now)
	return 1 
end

三 Java 实现

public class RedisRateLimiter {

  private JedisPool jedisPool;

  private String LUA_SCRIPT = "上面的lua脚本替换到此处";

  public RedisRateLimiter(JedisPool jedisPool) {
    this.jedisPool = jedisPool;
  }

  public boolean tryAcquire(String key, int rate, int capacity) {
        try (Jedis jedis = jedisPool.getResource()) {
      Object result = jedis.eval(LUA_SCRIPT,
          Arrays.asList(key, key + ":timestamp"), Arrays.asList(capacity, rate));
      long response = (Long) result;
      return response != 0;
    } catch (Exception e) {
        return false;
    }
  }
}