redis 实现令牌桶 限速器

redis 实现令牌桶 限速器

下边实现了一个基于Redis 令牌桶算法的速率限制器，并使用lua保证对Redis的操作是原子性的：

@Component
public class RedisRateLimiter {
    private final RedisTemplate<String, String> redisTemplate;
    // Redis中保存令牌桶的key的前缀
    private static final String KEY_PREFIX = "RATE_LIMITER:";
    // 令牌桶的key
    private final String key;
    // 令牌桶的容量
    private final int capacity;
    // 每秒新增的令牌数
    private final int tokensPerSecond;
    // 令牌桶的填充时间
    private final int refillTime;
    // 每次填充的令牌数
    private final double refillAmount;
    // 用于执行Lua脚本的常量
//    1. 首先获取当前令牌桶中的令牌数。
//    2. 如果令牌桶中还没有令牌，那么初始化令牌桶，并将初始令牌数设定为容量参数。
//    3. 获取令牌桶的容量、填充时间和每次填充的令牌数等参数。
//    4. 计算当前请求到上次填充时间间隔的时间，并根据时间间隔和不能超过容量的新增令牌数量计算出新增的令牌数。
//    5. 根据新增的令牌数和当前令牌桶中的令牌数计算新的令牌数。
//    6. 如果新的令牌数小于1，则表示当前令牌桶中已经没有令牌了，返回0表示请求未通过。
//    7. 如果新的令牌数大于等于1，则更新令牌桶中的令牌数，并为令牌桶设定过期时间。
//    8. 返回1表示请求通过。
    private static final String SCRIPT =
            " local currentTokens = tonumber(redis.call('get', KEYS[1])) -- 获取当前令牌桶中的令牌数\n" +
                    "if currentTokens == nil then -- 如果令牌桶中还没有令牌\n" +
                    "    redis.call('set', KEYS[1], ARGV[1]) -- 则初始化令牌桶，并设定初始令牌数为容量参数\n" +
                    "    redis.call('pexpire', KEYS[1], ARGV[2]) -- 并为令牌桶设定过期时间，防止占用内存\n" +
                    "    return 1 -- 返回1，表示请求通过\n" +
                    "end\n" +
                    "local maxTokens = tonumber(ARGV[1]) -- 获取令牌桶的容量\n" +
                    "local refillTime = tonumber(ARGV[2]) -- 获取填充时间\n" +
                    "local refillAmount = tonumber(ARGV[3]) -- 获取每次填充的令牌数\n" +
                    "local timePassed = tonumber(redis.call('pttl', KEYS[1])) / 1000 -- 获取当前与上次填充时间间隔\n" +
                    "local tokensToAdd = math.floor(timePassed * refillAmount / refillTime) -- 计算新增的令牌数\n" +
                    "local newTokens = math.min(currentTokens + tokensToAdd, maxTokens) -- 计算新的令牌数\n" +
                    "if newTokens < 1 then -- 如果新的令牌数小于1，则表示当前令牌桶中已经没有令牌了\n" +
                    "    return 0 -- 返回0，表示请求未通过\n" +
                    "end\n" +
                    "redis.call('set', KEYS[1], newTokens) -- 更新令牌桶中的令牌数\n" +
                    "redis.call('pexpire', KEYS[1], ARGV[2]) -- 为令牌桶设定过期时间\n" +
                    "return 1 -- 返回1，表示请求通过";

    // 构造函数，使用Spring的依赖注入注入RedisTemplate和一些配置
    public RedisRateLimiter(RedisTemplate<String, String> redisTemplate,
                            @Value("${rate-limiter.key}") String key,
                            @Value("${rate-limiter.capacity}") int capacity,
                            @Value("${rate-limiter.tokens-per-second}") int tokensPerSecond) {
        this.redisTemplate = redisTemplate;
        this.key = KEY_PREFIX + key;
        this.capacity = capacity;
        this.tokensPerSecond = tokensPerSecond;
        this.refillTime = 1000 / tokensPerSecond;
        this.refillAmount = (double) refillTime / 1000 * capacity;
    }

    /**
     * 尝试获取令牌
     *
     * @return 是否获取到令牌
     */
    public boolean tryAcquire() {
        // 传递Lua脚本和参数执行Redis命令，返回执行结果
        List<String> keys = Collections.singletonList(key);
        long result = (long) redisTemplate.execute(new DefaultRedisScript<>(SCRIPT, Long.class),
                keys, capacity, refillTime, refillAmount);
        return result == 1;
    }
}

在上述代码中，我们定义了一个RedisRateLimiter类，它使用Redis的Sorted Set来实现令牌桶算法的速率限制器。使用redisTemplate.execute()方法来执行一个Lua脚本，实现对Redis的原子操作。在tryAcquire()方法中，我们传递了一个Lua脚本，它定义了一个名为SCRIPT的常量。在这个脚本中，我们使用Redis的get()、set()、pexpire()等命令来实现令牌桶算法的速率限制器，并保证了对Redis的操作是原子性的。

• 我们使用了Lua脚本来执行Redis命令，实现了对Redis的原子操作。
• 对于每次请求，我们都会计算出当前请求能否被处理，如果可以处理，那么就更新令牌桶的状态，并返回true；否则，返回false。
  还需要注意，我们使用了Spring的@Value注解来从application.properties文件中读取配置，这样可以使代码更加灵活。

pexpire 命令是Redis的一个键命令，用于为指定的键设置过期时间（以毫秒为单位）。如果键不存在，则该命令不执行任何操作。该命令的语法如下：
pexpire key milliseconds
其中， key 是要设置过期时间的键， milliseconds 是过期时间（以毫秒为单位）。如果过期时间设置为0，则键将立即过期。如果键已经有过期时间了， pexpire 会用新的过期时间覆盖原有的过期时间。该命令返回一个整数值，表示成功设置过期时间的键的数量。