Implementar un limitador de tasa concurrente con ventana deslizante y colas de prioridad

Respuesta: ```python import time import threading import collections import heapq import uuid from dataclasses import dataclass, field from typing import Dict, List, Tuple, Any, Optional # --- Configuración --- # Define niveles con sus respectivos límites de tasa y tamaños de ventana. # window_seconds más bajo significa una limitación de tasa más estricta. TIER_CONFIG = { "free": {"max_requests": 5, "window_seconds": 60}, "standard": {"max_requests": 50, "window_seconds": 60}, "premium": {"max_requests": 200, "window_seconds": 60}, } # Cuánto tiempo mantener los datos del cliente antes de limpiarlos si no hay actividad. CLEANUP_THRESHOLD_SECONDS = 300 # --- Estructuras de Datos --- @dataclass(order=True) class Request: """Representa una solicitud con prioridad, para la cola de prioridad.""" # Usamos la prioridad como el primer elemento para que heapq la ordene. # Un número más bajo significa mayor prioridad. priority: int # Usamos un ID único para desempatar prioridades y asegurar una ordenación estable. request_id: str = field(default_factory=lambda: str(uuid.uuid4())) client_id: str timestamp: float payload: Any def __lt__(self, other): # Comparación personalizada para heapq: prioriza por prioridad, luego por request_id para estabilidad. if self.priority != other.priority: return self.priority < other.priority return self.request_id < other.request_id class SlidingWindowCounter: """Implementa un contador de ventana deslizante para la limitación de tasa.""" def __init__(self, max_requests: int, window_seconds: float): self.max_requests = max_requests self.window_seconds = window_seconds # Usamos una deque para almacenar las marcas de tiempo de las solicitudes dentro de la ventana. # Las deques ofrecen adiciones y eliminaciones eficientes desde ambos extremos. self.request_timestamps: collections.deque[float] = collections.deque() # Bloqueo para acceso seguro a hilos a request_timestamps. self.lock = threading.Lock() def record_request(self) -> bool: """Registra una solicitud y devuelve True si está dentro del límite, False en caso contrario.""" with self.lock: current_time = time.time() # Elimina las marcas de tiempo más antiguas que la ventana. while self.request_timestamps and self.request_timestamps[0] <= current_time - self.window_seconds: self.request_timestamps.popleft() # Comprueba si agregar una nueva solicitud excede el límite. 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: """Devuelve el número actual de solicitudes dentro de la ventana.""" with self.lock: current_time = time.time() # Limpia las marcas de tiempo antiguas antes de devolver el recuento. 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: """Un limitador de tasa seguro para hilos con múltiples niveles, ventana deslizante y colas de prioridad.""" 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 # Almacena los contadores de ventana deslizante para cada nivel. self.tier_counters: Dict[str, SlidingWindowCounter] = {} # Almacena la cola de prioridad para solicitudes diferidas para cada nivel. self.tier_queues: Dict[str, List[Request]] = collections.defaultdict(list) # Almacena la marca de tiempo de la última actividad para cada cliente para la limpieza. self.client_last_activity: Dict[str, float] = {} # Almacena el nivel asignado a cada cliente. self.client_tiers: Dict[str, str] = {} # Bloqueo principal para proteger las estructuras de datos compartidas como client_tiers, client_last_activity y tier_queues. # SlidingWindowCounter tiene su propio bloqueo interno. self.global_lock = threading.Lock() # Inicializa los contadores de nivel. for tier_name, config in tier_config.items(): self.tier_counters[tier_name] = SlidingWindowCounter( max_requests=config["max_requests"], window_seconds=config["window_seconds"] ) # Inicia un hilo en segundo plano para la limpieza. 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: """Registra un cliente con un nivel específico. Devuelve True si tiene éxito, False si el cliente ya existe o el nivel no es válido.""" if tier not in self.tier_config: print(f"[!] Error: El nivel '{tier}' no existe.") return False with self.global_lock: if client_id in self.client_tiers: print(f"[!] Advertencia: El cliente '{client_id}' ya está registrado.") return False self.client_tiers[client_id] = tier self.client_last_activity[client_id] = time.time() print(f"[*] Cliente '{client_id}' registrado con el nivel '{tier}'.") return True def _get_client_tier(self, client_id: str) -> Optional[str]: """Recupera de forma segura el nivel para un client_id dado.""" with self.global_lock: return self.client_tiers.get(client_id) def _update_client_activity(self, client_id: str): """Actualiza la marca de tiempo de la última actividad para un cliente.""" with self.global_lock: self.client_last_activity[client_id] = time.time() def allow_request(self, client_id: str) -> bool: """Comprueba si se permite una solicitud de client_id según el límite de tasa de su nivel.""" tier = self._get_client_tier(client_id) if not tier: print(f"[!] Error: Cliente '{client_id}' no registrado.") return False counter = self.tier_counters.get(tier) if not counter: # Esto no debería suceder si tier_config es válido y la inicialización es correcta. print(f"[!] Error interno: Contador no encontrado para el nivel '{tier}'.") return False if counter.record_request(): self._update_client_activity(client_id) # print(f"[*] Solicitud PERMITIDA para el cliente '{client_id}' (Nivel: {tier}). Recuento actual: {counter.get_current_count()}") return True else: # print(f"[*] Solicitud DENEGADA para el cliente '{client_id}' (Nivel: {tier}). Recuento actual: {counter.get_current_count()}") return False def enqueue_request(self, client_id: str, priority: int, payload: Any): """Pone en cola una solicitud limitada por tasa en la cola de prioridad apropiada.""" tier = self._get_client_tier(client_id) if not tier: print(f"[!] Error: Cliente '{client_id}' no registrado. No se puede poner en cola.") 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"[*] Solicitud ENCOLADA para el cliente '{client_id}' (Nivel: {tier}, Prioridad: {priority}). Tamaño de cola: {len(self.tier_queues[tier])}") self._update_client_activity(client_id) def dequeue_and_process(self, client_id: str) -> Optional[Tuple[Request, bool]]: """Intenta sacar de la cola y procesar la solicitud de mayor prioridad para el nivel de un cliente si hay capacidad disponible.""" tier = self._get_client_tier(client_id) if not tier: print(f"[!] Error: Cliente '{client_id}' no registrado. No se puede sacar de cola.") return None counter = self.tier_counters.get(tier) if not counter: print(f"[!] Error interno: Contador no encontrado para el nivel '{tier}'.") return None # Comprueba si hay capacidad disponible ANTES de intentar sacar de la cola. # Esto evita bloqueos y operaciones de heap innecesarias si no hay capacidad. if counter.get_current_count() < counter.max_requests: with self.global_lock: queue = self.tier_queues[tier] if queue: # Revisa la capacidad dentro del bloqueo para manejar condiciones de carrera. # Es posible que otro hilo acaba de llenar la capacidad. if counter.get_current_count() < counter.max_requests: highest_priority_request = heapq.heappop(queue) # Intenta registrar la solicitud nuevamente. Esto debería tener éxito si hay capacidad disponible. if counter.record_request(): self._update_client_activity(highest_priority_request.client_id) print(f"[*] Solicitud SACADA DE COLA y PROCESADA para el cliente '{highest_priority_request.client_id}' (Nivel: {tier}, Prioridad: {highest_priority_request.priority}). Tamaño de cola: {len(queue)}") return highest_priority_request, True else: # Esta es una condición de carrera rara: había capacidad disponible, pero record_request falló. # La vuelve a poner en cola y registra un error. heapq.heappush(queue, highest_priority_request) print(f"[!] Condición de carrera: Falló el registro de la solicitud sacada de cola para {highest_priority_request.client_id}. Vuelve a poner en cola.") return highest_priority_request, False else: # La capacidad fue tomada por otro hilo entre la comprobación externa y esta comprobación interna. # Vuelve a poner la solicitud en cola si fue sacada. heapq.heappush(queue, highest_priority_request) # print(f"[*] Intento de sacar de cola fallido para {client_id}: Capacidad llena por otro hilo.") return highest_priority_request, False else: # La cola está vacía. # print(f"[*] Intento de sacar de cola para {client_id}: Cola vacía.") return None, False else: # No hay capacidad disponible. # print(f"[*] Intento de sacar de cola para {client_id}: No hay capacidad disponible.") return None, False def _cleanup_task(self): """Tarea en segundo plano para eliminar datos de clientes inactivos.""" while True: time.sleep(self.cleanup_threshold_seconds / 2) # Comprueba periódicamente 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) # Opcionalmente, vaciar colas para este cliente si es necesario, aunque normalmente no se hace. # Para simplificar, dejamos las solicitudes encoladas como están. print(f"[*] Cliente inactivo '{client_id}' limpiado.") # --- Script de Demostración --- def simulate_requests(rate_limiter: RateLimiter, client_id: str, num_requests: int, delay: float, priority_offset: int = 0): """Simula una ráfaga de solicitudes de un solo cliente.""" print(f"--- Simulando {num_requests} solicitudes del cliente '{client_id}' ---") for i in range(num_requests): # Simula prioridades variables para las solicitudes encoladas priority = i % 5 + priority_offset # Número más bajo = mayor prioridad payload = f"data_{client_id}_{i}" if rate_limiter.allow_request(client_id): print(f"[Hilo {threading.current_thread().name}] Solicitud {i+1}/{num_requests} PERMITIDA para '{client_id}'. Payload: {payload}") else: print(f"[Hilo {threading.current_thread().name}] Solicitud {i+1}/{num_requests} DENEGADA para '{client_id}'. Poniendo en cola con prioridad {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): """Intenta procesar periódicamente las solicitudes diferidas de un cliente.""" print(f"--- Intentando procesar solicitudes diferidas para el cliente '{client_id}' ---") for i in range(num_attempts): time.sleep(1.5) # Espera un poco para que la capacidad se libere potencialmente print(f"[Hilo Procesador] Intento {i+1}/{num_attempts} de sacar de cola para '{client_id}'...") result = rate_limiter.dequeue_and_process(client_id) if result: request, success = result if success: print(f"[Hilo Procesador] Procesado con éxito: {request.payload} (Prioridad: {request.priority})") else: print(f"[Hilo Procesador] Falló el procesamiento (¿problema de capacidad?): {request.payload} (Prioridad: {request.priority}). Vuelve a poner en cola.") else: print(f"[Hilo Procesador] No se procesaron solicitudes diferidas para '{client_id}' en este intento (cola vacía o sin capacidad).") if __name__ == "__main__": print("--- Demostración del Limitador de Tasa ---") # Inicializa el limitador de tasa limiter = RateLimiter(tier_config=TIER_CONFIG, cleanup_threshold_seconds=CLEANUP_THRESHOLD_SECONDS) # Registra clientes 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--- Simulando solicitudes concurrentes ---") # Simula una ráfaga de solicitudes de múltiples hilos threads = [] # Cliente A (nivel gratuito) - es probable que alcance el límite de tasa t1 = threading.Thread(target=simulate_requests, args=(limiter, "client_A", 15, 0.1, 0), name="Thread-A") threads.append(t1) # Cliente B (nivel estándar) - más capacidad t2 = threading.Thread(target=simulate_requests, args=(limiter, "client_B", 60, 0.05, 1), name="Thread-B") threads.append(t2) # Cliente C (nivel premium) - mayor capacidad t3 = threading.Thread(target=simulate_requests, args=(limiter, "client_C", 250, 0.01, 2), name="Thread-C") threads.append(t3) # Cliente D (nivel gratuito) - otro cliente que alcanza el mismo límite t4 = threading.Thread(target=simulate_requests, args=(limiter, "client_D", 10, 0.15, 3), name="Thread-D") threads.append(t4) # Inicia los hilos de simulación for t in threads: t.start() # Da tiempo para las solicitudes iniciales y la puesta en cola time.sleep(2) print("\n--- Intentando procesar solicitudes diferidas ---") # Inicia hilos para procesar solicitudes diferidas para clientes que podrían haber sido limitados por tasa # Intentaremos procesar para los clientes A y D (nivel gratuito) varias veces. 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) # También intentamos para el cliente B, por si acaso 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() # Espera a que todos los hilos de simulación completen for t in threads: t.join() # Espera a que los hilos del procesador completen sus intentos for tp in processor_threads: tp.join() print("\n--- Demostración Completada ---") print("Comprobación final del estado (los recuentos son aproximados debido a la sincronización):") for tier_name, counter in limiter.tier_counters.items(): print(f"Nivel '{tier_name}': Solicitudes actuales en ventana: {counter.get_current_count()} (Máx: {counter.max_requests})") for tier_name, queue in limiter.tier_queues.items(): print(f"Nivel '{tier_name}': Solicitudes diferidas en cola: {len(queue)}") # --- Explicación de las Opciones de Diseño --- # 1. Implementación de Ventana Deslizante: # - Usamos `collections.deque` para almacenar las marcas de tiempo de las solicitudes dentro de la ventana. # - `deque` proporciona complejidad O(1) para agregar nuevas marcas de tiempo y eliminar las antiguas desde la izquierda. # - Cuando se llama a `record_request`, primero eliminamos todas las marcas de tiempo más antiguas que `current_time - window_seconds`. # - Esto asegura que `len(self.request_timestamps)` refleje con precisión el número de solicitudes dentro de la ventana deslizante actual. # - Compromiso: Este enfoque es preciso pero puede implicar múltiples pops en `record_request` si llegan muchas solicitudes en rápida sucesión y luego se detienen. Sin embargo, con el tiempo, el tamaño del deque se gestiona. # - Seguridad de Hilos: Cada instancia de `SlidingWindowCounter` tiene su propio `threading.Lock` para proteger su deque `request_timestamps`, asegurando que las llamadas concurrentes a `record_request` o `get_current_count` para el mismo nivel se serialicen. # 2. Múltiples Niveles: # - La clase `RateLimiter` mantiene un diccionario `tier_config` que mapea nombres de nivel a sus parámetros de límite de tasa (`max_requests`, `window_seconds`). # - También mantiene instancias separadas de `SlidingWindowCounter` y colas de prioridad (`tier_queues`) para cada nivel. # - La asignación de nivel del cliente se almacena en `client_tiers`. # - Esto permite una configuración flexible y el aislamiento de los límites de tasa entre diferentes tipos de clientes. # 3. Cola de Prioridad para Solicitudes Diferidas: # - Para cada nivel, se utiliza una lista estándar de Python como min-heap, gestionada por el módulo `heapq`. # - La clase `Request` está diseñada con `priority` como clave de ordenación principal (valor más bajo = mayor prioridad). # - Se incluye un `request_id` único para asegurar una ordenación estable (FIFO para prioridades iguales) y evitar errores de comparación si las marcas de tiempo son idénticas. # - `enqueue_request` usa `heapq.heappush` para agregar solicitudes, manteniendo la propiedad del heap. # - `dequeue_and_process` usa `heapq.heappop` para recuperar la solicitud de mayor prioridad. # - Crucialmente, `dequeue_and_process` primero comprueba la capacidad disponible *antes* de intentar sacar de la cola del heap. Si hay capacidad disponible, adquiere el bloqueo global, vuelve a comprobar la capacidad (para manejar condiciones de carrera), saca la solicitud e intenta `record_request` nuevamente. Si `record_request` tiene éxito, la solicitud se procesa. Si falla (debido a una condición de carrera donde la capacidad se acaba de llenar), la solicitud se vuelve a poner en el heap. # - Compromiso: Usar la lista + heapq de Python es eficiente para un número moderado de solicitudes diferidas. Para volúmenes extremadamente altos, se podría considerar una implementación de cola de prioridad concurrente más especializada, pero esto agrega una complejidad significativa. # 4. Seguridad de Hilos: # - `SlidingWindowCounter`: Usa un `threading.Lock` interno para su deque `request_timestamps`. # - `RateLimiter`: # - `global_lock` (`threading.Lock`): Protege el estado compartido del cliente (`client_tiers`, `client_last_activity`) y las colas de prioridad de nivel (`tier_queues`). Este bloqueo se adquiere para operaciones como `register_client`, `enqueue_request` y dentro de `dequeue_and_process` al acceder/modificar las colas. # - `allow_request`: Llama al `counter.record_request()` seguro para hilos. También actualiza la actividad del cliente bajo el bloqueo global. # - `dequeue_and_process`: Gestiona cuidadosamente los bloqueos. Primero comprueba la capacidad *sin* el bloqueo global. Si la capacidad *podría* estar disponible, adquiere el bloqueo global, vuelve a comprobar la capacidad, realiza la operación del heap y luego llama nuevamente al `counter.record_request()` seguro para hilos. # - Compromiso: El uso de múltiples bloqueos (uno por contador + un bloqueo global) tiene como objetivo reducir la contención en comparación con un único bloqueo global para todas las operaciones. Sin embargo, las operaciones que abarcan el estado del cliente y el estado del nivel (como sacar de cola) requieren un ordenamiento cuidadoso de los bloqueos o la adquisición del bloqueo global. # 5. Limpieza: # - Un hilo en segundo plano (`_cleanup_task`) escanea periódicamente `client_last_activity`. # - Identifica a los clientes cuya marca de tiempo de la última actividad es anterior a `cleanup_threshold_seconds`. # - Estos clientes se eliminan de `client_tiers` y `client_last_activity` bajo el `global_lock`. # - Esto evita que `RateLimiter` acumule indefinidamente grandes cantidades de estado para clientes inactivos. # - Compromiso: La limpieza no es instantánea; se ejecuta periódicamente. `cleanup_threshold_seconds` determina el compromiso entre el uso de memoria y la capacidad de respuesta de la limpieza. # 6. Casos Límite Manejados: # - Registro duplicado de clientes: `register_client` devuelve False y registra una advertencia. # - Clientes no registrados: `allow_request` y `enqueue_request` comprueban `client_tiers` y devuelven/registran errores. # - Colas de prioridad vacías: `dequeue_and_process` maneja casos donde `self.tier_queues[tier]` está vacío. # - Modificaciones concurrentes: Abordadas mediante el uso adecuado de `threading.Lock`. # - Precisión del reloj: Se utiliza `time.time()`, que generalmente es suficiente. Para requisitos de precisión extremadamente alta, se podrían considerar relojes monotónicos específicos de la plataforma, pero `time.time()` es estándar para este tipo de aplicación. # - Condiciones de carrera en la extracción de cola: Manejadas volviendo a comprobar la capacidad y volviendo a registrar la solicitud dentro de la sección crítica de `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")