refactor: better queue

Author: Anton

taskiq/receiver/receiver.py

@@ -17,7 +17,7 @@
 from taskiq.result import TaskiqResult
 from taskiq.semaphore import PrioritySemaphore
 from taskiq.state import TaskiqState
-from taskiq.utils import DequeQueue, maybe_awaitable
+from taskiq.utils import PriorityQueue, maybe_awaitable
 
 logger = getLogger(__name__)
 QUEUE_DONE = b"-1"
@@ -71,7 +71,7 @@ def __init__(  # noqa: WPS211
                 + "can result in undefined behavior",
             )
         self.sem_prefetch = PrioritySemaphore(max_prefetch)
-        self.queue: DequeQueue[bytes] = DequeQueue()
+        self.queue: PriorityQueue[bytes] = PriorityQueue()
 
         self.sem_sleeping: Optional[asyncio.Semaphore] = None
         if max_sleeping_tasks is not None and max_sleeping_tasks <= 0:
@@ -270,7 +270,7 @@ async def listen(self) -> None:  # pragma: no cover
             gr.start_soon(self.prefetcher, self.queue)
             gr.start_soon(self.runner, self.queue)
 
-    async def prefetcher(self, queue: "asyncio.Queue[Any]") -> None:
+    async def prefetcher(self, queue: "PriorityQueue[Any]") -> None:
         """
         Prefetch tasks data.
 
@@ -282,14 +282,14 @@ async def prefetcher(self, queue: "asyncio.Queue[Any]") -> None:
             try:
                 await self.sem_prefetch.acquire()
                 message = await iterator.__anext__()  # noqa: WPS609
-                await queue.put(message)
+                await queue.put_last(message)
 
             except StopAsyncIteration:
                 break
 
-        await queue.put(QUEUE_DONE)
+        await queue.put_last(QUEUE_DONE)
 
-    async def runner(self, queue: "asyncio.Queue[bytes]") -> None:  # noqa: C901, WPS213
+    async def runner(self, queue: "PriorityQueue[bytes]") -> None:  # noqa: C901, WPS213
         """
         Run tasks.
 
@@ -317,7 +317,7 @@ def task_cb(task: "asyncio.Task[Any]") -> None:
                 await self.sem.acquire()
 
             self.sem_prefetch.release()
-            message = await queue.get()
+            _, _, message = await queue.get()
             if message is QUEUE_DONE:
                 if self.sem is not None:
                     self.sem.release()

taskiq/utils.py

@@ -1,6 +1,6 @@
 import asyncio
 import inspect
-from typing import TYPE_CHECKING, Any, Coroutine, Deque, Generic, TypeVar, Union
+from typing import TYPE_CHECKING, Any, Coroutine, Generic, Tuple, TypeVar, Union
 
 _T = TypeVar("_T")  # noqa: WPS111
 
@@ -38,70 +38,52 @@ def remove_suffix(text: str, suffix: str) -> str:
     return text
 
 
-class DequeQueue(
-    asyncio.Queue,  # type: ignore
-    Generic[_T],
-):
-    """Deque based Queue."""
-
-    if TYPE_CHECKING:  # noqa: WPS604
-        _loop: asyncio.BaseEventLoop
-        _queue: Deque[_T]
-        _putters: Deque[Any]
-        _getters: Deque[Any]
-        _unfinished_tasks: int
-        _finished: asyncio.Event
-        _wakeup_next: Any
+class PriorityQueue(asyncio.PriorityQueue, Generic[_T]):  # type: ignore
+    """PriorityQueue based Queue."""
 
     async def put_first(self, item: _T) -> None:
         """
-        Wait till queue is not full. Put item in Queue as soon as possible. LIFO style.
+        Put item in Queue with highest priority.
 
         :param item: value to prepend
-        :raises BaseException: something goes wrong
-        :returns: nothing
         """
-        while self.full():
-            putter = self._loop.create_future()
-            self._putters.appendleft(putter)
-            try:
-                await putter
-            except BaseException:  # noqa: WPS424
-                putter.cancel()  # Just in case putter is not done yet.
-                try:  # noqa: WPS505
-                    # Clean self._putters from canceled putters.
-                    self._putters.remove(putter)
-                except ValueError:
-                    # The putter could be removed from self._putters by a
-                    # previous get_nowait call.
-                    pass  # noqa: WPS420
-                if not self.full() and not putter.cancelled():
-                    # We were woken up by get_nowait(), but can't take
-                    # the call.  Wake up the next in line.
-                    self._wakeup_next(self._putters)
-                raise
-
-        return self.put_first_nowait(item)
-
-    def put_first_nowait(self, item: _T) -> None:
+        self.counter += 1
+        await self.put((0, self.counter, item))
+
+    async def put_last(self, item: _T) -> None:
         """
-        Put item in Queue as soon as possible. LIFO style.
+        Put item in Queue with lowest priority.
 
-        :param item: value to prepend
-        :raises QueueFull: queue is full
+        :param item: value to append
         """
-        if self.full():
-            raise asyncio.QueueFull()
+        self.counter += 1
+        await self.put((1, self.counter, item))
 
-        self._put_first(item)
-        self._unfinished_tasks += 1
-        self._finished.clear()
-        self._wakeup_next(self._getters)
+    def _init(self, maxsize: int) -> None:
+        super()._init(maxsize)
+        self.counter = 0
 
-    def _put_first(self, item: _T) -> None:
-        """
-        Prepend item.
 
-        :param item: value to prepend
-        """
-        self._queue.appendleft(item)
+if TYPE_CHECKING:  # pragma: no cover
+
+    class PriorityQueue(  # type: ignore  # noqa: F811
+        asyncio.PriorityQueue[Tuple[int, int, _T]],
+        Generic[_T],
+    ):
+        """PriorityQueue based Queue."""
+
+        async def put_first(self, item: _T) -> None:
+            """
+            Put item in Queue with highest priority.
+
+            :param item: value to prepend
+            """
+            ...
+
+        async def put_last(self, item: _T) -> None:
+            """
+            Put item in Queue with lowest priority.
+
+            :param item: value to append
+            """
+            ...