Squashed 'utils_v2/' content from commit ddefb8fe
git-subtree-dir: utils_v2 git-subtree-split: ddefb8fec3a72ccff2cd85e75bd6687d0067c37b
This commit is contained in:
@@ -0,0 +1,212 @@
|
||||
"""
|
||||
|
||||
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!")
|
||||
Reference in New Issue
Block a user