Springboot如何实现限流

项目中我们有很多地方需要限流  比如接口限流 又分为全局接口限流 或者只针对每个用户的接口限流  可以根据不同的业务情况选择不同的方式即可

本次为了更加贴近项目实战我们采用的是redis配合lua脚本 来保证操作的原子性 

而通过注解和Spring aop 能使得我们的操作更加便捷 需要限流的接口加上注解即可

首先声明一个注解

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

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

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

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

从这个接口中可以去设置一些自定义参数  比如多少秒内可以执行多少次操作，还有具体是什么类型的限流 全局还是只针对每个用户的限流

@Aspect
@Component
public class RateLimiterAspect
{
    private RedisTemplate<Object, Object> redisTemplate;

    private RedisScript<Long> limitScript; 

    @Autowired
    public void setRedisTemplate(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
    {
        int time = rateLimiter.time();
        int count = rateLimiter.count();

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

    public String getCombineKey(RateLimiter rateLimiter, JoinPoint point)
    {
        StringBuffer stringBuffer = new StringBuffer(rateLimiter.key());
        if (rateLimiter.limitType() == LimitType.IP)
        {
     // 这里我们可以通过 实现一个Threadlocal 来获取当前登录用户的userid
            stringBuffer.append(threadlocal.get()).append("-");
        }
        MethodSignature signature = (MethodSignature) point.getSignature();
        Method method = signature.getMethod();
     //获取到当前接口所在类和方法名 用来确定唯一标识
        Class<?> targetClass = method.getDeclaringClass();
        stringBuffer.append(targetClass.getName()).append("-").append(method.getName());
        return stringBuffer.toString();
    }
}

通过String combineKey = getCombineKey(rateLimiter, point); 这个方法 你就可以拿到一个接口唯一标识 当做  key了

通过 redistemplate  的execute 方法去调用相应的 lua 脚本并传入参数

    @Bean
    public DefaultRedisScript<Long> limitScript()
    {
        DefaultRedisScript<Long> redisScript = new DefaultRedisScript<>();
        redisScript.setScriptText(limitScriptText());
        redisScript.setResultType(Long.class);
        return redisScript;
    }

通过  new DefaultRedisScript 这个对象去构建lua 脚本 我们需要给它设定脚本文本 和返回值类型即可 

 

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);";
    }
}

 Long number = redisTemplate.execute(limitScript, keys, count, time); 

这行代码便是调用 刚才构造好的lua脚本了  keys 是一个List 集合所以可以保证key的顺序 脚本中用  KEYS[1]  KEYS[2]  可以选择具体是第几个key 当然当前例子中只有一个key 所以无所谓

count 和 time 的顺序是要注意的 它的顺序就对应着脚本中 ARGV[1] ARGV[2] 的顺序 所以注意不要颠倒顺序   然后通过返回的 number去判断是否大于你设置的最大访问次数即可 大于抛出异常由前端提醒用户  

这样我们就实现了通过redis 去完成限流的要求了 

当然这个方式还可以去做一些其他的要求 比如秒杀这种并发量比较高的场景 ，因为redis的命令是串行化执行，再加上lua脚本的原子性 就保证了 秒杀场景下对原子性的要求 因此是非常合适去做的。

如果当前文章帮到了你，请点赞，收藏，关注 谢谢！