Repository navigation
gh-83055: Reserve freed asyncio.Queue slots for woken putters #160967
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: main
Are you sure you want to change the base?
Changes from all commits
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -49,6 +49,10 @@ def __init__(self, maxsize=0): | |
| self._getters = collections.deque() | ||
| # Futures. | ||
| self._putters = collections.deque() | ||
| # Free slots promised to putters that were woken up but have not | ||
| # run yet. Counting them in full() keeps a newcomer from taking | ||
| # the slot and overtaking the waiting putters. | ||
| self._reserved = 0 | ||
| self._unfinished_tasks = 0 | ||
| self._finished = locks.Event() | ||
| self._finished.set() | ||
|
|
@@ -120,7 +124,16 @@ def full(self): | |
| if self._maxsize <= 0: | ||
| return False | ||
| else: | ||
| return self.qsize() >= self._maxsize | ||
| return self.qsize() + self._reserved >= self._maxsize | ||
|
|
||
| def _wakeup_next_putter(self): | ||
| if not self._is_shutdown and not self.full(): | ||
| while self._putters: | ||
| putter = self._putters.popleft() | ||
| if not putter.done(): | ||
| putter.set_result(None) | ||
| self._reserved += 1 | ||
| break | ||
|
|
||
| async def put(self, item): | ||
| """Put an item into the queue. | ||
|
|
@@ -146,11 +159,15 @@ async def put(self, item): | |
| # The putter could be removed from self._putters by a | ||
| # previous get_nowait call or a shutdown call. | ||
| pass | ||
| if not self.full() and not putter.cancelled(): | ||
| if not putter.cancelled() and not self._is_shutdown: | ||
| # We were woken up by get_nowait(), but can't take | ||
| # the call. Wake up the next in line. | ||
| self._wakeup_next(self._putters) | ||
| # the call. Pass the reserved slot to the next in line. | ||
| self._reserved -= 1 | ||
| self._wakeup_next_putter() | ||
| raise | ||
| if not self._is_shutdown: | ||
| # Take the slot reserved for us. | ||
| self._reserved -= 1 | ||
| return self.put_nowait(item) | ||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more.
Contributor
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Correct, but this line is not affected by the PR. I prefer to keep it untouched here |
||
|
|
||
| def put_nowait(self, item): | ||
|
|
@@ -214,7 +231,7 @@ def get_nowait(self): | |
| raise QueueShutDown | ||
| raise QueueEmpty | ||
| item = self._get() | ||
| self._wakeup_next(self._putters) | ||
| self._wakeup_next_putter() | ||
| return item | ||
|
|
||
| def task_done(self): | ||
|
|
@@ -262,6 +279,7 @@ def shutdown(self, immediate=False): | |
| is reduced to zero, callers of Queue.join are unblocked. | ||
| """ | ||
| self._is_shutdown = True | ||
| self._reserved = 0 | ||
| if immediate: | ||
| while not self.empty(): | ||
| self._get() | ||
|
|
||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -431,6 +431,98 @@ async def test_cancelled_put_silence_value_error_exception(self): | |
| with self.assertRaises(asyncio.CancelledError): | ||
| await put_task | ||
|
|
||
| async def test_put_nowait_does_not_steal_slot_from_woken_putter(self): | ||
| # gh-83055: the slot freed by get_nowait() belongs to the | ||
| # woken putter until it runs. | ||
| q = asyncio.Queue(1) | ||
| q.put_nowait(0) | ||
| put_1 = asyncio.create_task(q.put(1)) | ||
| put_2 = asyncio.create_task(q.put(2)) | ||
| await asyncio.sleep(0) | ||
|
|
||
| self.assertEqual(q.get_nowait(), 0) | ||
| self.assertTrue(q.full()) | ||
| with self.assertRaises(asyncio.QueueFull): | ||
| q.put_nowait(3) | ||
|
|
||
| self.assertEqual(await q.get(), 1) | ||
| self.assertEqual(await q.get(), 2) | ||
| await put_1 | ||
| await put_2 | ||
| self.assertEqual(q._reserved, 0) | ||
|
|
||
| async def test_put_nowait_uses_unreserved_slots(self): | ||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. This test makes me deeply uncomfortable. Because, IMO, for an |
||
| q = asyncio.Queue(2) | ||
| q.put_nowait(1) | ||
| q.put_nowait(2) | ||
| put_3 = asyncio.create_task(q.put(3)) | ||
| await asyncio.sleep(0) | ||
|
|
||
| # Two free slots, only one of them is reserved for put_3. | ||
| q.get_nowait() | ||
| q.get_nowait() | ||
| self.assertFalse(q.full()) | ||
| q.put_nowait(4) | ||
| self.assertTrue(q.full()) | ||
|
|
||
| await put_3 | ||
| self.assertEqual(sorted([q.get_nowait(), q.get_nowait()]), [3, 4]) | ||
| self.assertEqual(q._reserved, 0) | ||
|
|
||
| async def test_put_order_with_many_putters(self): | ||
| # gh-83055: items must be got in the order put() was called. | ||
| q = asyncio.Queue(3) | ||
| counter = 0 | ||
|
|
||
| async def putter(n): | ||
| nonlocal counter | ||
| for i in range(50): | ||
| item = counter | ||
| counter += 1 | ||
| await q.put(item) | ||
| if (i + n) % 3 == 0: | ||
| await asyncio.sleep(0) | ||
|
|
||
| async with asyncio.TaskGroup() as tg: | ||
| for n in range(10): | ||
| tg.create_task(putter(n)) | ||
| items = [] | ||
| for i in range(500): | ||
| items.append(await q.get()) | ||
| if i % 2: | ||
| await asyncio.sleep(0) | ||
|
|
||
| self.assertEqual(items, sorted(items)) | ||
| self.assertEqual(q._reserved, 0) | ||
|
|
||
| async def test_cancelled_woken_putter_passes_slot_to_next(self): | ||
| q = asyncio.Queue(1) | ||
| q.put_nowait(0) | ||
| put_1 = asyncio.create_task(q.put(1)) | ||
| put_2 = asyncio.create_task(q.put(2)) | ||
| await asyncio.sleep(0) | ||
|
|
||
| q.get_nowait() | ||
| put_1.cancel() | ||
| with self.assertRaises(asyncio.CancelledError): | ||
| await put_1 | ||
| await put_2 | ||
| self.assertEqual(q.get_nowait(), 2) | ||
| self.assertEqual(q._reserved, 0) | ||
| self.assertFalse(q.full()) | ||
|
|
||
| async def test_shutdown_with_woken_putter(self): | ||
| q = asyncio.Queue(1) | ||
| q.put_nowait(0) | ||
| put_1 = asyncio.create_task(q.put(1)) | ||
| await asyncio.sleep(0) | ||
|
|
||
| q.get_nowait() | ||
| q.shutdown() | ||
| with self.assertRaises(asyncio.QueueShutDown): | ||
| await put_1 | ||
| self.assertEqual(q._reserved, 0) | ||
|
|
||
|
|
||
| class LifoQueueTests(unittest.IsolatedAsyncioTestCase): | ||
|
|
||
|
|
||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,5 @@ | ||
| :meth:`asyncio.Queue.get` and :meth:`asyncio.Queue.get_nowait` now reserve | ||
| the freed slot for the longest waiting :meth:`asyncio.Queue.put` call, so a | ||
| later :meth:`~asyncio.Queue.put` or :meth:`~asyncio.Queue.put_nowait` can no | ||
| longer take it and put items out of order. :meth:`asyncio.Queue.full` counts | ||
| reserved slots. |
Uh oh!
There was an error while loading. Please reload this page.