fix: more tests, better lock

Author: Anton

taskiq/receiver/receiver.py

@@ -330,7 +330,7 @@ def task_cb(task: "asyncio.Task[Any]") -> None:
             # https://textual.textualize.io/blog/2023/02/11/the-heisenbug-lurking-in-your-async-code/
             task.add_done_callback(task_cb)
 
-    async def task_idler(self, wait: float) -> None:
+    async def task_idler(self, wait: float) -> None:  # noqa: WPS213, WPS217
         """
         Temporary increasing `max_async_tasks` for at least `wait` amount of time.
 
@@ -346,7 +346,13 @@ async def task_idler(self, wait: float) -> None:
             return
 
         start_time = time()
-        async with self.sem_idle:
+        with anyio.move_on_after(wait) as scope:
+            await self.sem_idle.acquire()
+
+        if scope.cancel_called:  # noqa: WPS441
+            return
+
+        try:  # noqa: WPS501
             # Increase max_tasks
             # Increase max_prefetch in runner
             self.sem.release()
@@ -357,6 +363,8 @@ async def task_idler(self, wait: float) -> None:
             # Decrease max_prefetch in runner
             task = asyncio.create_task(self.queue.put_first(QUEUE_SKIP))
             # Decrease max_tasks
-
             await self.sem.acquire_first()
             await task
+
+        finally:
+            self.sem_idle.release()

tests/cli/worker/test_receiver.py

@@ -2,6 +2,7 @@
 from concurrent.futures import ThreadPoolExecutor
 from typing import Any, AsyncGenerator, List, Optional, TypeVar
 
+import anyio
 import pytest
 from taskiq_dependencies import Depends
 
@@ -388,3 +389,30 @@ async def wait_for_task(
 
     await broker.shutdown()
     await listen_task
+
+
+@pytest.mark.anyio
+async def test_tasks_sleep() -> None:
+    """"""
+    broker = InMemoryQueueBroker()
+
+    @broker.task
+    async def task_run(ind: int, ctx: Context = Depends()) -> int:
+        await ctx.task_idler(0.1)
+        return ind
+
+    receiver = get_receiver(broker, max_async_tasks=1, max_idle_tasks=20)
+    listen_task = asyncio.create_task(receiver.listen())
+
+    with anyio.fail_after(1):
+        tasks_tasks = [asyncio.create_task(task_run.kiq(ind)) for ind in range(100)]
+        tasks = await asyncio.gather(*tasks_tasks)
+        resps_tasks = [
+            asyncio.create_task(task.wait_result(timeout=1)) for task in tasks
+        ]
+        resps = await asyncio.gather(*resps_tasks)
+        value = [resp.return_value for resp in resps]
+        assert value == list(range(100))
+
+    await broker.shutdown()
+    await listen_task