bpo-37658: Actually return result in race condition (GH-29202)

3 files changed, 13 insertions, 36 deletions

diff --git a/Lib/asyncio/tasks.py b/Lib/asyncio/tasks.py
index 9a9d0d6..53eef84 100644
--- a/Lib/asyncio/tasks.py
+++ b/Lib/asyncio/tasks.py
@@ -415,11 +415,9 @@ async def wait_for(fut, timeout):
 
         await _cancel_and_wait(fut, loop=loop)
         try:
-            fut.result()
+            return fut.result()
         except exceptions.CancelledError as exc:
             raise exceptions.TimeoutError() from exc
-        else:
-            raise exceptions.TimeoutError()
 
     waiter = loop.create_future()
     timeout_handle = loop.call_later(timeout, _release_waiter, waiter)
@@ -455,11 +453,9 @@ async def wait_for(fut, timeout):
             # exception, we should re-raise it
             # See https://bugs.python.org/issue40607
             try:
-                fut.result()
+                return fut.result()
             except exceptions.CancelledError as exc:
                 raise exceptions.TimeoutError() from exc
-            else:
-                raise exceptions.TimeoutError()
     finally:
         timeout_handle.cancel()
 
diff --git a/Lib/test/test_asyncio/test_tasks.py b/Lib/test/test_asyncio/test_tasks.py
index 362fbf8..a88cb89 100644
--- a/Lib/test/test_asyncio/test_tasks.py
+++ b/Lib/test/test_asyncio/test_tasks.py
@@ -1009,20 +1009,16 @@ class BaseTaskTests:
         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
+        async def inner():
+            with contextlib.suppress(asyncio.CancelledError):
+                await asyncio.sleep(1)
+            return 1
 
-        loop = self.new_test_loop(gen)
+        async def main():
+            result = await asyncio.wait_for(inner(), timeout=.01)
+            assert result == 1
 
-        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")
+        asyncio.run(main())
 
     def test_wait_for_waits_for_task_cancellation(self):
         loop = asyncio.new_event_loop()
@@ -1101,24 +1097,6 @@ class BaseTaskTests:
         with self.assertRaises(FooException):
             loop.run_until_complete(foo())
 
-    def test_wait_for_raises_timeout_error_if_returned_during_cancellation(self):
-        loop = asyncio.new_event_loop()
-        self.addCleanup(loop.close)
-
-        async def foo():
-            async def inner():
-                try:
-                    await asyncio.sleep(0.2)
-                except asyncio.CancelledError:
-                    return 42
-
-            inner_task = self.new_task(loop, inner())
-
-            await asyncio.wait_for(inner_task, timeout=_EPSILON)
-
-        with self.assertRaises(asyncio.TimeoutError):
-            loop.run_until_complete(foo())
-
     def test_wait_for_self_cancellation(self):
         loop = asyncio.new_event_loop()
         self.addCleanup(loop.close)
diff --git a/Misc/NEWS.d/next/Library/2021-11-28-15-30-34.bpo-37658.8Hno7d.rst b/Misc/NEWS.d/next/Library/2021-11-28-15-30-34.bpo-37658.8Hno7d.rst
new file mode 100644
index 0000000..97d1e96
--- /dev/null
+++ b/Misc/NEWS.d/next/Library/2021-11-28-15-30-34.bpo-37658.8Hno7d.rst
@@ -0,0 +1,3 @@
+Fix issue when on certain conditions ``asyncio.wait_for()`` may allow a
+coroutine to complete successfully, but fail to return the result,
+potentially causing memory leaks or other issues.