Add a client-side rate limiting middleware

Add RateLimitMiddleware, a token-bucket client middleware that throttles
outgoing requests to a configurable rate and burst, with optional
per-domain buckets and numeric Retry-After handling on HTTP 429.

Promoted from the example in aiohttp#11969 (which was moved here): drop the
demo and module-level logging config, keep a strong reference to the
scheduler task so it cannot be garbage collected mid-run, and add tests
and documentation.

Author: rodrigobnogueira

CHANGES/12.feature.rst

@@ -0,0 +1,3 @@
+Added :class:`~aiohttp_client_middlewares.RateLimitMiddleware`, a client-side
+token-bucket rate limiter with optional per-domain buckets and ``Retry-After``
+handling -- by :user:`rodrigobnogueira`.

aiohttp_client_middlewares/__init__.py

@@ -1,11 +1,12 @@
 """Client middlewares for :mod:`aiohttp`.
 
 This package is the canonical home for reusable aiohttp *client* middlewares,
-starting with HTTP Digest authentication.
+starting with HTTP Digest authentication and client-side rate limiting.
 """
 
 from .digest_auth import DigestAuthMiddleware
+from .rate_limit import RateLimitMiddleware
 
 __version__ = "0.1.0"
 
-__all__ = ("DigestAuthMiddleware",)
+__all__ = ("DigestAuthMiddleware", "RateLimitMiddleware")

aiohttp_client_middlewares/rate_limit.py

@@ -0,0 +1,146 @@
+"""Client-side rate-limiting middleware for aiohttp.
+
+This middleware throttles outgoing requests using a token-bucket algorithm.
+It is *not* server-side rate limiting -- it limits how fast the client sends
+requests so it does not overwhelm upstream servers or exceed API quotas.
+
+Features:
+- Configurable rate and burst size
+- Optional per-domain buckets
+- Automatic ``Retry-After`` header handling
+"""
+
+import asyncio
+import logging
+import time
+from collections import defaultdict, deque
+from http import HTTPStatus
+
+from aiohttp import ClientHandlerType, ClientRequest, ClientResponse
+
+_LOGGER = logging.getLogger(__name__)
+
+
+class TokenBucket:
+    """FIFO token-bucket using an ``asyncio.Event`` queue.
+
+    Each caller appends its own event to a FIFO queue and waits. A single
+    ``_schedule`` coroutine services the queue front-to-back, sleeping until
+    each slot's send time arrives and then unblocking the corresponding
+    caller. This guarantees strict FIFO ordering even under high concurrency.
+    """
+
+    def __init__(self, rate: float, burst: int) -> None:
+        self._interval = 1.0 / rate
+        self._burst = burst
+        # Start *burst* intervals in the past so the first ``burst`` acquires
+        # are instant.
+        self._next_send = time.monotonic() - burst * self._interval
+        self._waiters: deque[asyncio.Event] = deque()
+        self._scheduling = False
+        # Keep a strong reference to the scheduler task: the event loop only
+        # holds a weak reference to it, so an otherwise-unreferenced task can
+        # be garbage collected mid-run.
+        self._scheduler_task: asyncio.Task[None] | None = None
+
+    async def acquire(self) -> None:
+        """Reserve the next send slot and wait until it arrives."""
+        event = asyncio.Event()
+        self._waiters.append(event)
+        self._ensure_scheduling()
+        await event.wait()
+
+    def _ensure_scheduling(self) -> None:
+        """Start the scheduler loop if it is not already running."""
+        if not self._scheduling:
+            self._scheduling = True
+            self._scheduler_task = asyncio.ensure_future(self._schedule())
+
+    async def _schedule(self) -> None:
+        """Service waiters in FIFO order, one slot at a time."""
+        while self._waiters:
+            now = time.monotonic()
+            # Cap drift so idle periods never accumulate more than *burst*
+            # free slots.
+            self._next_send = max(self._next_send, now - self._burst * self._interval)
+            self._next_send += self._interval
+            delay = self._next_send - now
+            if delay > 0:
+                await asyncio.sleep(delay)
+            self._waiters.popleft().set()
+        self._scheduling = False
+
+
+class RateLimitMiddleware:
+    """Client middleware that throttles requests with a token bucket.
+
+    The middleware delays each outgoing request until the bucket grants it a
+    slot, so the client never sends faster than ``rate`` requests per second
+    (allowing short bursts of up to ``burst`` requests).
+
+    :param float rate: Sustained request rate, in requests per second.
+    :param int burst: Number of requests allowed to go out back-to-back before
+        throttling kicks in.
+    :param bool per_domain: When ``True``, keep an independent bucket per
+        target host instead of a single global bucket.
+    :param bool respect_retry_after: When ``True``, sleep for the duration of a
+        numeric ``Retry-After`` header on an HTTP 429 response before returning
+        it to the caller.
+    """
+
+    rate: float
+    burst: int
+    per_domain: bool
+    respect_retry_after: bool
+
+    def __init__(
+        self,
+        rate: float = 10.0,
+        burst: int = 10,
+        per_domain: bool = False,
+        respect_retry_after: bool = True,
+    ) -> None:
+        self.rate = rate
+        self.burst = burst
+        self.per_domain = per_domain
+        self.respect_retry_after = respect_retry_after
+        self._global_bucket = TokenBucket(rate, burst)
+        self._domain_buckets: dict[str, TokenBucket] = defaultdict(
+            lambda: TokenBucket(rate, burst)
+        )
+
+    def _get_bucket(self, request: ClientRequest) -> TokenBucket:
+        if self.per_domain:
+            domain = request.url.host or "unknown"
+            return self._domain_buckets[domain]
+        return self._global_bucket
+
+    async def _handle_retry_after(self, response: ClientResponse) -> None:
+        if response.status != HTTPStatus.TOO_MANY_REQUESTS:
+            return
+        retry_after = response.headers.get("Retry-After")
+        if retry_after:
+            try:
+                wait_seconds = float(retry_after)
+                _LOGGER.info("Server requested Retry-After: %ss", wait_seconds)
+                await asyncio.sleep(wait_seconds)
+            except ValueError:
+                _LOGGER.debug(
+                    "Retry-After is not a number (likely HTTP-date): %s", retry_after
+                )
+
+    async def __call__(
+        self,
+        request: ClientRequest,
+        handler: ClientHandlerType,
+    ) -> ClientResponse:
+        """Run the request through the rate limiter."""
+        bucket = self._get_bucket(request)
+        await bucket.acquire()
+
+        response = await handler(request)
+
+        if self.respect_retry_after:
+            await self._handle_retry_after(response)
+
+        return response

docs/api.rst

@@ -75,3 +75,38 @@ Digest authentication
            # The middleware automatically handles the digest auth handshake.
            async with session.get("http://protected.example.com") as resp:
                assert resp.status == 200
+
+
+Rate limiting
+-------------
+
+.. class:: RateLimitMiddleware(rate=10.0, burst=10, per_domain=False, respect_retry_after=True)
+
+   Client middleware that throttles outgoing requests with a token bucket.
+
+   :param float rate: Sustained request rate, in requests per second.
+   :param int burst: Number of requests allowed to go out back-to-back before
+      throttling kicks in.
+   :param bool per_domain: Keep an independent bucket per target host instead
+      of a single global bucket.
+   :param bool respect_retry_after: Sleep for the duration of a numeric
+      ``Retry-After`` header on an HTTP 429 response before returning it to the
+      caller.
+
+   The middleware delays each request until the bucket grants it a slot, so the
+   client never sends faster than ``rate`` requests per second while still
+   allowing short bursts of up to ``burst`` requests. Slots are served in strict
+   FIFO order.
+
+   **Usage**
+
+   ::
+
+       from aiohttp import ClientSession
+       from aiohttp_client_middlewares import RateLimitMiddleware
+
+       # At most 5 requests/second, bursting up to 2.
+       rate_limit = RateLimitMiddleware(rate=5.0, burst=2)
+       async with ClientSession(middlewares=(rate_limit,)) as session:
+           async with session.get("http://example.com") as resp:
+               assert resp.status == 200

docs/code/rate_limit.py

@@ -0,0 +1,19 @@
+"""Quickstart example for the rate-limiting middleware."""
+
+import asyncio
+
+from aiohttp import ClientSession
+
+from aiohttp_client_middlewares import RateLimitMiddleware
+
+
+async def main() -> None:
+    # Throttle to at most 5 requests per second, allowing bursts of up to 2.
+    rate_limit = RateLimitMiddleware(rate=5.0, burst=2)
+    async with ClientSession(middlewares=(rate_limit,)) as session:
+        for _ in range(10):
+            async with session.get("https://httpbin.org/get") as resp:
+                print("Status:", resp.status)
+
+
+asyncio.run(main())

docs/index.rst

@@ -8,6 +8,8 @@ This package collects ready-to-use middlewares for
 
 - :class:`~aiohttp_client_middlewares.DigestAuthMiddleware` -- HTTP Digest
   authentication.
+- :class:`~aiohttp_client_middlewares.RateLimitMiddleware` -- client-side
+  token-bucket rate limiting.
 
 
 Installation
@@ -21,11 +23,15 @@ Installation
 Quickstart
 ----------
 
-Attach a middleware to a session through the ``middlewares`` argument and
-let it handle authentication for every request:
+Attach one or more middlewares to a session through the ``middlewares``
+argument. For HTTP Digest authentication:
 
 .. literalinclude:: code/digest_auth.py
 
+For client-side rate limiting:
+
+.. literalinclude:: code/rate_limit.py
+
 
 Contents
 --------

tests/test_rate_limit.py

@@ -0,0 +1,115 @@
+"""Tests for the rate-limiting middleware."""
+
+import asyncio
+import time
+
+from aiohttp import web
+from pytest_aiohttp import AiohttpClient
+
+from aiohttp_client_middlewares.rate_limit import RateLimitMiddleware, TokenBucket
+
+
+async def _ok_handler(request: web.Request) -> web.Response:
+    return web.Response(text="OK")
+
+
+def _make_app() -> web.Application:
+    app = web.Application()
+    app.router.add_get("/api", _ok_handler)
+    return app
+
+
+async def test_token_bucket_allows_burst() -> None:
+    """Tokens up to burst size should be available immediately."""
+    bucket = TokenBucket(rate=10.0, burst=3)
+    start = time.monotonic()
+    for _ in range(3):
+        await bucket.acquire()
+    elapsed = time.monotonic() - start
+    # All three should be near-instant (within the burst allowance).
+    assert elapsed < 0.05
+
+
+async def test_token_bucket_refills_after_idle() -> None:
+    """After draining, idle time should replenish burst slots."""
+    bucket = TokenBucket(rate=100.0, burst=1)
+    await bucket.acquire()
+    await asyncio.sleep(0.05)
+    start = time.monotonic()
+    await bucket.acquire()
+    elapsed = time.monotonic() - start
+    # Should be near-instant because idle time refilled the slot.
+    assert elapsed < 0.05
+
+
+async def test_token_bucket_fifo_ordering() -> None:
+    """Concurrent acquires should be served in FIFO order."""
+    bucket = TokenBucket(rate=100.0, burst=1)
+    order: list[int] = []
+
+    async def numbered_acquire(n: int) -> None:
+        await bucket.acquire()
+        order.append(n)
+
+    tasks = [asyncio.create_task(numbered_acquire(i)) for i in range(3)]
+    await asyncio.gather(*tasks)
+    assert order == [0, 1, 2]
+
+
+async def test_rate_limit_middleware_throttles(aiohttp_client: AiohttpClient) -> None:
+    """Global middleware should throttle requests beyond burst."""
+    middleware = RateLimitMiddleware(rate=50.0, burst=2)
+    client = await aiohttp_client(_make_app(), middlewares=(middleware,))
+
+    start = time.monotonic()
+    for _ in range(4):
+        resp = await client.get("/api")
+        assert resp.status == 200
+    elapsed = time.monotonic() - start
+
+    # 2 burst + 2 throttled at 50/s ~= 0.04s minimum wait. The upper bound
+    # catches hangs or accidental double-sleeps while staying generous for CI.
+    assert 0.02 <= elapsed < 0.5
+
+
+async def test_rate_limit_middleware_per_domain(aiohttp_client: AiohttpClient) -> None:
+    """Per-domain buckets should still throttle requests to the same host."""
+    middleware = RateLimitMiddleware(rate=100.0, burst=1, per_domain=True)
+    client = await aiohttp_client(_make_app(), middlewares=(middleware,))
+
+    start = time.monotonic()
+    # Same host, so the two requests share a bucket and the second one waits.
+    resp1 = await client.get("/api")
+    resp2 = await client.get("/api")
+    elapsed = time.monotonic() - start
+
+    assert resp1.status == 200
+    assert resp2.status == 200
+    assert 0.005 <= elapsed < 0.5
+
+
+async def test_rate_limit_middleware_respects_retry_after(
+    aiohttp_client: AiohttpClient,
+) -> None:
+    """The middleware should sleep on a 429 with a numeric ``Retry-After``."""
+    call_count = 0
+
+    async def rate_limited_handler(request: web.Request) -> web.Response:
+        nonlocal call_count
+        call_count += 1
+        if call_count <= 1:
+            return web.Response(status=429, headers={"Retry-After": "0.1"})
+        return web.Response(text="OK")
+
+    app = web.Application()
+    app.router.add_get("/api", rate_limited_handler)
+
+    middleware = RateLimitMiddleware(rate=100.0, burst=10, respect_retry_after=True)
+    client = await aiohttp_client(app, middlewares=(middleware,))
+
+    start = time.monotonic()
+    resp = await client.get("/api")
+    elapsed = time.monotonic() - start
+
+    assert resp.status == 429
+    assert 0.08 <= elapsed < 0.5