核心思想:通过redis的zset来实现窗口滑动,从而达到限流的作用。
yml核心配置如下:
spring:
application:
name: demo
redis:
host: 127.0.0.1
port: 6379
database: 1
server:
port: 8181
servlet:
context-path: /demo
核心pom文件如下:
<!-- redis相关 -->
<dependency>
<groupId>org.springframework.boot</groupId>
<artifactId>spring-boot-starter-data-redis</artifactId>
</dependency>
<!-- aop相关 -->
<dependency>
<groupId>org.springframework.boot</groupId>
<artifactId>spring-boot-starter-aop</artifactId>
</dependency>
<!-- lombak相关-->
<dependency>
<groupId>org.projectlombok</groupId>
<artifactId>lombok</artifactId>
</dependency>
1、自定义注解
@Target(ElementType.METHOD)
@Retention(RetentionPolicy.RUNTIME)
public @interface Limit {
/**
* 资源描述(日志打印使用无实际意义)
*/
String name() default "";
/**
* 资源前缀
*/
String prefix() default "";
/**
* 资源的key
*/
String key() default "";
/**
* 时间窗口 单位 s
*/
int period();
/**
* 限制访问次数
*/
int count();
}
2、自定义统一返回 + 自定义异常 + 全局异常拦截
- 统一结构返回
@Data
public class BaseResponse<T> {
private int code;
private T data;
private String message;
public BaseResponse() {
this.code = HttpStatus.INTERNAL_SERVER_ERROR.value();
this.message = HttpStatus.INTERNAL_SERVER_ERROR.getReasonPhrase();
}
public BaseResponse(int code, String message, T data) {
this.code = code;
this.message = message;
this.data = data;
}
public static <T> BaseResponse<T> with(int code, String msg, T data) {
return new BaseResponse<>(code, msg, data);
}
public static <T> BaseResponse<T> success(T data) {
return with(HttpStatus.OK.value(), "success", data);
}
}
- 自定义异常
public class LimitException extends RuntimeException{
public LimitException(String msg) {
super(msg);
}
}
- 全局异常拦截器
@Slf4j
@RestControllerAdvice
public class GlobalExceptionHandler {
@ExceptionHandler(LimitException.class)
public BaseResponse<?> limitException(LimitException limitException) {
log.error("【限流发生异常】");
return BaseResponse.with(HttpStatus.INTERNAL_SERVER_ERROR.value(), limitException.getMessage());
}
}
3、自定义切面 + lua脚本
- 自定义切面处理限流
@Aspect
@Component
public class LimitAspect {
@Resource
private RedisTemplate<String, Object> redisTemplate;
private DefaultRedisScript<Long> redisScript;
@PostConstruct
public void init() {
// 仅加载一次lua脚本资源
redisScript = new DefaultRedisScript<>();
redisScript.setLocation(new ClassPathResource("lua/limit.lua"));
redisScript.setResultType(Long.class);
}
@Before("@annotation(com.learn.demo.annotation.Limit)")
public void handle() {
// 先去获取资源信息
Method method = resolveMethod(point);
Limit limit = method.getAnnotation(Limit.class);
if (limit != null) {
// 去执行脚本 然后判断是否超过限流的时间了
String key = limit.prefix() + "_" + limit.key();
log.info("{}访问了系统{}", LocalDateTime.now().format(DateTimeFormatter.ofPattern("yyyy-MM-dd HH:mm:ss")), limit.name());
Long number = redisTemplate.execute(redisScript, Collections.singletonList(key), limit.period(), System.currentTimeMillis(), limit.count());
if (number != null && number.intValue() >= limit.count()) {
throw new LimitException("超过接口访问限制,请稍后重试");
}
}
}
private Method resolveMethod(JoinPoint point) {
// 获取方法签名
MethodSignature signature = (MethodSignature) point.getSignature();
Class<?> targetClass = point.getTarget().getClass();
Method method = getDeclaredMethod(targetClass, signature.getName(), signature.getMethod().getParameterTypes());
if (method == null) {
throw new IllegalStateException("无法解析目标方法: " + signature.getMethod().getName());
}
return method;
}
private Method getDeclaredMethod(Class<?> clazz, String name, Class<?>... parameterTypes) {
try {
return clazz.getDeclaredMethod(name, parameterTypes);
} catch (NoSuchMethodException e) {
Class<?> superClass = clazz.getSuperclass();
if (superClass != null) {
return getDeclaredMethod(superClass, name, parameterTypes);
}
}
return null;
}
}
- lua脚本
--获取KEY
local key = KEYS[1];
--获取ARGV内的参数
-- 缓存时间
local expire = tonumber(ARGV[1]);
-- 当前时间
local currentMs = tonumber(ARGV[2]);
-- 最大次数
local count = tonumber(ARGV[3]);
--窗口开始时间
local windowStartMs = currentMs - expire * 1000;
--获取key的次数
local current = redis.call('zcount', key, windowStartMs, currentMs);
--如果key的次数存在且大于预设值直接返回当前key的次数
if current and tonumber(current) >= count then
return tonumber(current);
end
-- 清除上一个窗口的数据
redis.call("ZREMRANGEBYSCORE", key, 0, windowStartMs);
-- 添加当前成员
redis.call("zadd", key, tostring(currentMs), currentMs);
redis.call("expire", key, expire);
--返回key的次数
return tonumber(current)
4、应用和测试
@RestController
@RequestMapping("/test")
public class TestController {
@GetMapping("/limit")
@Limit(name = "测试限流", prefix = "test", key = "limit", period = 10, count = 5)
public BaseResponse<?> limit() {
return BaseResponse.success("success");
}
}
可以看到 在第6次访问的时候 发生了异常
并且通过postman测试得到的返回也是限流的返回
5、踩坑记录
在使用过程中,发现执行lua脚本时一直获取不到ARGV的值,报错信息为:@user_script:20: user_script:20: attempt to perform arithmetic on local ‘expire’ (a nil value)。这个是因为lua脚本中的值为空,并且还使用该值。
解决方案:指定RedisTemplate的序列化
@Configuration
public class RedisConfig {
@Bean(name = "redisTemplate")
@ConditionalOnClass(RedisOperations.class)
public RedisTemplate<String, Object> redisTemplate(RedisConnectionFactory factory) {
RedisTemplate<String, Object> template = new RedisTemplate<>();
template.setConnectionFactory(factory);
Jackson2JsonRedisSerializer<Object> jackson2JsonRedisSerializer = new Jackson2JsonRedisSerializer<>(Object.class);
ObjectMapper mapper = new ObjectMapper();
mapper.setVisibility(PropertyAccessor.ALL, JsonAutoDetect.Visibility.ANY);
mapper.activateDefaultTyping(mapper.getPolymorphicTypeValidator(), ObjectMapper.DefaultTyping.NON_FINAL);
jackson2JsonRedisSerializer.setObjectMapper(mapper);
StringRedisSerializer stringRedisSerializer = new StringRedisSerializer();
// 指定 key 的序列化格式
template.setKeySerializer(stringRedisSerializer);
// 指定 hash 的 key 的序列化格式
template.setHashKeySerializer(stringRedisSerializer);
// 指定 value 的序列化格式
template.setValueSerializer(jackson2JsonRedisSerializer);
// 指定 hash 的 value 的序列化格式
template.setHashValueSerializer(jackson2JsonRedisSerializer);
template.afterPropertiesSet();
return template;
}
}
PS:这个只是最基础版的窗口窗口限流,无用户、ip等限制,支持的窗口时间也为秒,看官们可以通过根据修改Limit注解来实现功能的扩展。如果写的有什么问题和错误欢迎大家指出,如果有什么好的建议也欢迎大家提出,一同进步!!!