15行Python程式碼，幫你理解令牌桶演算法

在網路中傳輸資料時，為了防止網路擁塞，需限制流出網路的流量，使流量以比較均勻的速度向外傳送，令牌桶演算法就實現了這個功能，可控制傳送到網路上資料的數目，並允許突發資料的傳送。

什麼是令牌

從名字上看令牌桶，大概就是一個裝有令牌的桶吧，那麼什麼是令牌呢？

紫薇格格拿的令箭，可以發號施令，令行禁止。在計算機的世界中，令牌也有令行禁止的意思，有令牌，則相當於得到了進行操作的授權，沒有令牌，就什麼都不能做。

用令牌實現限速器

我們用1塊令牌來代表傳送1位元組資料的資格，假設我們源源不斷的發放令牌給程式，程式就有資格源源不斷的傳送資料，當我們不發放令牌給程式，程式就相當於被限流，無法傳送資料了。接下來我們說說限速器，所謂限速器，就是讓程式在單位時間內，最多隻能傳送一定大小的資料。假設在1秒發放10塊令牌，那麼程式傳送資料的速度就會被限制在10bytes/s。如果1秒內有大於10bytes的資料需要傳送，就會因為沒有令牌而被丟棄。

改進限速器——加個桶

我們實現的限速器，速度是恆定的，但是在實際的應用中，往往會有突發的傳輸需求（需要更快速的傳送，但是不會持續太久，也不會引起網路擁塞），這種資料碰上我們的限速器，就因為拿不到令牌而無法傳送。

對限速器進行一下改動，依然1秒產生10塊令牌，但是我們把產生出來的令牌先放到一個桶裡，當程式需要傳送的時候，從桶裡取令牌，不需要的時候，令牌就會在桶裡沉澱下來，假設桶裡沉澱了10塊令牌，程式最多就可以在1秒內傳送20bytes的資料，滿足了突發資料傳輸的要求，並且由於桶裡的令牌被用完，下一秒最多依然只能發10bytes的資料，不會因為持續傳送大量資料，對網路造成壓力。

15行Python程式碼實踐令牌桶

令牌桶需要以一定的速度生成令牌放入桶中，當程式要傳送資料時，再從桶中取出令牌。這裡似乎有點問題，如果我們使用一個死迴圈，來不停地發放令牌，程式就被阻塞住了，有沒有更好的辦法？

我們可以在取令牌的時候，用現在的時間減去上次取令牌的時間，乘以令牌的發放速度，計算出桶裡可以取的令牌數量（當然不能超過桶的大小），從而避免迴圈發放的邏輯。

接下來看程式碼：

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
複製程式碼

掃碼關注Python私房菜