Introduction to Throttling and Token Bucket

Welcome to the second lesson of our course on Throttling API Requests. In this lesson, we will delve into the Token Bucket algorithm, a powerful method for implementing throttling in Python-based APIs.

What is the Token Bucket Algorithm?

The Token Bucket algorithm is a rate limiting method that allows for controlled bursts of activity while maintaining a consistent average rate. Here's how it works:

  1. You have a "bucket" that holds tokens (representing request capacity)
  2. Tokens are added to the bucket at a fixed rate
  3. When a request arrives, it needs to consume a token to proceed
  4. If the bucket is empty, the request must either wait or be rejected

Advantages:

  • Allows for bursts of traffic (unlike fixed window limiters)
  • Simple to implement and understand
  • Low memory footprint
  • Configurable parameters for different scenarios

Disadvantages:

  • Requires ongoing token management (via async tasks)
  • May introduce slight latency for token checks
  • Needs careful tuning to balance performance and protection
Core Components of Token Bucket Implementation: 1. Token Management

Let's look at the key components needed to implement a token bucket throttle.

The heart of the algorithm is token management - tracking available tokens and replenishing them:

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

This simplified implementation shows the two essential operations:

  • Token replenishment: Adding tokens back to the bucket using async tasks
  • Token consumption: Checking for and using available tokens

The critical parameters that control throttling behavior are:

  • capacity: Maximum tokens (requests) that can be processed at once
  • refill_interval: How often tokens are added (in seconds)
  • refill_amount: Number of tokens added each interval
2. Request Handling with Context Managers

To integrate with FastAPI, we create a context manager that uses our token bucket:

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

The context manager performs a simple check - if a token is available, the request proceeds; otherwise, it's handled with a backoff strategy.

3. Exponential Backoff Strategy

A sophisticated throttling implementation doesn't just reject excess requests - it can attempt to process them when capacity becomes available:

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

The key insight here is the exponential backoff formula: 2 ** attempt * 0.1. This creates increasingly longer delays between retries (0.1s, 0.2s, 0.4s, 0.8s, etc.) up to a maximum of 2 seconds. This approach prevents overwhelming the server with retry attempts, spreading the load over time.

4. Resource Management (The Hard Part)

The most challenging aspect of implementing a token bucket is proper resource management. Issues to handle include:

  • Tracking pending requests: Each delayed request creates an asyncio task that needs to be tracked and potentially canceled.

  • Client disconnection handling: When a client disconnects while waiting for a retry, we need to clean up associated resources:

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()
  • Application shutdown: When the application shuts down, we need to clear all tasks:
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()
  • Task cancellation handling: Before processing a delayed request, handle cancellation gracefully:
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
Real-World Considerations

When implementing token bucket throttling in production, consider:

  1. Distributed systems: For APIs running on multiple servers, you'll need a shared token bucket, often implemented using Redis with Python's redis-py library.

  2. User identification: Instead of a global bucket, create buckets per user, API key, or IP address to prevent one user from consuming all capacity.

  3. Informative responses: Use headers to inform clients about rate limits:

    • X-RateLimit-Limit: Maximum capacity
    • X-RateLimit-Remaining: Current tokens available
    • X-RateLimit-Reset: When the bucket will refill

  4. Client guidance: Return clear error messages with retry recommendations:

from fastapi import HTTPException

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

To observe throttling in action, send a burst of requests using aiohttp:

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())

This will demonstrate the throttling behavior:

  • Initial requests succeed immediately (using available tokens)
  • Subsequent requests succeed with delays (as tokens replenish)
  • Final requests may fail with 429 status (after maximum retries)
Summary and Next Steps

In this lesson, we explored the Token Bucket algorithm for throttling API requests in Python. We focused on the key concepts and implementation challenges:

  1. Token management: Tracking and replenishing tokens using asyncio tasks
  2. Request handling: Processing or delaying requests using context managers
  3. Exponential backoff: Intelligently spacing retry attempts to reduce server load using asyncio.sleep
  4. Resource management: The hardest part - properly tracking and cleaning up asyncio tasks

As you move to the practice exercises, experiment with different configurations to see how changing parameters affects throttling behavior. This hands-on experience will help you understand how to apply throttling effectively in real-world Python API scenarios.

Previous
Next
Sign up
Join the 1M+ learners on CodeSignal
Be a part of our community of 1M+ users who develop and demonstrate their skills on CodeSignal