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

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

一、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 then
    currentValue=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 then
    last_tokens=capacity
end
local last_time=tonumber(redis.call('get',time_key))
if last_time==nil then
    last_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) then
    acuire=1
    current_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;
    @Override
    public 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 ,请求令牌数：1
        return Arrays.asList("1", "5", now, "1");
    }

    @Override
    public 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 {

    @Autowired
    private LimiterIntercepter limiterIntercepter;
    @Override
    public void addInterceptors(InterceptorRegistry registry) {
        registry.addInterceptor(limiterIntercepter);
    }
}

三、Jemeter压测工具测试

使用Jemeter压测工具进行测试

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