"""Limiter class implementation""" import asyncio import logging from contextlib import contextmanager from functools import wraps from inspect import isawaitable, iscoroutine, iscoroutinefunction from threading import RLock, local from time import monotonic, sleep from typing import Any, Awaitable, Callable, Iterable, List, Optional, Protocol, Tuple, Union from .abstracts import AbstractBucket, BucketFactory, Rate, RateItem from .buckets import InMemoryBucket logger = logging.getLogger(__name__) ItemMapping = Callable[[Any], Tuple[str, int]] DecoratorWrapper = Callable[[Callable[[Any], Any]], Callable[[Any], Any]] class LockLike(Protocol): def acquire(self, blocking: bool = ..., timeout: Union[float, int, None] = ...) -> bool: ... def release(self) -> None: ... class SingleBucketFactory(BucketFactory): """Single-bucket factory for quick use with Limiter""" bucket: AbstractBucket def __init__(self, bucket: AbstractBucket, schedule_leak: bool = True): """ Initialize the SingleBucketFactory with a bucket and an optional leak scheduling flag. schedule_leak (bool): If True, the factory will schedule periodic leaks for the bucket. Default is True. Disable only if you plan to handle leaking manually. """ self.bucket = bucket if schedule_leak: self.schedule_leak(bucket) def wrap_item(self, name: str, weight: int = 1): now = self.bucket.now() async def wrap_async(): return RateItem(name, await now, weight=weight) def wrap_sync(): return RateItem(name, now, weight=weight) return wrap_async() if isawaitable(now) else wrap_sync() def get(self, _: RateItem) -> AbstractBucket: return self.bucket @contextmanager def combined_lock(locks: Union[Iterable[LockLike], RLock], blocking: bool, timeout: int | float = -1): if not isinstance(locks, Iterable): acquired_ok = locks.acquire(blocking=blocking, timeout=timeout) if not acquired_ok: raise TimeoutError("acquire failed") try: yield finally: locks.release() else: acquired_locks = [] deadline = monotonic() + timeout if blocking and timeout != -1 else None try: for lock in locks: if not blocking: ok = lock.acquire(False) elif timeout == -1: ok = lock.acquire() else: assert deadline is not None remaining = max(0.0, deadline - monotonic()) ok = lock.acquire(timeout=remaining) if not ok: raise TimeoutError("Timeout while acquiring combined lock.") acquired_locks.append(lock) yield finally: for lock in reversed(acquired_locks): lock.release() class Limiter: """This class responsibility is to sum up all underlying logic and make working with async/sync functions easily """ bucket_factory: BucketFactory lock: Union[RLock, Iterable] buffer_ms: int # async_lock is thread local, created on first use _thread_local: local def __init__( self, argument: Union[BucketFactory, AbstractBucket, Rate, List[Rate]], buffer_ms: int = 50, ): """Init Limiter using either a single bucket / multiple-bucket factory / single rate / rate list. Parameters: argument (Union[BucketFactory, AbstractBucket, Rate, List[Rate]]): The bucket or rate configuration. """ self.buffer_ms = buffer_ms self.bucket_factory = self._init_bucket_factory(argument) self.lock = RLock() self._thread_local = local() if isinstance(argument, AbstractBucket): limiter_lock = argument.limiter_lock() if limiter_lock is not None: self.lock = (limiter_lock, self.lock) def buckets(self) -> List[AbstractBucket]: """Get list of active buckets""" return self.bucket_factory.get_buckets() def dispose(self, bucket: Union[int, AbstractBucket]) -> bool: """Dispose/Remove a specific bucket, using bucket-id or bucket object as param """ return self.bucket_factory.dispose(bucket) def _init_bucket_factory( self, argument: Union[BucketFactory, AbstractBucket, Rate, List[Rate]], ) -> BucketFactory: if isinstance(argument, Rate): argument = [argument] if isinstance(argument, list): assert len(argument) > 0, "Rates must not be empty" assert isinstance(argument[0], Rate), "Not valid rates list" rates = argument logger.info("Initializing default bucket(InMemoryBucket) with rates: %s", rates) argument = InMemoryBucket(rates) if isinstance(argument, AbstractBucket): argument = SingleBucketFactory(argument) assert isinstance(argument, BucketFactory), "Not a valid bucket/bucket-factory" return argument def _delay_waiter( self, bucket: AbstractBucket, item: RateItem, blocking: bool, _force_async: bool = False, deadline: Optional[float] = None ) -> Union[bool, Awaitable[bool]]: """On `try_acquire` failed, handle delay""" assert bucket.failing_rate is not None if not blocking: return False delay = bucket.waiting(item) if _force_async or isawaitable(delay): async def _handle_async(delay): while True: if deadline is not None: remaining_ms = (deadline - monotonic()) * 1000 if remaining_ms <= 0: raise TimeoutError() d = await self._handle_async_result(delay, deadline=deadline) if isawaitable(delay) else delay assert isinstance(d, int) if d == -1: return False assert d >= 0 d += self.buffer_ms if deadline is not None: remaining_ms = (deadline - monotonic()) * 1000 if remaining_ms <= 0: raise TimeoutError() if remaining_ms < d: await asyncio.sleep(remaining_ms / 1000) raise TimeoutError() await asyncio.sleep(d / 1000) else: await asyncio.sleep(d / 1000) item.timestamp += d r = bucket.put(item) r = await r if isawaitable(r) else r if r: return True delay = bucket.waiting(item) return _handle_async(delay) else: total_delay = 0 while True: assert not isawaitable(delay) logger.debug("delay=%d, total_delay=%s", delay, total_delay) if delay == -1: return False delay += self.buffer_ms total_delay += delay if deadline is not None: remaining_ms = (deadline - monotonic()) * 1000 if remaining_ms <= 0: raise TimeoutError() if remaining_ms < delay: sleep(remaining_ms / 1000) raise TimeoutError() sleep(delay / 1000) else: sleep(delay / 1000) item.timestamp += delay re_acquire = bucket.put(item) if isawaitable(re_acquire): async def _resume_async(re_acquire): reacquired = await self._handle_async_result(re_acquire, deadline=deadline) if reacquired: return True continued = self._delay_waiter( bucket, item, blocking=blocking, _force_async=True, deadline=deadline, ) return await self._handle_async_result(continued, deadline=deadline) return _resume_async(re_acquire) assert isinstance(re_acquire, bool) if re_acquire: return True delay = bucket.waiting(item) def handle_bucket_put( self, bucket: AbstractBucket, item: RateItem, blocking: bool, _force_async: bool = False, deadline: Optional[float] = None ) -> Union[bool, Awaitable[bool]]: """Putting item into bucket""" def _handle_result(is_success: bool): if not is_success: return self._delay_waiter(bucket, item, blocking=blocking, _force_async=_force_async, deadline=deadline) return True acquire = bucket.put(item) if isawaitable(acquire): async def _put_async(acquire): acquire_result = await self._handle_async_result(acquire, deadline=deadline) if acquire_result: return True result = self._delay_waiter( bucket, item, blocking=blocking, _force_async=True, deadline=deadline, ) return await self._handle_async_result(result, deadline=deadline) return _put_async(acquire) return _handle_result(acquire) # type: ignore def _get_async_lock(self): """Returns thread_local, loop-specific lock""" loop = asyncio.get_running_loop() try: # The async loop *can* change in a given thread lock = self._thread_local.async_lock if self._thread_local.async_lock_loop is loop: return lock except AttributeError: pass lock = asyncio.Lock() self._thread_local.async_lock = lock self._thread_local.async_lock_loop = loop return lock def try_acquire(self, name: str = "pyrate", weight: int = 1, blocking: bool = True, timeout: int | float = -1) -> Union[bool, Awaitable[bool]]: """ Attempt to acquire a permit from the limiter. Parameters ---------- name : str, default "pyrate" The bucket key to acquire from. weight : int, default 1 Number of permits to consume. timeout : int | float, default -1 Maximum time (in seconds) to wait; -1 means wait indefinitely. blocking : bool, default True If True, block until a permit is available (subject to timeout); if False, return immediately. Returns ------- bool or Awaitable[bool] True if the permit was acquired, False otherwise. Async limiters return an awaitable resolving to the same. """ if timeout < 0 and timeout != -1: raise ValueError("timeout must be -1 or >= 0") if not blocking and timeout != -1: raise RuntimeError("Can't set timeout with non-blocking") try: result = self._try_acquire(name=name, weight=weight, timeout=timeout, blocking=blocking) except TimeoutError: logger.debug("Acquisition TimeoutError") return False if not isawaitable(result): return result async def _resolve_result(async_result: Awaitable[bool]) -> bool: try: return await self._handle_async_result(async_result) except TimeoutError: logger.debug("Acquisition TimeoutError") return False return _resolve_result(result) async def _acquire_async(self, blocking, name, weight): return await self._handle_async_result(self._try_acquire(name, weight, blocking=blocking, _force_async=True)) async def try_acquire_async(self, name: str = "pyrate", weight: int = 1, blocking: bool = True, timeout: int | float = -1) -> bool: """ Attempt to asynchronously acquire a permit from the limiter. Parameters ---------- name : str, default "pyrate" The bucket key to acquire from. weight : int, default 1 Number of permits to consume. blocking : bool, default True If True, wait until a permit is available (subject to timeout); if False, return immediately. timeout : int | float, default -1 Maximum time (in seconds) to wait; -1 means wait indefinitely. Returns ------- bool True if the permit was acquired, False otherwise. Notes ----- This is the async variant of ``try_acquire``. A top-level, thread-local async lock is used to prevent blocking the event loop. """ if weight == 0: return True if timeout < 0 and timeout != -1: raise ValueError("timeout must be -1 or >= 0") if not blocking and timeout != -1: raise RuntimeError("Can't set timeout with non-blocking") async def run(): lock = self._get_async_lock() async with lock: return await self._acquire_async(blocking=blocking, name=name, weight=weight) if timeout == -1: return await run() try: return await asyncio.wait_for(run(), timeout=timeout) except asyncio.TimeoutError: return False async def _handle_async_acquire( self, item: Awaitable[RateItem], blocking: bool, _force_async: bool = False, deadline: Optional[float] = None, ): this_item = await self._handle_async_result(item, deadline=deadline) assert isinstance(this_item, RateItem) bucket = self.bucket_factory.get(this_item) if isawaitable(bucket): bucket = await self._handle_async_result(bucket, deadline=deadline) assert isinstance(bucket, AbstractBucket), f"Invalid bucket: item: {this_item.name}" result = self.handle_bucket_put(bucket, this_item, blocking=blocking, _force_async=_force_async, deadline=deadline) return await self._handle_async_result(result, deadline=deadline) async def _handle_async_bucket( self, bucket: Awaitable[AbstractBucket], item: RateItem, blocking: bool, _force_async: bool = False, deadline: Optional[float] = None, ): this_bucket = await self._handle_async_result(bucket, deadline=deadline) assert isinstance(this_bucket, AbstractBucket), f"Invalid bucket: item: {item.name}" result = self.handle_bucket_put(this_bucket, item, blocking=blocking, _force_async=_force_async, deadline=deadline) return await self._handle_async_result(result, deadline=deadline) async def _handle_async_result(self, result, deadline: Optional[float] = None): timed_out = False def _cancel_if_pending(awaitable: Any) -> None: if asyncio.isfuture(awaitable) and not awaitable.done(): awaitable.cancel() try: while isawaitable(result): if deadline is None: result = await result continue remaining = deadline - monotonic() if remaining <= 0: timed_out = True _cancel_if_pending(result) raise TimeoutError() try: result = await asyncio.wait_for(result, timeout=remaining) except asyncio.TimeoutError as exc: timed_out = True _cancel_if_pending(result) raise TimeoutError() from exc return result finally: if timed_out: _cancel_if_pending(result) if iscoroutine(result): result.close() def _cleanup_awaitable(self, awaitable: Any) -> None: if asyncio.isfuture(awaitable): awaitable.cancel() return if iscoroutine(awaitable): awaitable.close() return cancel = getattr(awaitable, "cancel", None) if callable(cancel): cancel() close = getattr(awaitable, "close", None) if callable(close): close() def _try_acquire( self, name: str, weight: int, blocking: bool, timeout: int | float = -1, _force_async: bool = False, _allow_async_result: bool = True, ) -> Union[bool, Awaitable[bool]]: """Try acquiring an item with name & weight Return true on success, false on failure """ deadline: Optional[float] = monotonic() + timeout if timeout != -1 else None with combined_lock(self.lock, blocking=blocking, timeout=timeout): assert weight >= 0, "item's weight must be >= 0" if weight == 0: # NOTE: if item is weightless, just let it go through # NOTE: this might change in the future return True item = self.bucket_factory.wrap_item(name, weight) if isawaitable(item): if not _force_async and not _allow_async_result: self._cleanup_awaitable(item) raise RuntimeError("Can't use async bucket with sync decorator") return self._handle_async_acquire(item, blocking=blocking, _force_async=_force_async, deadline=deadline) assert isinstance(item, RateItem) bucket = self.bucket_factory.get(item) if isawaitable(bucket): if not _force_async and not _allow_async_result: self._cleanup_awaitable(bucket) raise RuntimeError("Can't use async bucket with sync decorator") return self._handle_async_bucket(bucket=bucket, item=item, blocking=blocking, _force_async=_force_async, deadline=deadline) assert isinstance(bucket, AbstractBucket), f"Invalid bucket: item: {name}" if ( not _force_async and not _allow_async_result and self.bucket_factory._leaker and id(bucket) in self.bucket_factory._leaker.async_buckets ): raise RuntimeError("Can't use async bucket with sync decorator") result = self.handle_bucket_put(bucket, item, blocking=blocking, _force_async=_force_async, deadline=deadline) if isawaitable(result): if not _force_async and not _allow_async_result: self._cleanup_awaitable(result) raise RuntimeError("Can't use async bucket with sync decorator") return self._handle_async_result(result, deadline=deadline) return result def as_decorator(self, *, name="ratelimiter", weight=1): def deco(func: Callable[..., Any]) -> Callable[..., Any]: if iscoroutinefunction(func): @wraps(func) async def wrapper(*args, **kwargs): r = await self.try_acquire_async(name=name, weight=weight) while isawaitable(r): r = await r return await func(*args, **kwargs) return wrapper else: @wraps(func) def wrapper(*args, **kwargs): try: r = self._try_acquire(name=name, weight=weight, blocking=True, _allow_async_result=False) except TimeoutError: r = False if isawaitable(r): try: self._cleanup_awaitable(r) finally: raise RuntimeError("Can't use async bucket with sync decorator") return func(*args, **kwargs) return wrapper return deco def close(self) -> None: self.bucket_factory.close() def __enter__(self): return self def __exit__(self, exc_type, exc, tb) -> None: self.close() def __getstate__(self): """Get state for pickling""" state = self.__dict__.copy() state.pop("lock", None) state.pop("_thread_local", None) return state def __setstate__(self, state): """Restore state after unpickling""" self.__dict__.update(state) self.lock = RLock() self._thread_local = local()