From 6e1954cd8286e083e7f8d09516d91b6b15769a4e Mon Sep 17 00:00:00 2001
From: "Miss Islington (bot)"
 <31488909+miss-islington@users.noreply.github.com>
Date: Wed, 26 Aug 2020 11:26:28 -0700
Subject: [PATCH] bpo-37658: Fix asyncio.wait_for() to respect waited task
 status (GH-21894) (#21965)

Currently, if `asyncio.wait_for()` itself is cancelled it will always
raise `CancelledError` regardless if the underlying task is still
running.  This is similar to a race with the timeout, which is handled
already.
(cherry picked from commit a2118a14627256197bddcf4fcecad4c264c1e39d)

Co-authored-by: Elvis Pranskevichus <elvis@magic.io>
---
 Lib/asyncio/tasks.py                             |  9 ++++++---
 Lib/test/test_asyncio/test_tasks.py              | 16 ++++++++++++++++
 .../2020-08-15-15-21-40.bpo-37658.f9nivB.rst     |  2 ++
 3 files changed, 24 insertions(+), 3 deletions(-)
 create mode 100644 Misc/NEWS.d/next/Library/2020-08-15-15-21-40.bpo-37658.f9nivB.rst

diff --git a/Lib/asyncio/tasks.py b/Lib/asyncio/tasks.py
index 66e81f9200a..373be37fb4b 100644
--- a/Lib/asyncio/tasks.py
+++ b/Lib/asyncio/tasks.py
@@ -475,9 +475,12 @@ async def wait_for(fut, timeout, *, loop=None):
         try:
             await waiter
         except exceptions.CancelledError:
-            fut.remove_done_callback(cb)
-            fut.cancel()
-            raise
+            if fut.done():
+                return fut.result()
+            else:
+                fut.remove_done_callback(cb)
+                fut.cancel()
+                raise
 
         if fut.done():
             return fut.result()
diff --git a/Lib/test/test_asyncio/test_tasks.py b/Lib/test/test_asyncio/test_tasks.py
index c21b06938fe..a70dc0a87e1 100644
--- a/Lib/test/test_asyncio/test_tasks.py
+++ b/Lib/test/test_asyncio/test_tasks.py
@@ -861,6 +861,22 @@ class BaseTaskTests:
         res = loop.run_until_complete(task)
         self.assertEqual(res, "ok")
 
+    def test_wait_for_cancellation_race_condition(self):
+        def gen():
+            yield 0.1
+            yield 0.1
+            yield 0.1
+            yield 0.1
+
+        loop = self.new_test_loop(gen)
+
+        fut = self.new_future(loop)
+        loop.call_later(0.1, fut.set_result, "ok")
+        task = loop.create_task(asyncio.wait_for(fut, timeout=1))
+        loop.call_later(0.1, task.cancel)
+        res = loop.run_until_complete(task)
+        self.assertEqual(res, "ok")
+
     def test_wait_for_waits_for_task_cancellation(self):
         loop = asyncio.new_event_loop()
         self.addCleanup(loop.close)
diff --git a/Misc/NEWS.d/next/Library/2020-08-15-15-21-40.bpo-37658.f9nivB.rst b/Misc/NEWS.d/next/Library/2020-08-15-15-21-40.bpo-37658.f9nivB.rst
new file mode 100644
index 00000000000..694fbbbf346
--- /dev/null
+++ b/Misc/NEWS.d/next/Library/2020-08-15-15-21-40.bpo-37658.f9nivB.rst
@@ -0,0 +1,2 @@
+:meth:`asyncio.wait_for` now properly handles races between cancellation of
+itself and the completion of the wrapped awaitable.