khushalps
70a020f9b2
git-subtree-dir: utils_v2 git-subtree-split: 6c32fa9a3d50c06a5487783423147b2bf9ce7549
213 lines
8.3 KiB
Python
213 lines
8.3 KiB
Python
"""
|
|
|
|
AUTHOR:
|
|
|
|
Khushal P Soonderji
|
|
|
|
DATE:
|
|
|
|
Wednesday, 15th May, 2024
|
|
|
|
OBJECTIVE:
|
|
|
|
To provide a rate-limiting mechanism using the simple token bucket algorithm.
|
|
|
|
REFERENCES:
|
|
|
|
N/A
|
|
|
|
DOWNLOADS:
|
|
|
|
N/A
|
|
|
|
"""
|
|
|
|
|
|
# *****************************************************************************************************************
|
|
# ***** ****
|
|
# *** IMPORT ***
|
|
# ***** ****
|
|
# *****************************************************************************************************************
|
|
|
|
|
|
# To make sibling directories accessible for imports:
|
|
import sys
|
|
sys.path.append(".")
|
|
sys.path.append("..")
|
|
|
|
# For date and time keeping:
|
|
import datetime
|
|
import time
|
|
|
|
# For asynchronous activities:
|
|
import asyncio
|
|
|
|
# For mathematical operations:
|
|
import math
|
|
|
|
|
|
# *****************************************************************************************************************
|
|
# ***** ****
|
|
# *** MACROS / ONE-TIME INIT ***
|
|
# ***** ****
|
|
# *****************************************************************************************************************
|
|
|
|
|
|
# --- Nothing Yet
|
|
|
|
|
|
# *****************************************************************************************************************
|
|
# ***** ****
|
|
# *** VARIABLES ***
|
|
# ***** ****
|
|
# *****************************************************************************************************************
|
|
|
|
|
|
# --- Nothing Yet
|
|
|
|
|
|
# *****************************************************************************************************************
|
|
# ***** ****
|
|
# *** CLASSES ***
|
|
# ***** ****
|
|
# *****************************************************************************************************************
|
|
|
|
|
|
class TokenBucket:
|
|
|
|
"""
|
|
MODES: "add" - Add 'rate_limit' no. of tokens to the bucket every interval.
|
|
"reset" - Reset the token count to 'rate_limit' every interval
|
|
"""
|
|
|
|
MODE_RESET = 0
|
|
MODE_ADD = 1
|
|
|
|
ONE_YEAR = 3_15_36_000.0
|
|
ONE_MONTH = 26_78_400.0
|
|
ONE_WEEK = 6_04_800.0
|
|
ONE_DAY = 86_400.0
|
|
ONE_HOUR = 3_600.0
|
|
ONE_MINUTE = 60.0
|
|
ONE_SECOND = 1.0
|
|
|
|
def __init__(self, rate_limit, seconds = 1.0, mode = "reset", sleep = 0.1, in_sequence = False):
|
|
|
|
"""
|
|
Initialize the rate controller.
|
|
:param rate_limit: The no. of operations allowed per unit of time.
|
|
:param seconds: The time period in seconds in which the tokens get reset or added.
|
|
:param mode: To select what happens when the period is over. 'reset' mode means that the remaining tokens from
|
|
the previous period are discarded and the counter is set back to the rate limit, and 'add' mode means that
|
|
new tokens are added on top of the exiting ones that were unused in the previous period.
|
|
:param sleep: The delay to add before checking back to see if tokens are available. Keep it longer for longer
|
|
periods (which is adjusted by the 'seconds' parameter).
|
|
:param in_sequence: Whether, or not, you want to maintain the sequence in which the requests were made
|
|
(experimental).
|
|
"""
|
|
|
|
if rate_limit is not None:
|
|
rate_limit = int(rate_limit)
|
|
self.__semaphore = asyncio.Semaphore(1 if in_sequence else rate_limit)
|
|
mode = mode.lower()
|
|
if mode not in [self.MODE_RESET, self.MODE_ADD]: mode = self.MODE_RESET
|
|
self.__mode = mode
|
|
self.__rate_limit = rate_limit
|
|
self.__token_count = rate_limit
|
|
self.__last_token_generation_time = time.time()
|
|
self.__sleep = sleep
|
|
self.__seconds = seconds
|
|
|
|
async def has_turn(self):
|
|
|
|
"""
|
|
Just to check if a token is available.
|
|
Should be used only if you don't want to wait for turn in case it is unavailable.
|
|
:return: True if the token is available, False if not.
|
|
"""
|
|
|
|
# If the user doesn't want any rate-limits, we return immediately.
|
|
# Otherwise, we check if the user has any tokens available:
|
|
if self.__rate_limit is None: return True
|
|
if self.__token_count > 0: return True
|
|
else: return False
|
|
|
|
async def get_turn(self, timeout = None):
|
|
|
|
"""
|
|
To wait till either you get your turn or the wait gets timed-out.
|
|
:param timeout: The time (in seconds) to wait to get a turn before exiting with failure.
|
|
:return: True if a turn was received, else False if timed-out.
|
|
"""
|
|
|
|
# If the user doesn't want any rate-limits, we return immediately:
|
|
if self.__rate_limit is None: return
|
|
|
|
# Start by making variables:
|
|
got_turn = False
|
|
wait_start = time.time()
|
|
|
|
# Otherwise, we wait for the semaphore:
|
|
async with self.__semaphore:
|
|
|
|
# Wait till you get your turn or the attempt gets timed-out:
|
|
while True:
|
|
|
|
# If the timeout has been exceeded:
|
|
if (
|
|
timeout is not None and
|
|
time.time() - wait_start >= timeout
|
|
): break
|
|
|
|
# If it is time to generate new tokens:
|
|
time_delta = time.time() - self.__last_token_generation_time
|
|
if time_delta >= self.__seconds:
|
|
if self.__mode == self.MODE_RESET: self.__token_count = self.__rate_limit
|
|
else: self.__token_count += int(self.__rate_limit * math.floor(time_delta / self.__seconds))
|
|
self.__last_token_generation_time = time.time()
|
|
|
|
# If a token is available (or not):
|
|
if self.__token_count > 0:
|
|
self.__token_count -= 1
|
|
got_turn = True
|
|
break
|
|
else: await asyncio.sleep(self.__sleep)
|
|
|
|
# Done here:
|
|
return got_turn
|
|
|
|
|
|
# *****************************************************************************************************************
|
|
# ***** ****
|
|
# *** MAIN PROGRAM ***
|
|
# ***** ****
|
|
# *****************************************************************************************************************
|
|
|
|
|
|
if __name__ == "__main__":
|
|
|
|
import random
|
|
|
|
async def single_test(rate_controller, count):
|
|
has_turn = await rate_controller.has_turn()
|
|
got_turn = await rate_controller.get_turn()
|
|
print(f"TURN: {count:.<5} {'Y' if has_turn else '-'} / {'Y' if got_turn else '-'} ({datetime.datetime.now()})")
|
|
if got_turn: await asyncio.sleep(random.random())
|
|
|
|
async def multi_test(max_count):
|
|
rate_controller = TokenBucket(
|
|
rate_limit = 10,
|
|
seconds = 1,
|
|
in_sequence = False,
|
|
sleep = 0.025,
|
|
mode = "reset",
|
|
timeout = 3
|
|
)
|
|
tasks = [single_test(rate_controller, count + 1) for count in range(max_count)]
|
|
await asyncio.gather(*tasks)
|
|
|
|
|
|
start_time = time.time()
|
|
asyncio.run(multi_test(100))
|
|
print(f"FINISHED IN {time.time() - start_time} SECONDS!")
|