Skip to content

Commit da742ba

Browse files
cjerdonekcjerdonek
authored
bpo-31033: Improve the traceback for cancelled asyncio tasks (pythonGH-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."
1 parent d17f3d8 commit da742ba

File tree

10 files changed

+291
-81
lines changed

10 files changed

+291
-81
lines changed

‎Include/internal/pycore_pyerrors.h‎

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,20 @@ static inline PyObject* _PyErr_Occurred(PyThreadState *tstate)
1414
return tstate->curexc_type;
1515
}
1616

17+
static inline void _PyErr_ClearExcState(_PyErr_StackItem *exc_state)
18+
{
19+
PyObject *t, *v, *tb;
20+
t = exc_state->exc_type;
21+
v = exc_state->exc_value;
22+
tb = exc_state->exc_traceback;
23+
exc_state->exc_type = NULL;
24+
exc_state->exc_value = NULL;
25+
exc_state->exc_traceback = NULL;
26+
Py_XDECREF(t);
27+
Py_XDECREF(v);
28+
Py_XDECREF(tb);
29+
}
30+
1731

1832
PyAPI_FUNC(void) _PyErr_Fetch(
1933
PyThreadState *tstate,
@@ -36,6 +50,9 @@ PyAPI_FUNC(void) _PyErr_SetObject(
3650
PyObject *type,
3751
PyObject *value);
3852

53+
PyAPI_FUNC(void) _PyErr_ChainStackItem(
54+
_PyErr_StackItem *exc_state);
55+
3956
PyAPI_FUNC(void) _PyErr_Clear(PyThreadState *tstate);
4057

4158
PyAPI_FUNC(void) _PyErr_SetNone(PyThreadState *tstate, PyObject *exception);

‎Lib/asyncio/futures.py‎

Lines changed: 21 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -52,6 +52,8 @@ class Future:
5252
_loop = None
5353
_source_traceback = None
5454
_cancel_message = None
55+
# A saved CancelledError for later chaining as an exception context.
56+
_cancelled_exc = None
5557

5658
# This field is used for a dual purpose:
5759
# - Its presence is a marker to declare that a class implements
@@ -124,6 +126,21 @@ def get_loop(self):
124126
raise RuntimeError("Future object is not initialized.")
125127
return loop
126128

129+
def _make_cancelled_error(self):
130+
"""Create the CancelledError to raise if the Future is cancelled.
131+
132+
This should only be called once when handling a cancellation since
133+
it erases the saved context exception value.
134+
"""
135+
if self._cancel_message is None:
136+
exc = exceptions.CancelledError()
137+
else:
138+
exc = exceptions.CancelledError(self._cancel_message)
139+
exc.__context__ = self._cancelled_exc
140+
# Remove the reference since we don't need this anymore.
141+
self._cancelled_exc = None
142+
return exc
143+
127144
def cancel(self, msg=None):
128145
"""Cancel the future and schedule callbacks.
129146
@@ -175,9 +192,8 @@ def result(self):
175192
the future is done and has an exception set, this exception is raised.
176193
"""
177194
if self._state == _CANCELLED:
178-
raise exceptions.CancelledError(
179-
'' if self._cancel_message is None else self._cancel_message)
180-
195+
exc = self._make_cancelled_error()
196+
raise exc
181197
if self._state != _FINISHED:
182198
raise exceptions.InvalidStateError('Result is not ready.')
183199
self.__log_traceback = False
@@ -194,8 +210,8 @@ def exception(self):
194210
InvalidStateError.
195211
"""
196212
if self._state == _CANCELLED:
197-
raise exceptions.CancelledError(
198-
'' if self._cancel_message is None else self._cancel_message)
213+
exc = self._make_cancelled_error()
214+
raise exc
199215
if self._state != _FINISHED:
200216
raise exceptions.InvalidStateError('Exception is not set.')
201217
self.__log_traceback = False

‎Lib/asyncio/tasks.py‎

Lines changed: 12 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -270,8 +270,7 @@ def __step(self, exc=None):
270270
f'_step(): already done: {self!r}, {exc!r}')
271271
if self._must_cancel:
272272
if not isinstance(exc, exceptions.CancelledError):
273-
exc = exceptions.CancelledError(''
274-
if self._cancel_message is None else self._cancel_message)
273+
exc = self._make_cancelled_error()
275274
self._must_cancel = False
276275
coro = self._coro
277276
self._fut_waiter = None
@@ -293,11 +292,9 @@ def __step(self, exc=None):
293292
else:
294293
super().set_result(exc.value)
295294
except exceptions.CancelledError as exc:
296-
if exc.args:
297-
cancel_msg = exc.args[0]
298-
else:
299-
cancel_msg = None
300-
super().cancel(msg=cancel_msg) # I.e., Future.cancel(self).
295+
# Save the original exception so we can chain it later.
296+
self._cancelled_exc = exc
297+
super().cancel() # I.e., Future.cancel(self).
301298
except (KeyboardInterrupt, SystemExit) as exc:
302299
super().set_exception(exc)
303300
raise
@@ -787,8 +784,7 @@ def _done_callback(fut):
787784
# Check if 'fut' is cancelled first, as
788785
# 'fut.exception()' will *raise* a CancelledError
789786
# instead of returning it.
790-
exc = exceptions.CancelledError(''
791-
if fut._cancel_message is None else fut._cancel_message)
787+
exc = fut._make_cancelled_error()
792788
outer.set_exception(exc)
793789
return
794790
else:
@@ -804,9 +800,12 @@ def _done_callback(fut):
804800

805801
for fut in children:
806802
if fut.cancelled():
807-
# Check if 'fut' is cancelled first, as
808-
# 'fut.exception()' will *raise* a CancelledError
809-
# instead of returning it.
803+
# Check if 'fut' is cancelled first, as 'fut.exception()'
804+
# will *raise* a CancelledError instead of returning it.
805+
# Also, since we're adding the exception return value
806+
# to 'results' instead of raising it, don't bother
807+
# setting __context__. This also lets us preserve
808+
# calling '_make_cancelled_error()' at most once.
810809
res = exceptions.CancelledError(
811810
'' if fut._cancel_message is None else
812811
fut._cancel_message)
@@ -820,8 +819,7 @@ def _done_callback(fut):
820819
# If gather is being cancelled we must propagate the
821820
# cancellation regardless of *return_exceptions* argument.
822821
# See issue 32684.
823-
exc = exceptions.CancelledError(''
824-
if fut._cancel_message is None else fut._cancel_message)
822+
exc = fut._make_cancelled_error()
825823
outer.set_exception(exc)
826824
else:
827825
outer.set_result(results)

‎Lib/test/test_asyncio/test_tasks.py‎

Lines changed: 114 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@
1010
import re
1111
import sys
1212
import textwrap
13+
import traceback
1314
import types
1415
import unittest
1516
import weakref
@@ -57,6 +58,22 @@ def format_coroutine(qualname, state, src, source_traceback, generator=False):
5758
return 'coro=<%s() %s at %s>' % (qualname, state, src)
5859

5960

61+
def get_innermost_context(exc):
62+
"""
63+
Return information about the innermost exception context in the chain.
64+
"""
65+
depth = 0
66+
while True:
67+
context = exc.__context__
68+
if context is None:
69+
break
70+
71+
exc = context
72+
depth += 1
73+
74+
return (type(exc), exc.args, depth)
75+
76+
6077
class Dummy:
6178

6279
def __repr__(self):
@@ -111,9 +128,10 @@ async def coro():
111128
self.assertEqual(t._cancel_message, None)
112129

113130
t.cancel('my message')
131+
self.assertEqual(t._cancel_message, 'my message')
132+
114133
with self.assertRaises(asyncio.CancelledError):
115134
self.loop.run_until_complete(t)
116-
self.assertEqual(t._cancel_message, 'my message')
117135

118136
def test_task_cancel_message_setter(self):
119137
async def coro():
@@ -123,10 +141,8 @@ async def coro():
123141
t._cancel_message = 'my new message'
124142
self.assertEqual(t._cancel_message, 'my new message')
125143

126-
# Also check that the value is used for cancel().
127144
with self.assertRaises(asyncio.CancelledError):
128145
self.loop.run_until_complete(t)
129-
self.assertEqual(t._cancel_message, 'my new message')
130146

131147
def test_task_del_collect(self):
132148
class Evil:
@@ -548,8 +564,8 @@ async def task():
548564
def test_cancel_with_message_then_future_result(self):
549565
# Test Future.result() after calling cancel() with a message.
550566
cases = [
551-
((), ('',)),
552-
((None,), ('',)),
567+
((), ()),
568+
((None,), ()),
553569
(('my message',), ('my message',)),
554570
# Non-string values should roundtrip.
555571
((5,), (5,)),
@@ -573,13 +589,17 @@ async def coro():
573589
with self.assertRaises(asyncio.CancelledError) as cm:
574590
loop.run_until_complete(task)
575591
exc = cm.exception
576-
self.assertEqual(exc.args, expected_args)
592+
self.assertEqual(exc.args, ())
593+
594+
actual = get_innermost_context(exc)
595+
self.assertEqual(actual,
596+
(asyncio.CancelledError, expected_args, 2))
577597

578598
def test_cancel_with_message_then_future_exception(self):
579599
# Test Future.exception() after calling cancel() with a message.
580600
cases = [
581-
((), ('',)),
582-
((None,), ('',)),
601+
((), ()),
602+
((None,), ()),
583603
(('my message',), ('my message',)),
584604
# Non-string values should roundtrip.
585605
((5,), (5,)),
@@ -603,7 +623,11 @@ async def coro():
603623
with self.assertRaises(asyncio.CancelledError) as cm:
604624
loop.run_until_complete(task)
605625
exc = cm.exception
606-
self.assertEqual(exc.args, expected_args)
626+
self.assertEqual(exc.args, ())
627+
628+
actual = get_innermost_context(exc)
629+
self.assertEqual(actual,
630+
(asyncio.CancelledError, expected_args, 2))
607631

608632
def test_cancel_with_message_before_starting_task(self):
609633
loop = asyncio.new_event_loop()
@@ -623,7 +647,11 @@ async def coro():
623647
with self.assertRaises(asyncio.CancelledError) as cm:
624648
loop.run_until_complete(task)
625649
exc = cm.exception
626-
self.assertEqual(exc.args, ('my message',))
650+
self.assertEqual(exc.args, ())
651+
652+
actual = get_innermost_context(exc)
653+
self.assertEqual(actual,
654+
(asyncio.CancelledError, ('my message',), 2))
627655

628656
def test_cancel_yield(self):
629657
with self.assertWarns(DeprecationWarning):
@@ -805,6 +833,66 @@ async def coro():
805833
self.assertTrue(nested_task.cancelled())
806834
self.assertTrue(fut.cancelled())
807835

836+
def assert_text_contains(self, text, substr):
837+
if substr not in text:
838+
raise RuntimeError(f'text {substr!r} not found in:\n>>>{text}<<<')
839+
840+
def test_cancel_traceback_for_future_result(self):
841+
# When calling Future.result() on a cancelled task, check that the
842+
# line of code that was interrupted is included in the traceback.
843+
loop = asyncio.new_event_loop()
844+
self.set_event_loop(loop)
845+
846+
async def nested():
847+
# This will get cancelled immediately.
848+
await asyncio.sleep(10)
849+
850+
async def coro():
851+
task = self.new_task(loop, nested())
852+
await asyncio.sleep(0)
853+
task.cancel()
854+
await task # search target
855+
856+
task = self.new_task(loop, coro())
857+
try:
858+
loop.run_until_complete(task)
859+
except asyncio.CancelledError:
860+
tb = traceback.format_exc()
861+
self.assert_text_contains(tb, "await asyncio.sleep(10)")
862+
# The intermediate await should also be included.
863+
self.assert_text_contains(tb, "await task # search target")
864+
else:
865+
self.fail('CancelledError did not occur')
866+
867+
def test_cancel_traceback_for_future_exception(self):
868+
# When calling Future.exception() on a cancelled task, check that the
869+
# line of code that was interrupted is included in the traceback.
870+
loop = asyncio.new_event_loop()
871+
self.set_event_loop(loop)
872+
873+
async def nested():
874+
# This will get cancelled immediately.
875+
await asyncio.sleep(10)
876+
877+
async def coro():
878+
task = self.new_task(loop, nested())
879+
await asyncio.sleep(0)
880+
task.cancel()
881+
done, pending = await asyncio.wait([task])
882+
task.exception() # search target
883+
884+
task = self.new_task(loop, coro())
885+
try:
886+
loop.run_until_complete(task)
887+
except asyncio.CancelledError:
888+
tb = traceback.format_exc()
889+
self.assert_text_contains(tb, "await asyncio.sleep(10)")
890+
# The intermediate await should also be included.
891+
self.assert_text_contains(tb,
892+
"task.exception() # search target")
893+
else:
894+
self.fail('CancelledError did not occur')
895+
808896
def test_stop_while_run_in_complete(self):
809897

810898
def gen():
@@ -2391,15 +2479,14 @@ def cancelling_callback(_):
23912479

23922480
def test_cancel_gather_2(self):
23932481
cases = [
2394-
((), ('',)),
2395-
((None,), ('',)),
2482+
((), ()),
2483+
((None,), ()),
23962484
(('my message',), ('my message',)),
23972485
# Non-string values should roundtrip.
23982486
((5,), (5,)),
23992487
]
24002488
for cancel_args, expected_args in cases:
24012489
with self.subTest(cancel_args=cancel_args):
2402-
24032490
loop = asyncio.new_event_loop()
24042491
self.addCleanup(loop.close)
24052492

@@ -2417,15 +2504,20 @@ async def main():
24172504
qwe = self.new_task(loop, test())
24182505
await asyncio.sleep(0.2)
24192506
qwe.cancel(*cancel_args)
2420-
try:
2421-
await qwe
2422-
except asyncio.CancelledError as exc:
2423-
self.assertEqual(exc.args, expected_args)
2424-
else:
2425-
self.fail('gather did not propagate the cancellation '
2426-
'request')
2427-
2428-
loop.run_until_complete(main())
2507+
await qwe
2508+
2509+
try:
2510+
loop.run_until_complete(main())
2511+
except asyncio.CancelledError as exc:
2512+
self.assertEqual(exc.args, ())
2513+
exc_type, exc_args, depth = get_innermost_context(exc)
2514+
self.assertEqual((exc_type, exc_args),
2515+
(asyncio.CancelledError, expected_args))
2516+
# The exact traceback seems to vary in CI.
2517+
self.assertIn(depth, (2, 3))
2518+
else:
2519+
self.fail('gather did not propagate the cancellation '
2520+
'request')
24292521

24302522
def test_exception_traceback(self):
24312523
# See http://bugs.python.org/issue28843

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

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,2 @@
1+
When a :class:`asyncio.Task` is cancelled, the exception traceback
2+
now chains all the way back to where the task was first interrupted.

0 commit comments

Comments
 (0)