From 382a5635bd10c237c3e23e346b21cde27e48d7fa Mon Sep 17 00:00:00 2001
From: romasku <romasku135@gmail.com>
Date: Fri, 15 May 2020 23:12:05 +0300
Subject: [PATCH] bpo-40607: Reraise exception during task cancelation in
 asyncio.wait_for() (GH-20054)

Currently, if asyncio.wait_for() timeout expires, it cancels
inner future and then always raises TimeoutError. In case
those future is task, it can handle cancelation mannually,
and those process can lead to some other exception. Current
implementation silently loses thoses exception.

To resolve this, wait_for will check was the cancelation
successfull or not. In case there was exception, wait_for
will reraise it.

Co-authored-by: Roman Skurikhin <roman.skurikhin@cruxlab.com>
---
 Doc/library/asyncio-task.rst                  |  3 +-
 Lib/asyncio/tasks.py                          | 10 +++-
 Lib/test/test_asyncio/test_tasks.py           | 55 +++++++++++++++++--
 Misc/ACKS                                     |  1 +
 .../2020-05-13-15-32-13.bpo-40607.uSPFCi.rst  |  3 +
 5 files changed, 66 insertions(+), 6 deletions(-)
 create mode 100644 Misc/NEWS.d/next/Library/2020-05-13-15-32-13.bpo-40607.uSPFCi.rst

diff --git a/Doc/library/asyncio-task.rst b/Doc/library/asyncio-task.rst
index 42e2b4e2fc5..bc8a2722bcc 100644
--- a/Doc/library/asyncio-task.rst
+++ b/Doc/library/asyncio-task.rst
@@ -453,7 +453,8 @@ Timeouts
    wrap it in :func:`shield`.
 
    The function will wait until the future is actually cancelled,
-   so the total wait time may exceed the *timeout*.
+   so the total wait time may exceed the *timeout*. If an exception
+   happens during cancellation, it is propagated.
 
    If the wait is cancelled, the future *aw* is also cancelled.
 
diff --git a/Lib/asyncio/tasks.py b/Lib/asyncio/tasks.py
index 717837d8568..f5de1a2eea9 100644
--- a/Lib/asyncio/tasks.py
+++ b/Lib/asyncio/tasks.py
@@ -496,7 +496,15 @@ async def wait_for(fut, timeout, *, loop=None):
             # after wait_for() returns.
             # See https://bugs.python.org/issue32751
             await _cancel_and_wait(fut, loop=loop)
-            raise exceptions.TimeoutError()
+            # In case task cancellation failed with some
+            # exception, we should re-raise it
+            # See https://bugs.python.org/issue40607
+            try:
+                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 6eb6b46ec8a..0f8d921c5bc 100644
--- a/Lib/test/test_asyncio/test_tasks.py
+++ b/Lib/test/test_asyncio/test_tasks.py
@@ -80,6 +80,12 @@ def __await__(self):
         return self
 
 
+# The following value can be used as a very small timeout:
+# it passes check "timeout > 0", but has almost
+# no effect on the test performance
+_EPSILON = 0.0001
+
+
 class BaseTaskTests:
 
     Task = None
@@ -904,12 +910,53 @@ async def inner():
 
             inner_task = self.new_task(loop, inner())
 
-            with self.assertRaises(asyncio.TimeoutError):
-                await asyncio.wait_for(inner_task, timeout=0.1)
+            await asyncio.wait_for(inner_task, timeout=_EPSILON)
 
-            self.assertTrue(task_done)
+        with self.assertRaises(asyncio.TimeoutError) as cm:
+            loop.run_until_complete(foo())
 
-        loop.run_until_complete(foo())
+        self.assertTrue(task_done)
+        chained = cm.exception.__context__
+        self.assertEqual(type(chained), asyncio.CancelledError)
+
+    def test_wait_for_reraises_exception_during_cancellation(self):
+        loop = asyncio.new_event_loop()
+        self.addCleanup(loop.close)
+
+        class FooException(Exception):
+            pass
+
+        async def foo():
+            async def inner():
+                try:
+                    await asyncio.sleep(0.2)
+                finally:
+                    raise FooException
+
+            inner_task = self.new_task(loop, inner())
+
+            await asyncio.wait_for(inner_task, timeout=_EPSILON)
+
+        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()
diff --git a/Misc/ACKS b/Misc/ACKS
index b479aa5d807..fad920b0510 100644
--- a/Misc/ACKS
+++ b/Misc/ACKS
@@ -1593,6 +1593,7 @@ J. Sipprell
 Ngalim Siregar
 Kragen Sitaker
 Kaartic Sivaraam
+Roman Skurikhin
 Ville Skyttä
 Michael Sloan
 Nick Sloan
diff --git a/Misc/NEWS.d/next/Library/2020-05-13-15-32-13.bpo-40607.uSPFCi.rst b/Misc/NEWS.d/next/Library/2020-05-13-15-32-13.bpo-40607.uSPFCi.rst
new file mode 100644
index 00000000000..8060628b595
--- /dev/null
+++ b/Misc/NEWS.d/next/Library/2020-05-13-15-32-13.bpo-40607.uSPFCi.rst
@@ -0,0 +1,3 @@
+When cancelling a task due to timeout, :meth:`asyncio.wait_for` will now
+propagate the exception if an error happens during cancellation.
+Patch by Roman Skurikhin.