Python实现常见限流算法的示例代码

 brew install redis && brew services start redis

 def get_redis_con():
     pool = redis.ConnectionPool(max_connections=4, decode_responses=True)
     return redis.Redis(connection_pool=pool)

 def fixed_window(user, action, time_zone=60, max_times=30):
     key = f'{action}'
     count = get_redis_con().get(key)
     if not count:
         count = 1
         get_redis_con().setex(key, time_zone, count)
     if int(count) < max_times:
         get_redis_con().incr(key)
         return True
     return False

 def silde_window(user, action, time_zone=60, max_times=30):
     key = f'{action}'
     now_ts = time.time() * 1000
     # ms级时间戳，保证唯一
     value = now_ts
     # 时间窗口的左边界
     old_ts = now_ts - time_zone * 1000
     # 记录 {成员元素：分数值}
     mapping = {
         value: now_ts,
     }
     get_redis_con().zadd(key, mapping)
     # 删除时间窗口之前的数据
     get_redis_con().zremrangebyscore(key, -1, old_ts)
     # 获得窗口内的行为数量
     count = get_redis_con().zcard(key)
     get_redis_con().expire(key, time_zone + 1)
     if not count or int(count) < max_times:
         return True
     return False

 def pass_token_bucket(user, action, time_zone=60, max_times=30):
     key = f'{user}:{action}'
     # 令牌生成速度
     rate = max_times / time_zone
     capacity = max_times
     token_count = get_redis_con().hget(key, 'tokens')
     last_time = get_redis_con().hget(key, 'last_time')
     now = time.time()
     token_count = int(token_count) if token_count else capacity
     last_time = last_time if last_time else now
     # 经过一段时间之后生成的令牌数量
     new_token_count = int((now - last_time) * rate)
     if new_token_count > 1:
         token_count += new_token_count
         if token_count > capacity:
             token_count = capacity
         last_time = time.time()
         get_redis_con().hset(key, 'last_time', last_time)
 ​
     if token_count >= 1:
         token_count -= 1
         get_redis_con().hset(key, 'tokens', token_count)
         return True
     return False

 executor = ThreadPoolExecutor(max_workers=4)
 APIS = ['/api/a', '/get/user/1', '/get/user/2', '/get/user/3']
 ​
 ​
 def task() -> bool:
     user = random.randint(1000, 1010)
     status = pass_token_bucket(user, random.choice(APIS))
     if not status:
         raise SystemError('{}被限制'.format(user))
     return status
   
   if __name__ == '__main__':
     jobs = [executor.submit(task) for _ in range(1000)]
     for job in jobs:
         print(job.result())