簡單限流器封裝

開發(fā)過程中有時(shí)候 我們會做一些簡單的限流 操作，比如 告警提醒，發(fā)送驗(yàn)證碼 等，希望在 一段時(shí)間 只許調(diào)用幾次。

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

@RequiredArgsConstructor
@Getter
public enum LimitTypeEnum {

    SECOND(1,"秒"),MINUTE(60,"分"),HOUR(60*60,"小時(shí)"),DAY(60*60*24,"天");
    /**
     * 類型
     */
    private final int time;

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

 * 限制調(diào)用次數(shù)
 */
@Target(ElementType.METHOD)
@Retention(RetentionPolicy.RUNTIME)
public @interface Limit {
    /**
     * 限制次數(shù)
     * @return
     */
    long count() default 1L;

    /**
     * 提示消息
     * @return
     */
    String msg() default "請求過于頻繁";

    /**
     * 是否帶方法參數(shù) md5(類名+方法名+方法參數(shù))
     * @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 {
            //獲取當(dāng)前執(zhí)行類
            String className = pjp.getSignature().getDeclaringTypeName();
            //獲取當(dāng)前執(zhí)行類中的執(zhí)行方法
            String methodName = pjp.getSignature().getName();
            String key = className + methodName;
            if(limit.param()){
                // 參數(shù)
                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;
        }
    }


    /**
     * 獲取入?yún)?     * @param proceedingJoinPoint
     *
     * @return
     * */
    private Map<String, Object> getRequestParams(ProceedingJoinPoint proceedingJoinPoint) {
        Map<String, Object> requestParams = new HashMap<>();
        //參數(shù)名
        String[] paramNames =
                ((MethodSignature)proceedingJoinPoint.getSignature()).getParameterNames();
        //參數(shù)值
        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

可見并發(fā)測試的時(shí)候，有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 （令牌桶）