lua脚本+redis实现限流

定义注解 @RateLimiter

@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;
}
 
 
public enum LimitType
{
    /**
     * 默认策略全局限流
     */
    DEFAULT,
 
    /**
     * 根据请求者IP进行限流
     */
    IP
}

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

key：存储在redis里用到的key
time：限流时间，相当于redis里的有效期
count：限流次数
limitType: 限流类型，点开枚举发现有默认和IP两种限流方式，这两种方式的实现只是存储在redis里的key不同
2. 切面
我们来看一看@RateLimiter这个注解的切面RateLimiterAspect.java

@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

@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传入，判断返回值是否大于设定的阈值，大于则抛出异常不再向下执行，异常由全局异常处理器处理。