mirror of
https://github.com/python/cpython.git
synced 2025-07-08 03:45:36 +00:00
Issue #23812: Fix asyncio.Queue.get() to avoid loosing items on cancellation.
Patch by Gustavo J. A. M. Carneiro.
This commit is contained in:
parent
91e561aa77
commit
3fc0f2d288
3 changed files with 101 additions and 10 deletions
|
@ -47,7 +47,7 @@ class Queue:
|
|||
|
||||
# Futures.
|
||||
self._getters = collections.deque()
|
||||
# Pairs of (item, Future).
|
||||
# Futures
|
||||
self._putters = collections.deque()
|
||||
self._unfinished_tasks = 0
|
||||
self._finished = locks.Event(loop=self._loop)
|
||||
|
@ -98,7 +98,7 @@ class Queue:
|
|||
|
||||
def _consume_done_putters(self):
|
||||
# Delete waiters at the head of the put() queue who've timed out.
|
||||
while self._putters and self._putters[0][1].done():
|
||||
while self._putters and self._putters[0].done():
|
||||
self._putters.popleft()
|
||||
|
||||
def qsize(self):
|
||||
|
@ -148,8 +148,9 @@ class Queue:
|
|||
elif self._maxsize > 0 and self._maxsize <= self.qsize():
|
||||
waiter = futures.Future(loop=self._loop)
|
||||
|
||||
self._putters.append((item, waiter))
|
||||
self._putters.append(waiter)
|
||||
yield from waiter
|
||||
self._put(item)
|
||||
|
||||
else:
|
||||
self.__put_internal(item)
|
||||
|
@ -186,8 +187,7 @@ class Queue:
|
|||
self._consume_done_putters()
|
||||
if self._putters:
|
||||
assert self.full(), 'queue not full, why are putters waiting?'
|
||||
item, putter = self._putters.popleft()
|
||||
self.__put_internal(item)
|
||||
putter = self._putters.popleft()
|
||||
|
||||
# When a getter runs and frees up a slot so this putter can
|
||||
# run, we need to defer the put for a tick to ensure that
|
||||
|
@ -201,9 +201,39 @@ class Queue:
|
|||
return self._get()
|
||||
else:
|
||||
waiter = futures.Future(loop=self._loop)
|
||||
|
||||
self._getters.append(waiter)
|
||||
return (yield from waiter)
|
||||
try:
|
||||
return (yield from waiter)
|
||||
except futures.CancelledError:
|
||||
# if we get CancelledError, it means someone cancelled this
|
||||
# get() coroutine. But there is a chance that the waiter
|
||||
# already is ready and contains an item that has just been
|
||||
# removed from the queue. In this case, we need to put the item
|
||||
# back into the front of the queue. This get() must either
|
||||
# succeed without fault or, if it gets cancelled, it must be as
|
||||
# if it never happened.
|
||||
if waiter.done():
|
||||
self._put_it_back(waiter.result())
|
||||
raise
|
||||
|
||||
def _put_it_back(self, item):
|
||||
"""
|
||||
This is called when we have a waiter to get() an item and this waiter
|
||||
gets cancelled. In this case, we put the item back: wake up another
|
||||
waiter or put it in the _queue.
|
||||
"""
|
||||
self._consume_done_getters()
|
||||
if self._getters:
|
||||
assert not self._queue, (
|
||||
'queue non-empty, why are getters waiting?')
|
||||
|
||||
getter = self._getters.popleft()
|
||||
self._put_internal(item)
|
||||
|
||||
# getter cannot be cancelled, we just removed done getters
|
||||
getter.set_result(item)
|
||||
else:
|
||||
self._queue.appendleft(item)
|
||||
|
||||
def get_nowait(self):
|
||||
"""Remove and return an item from the queue.
|
||||
|
@ -213,8 +243,7 @@ class Queue:
|
|||
self._consume_done_putters()
|
||||
if self._putters:
|
||||
assert self.full(), 'queue not full, why are putters waiting?'
|
||||
item, putter = self._putters.popleft()
|
||||
self.__put_internal(item)
|
||||
putter = self._putters.popleft()
|
||||
# Wake putter on next tick.
|
||||
|
||||
# getter cannot be cancelled, we just removed done putters
|
||||
|
|
Loading…
Add table
Add a link
Reference in a new issue