Java Redisson多策略註解限流

/**
 * aop限流註解
 */
@Target({ElementType.METHOD, ElementType.TYPE})
@Retention(RetentionPolicy.RUNTIME)
@Inherited
@Documented
public @interface RedisLimit {


    String prefix() default "rateLimit:";

    //限流唯一標示
    String key() default "";

    //限流單位時間（單位為s）
    int time() default 1;

    //單位時間內限制的訪問次數
    int count();

    //限流類型
    LimitType type() default LimitType.CUSTOM;

}

public enum LimitType {

    /**
     * 自定義key
     */
    CUSTOM,

    /**
     * 請求者IP
     */
    IP,

    /**
     * 方法級別限流
     * key = ClassName+MethodName
     */
    METHOD,

    /**
     * 參數級別限流
     * key = ClassName+MethodName+Params
     */
    PARAMS,

    /**
     * 用戶級別限流
     * key = ClassName+MethodName+Params+UserId
     */
    USER,

    /**
     * 根據request的uri限流
     * key = Request_uri
     */
    REQUEST_URI,

    /**
     * 對requesturi+userId限流
     * key = Request_uri+UserId
     */
    REQUESTURI_USERID,


    /**
     * 對userId限流
     * key = userId
     */
    SINGLEUSER,

    /**
     * 對方法限流
     * key = ClassName+MethodName
     */
    SINGLEMETHOD,

    /**
     * 對uri+params限流
     * key = uri+params
     */
    REQUEST_URI_PARAMS,

    /**
     * 對uri+params+userId限流
     * key = uri+params+userId
     */
    REQUEST_URI_PARAMS_USERID;
    
}

@Component
public class ProceedingJoinPointUtil {
    @Autowired
    private HttpServletRequest request;

    private Map<LimitType, Function<ProceedingJoinPoint,String>> functionMap = new HashMap<>(9);

    @PostConstruct
    void initMap(){
        //初始化策略
        functionMap.put(LimitType.METHOD, this::getMethodTypeKey);
        functionMap.put(LimitType.PARAMS, this::getParamsTypeKey);
        functionMap.put(LimitType.USER, this::getUserTypeKey);
        functionMap.put(LimitType.REQUEST_URI,proceedingJoinPoint ->
                request.getRequestURI());
        functionMap.put(LimitType.REQUESTURI_USERID, proceedingJoinPoint ->
                request.getRequestURI()+BaseContext.getUserId());
        functionMap.put(LimitType.REQUEST_URI_PARAMS,proceedingJoinPoint ->
                request.getRequestURI()+getParams(proceedingJoinPoint));
        functionMap.put(LimitType.REQUEST_URI_PARAMS_USERID,proceedingJoinPoint ->
                request.getRequestURI()+getParams(proceedingJoinPoint)+BaseContext.getUserId());
        functionMap.put(LimitType.SINGLEUSER,(proceedingJoinPoint)->
                String.valueOf(BaseContext.getUserId()));
        functionMap.put(LimitType.SINGLEMETHOD,(proceedingJoinPoint -> {
            StringBuilder sb = new StringBuilder();
            appendMthodName(proceedingJoinPoint,sb);
            return sb.toString();
        }));
    }

    public Object getKey(ProceedingJoinPoint joinPoint, RedisLimit redisLimit) {
        //根據限制類型生成key
        Object generateKey = "";
        //自定義
        if(redisLimit.type() != LimitType.CUSTOM){
            generateKey = generateKey(redisLimit.type(), joinPoint);
        }else {
            //非自定義
            generateKey = redisLimit.key();
        }
        return generateKey;
    }

    /**
     * 根據LimitType生成key
     * @param type
     * @param joinPoint
     * @return
     */
    private Object generateKey(LimitType type , ProceedingJoinPoint joinPoint) {
        Function function = functionMap.get(type);
        Object result = function.apply(joinPoint);
        return result;
    }

    /**
     * 方法級別
     * key = ClassName+MethodName
     * @param joinPoint
     * @return
     */
    private String getMethodTypeKey(ProceedingJoinPoint joinPoint){
        StringBuilder sb = new StringBuilder();
        appendMthodName(joinPoint, sb);
        return sb.toString();
    }



    /**
     * 參數級別
     * key = ClassName+MethodName+Params
     * @param joinPoint
     * @return
     */
    private String getParamsTypeKey(ProceedingJoinPoint joinPoint){
        StringBuilder sb = new StringBuilder();
        appendMthodName(joinPoint, sb);
        appendParams(joinPoint, sb);
        return sb.toString();
    }



    /**
     * 用戶級別
     * key = ClassName+MethodName+Params+UserId
     */
    private String getUserTypeKey(ProceedingJoinPoint joinPoint){
        StringBuilder sb = new StringBuilder();
        appendMthodName(joinPoint, sb);
        appendParams(joinPoint, sb);
        //獲取userId
        appendUserId(sb);
        return sb.toString();
    }


    /**
     * StringBuilder添加類名和方法名
     * @param joinPoint
     * @param sb
     */
    private void appendMthodName(ProceedingJoinPoint joinPoint, StringBuilder sb) {
        Signature signature = joinPoint.getSignature();
        MethodSignature methodSignature = (MethodSignature) signature;
        Method method = methodSignature.getMethod();
        sb.append(joinPoint.getTarget().getClass().getName())//類名
                .append(method.getName());//方法名
    }

    /**
     * StringBuilder添加方法參數值
     * @param joinPoint
     * @param sb
     */
    private void appendParams(ProceedingJoinPoint joinPoint, StringBuilder sb) {
        for (Object o : joinPoint.getArgs()) {
            sb.append(o.toString());
        }
    }
    private String getParams(ProceedingJoinPoint joinPoint) {
        StringBuilder sb = new StringBuilder();
        for (Object o : joinPoint.getArgs()) {
            if(o instanceof MultipartFile){
                try {
                    ImageTypeCheck.getImgHeightAndWidth(((MultipartFile) o).getInputStream());
                } catch (IOException e) {
                    throw new BusinessException("MultipartFile輸入流獲取失敗,source:ProceedingJoinPointUtils.149",USER_PRINCIPAL_EMAIL);
                }
            }else {
                sb.append(o.toString());
            }
        }
        return sb.toString();
    }

    /**
     * StringBuilder添加UserId
     * @param sb
     */
    private void appendUserId(StringBuilder sb) {
        sb.append(BaseContext.getUserId());
    }
}

@Aspect
@Component
@Slf4j
public class RedisLimitAspect {
    @Autowired
    private RedissonClient redissonClient;

    @Autowired
    private ProceedingJoinPointUtil proceedingJoinPointUtil;

    @Pointcut("@annotation(com.cat.www.aop.limit.anno.RedisLimit)")
    private void pointCut() {
    }

    @Around("pointCut() && @annotation(redisLimit)")
    private Object around(ProceedingJoinPoint joinPoint, RedisLimit redisLimit) {
        Object generateKey = proceedingJoinPointUtil.getKey(joinPoint, redisLimit);
        //redis key
        String key = redisLimit.prefix() +generateKey.toString();
        //聲明一個限流器
        RRateLimiter rateLimiter = redissonClient.getRateLimiter(key);

        //設置速率，time秒中產生count個令牌
        rateLimiter.trySetRate(RateType.OVERALL, redisLimit.count(), redisLimit.time(), RateIntervalUnit.SECONDS);

        // 試圖獲取一個令牌，獲取到返回true
        boolean tryAcquire = rateLimiter.tryAcquire();
        if (!tryAcquire) {
            return new ResultData<>().FAILED().setResultIns("訪問過於頻繁");
        }
        Object obj = null;
        try {
            obj = joinPoint.proceed();
        } catch (Throwable e) {
            throw new RuntimeException();
        }

        return obj;
    }
}