スライディングウィンドウと優先度付きキューを備えた並行レートリミッタを実装する

```python import time import threading import collections import heapq import uuid from dataclasses import dataclass, field from typing import Dict, List, Tuple, Any, Optional # --- 設定 --- # 各ティアのレート制限とウィンドウサイズを定義します。 # window_seconds が小さいほど、レート制限は厳しくなります。 TIER_CONFIG = { "free": {"max_requests": 5, "window_seconds": 60}, "standard": {"max_requests": 50, "window_seconds": 60}, "premium": {"max_requests": 200, "window_seconds": 60}, } # アクティビティがない場合にクライアントデータを削除するまでの保持期間。 CLEANUP_THRESHOLD_SECONDS = 300 # --- データ構造 --- @dataclass(order=True) class Request: """優先度付きリクエストを表します。優先度キュー用です。""" # heapq が優先度でソートできるように、優先度を最初の要素として使用します。 # 数値が小さいほど優先度が高くなります。 priority: int # 優先度のタイブレークと安定ソートを保証するために、一意のIDを使用します。 request_id: str = field(default_factory=lambda: str(uuid.uuid4())) client_id: str timestamp: float payload: Any def __lt__(self, other): """heapq のカスタム比較: 優先度でソートし、次に request_id で安定ソートします。""" if self.priority != other.priority: return self.priority < other.priority return self.request_id < other.request_id class SlidingWindowCounter: """スライディングウィンドウカウンターをレート制限のために実装します。""" def __init__(self, max_requests: int, window_seconds: float): self.max_requests = max_requests self.window_seconds = window_seconds # ウィンドウ内のリクエストのタイムスタンプを格納するために deque を使用します。 # deque は両端からの追加と削除が効率的です。 self.request_timestamps: collections.deque[float] = collections.deque() # request_timestamps へのスレッドセーフなアクセス用のロック。 self.lock = threading.Lock() def record_request(self) -> bool: """リクエストを記録し、制限内であれば True を、そうでなければ False を返します。""" with self.lock: current_time = time.time() # ウィンドウより古いタイムスタンプを削除します。 while self.request_timestamps and self.request_timestamps[0] <= current_time - self.window_seconds: self.request_timestamps.popleft() # 新しいリクエストを追加すると制限を超えるかどうかを確認します。 if len(self.request_timestamps) < self.max_requests: self.request_timestamps.append(current_time) return True else: return False def get_current_count(self) -> int: """ウィンドウ内の現在リクエスト数を返します。""" with self.lock: current_time = time.time() # カウントを返す前に古いタイムスタンプをクリーンアップします。 while self.request_timestamps and self.request_timestamps[0] <= current_time - self.window_seconds: self.request_timestamps.popleft() return len(self.request_timestamps) class RateLimiter: """複数のティア、スライディングウィンドウ、優先度付きキューを備えたスレッドセーフなレートリミッターです。""" def __init__(self, tier_config: Dict[str, Dict[str, Any]], cleanup_threshold_seconds: float): self.tier_config = tier_config self.cleanup_threshold_seconds = cleanup_threshold_seconds # 各ティアのリクエストキュー（遅延リクエスト用）を格納します。 self.tier_queues: Dict[str, List[Request]] = collections.defaultdict(list) # クリーンアップ用の各クライアントの最終アクティビティタイムスタンプを格納します。 self.client_last_activity: Dict[str, float] = {} # 各クライアントに割り当てられたティアを格納します。 self.client_tiers: Dict[str, str] = {} # スライディングウィンドウカウンターを格納します。 self.tier_counters: Dict[str, SlidingWindowCounter] = {} # クライアントティア、クライアント最終アクティビティ、ティアキューなどの共有データ構造を保護するためのメインロック。 # SlidingWindowCounter は独自の内部ロックを持っています。 self.global_lock = threading.Lock() # ティアカウンターを初期化します。 for tier_name, config in tier_config.items(): self.tier_counters[tier_name] = SlidingWindowCounter( max_requests=config["max_requests"], window_seconds=config["window_seconds"] ) # クリーンアップ用のバックグラウンドスレッドを開始します。 self.cleanup_thread = threading.Thread(target=self._cleanup_task, daemon=True) self.cleanup_thread.start() def register_client(self, client_id: str, tier: str) -> bool: """クライアントを指定されたティアで登録します。成功した場合は True を、クライアントが既に存在するかティアが無効な場合は False を返します。""" if tier not in self.tier_config: print(f"[!] エラー: ティア '{tier}' は存在しません。") return False with self.global_lock: if client_id in self.client_tiers: print(f"[!] 警告: クライアント '{client_id}' は既に登録されています。") return False self.client_tiers[client_id] = tier self.client_last_activity[client_id] = time.time() print(f"[*] クライアント '{client_id}' がティア '{tier}' で登録されました。") return True def _get_client_tier(self, client_id: str) -> Optional[str]: """指定された client_id のティアを安全に取得します。""" with self.global_lock: return self.client_tiers.get(client_id) def _update_client_activity(self, client_id: str): """クライアントの最終アクティビティタイムスタンプを更新します。""" with self.global_lock: self.client_last_activity[client_id] = time.time() def allow_request(self, client_id: str) -> bool: """クライアントのティアのレート制限に基づいて、client_id からのリクエストが許可されるかどうかを確認します。""" tier = self._get_client_tier(client_id) if not tier: print(f"[!] エラー: クライアント '{client_id}' は登録されていません。") return False counter = self.tier_counters.get(tier) if not counter: # tier_config が有効で初期化が正しければ、これは発生しないはずです。 print(f"[!] 内部エラー: ティア '{tier}' のカウンターが見つかりません。") return False if counter.record_request(): self._update_client_activity(client_id) # print(f"[*] リクエスト許可: クライアント '{client_id}' (ティア: {tier})。現在のカウント: {counter.get_current_count()}") return True else: # print(f"[*] リクエスト拒否: クライアント '{client_id}' (ティア: {tier})。現在のカウント: {counter.get_current_count()}") return False def enqueue_request(self, client_id: str, priority: int, payload: Any): """レート制限されたリクエストを適切な優先度付きキューにエンキューします。""" tier = self._get_client_tier(client_id) if not tier: print(f"[!] エラー: クライアント '{client_id}' は登録されていません。エンキューできません。") return request = Request(priority=priority, client_id=client_id, timestamp=time.time(), payload=payload) with self.global_lock: heapq.heappush(self.tier_queues[tier], request) print(f"[*] リクエストエンキュー: クライアント '{client_id}' (ティア: {tier}, 優先度: {priority})。キューサイズ: {len(self.tier_queues[tier])}") self._update_client_activity(client_id) def dequeue_and_process(self, client_id: str) -> Optional[Tuple[Request, bool]]: """クライアントのティアで利用可能な容量がある場合、最優先度リクエストをデキューして処理しようとします。""" tier = self._get_client_tier(client_id) if not tier: print(f"[!] エラー: クライアント '{client_id}' は登録されていません。デキューできません。") return None counter = self.tier_counters.get(tier) if not counter: print(f"[!] 内部エラー: ティア '{tier}' のカウンターが見つかりません。") return None # デキューを試みる前に、容量が利用可能かどうかを確認します。 # 容量がない場合に不要なロックとヒープ操作を防ぎます。 if counter.get_current_count() < counter.max_requests: with self.global_lock: queue = self.tier_queues[tier] if queue: # ロック内で容量を再確認します。競合状態を処理するためです。 # 他のスレッドがちょうど容量を使い切った可能性があります。 if counter.get_current_count() < counter.max_requests: highest_priority_request = heapq.heappop(queue) # リクエストを再度記録しようとします。容量があれば成功するはずです。 if counter.record_request(): self._update_client_activity(highest_priority_request.client_id) print(f"[*] リクエストデキューと処理完了: クライアント '{highest_priority_request.client_id}' (ティア: {tier}, 優先度: {highest_priority_request.priority})。キューサイズ: {len(queue)}") return highest_priority_request, True else: # これはまれな競合状態です: 容量は利用可能でしたが、record_request が失敗しました。 # リクエストを戻し、エラーを記録します。 heapq.heappush(queue, highest_priority_request) print(f"[!] 競合状態: {highest_priority_request.client_id} のデキューされたリクエストの記録に失敗しました。戻しました。") return highest_priority_request, False else: # 外側のチェックと内側のチェックの間に、別のスレッドが容量を占有しました。 # ポップされた場合、リクエストをキューに戻します。 heapq.heappush(queue, highest_priority_request) # print(f"[*] {client_id} のデキュー試行失敗: 容量が他のスレッドによって満たされました。") return highest_priority_request, False else: # キューは空です。 # print(f"[*] {client_id} のデキュー試行: キューは空です。") return None, False else: # 容量が利用できません。 # print(f"[*] {client_id} のデキュー試行: 容量が利用できません。") return None, False def _cleanup_task(self): """非アクティブなクライアントデータを削除するバックグラウンドタスクです。""" while True: time.sleep(self.cleanup_threshold_seconds / 2) # 定期的にチェック current_time = time.time() clients_to_remove = [] with self.global_lock: for client_id, last_activity in self.client_last_activity.items(): if current_time - last_activity > self.cleanup_threshold_seconds: clients_to_remove.append(client_id) for client_id in clients_to_remove: tier = self.client_tiers.pop(client_id, None) if tier: self.client_last_activity.pop(client_id, None) # 必要に応じて、このクライアントのキューをクリアすることもできますが、通常は行われません。 # 簡単のため、エンキューされたリクエストはそのままにしておきます。 print(f"[*] 非アクティブなクライアント '{client_id}' をクリーンアップしました。") # --- デモンストレーションスクリプト --- def simulate_requests(rate_limiter: RateLimiter, client_id: str, num_requests: int, delay: float, priority_offset: int = 0): """単一クライアントからのリクエストのバーストをシミュレートします。""" print(f"--- クライアント '{client_id}' から {num_requests} 件のリクエストをシミュレート中 ---") for i in range(num_requests): # エンキューされたリクエストの優先度を変化させます priority = i % 5 + priority_offset # 数値が小さいほど優先度が高い payload = f"data_{client_id}_{i}" if rate_limiter.allow_request(client_id): print(f"[スレッド {threading.current_thread().name}] リクエスト {i+1}/{num_requests} 許可: '{client_id}'。ペイロード: {payload}") else: print(f"[スレッド {threading.current_thread().name}] リクエスト {i+1}/{num_requests} 拒否: '{client_id}'。優先度 {priority} でエンキューします。") rate_limiter.enqueue_request(client_id, priority, payload) time.sleep(delay) def process_deferred_requests(rate_limiter: RateLimiter, client_id: str, num_attempts: int): """クライアントの遅延リクエストを処理しようと定期的に試みます。""" print(f"--- クライアント '{client_id}' の遅延リクエストの処理を試行中 ---") for i in range(num_attempts): time.sleep(1.5) # 容量が解放されるのを少し待ちます print(f"[プロセッサースレッド] '{client_id}' のデキュー試行 {i+1}/{num_attempts}...") result = rate_limiter.dequeue_and_process(client_id) if result: request, success = result if success: print(f"[プロセッサースレッド] 正常に処理されました: {request.payload} (優先度: {request.priority})") else: print(f"[プロセッサースレッド] 処理に失敗しました (容量の問題？): {request.payload} (優先度: {request.priority})。再エンキューしました。") else: print(f"[プロセッサースレッド] '{client_id}' の遅延リクエストはこの試行では処理されませんでした (キュー空または容量なし)。") if __name__ == "__main__": print("--- レートリミッターデモンストレーション ---") # レートリミッターを初期化します limiter = RateLimiter(tier_config=TIER_CONFIG, cleanup_threshold_seconds=CLEANUP_THRESHOLD_SECONDS) # クライアントを登録します clients = { "client_A": "free", "client_B": "standard", "client_C": "premium", "client_D": "free", "client_E": "standard", } for cid, tier in clients.items(): limiter.register_client(cid, tier) print("\n--- 同時リクエストのシミュレーション ---") # 複数のスレッドからリクエストのバーストをシミュレートします threads = [] # クライアント A (無料ティア) - レート制限に達する可能性が高い t1 = threading.Thread(target=simulate_requests, args=(limiter, "client_A", 15, 0.1, 0), name="Thread-A") threads.append(t1) # クライアント B (標準ティア) - より多くの容量 t2 = threading.Thread(target=simulate_requests, args=(limiter, "client_B", 60, 0.05, 1), name="Thread-B") threads.append(t2) # クライアント C (プレミアムティア) - 最も高い容量 t3 = threading.Thread(target=simulate_requests, args=(limiter, "client_C", 250, 0.01, 2), name="Thread-C") threads.append(t3) # クライアント D (無料ティア) - 同じ制限に達する別のクライアント t4 = threading.Thread(target=simulate_requests, args=(limiter, "client_D", 10, 0.15, 3), name="Thread-D") threads.append(t4) # シミュレーションスレッドを開始します for t in threads: t.start() # 初期リクエストとエンキューに時間を置きます time.sleep(2) print("\n--- 遅延リクエストの処理を試行中 ---") # レート制限された可能性のあるクライアントの遅延リクエストを処理するスレッドを開始します # クライアント A と D (無料ティア) について複数回試行します。 processor_threads = [] tp1 = threading.Thread(target=process_deferred_requests, args=(limiter, "client_A", 5), name="Processor-A") processor_threads.append(tp1) tp2 = threading.Thread(target=process_deferred_requests, args=(limiter, "client_D", 5), name="Processor-D") processor_threads.append(tp2) # クライアント B についても試行します。 tp3 = threading.Thread(target=process_deferred_requests, args=(limiter, "client_B", 3), name="Processor-B") processor_threads.append(tp3) for tp in processor_threads: tp.start() # すべてのシミュレーションスレッドが完了するのを待ちます for t in threads: t.join() # プロセッサースレッドが試行を完了するのを待ちます for tp in processor_threads: tp.join() print("\n--- デモンストレーション完了 ---") print("最終状態チェック (カウントはタイミングにより概算です):") for tier_name, counter in limiter.tier_counters.items(): print(f"ティア '{tier_name}': 現在のウィンドウ内のリクエスト数: {counter.get_current_count()} (最大: {counter.max_requests})") for tier_name, queue in limiter.tier_queues.items(): print(f"ティア '{tier_name}': キュー内の遅延リクエスト数: {len(queue)}") # --- 設計上の選択に関する説明 --- # 1. スライディングウィンドウの実装: # - ウィンドウ内のリクエストのタイムスタンプを格納するために `collections.deque` を使用します。 # - `deque` は、新しいタイムスタンプの追加と古いタイムスタンプの左からの削除に O(1) の計算量を提供します。 # - `record_request` が呼び出されると、まず `current_time - window_seconds` より古いすべてのタイムスタンプを削除します。 # - これにより、`len(self.request_timestamps)` が現在のスライディングウィンドウ内のリクエスト数を正確に反映することが保証されます。 # - トレードオフ: このアプローチは正確ですが、多数のリクエストが短時間で到着し、その後停止した場合、`record_request` で複数のポップが発生する可能性があります。ただし、時間の経過とともに deque のサイズは管理されます。 # - スレッドセーフティ: 各 `SlidingWindowCounter` インスタンスは、`request_timestamps` deque を保護するための独自の `threading.Lock` を持ち、同じティアに対する `record_request` または `get_current_count` の同時呼び出しがシリアル化されることを保証します。 # 2. 複数のティア: # - `RateLimiter` クラスは、ティア名をレート制限パラメータ (`max_requests`, `window_seconds`) にマッピングする辞書 `tier_config` を保持します。 # - また、各ティアに対して個別の `SlidingWindowCounter` インスタンスと優先度付きキュー (`tier_queues`) を保持します。 # - クライアントティアの割り当ては `client_tiers` に格納されます。 # - これにより、異なるクライアントタイプ間で柔軟な設定とレート制限の分離が可能になります。 # 3. 遅延リクエスト用の優先度付きキュー: # - 各ティアについて、標準の Python リストが `heapq` モジュールによって管理される最小ヒープとして使用されます。 # - `Request` dataclass は、`priority` をプライマリソートキーとして設計されています (値が小さいほど優先度が高い)。 # - 一意の `request_id` が含まれており、安定ソート (同じ優先度の場合は FIFO) を保証し、タイムスタンプが同一の場合の比較エラーを防ぎます。 # - `enqueue_request` は `heapq.heappush` を使用してリクエストを追加し、ヒーププロパティを維持します。 # - `dequeue_and_process` は `heapq.heappop` を使用して最優先度リクエストを取得します。 # - 重要なのは、`dequeue_and_process` はヒープからポップしようとする前に、まず利用可能な容量を確認することです。容量が利用可能な場合、グローバルロックを取得し、容量を再確認し (競合状態を処理するため)、リクエストをポップし、`record_request` を再度呼び出します。`record_request` が成功すると、リクエストが処理されます。失敗した場合 (ちょうど容量が埋められた競合状態のため)、リクエストはヒープに戻されます。 # - トレードオフ: Python のリスト + heapq を使用することは、適度な数の遅延リクエストに対して効率的です。非常に高いボリュームの場合、より専門的な並列優先度キュー実装が検討されるかもしれませんが、それは複雑さを増します。 # 4. スレッドセーフティ: # - `SlidingWindowCounter`: 内部の `threading.Lock` を使用して `request_timestamps` deque を保護します。 # - `RateLimiter`: # - `global_lock` (`threading.Lock`): 共有クライアント状態 (`client_tiers`, `client_last_activity`) とティア優先度キュー (`tier_queues`) を保護します。このロックは、`register_client`、`enqueue_request`、および `dequeue_and_process` 内でキューにアクセス/変更する操作のために取得されます。 # - `allow_request`: スレッドセーフな `counter.record_request()` を呼び出します。また、グローバルロックの下でクライアントアクティビティを更新します。 # - `dequeue_and_process`: ロックを慎重に管理します。まずグローバルロックなしで容量を確認します。容量が利用可能である可能性がある場合、グローバルロックを取得し、容量を再確認し、ヒープ操作を実行し、次にスレッドセーフな `counter.record_request()` を再度呼び出します。 # - トレードオフ: 複数のロック (カウンターごとに 1 つ + グローバルロック 1 つ) の使用は、すべての操作に対して単一のグローバルロックを使用する場合と比較して、競合を減らすことを目的としています。ただし、クライアント状態とティア状態にまたがる操作 (デキューなど) では、慎重なロック順序付けまたはグローバルロックの取得が必要です。 # 5. クリーンアップ: # - バックグラウンドスレッド (`_cleanup_task`) が `client_last_activity` を定期的にスキャンします。 # - 最終アクティビティタイムスタンプが `cleanup_threshold_seconds` より古いクライアントを特定します。 # - これらのクライアントは、`global_lock` の下で `client_tiers` と `client_last_activity` から削除されます。 # - これにより、`RateLimiter` が非アクティブなクライアントの状態を無限に蓄積するのを防ぎます。 # - トレードオフ: クリーンアップは即時ではありません。定期的に実行されます。`cleanup_threshold_seconds` は、メモリ使用量とクリーンアップの応答性の間のトレードオフを決定します。 # 6. エッジケースの処理: # - クライアントの重複登録: `register_client` は False を返し、警告をログに記録します。 # - 未登録クライアント: `allow_request` および `enqueue_request` は `client_tiers` をチェックし、エラーを返します/ログに記録します。 # - 空の優先度付きキュー: `dequeue_and_process` は、`self.tier_queues[tier]` が空の場合のケースを処理します。 # - 同時変更: 適切な `threading.Lock` の使用によって対処されます。 # - 時計の精度: `time.time()` が使用されますが、これは一般的に十分です。非常に高い精度が要求される場合は、プラットフォーム固有の単調クロックが検討されるかもしれませんが、`time.time()` はこの種のアプリケーションでは標準です。 # - デキュー時の競合状態: 容量の再確認と、`dequeue_and_process` のクリティカルセクション内でのリクエストの再記録によって処理されます。 ```

import threading import time import heapq from collections import deque from dataclasses import dataclass from typing import Any, Dict, Optional, Tuple import itertools import random # Request dataclass holds minimal required fields @dataclass(order=True) class Request: # ordering will be driven primarily by the priority in heaps priority: int client_id: str timestamp: float payload: Any = None class RateLimiter: """ Thread-safe rate limiter with sliding window per client and per-tier priority queues. Design choices (high-level): - Sliding window implemented per client using a deque of request timestamps. This gives exact counts for the sliding window (precision) at the cost of storing one timestamp per request. - Per-tier priority queues implemented with heapq. Each queued Request is pushed as (priority, counter, request) to preserve FIFO among same-priority items. Lower priority number = higher priority. - Concurrency: use a combination of locks: * self.clients_lock protects the client registry structure. * each client has its own lock to protect its timestamp deque and last_seen metadata (reduces contention). * each tier has its own lock to protect its priority queue. Lock ordering is important to avoid deadlocks: when both client lock and tier lock need to be acquired, the implementation always acquires the client lock first, then the tier lock. All methods follow this order. - For dequeue_for_client, we must find the highest-priority waiting request for that client in the per-tier heap. Because heapq does not support efficient arbitrary-item removal, we pop items until we find one for the client, buffering others and pushing them back afterwards. This can be O(n) in the queue size in the worst case — a trade-off favoring simpler data structures and correctness over per-operation worst-case performance. - Time functions use time.monotonic() for durations to avoid system clock jumps affecting rate limiting. Trade-offs: - Precision vs space: deque per client is precise for sliding window but stores each request timestamp. - Performance vs complexity: scanning heaps for dequeue_for_client is simple but can be costly if queues are large. """ def __init__(self, tiers: Dict[str, Dict[str, float]]): """ tiers: mapping tier_name -> { 'max_requests': int, 'window_seconds': float } """ self.tiers = {} self.tiers_lock = threading.Lock() self._counter = itertools.count() # Initialize tier structures for name, conf in tiers.items(): if 'max_requests' not in conf or 'window_seconds' not in conf: raise ValueError(f"Tier {name} missing configuration") self.tiers[name] = { 'max_requests': int(conf['max_requests']), 'window_seconds': float(conf['window_seconds']), 'queue': [], # heap: elements (priority, counter, Request) 'lock': threading.Lock(), } # Clients registry: client_id -> { 'tier': str, 'timestamps': deque, 'lock': Lock, 'last_seen': float } self.clients: Dict[str, Dict[str, Any]] = {} self.clients_lock = threading.Lock() def register_client(self, client_id: str, tier_name: str): """Register a client into a tier. Raises ValueError for duplicate registration or unknown tier. Thread-safe. """ with self.clients_lock: if client_id in self.clients: raise ValueError(f"Client {client_id} already registered") if tier_name not in self.tiers: raise ValueError(f"Unknown tier {tier_name}") self.clients[client_id] = { 'tier': tier_name, 'timestamps': deque(), 'lock': threading.Lock(), 'last_seen': time.monotonic(), } def _remove_old_timestamps(self, client_entry: Dict[str, Any], window_seconds: float, now: float): # remove timestamps older than sliding window timestamps: deque = client_entry['timestamps'] cutoff = now - window_seconds while timestamps and timestamps[0] <= cutoff: timestamps.popleft() def allow_request(self, client_id: str, priority: int = 100, payload: Any = None, timestamp: Optional[float] = None) -> Tuple[bool, Optional[Request]]: """ Attempt to allow a request for client_id. If under the sliding window limit, the request is recorded and allowed (returns (True, None)). If the client is at capacity, the request is enqueued into the client's tier priority queue and (False, request) is returned. Raises KeyError if client is not registered. Thread-safe. """ if timestamp is None: now = time.monotonic() else: # If caller supplied an external timestamp, convert it to monotonic-relative by using monotonic() now # and assuming external timestamp is in same epoch; for demonstration we accept timestamp as-is. now = timestamp # Get client entry with self.clients_lock: client_entry = self.clients.get(client_id) if client_entry is None: raise KeyError(f"Client {client_id} is not registered") tier_name = client_entry['tier'] tier = self.tiers[tier_name] # Acquire client lock first (lock ordering rule), then possibly tier lock. with client_entry['lock']: self._remove_old_timestamps(client_entry, tier['window_seconds'], now) if len(client_entry['timestamps']) < tier['max_requests']: # allow request client_entry['timestamps'].append(now) client_entry['last_seen'] = now return True, None else: # at capacity -> enqueue into tier priority queue req = Request(priority=priority, client_id=client_id, timestamp=now, payload=payload) # Acquire tier lock to push into heap with tier['lock']: count = next(self._counter) heapq.heappush(tier['queue'], (req.priority, count, req)) client_entry['last_seen'] = now return False, req def enqueue_request(self, req: Request): """Directly enqueue a Request into its client's tier queue. Raises if client not registered. Thread-safe. """ with self.clients_lock: client_entry = self.clients.get(req.client_id) if client_entry is None: raise KeyError(f"Client {req.client_id} is not registered") tier_name = client_entry['tier'] tier = self.tiers[tier_name] with tier['lock']: count = next(self._counter) heapq.heappush(tier['queue'], (req.priority, count, req)) def dequeue_for_client(self, client_id: str, timestamp: Optional[float] = None) -> Optional[Request]: """ If capacity is available for client_id, dequeue and return the highest-priority queued Request for the client's tier. The returned Request is considered processed and will be counted toward the window. If no queued request is present for the client or capacity isn't available, return None. Raises KeyError if client not registered. Thread-safe. """ if timestamp is None: now = time.monotonic() else: now = timestamp with self.clients_lock: client_entry = self.clients.get(client_id) if client_entry is None: raise KeyError(f"Client {client_id} is not registered") tier_name = client_entry['tier'] tier = self.tiers[tier_name] # Acquire client lock first, then tier lock (consistent ordering) with client_entry['lock']: self._remove_old_timestamps(client_entry, tier['window_seconds'], now) if len(client_entry['timestamps']) >= tier['max_requests']: # no capacity return None # There is capacity; try to find highest-priority queued request for this client with tier['lock']: queue = tier['queue'] if not queue: return None temp = [] found_req = None try: while queue: priority, count, req = heapq.heappop(queue) if req.client_id == client_id: found_req = req break else: temp.append((priority, count, req)) finally: # push back any popped non-matching requests for item in temp: heapq.heappush(queue, item) if found_req is None: return None # mark the request as processed: append timestamp to client's timestamps client_entry['timestamps'].append(now) client_entry['last_seen'] = now return found_req def cleanup(self, cleanup_threshold_seconds: float): """ Remove tracking data for clients who have not made requests in the last cleanup_threshold_seconds. Does not unregister the client (keeps registration), but removes timestamp history to free memory. If a client has queued deferred requests, we keep their tracking data. Thread-safe. """ now = time.monotonic() cutoff = now - cleanup_threshold_seconds removed = [] # iterate clients safely with self.clients_lock: for client_id, client_entry in list(self.clients.items()): # acquire client lock to inspect safely with client_entry['lock']: # if the client has no activity recently and no queued requests, clear timestamps last_seen = client_entry.get('last_seen', 0) tier_name = client_entry['tier'] tier = self.tiers[tier_name] has_queued = False # check if any queued requests exist for this client in tier queue with tier['lock']: for _, _, req in tier['queue']: if req.client_id == client_id: has_queued = True break if last_seen < cutoff and not has_queued: # clear timestamps to free memory client_entry['timestamps'].clear() removed.append(client_id) return removed # Demonstration script if __name__ == "__main__": # Create a rate limiter with two tiers: free and premium tiers_conf = { 'free': {'max_requests': 3, 'window_seconds': 5.0}, 'premium': {'max_requests': 10, 'window_seconds': 5.0}, } rl = RateLimiter(tiers_conf) # Register clients clients = [ ('alice', 'free'), ('bob', 'free'), ('carol', 'premium'), ] for cid, tier in clients: rl.register_client(cid, tier) print(f"Registered client {cid} as tier {tier}") stop_event = threading.Event() # Worker threads simulate bursts of requests def client_worker(client_id: str, burst: int, sleep_between: float): for i in range(burst): allowed, req = rl.allow_request(client_id, priority=random.randint(1, 10), payload={'i': i}) ts = time.monotonic() if allowed: print(f"[{ts:.3f}] {client_id}: request {i} ALLOWED") else: print(f"[{ts:.3f}] {client_id}: request {i} ENQUEUED (priority={req.priority})") time.sleep(sleep_between) # Background processor that periodically tries to dequeue queued requests for all clients def deferred_processor(): while not stop_event.is_set(): # iterate all clients and try to dequeue one request if capacity exists with rl.clients_lock: client_ids = list(rl.clients.keys()) progressed = False for cid in client_ids: try: req = rl.dequeue_for_client(cid) except KeyError: continue if req is not None: ts = time.monotonic() print(f"[{ts:.3f}] Processed deferred request for {cid} (priority={req.priority}, payload={req.payload})") progressed = True if not progressed: # nothing processed; sleep a bit time.sleep(0.2) # Start threads: bursty clients threads = [] t1 = threading.Thread(target=client_worker, args=('alice', 8, 0.4)) t2 = threading.Thread(target=client_worker, args=('bob', 6, 0.3)) t3 = threading.Thread(target=client_worker, args=('carol', 12, 0.1)) threads.extend([t1, t2, t3]) processor_thread = threading.Thread(target=deferred_processor) processor_thread.daemon = True processor_thread.start() for t in threads: t.start() # Let threads finish for t in threads: t.join() # Allow background processor to handle remaining queued items for some time time.sleep(3.0) stop_event.set() processor_thread.join(timeout=1.0) # Show cleanup removed = rl.cleanup(cleanup_threshold_seconds=1.0) print(f"Cleanup removed timestamp history for clients: {removed}") # Show queued items left in each tier for tier_name, tier in rl.tiers.items(): with tier['lock']: qsize = len(tier['queue']) print(f"Tier '{tier_name}' has {qsize} queued requests left") print("Demo complete")