需求

现在有一个API开放平台，要求实现多维度的QPS控制，例如：1. 根据APP ID控制QPS；2. 根据APP ID+API控制QPS；3. 根据套餐进行QPS控制等。

原理

使用Redis Sorted Set数据结构存储请求，Sorted Set是有序不重复的，内部集合包含值和权重，权重设置为毫秒时间戳，这样做是为了根据毫秒时间戳范围查询当前时间窗口的请求数量，值可以使用UUID。
image.png
存储数据命令：zAdd $限流维度Key $毫秒时间戳 $UUID
查询当前时间窗口请求数量：zCount $限流维度Key $开始时间 $结束时间
这里需要注意的是，必须通过Redis获取时间，因为集群部署微服务的时候多台机器的时间可能有偏差，如果获取的是Linux时间将导致限流控制不准确。
Redis中Sorted Set存储的数据预览：
image.png

代码实现

  1. /**
  2.  * qps 限流方法，直接调用此方法，如果没有抛出异常则说明QPS没有超出限制
  3.  *
  4.  * @param qpsKey 限流对象，根据业务拼装字符串，实现多维度的qps控制
  5.  * @param qps    qps限制，例如：5，代表一秒只能接收5个请求
  6.  * @throws QPSException
  7.  */
  8. public void qpsLimit(String qpsKey, Long qps) throws QPSException {
  9.     // 通过Lua脚本保证多个Redis操作的原子性
  10.     String luaScript = 
  11.             // 获取Redis服务器时间
  12.             "local curTime = redis.call('time'); \n" + 
  13.             // 进行精确到毫秒
  14.             "curTime = (curTime[1] * 1000000 + curTime[2]) / 1000; \n" +
  15.             // 删除一秒之前的请求记录，防止set过大
  16.             String.format("redis.call('zRemRangeByScore', '%s', 0, curTime - 1001); \n", qpsKey) +
  17.             // 获取前一秒的请求数量
  18.             String.format("local curQps = redis.call('zCount', '%s', (curTime - 1000), curTime); \n", qpsKey) +
  19.             // 如果前一秒的请求数量大于QPS限制，返回1，代表QPS达到上限
  20.             String.format("if (curQps >= %s) then return 1; end \n", qps) +
  21.             // 否则将请求加入set中，返回0，代表QPS没有达到上限
  22.             String.format("redis.call('zAdd', '%s', curTime, '%s'); return 0; \n", qpsKey, UUID.randomUUID());
  24.     DefaultRedisScript<Long> script = new DefaultRedisScript<>(luaScript, Long.class);
  25.     Long execute = (Long) redisTemplate.execute(script, Collections.EMPTY_LIST);
  26.     if (Long.valueOf(1L).equals(execute)) {
  27.         throw new QPSException(String.format("QPS超出限制，QPS限制：%d/s", qps));
  28.     }
  29. }

结果验证

配置QPS限制为2/s，即一秒最多接收两个请求：
image.png

参考：https://www.jb51.net/article/215286.htm