SpringBoot 如何进行限流？老鸟们还可以这样玩！

Posted 2022-12-04 飘渺Jam

大家好，我是飘渺。

在之前一篇文章中我们详细介绍了为什么需要对接口进行限流，也介绍了常见的限流算法，最后还基于Guava工具类实现了接口限流。但是这种方式有个问题，无法实现分布式限流。那今天我们来利用Redis + Lua 来实现分布式限流。

Lua 脚本和 mysql 数据库的存储过程比较相似，他们执行一组命令，所有命令的执行要么全部成功或者失败，以此达到原子性。也可以把 Lua 脚本理解为，一段具有业务逻辑的代码块。

实现过程

第一步：引入Redis依赖包

<dependency>  
    <groupId>org.springframework.boot</groupId>  
    <artifactId>spring-boot-starter-data-redis</artifactId>  
</dependency>

第二步：配置Redis

/**
 * @author  JAVA日知录
 * @date 2022/5/2 22:35
 */
@Configuration
public class RedisConfig 

    @Bean
    public RedisTemplate<String, Object> redisTemplate(RedisConnectionFactory factory) 
        RedisTemplate<String, Object> template = new RedisTemplate<>();
        template.setConnectionFactory(factory);

        // 使用Jackson2JsonRedisSerializer来序列化和反序列化redis的value值（默认使用JDK的序列化方式）
        Jackson2JsonRedisSerializer<Object> serializer = new Jackson2JsonRedisSerializer<>(Object.class);
        ObjectMapper mapper = new ObjectMapper();
        mapper.setVisibility(PropertyAccessor.ALL, JsonAutoDetect.Visibility.ANY);
        mapper.activateDefaultTyping(LaissezFaireSubTypeValidator.instance,
                ObjectMapper.DefaultTyping.NON_FINAL, JsonTypeInfo.As.PROPERTY);
        serializer.setObjectMapper(mapper);

        template.setValueSerializer(serializer);
        // 使用StringRedisSerializer来序列化和反序列化redis的key值
        template.setKeySerializer(new StringRedisSerializer());
        template.afterPropertiesSet();

        return template;
    

第二步：自定义限流注解

@Retention(RetentionPolicy.RUNTIME)
@Target(ElementType.METHOD)
@Documented
public @interface RedisLimit 
    /**
     * 资源的key,唯一
     * 作用：不同的接口，不同的流量控制
     */
    String key() default "";

    /**
     * 最多的访问限制次数
     */
    long permitsPerSecond() default 2;

    /**
     * 过期时间也可以理解为单位时间，单位秒，默认60
     */
    long expire() default 60;


    /**
     * 得不到令牌的提示语
     */
    String msg() default "系统繁忙,请稍后再试.";

第三步：创建限流异常

/**
 * @author JAVA日知录
 * Redis限流自定义异常
 * @date 2022/5/2 21:43
 */
public class RedisLimitException extends RuntimeException
    public RedisLimitException(String msg) 
        super( msg );
    

第四步：使用AOP切面拦截限流注解

/**
 * Limit AOP
 * @author JAVA日知录
 * @date 2021/9/24 3:07 下午
 */
@Slf4j
@Aspect
@Component
public class RedisLimitAop 

    @Autowired
    private StringRedisTemplate stringRedisTemplate;


    @Pointcut("@annotation(com.jianzh5.blog.limit.redis.RedisLimit)")
    private void check() 

    

    @Before("check()")
    public void before(JoinPoint joinPoint) 
        MethodSignature signature = (MethodSignature) joinPoint.getSignature();
        Method method = signature.getMethod();

        //拿到RedisLimit注解，如果存在则说明需要限流
        RedisLimit redisLimit = method.getAnnotation(RedisLimit.class);

        if(redisLimit != null)
            //获取redis的key
            String key  = redisLimit.key();
            String className = method.getDeclaringClass().getName();
            String name = method.getName();

            String limitKey = key + className + method.getName();

            log.info(limitKey);

            if(StringUtils.isEmpty(key))
                throw new RedisLimitException( "key cannot be null" );
            

            long limit = redisLimit.permitsPerSecond();

            long expire = redisLimit.expire();

            List<String> keys = new ArrayList<>();
            keys.add( key );
            String luaScript = buildLuaScript();
            RedisScript<Long> redisScript = new DefaultRedisScript<>( luaScript, Long.class );

            Long count = stringRedisTemplate.execute( redisScript, keys, String.valueOf(limit), String.valueOf(expire) );

            log.info( "Access try count is  for key=", count, key );

            if (count != null && count == 0) 
                log.debug("令牌桶=，获取令牌失败",key);
                throw new RedisLimitException(redisLimit.msg());
            
        

    

    /**
     * 构建redis lua脚本
     * @return
     */
    private String buildLuaScript() 
        StringBuilder luaString = new StringBuilder();
        luaString.append( "local key = KEYS[1]" );
        //获取ARGV内参数Limit
        luaString.append( "\\nlocal limit = tonumber(ARGV[1])" );
        //获取key的次数
        luaString.append( "\\nlocal curentLimit = tonumber(redis.call('get', key) or \\"0\\")" );
        luaString.append( "\\nif curentLimit + 1 > limit then" );
        luaString.append( "\\nreturn 0" );
        luaString.append( "\\nelse" );
        //自增长 1
        luaString.append( "\\n redis.call(\\"INCRBY\\", key, 1)" );
        //设置过期时间
        luaString.append( "\\nredis.call(\\"EXPIRE\\", key, ARGV[2])" );
        luaString.append( "\\nreturn curentLimit + 1" );
        luaString.append( "\\nend" );
        return luaString.toString();
    

第五步：给需要限流的接口加上注解

/**
 * 公众号：JAVA日知录
 * 限流测试类基于Redis限流
 */
@Slf4j
@RestController
@RequestMapping("/limit/redis")
public class LimitRedisController 

    /**
     * 基于Redis AOP限流
     */
    @GetMapping("/test")
    @RedisLimit(key = "redis-limit:test", permitsPerSecond = 2, expire = 1, msg = "当前排队人数较多，请稍后再试！")
    public String test() 
        log.info("限流成功。。。");
        return "ok";
    

第六步：体验效果

通过访问测试地址：http://127.0.0.1:8080/limit/redis/test，反复刷新并观察输出结果：

正常响应时：

"status":100,"message":"操作成功","data":"ok","timestamp":1652343229643

触发限流时：

"status":500,"message":"当前排队人数较多，请稍后再试！","data":null,"timestamp":1652343239035

通过观察得之，基于自定义注解同样实现了接口限流的效果。

优化

程序每次执行每次都需要通过buildLuaScript()方法构建lua执行脚本，比较 low，我们可以生成一个lua文件放在resources目录下，利用@PostConstruct注解提前加载。

1. 在resouces文件夹下创建lua文件 rateLimiter.lua

--获取KEY
local key = KEYS[1]

local limit = tonumber(ARGV[1])

local curentLimit = tonumber(redis.call('get', key) or "0")

if curentLimit + 1 > limit
    then return 0
else
    -- 自增长 1
    redis.call('INCRBY', key, 1)
    -- 设置过期时间
    redis.call('EXPIRE', key, ARGV[2])
    return curentLimit + 1
end

1. 修改RedisLimitAop，通过@PostConstruct注入DefaultRedisScript

@Slf4j
@Aspect
@Component
public class RedisLimitAop 

    @Autowired
    private StringRedisTemplate stringRedisTemplate;

    private DefaultRedisScript<Long> redisScript;

    @PostConstruct
    public void init()
        redisScript = new DefaultRedisScript<>();
        redisScript.setResultType(Long.class);
        redisScript.setScriptSource(new ResourceScriptSource(new ClassPathResource("rateLimiter.lua")));
    
  ...

小结

基于Redis + Lua 可以很方便地实现分布式限流，算是SpringBoot老鸟系列限流文章的补充扩展。

那么现在问题来了，我们现在有基于Guava实现的单机限流，又有基于Redis+Lua实现的分布式限流，那能不能将两种限流功能做成一个独立的公共组件，让使用方根据实际情况选择对应的限流功能呢？

好了，今天的文章就到这里了，我是飘渺，咱们下期见～

老鸟系列源码已经上传至GitHub，需要的关注本公号并回复关键字 0923 获取源码地址。

最后说一句（别白嫖，求关注）新开了一个纯技术交流群（一群已满），群里氛围还不错，无广告，无套路，单纯的吹牛逼，侃人生，想进的可以通过下方二维码加我微信，备注进群！