fix: semaphore for different python versions

Author: Anton

taskiq/receiver/receiver.py

@@ -14,8 +14,9 @@
 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, DequeSemaphore, maybe_awaitable
+from taskiq.utils import DequeQueue, maybe_awaitable
 
 logger = getLogger(__name__)
 QUEUE_DONE = b"-1"

taskiq/semaphore.py

@@ -0,0 +1,65 @@
+import asyncio
+import sys
+from typing import TYPE_CHECKING, Any, Deque, Literal
+
+
+class DequeSemaphore(asyncio.Semaphore):
+    """Deque based Semaphore."""
+
+    if TYPE_CHECKING:  # noqa: WPS604
+        _loop: asyncio.BaseEventLoop
+        _waiters: Deque[Any]
+        _wakeup_scheduled: bool
+        _get_loop: Any
+
+    if sys.version_info[1] < 10:  # noqa: C901, WPS604
+
+        async def acquire_first(self) -> Literal[True]:
+            """
+            Acquire as soon as possible. LIFO style.
+
+            :param self: DequeSemaphore
+            :raises BaseException: exception
+            :return: true
+            """
+            while self._value < 0:
+                fut = self._loop.create_future()
+                self._waiters.appendleft(fut)
+
+                try:
+                    await fut
+                except BaseException:  # noqa: WPS424
+                    self._value += 1
+
+                    fut.cancel()
+                    if not self.locked() and not fut.cancelled():
+                        self._wake_up_next()
+                    raise
+
+            return True
+
+    else:
+
+        async def acquire_first(self) -> Literal[True]:
+            """
+            Acquire as soon as possible. LIFO style.
+
+            :param self: DequeSemaphore
+            :raises asyncio.exceptions.CancelledError: exception
+            :return: true
+            """
+            # _wakeup_scheduled is set if *another* task is scheduled to wakeup
+            # but its acquire() is not resumed yet
+            while self._wakeup_scheduled or self._value <= 0:
+                fut = self._get_loop().create_future()
+                self._waiters.appendleft(fut)
+                try:
+                    await fut
+                    # reset _wakeup_scheduled *after* waiting for a future
+                    self._wakeup_scheduled = False
+                except asyncio.exceptions.CancelledError:
+                    self._wake_up_next()
+                    raise
+
+            self._value -= 1
+            return True

taskiq/utils.py

@@ -2,8 +2,6 @@
 import inspect
 from typing import TYPE_CHECKING, Any, Coroutine, Deque, Generic, TypeVar, Union
 
-from typing_extensions import Literal
-
 _T = TypeVar("_T")  # noqa: WPS111
 
 
@@ -40,37 +38,6 @@ def remove_suffix(text: str, suffix: str) -> str:
     return text
 
 
-class DequeSemaphore(asyncio.Semaphore):
-    """Deque based Semaphore."""
-
-    if TYPE_CHECKING:  # noqa: WPS604
-        _loop: asyncio.BaseEventLoop
-
-    async def acquire_first(self) -> Literal[True]:
-        """
-        Acquire as soon as possible. LIFO style.
-
-        :raises BaseException: exception
-        :return: true
-        """
-        self._value -= 1
-
-        while self._value < 0:
-            fut: asyncio.Future[Any] = self._loop.create_future()
-            self._waiters.appendleft(fut)
-            try:
-                await fut
-            except BaseException:  # noqa: WPS424
-                self._value += 1
-
-                fut.cancel()
-                if not self.locked() and not fut.cancelled():
-                    self._wake_up_next()
-                raise
-
-        return True
-
-
 class DequeQueue(
     asyncio.Queue,  # type: ignore
     Generic[_T],