Welcome, guest | Sign In | My Account | Store | Cart

Implementation of the token bucket algorithm (Python recipe) by Alec Thomas
ActiveState Code (http://code.activestate.com/recipes/511490/)

An implementation of the token bucket algorithm in Python.

Python, 49 lines
Download
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
from time import time

class TokenBucket(object):
    """An implementation of the token bucket algorithm.
    
    >>> bucket = TokenBucket(80, 0.5)
    >>> print bucket.consume(10)
    True
    >>> print bucket.consume(90)
    False
    """
    def __init__(self, tokens, fill_rate):
        """tokens is the total tokens in the bucket. fill_rate is the
        rate in tokens/second that the bucket will be refilled."""
        self.capacity = float(tokens)
        self._tokens = float(tokens)
        self.fill_rate = float(fill_rate)
        self.timestamp = time()

    def consume(self, tokens):
        """Consume tokens from the bucket. Returns True if there were
        sufficient tokens otherwise False."""
        if tokens <= self.tokens:
            self._tokens -= tokens
        else:
            return False
        return True

    def get_tokens(self):
        if self._tokens < self.capacity:
            now = time()
            delta = self.fill_rate * (now - self.timestamp)
            self._tokens = min(self.capacity, self._tokens + delta)
            self.timestamp = now
        return self._tokens
    tokens = property(get_tokens)

if __name__ == '__main__':
    from time import sleep
    bucket = TokenBucket(80, 1)
    print "tokens =", bucket.tokens
    print "consume(10) =", bucket.consume(10)
    print "consume(10) =", bucket.consume(10)
    sleep(1)
    print "tokens =", bucket.tokens
    sleep(1)
    print "tokens =", bucket.tokens
    print "consume(90) =", bucket.consume(90)
    print "tokens =", bucket.tokens

The token bucket algorithm is a very simple (and thus, hopefully I haven't screwed up this implementation) but useful method of rate limiting.

The algorithm consists of a bucket with a maximum capacity of N tokens which refills at a rate R tokens per second. Each token typically represents a quantity of whatever resource is being rate limited (network bandwidth, connections, etc.).

This allows for a fixed rate of flow R with bursts up to N without impacting the consumer.

This implementation is not thread-safe.

WP article on this algorithm: http://en.wikipedia.org/wiki/Token_bucket

Tags: algorithms

2 comments

Phillip Oertel 13 years, 3 months ago  # | flag

Alec, great sample, thank you.

I think your analysis is a little off though. I'll demonstrate:

bucket = TokenBucket(10, 10)
sleep(1)
bucket.consume(10)
print "tokens = ", bucket.tokens

>>> tokens = 10

The average rate will not exceed R, like you say, but the allowed burstiness is actually R+N. In my application (a simple QPS limiter) this was a bug. The solution is to do lines 32 and 35 before the condition in line 31.

Phillip Oertel 13 years, 3 months ago  # | flag

Sorry, typo! Line 32 should appear before the condition, line 35 should appear after. The revised code:

def get_tokens(self):
    now = time()
    if self._tokens < self.capacity:
        delta = self.fill_rate * (now - self.timestamp)
        self._tokens = min(self.capacity, self._tokens + delta)
    self.timestamp = now
    return self._tokens
tokens = property(get_tokens)