使用RateLimiter限流（AOP + 注解实现）_aopratelimiter 限流

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

引入依赖

        <!-- guava -->
        <dependency>
            <groupId>com.google.guava</groupId>
            <artifactId>guava</artifactId>
            <version>27.0.1-jre</version>
        </dependency>
        <!-- spring-boot-starter-aop -->
        <dependency>
            <groupId>org.springframework.boot</groupId>
            <artifactId>spring-boot-starter-aop</artifactId>
        </dependency>
1
2
3
4
5
6
7
8
9
10
11

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

• 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;

}
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24

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

package com.dfyang.ratelimiteraop.aop;

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

}
1
2
3
4
5
6
7
8
9
10

用法如下，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 "访问成功";
    }

}
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26

如果用户注解校验的应该很熟悉，只是这里的Success并没有错误信息，只作为用户成功获取令牌的标识。

创建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
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126

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

如果有错或有改进的地方，欢迎指正！