Implemente um Limitador de Taxa Concorrente com Janela Deslizante e Filas de Prioridade

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)