Give clone() the host it is being built for (#27)

Author: rodrigobnogueira

aiohttp_client_middlewares/rate_limit.py

@@ -13,7 +13,6 @@
 import math
 import time
 from abc import ABC, abstractmethod
-from collections import defaultdict
 
 from aiohttp import ClientHandlerType, ClientRequest, ClientResponse, ClientTimeout
 
@@ -46,11 +45,15 @@ async def acquire(self) -> float:
         """
 
     @abstractmethod
-    def clone(self) -> "RateLimiter":
-        """Return a fresh limiter with the same configuration.
-
-        Per-domain mode clones the configured limiter once per target
-        host, so state (queued slots, accrued tokens) must not carry over.
+    def clone(self, host: str, /) -> "RateLimiter":
+        """Return a fresh limiter, configured the same, scoped to *host*.
+
+        Per-domain mode calls this the first time it meets a host, so state
+        (queued slots, accrued tokens) must not carry over. Threads racing on
+        that first contact may each build one and only one is kept, so the
+        call itself should have no side effects. An algorithm that keeps its
+        state in-process can ignore *host*; one that keeps it in a shared
+        backend needs it in the key, or every host draws on one limit.
         """
 
     def release(self) -> None:
@@ -147,8 +150,11 @@ async def acquire(self) -> float:
         self._tokens -= 1.0
         return max(0.0, -self._tokens) * self._interval
 
-    def clone(self) -> "TokenBucket":
-        """Return a fresh, full bucket with the same rate and burst."""
+    def clone(self, host: str, /) -> "TokenBucket":
+        """Return a fresh, full bucket with the same rate and burst.
+
+        The bucket's state is per-object, so *host* needs no part in it.
+        """
         return TokenBucket(rate=self._rate, burst=int(self._burst))
 
     def release(self) -> None:
@@ -182,7 +188,7 @@ class RateLimitMiddleware:
     :param RateLimiter limiter: The :class:`RateLimiter` to throttle with --
         for example ``TokenBucket(rate=5.0, burst=2)``. With
         ``per_domain=True`` it acts as a template: each target host gets
-        ``limiter.clone()`` the first time that host is seen.
+        ``limiter.clone(host)`` the first time that host is seen.
     :param bool per_domain: When ``True``, keep an independent limiter per
         target host instead of a single global one. Limiters are keyed on the
         URL host only (port and scheme are not distinguished) and are never
@@ -198,31 +204,39 @@ def __init__(
     ) -> None:
         if not isinstance(limiter, RateLimiter):
             raise TypeError(f"limiter must be a RateLimiter, got {limiter!r}")
-        self._per_domain = per_domain
-        self._global_limiter: RateLimiter | None = None
-        if per_domain:
-            self._domain_limiters: dict[str, RateLimiter] = defaultdict(limiter.clone)
-        else:
-            self._global_limiter = limiter
+        # The one limiter in global mode; the template to clone in per-domain
+        # mode. Whether the per-host dict exists is what says which.
+        self._limiter = limiter
+        self._domain_limiters: dict[str, RateLimiter] | None = (
+            {} if per_domain else None
+        )
 
     @property
     def per_domain(self) -> bool:
         """Whether this middleware keeps one limiter per target host.
 
-        Read-only: the limiters are built once in ``__init__``, so flipping
-        this afterwards could not take effect.
+        Read-only: which limiter a request gets is decided in ``__init__``,
+        so flipping this afterwards could not take effect.
         """
-        return self._per_domain
+        return self._domain_limiters is not None
 
     def _get_limiter(self, request: ClientRequest) -> RateLimiter:
-        if self._global_limiter is not None:
-            return self._global_limiter
+        limiters = self._domain_limiters
+        if limiters is None:
+            return self._limiter
         # aiohttp raises InvalidUrlClientError for host-less URLs before
         # any middleware runs (on redirects too), so ``host`` is only
         # ``None`` in the type; the assert narrows it for mypy.
         domain = request.url.host
         assert domain is not None
-        return self._domain_limiters[domain]
+        limiter = limiters.get(domain)
+        if limiter is None:
+            # setdefault, not an assignment: threads racing for a host they
+            # have not seen before must all leave with the limiter that was
+            # stored, or each gets a private full budget and the burst
+            # allowance is briefly multiplied by the number of racers.
+            limiter = limiters.setdefault(domain, self._limiter.clone(domain))
+        return limiter
 
     async def __call__(
         self,

docs/api.rst

@@ -86,8 +86,10 @@ Rate limiting
    :class:`RateLimitMiddleware` accepts. Implementations provide async
    ``acquire()``, which reserves a slot and returns its delay as a non-negative
    finite number of seconds -- ``wait()`` takes that on trust -- and
-   ``clone()``, which returns a fresh limiter with the same configuration
-   (used once per host by per-domain mode)::
+   ``clone(host)``, which returns a fresh limiter with the same configuration
+   scoped to one host (called when per-domain mode first meets a host; a
+   racing thread's extra clone is discarded, so it should have no side
+   effects)::
 
        class RedisLimiter(RateLimiter):
            def __init__(self, redis, key, script):
@@ -100,14 +102,16 @@ Rate limiting
                ms = await self._redis.evalsha(self._script, 1, self._key)
                return ms / 1000
 
-           def clone(self):
-               return RedisLimiter(self._redis, self._key, self._script)
+           def clone(self, host):
+               return RedisLimiter(self._redis, f"{self._key}:{host}", self._script)
 
-   The sketch keeps one key across clones, and ``clone()`` takes no arguments
-   so it cannot learn the host; give each host its own limiter rather than
-   using ``per_domain=True``. It also leaves ``release()`` at the default
-   no-op, and a cancelled round trip can leave a reservation nobody holds;
-   an expiry on each reservation covers both.
+   Putting *host* in the key is what makes ``per_domain=True`` mean a budget
+   per host for a shared backend; a limiter that keeps its state in-process,
+   like :class:`TokenBucket`, has nothing to key and can ignore it. Redirects
+   pick hosts too, so give those keys an expiry of their own rather than let a
+   shared backend keep one for every host ever seen. The sketch also leaves
+   ``release()`` at the default no-op, and a cancelled round trip can leave a
+   reservation nobody holds; an expiry on each reservation covers both.
 
    ``wait(timeout=None)`` is supplied by the base class. It charges async
    acquisition against *timeout* once ``acquire()`` returns -- without bounding
@@ -150,7 +154,7 @@ Rate limiting
 
    :param RateLimiter limiter: The :class:`RateLimiter` to throttle with --
       for example ``TokenBucket(rate=5.0, burst=2)``. With ``per_domain=True``
-      it acts as a template: each target host gets ``limiter.clone()`` the
+      it acts as a template: each target host gets ``limiter.clone(host)`` the
       first time that host is seen.
    :param bool per_domain: Keep an independent limiter per target host instead
       of a single global one. Limiters are keyed on the URL host only (port

tests/test_rate_limit.py

@@ -1,6 +1,7 @@
 """Tests for the rate-limiting middleware."""
 
 import asyncio
+import threading
 import time
 from types import SimpleNamespace
 from unittest import mock
@@ -215,25 +216,111 @@ def test_per_domain_uses_distinct_limiters() -> None:
 
     assert limiter_a is not limiter_b  # distinct hosts -> isolated limiters
     assert limiter_a is limiter_a_again  # same host -> same limiter
+    assert middleware._domain_limiters is not None
     assert len(middleware._domain_limiters) == 2
 
 
+class _KeyedByHost(RateLimiter):
+    """A limiter whose state lives under a key, as a Redis-backed one would.
+
+    Every host it is cloned for is recorded on a list the clones share with
+    their template, so a test can see how often the middleware asks.
+    """
+
+    def __init__(
+        self, prefix: str, host: str = "", clones: list[str] | None = None
+    ) -> None:
+        self._prefix = prefix
+        self.key = f"{prefix}:{host}" if host else prefix
+        self.clones = [] if clones is None else clones
+
+    async def acquire(self) -> float:
+        return 0.0
+
+    def clone(self, host: str, /) -> "_KeyedByHost":
+        self.clones.append(host)
+        return _KeyedByHost(self._prefix, host, self.clones)
+
+
+def test_clone_is_given_the_host_it_is_built_for() -> None:
+    """A shared backend can only scope per host if clone() is told the host.
+
+    A host seen again gets the stored limiter back without another clone.
+    """
+    template = _KeyedByHost("rl")
+    middleware = RateLimitMiddleware(template, per_domain=True)
+
+    a = middleware._get_limiter(_fake_request("a.example"))
+    b = middleware._get_limiter(_fake_request("b.example"))
+    a_again = middleware._get_limiter(_fake_request("a.example"))
+
+    assert isinstance(a, _KeyedByHost) and isinstance(b, _KeyedByHost)
+    assert (a.key, b.key) == ("rl:a.example", "rl:b.example")
+    assert a_again is a
+    assert template.clones == ["a.example", "b.example"]
+
+
+class _BlocksInClone(RateLimiter):
+    """Limiter whose ``clone()`` holds each caller until they have all arrived."""
+
+    def __init__(self, barrier: threading.Barrier, host: str = "") -> None:
+        self._barrier = barrier
+        self.host = host
+
+    async def acquire(self) -> float:
+        return 0.0
+
+    def clone(self, host: str, /) -> "_BlocksInClone":
+        self._barrier.wait()
+        return _BlocksInClone(self._barrier, host)
+
+
+def test_first_contact_hands_every_racer_the_stored_limiter() -> None:
+    """Threads meeting a new host at once must all leave with the same limiter.
+
+    Each keeping the one it built would hand every racer a private, full
+    budget, multiplying the burst allowance by the number of racers for that
+    instant. ``dict.setdefault`` is what makes the winner the one everybody
+    gets; a plain assignment does not. The barrier sits inside ``clone()``,
+    so this also pins that the miss path stays lock-free: every racer builds
+    one and the extra clones are discarded.
+    """
+    racers = 4
+    barrier = threading.Barrier(racers, timeout=10)
+    middleware = RateLimitMiddleware(_BlocksInClone(barrier), per_domain=True)
+    handed_out: list[RateLimiter] = []
+
+    def race() -> None:
+        handed_out.append(middleware._get_limiter(_fake_request("new.example")))
+
+    threads = [threading.Thread(target=race) for _ in range(racers)]
+    for thread in threads:
+        thread.start()
+    for thread in threads:
+        thread.join(timeout=10)
+
+    assert len(handed_out) == racers
+    assert middleware._domain_limiters is not None
+    stored = middleware._domain_limiters["new.example"]
+    assert all(limiter is stored for limiter in handed_out)
+
+
 def test_global_mode_shares_one_limiter() -> None:
     """Without ``per_domain`` every host goes through the one limiter."""
     middleware = RateLimitMiddleware(TokenBucket(rate=10.0, burst=1))
 
     limiter_a = middleware._get_limiter(_fake_request("a.example"))
     limiter_b = middleware._get_limiter(_fake_request("b.example"))
 
-    assert limiter_a is limiter_b is middleware._global_limiter
+    assert limiter_a is limiter_b is middleware._limiter
 
 
 @pytest.mark.parametrize("per_domain", (False, True))
 def test_per_domain_is_readable_and_read_only(per_domain: bool) -> None:
     """``per_domain`` reports the configured mode and cannot be reassigned.
 
-    The limiters are built once in ``__init__``, so a writable attribute
-    would silently do nothing.
+    Which limiter a request gets is decided in ``__init__``, so a writable
+    attribute would silently do nothing.
     """
     middleware = RateLimitMiddleware(
         TokenBucket(rate=10.0, burst=1), per_domain=per_domain
@@ -276,7 +363,7 @@ async def test_middleware_bails_before_sleeping_when_timeout_known() -> None:
     async def handler(req: ClientRequest) -> ClientResponse:
         raise AssertionError("a doomed request must never be sent")
 
-    bucket = middleware._global_limiter
+    bucket = middleware._limiter
     assert isinstance(bucket, TokenBucket)
     assert await bucket.acquire() == 0.0  # drain the burst slot
 
@@ -296,7 +383,7 @@ async def test_middleware_cancel_during_sleep_releases_slot() -> None:
     async def handler(req: ClientRequest) -> ClientResponse:
         raise AssertionError("the cancelled request must never be sent")
 
-    bucket = middleware._global_limiter
+    bucket = middleware._limiter
     assert isinstance(bucket, TokenBucket)
     assert await bucket.acquire() == 0.0  # drain the burst slot
 
@@ -314,7 +401,7 @@ def test_limiter_injection_is_used_directly() -> None:
     """The caller-provided limiter is the one throttling, not a copy of it."""
     bucket = TokenBucket(rate=100.0, burst=1)
     middleware = RateLimitMiddleware(bucket)
-    assert middleware._global_limiter is bucket
+    assert middleware._limiter is bucket
 
 
 def test_non_limiter_rejected() -> None:
@@ -330,7 +417,7 @@ async def test_token_bucket_clone_is_fresh(clock: _FakeClock) -> None:
     await template.acquire()
     assert await template.acquire() > 0.0  # template drained into debt
 
-    fresh = template.clone()
+    fresh = template.clone("example.com")
     assert await fresh.acquire() == 0.0  # full burst again
     assert await fresh.acquire() == 0.0
     assert await fresh.acquire() == pytest.approx(0.1)  # same rate as the template
@@ -355,7 +442,7 @@ def __init__(self, delay: float) -> None:
     async def acquire(self) -> float:
         return self._delay
 
-    def clone(self) -> "_FixedDelay":
+    def clone(self, host: str, /) -> "_FixedDelay":
         return _FixedDelay(self._delay)
 
 
@@ -405,7 +492,7 @@ async def acquire(self) -> float:
         await asyncio.sleep(0)
         return 0.0
 
-    def clone(self) -> "_AwaitsToReserve":
+    def clone(self, host: str, /) -> "_AwaitsToReserve":
         return _AwaitsToReserve()
 
 
@@ -442,7 +529,7 @@ async def acquire(self) -> float:
     def release(self) -> None:
         self.releases += 1
 
-    def clone(self) -> "_ElapsedAcquire":
+    def clone(self, host: str, /) -> "_ElapsedAcquire":
         return _ElapsedAcquire(self._clock, self._spend, self._delay)
 
 
@@ -489,7 +576,7 @@ async def acquire(self) -> float:
     def release(self) -> None:
         self.releases += 1
 
-    def clone(self) -> "_CancelledAcquire":
+    def clone(self, host: str, /) -> "_CancelledAcquire":
         return _CancelledAcquire()
 
 
@@ -572,7 +659,7 @@ async def acquire(self) -> float:
     def release(self) -> None:
         self.releases += 1
 
-    def clone(self) -> "_RecordsReleases":
+    def clone(self, host: str, /) -> "_RecordsReleases":
         return _RecordsReleases(self._delay, self._error)