The Token-Bucket Algorithm: Smart Rate-Limiting for Modern API Workloads

╭───────────────── Bucket (Capacity: 10) ──────────────────╮
│ Tokens: 🟢🟢🟢🟢🟢⚪⚪⚪⚪⚪ (5 remaining)                   │
╰──────────────────────────────────────────────────────────╯
▲                                      │
│ Consume 1 token per request          │ Refills at a slow,
▼                                      │ steady rate (e.g., 1 token / 5 sec)
┌──────────┐                            ┌────────────┐
│ API Call │ ─ ─ ─ ─ Needs Token ─ ─ \> │ Is Token   │
└──────────┘                            │ Available? │─ ─ YES ─ \> Process Request
└─────┬──────┘
│
└─ NO ─ ─ ─ \> Reject (HTTP 429)

import time
import threading
from typing import final

@final
class TokenBucket:
    """
    Thread-safe Token Bucket for regulating request rates.

    Attributes:
        capacity (float): Maximum tokens in the bucket (burst limit).
        _refill_rate (float): Tokens replenished per second.
        _tokens (float): Current token count.
        _last_refill_time (float): Last timestamp tokens were refilled.
        _lock (threading.Lock): Lock to ensure thread-safe operations.
    """

    def __init__(self, capacity: float, refill_rate: float):
        """
        Initialize a TokenBucket.

        Args:
            capacity (float): >0, max tokens bucket can hold.
            refill_rate (float): >0, tokens added each second.

        Raises:
            ValueError: if capacity or refill_rate are non-positive.
        """
        # Validate inputs
        if not (capacity > 0 and refill_rate > 0):
            raise ValueError("Capacity and refill rate must be positive")

        # Core parameters
        self.capacity = float(capacity)
        self._refill_rate = float(refill_rate)

        # Start full
        self._tokens = float(capacity)
        # Timestamp for next refill calculation
        self._last_refill_time = time.monotonic()

        # Lock for thread-safe token modifications
        self._lock = threading.Lock()

    def _refill(self) -> None:
        """
        Refill tokens based on elapsed time since last refill.
        Must be called while holding the lock.

        Calculates:
            elapsed = now - _last_refill_time
            tokens_to_add = elapsed * _refill_rate
        Applies cap: min(capacity, _tokens + tokens_to_add)
        Updates _last_refill_time to now.
        """
        now = time.monotonic()
        elapsed = now - self._last_refill_time

        # Only refill if time has advanced
        if elapsed > 0:
            tokens_to_add = elapsed * self._refill_rate
            # Increase tokens, but do not exceed capacity
            self._tokens = min(self.capacity, self._tokens + tokens_to_add)
            # Reset refill timestamp
            self._last_refill_time = now

    def consume(self, tokens: float = 1.0) -> bool:
        """
        Attempt to remove tokens from the bucket.

        Args:
            tokens (float): Number of tokens to consume (>0).

        Returns:
            bool: True if tokens were available and consumed, False otherwise.

        Raises:
            ValueError: if requested token count is non-positive.
        """
        if tokens <= 0:
            raise ValueError("Number of tokens to consume must be positive")

        with self._lock:
            # Refill before attempting consumption
            self._refill()

            # If enough tokens, deduct and allow
            if self._tokens >= tokens:
                self._tokens -= tokens
                return True

            # Not enough tokens, rate limit exceeded
            return False

    @property
    def tokens(self) -> float:
        """
        Get current token count, after refilling.

        Returns:
            float: current available tokens (<= capacity).
        """
        with self._lock:
            # Ensure up-to-date token count
            self._refill()
            return self._tokens

import unittest
import time
from token_bucket import TokenBucket

class TestTokenBucketRealistic(unittest.TestCase):
    """
    Test suite for the TokenBucket algorithm, as used in an API rate limiter.

    Scenario:
      - Burst capacity: 10 tokens (max immediate requests)
      - Refill rate:    5 tokens per second (sustained rate)

    Each test prints actual vs. expected outcomes for clarity in a blog post.
    """

    def test_burst_limit(self):
        """
        Burst Test:
        Simulates sending 11 rapid API requests. The bucket should allow
        the first 10 (up to capacity) and deny the 11th.
        """
        bucket = TokenBucket(capacity=10, refill_rate=5)
        print("\n--- Burst Test (capacity=10) ---")
        allowed_count = 0

        for i in range(1, 12):
            # Attempt to consume 1 token per request
            allowed = bucket.consume()
            expected = (i <= 10)
            status = "ALLOWED" if allowed else "DENIED"
            exp_status = "ALLOWED" if expected else "DENIED"
            print(f" Request #{i:2d}: actual={status:7s}  expected={exp_status}")
            if allowed:
                allowed_count += 1

        # Verify exactly 10 requests were permitted
        self.assertEqual(allowed_count, 10, f"Allowed {allowed_count} requests; expected 10")
        # The next immediate request must be denied
        self.assertFalse(bucket.consume(), "11th request must be DENIED")

    def test_sustained_rate(self):
        """
        Sustained Throughput Test:
        After draining the bucket, wait 1 second to allow the refill rate
        to add tokens. Exactly 5 requests should succeed (5 tokens/sec).
        """
        bucket = TokenBucket(capacity=10, refill_rate=5)
        # Drain the bucket fully (simulate a sudden burst)
        for _ in range(10):
            bucket.consume()

        print("\n--- Sustained Rate Test (5 tokens/sec) ---")
        print(" Waiting 1 second to refill...")
        time.sleep(1)

        # Check how many tokens were refilled
        available = bucket.tokens
        print(f" Tokens refilled: actual={available:.2f}  expected≈5.00")
        self.assertGreaterEqual(available, 4.9,
                                f"Expected ≥4.9 tokens, got {available:.2f}")

        successes = 0
        # Attempt 6 subsequent requests: only first 5 should pass
        for i in range(1, 7):
            allowed = bucket.consume()
            expected = (i <= 5)
            status = "ALLOWED" if allowed else "DENIED"
            exp_status = "ALLOWED" if expected else "DENIED"
            print(f" Request #{i:2d} after refill: actual={status:7s}  expected={exp_status}")
            if allowed:
                successes += 1

        self.assertEqual(successes, 5, f"Allowed {successes} after refill; expected 5")
        self.assertFalse(bucket.consume(), "6th request after refill must be DENIED")

    def test_no_overfill(self):
        """
        Idle Capping Test:
        Consume a few tokens, then wait long enough that the theoretical
        refill would exceed capacity. The bucket should never exceed its max.
        """
        bucket = TokenBucket(capacity=10, refill_rate=5)
        # Consume 3 tokens (immediate requests)
        bucket.consume(3)
        current = bucket.tokens
        print("\n--- No Overfill Test ---")
        print(f" After consuming 3: actual={current:.2f}  expected≈7.00")

        print(" Waiting 3 seconds (would refill 15 tokens) ...")
        time.sleep(3)

        final = bucket.tokens
        print(f" After idle: actual={final:.2f}  expected=10.00 (cap)")
        # Ensure the bucket caps at its capacity
        self.assertAlmostEqual(final, 10.0, delta=1e-3, msg=f"Final tokens {final:.2f}; should cap at 10.00")

if __name__ == '__main__':
    # Default verbosity=1: clean pass/fail summary
    unittest.main()

{
  "userId": "user-12345",
  "plan": "pro",
  "buckets": {
    "premium_models": {
      "capacity": 200,
      "refillRate": 0.0023148, // 200 / 86400
      "tokens": 145.7,
      "lastRefillTime": 1718992345.678
    }
  }
}