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

(cherry picked from commit 934a82623793e9d52b85f74d5395d65927a52205)

Co-authored-by: Sam Bull <aa6bs0@sambull.org>
Co-authored-by: Sam Bull <aa6bs0@sambull.org>

diff --git a/Lib/asyncio/tasks.py b/Lib/asyncio/tasks.py
index 22ed3289edf3cd521ff27b418470fc783c9a8c0c..27a3c8c5a88df45cdfa8871b405660bfb189b87f 100644 (file)
--- a/Lib/asyncio/tasks.py
+++ b/Lib/asyncio/tasks.py
@@ -449,11 +449,9 @@ async def wait_for(fut, timeout, *, loop=None):
 
         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)
@@ -489,11 +487,9 @@ async def wait_for(fut, timeout, *, loop=None):
             # 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 402edaec0d8f9870dce434a3040a1cb99eab4cc3..a6afce0766493dec8f685612fbacee5c6f829f72 100644 (file)
--- a/Lib/test/test_asyncio/test_tasks.py
+++ b/Lib/test/test_asyncio/test_tasks.py
@@ -1121,20 +1121,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()
@@ -1213,24 +1209,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 (file)
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.