summaryrefslogtreecommitdiffstats
path: root/Lib/asyncio/tasks.py
diff options
context:
space:
mode:
Diffstat (limited to 'Lib/asyncio/tasks.py')
-rw-r--r--Lib/asyncio/tasks.py23
1 files changed, 21 insertions, 2 deletions
diff --git a/Lib/asyncio/tasks.py b/Lib/asyncio/tasks.py
index 6cef33d..72792a2 100644
--- a/Lib/asyncio/tasks.py
+++ b/Lib/asyncio/tasks.py
@@ -412,14 +412,17 @@ async def wait_for(fut, timeout, *, loop=None):
return fut.result()
else:
fut.remove_done_callback(cb)
- fut.cancel()
+ # We must ensure that the task is not running
+ # after wait_for() returns.
+ # See https://bugs.python.org/issue32751
+ await _cancel_and_wait(fut, loop=loop)
raise futures.TimeoutError()
finally:
timeout_handle.cancel()
async def _wait(fs, timeout, return_when, loop):
- """Internal helper for wait() and wait_for().
+ """Internal helper for wait().
The fs argument must be a collection of Futures.
"""
@@ -461,6 +464,22 @@ async def _wait(fs, timeout, return_when, loop):
return done, pending
+async def _cancel_and_wait(fut, loop):
+ """Cancel the *fut* future or task and wait until it completes."""
+
+ waiter = loop.create_future()
+ cb = functools.partial(_release_waiter, waiter)
+ fut.add_done_callback(cb)
+
+ try:
+ fut.cancel()
+ # We cannot wait on *fut* directly to make
+ # sure _cancel_and_wait itself is reliably cancellable.
+ await waiter
+ finally:
+ fut.remove_done_callback(cb)
+
+
# This is *not* a @coroutine! It is just an iterator (yielding Futures).
def as_completed(fs, *, loop=None, timeout=None):
"""Return an iterator whose values are coroutines.