Lua限流器的3种写法

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

引言

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

固定窗口限流

固定窗口限流也是最简单的限流算法，它将时间划分为固定大小的窗口，每个窗口内允许的请求次数是固定的，也是最好理解的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 thenreturn 0
else-- 否则请求次数加 1local count = redis.call('incr', key)-- 如果是第一次请求，设置过期时间，保证窗口的有效性if count == 1 thenredis.call('expire', key, window)end-- 返回 1 表示允许请求return 1
end

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

滑动窗口限流

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

-- 获取限流的键名
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 thenreturn 0
else-- 否则将当前请求的时间戳添加到有序集合中redis.call('zadd', key, now, now)-- 返回 1 表示允许请求return 1
end

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

令牌桶限流

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

-- 获取令牌桶的键名
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命令框架

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) {// 连接 RedisJedis 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();}
}

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