分布式限流解决三方接口QPS限制异常

注：记录开发，自己总结，随便写写，不喜勿喷。

问题描述

之前出现过调三方接口qps异常，我还记录过日记：https://blog.csdn.net/weixin_42459814/article/details/127465053，这种问题经常出现，出现的原因还不止一种，有时候产品放量，有时候集中缓存失效，不同场景用同一appkey等等（三方是根据请求的appkey限制QPS的）。

我主要负责这块业务，只能去寻找解决方案，百度了一波，主要是采用分布式限流来解决。

解决方案

常见的分布式限流方案有滑动窗口算法、漏桶算法、令牌桶算法等等，接下来先简单介绍一下。

滑动窗口算法：

（1）将整个时间划分为更小的多个时间区间

（2）一个时间窗口占用固定的多个时间区间，每有一次请求，就给一个时间区间计数

（3）每经过一个时间区间，就抛弃最老的一个时间区间，加入一个最新的时间区间

（4）如果当前窗口内区间的请求计数总和超过了限制数量，则本窗口内所有请求都会被丢弃

如上图，整个窗口内的时间长度是固定的，记录好每个窗口的请求数并求和，将和与限流量比较，判断请求是否继续。整个窗口在我们场景就是一秒，更小的窗口按1毫秒划分，也可以考虑划的更细。

漏桶算法：

（1）将每个请求视为“水滴”放入漏桶进行存储

（2）漏桶以固定速率漏出水滴（处理请求）

（3）漏桶满了，多余的水滴就丢弃

简单说来就是：如果当前速率小于阈值则直接处理请求，否则不直接处理请求，进入缓冲区，并增加当前水位

漏桶算法的缺陷也很明显，当短时间内有大量的突发请求时，即便此时服务器没有任何负载，每个请求也都得在队列中等待一段时间才能被响应。

这种方案在我们的场景里好像不是很好用。

令牌桶算法：

（1）令牌以固定速率生成

（2）生成的令牌放入令牌桶中存放，如果令牌桶满了则多余的令牌直接丢弃，当请求到达时，会尝试从令牌桶中取令牌，得到令牌的请求可以执行

（3）如果桶空了，则丢弃取令牌的请求

令牌桶的容量大小理论上就是程序需要支撑的最大并发数。令牌桶算法既能够将所有的请求平均分布到时间区间内，又能接受服务器能够承受范围内的突发请求，因此是目前使用较为广泛的一种限流算法。

方案落地

思考一下，在分布式场景，我们要计数的话，就需要在同一个服务或者中间件计数。计数用个服务属实没必要，而且涉及高可用集群以及集群间计数信息同步的问题。这样的话就只能考虑中间件了，redis是一个值得依赖的中间件。

不管是那种算法实现都会有一个问题，这里会产生多个redis命令的执行。如滑动窗口中需要获取单位时间内的请求数、删除之前窗口的记录、添加请求计数等等；令牌桶也需要获取令牌数，扣令牌、恢复令牌数等等，在并发场景下都需要保证原子性。

因此考虑采用redis+lua+aop方案解决分布式限流的问题，接下来看代码。

以滑动窗口为例：

lua脚本如下

-- 获取zset的key
local key = KEYS[1]
-- 脚本传入的限流大小
local limit = tonumber(ARGV[1])
-- 脚本传入的限流起始时间戳
local start = tonumber(ARGV[2])
-- 脚本传入的限流当前时间戳
local now = tonumber(ARGV[3])
-- 脚本传入的限流当前时间戳
local uuid = ARGV[4]
-- 获取当前流量总数
local count = tonumber(redis.call('zcount',key, start, now))
--是否超出限流值
if count + 1 >limit then
    return false
-- 不需要限流
else
    -- 添加当前访问时间戳到zset
    redis.call('zadd', key, now, uuid)
    -- 移除时间区间以外不用的数据，不然会导致zset过大
    redis.call('zremrangebyscore',key, 0, start)
    return true
end

定义注解：

@Target({ElementType.METHOD})
@Retention(RetentionPolicy.RUNTIME)
@Documented
public @interface RateLimit {

    // 限流大小
    int limit();

    // 限流资源名称
    String rateName() default "";
}

aop代码如下：

/**
 * @author shunsheng
 * @Title: RatelimitAspect.java
 * @Description
 * @date 2023 02-08 17:24.
 */
@Aspect
@Component
public class RatelimitAspect {

    @Autowired
    private StringRedisTemplate stringRedisTemplate;

    @Autowired
    private RedisScript<Boolean> rateLimitScript;


    @Pointcut("@annotation(com.shunsheng.study.anno.RateLimit)")
    public void pointCut(){}

    @Before("pointCut() && @annotation(rateLimit)")
    public void before(RateLimit rateLimit) throws Throwable {
        //注解上的参数信息
        int limit = rateLimit.limit();
        String name = rateLimit.rateName();
        //当前时间戳
        long now = System.currentTimeMillis();
        //调用lua脚本获取限流结果
        Boolean isAccess = stringRedisTemplate.execute(
                //lua限流脚本
                rateLimitScript,
                //限流资源名称
                Collections.singletonList(name),
                //限流大小                
                String.valueOf(limit),
                //限流窗口的左区间
                String.valueOf(now - 1000),
                //限流窗口的左区间
                String.valueOf(now),
                //id值，保证zset集合里面不重复，不然会覆盖
                UUID.randomUUID().toString()
        );

        if (!isAccess){
            throw new BusinessException(123, "限流了");
        }
    }
}

脚本读取配置类

@Configuration
public class rateConfig {

/*脚本字符串
    private String lua =
    "local key = KEYS[1]"   +
    "local limit = tonumber(ARGV[1])"   +
    "local start = tonumber(ARGV[2])"   +
    "local now = tonumber(ARGV[3])" +
    "local uuid = ARGV[4]"  +
    "local count = tonumber(redis.call('zcount',key, start, now))"  +
    "if count + 1 >limit then"  +
    "  return false;"  +
    "else"  +
    "  redis.call('zadd', key, now, uuid)"    +
    "  redis.call('zremrangebyscore',key, 0, start)"  +
    "  return true;"    +
    "end";
*/
    @Bean
    public RedisScript<Boolean> loadRedisScript(){
        DefaultRedisScript<Boolean> redisScript = new DefaultRedisScript<>();
        //lua脚本路径
        redisScript.setLocation(new ClassPathResource("LuaScript/token_rate_limit.lua"));
//        redisScript.setScriptText(lua);
        //lua脚本返回值
        redisScript.setResultType(java.lang.Boolean.class);
        return redisScript;
    }
}

统一异常处理类

@RestControllerAdvice
public class ExceptionController {

    @ExceptionHandler(BusinessException.class)
    public Object rateLimitExceptionHandle(BusinessException e) {
        //TODO
        return "限流了";
    }
}

需要限流的方法加个注解

@RateLimit(limit = 1, rateName = "testtoken")
    public ActivityInfoDTO getByLocalCache(String activityId) {
        String key = String.join(CacheKey.cacheKeyDelimiter, CacheKey.activityLocalCacheKeyPrefix, activityId);
        //如果一个key不存在，那么会进入指定的函数生成value
        return activityCaffeineCache.get(key);
    }

写个测试类试一下

@RestController
@RequestMapping("/v1/test")
public class TestController {

    @Resource
    ActivityDaoClient activityDaoClient;


    @PostMapping(value = "/test1", name = "限流测试")
    public Object Test(@RequestHeader HttpHeaders headers, @RequestBody LimitTestDTO request) {
        return activityDaoClient.selectActivityInfoById(request.getId());
    }
}

可以试一下哦，亲测可用，我已经上生产了。

如果想用令牌桶算法的话，只需要以下修改

lua脚本如下：

-- 令牌桶限流: 不支持预消费, 初始桶是满的
-- KEYS[1] string 限流的key
-- ARGV[1] int  桶最大容量
-- ARGV[2] int  每次添加令牌数
-- ARGV[3] int  令牌添加间隔(豪秒)
-- ARGV[4] int  当前时间戳(毫秒)
local bucket_capacity = tonumber(ARGV[1])
local add_token = tonumber(ARGV[2])
local add_interval = tonumber(ARGV[3])
local now = tonumber(ARGV[4])
-- 保存上一次更新桶的时间的key
local LAST_TIME_KEY = KEYS[1].."time";
-- 获取当前桶中令牌数
local token_cnt = tonumber(redis.call("get", KEYS[1]))
-- 桶完全恢复需要的最大时长,秒
local reset_time = math.ceil(bucket_capacity / add_token) * add_interval / 1000;
if token_cnt then -- 令牌桶存在
 -- 上一次更新桶的时间
 local last_time = tonumber(redis.call('get', LAST_TIME_KEY))
 -- 恢复倍数
 local multiple = math.floor((now - last_time) / add_interval)
 -- 恢复令牌数
 local recovery_cnt = multiple * add_token
 -- 确保不超过桶容量
 local token_cnt = math.min(bucket_capacity, token_cnt + recovery_cnt) - 1
 if token_cnt < 0 then
  return false;
 end
 -- 重新设置过期时间, 避免key过期
 redis.call('set', KEYS[1], token_cnt, 'EX', reset_time)
 redis.call('set', LAST_TIME_KEY, last_time + multiple * add_interval, 'EX', reset_time)
 return true;
else -- 令牌桶不存在
 token_cnt = bucket_capacity - 1
 -- 设置过期时间避免key一直存在
 redis.call('set', KEYS[1], token_cnt, 'EX', reset_time);
 redis.call('set', LAST_TIME_KEY, now, 'EX', reset_time + 1);
 return true;
end

注解定义如下

@Target({ElementType.METHOD})
@Retention(RetentionPolicy.RUNTIME)
@Documented
public @interface RateLimit {

    // 令牌桶大小
    int bucketCapacity();
    // 每次添加令牌数
    int addToken();
    // 令牌添加间隔（毫秒）
    int addInterval();

    // 限流资源名称
    String rateName() default "";
}

aop如下

/**
 * @author shunsheng
 * @Title: RatelimitAspect.java
 * @Description
 * @date 2023 02-08 17:24.
 */
@Aspect
@Component
public class RatelimitAspect {

    @Autowired
    private StringRedisTemplate stringRedisTemplate;

    @Autowired
    private RedisScript<Boolean> rateLimitScript;


    @Pointcut("@annotation(com.shunsheng.study.anno.RateLimit)")
    public void pointCut(){}

    @Before("pointCut() && @annotation(rateLimit)")
    public void before(RateLimit rateLimit) throws Throwable {
        //注解上的参数信息
        int bucketCapacity = rateLimit.bucketCapacity();
        int addToken = rateLimit.addToken();
        int addInterval = rateLimit.addInterval();
        String name = rateLimit.rateName();
        //当前时间戳
        long now = System.currentTimeMillis();
        //调用lua脚本获取限流结果
        Boolean isAccess = stringRedisTemplate.execute(
                //lua限流脚本
                rateLimitScript,
                //限流资源名称
                Collections.singletonList(name),
                //桶最大容量
                String.valueOf(bucketCapacity),
                //每次添加令牌数
                String.valueOf(addToken),
                //令牌添加间隔(豪秒)
                String.valueOf(addInterval),
                //当前时间戳
                String.valueOf(System.currentTimeMillis())
        );

        if (!isAccess){
            throw new BusinessException(123, "限流了");
        }
    }
}

需要限流的方法改下注解

@RateLimit(bucketCapacity = 100, addToken = 1, addInterval = 500, rateName = "testtoken")
    public ActivityInfoDTO getByLocalCache(String activityId) {
        String key = String.join(CacheKey.cacheKeyDelimiter, CacheKey.activityLocalCacheKeyPrefix, activityId);
        //如果一个key不存在，那么会进入指定的函数生成value
        return activityCaffeineCache.get(key);
    }