Jedis使用lua脚本完成令牌桶限流

文章目录

Jedis使用lua脚本完成令牌桶限流
一、lua脚本的简单语法
二、令牌桶限流
- 1. 构思
- 2. 实现
三、Jemeter压测工具测试

一、lua脚本的简单语法

KEYS[1]
ARGV[1]
这两个参数分别代表了我们传入的key数组的一号元素和arg数组的一号元素
下面来看一个简单的使用
eval “return redis.call(‘get’,KEYS[1])” 1 one

首先eval 是解析Lua脚本的关键字，字符串中的就是写的lua脚本，其中
redis.call()表示调用redis。 1 表示数组的第一位，one表示的就是KEY[1]传入one。

下面演示jedis调用lua脚本实例

    public static void main(String[] args) throws IOException {Jedis jedis = new Jedis("127.0.0.1", 6379);ClassPathResource classPathResource = new ClassPathResource("/META-INF/script/redis_limiter.lua");byte[] buffer = new byte[(int)classPathResource.getFile().length()];classPathResource.getInputStream().read(buffer);String script = new String(buffer);System.out.println(script);String lua = jedis.scriptLoad(script);Object evalsha = jedis.evalsha(lua, getKeys(), getArgs());System.out.println(evalsha);}public static List<String> getKeys() {return Arrays.asList("one");}public static List<String> getArgs() {return Arrays.asList("50");}

local key=KEYS[1]
local arg=ARGV[1]local value=tonumber(arg)
local currentValue=tonumber(redis.call('get', key))
if currentValue==nil thencurrentValue=0
end
local new_value=(currentValue+value)redis.call('set',key,new_value)

上面的Lue脚本中，定义了key值和参数值，tonumber()表示转成数字。
这段的大概意思就是获取key的值，判断是否为空，为空的话表示该值从0开始，然后调用redis函数给key赋新值。

二、令牌桶限流

1. 构思

首先要明确什么是令牌桶限流，其实就是单位时间内向桶中发放令牌，如果有请求进来，请求申请令牌，如果令牌有的话，那么就可以拿到令牌，进行正常的操作，如果获取令牌失败，那么就拒绝访问。（桶中容量要做限制，不然有一段时间没人访问，一直方法令牌，当流量高峰达到，桶中已经屯了很多令牌，没有起到限流的作用）

基于这一目的，我们可以明确二点：

恒定速率发放令牌
令牌桶有最大容量

2. 实现

借助于Lua脚本我们可以实现原子性，因为redis处理业务的线程是单线程模式（主从复制时会另起线程）
那么我们就可以这样设计
首先在redis中存储一个key-value值表示令牌桶和令牌桶中的令牌数量
然后再存储一个上次访问的时间戳
在lua脚本中我们传入四个参数

令牌桶的速率
令牌桶的容量
请求时时间戳
请求得令牌数量
在使用redis.call()函数调用得到当前令牌桶数和上次刷新得时间戳
我们就可以计算出来这次时间和上次时间之间之差，算出来这两次访问之间产生多少令牌，然后使用上次访问得令牌桶数+中间访问产生得容量=当前令牌桶中数量。
那么我们就可以判断当前请求得令牌数量是否>当前桶中数量，如果大于return false,如果小于returen ture,并且在返回之前重新刷新redis中存储得令牌桶数和时间戳
下面我们来实战一下

--传入令牌桶得key和时间戳得key
local token_key=KEYS[1]
local time_key=KEYS[2]-- 分别传入 速率，容量，现在得时间，请求得令牌数
local rate=tonumber(ARGV[1])
local capacity=tonumber(ARGV[2])
local now_time=tonumber(ARGV[3])
local requestNum=tonumber(ARGV[4])
--获取上次访问 得令牌树和时间绰
local last_tokens=tonumber(redis.call('get',token_key))
if last_tokens==nil thenlast_tokens=capacity
end
local last_time=tonumber(redis.call('get',time_key))
if last_time==nil thenlast_time=0
end
--1，计算时间差
local time_del=math.max(0,(now_time-last_time))
--2,计算当前令牌桶中应该有多少令牌数
local current_tokens=math.min(capacity,(last_tokens+time_del*rate))
--3,判断是否允许访问
local acuire=0
if (requestNum <= current_tokens) thenacuire=1current_tokens=(current_tokens-requestNum)
end-- 计算一下过期时间 （默认就是填满令牌得时间*2）
local ttl = 60
--4，刷新记录redis.call('setex',token_key,ttl,current_tokens)redis.call('setex',time_key,ttl,now_time)return { acuire }

java代码
限流拦截器

package com.xzq.config;@Component
public class LimiterIntercepter implements HandlerInterceptor, InitializingBean {private Logger logger = LoggerFactory.getLogger(LimiterIntercepter.class);private static String TOKEN_BUCKET = "TOKEN_BUCKET";private static String REFRESH_TIME = "REFRESH_TIME";private String scriptLua;@Overridepublic boolean preHandle(HttpServletRequest request, HttpServletResponse response, Object handler) throws Exception {if (isallow()) {logger.info("允许进入");return true;}else{logger.info("限制进入");response.setStatus(500);return false;}}public boolean isallow() {Jedis jedis = new Jedis("127.0.0.1", 6379);long result = (long)((List) jedis.evalsha(scriptLua, getKeys(), getArgs())).get(0);return result == 1L;}public static List<String> getKeys() {return Arrays.asList(TOKEN_BUCKET, REFRESH_TIME);}public static List<String> getArgs() {String now = String.valueOf(System.currentTimeMillis() / 1000);// 速率： 1 ,  容量： 5  , 现在时间秒值: now ,请求令牌数：1return Arrays.asList("1", "5", now, "1");}@Overridepublic void afterPropertiesSet() throws Exception {Jedis jedis = new Jedis("127.0.0.1", 6379);ClassPathResource classPathResource = new ClassPathResource("/META-INF/script/redis_limiter.lua");byte[] buffer = new byte[(int)classPathResource.getFile().length()];classPathResource.getInputStream().read(buffer);scriptLua = jedis.scriptLoad(new String(buffer));}
}

mvc配置

package com.xzq.config;
@Component
public class MvcConfig implements WebMvcConfigurer {@Autowiredprivate LimiterIntercepter limiterIntercepter;@Overridepublic void addInterceptors(InterceptorRegistry registry) {registry.addInterceptor(limiterIntercepter);}
}

三、Jemeter压测工具测试

使用Jemeter压测工具进行测试

可以看到一秒内只允许进入五个请求，因为令牌桶容量为5，初始是值就是5，后面就是一秒进一个，因为我们设置得速率是1。