# token_bucket.py
#! /usr/bin/python3
from time import time
from threading import Lock
class TokenBucket:
"""
An implementation of the token bucket algorithm.
"""
def __init__(self):
self.tokens = 0
self.rate = 0
self.last = time()
self.lock = Lock()
def set_rate(self, rate):
with self.lock:
self.rate = rate
self.tokens = self.rate
def consume(self, tokens):
with self.lock:
if not self.rate:
return 0
now = time()
lapse = now - self.last
self.last = now
self.tokens += lapse * self.rate
if self.tokens > self.rate:
self.tokens = self.rate
self.tokens -= tokens
if self.tokens >= 0:
return 0
else:
return -self.tokens / self.rate
if __name__ == '__main__':
import sys
from time import sleep
one_kb = 1024 * 1024
bucket = TokenBucket()
bucket.set_rate(one_kb)
for _ in range(10):
nap = bucket.consume(one_kb)
sleep(nap)
print(".")
sys.exit(0)
#######################################################
#! /usr/bin/python3
import logging
import time
import unittest
from .token_bucket import TokenBucket
logging.disable(logging.CRITICAL)
class UtilsHTTPRateLimitTest(unittest.TestCase):
def setUp(self):
self.bucket = TokenBucket()
def tearDown(self):
pass
def test_set_rate(self):
self.bucket.set_rate(1024*1024)
self.assertEqual(self.bucket.rate, 1024*1024)
self.assertEqual(self.bucket.tokens, self.bucket.rate)
def test_consume(self):
self.assertEqual(self.bucket.consume(0), 0)
self.assertEqual(self.bucket.consume(999), 0)
self.bucket.set_rate(1024*1024)
self.assertEqual(self.bucket.consume(1024*1024), 0)
self.assertLessEqual(self.bucket.consume(1024*1024*2), 2)
self.assertLessEqual(self.bucket.consume(1024*1024*4), 6)
self.assertLessEqual(self.bucket.consume(1024*1024*6), 12)
self.bucket.set_rate(1024*1024)
time.sleep(1)
self.assertEqual(self.bucket.consume(1024*1024), 0)
self.assertGreater(self.bucket.consume(1024*1024), 0)
