Redisson多策略注解限流 - ToB企服应用市场:ToB评测及商务社交产业平台

/**
* 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;
}
}

复制代码