Last active
December 2, 2021 19:25
-
-
Save yeus/dff02dce88c6da9073425b5309f524dd to your computer and use it in GitHub Desktop.
async rate limiting function using token bucket algorithm
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
""" | |
author: Thomas Meschede | |
license: MIT | |
""" | |
def rate_limit_burst(max_calls: float, interval: float): | |
""" | |
Rate-limits the decorated function locally, for one process. Using Token Bucket algorithm. | |
max_calls: maximum number of calls of function in interval | |
interval: interval in [s] | |
""" | |
token_fill_rate: float = max_calls/interval | |
def decorate(func): | |
last_time_called = time.perf_counter() | |
last_time_token_added = time.perf_counter() | |
token_bucket: float = max_calls # initialize with max tokens | |
@wraps(func) | |
async def rate_limited_function(*args, **kwargs): | |
nonlocal last_time_token_added | |
nonlocal token_bucket | |
try: | |
elapsed = time.perf_counter() - last_time_token_added | |
# refill token_bucket to a maximum of max_calls | |
token_bucket = min(token_bucket + elapsed * token_fill_rate, max_calls) | |
last_time_token_added = time.perf_counter() | |
# check if have to wait for a function call (min 1 token in order | |
# to make a call) | |
if token_bucket < 1.0: | |
left_to_wait = (1 - token_bucket) / token_fill_rate | |
await asyncio.sleep(left_to_wait) | |
return await func(*args, **kwargs) | |
finally: | |
# for every call we can consume one token, if the bucket is empty, | |
# we have to wait | |
token_bucket -= 1.0 | |
return rate_limited_function | |
return decorate |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment