import time


class TokenBucket(object):

    # rate是令牌發放速度，capacity是桶的大小
    def __init__(self, rate, capacity):
        self._rate = rate
        self._capacity = capacity
        self._current_amount = 0
        self._last_consume_time = int(time.time())

    # token_amount是傳送資料需要的令牌數
    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

作者：simpleapples
連結：https://juejin.im/post/5ab10045518825557005db65
來源：掘金
著作權歸作者所有。商業轉載請聯絡作者獲得授權，非商業轉載請註明出處。