通过Redis+自定义注解实现接口限流策略_自定义注解+redis实现接口限流

1、前言

通过自定义注解+reids+lua实现，接口限流策略，其实质就是对redis的分布式锁的应用。

流程基本如下：

• 1、Controller接口的方法，实现自定义注解@RateLimiter。
• 2、自定义拦截RateLimiterHandlerInterceptor，拦截包含注解@RateLimiter的接口，进行验证。
• 3、为了保证在并发请求下的精确性，使用redis+lua脚本进行加锁。
• 4、如果窗口时间内请求没有达到上限则放行，如果达到了上限，则返回错误。

2、代码实现

2.1 自定义注解

@Target(ElementType.METHOD)
@Retention(RetentionPolicy.RUNTIME)
@Documented
public @interface IpRateLimiter {

    //一个IP下请求的并发限制

    int period() default 1;  //限流时间周期，默认 1S

    int count() default 10; //周期内限制次数，默认  10次

    boolean rateIP() default true; //默认限制IP，设置为false表示只限制接口请求次数
}
1
2
3
4
5
6
7
8
9
10
11
12
13

rateIP表示更加细致的控制，如果不需要限制ip，而是对接口设置一个统一的访问上限，则将rateIP设置为flase，比如：@RateLimiter(rateIP = false)

2.2 lua脚本配置

为了保证适应高并发，通过reids+lua脚本来实现加锁与解锁，达到一致性的目的。

lua脚本：
在resources目录下创建lua文件夹，在lua创建rateLimit.lua文件，内容如下：

local key = KEYS[1]
local limit = tonumber(KEYS[2])
local length = tonumber(KEYS[3])
--redis.log(redis.LOG_NOTICE,' length: '..length)
local current = redis.call('GET', key)
if current == false then
   --redis.log(redis.LOG_NOTICE,key..' is nil ')
   redis.call('SET', key,1)
   redis.call('EXPIRE',key,length)
   --redis.log(redis.LOG_NOTICE,' set expire end')
   return '1'
else
   --redis.log(redis.LOG_NOTICE,key..' value: '..current)
   local num_current = tonumber(current)
   if num_current+1 > limit then
       return '0'
   else
       redis.call('INCRBY',key,1)
       return '1'
   end
end
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21

其中key就是请求的接口信息，length为窗口时间，limit为窗口时间内允许访问的最大并发数量，比如：key=test，length=1，limit=10，就表示接口test在1秒内最大的并发为10，脚本中返回1表示成功，0表示失败。

rateLimit.lua逻辑如下：

• 1、根据key获取当前的请求数量；
• 2、如果为空，则将key的值设置为1，并且设置过期时间也就是窗口时间，并且返回1;
• 3、如果不为空，则判断当前请求数+1是否大于limit(最大并发数)，如果是则返回0，如果不是则将key的值自加1，并且返回1；

RedisLuaConfig：

创建Redis操作Lua的配置类，内容如下：

@Configuration
public class RedisLuaConfig {
    @Resource
    private StringRedisTemplate stringRedisTemplate;

    /**
     * @param keyList redis得Lua脚本中的key列表，我们把参数放在这里面传递
     * @return result 返回 1表示，正常，0表示限制访问
     */
    public String runLuaScript(List<String> keyList) {
        DefaultRedisScript<String> redisScript = new DefaultRedisScript<>();
        redisScript.setScriptSource(new ResourceScriptSource(new ClassPathResource("lua/rateLimit.lua")));
        redisScript.setResultType(String.class);
        String args = "none";
        try {
            return stringRedisTemplate.execute(redisScript, keyList);
        } catch (Exception e) {
            e.printStackTrace();
            return "0";
        }
    }
}
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22

2.3 拦截器配置

创建RateLimiterHandlerInterceptor类并且实现HandlerInterceptor接口，在该类中获取含有RateLimiter的请求，并且取出注解的配置，然后通过ip + "@" + httpMethod + "@" + path的形式生成key，再根据窗口时间和最大并发数，请求lua脚本，实现限流判断。

为了防止被暴力请求，默认情况下所有的接口进行限流，默认最大并发的10。

@Component
@Slf4j
public class RateLimiterHandlerInterceptor implements HandlerInterceptor {

    @Resource
    private RedisLuaConfig redisLuaConfig;

    private static final int DEFAULT_PERIOD = 1;
    private static final int DEFAULT_COUNT = 10;

    private static final String LIMITER_KEY = "limiter:";
    private static final String LIMITER_IP_KEY = "limiter-ip:";

    private static final String SUCCESS_CODE = "1";

    @Override
    public boolean preHandle(HttpServletRequest request, HttpServletResponse response, Object handler) throws Exception {

        Method method = ((HandlerMethod) handler).getMethod();

        String path = request.getServletPath();
        String httpMethod = request.getMethod(); //GET、PUT、DELETE、POST

        //获取当前请求IP
        String ip = IpUtils.getIpAddr(request);

        //如果没有开启开启注解,也进行默认限流操作
        int period = DEFAULT_PERIOD;  //限流时间周期，默认 1S
        int count = DEFAULT_COUNT; //周期内限制次数，默认  10次
        boolean rateIP = true;//默认限制IP，设置为false表示只限制接口请求次数

        IpRateLimiter ipRateLimiter = method.getAnnotation(IpRateLimiter.class);
        if (ipRateLimiter != null) {
            period = ipRateLimiter.period();
            count = ipRateLimiter.count();
            //设置当前key,对同一个ip下同一个请求进行限流操作
            rateIP = ipRateLimiter.rateIP();
        }

        //不对ip进行限制
        String key = LIMITER_KEY + httpMethod + "@" + path;
        if (rateIP) {
            //对ip进行次数限制
            key = LIMITER_IP_KEY + ip + "@" + httpMethod + "@" + path;
        }

        String res = "";
        try {
            List<String> keyList = new ArrayList<>();
            keyList.add(key);
            //表示时间周期内运行访问得次数
            keyList.add(String.valueOf(count)); //count
            keyList.add(String.valueOf(period));//period
            res = redisLuaConfig.runLuaScript(keyList);
        } catch (Exception e) {
			e.printStackTrace();
			throw new RuntimeException("你被限流了");
        }
        //正常执行RedisLua
        if (!SUCCESS_CODE.equals(res)) {
           log.error("你被限流了");
           throw new RuntimeException("你被限流了");
        }
        return true;
    }
}
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

3、测试

将Controller中的一个接口，加上RateLimiter注解，为了更好的验证限流，将请求的最大并发数设置为1

@RestController
@RequestMapping("test")
public class TestController {
    @GetMapping("/start")
    @RateLimiter(count = 1)
    public String start() {
        logger.info("成功访问接口");
        return "操作成功";
    }
}
1
2
3
4
5
6
7
8
9
10

通过浏览器快速刷新访问接口，就会触发限流，结果如下：