I've written a class that implements something akin to the token bucket algorithm so that I may rate limit aysnc HTTP requests made from my application.
The code works but I'm still not sure if it feels 'right', maybe just because I am new to async programming in Python.
- It feels weird/bad how the state (
self.tokens) is managed (or not managed) meaning checks in the
- Maybe I shouldn't be using a context manager (
__aexit__) for this.
- I'm not sure if I should be doing any cleaning-up in
- Implementation of the algorithm may not be 100% faithful as I am waiting until the bucket is empty before refilling as a batch . I think this may be OK because only one token can be consumed per request.
""" Typical usage example: async def rateLimitedTask(bucket): await bucket.removeToken() #wait for token #do the task which has now been rate limited ... async with TokenBucket(2, 1/5) as bucket: #make arbitrary async calls to rateLimitedTask() ... """ import asyncio import time from collections import deque class TokenBucket(): """Context manager which provides a token bucket. Attributes: tokens (collections.deque): The tokens (which manifest as timestamps) in reverse chronological order. last (float): The most recent token consumed. rate (float): See __init__ arg `tokensPerS`. capacity (int): See __init__ arg `capacity`. sleepDuration (float): See __init__ arg `refillSleepS`. """ def __init__(self, capacity=1, tokensPerS=1.0, refillSleepS=0.1): """Initialises the TokenBucket context manager. Args: capacity (int, optional): The maxiumum tokens the bucket can hold. Larger is burstier. Defaults to 1. tokensPerS (float, optional): The average replenishment rate in tokens per second. Defaults to 1. refillSleepS (float, optional): If not enough time has passed to refill; how many seconds to sleep before trying again. Defaults to 0.1. """ self.capacity = capacity self.rate = tokensPerS self.last = None self.tokens = deque() self.sleepDuration = refillSleepS def hasTokens(self): return len(self.tokens) > 0 def isEmpty(self): return len(self.tokens) < 1 async def _refill(self): """Refill the empty bucket back up to capacity. Is called only when the bucket is empty. The tokens themselves are timestamps of when the token was created. Uses monotonic time to safeguard against system clock changes. """ if self.last: #prevents unnecessary run on virgin execution while self.capacity / (time.monotonic() - self.last) > self.rate: await asyncio.sleep(self.sleepDuration) if self.hasTokens(): return #another call filled the bucket already for _ in range(self.capacity): self.tokens.appendleft(time.monotonic()) async def removeToken(self): """Removes a token from the bucket and returns void. If the bucket is empty it will await its refill before returning. """ if self.isEmpty(): await self._refill() self.last = self.tokens.pop() async def __aenter__(self): return self async def __aexit__(self, exc_type, exc, tb): pass
Any help would be much appreciated.