Lua限流器的3种写法

学而不思则罔，思而不学则殆

引言

上篇文章讲解了Lua脚本，事务和Pipline之间的使用方式和性能差距，本篇文章将聚焦Lua脚本，我将用三种写法来展现如何实现一个Redis限流器

固定窗口限流

固定窗口限流也是最简单的限流算法，它将时间划分为固定大小的窗口，每个窗口内允许的请求次数是固定的，也是最好理解的lua脚本

Lua 复制代码

-- 获取限流的键名，通常为一个唯一标识，如用户 ID、IP 等
local key = KEYS[1]
-- 获取该窗口内允许的最大请求次数
local limit = tonumber(ARGV[1])
-- 获取当前窗口的时间长度（秒）
local window = tonumber(ARGV[2])

-- 获取当前键对应的请求次数
local current = tonumber(redis.call('get', key) or "0")

-- 如果当前请求次数超过了限制，返回 0 表示限流
if current + 1 > limit then
    return 0
else
    -- 否则请求次数加 1
    local count = redis.call('incr', key)
    -- 如果是第一次请求，设置过期时间，保证窗口的有效性
    if count == 1 then
        redis.call('expire', key, window)
    end
    -- 返回 1 表示允许请求
    return 1
end

本质是一个时间长度也就是窗口中我只允许固定次数通过，所以由此我们可以延伸出滑动窗口

滑动窗口限流

滑动窗口限流算法在固定窗口的基础上进行了改进，它将时间窗口划分为更小的时间片，每个时间片记录该时间段内的请求次数，通过滑动窗口的方式统计一定时间内的总请求次数

Lua 复制代码

-- 获取限流的键名
local key = KEYS[1]
-- 获取该窗口内允许的最大请求次数
local limit = tonumber(ARGV[1])
-- 获取当前时间戳（毫秒）
local now = tonumber(ARGV[2])
-- 获取滑动窗口的时间长度（毫秒）
local window = tonumber(ARGV[3])

-- 移除窗口外的请求记录
redis.call('zremrangebyscore', key, 0, now - window)
-- 获取当前窗口内的请求数量
local current = tonumber(redis.call('zcard', key))

-- 如果当前请求次数超过了限制，返回 0 表示限流
if current + 1 > limit then
    return 0
else
    -- 否则将当前请求的时间戳添加到有序集合中
    redis.call('zadd', key, now, now)
    -- 返回 1 表示允许请求
    return 1
end

原本的维度是一个窗口，而现在是将窗口放大成时间单位，然后用颗粒度更小的窗口去获取次数，最后统计这个大窗口的流量不变就行

令牌桶限流

令牌桶算法可以说是一种比较常用的限流算法，它以固定的速率向桶中添加令牌，每个请求需要从桶中获取一个或多个令牌才能被处理

Lua 复制代码

-- 获取令牌桶的键名
local key = KEYS[1]
-- 获取桶的容量
local capacity = tonumber(ARGV[1])
-- 获取令牌生成的速率（每秒生成的令牌数）
local rate = tonumber(ARGV[2])
-- 获取当前时间戳（秒）
local now = tonumber(ARGV[3])
-- 获取每个请求需要的令牌数
local tokens_needed = tonumber(ARGV[4])

-- 获取上次更新的时间和剩余的令牌数
local last_update = tonumber(redis.call('hget', key, 'last_update') or now)
local tokens = tonumber(redis.call('hget', key, 'tokens') or capacity)

-- 计算从上次更新到现在应该生成的令牌数
local new_tokens = math.min(capacity, tokens + (now - last_update) * rate)

-- 如果剩余的令牌数足够，处理请求
if new_tokens >= tokens_needed then
    -- 更新剩余的令牌数
    local remaining_tokens = new_tokens - tokens_needed
    -- 更新上次更新的时间和剩余的令牌数
    redis.call('hset', key, 'last_update', now)
    redis.call('hset', key, 'tokens', remaining_tokens)
    -- 返回 1 表示允许请求
    return 1
else
    -- 否则返回 0 表示限流
    return 0
end

实战Java

一般我们的Lua脚本都是放在配置文件中管理，无论是阅读性还是维护性都是最优解

例如把上面的任一脚本创建于rate_limiter.lua文件中（自定义名称）

配合Jedis我们常用的redis命令框架

java 复制代码

import redis.clients.jedis.Jedis;
import java.io.BufferedReader;
import java.io.FileReader;
import java.io.IOException;
import java.util.Arrays;
import java.util.List;

public class RedisRateLimiterFromConfigJedis {
    public static void main(String[] args) {
        // 连接 Redis
        Jedis jedis = new Jedis("localhost", 6379);

        // 限流键名
        String key = "rate_limit:user:1";
        // 窗口内允许的最大请求次数
        int limit = 10;
        // 窗口时间长度（秒）
        int window = 60;

        // 读取 Lua 脚本文件
        String script = readScriptFromFile("rate_limiter.lua");

        // 键名参数列表
        List<String> keys = Arrays.asList(key);
        // 其他参数列表
        List<String> argsList = Arrays.asList(String.valueOf(limit), String.valueOf(window));

        // 执行 Lua 脚本
        Object result = jedis.eval(script, keys, argsList);

        if (result instanceof Long && (Long) result == 1) {
            System.out.println("允许请求");
        } else {
            System.out.println("请求被限流");
        }

        // 关闭连接
        jedis.close();
    }

    private static String readScriptFromFile(String filePath) {
        StringBuilder script = new StringBuilder();
        try (BufferedReader reader = new BufferedReader(new FileReader(filePath))) {
            String line;
            while ((line = reader.readLine()) != null) {
                script.append(line).append("\n");
            }
        } catch (IOException e) {
            e.printStackTrace();
        }
        return script.toString();
    }
}

其中端口和文件名可以自定义