限流算法 - 令牌桶算法(python)

class TokenBucket:
    '''
        令牌桶算法：出
            令牌出桶速率恒定，当桶中无剩余令牌，则不能取出
    '''
    def __init__(self, rate, capacity):
        '''
            rate:出桶速率
            volume: 最大容积
            current：桶中现有令牌
            times：计时
        '''
        self._rate = rate
        self._capacity = capacity
        self._current_amount = 0
        self._last_consume_time = int(time.time())

    def consume(self, token_amount):
        # 从上次发送到这次发送，新取出的令牌数量
        increment = (int(time.time()) - self._last_consume_time) * self._rate
        # 桶中当前剩余令牌数量
        self._current_amount = min(increment + self._current_amount, self._capacity)
        # 如果需要的令牌超过剩余的令牌，则不能发送数据
        if token_amount > self._current_amount:
            return False
        self._last_consume_time = int(time.time())
        # 可以取出令牌，取出后桶中剩余令牌数量
        self._current_amount -= token_amount
        return True