使用RateLimiter限流(AOP + 注解实现)

注意:这里不会对RateLimiter做介绍

引入依赖

        
        <dependency>
            <groupId>com.google.guavagroupId>
            <artifactId>guavaartifactId>
            <version>27.0.1-jreversion>
        dependency>
        
        <dependency>
            <groupId>org.springframework.bootgroupId>
            <artifactId>spring-boot-starter-aopartifactId>
        dependency>

创建限流注解,加在需要进行限流的方法上
三个参数的基本含义为:

  • rate:每秒生成多少个令牌
  • timeout:申请令牌超时时间
  • timeunit:超时时间单位
package com.dfyang.ratelimiteraop.aop;

import java.lang.annotation.ElementType;
import java.lang.annotation.Retention;
import java.lang.annotation.RetentionPolicy;
import java.lang.annotation.Target;
import java.util.concurrent.TimeUnit;

/**
 * @Auther: 55411
 * @Date: 2019/6/24 21:00
 * @Description: 平滑突发限流注解
 */
@Target(ElementType.METHOD)
@Retention(RetentionPolicy.RUNTIME)
public @interface SmoothBurstyLimit {

    double rate();

    long timeout();

    TimeUnit timeunit() default TimeUnit.SECONDS;

}

创建Success类,在需要限流的方法上添加Success参数(必须),以便在客户端失败获取令牌时进行降级处理。

package com.dfyang.ratelimiteraop.aop;

/**
 * @Auther: 55411
 * @Date: 2019/6/24 21:43
 * @Description: 作为申请令牌成功的标识
 */
public class Success {

}

用法如下,Spring会自动为我们初始化Success对象,因此当客户端失败获取令牌时,我们只需要将Success对象设置为null即可

package com.dfyang.ratelimiteraop.controller;

import com.dfyang.ratelimiteraop.aop.SmoothBurstyLimit;
import com.dfyang.ratelimiteraop.aop.Success;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RestController;

/**
 * @Auther: 55411
 * @Date: 2019/6/24 20:37
 * @Description: 测试Controller
 */
@RestController
public class LimitController {

    @GetMapping("/limit")
    @SmoothBurstyLimit(rate = 2, timeout = 2)
    public String limit(Success success) {
        if (success == null) {
            // TODO 进行降级处理
            return "降级";
        }
        return "访问成功";
    }

}

如果用户注解校验的应该很熟悉,只是这里的Success并没有错误信息,只作为用户成功获取令牌的标识。
使用RateLimiter限流(AOP + 注解实现)_第1张图片
创建Aop
1)对加了@SmoothBurstyLimit注解的类进行处理。
2)首先获取注解上的参数,对rate进行合法校验。
3)获取该限流方法对应的RateLimiter,这里使用了ConcurrentHashMap保证线程安全
4)尝试获取令牌
5)失败将方法上的Success参数设为null

package com.dfyang.ratelimiteraop.aop;

import com.google.common.util.concurrent.RateLimiter;
import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.annotation.Around;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.annotation.Pointcut;
import org.aspectj.lang.reflect.MethodSignature;
import org.springframework.stereotype.Component;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.TimeUnit;

/**
 * @Auther: 55411
 * @Date: 2019/6/24 21:06
 * @Description: 平滑突发限流注解AOP
 */
@Aspect
@Component
public class SmoothBurstyAop {

    private Map<String, RateLimiter> rateLimiterMap = new ConcurrentHashMap<>();

    @Pointcut("execution(public * com.dfyang.ratelimiteraop.controller.*.*(..))")
    public void pointCut(){}

    @Around("pointCut()")
    public Object around(ProceedingJoinPoint joinPoint) throws Throwable {
        MethodSignature signature = (MethodSignature) joinPoint.getSignature();
        // 获取方法上的@SmoothBurstyLimit注解
        SmoothBurstyLimit smoothBurstyLimit = signature.getMethod().getDeclaredAnnotation(SmoothBurstyLimit.class);
        Object[] args = null;
        if (smoothBurstyLimit != null) {
            // 获取注解的rate参数,并进行合法判断
            double rate = smoothBurstyLimit.rate();
            checkRateLegality(rate);
            // 获取注解的timeout参数,并与0进行比较
            long timeout = Math.max(smoothBurstyLimit.timeout(), 0L);
            // 获取注解的超时时间单位
            TimeUnit timeunit = smoothBurstyLimit.timeunit();
            // 根据类名 + 方法名作为Map的key值
            Class clazz = joinPoint.getTarget().getClass();
            String rateLimiterKey = clazz.getSimpleName() + signature.getName();
            // 从map中获取RateLimiter对象,如果没有,创建并添加的map中
            RateLimiter rateLimiter = getRateLimiter(rateLimiterKey, clazz, rate);
            // 获取令牌
            boolean success = rateLimiter.tryAcquire(timeout, timeunit);
            // 如果失败,将限流方法上的Success注解设置为null
            if (!success) {
                args = joinPoint.getArgs();
                handleFailure(args);
                // 将参数回填
                return joinPoint.proceed(args);
            }
        }
        return joinPoint.proceed();
    }

    /**
     * 从map中获取RateLimiter对象,如果没有,创建并添加的map中
     * @param key 键
     * @param clazz 类锁
     * @param rate 生成令牌的速率
     * @return RateLimiter对象
     */
    private RateLimiter getRateLimiter(String key, Class clazz, double rate) {
        RateLimiter rateLimiter = rateLimiterMap.get(key);
        if (rateLimiter == null) {
            rateLimiter = rateLimiter.create(rate);
            rateLimiterMap.put(key, rateLimiter);
        }
        return rateLimiter;
    }

    /**
     * 从map中获取RateLimiter对象,如果没有,创建并添加的map中
     * (使用Controller类作为锁,加锁可以保证每个key对应的RateLimiter只创建一次)
     * (由于ConcurrentHashMap只会对value进行覆盖,所以这里加锁意义不大)
     * @param key 键
     * @param clazz 类锁
     * @param rate 生成令牌的速率
     * @return RateLimiter对象
     */
    private RateLimiter getRateLimiterSync(String key, Class clazz, double rate) {
        RateLimiter rateLimiter = rateLimiterMap.get(key);
        if (rateLimiter == null) {
            synchronized (clazz) {
                rateLimiter = rateLimiterMap.get(key);
                if (rateLimiter == null) {
                    rateLimiter = rateLimiter.create(rate);
                    rateLimiterMap.put(key, rateLimiter);
                }
            }
        }
        return rateLimiter;
    }

    /**
     * 获取令牌失败时进行处理处理
     * @param args 方法中的参数
     * @throws IllegalArgumentException 非法参数异常
     */
    private void handleFailure(Object[] args) throws IllegalArgumentException {
        boolean canLimit = false;
        for (int i = 0,length = args.length; i < length; i++) {
            if (args[i] instanceof Success) {
                args[i] = null;
                canLimit = true;
                break;
            }
        }
        if (!canLimit)
            throw new IllegalArgumentException("无法完成限流,原因是限流方法上未加Success对象");
    }

    /**
     * 如果@SmoothBurstyLimit注解rate值错误,抛出异常
     * @param rate 生成令牌的速率
     * @throws IllegalAccessException 非法参数异常
     */
    private void checkRateLegality(double rate) throws IllegalArgumentException {
        if (rate <= 0)
            throw new IllegalArgumentException("@SmoothBurstyLimit注解rate值错误,rate = " + rate);
    }
}

改进之处——
1)可以在@SmoothBurstyLimit注解name属性,作为ConcurrentHashMap的key值,这样一个RateLimiter可以对多个方法进行限流。
2)在失败获取令牌时这里是将Success设置为null,也可以在Aop里面直接对其进行处理。
3)这里使用的guava的平滑突发限流(还有一种平滑预热限流),但这里突发没有显示出来,可以给@SmoothBurstyLimit注解增加一个number参数表示一次申请多少个令牌

如果有错或有改进的地方,欢迎指正!

你可能感兴趣的:(使用RateLimiter限流(AOP + 注解实现))