SpringbootInterceptor+Redis接口限流

SpringbootInterceptor+Redis接口限流

Springboot Intercepter拦截器 + Redis 限流

1. 添加依赖

<dependency>
    <groupId>org.springframework.boot</groupId>
    <artifactId>spring-boot-starter-data-redis</artifactId>
</dependency>

2. 创建注解

/**
 * @author devcxl
 */
@Target({ElementType.METHOD})
@Retention(RetentionPolicy.RUNTIME)
public @interface Limit {
    /**
     * 次数
     *
     * @return
     */
    int max() default 1;

    /**
     * 时间
     *
     * @return
     */
    int time() default 5;

    /**
     * 时间单位
     *
     * @return
     */
    TimeUnit timeUnit() default TimeUnit.SECONDS;

    /**
     * 提示
     *
     * @return
     */
    String msg() default "系统繁忙,请稍后再试。";
}

3. 拦截器实现

/**
 * 接口限流
 *
 * @author devcxl
 */
public class LimitInterceptor implements HandlerInterceptor {
    /**
     * redisKey模板
     */
    private static final  String LIMIT_KEY_TEMPLATE = "limit_%s_%s";

    @Resource
    private RedisTemplate<String, Integer> redisTemplate;

    /**
     * 检查API限制
     *
     * @param limit    注解
     * @param limitKey RedisKey
     * @return 是否拦截
     */
    private boolean checkLimit(Limit limit, String limitKey) {
        int max = limit.max();
        int time = limit.time();
        TimeUnit timeUnit = limit.timeUnit();
        Integer count = redisTemplate.opsForValue().get(limitKey);
        if (count != null) {
            if (count < max) {
                Long expire = redisTemplate.getExpire(limitKey);
                if (expire != null && expire <= 0) {
                    redisTemplate.opsForValue().set(limitKey, 1, time, timeUnit);
                } else {
                    redisTemplate.opsForValue().set(limitKey, ++count, time, timeUnit);
                }
            } else {
                // RuntimeException 可以改为自定义异常 可在全局异常捕获中处理并返回给前端
                throw new RuntimeException(limit.msg());
            }
        } else {
            redisTemplate.opsForValue().set(limitKey, 1, time, timeUnit);
        }
        return true;
    }

    /**
     * default preHandle
     *
     * @param request  请求
     * @param response 返回
     * @param handler  请求方法
     * @return 是否拦截
     */
    @Override
    public boolean preHandle(HttpServletRequest request, HttpServletResponse response, Object handler) {
        if (handler instanceof HandlerMethod) {
            HandlerMethod handlerMethod = (HandlerMethod) handler;
            Method method = handlerMethod.getMethod();
            Limit limit = method.getAnnotation(Limit.class);
            if (limit == null) {
                return true;
            } else {
                // 将请求路径和ip地址设为唯一标识 也可以自定义
                String limitKey = String.format(
                        LIMIT_KEY_TEMPLATE,
                        request.getRequestURI(),
                        ServletUtil.getClientIP(request)
                );
                return checkLimit(limit, limitKey);
            }
        } else {
            return false;
        }
    }

}

4.使用

@Slf4j
@RestController
@RequestMapping("/test")
public class TestController {
    @GetMapping("/limit")
    @Limit(max = 1, timeout = 300, timeunit = TimeUnit.MILLISECONDS, msg = "系统繁忙,请稍后再试。。。")
    public Resp<String> test1() {
        log.info("test1");
        return Resp.ok();
    }
}