Token Bucket Throttling

import asyncio
from typing import Optional

class TokenBucket:
    def __init__(
        self,
        capacity: int = 5,
        refill_interval: float = 1.0,
        refill_amount: int = 1
    ):
        self.capacity = capacity
        self.refill_interval = refill_interval
        self.refill_amount = refill_amount
        
        self.tokens = capacity
        self.refill_task: Optional[asyncio.Task] = None
    
    async def initialize(self):
        """Initialize the token bucket and start the refill task"""
        if self.refill_task is None:
            self.refill_task = asyncio.create_task(self._refill_tokens())
    
    async def _refill_tokens(self):
        """Core token replenishment logic"""
        try:
            while True:
                await asyncio.sleep(self.refill_interval)
                self.tokens = min(self.tokens + self.refill_amount, self.capacity)
        except asyncio.CancelledError:
            pass
    
    def consume_token(self) -> bool:
        """Basic token consumption"""
        if self.tokens > 0:
            self.tokens -= 1
            return True
        return False

from contextlib import asynccontextmanager
from fastapi import Request, Response

@asynccontextmanager
async def middleware(self, request: Request, response: Response):
    """Context manager for token bucket throttling"""
    if self.consume_token():
        # Token available, process request
        response.headers["X-RateLimit-Remaining"] = str(self.tokens)
        yield
    else:
        # No tokens, handle with backoff strategy
        await self.retry_with_backoff(request, response)
        yield

async def retry_with_backoff(
    self,
    request: Request, 
    response: Response, 
    attempt: int = 1
):
    """Retry the request with exponential backoff"""
    # Calculate delay with exponential growth
    delay = min(2 ** attempt * 0.1, 2.0)  # Cap at 2 seconds
    
    async def retry_callback():
        await asyncio.sleep(delay)
        
        if self.consume_token():
            # Got a token on retry
            response.headers["X-RateLimit-Remaining"] = str(self.tokens)
            response.headers["X-RateLimit-Retry-Count"] = str(attempt)
            return
        elif attempt < self.max_retry_attempts:
            # Try again with increased delay
            await self.retry_with_backoff(request, response, attempt + 1)
        else:
            # Exceeded retry attempts
            from fastapi import HTTPException
            raise HTTPException(
                status_code=429,
                detail={
                    "error": "Too Many Requests",
                    "message": "Rate limit exceeded. Try again later."
                }
            )
    
    # Create and track the retry task
    task = asyncio.create_task(retry_callback())
    await task

from typing import Dict
import uuid
import time

class TokenBucket:
    def __init__(self, ...):
        # ... other initialization
        self.pending_tasks: Dict[str, asyncio.Task] = {}
    
    def _generate_request_id(self, request: Request) -> str:
        """Generate a unique request ID"""
        client_ip = request.client.host if request.client else "unknown"
        return f"{client_ip}:{int(time.time() * 1000)}:{uuid.uuid4().hex[:7]}"
    
    @asynccontextmanager
    async def middleware(self, request: Request, response: Response):
        request_id = self._generate_request_id(request)
        
        try:
            # Process the request...
            yield
        finally:
            # Clean up any pending task for this request
            if request_id in self.pending_tasks:
                task = self.pending_tasks.pop(request_id)
                if not task.done():
                    task.cancel()

async def shutdown(self):
    """Clean up resources when shutting down"""
    if self.refill_task:
        self.refill_task.cancel()
        try:
            await self.refill_task
        except asyncio.CancelledError:
            pass
        self.refill_task = None
    
    # Clean up any pending tasks
    for task in list(self.pending_tasks.values()):
        task.cancel()
    
    # Wait for tasks to complete cancellation
    if self.pending_tasks:
        await asyncio.gather(*self.pending_tasks.values(), return_exceptions=True)
    
    self.pending_tasks.clear()

async def retry_callback():
    try:
        # Remove this task from tracking once it executes
        self.pending_tasks.pop(request_id, None)
        
        await asyncio.sleep(delay)
        
        # Process retry logic...
    except asyncio.CancelledError:
        # Request was cancelled (e.g., client disconnected)
        pass

from fastapi import HTTPException

raise HTTPException(
    status_code=429,
    detail={
        "error": "Too Many Requests",
        "message": "Rate limit exceeded",
        "retryAfter": 5
    }
)

import asyncio
import aiohttp

async def make_request(session, request_id):
    try:
        async with session.get('http://localhost:3000/api/test') as response:
            print(f"Request {request_id}: Success - Status: {response.status}")
            return response.status
    except Exception as err:
        print(f"Request {request_id}: Failed - {str(err)}")
        return None

async def test_throttling():
    async with aiohttp.ClientSession() as session:
        # Send 30 concurrent requests
        tasks = [make_request(session, i+1) for i in range(30)]
        results = await asyncio.gather(*tasks)
        
        print(f"Results: {results}")

# Run the test
asyncio.run(test_throttling())