基于Spring AOP与Redisson的令牌桶限流注解实践

1. 什么是限流

举个例子🌰

整体限流：比如说一个系统只有1万QPS，一下子来了2万，系统就会出现问题，所以要将后到达的1万请求给限制住，就把流量给限制在1万以内了，被限流的用户可以给他返回一个”系统繁忙请稍后重试“之类，请重试的提示，这样起码可以保住1万的请求能正常返回。

局部限流：也可以叫接口限流，比如说用户发送短信验证码的场景，因为短信服务是收费，不可能让用户可以无限发，要根据用户手机号或者ip进行限制，可以设置60秒内可以发送2次，超过两次就提示”操作频繁请稍后重试“的提示。

2. 令牌桶算法

令牌桶算法是一种在网络流量整形和速率限制中广泛使用的算法。其基本原理是通过一个虚拟的“桶”来控制数据的传输速率，桶中存放着一定数量的令牌，每个令牌代表了一个单位的数据传输权限。

基本概念

• 令牌桶：一个虚拟的容器，用于存放令牌。桶的大小固定，表示桶中最多可以存放的令牌数量。
• 令牌：每个令牌代表了一个单位的数据传输权限，通常代表一个字节或数据包。
• 令牌生成速率：系统以固定的速率向桶中添加令牌，这个速率决定了长期平均传输速率。

工作流程

令牌桶算法的工作流程大致可以分为以下几个步骤：

1. 令牌生成：系统按照设定的速率（如每秒生成一定数量的令牌）周期性地向桶中添加令牌。如果桶已满，则新生成的令牌会被丢弃或拒绝。
2. 请求处理：当一个数据包或请求到达时，它需要从桶中取出一个令牌才能被处理。如果桶中有足够的令牌，请求可以立即被处理；如果桶中令牌不足，则请求必须等待，直到桶中再次有令牌可用。
3. 突发传输：由于桶可以存放一定数量的令牌，系统可以在短时间内处理等于桶容量大小的突发流量，而不会因为短暂的流量高峰而完全阻塞。
   

Redisson 框架已经基于 Redis 内置了令牌桶算法的实现，因此用户无需自行定义，可以直接利用 Redisson 提供的这一功能来编写相关代码。

3. 限流设计

基于上面的场景与算法，我们可以大致想想这个限流的函数应该怎么实现，首先得先知道是整体限流还是局部限流吧，那局部限流又得知道根据什么来限流，是实体ID/用户IP还是其他的，所以可以确定该函数的第一个参数： 限流的标志 key，有了限流标志后，还需要设置限流的规则，60秒内只能有两次操作，1秒内只允许操作1万次，所以还得有两个参数，第一是参数是限流的时间 time 第二个是限流的次数 count，当然被限流的用户要给他返回友好的提示，所有还有message参数，而这种通用的操作可以使用 AOP 来操作。

4. 代码实现

下面将基于自定义注解、Spring AOP、Redisson 来实现限流的功能，其中的代码示例将展示如何将这些组件整合起来，实现一个开箱即用的限流解决方案

4.1 引入依赖编写配置

在 pom 文件添加以下标签，引入 redis、redisson 和 hutool 的依赖

<!--redis-->
<dependency>
    <groupId>org.redisson</groupId>
    <artifactId>redisson-spring-boot-starter</artifactId>
</dependency>

<!--redisson-->
<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>5.8.20</version>
</dependency>

<!--lombok-->
<dependency>
    <groupId>org.projectlombok</groupId>
    <artifactId>lombok</artifactId>
</dependency>

在 yml/properties 文件中填写 redis 相关的配置

spring:
  data:
    redis:
      host: 127.0.0.1
      password: # 有设置密码就填写
      port: 6379
      database: 1

4.2 定义限流类型枚举

/**
 * 限流类型
 */

public enum LimitType {
    /**
     * 默认策略全局限流
     */
    DEFAULT,

    /**
     * 根据请求者IP进行限流
     */
    IP,
}

4.3 定义注解

import java.lang.annotation.*;

/**
 * 限流注解
 */
@Target(ElementType.METHOD)
@Retention(RetentionPolicy.RUNTIME)
@Documented
public @interface RateLimiter {
    /**
     * 限流key,支持使用Spring el表达式来动态获取方法上的参数值
     * 格式类似于  #code.id #{#code}
     */
    String key() default "";

    /**
     * 限流时间,单位秒
     */
    int time() default 60;

    /**
     * 限流次数
     */
    int count() default 100;

    /**
     * 限流类型
     */
    LimitType limitType() default LimitType.DEFAULT;

    /**
     * 提示消息
     */
    String message() default "请求频繁,请稍后重试";
}

4.4 定义切面

• @Before注解：处理请求前执行，此时未到达controller

import cn.hutool.core.util.ArrayUtil;
import cn.hutool.core.util.StrUtil;
import cn.hutool.extra.servlet.JakartaServletUtil;
import jakarta.servlet.http.HttpServletRequest;
import org.aspectj.lang.JoinPoint;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.annotation.Before;
import org.aspectj.lang.reflect.MethodSignature;
import org.redisson.api.RRateLimiter;
import org.redisson.api.RateIntervalUnit;
import org.redisson.api.RateType;
import org.redisson.api.RedissonClient;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.core.DefaultParameterNameDiscoverer;
import org.springframework.core.ParameterNameDiscoverer;
import org.springframework.expression.EvaluationContext;
import org.springframework.expression.Expression;
import org.springframework.expression.ExpressionParser;
import org.springframework.expression.ParserContext;
import org.springframework.expression.common.TemplateParserContext;
import org.springframework.expression.spel.standard.SpelExpressionParser;
import org.springframework.expression.spel.support.StandardEvaluationContext;
import org.springframework.web.context.request.RequestAttributes;
import org.springframework.web.context.request.RequestContextHolder;
import org.springframework.web.context.request.ServletRequestAttributes;

import java.lang.reflect.Method;

/**
 * 限流处理
 */
@Aspect
public class RateLimiterAspect {

    @Autowired
    private RedissonClient redissonClient;

    /**
     * 定义spel表达式解析器
     */
    private final ExpressionParser parser = new SpelExpressionParser();
    /**
     * 定义spel解析模版
     */
    private final ParserContext parserContext = new TemplateParserContext();
    /**
     * 定义spel上下文对象进行解析
     */
    private final EvaluationContext context = new StandardEvaluationContext();
    /**
     * 方法参数解析器
     */
    private final ParameterNameDiscoverer pnd = new DefaultParameterNameDiscoverer();

    @Before("@annotation(rateLimiter)")
    public void doBefore(JoinPoint point, RateLimiter rateLimiter) throws Throwable {
        int time = rateLimiter.time();
        int count = rateLimiter.count();
        String combineKey = getCombineKey(rateLimiter, point);
        try {
            RateType rateType = RateType.OVERALL;
            if (rateLimiter.limitType() == LimitType.CLUSTER) {
                rateType = RateType.PER_CLIENT;
            }
            long number = rateLimiter(combineKey, rateType, count, time);
            if (number == -1) {
                throw new BusinessException(rateLimiter.message());
            }
        } catch (Exception e) {
            if (e instanceof BusinessException) {
                throw e;
            } else {
                throw new RuntimeException("服务器限流异常，请稍候再试");
            }
        }
    }

    /**
     * 获取限流的 key
     */
    public String getCombineKey(RateLimiter rateLimiter, JoinPoint point) {
        String key = rateLimiter.key();
        // 获取方法(通过方法签名来获取)
        MethodSignature signature = (MethodSignature) point.getSignature();
        Method method = signature.getMethod();
        // 判断是否是spel格式
        if (StrUtil.containsAny(key, "#")) {
            // 获取参数值
            Object[] args = point.getArgs();
            // 获取方法上参数的名称
            String[] parameterNames = pnd.getParameterNames(method);
            if (ArrayUtil.isEmpty(parameterNames)) {
                throw new BusinessException("限流key解析异常!请联系管理员!");
            }
            for (int i = 0; i < parameterNames.length; i++) {
                context.setVariable(parameterNames[i], args[i]);
            }
            // 解析返回给key
            try {
                Expression expression;
                if (StrUtil.startWith(key, parserContext.getExpressionPrefix())
                        && StrUtil.endWith(key, parserContext.getExpressionSuffix())) {
                    expression = parser.parseExpression(key, parserContext);
                } else {
                    expression = parser.parseExpression(key);
                }
                key = expression.getValue(context, String.class) + ":";
            } catch (Exception e) {
                throw new BusinessException("限流key解析异常!请联系管理员!");
            }
        }
        StringBuilder stringBuffer = new StringBuilder("rate_limit:");
        // 获取请求的地址
        ServletRequestAttributes attributes =(ServletRequestAttributes) RequestContextHolder.getRequestAttributes();
        HttpServletRequest request = attributes.getRequest();
        stringBuffer.append(request.getRequestURI()).append(":");
        if (rateLimiter.limitType() == LimitType.IP) {
            // 获取请求的ip
            String ip = JakartaServletUtil.getClientIP(request);
            stringBuffer.append(ip).append(":");
        }
        return stringBuffer.append(key).toString();
    }

    /**
     * 限流
     *
     * @param key          限流key
     * @param rateType     限流类型
     * @param rate         速率
     * @param rateInterval 速率间隔
     * @return -1 表示失败
     */
    public long rateLimiter(String key, RateType rateType, int rate, int rateInterval) {
        RRateLimiter rateLimiter = redissonClient.getRateLimiter(key);
        rateLimiter.trySetRate(rateType, rate, rateInterval, RateIntervalUnit.SECONDS);
        if (rateLimiter.tryAcquire()) {
            return rateLimiter.availablePermits();
        } else {
            return -1L;
        }
    }
}

4.5 捕获异常

import lombok.AllArgsConstructor;
import lombok.Data;
import lombok.EqualsAndHashCode;
import lombok.NoArgsConstructor;

/**
 * 业务异常
 */
@Data
@EqualsAndHashCode(callSuper = true)
@NoArgsConstructor
@AllArgsConstructor
public final class BusinessException extends RuntimeException {

    
    private static final long serialVersionUID = 1L;

    /**
     * 错误码
     */
    private Integer code;

    /**
     * 错误提示
     */
    private String message;

    /**
     * 错误明细，内部调试错误
     */
    private String detailMessage;

    public BusinessException(String message) {
        this.message = message;
    }

    public BusinessException(String message, Integer code) {
        this.message = message;
        this.code = code;
    }

    public String getDetailMessage() {
        return detailMessage;
    }

    @Override
    public String getMessage() {
        return message;
    }

    public Integer getCode() {
        return code;
    }

    public BusinessException setMessage(String message) {
        this.message = message;
        return this;
    }

    public BusinessException setDetailMessage(String detailMessage) {
        this.detailMessage = detailMessage;
        return this;
    }
}

/**
 * 全局异常处理器
 */
@Slf4j
@RestControllerAdvice
public class GlobalExceptionHandler {

    /**
     * 业务异常
     */
    @ExceptionHandler(BusinessException.class)
    public ResponseResult<Void> handleServiceException(BusinessException e, HttpServletRequest request) {
        log.error(e.getMessage());
        Integer code = e.getCode();
        return ObjectUtil.isNotNull(code) ? ResponseResult.fail(code, e.getMessage()) : ResponseResult.fail(e.getMessage());
    }
}

4.6 使用限流注解

150秒内调用三次接口返回的信息：