# -*- coding: utf-8 -*- """ Token-Bucket-Algorithmus für Rate-Limiting. Implementiert den klassischen Token-Bucket-Algorithmus für gleichmäßige Durchsatzbegrenzung mit Burst-Unterstützung. """ from __future__ import annotations import threading import time from dataclasses import dataclass, field from datetime import datetime from typing import Any @dataclass class TokenBucketStats: """ Statistiken für einen Token-Bucket. Bietet Einblick in die Nutzung und Auslastung des Buckets. """ requests_allowed: int = 0 """Anzahl erlaubter Requests.""" requests_denied: int = 0 """Anzahl abgelehnter Requests.""" total_tokens_consumed: float = 0.0 """Gesamtzahl verbrauchter Tokens.""" last_request: datetime | None = None """Zeitpunkt des letzten Requests.""" @property def total_requests(self) -> int: """Gesamtzahl der Requests.""" return self.requests_allowed + self.requests_denied @property def allow_rate(self) -> float: """Erlaubnisrate in Prozent.""" if self.total_requests == 0: return 100.0 return (self.requests_allowed / self.total_requests) * 100 def record_allow(self, tokens: float = 1.0) -> None: """Zeichnet einen erlaubten Request auf.""" self.requests_allowed += 1 self.total_tokens_consumed += tokens self.last_request = datetime.now() def record_deny(self) -> None: """Zeichnet einen abgelehnten Request auf.""" self.requests_denied += 1 self.last_request = datetime.now() def reset(self) -> None: """Setzt alle Statistiken zurück.""" self.requests_allowed = 0 self.requests_denied = 0 self.total_tokens_consumed = 0.0 self.last_request = None def to_dict(self) -> dict[str, Any]: """Konvertiert die Statistiken in ein Dictionary.""" return { "requests_allowed": self.requests_allowed, "requests_denied": self.requests_denied, "total_requests": self.total_requests, "allow_rate": self.allow_rate, "total_tokens_consumed": self.total_tokens_consumed, "last_request": ( self.last_request.isoformat() if self.last_request else None ), } class TokenBucket: """ Token-Bucket-Implementierung für Rate-Limiting. Der Bucket füllt sich kontinuierlich mit Tokens bis zur maximalen Kapazität. Jeder Request verbraucht Tokens. Wenn nicht genügend Tokens verfügbar sind, wird der Request abgelehnt. Attributes: capacity: Maximale Anzahl von Tokens im Bucket. refill_rate: Tokens pro Sekunde, die hinzugefügt werden. tokens: Aktuelle Anzahl von Tokens. Example: bucket = TokenBucket(capacity=10, refill_rate=2.0) if bucket.consume(): # Request erlaubt process_request() else: # Rate limit erreicht raise TooManyRequests() """ def __init__( self, capacity: float, refill_rate: float, initial_tokens: float | None = None, ) -> None: """ Initialisiert den Token-Bucket. Args: capacity: Maximale Bucket-Kapazität. refill_rate: Tokens pro Sekunde. initial_tokens: Anfängliche Token-Anzahl (Default: capacity). """ self._capacity = capacity self._refill_rate = refill_rate self._tokens = initial_tokens if initial_tokens is not None else capacity self._last_refill = time.monotonic() self._lock = threading.Lock() self._stats = TokenBucketStats() @property def capacity(self) -> float: """Gibt die Bucket-Kapazität zurück.""" return self._capacity @property def refill_rate(self) -> float: """Gibt die Refill-Rate zurück.""" return self._refill_rate @property def tokens(self) -> float: """Gibt die aktuelle Token-Anzahl zurück (nach Refill).""" with self._lock: self._refill() return self._tokens @property def stats(self) -> TokenBucketStats: """Gibt die Statistiken zurück.""" return self._stats def _refill(self) -> None: """Füllt den Bucket basierend auf verstrichener Zeit.""" now = time.monotonic() elapsed = now - self._last_refill if elapsed > 0: new_tokens = elapsed * self._refill_rate self._tokens = min(self._capacity, self._tokens + new_tokens) self._last_refill = now def consume(self, tokens: float = 1.0) -> bool: """ Versucht, Tokens zu verbrauchen. Args: tokens: Anzahl zu verbrauchender Tokens. Returns: True wenn erfolgreich, False wenn nicht genug Tokens. """ with self._lock: self._refill() if self._tokens >= tokens: self._tokens -= tokens self._stats.record_allow(tokens) return True self._stats.record_deny() return False def try_consume(self, tokens: float = 1.0) -> tuple[bool, float]: """ Versucht, Tokens zu verbrauchen und gibt Wartezeit zurück. Args: tokens: Anzahl zu verbrauchender Tokens. Returns: Tuple (erfolgreich, Wartezeit in Sekunden). """ with self._lock: self._refill() if self._tokens >= tokens: self._tokens -= tokens self._stats.record_allow(tokens) return True, 0.0 # Berechne, wie lange auf genug Tokens gewartet werden muss needed = tokens - self._tokens wait_time = needed / self._refill_rate self._stats.record_deny() return False, wait_time def wait_and_consume( self, tokens: float = 1.0, max_wait: float | None = None, ) -> bool: """ Wartet bis Tokens verfügbar sind und verbraucht sie. Args: tokens: Anzahl zu verbrauchender Tokens. max_wait: Maximale Wartezeit in Sekunden. Returns: True wenn erfolgreich, False wenn Timeout. """ start = time.monotonic() while True: success, wait_time = self.try_consume(tokens) if success: return True # Timeout-Prüfung if max_wait is not None: elapsed = time.monotonic() - start remaining = max_wait - elapsed if remaining <= 0: return False wait_time = min(wait_time, remaining) time.sleep(wait_time) def reset(self) -> None: """Setzt den Bucket auf volle Kapazität zurück.""" with self._lock: self._tokens = self._capacity self._last_refill = time.monotonic() def drain(self) -> float: """ Leert den Bucket und gibt verbleibende Tokens zurück. Returns: Anzahl der Tokens, die im Bucket waren. """ with self._lock: self._refill() tokens = self._tokens self._tokens = 0 return tokens def __repr__(self) -> str: """String-Repräsentation.""" return ( f"TokenBucket(capacity={self._capacity}, " f"refill_rate={self._refill_rate}, " f"tokens={self.tokens:.2f})" ) class TokenBucketLimiter: """ Rate-Limiter basierend auf Token-Buckets. Verwaltet mehrere Buckets für verschiedene Schlüssel (z.B. User-IDs, IP-Adressen). """ def __init__( self, capacity: float, refill_rate: float, key_ttl: float = 3600.0, ) -> None: """ Initialisiert den Limiter. Args: capacity: Bucket-Kapazität pro Schlüssel. refill_rate: Tokens pro Sekunde pro Schlüssel. key_ttl: Zeit in Sekunden, bevor ungenutzte Buckets gelöscht werden. """ self._capacity = capacity self._refill_rate = refill_rate self._key_ttl = key_ttl self._buckets: dict[str, tuple[TokenBucket, float]] = {} self._lock = threading.Lock() def _get_bucket(self, key: str) -> TokenBucket: """Holt oder erstellt einen Bucket für einen Schlüssel.""" now = time.monotonic() with self._lock: if key in self._buckets: bucket, _ = self._buckets[key] self._buckets[key] = (bucket, now) return bucket bucket = TokenBucket( capacity=self._capacity, refill_rate=self._refill_rate, ) self._buckets[key] = (bucket, now) return bucket def is_allowed(self, key: str, tokens: float = 1.0) -> bool: """ Prüft, ob ein Request erlaubt ist. Args: key: Identifikator für den Rate-Limit-Kontext. tokens: Anzahl zu verbrauchender Tokens. Returns: True wenn erlaubt. """ bucket = self._get_bucket(key) return bucket.consume(tokens) def get_wait_time(self, key: str, tokens: float = 1.0) -> float: """ Gibt die Wartezeit bis zur Erlaubnis zurück. Args: key: Identifikator. tokens: Anzahl benötigter Tokens. Returns: Wartezeit in Sekunden (0 wenn sofort erlaubt). """ bucket = self._get_bucket(key) success, wait_time = bucket.try_consume(tokens) return 0.0 if success else wait_time def wait_and_proceed( self, key: str, tokens: float = 1.0, max_wait: float | None = None, ) -> bool: """ Wartet und verbraucht dann Tokens. Args: key: Identifikator. tokens: Anzahl zu verbrauchender Tokens. max_wait: Maximale Wartezeit. Returns: True wenn erfolgreich. """ bucket = self._get_bucket(key) return bucket.wait_and_consume(tokens, max_wait) def cleanup(self) -> int: """ Entfernt abgelaufene Buckets. Returns: Anzahl der entfernten Buckets. """ now = time.monotonic() expired_keys = [] with self._lock: for key, (_, last_access) in self._buckets.items(): if now - last_access > self._key_ttl: expired_keys.append(key) for key in expired_keys: del self._buckets[key] return len(expired_keys) def reset(self, key: str | None = None) -> None: """ Setzt einen oder alle Buckets zurück. Args: key: Optionaler spezifischer Schlüssel. """ with self._lock: if key: if key in self._buckets: bucket, last_access = self._buckets[key] bucket.reset() else: for bucket, _ in self._buckets.values(): bucket.reset() def get_stats(self, key: str) -> TokenBucketStats | None: """ Gibt Statistiken für einen Schlüssel zurück. Args: key: Identifikator. Returns: Statistiken oder None. """ with self._lock: if key in self._buckets: bucket, _ = self._buckets[key] return bucket.stats return None @property def bucket_count(self) -> int: """Gibt die Anzahl aktiver Buckets zurück.""" with self._lock: return len(self._buckets)