Make RateLimiter.acquire() async (#25)

Author: rodrigobnogueira

aiohttp_client_middlewares/rate_limit.py

@@ -21,31 +21,28 @@
 class RateLimiter(ABC):
     """Base class for rate-limit algorithms.
 
-    Implementations provide the synchronous :meth:`acquire` and
-    :meth:`clone`; the async sleeping and timeout logic live here in
-    :meth:`wait`, shared by every algorithm. Because :meth:`acquire` is
-    synchronous, callers on one event loop reserve slots atomically in
-    arrival order.
-
-    A limiter that needs I/O to reserve a slot -- one backed by Redis or a
-    database, say -- overrides :meth:`wait` rather than :meth:`acquire`:
-    ``wait`` is the only method the middleware calls and is already a
-    coroutine. Such an implementation owns what it takes over: ordering
-    between concurrent callers, charging the round trip against *timeout*,
-    and handing the slot back when the caller goes away.
-
-    :meth:`acquire` and :meth:`clone` stay abstract either way, so an
-    implementation that overrides :meth:`wait` still has to define both to
-    be instantiable. Its :meth:`acquire` is never called and can simply
-    raise; :meth:`clone` is called for real by ``per_domain=True``.
+    Implementations provide an async :meth:`acquire` and a synchronous
+    :meth:`clone`. The sleeping, timeout and post-reservation cancellation
+    logic lives in :meth:`wait`, shared by every algorithm, so reserving a
+    slot may perform I/O of its own -- against Redis or a database, say.
+
+    Until :meth:`acquire` returns, cleaning up a half-made reservation is
+    its own responsibility; once it returns, :meth:`wait` owns the slot and
+    calls :meth:`release` if it cannot be used.
+
+    An async method that contains no suspension point still runs atomically
+    when awaited directly. :class:`TokenBucket` relies on that property to
+    preserve arrival ordering on one event loop.
     """
 
     @abstractmethod
-    def acquire(self) -> float:
+    async def acquire(self) -> float:
         """Reserve a slot and return the delay to sleep before sending.
 
-        Must return without awaiting: the arrival-order guarantee above
-        holds precisely because there is no suspension point here.
+        The delay must be non-negative, finite seconds; :meth:`wait` takes that
+        on trust, and a NaN would send the request through unthrottled. If
+        cancellation or another exception prevents this method from
+        returning, it must not leave a reservation behind.
         """
 
     @abstractmethod
@@ -64,28 +61,35 @@ def release(self) -> None:
         cancelled while sleeping. The default is a no-op for algorithms
         that have nothing to return.
 
-        Runs from an ``except asyncio.CancelledError`` block, so it must
-        not await either: a second cancellation, or the loop shutting
-        down, would truncate it part-way and lose the slot for good.
+        Must neither await nor raise, since one of those calls is from an
+        ``except asyncio.CancelledError`` block: awaiting there can be
+        truncated part-way, and raising would replace the exception the
+        caller is owed. A limiter that has to reach its backend to hand a
+        slot back can schedule that round trip as a task from here.
         """
 
     async def wait(self, timeout: float | None = None) -> None:
-        """Reserve a slot and sleep out its delay.
+        """Reserve a slot and wait until the request may be sent.
 
-        When the delay would exceed *timeout*, the slot is handed back and
-        :exc:`asyncio.TimeoutError` is raised without sleeping, so a
-        request that could never be sent in time fails fast.
-
-        This is the method the middleware calls, and the one to override
-        when reserving a slot needs I/O of its own.
+        Time in :meth:`acquire` is charged against *timeout* once it
+        returns, though not bounded by it, so an implementation that can
+        hang needs its own deadline. When the delay exceeds what is left,
+        the slot is handed back and :exc:`asyncio.TimeoutError` raised
+        without sleeping.
         """
-        delay = self.acquire()
-        if timeout is not None and delay > timeout:
-            self.release()
-            raise asyncio.TimeoutError(
-                f"rate limiter would delay the request {delay:.3f}s, "
-                f"beyond the {timeout:.3f}s timeout"
-            )
+        started = time.monotonic()
+        delay = await self.acquire()
+
+        if timeout is not None:
+            # Goes negative when acquiring alone outlasted the timeout; the
+            # message reports it as such rather than clamping it to zero.
+            remaining = timeout - (time.monotonic() - started)
+            if delay > remaining:
+                self.release()
+                raise asyncio.TimeoutError(
+                    f"rate limiter would delay the request {delay:.3f}s, "
+                    f"beyond the {remaining:.3f}s remaining timeout"
+                )
         if delay > 0.0:
             try:
                 await asyncio.sleep(delay)
@@ -131,11 +135,13 @@ def _refill(self) -> None:
         )
         self._last_refill = now
 
-    def acquire(self) -> float:
+    async def acquire(self) -> float:
         """Take one token and return the delay to sleep before sending.
 
         The delay is the exact fractional deficit (not rounded to whole
         intervals), so a caller never waits longer than the bucket needs.
+        There is deliberately no suspension point: callers on one event
+        loop reserve slots atomically, in arrival order.
         """
         self._refill()
         self._tokens -= 1.0
@@ -155,14 +161,16 @@ class RateLimitMiddleware:
     """Client middleware that throttles requests through a :class:`RateLimiter`.
 
     The middleware waits on the limiter before sending, so the client never
-    sends faster than the limiter allows and slots are granted in arrival
-    order. Cancellation is the one exception: a slot handed back by
-    :meth:`RateLimiter.release` frees capacity that queued callers have
-    already been given fixed delays against, so two of them can briefly
-    send in the same instant. When aiohttp exposes the request's timeout
-    (aiohttp 3.15 and newer), a wait that would exceed it fails immediately
-    with :exc:`asyncio.TimeoutError` instead of sleeping toward a guaranteed
-    timeout.
+    sends faster than the limiter allows. What that ordering is worth is the
+    limiter's to say. :class:`TokenBucket` grants slots in arrival order
+    because its :meth:`~RateLimiter.acquire` has no suspension point.
+
+    For :class:`TokenBucket`, cancellation is the one exception to arrival
+    order: a handed-back slot frees capacity that queued callers have already
+    been given fixed delays against, so two of them can briefly send in the
+    same instant. When aiohttp exposes the request's timeout (aiohttp 3.15
+    and newer), a wait that would exceed it fails immediately with
+    :exc:`asyncio.TimeoutError` instead of sleeping toward a guaranteed timeout.
 
     Middleware order matters: middlewares listed earlier wrap the ones listed
     later, and a middleware that retries internally (for example,
@@ -231,5 +239,10 @@ async def __call__(
         # the getattr() is gated on raising the floor to 3.15; that change
         # is ready in #23 and waits only on the aiohttp release.
         client_timeout: ClientTimeout | None = getattr(request, "timeout", None)
-        await limiter.wait(None if client_timeout is None else client_timeout.total)
+        total = None if client_timeout is None else client_timeout.total
+        if total is not None and total <= 0.0:
+            # aiohttp arms its own deadline only for a positive total, so a
+            # zero or negative one means "no timeout", not "no budget left".
+            total = None
+        await limiter.wait(total)
         return await handler(request)

docs/api.rst

@@ -82,37 +82,61 @@ Rate limiting
 
 .. class:: RateLimiter()
 
-   Abstract base class for rate-limit algorithms. Implementations provide the
-   synchronous ``acquire()``, which reserves a slot and returns the delay, in
-   seconds, to sleep before sending, and ``clone()``, which returns a fresh
-   limiter with the same configuration (used once per host by per-domain
-   mode). The async ``wait(timeout=None)`` method is shared by all
-   implementations: it acquires a slot, fails fast with
-   :exc:`asyncio.TimeoutError` -- handing the slot back -- when the delay
-   would exceed *timeout*, sleeps out the delay otherwise, and hands the slot
-   back if the caller is cancelled mid-sleep. ``release()`` defaults to a
-   no-op for algorithms that have nothing to return.
-
-   Neither ``acquire()`` nor ``release()`` may await. Their being synchronous
-   is what reserves slots atomically, in arrival order, across concurrent
-   callers on one event loop, and ``release()`` additionally runs from a
-   cancellation handler, where an awaiting implementation can be truncated
-   and lose the slot. A limiter that needs I/O to reserve a slot -- one
-   backed by Redis or a database, say -- overrides ``wait()`` instead: it is
-   the only method the middleware calls and is already a coroutine. Such an
-   implementation takes on ordering between concurrent callers, charging its
-   round trip against *timeout*, and returning the slot when the caller goes
-   away. ``acquire()`` and ``clone()`` are abstract either way, so it still
-   has to define both to be instantiable; its ``acquire()`` is never called
-   and can simply raise, while ``clone()`` is used for real by
-   ``per_domain=True``.
+   Abstract base class for rate-limit algorithms, and the type
+   :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)::
+
+       class RedisLimiter(RateLimiter):
+           def __init__(self, redis, key, script):
+               self._redis, self._key, self._script = redis, key, script
+
+           async def acquire(self):
+               # An atomic script reserves the next slot and returns the delay
+               # in milliseconds: Redis turns Lua numbers into integers, so a
+               # fractional second cannot come back as one.
+               ms = await self._redis.evalsha(self._script, 1, self._key)
+               return ms / 1000
+
+           def clone(self):
+               return RedisLimiter(self._redis, self._key, 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.
+
+   ``wait(timeout=None)`` is supplied by the base class. It charges async
+   acquisition against *timeout* once ``acquire()`` returns -- without bounding
+   the call itself, so an implementation that can hang needs its own deadline --
+   then fails fast when the delay exceeds what is left, sleeps otherwise, and
+   calls ``release()`` if an acquired slot cannot be used. ``release()`` defaults
+   to a no-op for algorithms with nothing to return, and stays synchronous: one
+   of those calls is from a cancellation handler, where awaiting can be truncated
+   part-way and lose the slot for good, and raising would replace the exception
+   the caller is owed. A limiter that has to reach its backend to hand a slot
+   back can schedule that round trip as a task.
+
+   ``acquire()`` must be cancellation-safe: if cancellation or another
+   exception prevents it from returning, it must leave no reservation behind.
+   Once it returns successfully, ``wait()`` owns that cleanup. A backend whose
+   reservation can outlive a cancelled network operation should use an
+   idempotency key, transaction, or expiry so interrupted acquisition cannot
+   leak capacity.
+
+   An async ``acquire()`` with no suspension point still reserves atomically
+   across callers on one event loop. An implementation that performs I/O
+   determines its own ordering at those suspension points.
 
 .. class:: TokenBucket(rate=10.0, burst=10)
 
    A :class:`RateLimiter`: tokens accrue continuously at ``rate`` per second,
-   capped at ``burst``; ``acquire()`` takes one token and the count may go
-   negative, which is what queues callers up in arrival order. The bucket
-   holds no tasks or loop state.
+   capped at ``burst``; async ``acquire()`` takes one token and the count may
+   go negative, which is what queues callers up in arrival order. It contains
+   no suspension point and the bucket holds no tasks or loop state.
 
    :param float rate: Token accrual rate, in tokens per second. Must be a
       positive, finite number.
@@ -137,10 +161,13 @@ Rate limiting
    :raises TypeError: if ``limiter`` is not a :class:`RateLimiter`.
 
    The middleware waits on the limiter before sending, so the client never
-   sends faster than the limiter allows and slots are granted in arrival
-   order. Cancellation is the one exception: a slot handed back by
-   ``release()`` frees capacity that queued callers already hold fixed delays
-   against, so two of them can briefly send in the same instant. When aiohttp
+   sends faster than the limiter allows. What that ordering is worth is the
+   limiter's to say: :class:`TokenBucket` grants slots in arrival order because
+   its async ``acquire()`` does not suspend, while an I/O-backed limiter orders
+   callers according to its backend. For :class:`TokenBucket`, cancellation is
+   the one exception: a slot handed back by ``release()`` frees capacity that
+   queued callers already hold fixed delays against, so two of them can
+   briefly send in the same instant. When aiohttp
    exposes the request's total timeout to the middleware
    (aiohttp 3.15 and newer), a wait that would exceed it fails immediately
    with :exc:`asyncio.TimeoutError` instead of sleeping toward a guaranteed