Queue Based Throttling

import asyncio
from dataclasses import dataclass
from fastapi import HTTPException, Request, Response
from typing import List, Optional
import time

# Configuration constants
MAX_CONCURRENT = 3
MAX_QUEUE_SIZE = 10
QUEUE_TIMEOUT_MS = 5000  # 5 seconds timeout

current_requests = 0

@dataclass
class QueuedRequest:
    request: Request
    response: Response
    timestamp: float  # When the request was queued
    ready_event: asyncio.Event  # Signal when request is ready or timed out
    timed_out: bool = False  # Flag to indicate if request timed out

request_queue: List[QueuedRequest] = []

class QueueThrottler:
    def __init__(self):
        self.queue_task: Optional[asyncio.Task] = None
        self.running = False
    
    async def initialize(self):
        """Initialize the queue processing task"""
        if self.queue_task is None:
            self.running = True
            self.queue_task = asyncio.create_task(self._process_queue())
    
    async def _process_queue(self):
        """Continuously process the queue"""
        global current_requests, request_queue
        
        try:
            while self.running:
                await asyncio.sleep(0.1)  # Process every 100ms
                
                now = time.time() * 1000  # Convert to milliseconds
                
                # First, remove expired requests
                i = 0
                while i < len(request_queue):
                    queued_req = request_queue[i]
                    if now - queued_req.timestamp > QUEUE_TIMEOUT_MS:
                        # Request has timed out, signal the handler
                        queued_req.timed_out = True
                        queued_req.ready_event.set()
                        request_queue.pop(i)
                    else:
                        i += 1
                
                # Process eligible requests
                while current_requests < MAX_CONCURRENT and len(request_queue) > 0:
                    queued_req = request_queue.pop(0)  # FIFO - remove from front
                    current_requests += 1
                    
                    # Signal the request is ready to proceed
                    queued_req.ready_event.set()
                    
                    # Create a task to handle the request completion tracking
                    asyncio.create_task(self._track_request_completion())
        
        except asyncio.CancelledError:
            pass
    
    async def _track_request_completion(self):
        """Track when a request completes to decrement the counter"""
        global current_requests
        
        try:
            # Wait for the request to complete (this is a simplified approach)
            # In a real implementation, you'd integrate more closely with FastAPI's lifecycle
            await asyncio.sleep(0.1)  # Small delay to allow request processing
            
        finally:
            current_requests -= 1

from contextlib import asynccontextmanager

# Global throttler instance
throttler = QueueThrottler()

@asynccontextmanager
async def queue_throttle(request: Request, response: Response):
    """Queue-based throttling context manager"""
    global current_requests, request_queue
    
    # Initialize throttler if needed
    if throttler.queue_task is None:
        await throttler.initialize()
    
    # Check if queue is full
    if len(request_queue) >= MAX_QUEUE_SIZE:
        raise HTTPException(
            status_code=503,
            detail="Queue is full, please try again later."
        )
    
    # Check if we can process immediately
    if current_requests < MAX_CONCURRENT:
        current_requests += 1
        try:
            yield
        finally:
            current_requests -= 1
    else:
        # Add to queue
        queued_request = QueuedRequest(
            request=request,
            response=response,
            timestamp=time.time() * 1000,
            ready_event=asyncio.Event()
        )
        request_queue.append(queued_request)
        
        # Wait for request to be ready or timeout
        await queued_request.ready_event.wait()
        
        # Check if request timed out
        if queued_request.timed_out:
            raise HTTPException(
                status_code=408,
                detail="Request timeout while waiting in queue"
            )
        
        yield

# Cleanup function
async def cleanup_throttler():
    """Clean up the throttler on shutdown"""
    throttler.running = False
    if throttler.queue_task:
        throttler.queue_task.cancel()
        try:
            await throttler.queue_task
        except asyncio.CancelledError:
            pass
        throttler.queue_task = None

# Usage in FastAPI route
@router.get("/api/data")
async def get_data(request: Request, response: Response):
    async with queue_throttle(request, response):
        # Your API logic here
        await asyncio.sleep(1)  # Simulate processing
        return {"message": "Data processed successfully"}

# Key points from testing with 15 concurrent requests
# 1. First 3 requests (MAX_CONCURRENT) complete quickly
# 2. Next 10 requests (MAX_QUEUE_SIZE) are queued and complete with increasing delays
# 3. Remaining 2 requests immediately receive 503 errors (queue full)