gh-83055: Prevent get/put[_nowait] item in asyncio.Queue when there are woken tasks

Lib/asyncio/queues.py

@@ -54,6 +54,10 @@ def __init__(self, maxsize=0):
         self._finished.set()
         self._init(maxsize)
         self._is_shutdown = False
+        # See gh-83055
+        self._woken_getter = 0
+        self._woken_putter = 0
+        self._just_woken = False
 
     # These three are overridable in subclasses.
 
@@ -70,12 +74,20 @@ def _put(self, item):
 
     def _wakeup_next(self, waiters):
         # Wake up the next waiter (if any) that isn't cancelled.
+        woken = 0
         while waiters:
             waiter = waiters.popleft()
             if not waiter.done():
                 waiter.set_result(None)
+                woken = 1
                 break
 
+        if woken:
+            if waiters is self._getters:
+                self._woken_getter += woken
+            else:
+                self._woken_putter += woken
+
     def __repr__(self):
         return f'<{type(self).__name__} at {id(self):#x} {self._format()}>'
 
@@ -88,9 +100,9 @@ def _format(self):
         result = f'maxsize={self._maxsize!r}'
         if getattr(self, '_queue', None):
             result += f' _queue={list(self._queue)!r}'
-        if self._getters:
+        if self._getters or self._woken_getter:
             result += f' _getters[{len(self._getters)}]'
-        if self._putters:
+        if self._putters or self._woken_putter:
             result += f' _putters[{len(self._putters)}]'
         if self._unfinished_tasks:
             result += f' tasks={self._unfinished_tasks}'
@@ -130,13 +142,16 @@ async def put(self, item):
 
         Raises QueueShutDown if the queue has been shut down.
         """
+        self._just_woken = False
         while self.full():
             if self._is_shutdown:
                 raise QueueShutDown
             putter = self._get_loop().create_future()
             self._putters.append(putter)
             try:
                 await putter
+                self._just_woken = True
+                self._woken_putter -= 1
             except:
                 putter.cancel()  # Just in case putter is not done yet.
                 try:
@@ -149,21 +164,29 @@ async def put(self, item):
                 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._woken_putter -= 1
                     self._wakeup_next(self._putters)
                 raise
-        return self.put_nowait(item)
+        self.put_nowait(item)
 
     def put_nowait(self, item):
         """Put an item into the queue without blocking.
 
         If no free slot is immediately available, raise QueueFull.
 
         Raises QueueShutDown if the queue has been shut down.
+
+        Raises QueueWithPendingTasks if there are pending putters
+        or a woken putter is about to put an item into the queue.
         """
         if self._is_shutdown:
             raise QueueShutDown
         if self.full():
             raise QueueFull
+        if not self._just_woken and self._woken_putter > 0:
+            raise QueueFull("At least a woken putter task is waiting"
+                            " to insert an item")
+        self._just_woken = False
         self._put(item)
         self._unfinished_tasks += 1
         self._finished.clear()
@@ -177,13 +200,16 @@ async def get(self):
         Raises QueueShutDown if the queue has been shut down and is empty,
         or if the queue has been shut down immediately.
         """
+        self._just_woken = False
         while self.empty():
             if self._is_shutdown:
                 raise QueueShutDown
             getter = self._get_loop().create_future()
             self._getters.append(getter)
             try:
                 await getter
+                self._just_woken = True
+                self._woken_getter -= 1
             except:
                 getter.cancel()  # Just in case getter is not done yet.
                 try:
@@ -196,6 +222,7 @@ async def get(self):
                 if not self.empty() and not getter.cancelled():
                     # We were woken up by put_nowait(), but can't take
                     # the call.  Wake up the next in line.
+                    self._woken_getter -= 1
                     self._wakeup_next(self._getters)
                 raise
         return self.get_nowait()
@@ -213,6 +240,10 @@ def get_nowait(self):
             if self._is_shutdown:
                 raise QueueShutDown
             raise QueueEmpty
+        if not self._just_woken and self._woken_getter > 0:
+            raise QueueEmpty("At least a woken getter task is"
+                             " waiting to get an item")
+        self._just_woken = False
         item = self._get()
         self._wakeup_next(self._putters)
         return item

Lib/test/test_asyncio/test_base_events.py

@@ -181,6 +181,7 @@ def setUp(self):
         self.loop = base_events.BaseEventLoop()
         self.loop._selector = mock.Mock()
         self.loop._selector.select.return_value = ()
+        self.loop._process_events = mock.Mock()
         self.set_event_loop(self.loop)
 
     def test_not_implemented(self):
@@ -194,8 +195,8 @@ def test_not_implemented(self):
         self.assertRaises(
             NotImplementedError,
             self.loop._make_datagram_transport, m, m)
-        self.assertRaises(
-            NotImplementedError, self.loop._process_events, [])
+        #self.assertRaises(
+        #    NotImplementedError, self.loop._process_events, [])
         self.assertRaises(
             NotImplementedError, self.loop._write_to_self)
         self.assertRaises(

Lib/test/test_asyncio/test_queues.py

@@ -213,6 +213,62 @@ async def test_cancelled_getters_not_being_held_in_self_getters(self):
 
         self.assertEqual(len(queue._getters), 0)
 
+    async def test_get_item_when_exists_a_woken_getter(self):
+        # See gh-83055, prevents inappropriate get items because of
+        # a woken getter.
+        async def _get(q, results):
+            item = await q.get()
+            results.append(item)
+
+        maxsize = 1
+        q = asyncio.Queue(maxsize)
+        n = 0
+        results = []
+        g = 3
+        for _ in range(g):
+            asyncio.create_task(_get(q, results))
+        await asyncio.sleep(0.0)
+        # queue Full: 0 item in queue, 3 tasks into the _getters, 0 woken getter
+        self.assertEqual(q.qsize(), 0)
+        self.assertEqual(len(q._getters), g)
+        self.assertEqual(q._woken_getter, 0)
+
+        await q.put(n:=n+1)
+        g -= 1
+        # queue Full: 1 item in queue, 2 tasks into the _getters, 1 woken getter
+        self.assertEqual(q.qsize(), maxsize)
+        self.assertEqual(len(q._getters), g)
+        self.assertEqual(q._woken_getter, 1)
+
+        await q.put(n:=n+1)
+        g -= 1
+        # queue full: 1 item in queue, 1 task into the _getters,
+        # 1 item in transit
+        self.assertEqual(q.qsize(), maxsize)
+        self.assertEqual(len(q._getters), g)
+        self.assertEqual(q._woken_getter, 1)
+
+        msg = "At least a woken getter task is waiting"
+        with self.assertRaisesRegex(asyncio.QueueEmpty, expected_regex=msg):
+            await q.get()
+
+        await q.put(n:=n+1)
+        g -= 1
+        # queue full: 1 item in queue, 0 item waiting into the _getters,
+        # 1 item in transit
+        with self.assertRaisesRegex(asyncio.QueueEmpty, expected_regex=msg):
+            q.get_nowait()
+
+        self.assertEqual(q.qsize(), maxsize)
+        self.assertEqual(len(q._getters), 0)
+        self.assertEqual(q._woken_getter, 1)
+
+        # This is the end
+        await asyncio.sleep(0.0)
+        self.assertListEqual(results, list(range(1, n+1)))
+        self.assertEqual(q._woken_getter, 0)
+        self.assertTrue(q.empty())
+
 
 class QueuePutTests(unittest.IsolatedAsyncioTestCase):
 
@@ -431,6 +487,69 @@ async def test_cancelled_put_silence_value_error_exception(self):
         with self.assertRaises(asyncio.CancelledError):
             await put_task
 
+    async def test_put_can_break_fifo_order_with_pending_putters(self):
+        # See gh-83055, prevents inappropriate get items because of
+        # a woken getter. Can break the FIFO order of the queue.
+        maxsize = 1
+        q = asyncio.Queue(maxsize)
+        n = 3
+        results = []
+        for i in range(n):
+            asyncio.create_task(q.put(i+1))
+        await asyncio.sleep(0.0)
+        g = n
+
+        # Queue full: 1 item in queue, 2 items waiting into the _putters
+        self.assertEqual(q.qsize(), maxsize)
+        self.assertEqual(len(q._putters), g-maxsize)
+        self.assertEqual(q._woken_putter, 0)
+
+        item = await q.get()
+        results.append(item)
+        g -= 1
+        # queue empty: 0 item in queue, 1 item into the _putters,
+        # 1 item in transit
+        self.assertEqual(q.qsize(), 0)
+        self.assertEqual(len(q._putters), g-maxsize)
+        self.assertEqual(q._woken_putter, 1)
+
+        # queue is empty, but put_nowait fails,
+        self.assertTrue(q.empty())
+        msg = 'At least a woken putter task is waiting'
+        with self.assertRaisesRegex(asyncio.QueueFull, expected_regex=msg):
+            q.put_nowait(n+1)
+
+        # get an another item
+        item = await q.get()
+        results.append(item)
+        g -= 1
+
+        # Queue empty: 0 item in queue, 0 item into the _putters,
+        # 1 item in transit
+        self.assertEqual(q.qsize(), 0)
+        self.assertEqual(len(q._putters), g-maxsize)
+        self.assertEqual(q._woken_putter, 1)
+
+        # queue is empty, but put fails,
+        self.assertTrue(q.empty())
+        with self.assertRaisesRegex(asyncio.QueueFull, expected_regex=msg):
+            await q.put(n+1)
+
+        # get an another item
+        await asyncio.sleep(0)
+        item = await q.get()
+        results.append(item)
+        g -= 1
+        # Queue empty: 0 item in queue, 0 item into the _putters,
+        # 0 item in transit
+        self.assertEqual(q.qsize(), 0)
+        self.assertEqual(len(q._putters), 0)
+        self.assertEqual(q._woken_putter, 0)
+
+        # This is the end
+        self.assertListEqual(results, list(range(1, n+1)))
+        self.assertTrue(q.empty())
+
 
 class LifoQueueTests(unittest.IsolatedAsyncioTestCase):