57 lines
1.7 KiB
Python
57 lines
1.7 KiB
Python
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import time
|
||
|
|
from collections import defaultdict, deque
|
||
|
|
from collections.abc import Callable, Iterator
|
||
|
|
from contextlib import contextmanager
|
||
|
|
from threading import Lock
|
||
|
|
|
||
|
|
from qonto_assistant.errors import ConcurrencyLimitExceededError, RateLimitExceededError
|
||
|
|
|
||
|
|
|
||
|
|
class RateLimiter:
|
||
|
|
def __init__(
|
||
|
|
self,
|
||
|
|
*,
|
||
|
|
limit: int,
|
||
|
|
window_seconds: int,
|
||
|
|
clock: Callable[[], float] = time.monotonic,
|
||
|
|
) -> None:
|
||
|
|
self.limit = limit
|
||
|
|
self.window_seconds = window_seconds
|
||
|
|
self.clock = clock
|
||
|
|
self._events: dict[str, deque[float]] = defaultdict(deque)
|
||
|
|
self._lock = Lock()
|
||
|
|
|
||
|
|
def check(self, key: str) -> None:
|
||
|
|
now = self.clock()
|
||
|
|
with self._lock:
|
||
|
|
window_start = now - self.window_seconds
|
||
|
|
bucket = self._events[key]
|
||
|
|
while bucket and bucket[0] < window_start:
|
||
|
|
bucket.popleft()
|
||
|
|
if len(bucket) >= self.limit:
|
||
|
|
raise RateLimitExceededError(f"Rate limit exceeded for {key}")
|
||
|
|
bucket.append(now)
|
||
|
|
|
||
|
|
|
||
|
|
class ConcurrencyLimiter:
|
||
|
|
def __init__(self, *, limit: int) -> None:
|
||
|
|
self.limit = limit
|
||
|
|
self._counts: dict[str, int] = defaultdict(int)
|
||
|
|
self._lock = Lock()
|
||
|
|
|
||
|
|
@contextmanager
|
||
|
|
def slot(self, key: str) -> Iterator[None]:
|
||
|
|
with self._lock:
|
||
|
|
if self._counts[key] >= self.limit:
|
||
|
|
raise ConcurrencyLimitExceededError(f"Concurrency limit exceeded for {key}")
|
||
|
|
self._counts[key] += 1
|
||
|
|
try:
|
||
|
|
yield
|
||
|
|
finally:
|
||
|
|
with self._lock:
|
||
|
|
self._counts[key] -= 1
|
||
|
|
if self._counts[key] <= 0:
|
||
|
|
del self._counts[key]
|