Faster acquire (#218)

schoennenbeck · mjpieters · web-flow · commit 0b53b86cad0e · 2024-11-30T21:03:30.000Z
* Faster acquire

* Addressed PR comments

* Mock out loop time to make test faster

* Add changelog entry

---------

Co-authored-by: Martijn Pieters &lt;mj@zopatista.com&gt;
diff --git a/changelog.d/217.bugfix.md b/changelog.d/217.bugfix.md
@@ -0,0 +1 @@
+Fixed wait time calculation for waiting tasks, making acquisition faster (PR by @schoennenbeck)
diff --git a/src/aiolimiter/leakybucket.py b/src/aiolimiter/leakybucket.py
@@ -102,9 +102,13 @@ async def acquire(self, amount: float = 1) -> None:
             fut = loop.create_future()
             self._waiters[task] = fut
             try:
-                await asyncio.wait_for(
-                    asyncio.shield(fut), 1 / self._rate_per_sec * amount
+                # we need to wait until current_capacity is equal or greater than the
+                # requested amount, where current_capacity is
+                # (self.max_rate - self._level)
+                wait_time = (
+                    1 / self._rate_per_sec * (amount - self.max_rate + self._level)
                 )
+                await asyncio.wait_for(asyncio.shield(fut), wait_time)
             except asyncio.TimeoutError:
                 pass
             fut.cancel()
diff --git a/tests/test_aiolimiter.py b/tests/test_aiolimiter.py
@@ -3,7 +3,6 @@
 # Licensed under the MIT license as detailed in LICENSE.txt
 
 import asyncio
-import time
 from pathlib import Path
 from unittest import mock
 
@@ -12,7 +11,8 @@
 
 from aiolimiter import AsyncLimiter
 
-WAIT_LIMIT = 2  # seconds before we declare the test failed
+# max number of wait_for rounds when waiting for all events to settle
+MAX_WAIT_FOR_ITER = 5
 
 
 def test_version():
@@ -32,14 +32,16 @@ def test_version():
 
 async def wait_for_n_done(tasks, n):
     """Wait for n (or more) tasks to have completed"""
-    start = time.time()
-    pending = tasks
+    iteration = 0
     remainder = len(tasks) - n
-    while time.time() <= start + WAIT_LIMIT and len(pending) > remainder:
+    while iteration <= MAX_WAIT_FOR_ITER:
+        iteration += 1
         _, pending = await asyncio.wait(
-            tasks, timeout=WAIT_LIMIT, return_when=asyncio.FIRST_COMPLETED
+            tasks, timeout=0, return_when=asyncio.FIRST_COMPLETED
         )
-    assert len(pending) >= remainder
+        if len(pending) <= remainder:
+            break
+    assert len(pending) <= remainder
     return pending
 
 
@@ -73,33 +75,61 @@ async def async_contextmanager_task(limiter):
         pass
 
 
-@pytest.mark.parametrize("task", [acquire_task, async_contextmanager_task])
-async def test_acquire(task):
-    current_time = 0
+class MockLoopTime:
+    def __init__(self):
+        self.current_time = 0
+        event_loop = asyncio.get_running_loop()
+        self.patch = mock.patch.object(event_loop, "time", self.mocked_time)
+
+    def mocked_time(self):
+        return self.current_time
 
-    def mocked_time():
-        return current_time
+    def __enter__(self):
+        self.patch.start()
+        return self
 
+    def __exit__(self, *_):
+        self.patch.stop()
+
+
+@pytest.mark.parametrize("task", [acquire_task, async_contextmanager_task])
+async def test_acquire(task):
     # capacity released every 2 seconds
     limiter = AsyncLimiter(5, 10)
 
-    event_loop = asyncio.get_running_loop()
-    with mock.patch.object(event_loop, "time", mocked_time):
+    with MockLoopTime() as mocked_time:
         tasks = [asyncio.ensure_future(task(limiter)) for _ in range(10)]
 
         pending = await wait_for_n_done(tasks, 5)
         assert len(pending) == 5
 
-        current_time = 3  # releases capacity for one and some buffer
+        mocked_time.current_time = 3  # releases capacity for one and some buffer
         assert limiter.has_capacity()
 
         pending = await wait_for_n_done(pending, 1)
         assert len(pending) == 4
 
-        current_time = 7  # releases capacity for two more, plus buffer
+        mocked_time.current_time = 7  # releases capacity for two more, plus buffer
         pending = await wait_for_n_done(pending, 2)
         assert len(pending) == 2
 
-        current_time = 11  # releases the remainder
+        mocked_time.current_time = 11  # releases the remainder
         pending = await wait_for_n_done(pending, 2)
         assert len(pending) == 0
+
+
+async def test_acquire_wait_time():
+    limiter = AsyncLimiter(3, 3)
+
+    with MockLoopTime() as mocked_time:
+        # Fill the bucket with an amount of 1
+        await limiter.acquire(1)
+
+        # Acquiring an amount of 3 now should take 1 second
+        task = asyncio.ensure_future(limiter.acquire(3))
+        pending = await wait_for_n_done([task], 0)
+        assert pending
+
+        mocked_time.current_time = 1
+        pending = await wait_for_n_done([task], 1)
+        assert not pending

Original file line number	Diff line number	Diff line change
`@@ -0,0 +1 @@`
	`1`	`+Fixed wait time calculation for waiting tasks, making acquisition faster (PR by @schoennenbeck)`