# -*- coding: utf-8 -*- """ Rate-Limit-Decorators für Funktionen. Bietet einfache Decorators für Durchsatzbegrenzung von Funktionen und Methoden. """ from __future__ import annotations import asyncio import functools import hashlib from dataclasses import dataclass from typing import Any, Callable, ParamSpec, TypeVar from trixy_core.utils.ratelimit.token_bucket import TokenBucket, TokenBucketLimiter from trixy_core.utils.ratelimit.sliding_window import SlidingWindow, SlidingWindowLimiter P = ParamSpec("P") R = TypeVar("R") class RateLimitExceeded(Exception): """ Exception wenn Rate-Limit überschritten wurde. Enthält Informationen zur Wartezeit bis zur nächsten erlaubten Ausführung. """ def __init__( self, message: str = "Rate limit exceeded", retry_after: float | None = None, limit_key: str | None = None, ) -> None: """ Initialisiert die Exception. Args: message: Fehlermeldung. retry_after: Zeit in Sekunden bis zur nächsten Erlaubnis. limit_key: Identifikator des Rate-Limits. """ super().__init__(message) self.retry_after = retry_after self.limit_key = limit_key def __str__(self) -> str: """String-Repräsentation.""" msg = super().__str__() if self.retry_after: msg += f" (retry after {self.retry_after:.2f}s)" return msg @dataclass class RateLimitConfig: """ Konfiguration für Rate-Limiting. Definiert die Parameter für den Rate-Limiter. """ # Token-Bucket-Parameter capacity: float = 10.0 """Maximale Bucket-Kapazität (Token-Bucket).""" refill_rate: float = 1.0 """Tokens pro Sekunde (Token-Bucket).""" # Sliding-Window-Parameter window_size: float = 60.0 """Fenstergröße in Sekunden (Sliding-Window).""" max_requests: int = 60 """Maximale Requests pro Fenster (Sliding-Window).""" # Allgemeine Parameter algorithm: str = "token_bucket" """Algorithmus: 'token_bucket' oder 'sliding_window'.""" key_func: Callable[..., str] | None = None """Funktion zur Schlüsselgenerierung für Per-Key-Limiting.""" block: bool = False """Ob gewartet werden soll statt Exception zu werfen.""" max_wait: float | None = None """Maximale Wartezeit wenn block=True.""" on_limited: Callable[[str], None] | None = None """Callback wenn Limit erreicht.""" # Globale Limiter-Instanzen _token_bucket_limiters: dict[str, TokenBucketLimiter] = {} _sliding_window_limiters: dict[str, SlidingWindowLimiter] = {} def _get_key( func: Callable, args: tuple, kwargs: dict, key_func: Callable[..., str] | None, ) -> str: """Generiert einen Schlüssel für die Rate-Limit-Instanz.""" if key_func: return key_func(*args, **kwargs) # Default: Funktionsname als Schlüssel return func.__qualname__ def rate_limit( capacity: float = 10.0, refill_rate: float = 1.0, window_size: float = 60.0, max_requests: int = 60, algorithm: str = "token_bucket", key_func: Callable[..., str] | None = None, block: bool = False, max_wait: float | None = None, on_limited: Callable[[str], None] | None = None, config: RateLimitConfig | None = None, ) -> Callable[[Callable[P, R]], Callable[P, R]]: """ Decorator für Rate-Limiting von Funktionen. Args: capacity: Token-Bucket-Kapazität. refill_rate: Token-Bucket Refill-Rate. window_size: Sliding-Window-Größe. max_requests: Sliding-Window Maximum. algorithm: 'token_bucket' oder 'sliding_window'. key_func: Optionale Funktion für Per-Key-Limiting. block: Ob gewartet werden soll. max_wait: Maximale Wartezeit. on_limited: Callback bei Limit-Erreichen. config: Alternative Konfiguration als Objekt. Returns: Decorator-Funktion. Example: @rate_limit(capacity=10, refill_rate=2) def api_call(): return requests.get(url) @rate_limit( algorithm="sliding_window", window_size=60, max_requests=100, key_func=lambda user_id: str(user_id) ) def user_action(user_id): pass @rate_limit(block=True, max_wait=5.0) def blocking_call(): # Wartet bis zu 5 Sekunden pass """ if config: cfg = config else: cfg = RateLimitConfig( capacity=capacity, refill_rate=refill_rate, window_size=window_size, max_requests=max_requests, algorithm=algorithm, key_func=key_func, block=block, max_wait=max_wait, on_limited=on_limited, ) def decorator(func: Callable[P, R]) -> Callable[P, R]: limiter_name = func.__qualname__ # Limiter erstellen if cfg.algorithm == "token_bucket": if limiter_name not in _token_bucket_limiters: _token_bucket_limiters[limiter_name] = TokenBucketLimiter( capacity=cfg.capacity, refill_rate=cfg.refill_rate, ) limiter = _token_bucket_limiters[limiter_name] else: if limiter_name not in _sliding_window_limiters: _sliding_window_limiters[limiter_name] = SlidingWindowLimiter( window_size=cfg.window_size, max_requests=cfg.max_requests, ) limiter = _sliding_window_limiters[limiter_name] @functools.wraps(func) def wrapper(*args: P.args, **kwargs: P.kwargs) -> R: key = _get_key(func, args, kwargs, cfg.key_func) if cfg.block: # Warten und dann ausführen if isinstance(limiter, TokenBucketLimiter): success = limiter.wait_and_proceed(key, max_wait=cfg.max_wait) else: success = limiter.wait_and_proceed(key, max_wait=cfg.max_wait) if not success: if cfg.on_limited: cfg.on_limited(key) raise RateLimitExceeded( "Rate limit exceeded after waiting", limit_key=key, ) else: # Sofortige Prüfung if not limiter.is_allowed(key): if cfg.on_limited: cfg.on_limited(key) # Retry-After berechnen if isinstance(limiter, TokenBucketLimiter): retry_after = limiter.get_wait_time(key) else: retry_after = limiter.get_retry_after(key) raise RateLimitExceeded( "Rate limit exceeded", retry_after=retry_after, limit_key=key, ) return func(*args, **kwargs) # Zugriff auf Limiter und Stats wrapper.limiter = limiter # type: ignore wrapper.reset = lambda k=None: limiter.reset(k) # type: ignore return wrapper return decorator def async_rate_limit( capacity: float = 10.0, refill_rate: float = 1.0, window_size: float = 60.0, max_requests: int = 60, algorithm: str = "token_bucket", key_func: Callable[..., str] | None = None, block: bool = False, max_wait: float | None = None, on_limited: Callable[[str], None] | None = None, config: RateLimitConfig | None = None, ) -> Callable[[Callable[P, R]], Callable[P, R]]: """ Async-Decorator für Rate-Limiting. Wie @rate_limit, aber für async Funktionen. Args: Siehe rate_limit. Returns: Decorator-Funktion. Example: @async_rate_limit(capacity=10, refill_rate=2) async def api_call(): return await aiohttp_session.get(url) """ if config: cfg = config else: cfg = RateLimitConfig( capacity=capacity, refill_rate=refill_rate, window_size=window_size, max_requests=max_requests, algorithm=algorithm, key_func=key_func, block=block, max_wait=max_wait, on_limited=on_limited, ) def decorator(func: Callable[P, R]) -> Callable[P, R]: limiter_name = func.__qualname__ if cfg.algorithm == "token_bucket": if limiter_name not in _token_bucket_limiters: _token_bucket_limiters[limiter_name] = TokenBucketLimiter( capacity=cfg.capacity, refill_rate=cfg.refill_rate, ) limiter = _token_bucket_limiters[limiter_name] else: if limiter_name not in _sliding_window_limiters: _sliding_window_limiters[limiter_name] = SlidingWindowLimiter( window_size=cfg.window_size, max_requests=cfg.max_requests, ) limiter = _sliding_window_limiters[limiter_name] @functools.wraps(func) async def wrapper(*args: P.args, **kwargs: P.kwargs) -> R: key = _get_key(func, args, kwargs, cfg.key_func) if cfg.block: # Async warten start_wait = 0.0 while True: if limiter.is_allowed(key): break if isinstance(limiter, TokenBucketLimiter): wait_time = limiter.get_wait_time(key) else: wait_time = limiter.get_retry_after(key) if cfg.max_wait is not None: if start_wait + wait_time > cfg.max_wait: if cfg.on_limited: cfg.on_limited(key) raise RateLimitExceeded( "Rate limit exceeded after waiting", limit_key=key, ) await asyncio.sleep(min(wait_time, 0.1)) start_wait += min(wait_time, 0.1) else: if not limiter.is_allowed(key): if cfg.on_limited: cfg.on_limited(key) if isinstance(limiter, TokenBucketLimiter): retry_after = limiter.get_wait_time(key) else: retry_after = limiter.get_retry_after(key) raise RateLimitExceeded( "Rate limit exceeded", retry_after=retry_after, limit_key=key, ) return await func(*args, **kwargs) wrapper.limiter = limiter # type: ignore wrapper.reset = lambda k=None: limiter.reset(k) # type: ignore return wrapper return decorator def throttle( calls: int, period: float, key_func: Callable[..., str] | None = None, ) -> Callable[[Callable[P, R]], Callable[P, R]]: """ Vereinfachter Decorator für Throttling. Erlaubt eine bestimmte Anzahl von Aufrufen pro Zeitperiode. Args: calls: Maximale Anzahl Aufrufe. period: Zeitperiode in Sekunden. key_func: Optionale Funktion für Per-Key-Throttling. Returns: Decorator-Funktion. Example: @throttle(calls=10, period=60) # 10 Aufrufe pro Minute def api_call(): pass """ return rate_limit( algorithm="sliding_window", window_size=period, max_requests=calls, key_func=key_func, ) def debounce( wait: float, ) -> Callable[[Callable[P, R]], Callable[P, R]]: """ Decorator für Debouncing. Verhindert mehrfache Aufrufe innerhalb einer kurzen Zeitspanne. Nur der letzte Aufruf innerhalb des Wartezeitraums wird ausgeführt. Args: wait: Wartezeit in Sekunden. Returns: Decorator-Funktion. Example: @debounce(0.5) def on_input_change(value): # Wird nur ausgeführt wenn 0.5s keine weiteren Aufrufe process(value) """ import time import threading def decorator(func: Callable[P, R]) -> Callable[P, R]: last_call = [0.0] timer: list[threading.Timer | None] = [None] lock = threading.Lock() last_args: list[tuple] = [()] last_kwargs: list[dict] = [{}] @functools.wraps(func) def wrapper(*args: P.args, **kwargs: P.kwargs) -> None: with lock: last_args[0] = args last_kwargs[0] = kwargs if timer[0]: timer[0].cancel() def execute(): func(*last_args[0], **last_kwargs[0]) timer[0] = threading.Timer(wait, execute) timer[0].start() return wrapper # type: ignore return decorator