Keep Retrying until Max Delay Has Expired (#200)

* feat: implement retry_until_max_delay for limiter, to retry until the max delay has expired

* reduce log verbosity

* Update pyrate_limiter/limiter.py

Add docstring

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

* move delay increment outside loop and conditionally apply the raise-delay-exception

* condition to apply the right timeout

* fix fetch depth

* fetch-depth: was pulling wrong commit

* undo

* Merge and make typing happy again

* fix the delay: needs to be after we get waiting time

* Update pyrate_limiter/limiter.py

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

---------

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
Co-authored-by: Vu Tran <me@vutr.io>

Author: 3 people

benchmarks/stress_limiters.py

@@ -10,11 +10,13 @@
 from typing import Literal
 from typing import Optional
 
+from pyrate_limiter import AbstractClock
 from pyrate_limiter import Duration
 from pyrate_limiter import Limiter
 from pyrate_limiter import Rate
 from pyrate_limiter import SQLiteBucket
 from pyrate_limiter import SQLiteClock
+from pyrate_limiter import TimeClock
 
 logger = logging.getLogger(__name__)
 
@@ -34,11 +36,11 @@ def create_sqlite_limiter(rate: Rate, use_fileLock: bool, max_delay: int):
 
     # retry_until_max_delay=True
     if use_fileLock:
-        kwargs = dict(clock=SQLiteClock(bucket))
+        clock: AbstractClock = SQLiteClock(bucket)
     else:
-        kwargs = {}
+        clock = TimeClock()
 
-    return Limiter(bucket, raise_when_fail=False, max_delay=max_delay, **kwargs)
+    return Limiter(bucket, raise_when_fail=False, max_delay=max_delay, clock=clock, retry_until_max_delay=True)
 
 
 def create_rate_limiter_factory(

pyrate_limiter/limiter.py

@@ -64,6 +64,7 @@ class Limiter:
 
     bucket_factory: BucketFactory
     raise_when_fail: bool
+    retry_until_max_delay: bool
     max_delay: Optional[int] = None
     lock: RLock
 
@@ -73,13 +74,23 @@ def __init__(
         clock: AbstractClock = TimeClock(),
         raise_when_fail: bool = True,
         max_delay: Optional[Union[int, Duration]] = None,
+        retry_until_max_delay: bool = False
     ):
         """Init Limiter using either a single bucket / multiple-bucket factory
-        / single rate / rate list
+        / single rate / rate list.
+
+        Parameters:
+            argument (Union[BucketFactory, AbstractBucket, Rate, List[Rate]]): The bucket or rate configuration.
+            clock (AbstractClock, optional): The clock instance to use for rate limiting. Defaults to TimeClock().
+            raise_when_fail (bool, optional): Whether to raise an exception when rate limiting fails. Defaults to True.
+            max_delay (Optional[Union[int, Duration]], optional): The maximum delay allowed for rate limiting.
+            Defaults to None.
+            retry_until_max_delay (bool, optional): If True, retry operations until the maximum delay is reached.
+                Useful for ensuring operations eventually succeed within the allowed delay window. Defaults to False.
         """
         self.bucket_factory = self._init_bucket_factory(argument, clock=clock)
         self.raise_when_fail = raise_when_fail
-
+        self.retry_until_max_delay = retry_until_max_delay
         if max_delay is not None:
             if isinstance(max_delay, Duration):
                 max_delay = int(max_delay)
@@ -179,25 +190,40 @@ async def _handle_async():
                 nonlocal delay
                 delay = await delay
                 assert isinstance(delay, int), "Delay not integer"
-                delay += 50
-
-                if delay > self.max_delay:
-                    logger.error(
-                        "Required delay too large: actual=%s, expected=%s",
-                        delay,
-                        self.max_delay,
-                    )
-                    self._raise_delay_exception_if_necessary(bucket, item, delay)
-                    return False
 
-                await asyncio.sleep(delay / 1000)
-                item.timestamp += delay
-                re_acquire = bucket.put(item)
+                total_delay = 0
+                delay += 50
 
-                if isawaitable(re_acquire):
-                    re_acquire = await re_acquire
-
-                return _handle_reacquire(re_acquire)
+                while True:
+                    total_delay += delay
+
+                    if self.retry_until_max_delay:
+                        if self.max_delay is not None and total_delay > self.max_delay:
+                            logger.error("Total delay exceeded max_delay: total_delay=%s, max_delay=%s",
+                                         total_delay, self.max_delay)
+                            self._raise_delay_exception_if_necessary(bucket, item, total_delay)
+                            return False
+                    else:
+                        if self.max_delay is not None and delay > self.max_delay:
+                            logger.error(
+                                "Required delay too large: actual=%s, expected=%s",
+                                delay,
+                                self.max_delay,
+                            )
+                            self._raise_delay_exception_if_necessary(bucket, item, delay)
+                            return False
+
+                    await asyncio.sleep(delay / 1000)
+                    item.timestamp += delay
+                    re_acquire = bucket.put(item)
+
+                    if isawaitable(re_acquire):
+                        re_acquire = await re_acquire
+
+                    if not self.retry_until_max_delay:
+                        return _handle_reacquire(re_acquire)
+                    elif re_acquire:
+                        return True
 
             return _handle_async()
 
@@ -213,23 +239,40 @@ async def _handle_async():
             self._raise_bucket_full_if_necessary(bucket, item)
             return False
 
-        delay += 50
+        total_delay = 0
 
-        if delay > self.max_delay:
-            logger.error(
-                "Required delay too large: actual=%s, expected=%s",
-                delay,
-                self.max_delay,
-            )
-            self._raise_delay_exception_if_necessary(bucket, item, delay)
-            return False
+        while True:
+            logger.debug("delay=%d, total_delay=%s", delay, total_delay)
+            delay = bucket.waiting(item)
+            assert isinstance(delay, int)
+
+            delay += 50
+            total_delay += delay
+
+            if self.max_delay is not None and total_delay > self.max_delay:
+                logger.error(
+                    "Required delay too large: actual=%s, expected=%s",
+                    delay,
+                    self.max_delay,
+                )
 
-        sleep(delay / 1000)
-        item.timestamp += delay
-        re_acquire = bucket.put(item)
-        # NOTE: if delay is not Awaitable, then `bucket.put` is not Awaitable
-        assert isinstance(re_acquire, bool)
-        return _handle_reacquire(re_acquire)
+                if self.retry_until_max_delay:
+                    self._raise_delay_exception_if_necessary(bucket, item, total_delay)
+                else:
+                    self._raise_delay_exception_if_necessary(bucket, item, delay)
+
+                return False
+
+            sleep(delay / 1000)
+            item.timestamp += delay
+            re_acquire = bucket.put(item)
+            # NOTE: if delay is not Awaitable, then `bucket.put` is not Awaitable
+            assert isinstance(re_acquire, bool)
+
+            if not self.retry_until_max_delay:
+                return _handle_reacquire(re_acquire)
+            elif re_acquire:
+                return True
 
     def handle_bucket_put(
         self,