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

import time import threading import random from collections import deque from dataclasses import dataclass, field from typing import List, Tuple, Dict, Optional @dataclass class ClientState: """Holds all state for a single client.""" lock: threading.Lock = field(default_factory=threading.Lock, repr=False) requests: deque[Tuple[float, int]] = field(default_factory=deque) burst: int = 0 stats: Dict[str, int] = field(default_factory=lambda: {"admitted": 0, "rejected": 0, "revoked": 0}) revoked_log: List[Tuple[float, int]] = field(default_factory=list) class RateLimiter: """ A thread-safe rate limiter implementing a sliding window algorithm with priority queues and burst allowances. """ def __init__(self, max_requests: int, window_seconds: float, default_burst: int = 0): if max_requests <= 0 or window_seconds <= 0: raise ValueError("max_requests and window_seconds must be positive.") self.max_requests = max_requests self.window_seconds = window_seconds self.default_burst = default_burst self._clients: Dict[str, ClientState] = {} self._clients_lock = threading.Lock() def _get_or_create_client(self, client_id: str) -> ClientState: """Atomically get or create a client's state object.""" if client_id not in self._clients: with self._clients_lock: # Double-check locking pattern to avoid race conditions on creation if client_id not in self._clients: self._clients[client_id] = ClientState(burst=self.default_burst) return self._clients[client_id] def set_client_burst(self, client_id: str, burst: int) -> None: """Override burst allowance for a specific client.""" client = self._get_or_create_client(client_id) with client.lock: client.burst = burst def allow(self, client_id: str, priority: int = 3, timestamp: Optional[float] = None) -> bool: """ Check if a request is allowed. If timestamp is None, use current time. Returns True if the request is admitted, False if rejected. """ if not 1 <= priority <= 5: raise ValueError("Priority must be an integer between 1 and 5.") now = timestamp if timestamp is not None else time.time() client = self._get_or_create_client(client_id) with client.lock: # 1. Slide the window: remove expired requests window_start = now - self.window_seconds while client.requests and client.requests[0][0] < window_start: client.requests.popleft() # 2. Determine the effective request limit (base + burst) effective_max = self.max_requests if client.burst > 0 and client.requests: first_req_time = client.requests[0][0] if (now - first_req_time) >= (self.window_seconds / 2.0): effective_max += client.burst current_requests_count = len(client.requests) # 3. Decision: Admit, Revoke, or Reject if current_requests_count < effective_max: # Case A: Space is available (either base or burst) client.requests.append((now, priority)) client.stats['admitted'] += 1 return True else: # Case B: Window is full, check for priority revocation # Find the request with the highest priority number (lowest actual priority) worst_request_idx = -1 worst_priority = -1 for i, req in enumerate(client.requests): if req[1] > worst_priority: worst_priority = req[1] worst_request_idx = i # If a suitable candidate for revocation exists... if worst_request_idx != -1 and worst_priority > priority: # Revoke the lowest-priority request revoked_req = client.requests[worst_request_idx] del client.requests[worst_request_idx] client.revoked_log.append(revoked_req) client.stats['revoked'] += 1 # Admit the new, higher-priority request client.requests.append((now, priority)) client.stats['admitted'] += 1 return True else: # Case C: No space and no possible revocation -> Reject client.stats['rejected'] += 1 return False def get_stats(self, client_id: str) -> dict: """ Return a dict with keys: 'admitted', 'rejected', 'revoked', 'utilization' """ if client_id not in self._clients: return {'admitted': 0, 'rejected': 0, 'revoked': 0, 'utilization': 0.0} client = self._clients[client_id] with client.lock: # Slide window for accurate utilization calculation at the current time now = time.time() window_start = now - self.window_seconds while client.requests and client.requests[0][0] < window_start: client.requests.popleft() current_count = len(client.requests) utilization = min(1.0, current_count / self.max_requests) if self.max_requests > 0 else 0.0 stats_copy = client.stats.copy() stats_copy['utilization'] = utilization return stats_copy def get_revoked_log(self, client_id: str) -> list: """ Return a list of (timestamp, priority) tuples for revoked requests for the given client, in chronological order. """ if client_id not in self._clients: return [] client = self._clients[client_id] with client.lock: # The log is already in chronological order of revocation return list(client.revoked_log) # --- Demonstration Script --- def run_demonstration(): print("--- Rate Limiter Demonstration ---") limiter = RateLimiter(max_requests=5, window_seconds=10.0, default_burst=2) # --- Client 1: Standard, Rejection, and Priority Revocation --- client1 = "client-A" print(f"\n--- Testing Client: {client1} ---") # 1. Fill the window with medium-priority requests print("1. Filling the window (5 requests)...") for i in range(5): limiter.allow(client1, priority=4, timestamp=float(i)) # 2. Reject a request print("2. Sending a 6th request (should be rejected)...") is_allowed = limiter.allow(client1, priority=4, timestamp=5.0) print(f" Request at T=5.0 allowed: {is_allowed} (Expected: False)") # 3. Revoke a low-priority request with a high-priority one print("3. Sending a high-priority request (should revoke one)...") is_allowed = limiter.allow(client1, priority=1, timestamp=6.0) print(f" High-priority request at T=6.0 allowed: {is_allowed} (Expected: True)") # 4. Let the window slide print("4. Waiting for window to slide...") # At T=11.0, request at T=0.0 has expired. Window contains T=1,2,3,4,6 (5 requests) is_allowed = limiter.allow(client1, priority=5, timestamp=11.0) print(f" Request at T=11.0 allowed: {is_allowed} (Expected: False)") # At T=12.0, requests at T=0,1 have expired. Window contains T=2,3,4,6 (4 requests) is_allowed = limiter.allow(client1, priority=5, timestamp=12.0) print(f" Request at T=12.0 allowed: {is_allowed} (Expected: True)") # --- Client 2: Burst Allowance --- client2 = "client-B" print(f"\n--- Testing Client: {client2} (Burst) ---") limiter.set_client_burst(client2, 3) # Override burst to 3 for this client # 1. Send initial request print("1. Sending initial request at T=20.0...") limiter.allow(client2, priority=3, timestamp=20.0) # 2. Try to burst before half-window (should fail) print("2. Attempting to burst before half-window (T=24.0)...") # Window is 10s, half is 5s. T=24.0 is 4s after T=20.0. # Should allow up to max_requests (5) for i in range(4): # Requests 2, 3, 4, 5 limiter.allow(client2, priority=3, timestamp=24.0) is_allowed = limiter.allow(client2, priority=3, timestamp=24.1) print(f" 6th request at T=24.1 (before burst active) allowed: {is_allowed} (Expected: False)") # 3. Activate burst after half-window print("3. Activating burst after half-window (T=25.0)...") # T=25.0 is 5s after T=20.0. Burst is now active. # Limit is now 5 (base) + 3 (burst) = 8. # Currently 5 requests in window (from T=20, T=24). Should allow 3 more. allowed_count = 0 for _ in range(4): # Try to send 4 more if limiter.allow(client2, priority=3, timestamp=25.0): allowed_count += 1 print(f" Allowed {allowed_count} burst requests at T=25.0. (Expected: 3)") # --- Final Stats --- print("\n--- Final Stats ---") print(f"Client A Stats: {limiter.get_stats(client1)}") print(f"Client A Revoked Log: {limiter.get_revoked_log(client1)}") print(f"Client B Stats: {limiter.get_stats(client2)}") print(f"Client B Revoked Log: {limiter.get_revoked_log(client2)}") print(f"Unknown Client Stats: {limiter.get_stats('client-C')}") # --- Multithreaded Test --- def run_multithreaded_test(): print("\n--- Multithreaded Test ---") limiter = RateLimiter(max_requests=20, window_seconds=5.0) client_id = "mt-client" num_threads = 5 requests_per_thread = 10 total_requests = num_threads * requests_per_thread barrier = threading.Barrier(num_threads) def worker(): barrier.wait() for i in range(requests_per_thread): priority = random.randint(1, 5) limiter.allow(client_id, priority=priority) time.sleep(0.01) threads = [threading.Thread(target=worker) for _ in range(num_threads)] for t in threads: t.start() for t in threads: t.join() print(f"\nFinished multithreaded test with {num_threads} threads and {requests_per_thread} requests each.") stats = limiter.get_stats(client_id) print(f"Final Stats for {client_id}: {stats}") total_processed = stats['admitted'] + stats['rejected'] print(f"Total requests attempted: {total_requests}") print(f"Total requests processed (admitted + rejected): {total_processed}") if total_processed == total_requests: print("Test PASSED: All requests were accounted for, indicating thread safety.") else: print(f"Test FAILED: Mismatch in request counts. Expected {total_requests}, got {total_processed}.") if __name__ == '__main__': run_demonstration() run_multithreaded_test()

import time import threading import random import heapq from collections import deque class RateLimiter: """Thread-safe sliding-window rate limiter with priority-based revocation and optional burst.""" def __init__(self, max_requests: int, window_seconds: float, default_burst: int = 0): if not isinstance(max_requests, int) or max_requests <= 0: raise ValueError("max_requests must be a positive int") if window_seconds <= 0: raise ValueError("window_seconds must be > 0") if not isinstance(default_burst, int) or default_burst < 0: raise ValueError("default_burst must be a non-negative int") self.max_requests = max_requests self.window_seconds = float(window_seconds) self.default_burst = default_burst self._lock = threading.RLock() self._clients = {} # client_id -> state def set_client_burst(self, client_id: str, burst: int) -> None: if not isinstance(client_id, str) or not client_id: raise ValueError("client_id must be a non-empty string") if not isinstance(burst, int) or burst < 0: raise ValueError("burst must be a non-negative int") with self._lock: st = self._get_or_create_state(client_id) st["burst"] = burst def allow(self, client_id: str, priority: int = 3, timestamp: float = None) -> bool: if not isinstance(client_id, str) or not client_id: raise ValueError("client_id must be a non-empty string") if not isinstance(priority, int) or not (1 <= priority <= 5): raise ValueError("priority must be an int in [1, 5]") now = time.time() if timestamp is None else float(timestamp) with self._lock: st = self._get_or_create_state(client_id) self._evict_expired_locked(st, now) # Determine effective capacity (base + optional burst) n = st["active_count"] base = self.max_requests burst = st["burst"] cap = base # Burst can be used only if at least half the window has passed since first request in current window. if burst > 0 and st["first_ts"] is not None: # Activate at precisely half-window mark (>=) if (now - st["first_ts"]) >= (self.window_seconds / 2.0): cap = base + burst # Fast path: within capacity if n < cap: self._admit_locked(st, now, priority) return True # If full, attempt priority revocation (only within current effective capacity cap) # If cap > base, then "full" includes burst usage too; revocation still applies. # We can only revoke a strictly lower priority (numerically higher). # Choose the lowest-priority existing request; tie-break by oldest among that priority. # If no lower-priority exists, reject. # Make sure heap top points to a live active event. cand = self._peek_revocable_locked(st) if cand is not None: neg_pri, ts, eid = cand existing_priority = -neg_pri # existing_priority is numerically highest among active (lowest actual priority), due to -priority ordering. if existing_priority > priority: # Revoke candidate self._pop_specific_candidate_locked(st, eid) # Admit new self._admit_locked(st, now, priority) st["revoked"] += 1 st["revoked_log"].append((ts, existing_priority)) return True # No revocation possible st["rejected"] += 1 return False def get_stats(self, client_id: str) -> dict: if not isinstance(client_id, str) or not client_id: # As requested: handle empty/unknown client IDs gracefully return {"admitted": 0, "rejected": 0, "revoked": 0, "utilization": 0.0} with self._lock: st = self._clients.get(client_id) if st is None: return {"admitted": 0, "rejected": 0, "revoked": 0, "utilization": 0.0} now = time.time() self._evict_expired_locked(st, now) base = self.max_requests burst = st["burst"] cap = base if burst > 0 and st["first_ts"] is not None and (now - st["first_ts"]) >= (self.window_seconds / 2.0): cap = base + burst util = (st["active_count"] / cap) if cap > 0 else 0.0 # clamp if util < 0.0: util = 0.0 if util > 1.0: util = 1.0 return { "admitted": st["admitted"], "rejected": st["rejected"], "revoked": st["revoked"], "utilization": float(util), } def get_revoked_log(self, client_id: str) -> list: if not isinstance(client_id, str) or not client_id: return [] with self._lock: st = self._clients.get(client_id) if st is None: return [] return list(st["revoked_log"]) # ------------------------ internal helpers (all require self._lock held) ------------------------ def _get_or_create_state(self, client_id: str) -> dict: st = self._clients.get(client_id) if st is None: st = { "burst": self.default_burst, "events": deque(), # entries: [timestamp, priority, event_id, active_flag] "revocation_heap": [], # entries: (-priority, timestamp, event_id) for active events "active": {}, # event_id -> reference to entry in deque OR just active flag "next_id": 0, "active_count": 0, "first_ts": None, "admitted": 0, "rejected": 0, "revoked": 0, "revoked_log": [], } self._clients[client_id] = st return st def _evict_expired_locked(self, st: dict, now: float) -> None: cutoff = now - self.window_seconds # Sliding window: keep entries with timestamp > cutoff. # Requests exactly at boundary (timestamp == cutoff) are expired. ev = st["events"] while ev and ev[0][0] <= cutoff: ts, pri, eid, active = ev[0] ev.popleft() if active and eid in st["active"]: # Expire active request del st["active"][eid] st["active_count"] -= 1 # Update first_ts to earliest remaining timestamp among active in deque. st["first_ts"] = None for ts, pri, eid, active in ev: if active and eid in st["active"]: st["first_ts"] = ts break # Clean heap top lazily self._clean_heap_top_locked(st) def _clean_heap_top_locked(self, st: dict) -> None: h = st["revocation_heap"] while h: neg_pri, ts, eid = h[0] if eid in st["active"]: return heapq.heappop(h) def _peek_revocable_locked(self, st: dict): self._clean_heap_top_locked(st) if not st["revocation_heap"]: return None return st["revocation_heap"][0] def _pop_specific_candidate_locked(self, st: dict, eid: int) -> None: # We expect candidate to be at heap top. Pop until we pop that eid. h = st["revocation_heap"] while h: neg_pri, ts, heid = heapq.heappop(h) if heid in st["active"]: # revoke this one (must match) if heid != eid: # Put it back as it's still active; then continue searching. heapq.heappush(h, (neg_pri, ts, heid)) # fall back: revoke the one we just saw instead? No; keep deterministic behavior. # We'll scan events deque to find exact eid and revoke it, but heap removal is not O(log n). # Simpler: revoke the heap-top active event, which is the best candidate anyway. # So if mismatch happens (shouldn't), revoke heid. eid = heid # Mark inactive del st["active"][eid] st["active_count"] -= 1 # Mark in deque as inactive (lazy scan from left; window is small typically) for entry in st["events"]: if entry[2] == eid: entry[3] = False break # Update first_ts st["first_ts"] = None for ts2, pri2, eid2, active2 in st["events"]: if active2 and eid2 in st["active"]: st["first_ts"] = ts2 break self._clean_heap_top_locked(st) return # If heap is empty, nothing to do. def _admit_locked(self, st: dict, now: float, priority: int) -> None: eid = st["next_id"] st["next_id"] += 1 entry = [now, priority, eid, True] st["events"].append(entry) st["active"][eid] = True st["active_count"] += 1 heapq.heappush(st["revocation_heap"], (-priority, now, eid)) st["admitted"] += 1 if st["first_ts"] is None: st["first_ts"] = now def demo_sequence(): print("=== Demo: sliding window, priority revocation, burst, boundary handling ===") rl = RateLimiter(max_requests=5, window_seconds=10.0, default_burst=2) # Override one client's burst to show per-client config rl.set_client_burst("clientA", 2) rl.set_client_burst("clientB", 0) t0 = 1_000_000.0 # Fill base window for clientA with low priority (5) for i in range(5): ok = rl.allow("clientA", priority=5, timestamp=t0 + i * 0.1) print("clientA admit low", i, ok) # Now window is full. Higher-priority request should revoke one of the priority-5 requests. ok = rl.allow("clientA", priority=2, timestamp=t0 + 0.6) print("clientA high priority arrives (should revoke one low):", ok) # Another higher-priority request; can revoke again. ok = rl.allow("clientA", priority=1, timestamp=t0 + 0.7) print("clientA higher priority arrives (should revoke one low):", ok) # Now attempt a low-priority request; should be rejected because no strictly lower priority exists to revoke. ok = rl.allow("clientA", priority=5, timestamp=t0 + 0.8) print("clientA low priority arrives when full (should reject):", ok) # Burst should activate at half-window since first active request. # First request was at t0. Half-window mark is t0+5.0. At exactly +5.0, burst becomes available. ok = rl.allow("clientA", priority=3, timestamp=t0 + 5.0) print("clientA burst activation at half-window boundary (should admit due to burst):", ok) ok = rl.allow("clientA", priority=4, timestamp=t0 + 5.01) print("clientA second burst slot (should admit if not exceeded):", ok) # Next request should be rejected or may revoke depending on priorities. ok = rl.allow("clientA", priority=4, timestamp=t0 + 5.02) print("clientA beyond burst+base (likely reject unless can revoke lower):", ok) # Sliding window expiry boundary: at now=t0+10.0, requests at t0 are expired (<= cutoff). # Send at exactly t0+10.0: entries at t0 are expired. ok = rl.allow("clientA", priority=3, timestamp=t0 + 10.0) print("clientA at exact window boundary (t0+10.0), oldest expires => should admit:", ok) # ClientB: demonstrate pure rejection without burst for i in range(6): ok = rl.allow("clientB", priority=3, timestamp=t0 + i * 0.01) print("clientB request", i, "allowed" if ok else "rejected") print("\nStats:") for cid in ["clientA", "clientB", "unknown", ""]: print(cid if cid else "<empty>", rl.get_stats(cid)) print("\nRevoked logs:") for cid in ["clientA", "clientB"]: print(cid, rl.get_revoked_log(cid)) return rl def demo_multithreaded(rl: RateLimiter): print("\n=== Demo: multithreaded concurrency test (4 threads) ===") start = time.time() stop_at = start + 1.0 results_lock = threading.Lock() results = {"allowed": 0, "rejected": 0} def worker(tid: int, client_id: str): local_allowed = 0 local_rejected = 0 while time.time() < stop_at: pri = random.randint(1, 5) ok = rl.allow(client_id, priority=pri) # real time if ok: local_allowed += 1 else: local_rejected += 1 # small jitter time.sleep(random.random() * 0.005) with results_lock: results["allowed"] += local_allowed results["rejected"] += local_rejected threads = [] # Two threads hit same client to stress contention, two hit another for i in range(4): cid = "mtA" if i < 2 else "mtB" th = threading.Thread(target=worker, args=(i, cid), daemon=True) threads.append(th) th.start() for th in threads: th.join() print("Concurrent run totals:", results) print("mtA stats:", rl.get_stats("mtA")) print("mtB stats:", rl.get_stats("mtB")) if __name__ == "__main__": rl = demo_sequence() demo_multithreaded(rl)