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

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

文章目录

一、lua脚本的简单语法

KEYS[1]
ARGV[1]
这两个参数分别代表了我们传入的key数组的一号元素和arg数组的一号元素
下面来看一个简单的使用
eval “return redis.call(‘get’,KEYS[1])” 1 one
Jedis使用lua脚本完成令牌桶限流
首先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压测工具进行测试
Jedis使用lua脚本完成令牌桶限流
可以看到一秒内只允许进入五个请求,因为令牌桶容量为5,初始是值就是5,后面就是一秒进一个,因为我们设置得速率是1。

上一篇:c# 连接MySQL 提示SSL Connection error


下一篇:Redis | 第9章 Lua 脚本与排序《Redis设计与实现》