bpo-32751: Wait for task cancellation in asyncio.wait_for() (GH-7216)

Currently, asyncio.wait_for(fut), upon reaching the timeout deadline,
cancels the future and returns immediately.  This is problematic for
when *fut* is a Task, because it will be left running for an arbitrary
amount of time.  This behavior is iself surprising and may lead to
related bugs such as the one described in bpo-33638:

    condition = asyncio.Condition()
    async with condition:
        await asyncio.wait_for(condition.wait(), timeout=0.5)

Currently, instead of raising a TimeoutError, the above code will fail
with `RuntimeError: cannot wait on un-acquired lock`, because
`__aexit__` is reached _before_ `condition.wait()` finishes its
cancellation and re-acquires the condition lock.

To resolve this, make `wait_for` await for the task cancellation.
The tradeoff here is that the `timeout` promise may be broken if the
task decides to handle its cancellation in a slow way.  This represents
a behavior change and should probably not be back-patched to 3.6 and
earlier.
(cherry picked from commit e2b340ab4196e1beb902327f503574b5d7369185)

Co-authored-by: Elvis Pranskevichus <elvis@magic.io>

1 files changed, 21 insertions, 2 deletions

diff --git a/Lib/asyncio/tasks.py b/Lib/asyncio/tasks.py
index 6cef33d521..72792a25cf 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.