简单限流器封装

开发过程中有时候 我们会做一些简单的限流 操作，比如 告警提醒，发送验证码 等，希望在 一段时间 只许调用几次。

下面基于redis incr 命令通用封装

@RequiredArgsConstructor
@Getter
public enum LimitTypeEnum {

    SECOND(1,"秒"),MINUTE(60,"分"),HOUR(60*60,"小时"),DAY(60*60*24,"天");
    /**
     * 类型
     */
    private final int time;

    /**
     * 描述
     */
    private final String description;
}

 * 限制调用次数
 */
@Target(ElementType.METHOD)
@Retention(RetentionPolicy.RUNTIME)
public @interface Limit {
    /**
     * 限制次数
     * @return
     */
    long count() default 1L;

    /**
     * 提示消息
     * @return
     */
    String msg() default "请求过于频繁";

    /**
     * 是否带方法参数 md5(类名+方法名+方法参数)
     * @return
     */
    boolean param() default true;
    /**
     * 限制类型
     * @return
     */
    LimitTypeEnum limitType() default LimitTypeEnum.SECOND;
}

@Aspect
@Component
public class LimitAspect {

    private static final String LOCK_REPEATED_SUBMIT = "limit:";

    @Autowired
    private RedisTemplate redisTemplate;

    @Around("@annotation(limit)")
    public Object around(ProceedingJoinPoint pjp, Limit limit) throws Throwable {
        try {
            //获取当前执行类
            String className = pjp.getSignature().getDeclaringTypeName();
            //获取当前执行类中的执行方法
            String methodName = pjp.getSignature().getName();
            String key = className + methodName;
            if(limit.param()){
                // 参数
                Map<String, Object> params = getRequestParams(pjp);
                key = key + JSON.toJSONString(params);
            }
            String md5 = SecureUtil.md5(key);
            String redisKey = LOCK_REPEATED_SUBMIT + md5;
            redisTemplate.opsForValue().setIfAbsent(redisKey, CommonConstants.ZERO,limit.limitType().getTime(),TimeUnit.SECONDS);
            Long count = redisTemplate.opsForValue().increment(redisKey);
            if(count <= limit.count()){
                Object result = pjp.proceed();
                return result;
            }
            throw new BusinessException(limit.msg());
        } catch (Throwable e) {
            throw e;
        }
    }


    /**
     * 获取入参
     * @param proceedingJoinPoint
     *
     * @return
     * */
    private Map<String, Object> getRequestParams(ProceedingJoinPoint proceedingJoinPoint) {
        Map<String, Object> requestParams = new HashMap<>();
        //参数名
        String[] paramNames =
                ((MethodSignature)proceedingJoinPoint.getSignature()).getParameterNames();
        //参数值
        Object[] paramValues = proceedingJoinPoint.getArgs();

        for (int i = 0; i < paramNames.length; i++) {
            Object value = paramValues[i];
            //如果是文件对象
            if (value instanceof MultipartFile) {
                MultipartFile file = (MultipartFile) value;
                //获取文件名
                value = file.getOriginalFilename();
            }
            if(value instanceof HttpServletResponse){
                continue;
            }
            if(value instanceof HttpServletRequest){
                continue;
            }
            requestParams.put(paramNames[i], value);
        }

        return requestParams;
    }
}

使用

image.png

测试

image.png

image.png

可见并发测试的时候，有5 次请求 被拒绝了，限流成功。

限流还有其他 方式 令牌桶，漏桶，滑动窗口 等

有兴趣 参考 Guava， redisson https://github.com/redisson/redisson/wiki/6.-%E5%88%86%E5%B8%83%E5%BC%8F%E5%AF%B9%E8%B1%A1#612-%E9%99%90%E6%B5%81%E5%99%A8ratelimiter （令牌桶）