SpringBoot 自研分布式限流组件

发表于 阅读次数： 评论数： 本文字数： 945 阅读时长 ≈ 1 分钟

大纲

前言

本文将介绍在 SpringBoot 项目中，如何基于 Redis + Lua 脚本 + AOP + 反射 + 自定义注解自研分布式限流组件，且支持拔插式使用。由于篇幅有限，下面使用的是 Redis 单机服务，若是在生产环境，为了保证系统的可用性，建议部署 Redis 集群，并使用 Redisson 作为 Redis 的客户端，这里不再累述。

代码下载

完整的案例代码可以从 这里 下载得到。值得一提的是，本文的案例代码不仅适用于 SpringBoot 项目，理论上适用于任何 Spring 项目。

业务需求

公司内部使用了自研的 RPC 框架或者 Dubbo 来开发微服务项目，又或者为了安全性不允许项目引入过多的外部组件，导致无法使用 Spring Cloud 相关的限流组件（如 Sentinel），因此需要自研限流组件来满足业务需求。

  • 可配置
    • 在规定时间内，可以随意调整限流的时间和次数
    • 比如，支持设定 1 秒内最多允许访问 5 次接口，超过设定会启动限流功能，保护系统不过载
  • 可插拔
    • 可以按照促销活动、VIP 等级、方法使用频率等业务规则，要求 Controller 里面的业务方法有标识性的限流控制机制
      • 添加，方法自带限流功能
      • 不添加，方法没有限流功能
  • 可通用
    • 开发的自定义限流共用模块，可以给整个开发团队赋能公用
    • 不能和业务逻辑代码写死，支持独立出来，并可以配置
  • 高可用
    • 在高并发环境下，可以实时生效

解决方案

  • 通过自定义注解来实现业务解耦，可配置（在规定时间内，可以随意调整限流的时间和次数），可拔插使用，即一个注解就可以搞定限流的核心功能。
  • 底层使用 Redis + Lua 脚本实现，支持高并发，且满足事务的一致性与原子性要求。
  • 自定义 AOP 切面类，实现业务解耦。

案例代码

引入依赖

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
<parent>
    <groupId>org.springframework.boot</groupId>
    <artifactId>spring-boot-starter-parent</artifactId>
    <version>3.2.0</version>
    <relativePath/>
</parent>

<properties>
    <project.build.sourceEncoding>UTF-8</project.build.sourceEncoding>
    <maven.compiler.source>17</maven.compiler.source>
    <maven.compiler.target>17</maven.compiler.target>
    <hutool.version>5.8.27</hutool.version>
    <spring.cloud.version>2023.0.0</spring.cloud.version>
    <spring.cloud.alibaba.version>2023.0.0.0-RC1</spring.cloud.alibaba.version>
</properties>

<dependencyManagement>
    <dependencies>
        <dependency>
            <groupId>org.springframework.cloud</groupId>
            <artifactId>spring-cloud-dependencies</artifactId>
            <version>${spring.cloud.version}</version>
            <type>pom</type>
            <scope>import</scope>
        </dependency>
        <dependency>
            <groupId>com.alibaba.cloud</groupId>
            <artifactId>spring-cloud-alibaba-dependencies</artifactId>
            <version>${spring.cloud.alibaba.version}</version>
            <type>pom</type>
            <scope>import</scope>
        </dependency>
    </dependencies>
</dependencyManagement>

<dependencies>
    <!--Web-->
    <dependency>
        <groupId>org.springframework.boot</groupId>
        <artifactId>spring-boot-starter-web</artifactId>
    </dependency>
    <!--Aop-->
    <dependency>
        <groupId>org.springframework.boot</groupId>
        <artifactId>spring-boot-starter-aop</artifactId>
    </dependency>
    <!--Test-->
    <dependency>
        <groupId>org.springframework.boot</groupId>
        <artifactId>spring-boot-starter-test</artifactId>
        <scope>test</scope>
    </dependency>
    <!--Redis-->
    <dependency>
        <groupId>org.springframework.boot</groupId>
        <artifactId>spring-boot-starter-data-redis</artifactId>
    </dependency>
    <!--HuTool-->
    <dependency>
        <groupId>cn.hutool</groupId>
        <artifactId>hutool-all</artifactId>
        <version>${hutool.version}</version>
    </dependency>
    <!--Lombok-->
    <dependency>
        <groupId>org.projectlombok</groupId>
        <artifactId>lombok</artifactId>
        <optional>true</optional>
    </dependency>
</dependencies>

配置信息

1
2
3
4
5
6
7
8
9
10
11
12
spring:
  data:
    redis:
      database: 0
      host: 127.0.0.1
      port: 6379
      lettuce:
        pool:
          max-active: 8
          max-wait: -1ms
          max-idle: 8
          min-idle: 0

Lua 脚本

创建 Lua 脚本文件 RateLimiter.lua，并存放在项目的 /src/main/resources 目录下

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
-- 获取KEY，针对特定的接口进行限流
local key = KEYS[1]

-- 获取在限流注解上标注的限流次数
local limit = tonumber(ARGV[1])

-- 获取当前已访问接口的次数
local curentLimit = tonumber(redis.call('get', key) or "0")

-- 超过限流次数直接返回-1，否则递增访问次数
if curentLimit + 1 > limit
then return -1
else
    -- 自增长 1
    redis.call('INCRBY', key, 1)
    -- 设置过期时间
    redis.call('EXPIRE', key, ARGV[2])
    return curentLimit + 1
end

核心代码

  • 自定义异常类
1
2
3
4
5
6
7
8
9
10
11
12
13
14
public class RedisLimitException extends RuntimeException {

    public RedisLimitException() {
    }

    public RedisLimitException(String message) {
        super(message);
    }

    public RedisLimitException(String message, Throwable cause) {
        super(message, cause);
    }

}
  • Redis 配置类
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
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.data.redis.connection.lettuce.LettuceConnectionFactory;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.data.redis.serializer.GenericJackson2JsonRedisSerializer;
import org.springframework.data.redis.serializer.StringRedisSerializer;

@Configuration
public class RedisConfig {

    @Bean
    public RedisTemplate<String, Object> redisTemplate(LettuceConnectionFactory lettuceConnectionFactory) {
        RedisTemplate<String, Object> redisTemplate = new RedisTemplate<>();
        redisTemplate.setConnectionFactory(lettuceConnectionFactory);
        //设置key序列化方式String
        redisTemplate.setKeySerializer(new StringRedisSerializer());
        //设置value的序列化方式Json
        redisTemplate.setValueSerializer(new GenericJackson2JsonRedisSerializer());
        //设置key序列化方式String
        redisTemplate.setHashKeySerializer(new StringRedisSerializer());
        //设置value的序列化方式Json
        redisTemplate.setHashValueSerializer(new GenericJackson2JsonRedisSerializer());
        redisTemplate.afterPropertiesSet();
        return redisTemplate;
    }

}
  • 自定义注解
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
import java.lang.annotation.Documented;
import java.lang.annotation.ElementType;
import java.lang.annotation.Retention;
import java.lang.annotation.RetentionPolicy;
import java.lang.annotation.Target;

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

    /**
     * 资源的唯一 Key
     * <p> 作用：不同的接口，不同的流量控制
     */
    String key() default "";

    /**
     * 允许访问的最大次数
     */
    long limit() default 5;

    /**
     * 过期时间，也可以理解为单位时间或者滑动窗口时间
     * <p> 单位为秒，默认值为 60
     */
    long expire() default 60;

    /**
     * 限流的提示语
     */
    String msg() default "当前接口的负载过高，请稍后再试";

}
  • AOP 切面类
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
import cn.hutool.core.util.StrUtil;
import com.clay.limit.annotations.RedisLimitAnnotation;
import com.clay.limit.exception.RedisLimitException;
import jakarta.annotation.PostConstruct;
import jakarta.annotation.Resource;
import lombok.extern.slf4j.Slf4j;
import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.annotation.Around;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.reflect.MethodSignature;
import org.springframework.core.io.ClassPathResource;
import org.springframework.data.redis.core.StringRedisTemplate;
import org.springframework.data.redis.core.script.DefaultRedisScript;
import org.springframework.scripting.support.ResourceScriptSource;
import org.springframework.stereotype.Component;

import java.lang.reflect.Method;
import java.util.Collections;
import java.util.List;

@Slf4j
@Aspect
@Component
public class RedisLimitAspect {

    Object result = null;

    @Resource
    private StringRedisTemplate stringRedisTemplate;

    private DefaultRedisScript<Long> redisLuaScript;

    @PostConstruct
    public void init() {
        // 加载 Lua 脚本
        redisLuaScript = new DefaultRedisScript<>();
        redisLuaScript.setResultType(Long.class);
        redisLuaScript.setScriptSource(new ResourceScriptSource(new ClassPathResource("RateLimiter.lua")));
    }

    @Around("@annotation(com.clay.limit.annotations.RedisLimitAnnotation)")
    public Object around(ProceedingJoinPoint joinPoint) {
        System.out.println("---------@Around before");

        // 通过反射获取目标方法
        MethodSignature signature = (MethodSignature) joinPoint.getSignature();
        Method method = signature.getMethod();

        // 通过反射获取目标方法上的RedisLimitAnnotation注解，如果存在，则说明需要限流
        RedisLimitAnnotation redisLimitAnnotation = method.getAnnotation(RedisLimitAnnotation.class);

        if (redisLimitAnnotation != null) {
            // 获取Redis中的key
            String key = redisLimitAnnotation.key();
            String className = method.getDeclaringClass().getName();
            String methodName = method.getName();

            if (StrUtil.isBlank(key)) {
                throw new RedisLimitException("it's danger, limit key cannot be empty");
            }

            String limitInfo = key + "\t" + className + "." + methodName + "()";
            log.info(limitInfo);

            long limit = redisLimitAnnotation.limit();
            long expire = redisLimitAnnotation.expire();
            List<String> keys = Collections.singletonList(key);

            // 执行 Lua 脚本
            Long count =
                stringRedisTemplate.execute(redisLuaScript, keys, String.valueOf(limit), String.valueOf(expire));

            if (count != null && count == -1) {
                log.warn("启动限流功能, Key 为 " + key);
                return redisLimitAnnotation.msg();
            }

            log.info("Access try count is " + count + ", limit key is " + key);
        }

        try {
            // 执行目标方法
            result = joinPoint.proceed();
        } catch (Throwable e) {
            throw new RuntimeException(e);
        }

        System.out.println("---------@Around after");

        return result;
    }

}
  • 主启动类，添加了 @EnableAspectJAutoProxy 注解
1
2
3
4
5
6
7
8
9
@SpringBootApplication
@EnableAspectJAutoProxy
public class MainApplication {

    public static void main(String[] args) {
        SpringApplication.run(MainApplication.class, args);
    }

}

测试代码

  • 控制器类
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
import cn.hutool.core.util.IdUtil;
import com.clay.limit.annotations.RedisLimitAnnotation;
import lombok.extern.slf4j.Slf4j;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RestController;

@Slf4j
@RestController
@RequestMapping("/order")
public class OrderController {

    @GetMapping("/limit")
    @RedisLimitAnnotation(key = "orderLimit", limit = 5, expire = 10, msg = "当前排队人数较多，请稍后再试！")
    public String limit() {
        return "正常业务返回，订单流水：" + IdUtil.fastUUID();
    }

}
  • 当正常调用接口时，控制台输出的日志信息如下：
1
2
[http-nio-8080-exec-4] INFO  c.clay.limit.aspect.RedisLimitAspect - orderLimit	com.clay.limit.controller.OrderController.limit()
[http-nio-8080-exec-4] INFO  c.clay.limit.aspect.RedisLimitAspect - Access try count is 1, limit key is orderLimit
  • 当频繁调用接口时，控制台输出的日志信息如下：
1
2
[http-nio-8080-exec-9] INFO  c.clay.limit.aspect.RedisLimitAspect - orderLimit	com.clay.limit.controller.OrderController.limit()
[http-nio-8080-exec-9] WARN  c.clay.limit.aspect.RedisLimitAspect - 启动限流功能, Key 为 orderLimit

参考资料