fix: custom semaphore

Author: Anton

taskiq/receiver/receiver.py

@@ -14,6 +14,7 @@
 from taskiq.message import TaskiqMessage
 from taskiq.receiver.params_parser import parse_params
 from taskiq.result import TaskiqResult
+from taskiq.semaphore import DequeSemaphore
 from taskiq.state import TaskiqState
 from taskiq.utils import DequeQueue, maybe_awaitable
 
@@ -58,15 +59,15 @@ def __init__(  # noqa: WPS211
             self.task_signatures[task.task_name] = inspect.signature(task.original_func)
             self.task_hints[task.task_name] = get_type_hints(task.original_func)
             self.dependency_graphs[task.task_name] = DependencyGraph(task.original_func)
-        self.sem: "Optional[asyncio.Semaphore]" = None
+        self.sem: "Optional[DequeSemaphore]" = None
         if max_async_tasks is not None and max_async_tasks > 0:
-            self.sem = asyncio.Semaphore(max_async_tasks)
+            self.sem = DequeSemaphore(max_async_tasks)
         else:
             logger.warning(
                 "Setting unlimited number of async tasks "
                 + "can result in undefined behavior",
             )
-        self.sem_prefetch = asyncio.Semaphore(max_prefetch)
+        self.sem_prefetch = DequeSemaphore(max_prefetch)
         self.queue: DequeQueue[bytes] = DequeQueue()
 
         self.sem_idle: Optional[asyncio.Semaphore] = None
@@ -309,7 +310,7 @@ def task_cb(task: "asyncio.Task[Any]") -> None:
                 break
             if message is QUEUE_SKIP:
                 # Decrease max_prefetch
-                prefetch_dec = asyncio.create_task(self.sem_prefetch.acquire())
+                prefetch_dec = asyncio.create_task(self.sem_prefetch.acquire_first())
                 prefetch_dec.add_done_callback(tasks.discard)
                 tasks.add(prefetch_dec)
 
@@ -356,5 +357,5 @@ 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()
+            await self.sem.acquire_first()
             await task

taskiq/semaphore.py

@@ -0,0 +1,99 @@
+import asyncio
+import collections
+from typing import Any, Deque
+
+from typing_extensions import Literal
+
+
+class DequeSemaphore:
+    """Custom deque based semaphore."""
+
+    def __init__(self, value: int) -> None:
+        self._value = value
+        self._waiters: Deque[asyncio.Future[Any]] = collections.deque()
+
+        if self._value < 0:
+            raise ValueError("Value should be >= 0")
+
+    async def __aenter__(self) -> Literal[True]:
+        return await self.acquire()
+
+    async def __aexit__(self, exc_type: Any, exc: Any, tb: Any) -> None:
+        self.release()
+
+    def locked(self) -> bool:
+        """
+        Returns True if semaphore cannot be acquired immediately.
+
+        :returns: true or false
+        """
+        return self._value == 0
+
+    def release(self) -> None:
+        """Release a semaphore, incrementing the internal counter by one.
+
+        When it was zero on entry and another coroutine is waiting for it to
+        become larger than zero again, wake up that coroutine.
+        """
+        self._value += 1
+        self._wakeup_next()
+
+    async def acquire(self, first: bool = False) -> Literal[True]:  # noqa: C901
+        """
+        Acquire a semaphore.
+
+        :param first: acquire ASAP
+        :raises asyncio.exceptions.CancelledError: task cancelled
+        :returns: true
+        """
+        if not self.locked() and not self._waiters:
+            # No need to wait as the semaphore is not locked
+            # and no one is waiting
+            self._value -= 1
+            return True
+
+        # if there are waiters or the semaphore is locked
+        fut: asyncio.Future[Any] = asyncio.Future()
+
+        if first:
+            self._waiters.appendleft(fut)
+        else:
+            self._waiters.append(fut)
+
+        try:
+            try:  # noqa: WPS501, WPS505
+                await fut
+            finally:
+                self._waiters.remove(fut)
+
+        except asyncio.exceptions.CancelledError:
+            if not fut.cancelled():
+                self._value += 1
+                self._wakeup_next()
+            raise
+
+        if not self.locked():
+            # This is required for strict FIFO ordering
+            # otherwise it can cause starvation on the waiting tasks
+            # The next loop iteration will wake up the task and switch
+            self._wakeup_next()
+
+        return True
+
+    async def acquire_first(self) -> Literal[True]:
+        """
+        Acquire a semaphore ASAP.
+
+        :returns: true
+        """
+        return await self.acquire(True)
+
+    def _wakeup_next(self) -> None:
+        if not self._waiters:
+            return
+
+        for fut in self._waiters:
+            if not fut.done():
+                self._value -= 1
+                fut.set_result(True)
+                return

tests/test_semaphore.py

@@ -0,0 +1,38 @@
+import asyncio
+
+import anyio
+import pytest
+
+from taskiq.semaphore import DequeSemaphore
+
+
+@pytest.mark.anyio
+async def test_semaphore_exception() -> None:
+    with pytest.raises(ValueError):
+        DequeSemaphore(-1)
+
+
+@pytest.mark.anyio
+async def test_semaphore() -> None:
+    sem = DequeSemaphore(1)
+
+    async def c1() -> None:
+        await sem.acquire()
+
+    async def c2() -> None:
+        await sem.acquire()
+
+    async def c3() -> None:
+        await sem.acquire()
+
+    t1 = asyncio.create_task(c1())
+    t2 = asyncio.create_task(c2())
+    t3 = asyncio.create_task(c3())
+    await asyncio.sleep(0)
+
+    sem.release()
+    sem.release()
+    t2.cancel()
+
+    async with anyio.maybe_async_cm(anyio.move_on_after(1)):
+        await asyncio.gather(t1, t2, t3, return_exceptions=True)