若依源码分析--@RateLimiter

若依作为最近非常火的脚手架，分析它的源码，不仅可以更好的使用它，在出错时及时定位，也可以在需要个性化功能时轻车熟路的修改它以满足我们自己的需求，同时也可以学习人家解决问题的思路，提升自己的技术水平

若依提供了很多实用且不花哨的注解，本文记录了其中的一个注解@RateLimiter--限流注解的实现步骤

版本说明

以下源码内容是基于RuoYi-Vue-3.8.2版本，即前后端分离版本

主要思想

标注了@RateLimiter注解的方法，在执行前调用lua脚本，把一段时间内的访问次数存入redis并返回，判断返回值是否大于设定的阈值，大于则抛出异常，由全局异常处理器处理

具体步骤

1. 注解

我们先来看一看@RateLimiter注解，在src/main/java/com/ruoyi/common/annotation包下

@Target(ElementType.METHOD)
@Retention(RetentionPolicy.RUNTIME)
@Documented
public @interface RateLimiter 
{
    // 限流key
    public String key() default Constants.RATE_LIMIT_KEY;

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

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

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

一个作用在方法上的注解，有四个属性

key：存储在redis里用到的key
time：限流时间，相当于redis里的有效期
count：限流次数
limitType: 限流类型，点开枚举发现有默认和IP两种限流方式，这两种方式的实现只是存储在redis里的key不同

2. 切面

我们来看一看@RateLimiter这个注解的切面RateLimiterAspect.java，在src/main/java/com/ruoyi/framework/aspectj包里

@Aspect
@Component
public class RateLimiterAspect 
{
    private static final Logger log = LoggerFactory.getLogger(RateLimiterAspect.class);

    private RedisTemplate<Object, Object> redisTemplate;

    private RedisScript<Long> limitScript;

    @Autowired
    public void setRedisTemplate1(RedisTemplate<Object, Object> redisTemplate)
    {
        this.redisTemplate = redisTemplate;
    }

    @Autowired
    public void setLimitScript(RedisScript<Long> limitScript)
    {
        this.limitScript = limitScript;
    }

    @Before("@annotation(rateLimiter)")
    public void doBefore(JoinPoint point, RateLimiter rateLimiter) throws Throwable
    {
        String key = rateLimiter.key();
        int time = rateLimiter.time();
        int count = rateLimiter.count();

        String combineKey = getCombineKey(rateLimiter, point);
        List<Object> keys = Collections.singletonList(combineKey);
        try
        {
            // 调用lua脚本，传入三个参数
            Long number = redisTemplate.execute(limitScript, keys, count, time);
            if (StringUtils.isNull(number) || number.intValue() > count)
            {
                throw new ServiceException("访问过于频繁，请稍候再试");
            }
            log.info("限制请求'{}',当前请求'{}',缓存key'{}'", count, number.intValue(), key);
        }
        catch (ServiceException e)
        {
            throw e;
        }
        catch (Exception e)
        {
            throw new RuntimeException("服务器限流异常，请稍候再试");
        }
    }

    public String getCombineKey(RateLimiter rateLimiter, JoinPoint point)
    {
        // 获取注解中的key值
        StringBuffer stringBuffer = new StringBuffer(rateLimiter.key());
        // 判断限流类型，如果是IP限流，就在key后添加上IP（若依自己写了一个获取ip的方法类，大家可以自行查看）
        if (rateLimiter.limitType() == LimitType.IP)
        {
            stringBuffer.append(IpUtils.getIpAddr(ServletUtils.getRequest())).append("-");
        }
        // 获取方法
        MethodSignature signature = (MethodSignature) point.getSignature();
        Method method = signature.getMethod();
        // 获取类
        Class<?> targetClass = method.getDeclaringClass();
        // key中添加方法名-类名
        stringBuffer.append(targetClass.getName()).append("-").append(method.getName());
        return stringBuffer.toString();
    }
}

简单说明一下这个切面类：

使用了set的方式注入了RedisTemplate和RedisScript，RedisTemplate大家都很熟悉，RedisScript是用于加载和执行lua脚本的
定义了一个前置通知（废话，限流肯定是前置），通过getCombineKey方法获取应该存入redis中的key，getCombineKey方法每一步我都做了注解
将key、time、count作为参数传入lua脚本，执行脚本，判断返回值为空或者或者返回值大于设定的count，抛出异常，由全局异常处理器处理，方法不再往下执行，达到了限流的效果

3. lua脚本

最后，我们来看一看若依是怎么写lua脚本的，在脚本在redis的配置类RedisConfig.java里，该类在src/main/java/com/ruoyi/framework/config包下

@Configuration
@EnableCaching
public class RedisConfig extends CachingConfigurerSupport
{
    ……
    
    @Bean
    public DefaultRedisScript<Long> limitScript()
    {
        // 泛型是返回值的类型
        DefaultRedisScript<Long> redisScript = new DefaultRedisScript<>();
        // 设置脚本
        redisScript.setScriptText(limitScriptText());
        // 设置返回值类型
        redisScript.setResultType(Long.class);
        return redisScript;
    }

    /**
     * 限流脚本
     */
    private String limitScriptText()
    {
        return "local key = KEYS[1]\n" +
                "local count = tonumber(ARGV[1])\n" +
                "local time = tonumber(ARGV[2])\n" +
                "local current = redis.call('get', key);\n" +
                "if current and tonumber(current) > count then\n" +
                "    return tonumber(current);\n" +
                "end\n" +
                "current = redis.call('incr', key)\n" +
                "if tonumber(current) == 1 then\n" +
                "    redis.call('expire', key, time)\n" +
                "end\n" +
                "return tonumber(current);";
    }
}

我们主要看下lua脚本：

接收3个变量：key，阈值count，过期时间time
调用get(key)方法获取key中的值current，如果这个key存在并且current大于count，返回current
调用redis的自增函数赋值给current，当current=1时（即第一次访问该接口），调用redis的设置过期时间函数给当前key设置过期时间
返回current

使用lua脚本可以在并发的情况下更好的满足原子性，只是我不太明白若依为什么不把脚本文件单独拿出来写在resources文件夹下，这样阅读和维护都会更加方便。总之，这就是若依限流注解的全部内容

总结

标注了@RateLimiter注解的方法，在执行方法前调用lua脚本，把自己的类名+方法名当做key传入，判断返回值是否大于设定的阈值，大于则抛出异常不再向下执行，异常由全局异常处理器处理。