lua 令牌桶 源码_基于Redis+LUA脚本的令牌桶算法限流策略实现
/**
*
*/
package com.matrix.cloud.service.redis.ratelimit;
import java.util.Collections;
import javax.annotation.Resource;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.data.redis.core.RedisCallback;
import org.springframework.data.redis.core.StringRedisTemplate;
import org.springframework.data.redis.core.script.DefaultRedisScript;
import org.springframework.data.redis.core.script.RedisScript;
import org.springframework.data.redis.serializer.StringRedisSerializer;
import org.springframework.stereotype.Service;
/**
* 1 使用redis+lua脚本实现接口限流
* 基于Redis的令牌桶算法限流策略实现 测试通过
* @author lmc
* 2018年12月21日
*/
@Service
public class RateLimitClient {
@Autowired
StringRedisTemplate redisTemplate;
//这里需要到redis配置文件下配置相关bean
@Qualifier("ratelimitInitLua")
@Resource
RedisScript ratelimitInitLua;
public boolean initToken(String key){
boolean token;
//此序列化redis会把数字类字符串存为数字类型
redisTemplate.setKeySerializer(new StringRedisSerializer());
redisTemplate.setValueSerializer(new StringRedisSerializer());
Long currMillSecond = redisTemplate.execute(
(RedisCallback) redisConnection -> redisConnection.time()
);
/** 初始化接收到的参数
* redis.pcall("HMSET",KEYS[1], "last_mill_second",ARGV[1], "curr_permits",ARGV[2],
* "max_burst",ARGV[3], "rate",ARGV[4], "app",ARGV[5])
*/
StringBuffer ratelimitStr= new StringBuffer();
ratelimitStr.append(" local result=1 ");
ratelimitStr.append(" redis.pcall('HMSET',KEYS[1], ");
ratelimitStr.append(" 'last_mill_second',ARGV[1], ");
ratelimitStr.append(" 'curr_permits',ARGV[2], ");
ratelimitStr.append(" 'max_burst',ARGV[3], ");
ratelimitStr.append(" 'rate',ARGV[4], ");
ratelimitStr.append(" 'app',ARGV[5]) ");
ratelimitStr.append(" return result ");
DefaultRedisScript ratelimitLua = new DefaultRedisScript<>(ratelimitStr.toString(), Long.class);
String last_mill_second = String.valueOf(currMillSecond);//上一次添加令牌的毫秒数
String curr_permits = "3";//令牌桶的最少令牌数
String max_permits = "200";//令牌桶的最大令牌数
String rate = "100";//向令牌桶中添加令牌的速率 , 令牌消耗速率
String app = "skynet";//定义标记，比如哪些是被限流的
Long accquire = redisTemplate.execute(ratelimitLua,Collections.singletonList("ratelimit:"+key), currMillSecond.toString(), curr_permits, max_permits, rate, app);
if (accquire == 1) {
token = true;
} else if (accquire == 0) {
token = true;
} else {
token = false;
}
return token;
}
/*
* last_mill_second 最后时间毫秒
curr_permits 当前可用的令牌
max_burst 令牌桶最大值
rate 每秒生成几个令牌
app 应用
令牌桶内令牌生成借鉴Guava-RateLimiter类的设计
每次根据时间戳生成token，不超过最大值
permits 每次请求令牌数
*/
public boolean accquireToken(String key, Integer permits) {
boolean token;
redisTemplate.setKeySerializer(new StringRedisSerializer());
redisTemplate.setValueSerializer(new StringRedisSerializer());
Long currMillSecond = redisTemplate.execute(
(RedisCallback) redisConnection -> redisConnection.time()
);
Long accquire = redisTemplate.execute(ratelimitInitLua,
Collections.singletonList("ratelimit:"+key), permits.toString(), currMillSecond.toString());
if (accquire == 1) {
token = true;
} else {
token = false;
}
return token;
}
}
ratelimitInit.lua 配置文件加载进来
local ratelimit_info=redis.pcall("HMGET",KEYS[1],"last_mill_second","curr_permits","max_burst","rate","app")
local last_mill_second=ratelimit_info[1]
local curr_permits=tonumber(ratelimit_info[2])
local max_burst=tonumber(ratelimit_info[3])
local rate=tonumber(ratelimit_info[4])
local app=tostring(ratelimit_info[5])
if app == nil then
return 0
end
local local_curr_permits=max_burst;
if(type(last_mill_second) ~='boolean' and last_mill_second ~=nil) then
local reverse_permits=math.floor((ARGV[2]-last_mill_second)/1000)*rate
if(reverse_permits>0) then
redis.pcall("HMSET",KEYS[1],"last_mill_second",ARGV[2])
end
local expect_curr_permits=reverse_permits+curr_permits
local_curr_permits=math.min(expect_curr_permits,max_burst);
else
redis.pcall("HMSET",KEYS[1],"last_mill_second",ARGV[2])
end
local result=-1
if(local_curr_permits-ARGV[1]>0) then
result=1
redis.pcall("HMSET",KEYS[1],"curr_permits",local_curr_permits-ARGV[1])
else
redis.pcall("HMSET",KEYS[1],"curr_permits",local_curr_permits)
end
return result
调用 方式 方式简单点
@RequestMapping(value = "/getRatelimitlua", method = RequestMethod.POST)
@ResponseBody
public Object RatelimitRedisLua() {//@RequestParam(value="uid") String uid
String key="lmc168";
rateLimitClient.initToken(key);//初如化开始
if (!rateLimitClient.accquireToken(key, 1)) {
System.out.println("触发限流API:调用太忙了,请休息下");
//throw new Exception();
}else{
System.out.println("没有触发限流策略");
}
return key;
}