SpringBoot如何进行限流，老鸟们还可以这样玩！
大家好，我是飘渺。
在SpringBoot 如何进行限流，老鸟们都这么玩的!一文中我们详细介绍了为什么需要对接口进行限流，也介绍了常见的限流算法，最后还基于Guava工具类实现了接口限流。但是这种方式有个问题，无法实现分布式限流。那今天我们来利用Redis + Lua 来实现分布式限流。
Lua 脚本和 MySQL 数据库的存储过程比较相似，他们执行一组命令，所有命令的执行要么全部成功或者失败，以此达到原子性。也可以把 Lua 脚本理解为，一段具有业务逻辑的代码块。
实现过程
第一步：引入Redis依赖包
<dependency>
<groupId>org.springframework.boot</groupId>
<artifactId>spring-boot-starter-data-redis</artifactId>
</dependency>
第二步：配置Redis
/**
* @author JAVA日知录
* @date 2022/5/2 22:35
*/
@Configuration
public class RedisConfig {
@Bean
public RedisTemplate<String, Object> redisTemplate(RedisConnectionFactory factory) {
RedisTemplate<String, Object> template = new RedisTemplate<>();
template.setConnectionFactory(factory);
// 使用Jackson2JsonRedisSerializer来序列化和反序列化redis的value值（默认使用JDK的序列化方式）
Jackson2JsonRedisSerializer<Object> serializer = new Jackson2JsonRedisSerializer<>(Object.class);
ObjectMapper mapper = new ObjectMapper();
mapper.setVisibility(PropertyAccessor.ALL, JsonAutoDetect.Visibility.ANY);
mapper.activateDefaultTyping(LaissezFaireSubTypeValidator.instance,
ObjectMapper.DefaultTyping.NON_FINAL, JsonTypeInfo.As.PROPERTY);
serializer.setObjectMapper(mapper);
template.setValueSerializer(serializer);
// 使用StringRedisSerializer来序列化和反序列化redis的key值
template.setKeySerializer(new StringRedisSerializer());
template.afterPropertiesSet();
return template;
}
}
第二步：自定义限流注解
@Retention(RetentionPolicy.RUNTIME)
@Target({ElementType.METHOD})
@Documented
public @interface RedisLimit {
/**
* 资源的key,唯一
* 作用：不同的接口，不同的流量控制
*/
String key() default "";
/**
* 最多的访问限制次数
*/
long permitsPerSecond() default 2;
/**
* 过期时间也可以理解为单位时间，单位秒，默认60
*/
long expire() default 60;
/**
* 得不到令牌的提示语
*/
String msg() default "系统繁忙,请稍后再试.";
}
第三步：创建限流异常
/**
* @author JAVA日知录
* Redis限流自定义异常
* @date 2022/5/2 21:43
*/
public class RedisLimitException extends RuntimeException{
public RedisLimitException(String msg) {
super( msg );
}
}
第四步：使用AOP切面拦截限流注解
/**
* Limit AOP
* @author JAVA日知录
* @date 2021/9/24 3:07 下午
*/
@Slf4j
@Aspect
@Component
public class RedisLimitAop {
@Autowired
private StringRedisTemplate stringRedisTemplate;
@Pointcut("@annotation(com.jianzh5.blog.limit.redis.RedisLimit)")
private void check() {
}
@Before("check()")
public void before(JoinPoint joinPoint) {
MethodSignature signature = (MethodSignature) joinPoint.getSignature();
Method method = signature.getMethod();
//拿到RedisLimit注解，如果存在则说明需要限流
RedisLimit redisLimit = method.getAnnotation(RedisLimit.class);
if(redisLimit != null){
//获取redis的key
String key = redisLimit.key();
String className = method.getDeclaringClass().getName();
String name = method.getName();
String limitKey = key + className + method.getName();
log.info(limitKey);
if(StringUtils.isEmpty(key)){
throw new RedisLimitException( "key cannot be null" );
}
long limit = redisLimit.permitsPerSecond();
long expire = redisLimit.expire();
List<String> keys = new ArrayList<>();
keys.add( key );
String luaScript = buildLuaScript();
RedisScript<Long> redisScript = new DefaultRedisScript<>( luaScript, Long.class );
Long count = stringRedisTemplate.execute( redisScript, keys, String.valueOf(limit), String.valueOf(expire) );
log.info( "Access try count is {} for key={}", count, key );
if (count != null && count == 0) {
log.debug("令牌桶={}，获取令牌失败",key);
throw new RedisLimitException(redisLimit.msg());
}
}
}
/**
* 构建redis lua脚本
* @return
*/
private String buildLuaScript() {
StringBuilder luaString = new StringBuilder();
luaString.append( "local key = KEYS[1]" );
//获取ARGV内参数Limit
luaString.append( "\nlocal limit = tonumber(ARGV[1])" );
//获取key的次数
luaString.append( "\nlocal curentLimit = tonumber(redis.call('get', key) or \"0\")" );
luaString.append( "\nif curentLimit + 1 > limit then" );
luaString.append( "\nreturn 0" );
luaString.append( "\nelse" );
//自增长 1
luaString.append( "\n redis.call(\"INCRBY\", key, 1)" );
//设置过期时间
luaString.append( "\nredis.call(\"EXPIRE\", key, ARGV[2])" );
luaString.append( "\nreturn curentLimit + 1" );
luaString.append( "\nend" );
return luaString.toString();
}
}
第五步：给需要限流的接口加上注解
/**
* 公众号：JAVA日知录
* 限流测试类基于Redis限流
*/
@Slf4j
@RestController
@RequestMapping("/limit/redis")
public class LimitRedisController {
/**
* 基于Redis AOP限流
*/
@GetMapping("/test")
@RedisLimit(key = "redis-limit:test", permitsPerSecond = 2, expire = 1, msg = "当前排队人数较多，请稍后再试！")
public String test() {
log.info("限流成功。。。");
return "ok";
}
}
第六步：体验效果
通过访问测试地址： http://127.0.0.1:8080/limit/redis/test，反复刷新并观察输出结果：
正常响应时：
{"status":100,"message":"操作成功","data":"ok","timestamp":1652343229643}
触发限流时：
{"status":500,"message":"当前排队人数较多，请稍后再试！","data":null,"timestamp":1652343239035}
通过观察得之，基于自定义注解同样实现了接口限流的效果。
优化
程序每次执行每次都需要通过buildLuaScript()方法构建lua执行脚本，比较 low，我们可以生成一个lua文件放在resources目录下，利用@PostConstruct注解提前加载。
- 在resouces文件夹下创建lua文件 rateLimiter.lua
--获取KEY
local key = KEYS[1]
local limit = tonumber(ARGV[1])
local curentLimit = tonumber(redis.call('get', key) or "0")
if curentLimit + 1 > limit
then return 0
else
-- 自增长 1
redis.call('INCRBY', key, 1)
-- 设置过期时间
redis.call('EXPIRE', key, ARGV[2])
return curentLimit + 1
end
- 修改RedisLimitAop，通过@PostConstruct注入DefaultRedisScript
@Slf4j
@Aspect
@Component
public class RedisLimitAop {
@Autowired
private StringRedisTemplate stringRedisTemplate;
private DefaultRedisScript<Long> redisScript;
@PostConstruct
public void init(){
redisScript = new DefaultRedisScript<>();
redisScript.setResultType(Long.class);
redisScript.setScriptSource(new ResourceScriptSource(new ClassPathResource("rateLimiter.lua")));
}
...
}
小结
基于Redis + Lua 可以很方便地实现分布式限流，算是SpringBoot老鸟系列限流文章的补充扩展。
那么现在问题来了，我们现在有基于Guava实现的单机限流，又有基于Redis+Lua实现的分布式限流，那能不能将两种限流功能做成一个独立的公共组件，让使用方根据实际情况选择对应的限流功能呢？
老鸟系列源码已经上传至GitHub，需要的在公号【JAVA日知录】回复关键字 0923 获取源码地址。
