diff --git a/Lib/asyncio/queues.py b/Lib/asyncio/queues.py index 8871cdc16c68b3..6810197a89a857 100644 --- a/Lib/asyncio/queues.py +++ b/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,6 +142,7 @@ 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 @@ -137,6 +150,8 @@ async def put(self, item): 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,9 +164,10 @@ 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. @@ -164,6 +180,10 @@ def put_nowait(self, item): 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,6 +197,7 @@ 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 @@ -184,6 +205,8 @@ async def get(self): 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 +219,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 +237,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 diff --git a/Lib/test/test_asyncio/test_queues.py b/Lib/test/test_asyncio/test_queues.py index 7c1548f1948f77..cf8689cec8861b 100644 --- a/Lib/test/test_asyncio/test_queues.py +++ b/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):