简单限流器封装

简单限流器封装

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

下面基于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;
	}
}

使用

测试

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

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

有兴趣 参考 Guava， redisson https://github.com/redisson/redisson/wiki/6.-分布式对象#612-限流器ratelimiter （令牌桶）