gh-102780: Fix uncancel() call in asyncio timeouts (#102815)

Also use `raise TimeOut from <CancelledError instance>` so that the CancelledError is set
in the `__cause__` field rather than in the `__context__` field.

Co-authored-by: Guido van Rossum <[email protected]>
Co-authored-by: Alex Waygood <[email protected]>

Author: 3 people

Doc/library/asyncio-task.rst

@@ -300,13 +300,17 @@ in the task at the next opportunity.
 It is recommended that coroutines use ``try/finally`` blocks to robustly
 perform clean-up logic. In case :exc:`asyncio.CancelledError`
 is explicitly caught, it should generally be propagated when
-clean-up is complete. Most code can safely ignore :exc:`asyncio.CancelledError`.
+clean-up is complete. :exc:`asyncio.CancelledError` directly subclasses
+:exc:`BaseException` so most code will not need to be aware of it.
 
 The asyncio components that enable structured concurrency, like
 :class:`asyncio.TaskGroup` and :func:`asyncio.timeout`,
 are implemented using cancellation internally and might misbehave if
 a coroutine swallows :exc:`asyncio.CancelledError`. Similarly, user code
-should not call :meth:`uncancel <asyncio.Task.uncancel>`.
+should not generally call :meth:`uncancel <asyncio.Task.uncancel>`.
+However, in cases when suppressing :exc:`asyncio.CancelledError` is
+truly desired, it is necessary to also call ``uncancel()`` to completely
+remove the cancellation state.
 
 .. _taskgroups:
 
@@ -1148,7 +1152,9 @@ Task Object
       Therefore, unlike :meth:`Future.cancel`, :meth:`Task.cancel` does
       not guarantee that the Task will be cancelled, although
       suppressing cancellation completely is not common and is actively
-      discouraged.
+      discouraged.  Should the coroutine nevertheless decide to suppress
+      the cancellation, it needs to call :meth:`Task.uncancel` in addition
+      to catching the exception.
 
       .. versionchanged:: 3.9
          Added the *msg* parameter.
@@ -1238,6 +1244,10 @@ Task Object
       with :meth:`uncancel`.  :class:`TaskGroup` context managers use
       :func:`uncancel` in a similar fashion.
 
+      If end-user code is, for some reason, suppresing cancellation by
+      catching :exc:`CancelledError`, it needs to call this method to remove
+      the cancellation state.
+
    .. method:: cancelling()
 
       Return the number of pending cancellation requests to this Task, i.e.,

Lib/asyncio/timeouts.py

@@ -84,6 +84,7 @@ def __repr__(self) -> str:
     async def __aenter__(self) -> "Timeout":
         self._state = _State.ENTERED
         self._task = tasks.current_task()
+        self._cancelling = self._task.cancelling()
         if self._task is None:
             raise RuntimeError("Timeout should be used inside a task")
         self.reschedule(self._when)
@@ -104,10 +105,10 @@ async def __aexit__(
         if self._state is _State.EXPIRING:
             self._state = _State.EXPIRED
 
-            if self._task.uncancel() == 0 and exc_type is exceptions.CancelledError:
-                # Since there are no outstanding cancel requests, we're
+            if self._task.uncancel() <= self._cancelling and exc_type is exceptions.CancelledError:
+                # Since there are no new cancel requests, we're
                 # handling this.
-                raise TimeoutError
+                raise TimeoutError from exc_val
         elif self._state is _State.ENTERED:
             self._state = _State.EXITED
 

Lib/test/test_asyncio/test_timeouts.py

@@ -247,6 +247,36 @@ async def test_nested_timeout_in_finally(self):
                         async with asyncio.timeout(0.01):
                             await asyncio.sleep(10)
 
+    async def test_timeout_after_cancellation(self):
+        try:
+            asyncio.current_task().cancel()
+            await asyncio.sleep(1)  # work which will be cancelled
+        except asyncio.CancelledError:
+            pass
+        finally:
+            with self.assertRaises(TimeoutError):
+                async with asyncio.timeout(0.0):
+                    await asyncio.sleep(1)  # some cleanup
+
+    async def test_cancel_in_timeout_after_cancellation(self):
+        try:
+            asyncio.current_task().cancel()
+            await asyncio.sleep(1)  # work which will be cancelled
+        except asyncio.CancelledError:
+            pass
+        finally:
+            with self.assertRaises(asyncio.CancelledError):
+                async with asyncio.timeout(1.0):
+                    asyncio.current_task().cancel()
+                    await asyncio.sleep(2)  # some cleanup
+
+    async def test_timeout_exception_cause (self):
+        with self.assertRaises(asyncio.TimeoutError) as exc:
+            async with asyncio.timeout(0):
+                await asyncio.sleep(1)
+        cause = exc.exception.__cause__
+        assert isinstance(cause, asyncio.CancelledError)
+
 
 if __name__ == '__main__':
     unittest.main()

Misc/NEWS.d/next/Library/2023-03-22-16-15-18.gh-issue-102780.NEcljy.rst

@@ -0,0 +1,3 @@
+The :class:`asyncio.Timeout` context manager now works reliably even when performing cleanup due
+to task cancellation.  Previously it could raise a
+:exc:`~asyncio.CancelledError` instead of an :exc:`~asyncio.TimeoutError` in such cases.