python-learning/rate_limiters/fixed_window.py

46 lines
1.3 KiB
Python

import threading
from datetime import datetime, timedelta
from time import sleep
class FixedWindowLimiter:
def __init__(self, window_size: int, window_duration: int):
self.window_size = window_size
self.window_duration = window_duration
self.count = 0
self.last_reset_time = datetime.now()
self._lock = threading.Lock()
def __enter__(self):
with self._lock:
if datetime.now() - self.last_reset_time > timedelta(
seconds=self.window_duration
):
self.count = 0
self.last_reset_time = datetime.now()
if self.count < self.window_size:
self.count += 1
print("ALLOWED")
return self
else:
raise RuntimeError("Rate exceeded!")
def __exit__(self, exc_type, exc_value, exc_traceback):
return False
def make_requests(number: int, limiter: FixedWindowLimiter) -> None:
for i in range(1, number):
try:
with limiter as lim:
print(f"PRE-BUDGET: {lim.window_size - lim.count}")
print("REQUEST MADE!")
except RuntimeError:
print("Request limit reached!")
sleep(0.5)
# lim_1 = FixedWindowLimiter(window_duration=10, window_size=10)
# make_requests(100, lim_1)