Revenue Management Pricing & forecasting pipelines
4 min read Ingestion

Implementing a token bucket rate limiter in Python

OTA APIs publish a request budget, and the fastest way to lose access to one is to spend that budget in a burst. A token bucket smooths outbound traffic to a sustainable rate while still allowing short, controlled bursts — exactly the shape most quotas are designed around. This guide builds one, extending Rate Limiting & Retry Strategies in the Data Ingestion & OTA API Integration Workflows pillar. It pairs naturally with the backoff logic that handles the 429s a limiter is meant to prevent.

Prerequisites

  • Python 3.11+ with asyncio
  • Standard library time, logging, dataclasses
  • The provider’s documented sustained rate and burst allowance
  • A per-domain limiter instance, since quotas are per-provider
  • Backoff for the rare 429, as in Troubleshooting OTA integration failures

Step 1 — Model the bucket

The bucket holds tokens up to a capacity and refills at a steady rate. A request costs one token; when the bucket is empty, the caller waits for enough refill. Using a monotonic clock avoids the drift a wall clock introduces over a long-running pipeline.

python 
from __future__ import annotations

import logging
import time
from dataclasses import dataclass, field

logger = logging.getLogger("ota.ratelimit")


@dataclass
class TokenBucket:
    rate_per_sec: float          # sustained refill rate
    capacity: float              # max burst
    _tokens: float = field(init=False)
    _last: float = field(init=False)

    def __post_init__(self) -> None:
        if self.rate_per_sec <= 0 or self.capacity <= 0:
            raise ValueError("rate and capacity must be positive")
        self._tokens = self.capacity
        self._last = time.monotonic()

    def _refill(self, now: float) -> None:
        elapsed = now - self._last
        self._tokens = min(self.capacity, self._tokens + elapsed * self.rate_per_sec)
        self._last = now

Step 2 — Compute the wait for a token

Separate the decision (how long until a token is available) from the waiting, so the same logic is testable without sleeping. If a token is ready, the wait is zero; otherwise it is the time to refill the deficit.

python 
    def time_until_available(self, now: float, cost: float = 1.0) -> float:
        self._refill(now)
        if self._tokens >= cost:
            return 0.0
        return (cost - self._tokens) / self.rate_per_sec

Step 3 — Acquire asynchronously without blocking the loop

Acquisition waits the computed time, then deducts the token. Because it uses asyncio.sleep, many coroutines can share one limiter without blocking the event loop — the whole point of an async ingestion pipeline.

python 
import asyncio


class AsyncLimiter:
    def __init__(self, rate_per_sec: float, capacity: float) -> None:
        self._bucket = TokenBucket(rate_per_sec, capacity)
        self._lock = asyncio.Lock()

    async def acquire(self, cost: float = 1.0) -> None:
        async with self._lock:
            wait = self._bucket.time_until_available(time.monotonic(), cost)
            if wait > 0:
                logger.debug("rate limit: sleeping %.3fs", wait)
                await asyncio.sleep(wait)
                self._bucket._refill(time.monotonic())
            self._bucket._tokens -= cost

Holding the lock across the wait serializes token accounting so two coroutines cannot both believe the last token is theirs — the classic bug that lets a limiter briefly exceed its own rate.

Verification and testing

Test the pure decision function with an injected clock, so the suite runs instantly rather than sleeping.

python 
def test_bucket_enforces_rate() -> None:
    bucket = TokenBucket(rate_per_sec=2.0, capacity=2.0)
    t0 = bucket._last
    # Drain the burst capacity
    assert bucket.time_until_available(t0) == 0.0
    bucket._tokens -= 1
    assert bucket.time_until_available(t0) == 0.0
    bucket._tokens -= 1
    # Empty now; next token needs 0.5s at 2/sec
    assert abs(bucket.time_until_available(t0) - 0.5) < 1e-9

Common pitfalls and edge cases

  • Wall-clock refill. time.time() jumps on NTP correction; use time.monotonic() for elapsed-time math.
  • Global limiter across domains. One bucket for all providers throttles fast APIs to the slowest; instantiate per domain.
  • Unlocked accounting. Concurrent acquire without a lock double-spends the last token; serialize the check-and-deduct.
  • Capacity equals rate. Setting burst capacity to one removes all burst tolerance and stalls under normal jitter; size capacity to the provider’s documented burst.
  • Sleeping under the lock forever. A pathological cost larger than capacity waits indefinitely; validate that cost never exceeds capacity.