Python3 下利用 redis 对接口限速（qbit）

前言

• 技术栈

  Redis Server 5.0.7
  Python       3.11
  reids 库     5.0.3

• 利用 redis 的 zset 实现接口调用限速

案例

• 测试代码

  # encoding: utf-8
  # author: qbit
  # date: 2024-03-28
  # summary: 测试用 redis 做接口限流
  
  import redis
  import time
  
  Period = 10         # 窗口周期 10s
  MaxCount = 5        # 周期内最多运行调用 5 次
  RedisClient = redis.StrictRedis(host='192.168.21.152', port=6379, db=2)
  
  def IsCallAllowed(userId: str):
    r""" 是否允许调用 """
    period = int(Period * 1000)
    nowTime = int(time.time() * 1000)       # 时间精度设为毫秒
    pipe: redis.client.Pipeline = RedisClient.pipeline()
  
    # 删除窗口之外的记录
    # zset 为空时，会自动删除 key
    pipe.zremrangebyscore(userId, '-inf', nowTime - period)
    # 将当前请求的时间戳添加到 zset 中
    pipe.zadd(userId, {nowTime: nowTime})
    # 当前 userId 在 zset 中的次数, 查看范围可以用 zcount
    pipe.zcard(userId)
    #查看有序集合
    pipe.zrange(userId, 0, -1, withscores=True)
    # 设置过期时间，时间单位未秒
    # 这里是为了自动删除长期不被访问的 key
    pipe.expire(userId, Period)    
  
    results = pipe.execute()
    # results 代表上面几个命令的返回值
    # logger.debug(f"results: \n{results}")
    # delCount = results[0]
    # addCount = results[1]
    nowCount = results[2]    
    # zset = results[3]
    # expireFlag = results[4]
    # print(f"delCount: {delCount}")
    # print(f"addCount: {addCount}")
    # print(f"nowCount: {nowCount}")    
    # print(f"zset: {zset}")
    # print(f"expireFlag: {expireFlag}")
  
    return nowCount <= MaxCount
  
  
  if __name__ == "__main__":
    for i in range(1, 21):
        bRtn = IsCallAllowed('qbit')
        if bRtn:
            print(f"{time.strftime('%Y-%m-%d %H:%M:%S')} 第 {i} 次调用被允许")
        else:
            print(f"{time.strftime('%Y-%m-%d %H:%M:%S')} 第 {i} 次调用被拒绝")
        if i <= 10:
            time.sleep(1)
        else:
            time.sleep(2)       # 10s 后降低请求速率
        

• 结果输出

  2024-03-28 14:33:53 第 1 次调用被允许
  2024-03-28 14:33:54 第 2 次调用被允许
  2024-03-28 14:33:55 第 3 次调用被允许
  2024-03-28 14:33:56 第 4 次调用被允许
  2024-03-28 14:33:57 第 5 次调用被允许
  2024-03-28 14:33:58 第 6 次调用被拒绝
  2024-03-28 14:33:59 第 7 次调用被拒绝
  2024-03-28 14:34:00 第 8 次调用被拒绝
  2024-03-28 14:34:01 第 9 次调用被拒绝
  2024-03-28 14:34:02 第 10 次调用被拒绝
  2024-03-28 14:34:03 第 11 次调用被拒绝
  2024-03-28 14:34:05 第 12 次调用被拒绝
  2024-03-28 14:34:07 第 13 次调用被拒绝
  2024-03-28 14:34:09 第 14 次调用被拒绝
  2024-03-28 14:34:11 第 15 次调用被拒绝
  2024-03-28 14:34:13 第 16 次调用被允许
  2024-03-28 14:34:15 第 17 次调用被允许
  2024-03-28 14:34:17 第 18 次调用被允许
  2024-03-28 14:34:19 第 19 次调用被允许
  2024-03-28 14:34:21 第 20 次调用被允许

参考文献

• 实现接口的限流各种算法

本文出自 qbit snap