Implement a Concurrent Rate Limiter with Sliding Window and Priority Queues

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)