基于guava RateLimiter及spring aop 注解实现的访问速率限制

1、pom引入

        com.google.guavaguava14.0.1org.springframework.bootspring-boot-starter-aop

2、注解类

@Documented
@Retention(RetentionPolicy.RUNTIME)
@Target({ ElementType.METHOD, ElementType.TYPE })
public @interface SpeedLimit {/*** 开启只对特定ip限制*/boolean forIP() default false;/*** 请求次数*/double num() default 1;/*** 时间（单位 s）*/int time() default 1;/*** 阻塞等待*/boolean block() default true;}

3、主要服务类

@Slf4j
public class RequestCacheService {// 缓存最大数private static final int MAXIMIZE = 1000;//private static final int DEFAULT_RATE = 1;/*** 创建缓存器*/private static final LoadingCache requestCaches = CacheBuilder.newBuilder().maximumSize(MAXIMIZE).expireAfterWrite(1, TimeUnit.MINUTES).removalListener((notification) -> {log.info(notification.getKey() + " 令牌桶被移除了，原因： " + notification.getCause());}).build(new CacheLoader() {@Overridepublic RateLimiter load(String key){return RateLimiter.create(DEFAULT_RATE);}});/*** 访问速率限制** @param key* @param num* @param time* @return* @throws ExecutionException*/public static RateLimiter rateLimiter(String key, double num, int time) throws ExecutionException {double rate = Double.parseDouble(String.format("%.2f", num / time));if(0 == rate){rate = 0.01;}RateLimiter rateLimiter = requestCaches.get(key);if(rate != rateLimiter.getRate()){rateLimiter.setRate(rate);}return rateLimiter;}/*** 尝试获取令牌** @param key* @param num* @param time* @param block* @return* @throws ExecutionException*/public static boolean tryAcquire(String key, double num, int time, boolean block) throws ExecutionException {RateLimiter rateLimiter = rateLimiter(key, num, time);if(block){rateLimiter.acquire(1);return true;}return rateLimiter.tryAcquire();}
}

4、切面实现类

@Component
@Aspect
public class RateLimitAspect {@Resourceprivate HttpServletRequest request;@Pointcut(" @annotation(com.example.bootdemo.aop.SpeedLimit) || " +"@within(com.example.bootdemo.aop.SpeedLimit)")public void pointcut() {}@Around("pointcut()")public Object around(ProceedingJoinPoint joinPoint) throws Throwable {String ipAddr = ServletUtil.getClientIP(request);MethodSignature signature = (MethodSignature) joinPoint.getSignature();SpeedLimit annotation = signature.getMethod().getAnnotation(SpeedLimit.class);if(null == annotation){annotation = joinPoint.getTarget().getClass().getAnnotation(SpeedLimit.class);}boolean enableLimit = checkEnableLimit(ipAddr, annotation);try{if(enableLimit){return speedLimit(joinPoint, getKey(signature), annotation);}return joinPoint.proceed();}catch (Throwable e){throw e;}}/*** 速度限制** @param joinPoint* @param key* @param annotation* @return* @throws Throwable*/private Object speedLimit(ProceedingJoinPoint joinPoint, String key, SpeedLimit annotation) throws Throwable {double num = annotation.num();int time = annotation.time();boolean block = annotation.block();try{if (RequestCacheService.tryAcquire(key, num, time, block)) {// 获得令牌（不限制访问）return joinPoint.proceed();} else {// 未获得令牌（限制访问）return Result.FAIL();}}catch (Throwable e){throw e;}}/*** 检测是否开启限制** @param ip* @param annotation*/private boolean checkEnableLimit(String ip, SpeedLimit annotation) {if(annotation.forIP()){List ips = getLimitIPs();if(null != ips && 0 != ips.size() && ips.contains(ip)){// 只对存在的ip进行限制return true;}return false;}// 默认限制所有return true;}/*** 获取key** @param signature* @return*/private String getKey(MethodSignature signature){return signature.getDeclaringTypeName() + "." + signature.getMethod().getName();}/*** 获取ip集合** @return*/private List getLimitIPs() {return Lists.newArrayList("127.0.0.1");}}

5、测试controller

@RestController
@RequestMapping("/limit")
public class LimitController {/*** 限制三秒内访问五次* * @return*/@RequestMapping("/test1")@SpeedLimit(num = 5, time = 3)public Result test1(){System.out.println("访问成功" + new SimpleDateFormat("HH:mm:ss").format(new Date()));return Result.OK();}
}

本文来自互联网用户投稿，文章观点仅代表作者本人，不代表本站立场，不承担相关法律责任。如若转载，请注明出处。 如若内容造成侵权/违法违规/事实不符，请点击【内容举报】进行投诉反馈！