← 输入框|github.com|刷新
Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
13 changes: 13 additions & 0 deletions Doc/library/asyncio-queue.rst
Original file line number Diff line number Diff line change
Expand Up @@ -57,6 +57,13 @@ Queue
If the queue was initialized with ``maxsize=0`` (the default),
then :meth:`full` never returns ``True``.

A slot freed while :meth:`put` calls are waiting is reserved
for the first of them until it resumes, so :meth:`full` can return
``True`` while :meth:`qsize` is less than :attr:`maxsize`.

.. versionchanged:: next
Slots reserved for waiting :meth:`put` calls are counted as full.

.. method:: get()
:async:

Expand Down Expand Up @@ -90,6 +97,12 @@ Queue
Put an item into the queue. If the queue is full, wait until a
free slot is available before adding the item.

Waiting calls are served in first in, first out order: a freed
slot cannot be taken by a later :meth:`put` or :meth:`put_nowait`.

.. versionchanged:: next
A freed slot is reserved for the longest waiting call.

Raises :exc:`QueueShutDown` if the queue has been shut down.

.. method:: put_nowait(item)
Expand Down
28 changes: 23 additions & 5 deletions Lib/asyncio/queues.py
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand Down Expand Up @@ -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.
Expand All @@ -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)

def put_nowait(self, item):
Expand Down Expand Up @@ -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):
Expand Down Expand Up @@ -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()
Expand Down
92 changes: 92 additions & 0 deletions Lib/test/test_asyncio/test_queues.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
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):

Expand Down
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.
Loading