bpo-31033: Improve the traceback for cancelled asyncio tasks (GH-19951)

When an asyncio.Task is cancelled, the exception traceback now
starts with where the task was first interrupted.  Previously,
the traceback only had "depth one."

Author: cjerdonek

Include/internal/pycore_pyerrors.h

@@ -14,6 +14,20 @@ static inline PyObject* _PyErr_Occurred(PyThreadState *tstate)
     return tstate->curexc_type;
 }
 
+static inline void _PyErr_ClearExcState(_PyErr_StackItem *exc_state)
+{
+    PyObject *t, *v, *tb;
+    t = exc_state->exc_type;
+    v = exc_state->exc_value;
+    tb = exc_state->exc_traceback;
+    exc_state->exc_type = NULL;
+    exc_state->exc_value = NULL;
+    exc_state->exc_traceback = NULL;
+    Py_XDECREF(t);
+    Py_XDECREF(v);
+    Py_XDECREF(tb);
+}
+
 
 PyAPI_FUNC(void) _PyErr_Fetch(
     PyThreadState *tstate,
@@ -36,6 +50,9 @@ PyAPI_FUNC(void) _PyErr_SetObject(
     PyObject *type,
     PyObject *value);
 
+PyAPI_FUNC(void) _PyErr_ChainStackItem(
+    _PyErr_StackItem *exc_state);
+
 PyAPI_FUNC(void) _PyErr_Clear(PyThreadState *tstate);
 
 PyAPI_FUNC(void) _PyErr_SetNone(PyThreadState *tstate, PyObject *exception);

Lib/asyncio/futures.py

@@ -52,6 +52,8 @@ class Future:
     _loop = None
     _source_traceback = None
     _cancel_message = None
+    # A saved CancelledError for later chaining as an exception context.
+    _cancelled_exc = None
 
     # This field is used for a dual purpose:
     # - Its presence is a marker to declare that a class implements
@@ -124,6 +126,21 @@ def get_loop(self):
             raise RuntimeError("Future object is not initialized.")
         return loop
 
+    def _make_cancelled_error(self):
+        """Create the CancelledError to raise if the Future is cancelled.
+
+        This should only be called once when handling a cancellation since
+        it erases the saved context exception value.
+        """
+        if self._cancel_message is None:
+            exc = exceptions.CancelledError()
+        else:
+            exc = exceptions.CancelledError(self._cancel_message)
+        exc.__context__ = self._cancelled_exc
+        # Remove the reference since we don't need this anymore.
+        self._cancelled_exc = None
+        return exc
+
     def cancel(self, msg=None):
         """Cancel the future and schedule callbacks.
 
@@ -175,9 +192,8 @@ def result(self):
         the future is done and has an exception set, this exception is raised.
         """
         if self._state == _CANCELLED:
-            raise exceptions.CancelledError(
-                '' if self._cancel_message is None else self._cancel_message)
-
+            exc = self._make_cancelled_error()
+            raise exc
         if self._state != _FINISHED:
             raise exceptions.InvalidStateError('Result is not ready.')
         self.__log_traceback = False
@@ -194,8 +210,8 @@ def exception(self):
         InvalidStateError.
         """
         if self._state == _CANCELLED:
-            raise exceptions.CancelledError(
-                '' if self._cancel_message is None else self._cancel_message)
+            exc = self._make_cancelled_error()
+            raise exc
         if self._state != _FINISHED:
             raise exceptions.InvalidStateError('Exception is not set.')
         self.__log_traceback = False

Lib/asyncio/tasks.py

@@ -270,8 +270,7 @@ def __step(self, exc=None):
                 f'_step(): already done: {self!r}, {exc!r}')
         if self._must_cancel:
             if not isinstance(exc, exceptions.CancelledError):
-                exc = exceptions.CancelledError(''
-                    if self._cancel_message is None else self._cancel_message)
+                exc = self._make_cancelled_error()
             self._must_cancel = False
         coro = self._coro
         self._fut_waiter = None
@@ -293,11 +292,9 @@ def __step(self, exc=None):
             else:
                 super().set_result(exc.value)
         except exceptions.CancelledError as exc:
-            if exc.args:
-                cancel_msg = exc.args[0]
-            else:
-                cancel_msg = None
-            super().cancel(msg=cancel_msg)  # I.e., Future.cancel(self).
+            # Save the original exception so we can chain it later.
+            self._cancelled_exc = exc
+            super().cancel()  # I.e., Future.cancel(self).
         except (KeyboardInterrupt, SystemExit) as exc:
             super().set_exception(exc)
             raise
@@ -787,8 +784,7 @@ def _done_callback(fut):
                 # Check if 'fut' is cancelled first, as
                 # 'fut.exception()' will *raise* a CancelledError
                 # instead of returning it.
-                exc = exceptions.CancelledError(''
-                    if fut._cancel_message is None else fut._cancel_message)
+                exc = fut._make_cancelled_error()
                 outer.set_exception(exc)
                 return
             else:
@@ -804,9 +800,12 @@ def _done_callback(fut):
 
             for fut in children:
                 if fut.cancelled():
-                    # Check if 'fut' is cancelled first, as
-                    # 'fut.exception()' will *raise* a CancelledError
-                    # instead of returning it.
+                    # Check if 'fut' is cancelled first, as 'fut.exception()'
+                    # will *raise* a CancelledError instead of returning it.
+                    # Also, since we're adding the exception return value
+                    # to 'results' instead of raising it, don't bother
+                    # setting __context__.  This also lets us preserve
+                    # calling '_make_cancelled_error()' at most once.
                     res = exceptions.CancelledError(
                         '' if fut._cancel_message is None else
                         fut._cancel_message)
@@ -820,8 +819,7 @@ def _done_callback(fut):
                 # If gather is being cancelled we must propagate the
                 # cancellation regardless of *return_exceptions* argument.
                 # See issue 32684.
-                exc = exceptions.CancelledError(''
-                    if fut._cancel_message is None else fut._cancel_message)
+                exc = fut._make_cancelled_error()
                 outer.set_exception(exc)
             else:
                 outer.set_result(results)

Lib/test/test_asyncio/test_tasks.py

@@ -10,6 +10,7 @@
 import re
 import sys
 import textwrap
+import traceback
 import types
 import unittest
 import weakref
@@ -57,6 +58,22 @@ def format_coroutine(qualname, state, src, source_traceback, generator=False):
         return 'coro=<%s() %s at %s>' % (qualname, state, src)
 
 
+def get_innermost_context(exc):
+    """
+    Return information about the innermost exception context in the chain.
+    """
+    depth = 0
+    while True:
+        context = exc.__context__
+        if context is None:
+            break
+
+        exc = context
+        depth += 1
+
+    return (type(exc), exc.args, depth)
+
+
 class Dummy:
 
     def __repr__(self):
@@ -111,9 +128,10 @@ async def coro():
         self.assertEqual(t._cancel_message, None)
 
         t.cancel('my message')
+        self.assertEqual(t._cancel_message, 'my message')
+
         with self.assertRaises(asyncio.CancelledError):
             self.loop.run_until_complete(t)
-        self.assertEqual(t._cancel_message, 'my message')
 
     def test_task_cancel_message_setter(self):
         async def coro():
@@ -123,10 +141,8 @@ async def coro():
         t._cancel_message = 'my new message'
         self.assertEqual(t._cancel_message, 'my new message')
 
-        # Also check that the value is used for cancel().
         with self.assertRaises(asyncio.CancelledError):
             self.loop.run_until_complete(t)
-        self.assertEqual(t._cancel_message, 'my new message')
 
     def test_task_del_collect(self):
         class Evil:
@@ -548,8 +564,8 @@ async def task():
     def test_cancel_with_message_then_future_result(self):
         # Test Future.result() after calling cancel() with a message.
         cases = [
-            ((), ('',)),
-            ((None,), ('',)),
+            ((), ()),
+            ((None,), ()),
             (('my message',), ('my message',)),
             # Non-string values should roundtrip.
             ((5,), (5,)),
@@ -573,13 +589,17 @@ async def coro():
                 with self.assertRaises(asyncio.CancelledError) as cm:
                     loop.run_until_complete(task)
                 exc = cm.exception
-                self.assertEqual(exc.args, expected_args)
+                self.assertEqual(exc.args, ())
+
+                actual = get_innermost_context(exc)
+                self.assertEqual(actual,
+                    (asyncio.CancelledError, expected_args, 2))
 
     def test_cancel_with_message_then_future_exception(self):
         # Test Future.exception() after calling cancel() with a message.
         cases = [
-            ((), ('',)),
-            ((None,), ('',)),
+            ((), ()),
+            ((None,), ()),
             (('my message',), ('my message',)),
             # Non-string values should roundtrip.
             ((5,), (5,)),
@@ -603,7 +623,11 @@ async def coro():
                 with self.assertRaises(asyncio.CancelledError) as cm:
                     loop.run_until_complete(task)
                 exc = cm.exception
-                self.assertEqual(exc.args, expected_args)
+                self.assertEqual(exc.args, ())
+
+                actual = get_innermost_context(exc)
+                self.assertEqual(actual,
+                    (asyncio.CancelledError, expected_args, 2))
 
     def test_cancel_with_message_before_starting_task(self):
         loop = asyncio.new_event_loop()
@@ -623,7 +647,11 @@ async def coro():
         with self.assertRaises(asyncio.CancelledError) as cm:
             loop.run_until_complete(task)
         exc = cm.exception
-        self.assertEqual(exc.args, ('my message',))
+        self.assertEqual(exc.args, ())
+
+        actual = get_innermost_context(exc)
+        self.assertEqual(actual,
+            (asyncio.CancelledError, ('my message',), 2))
 
     def test_cancel_yield(self):
         with self.assertWarns(DeprecationWarning):
@@ -805,6 +833,66 @@ async def coro():
         self.assertTrue(nested_task.cancelled())
         self.assertTrue(fut.cancelled())
 
+    def assert_text_contains(self, text, substr):
+        if substr not in text:
+            raise RuntimeError(f'text {substr!r} not found in:\n>>>{text}<<<')
+
+    def test_cancel_traceback_for_future_result(self):
+        # When calling Future.result() on a cancelled task, check that the
+        # line of code that was interrupted is included in the traceback.
+        loop = asyncio.new_event_loop()
+        self.set_event_loop(loop)
+
+        async def nested():
+            # This will get cancelled immediately.
+            await asyncio.sleep(10)
+
+        async def coro():
+            task = self.new_task(loop, nested())
+            await asyncio.sleep(0)
+            task.cancel()
+            await task  # search target
+
+        task = self.new_task(loop, coro())
+        try:
+            loop.run_until_complete(task)
+        except asyncio.CancelledError:
+            tb = traceback.format_exc()
+            self.assert_text_contains(tb, "await asyncio.sleep(10)")
+            # The intermediate await should also be included.
+            self.assert_text_contains(tb, "await task  # search target")
+        else:
+            self.fail('CancelledError did not occur')
+
+    def test_cancel_traceback_for_future_exception(self):
+        # When calling Future.exception() on a cancelled task, check that the
+        # line of code that was interrupted is included in the traceback.
+        loop = asyncio.new_event_loop()
+        self.set_event_loop(loop)
+
+        async def nested():
+            # This will get cancelled immediately.
+            await asyncio.sleep(10)
+
+        async def coro():
+            task = self.new_task(loop, nested())
+            await asyncio.sleep(0)
+            task.cancel()
+            done, pending = await asyncio.wait([task])
+            task.exception()  # search target
+
+        task = self.new_task(loop, coro())
+        try:
+            loop.run_until_complete(task)
+        except asyncio.CancelledError:
+            tb = traceback.format_exc()
+            self.assert_text_contains(tb, "await asyncio.sleep(10)")
+            # The intermediate await should also be included.
+            self.assert_text_contains(tb,
+                "task.exception()  # search target")
+        else:
+            self.fail('CancelledError did not occur')
+
     def test_stop_while_run_in_complete(self):
 
         def gen():
@@ -2391,15 +2479,14 @@ def cancelling_callback(_):
 
     def test_cancel_gather_2(self):
         cases = [
-            ((), ('',)),
-            ((None,), ('',)),
+            ((), ()),
+            ((None,), ()),
             (('my message',), ('my message',)),
             # Non-string values should roundtrip.
             ((5,), (5,)),
         ]
         for cancel_args, expected_args in cases:
             with self.subTest(cancel_args=cancel_args):
-
                 loop = asyncio.new_event_loop()
                 self.addCleanup(loop.close)
 
@@ -2417,15 +2504,20 @@ async def main():
                     qwe = self.new_task(loop, test())
                     await asyncio.sleep(0.2)
                     qwe.cancel(*cancel_args)
-                    try:
-                        await qwe
-                    except asyncio.CancelledError as exc:
-                        self.assertEqual(exc.args, expected_args)
-                    else:
-                        self.fail('gather did not propagate the cancellation '
-                                  'request')
-
-                loop.run_until_complete(main())
+                    await qwe
+
+                try:
+                    loop.run_until_complete(main())
+                except asyncio.CancelledError as exc:
+                    self.assertEqual(exc.args, ())
+                    exc_type, exc_args, depth = get_innermost_context(exc)
+                    self.assertEqual((exc_type, exc_args),
+                        (asyncio.CancelledError, expected_args))
+                    # The exact traceback seems to vary in CI.
+                    self.assertIn(depth, (2, 3))
+                else:
+                    self.fail('gather did not propagate the cancellation '
+                              'request')
 
     def test_exception_traceback(self):
         # See http://bugs.python.org/issue28843

Misc/NEWS.d/next/Library/2020-05-06-02-33-00.bpo-31033.aX12pw.rst

@@ -0,0 +1,2 @@
+When a :class:`asyncio.Task` is cancelled, the exception traceback
+now chains all the way back to where the task was first interrupted.