python-令牌桶实现接口限流

令牌桶限流:
再用flask 做项目开发的时候。发现用自带的限流太麻烦了。然后百度到的又是各种问题的。

经测试调整过后。代码如下：
使用方式为：
1、装配器：

class Route(Resource)
	@sleep_and_retry
	@limit_rate_tool
	def post(self):
		pass

2、普通方法使用
		if not can_pass_token_bucket(user, "limit"):
			logger.info("拒绝通过，达到最大流量")
		logger.info("获得令牌，允许通过")
		pass
		

import functools
import math
import time
from datetime import datetime

from flask import request

from exceptions import RateLimitException
from extentions import redis
from utils.common import logger
from utils.date_time import now_datetime

hasinit = False


# 延迟机制
def sleep_and_retry(func):
    @functools.wraps(func)
    def wrapper(*args, **kargs):
        while True:
            try:
                return func(*args, **kargs)
            except RateLimitException as exception:
                time.sleep(exception.period_remaining)

    return wrapper


def limit_rate_tool(func):
    @functools.wraps(func)
    def wrap(*args):
        url = request.url
        period_remaining = 2  # AdogeConsts.LIMIT_RATE_PERIOD_REMAINING
        if not can_pass_token_bucket(url, "limit"):
            logger.info(f"接口{url}达到最大限流，等待{period_remaining}秒")
            raise RateLimitException(message="达到最大限流", period_remaining=period_remaining)

            pass  # if not can_pass_token_bucket(url, "limit", time_zone=5, times=2):
        else:
            return func(args)
            pass  # else-if not can_pass_token_bucket(url, "limit", time_zone=5, times=2):
        pass  # def wrap(*args, **kwargs):

    return wrap
    pass  # def limit_rate_tool(func):


def init_redis_key(key, times):
    redis.hset(key, 'tokens', times)
    last_time = time.mktime(datetime.now().timetuple())  # 记录令牌生成时间
    redis.hset(key, 'last_time', last_time)


#  令牌增加的时机为（当前时间-上一次令牌增加时间) >= (time_zone / times) ,且最短增加时间间隔为1s
#  每秒增加的个数为 times / time_zone，且最大增加个数为times
def can_pass_token_bucket(user, action, time_zone=1, times=3):  # AdogeConsts.LIMIT_RATE_MAX_NUM):
    """
    :param user: 用户唯一标识
    :param action: 用户访问的接口标识(即用户在客户端进行的动作)
    :param time_zone: 接口限制的时间段
    :param times: 限制的时间段内允许多少请求通过
    """
    # 请求来了就倒水，倒水速率有限制
    key = '{}:{}'.format(user, action)
    global hasinit
    if not hasinit:
        init_redis_key(key, times)
        hasinit = True
    rate = times / time_zone  # 令牌生成速度
    capacity = times  # 桶容量
    tokens = redis.hget(key, 'tokens')  # 看桶中有多少令牌
    last_time = redis.hget(key, 'last_time')  # 上次令牌生成时间
    now = time.mktime(datetime.now().timetuple())
    if tokens is not None:
        tokens = int(tokens)
    else:
        logger.info(f"当前令牌为空,重置令牌数量{capacity}")
        tokens = capacity
    if last_time is not None and int(float(last_time.decode())) is not None:
        last_time = int(float(last_time.decode()))
    else:
        last_time = now
    logger.info(f"tokens:{tokens}")
    # 时间戳相差60 代表现实生活相差一秒。所以，如果要求每秒生产的令牌，需要除与
    now_datetime2 = datetime.fromtimestamp(now)
    last_time_datetime = datetime.fromtimestamp(last_time)
    seconds = (now_datetime2 - last_time_datetime).seconds
    logger.info(f"seconds: {seconds}")
    delta_tokens = seconds * rate  # 经过一段时间后生成的令牌
    if delta_tokens >= 1:
        logger.info(f"{key}增加前令牌数量:{tokens}")
        tokens = tokens + math.ceil(delta_tokens)  # 增加令牌
        logger.info(f"当前增加令牌时间{now_datetime()} ,增加个数为:{delta_tokens},当前剩余个数：{tokens}")
        if tokens > capacity:
            logger.info(f"当前增加令牌后大于桶数量,重置令牌数量{capacity}")
            tokens = capacity
        last_time = time.mktime(datetime.now().timetuple())  # 记录令牌生成时间
        redis.hset(key, 'last_time', last_time)
        redis.hset(key, 'tokens', tokens)

    if tokens >= 1:
        logger.info(f"减少前令牌数量:{tokens}")
        tokens = tokens - 1  # 请求进来了，令牌就减少1
        logger.info(f"{key}当前减少令牌时间{now_datetime()} ,当前个数：{tokens}")
        redis.hset(key, 'tokens', tokens)
        return True
    return False