fix: more tests

Author: Anton

taskiq/receiver/receiver.py

@@ -357,5 +357,6 @@ 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

taskiq/semaphore.py

@@ -15,6 +15,7 @@ class DequeSemaphore:
     def __init__(self, value: int) -> None:
         self._value = value
         self._waiters: Deque[asyncio.Future[Any]] = collections.deque()
+        self._waiters_first: Deque[asyncio.Future[Any]] = collections.deque()
 
         if self._value < 0:
             raise ValueError("Value should be >= 0")
@@ -31,8 +32,10 @@ def locked(self) -> bool:
 
         :returns: true or false
         """
-        return self._value == 0 or (
-            any(not waiter.cancelled() for waiter in (self._waiters or ()))
+        return (
+            self._value == 0
+            or any(not waiter.cancelled() for waiter in (self._waiters or ()))
+            or any(not waiter.cancelled() for waiter in (self._waiters_first or ()))
         )
 
     def release(self) -> None:
@@ -62,15 +65,18 @@ async def acquire(self, first: bool = False) -> Literal[True]:  # noqa: C901
         fut: asyncio.Future[Any] = asyncio.Future()
 
         if first:
-            self._waiters.appendleft(fut)
+            self._waiters_first.append(fut)
         else:
             self._waiters.append(fut)
 
         try:
             try:  # noqa: WPS501, WPS505
                 await fut
             finally:
-                self._waiters.remove(fut)
+                if first:
+                    self._waiters_first.remove(fut)
+                else:
+                    self._waiters.remove(fut)
 
         except asyncio.exceptions.CancelledError:
             if not fut.cancelled():
@@ -95,9 +101,15 @@ async def acquire_first(self) -> Literal[True]:
         return await self.acquire(True)
 
     def _wakeup_next(self) -> None:
-        if not self._waiters:
+        if not self._waiters and not self._waiters_first:
             return
 
+        for fut in self._waiters_first:
+            if not fut.done():
+                self._value -= 1
+                fut.set_result(True)
+                return
+
         for fut in self._waiters:
             if not fut.done():
                 self._value -= 1

tests/cli/worker/test_receiver.py

@@ -14,6 +14,7 @@
 from taskiq.message import TaskiqMessage
 from taskiq.receiver import Receiver
 from taskiq.result import TaskiqResult
+from taskiq.task import AsyncTaskiqTask
 
 _T = TypeVar("_T")
 
@@ -346,3 +347,44 @@ async def task_map(vals: List[int], ctx: Context = Depends()) -> List[int]:
 
     await broker.shutdown()
     await listen_task
+
+
+@pytest.mark.anyio
+async def test_tasks_chain_deep() -> None:
+    """"""
+    broker = InMemoryQueueBroker()
+
+    @broker.task
+    async def task_run(depth: int, val: Any, ctx: Context = Depends()) -> Any:
+        if depth == 0:
+            return val
+
+        t = await task_run.kiq(depth - 1, val)
+        resp = await wait_for_task(t, interval=0.05, ctx=ctx)
+        return resp.return_value
+
+    async def wait_for_task(
+        task: AsyncTaskiqTask[Any],
+        interval: float,
+        ctx: Context,
+    ) -> TaskiqResult[Any]:
+        while True:
+            resp_task = asyncio.create_task(
+                task.wait_result(interval * 0.4, timeout=interval),
+            )
+            await ctx.task_idler(interval)
+
+            try:
+                return await resp_task
+            except TaskiqResultTimeoutError:
+                continue
+
+    receiver = get_receiver(broker, max_async_tasks=1, max_idle_tasks=10)
+    listen_task = asyncio.create_task(receiver.listen())
+
+    task = await task_run.kiq(10, "hello world!")
+    resp = await task.wait_result(timeout=1)
+    assert resp.return_value == "hello world!"
+
+    await broker.shutdown()
+    await listen_task