python 实现限流

固定窗口

固定窗口就是记录一个固定的时间窗口内的操作次数，操作次数超过阈值则进行限流。

def fix_window_limit(redis_obj, period, max_count):
    """
    固定窗口
    :param redis_obj:redis连接对象
    :param period: 周期 秒为单位
    :param max_count: 最大请求数量数量
    :return:
    """
    key = "global_limit"
    current = redis_obj.get(key)
    if current:
        current = int(current.decode('utf-8'))
        if current >= max_count:
            return False
    value = r.incr(key)
    if value == 1:
        r.expire(key, period)
    return True

滑动窗口

滑动窗口就是记录一个滑动的时间窗口内的操作次数，操作次数超过阈值则进行限流。

def rolling_window_limit(redis_obj, period, max_count):
    """
    滑动窗口
    :param redis_obj:redis连接对象
    :param period: 周期 秒为单位
   :param max_count: 最大请求数量数量
    :return:
    """
    key = "global_limit"
    now_time = int(time.time()*1000)
    # 使用管道
    pipe = redis_obj.pipeline()
    pipe.multi()
    # 添加当前操作当zset中
    pipe.zadd(key, {str(now_time): now_time})
    # 整理zset，删除时间窗口外的数据
    pipe.zremrangebyscore(key, 0, now_time - period * 1000)
    # 获取当前窗口的数量
    pipe.zcard(key)
    pipe.expire(key, period+1)
    result = pipe.execute()
    pipe.close()
    return result[-2] <= max_count