chore: refactor to back by TaskQueue instead of Event

Author: imnotjames

Diff for: asyncio/queue.py

@@ -8,7 +8,7 @@
 Exceptions and classes related to asyncio Queue implementations.
 """
 
-from . import event
+from . import core
 
 
 class QueueEmpty(Exception):
@@ -19,6 +19,19 @@ class QueueFull(Exception):
     """Raised when the Queue.put_nowait() method is called on a full Queue."""
 
 
+async def _wait_on_task_queue(task_queue: core.TaskQueue):
+    task_queue.push(core.cur_task)
+    # Set calling task's data to the TaskQueue so it can be removed if needed
+    core.cur_task.data = task_queue
+    # Send control back
+    await core._never()
+
+
+def _release_task_queue(task_queue: core.TaskQueue):
+    while task_queue.peek():
+        core._task_queue.push(task_queue.pop())
+
+
 class Queue:
     """
     A queue, useful for coordinating producer and consumer coroutines.
@@ -37,27 +50,21 @@ def __init__(self, maxsize=0):
 
         self._queue = []
 
-        self._join_counter = 0
-        self._join_event = event.Event()
-        self._join_event.set()
+        self._active_tasks = 0
 
-        self._put_event = event.Event()
-        self._get_event = event.Event()
+        self._waiting_for_completion = core.TaskQueue()
+        self._waiting_for_put = core.TaskQueue()
+        self._waiting_for_get = core.TaskQueue()
 
     def _get(self):
         value = self._queue.pop(0)
-        self._get_event.set()
-        self._get_event.clear()
+        _release_task_queue(self._waiting_for_get)
         return value
 
     def _put(self, val):
-        self._join_counter += 1
-        self._join_event.clear()
-
         self._queue.append(val)
-
-        self._put_event.set()
-        self._put_event.clear()
+        self._active_tasks += 1
+        _release_task_queue(self._waiting_for_put)
 
     async def get(self):
         """
@@ -66,7 +73,7 @@ async def get(self):
         If queue is empty, wait until an item is available.
         """
         while self.empty():
-            await self._put_event.wait()
+            await _wait_on_task_queue(self._waiting_for_put)
         return self._get()
 
     def get_nowait(self):
@@ -87,7 +94,7 @@ async def put(self, val):
         slot is available before adding item.
         """
         while self.full():
-            await self._get_event.wait()
+            await _wait_on_task_queue(self._waiting_for_get)
         self._put(val)
 
     def put_nowait(self, val):
@@ -129,17 +136,17 @@ def task_done(self):
         Raises ValueError if called more times than there were items placed in
         the queue.
         """
-        if self._join_counter == 0:
+        if self._active_tasks == 0:
             # Can't have less than 0
             raise ValueError("task_done() called too many times")
 
-        self._join_counter -= 1
+        self._active_tasks -= 1
 
-        if self._join_counter == 0:
-            self._join_event.set()
+        if self._active_tasks == 0:
+            _release_task_queue(self._waiting_for_completion)
 
     async def join(self):
         """
         Block until all items in the queue have been gotten and processed.
         """
-        await self._join_event.wait()
+        await _wait_on_task_queue(self._waiting_for_completion)