bpo-45924: Fix asyncio incorrect traceback when future's exception is raised multiple times (GH-30274)

4 files changed, 47 insertions, 2 deletions

diff --git a/Lib/asyncio/futures.py b/Lib/asyncio/futures.py
index 08c79e7..4bd9629 100644
--- a/Lib/asyncio/futures.py
+++ b/Lib/asyncio/futures.py
@@ -206,7 +206,7 @@ class Future:
             raise exceptions.InvalidStateError('Result is not ready.')
         self.__log_traceback = False
         if self._exception is not None:
-            raise self._exception
+            raise self._exception.with_traceback(self._exception_tb)
         return self._result
 
     def exception(self):
@@ -282,6 +282,7 @@ class Future:
             raise TypeError("StopIteration interacts badly with generators "
                             "and cannot be raised into a Future")
         self._exception = exception
+        self._exception_tb = exception.__traceback__
         self._state = _FINISHED
         self.__schedule_callbacks()
         self.__log_traceback = True
diff --git a/Lib/test/test_asyncio/test_futures2.py b/Lib/test/test_asyncio/test_futures2.py
index 60b5885..71279b6 100644
--- a/Lib/test/test_asyncio/test_futures2.py
+++ b/Lib/test/test_asyncio/test_futures2.py
@@ -1,13 +1,42 @@
 # IsolatedAsyncioTestCase based tests
 import asyncio
+import traceback
 import unittest
+from asyncio import tasks
 
 
 def tearDownModule():
     asyncio.set_event_loop_policy(None)
 
 
-class FutureTests(unittest.IsolatedAsyncioTestCase):
+class FutureTests:
+
+    async def test_future_traceback(self):
+
+        async def raise_exc():
+            raise TypeError(42)
+
+        future = self.cls(raise_exc())
+
+        for _ in range(5):
+            try:
+                await future
+            except TypeError as e:
+                tb = ''.join(traceback.format_tb(e.__traceback__))
+                self.assertEqual(tb.count("await future"), 1)
+            else:
+                self.fail('TypeError was not raised')
+
+@unittest.skipUnless(hasattr(tasks, '_CTask'),
+                       'requires the C _asyncio module')
+class CFutureTests(FutureTests, unittest.IsolatedAsyncioTestCase):
+    cls = tasks._CTask
+
+class PyFutureTests(FutureTests, unittest.IsolatedAsyncioTestCase):
+    cls = tasks._PyTask
+
+class FutureReprTests(unittest.IsolatedAsyncioTestCase):
+
     async def test_recursive_repr_for_pending_tasks(self):
         # The call crashes if the guard for recursive call
         # in base_futures:_future_repr_info is absent
diff --git a/Misc/NEWS.d/next/Library/2021-12-27-15-32-15.bpo-45924.0ZpHX2.rst b/Misc/NEWS.d/next/Library/2021-12-27-15-32-15.bpo-45924.0ZpHX2.rst
new file mode 100644
index 0000000..5cda227
--- /dev/null
+++ b/Misc/NEWS.d/next/Library/2021-12-27-15-32-15.bpo-45924.0ZpHX2.rst
@@ -0,0 +1 @@
+Fix :mod:`asyncio` incorrect traceback when future's exception is raised multiple times. Patch by Kumar Aditya.
diff --git a/Modules/_asynciomodule.c b/Modules/_asynciomodule.c
index bb7c5a7..46175cc 100644
--- a/Modules/_asynciomodule.c
+++ b/Modules/_asynciomodule.c
@@ -70,6 +70,7 @@ typedef enum {
     PyObject *prefix##_context0;                                            \
     PyObject *prefix##_callbacks;                                           \
     PyObject *prefix##_exception;                                           \
+    PyObject *prefix##_exception_tb;                                        \
     PyObject *prefix##_result;                                              \
     PyObject *prefix##_source_tb;                                           \
     PyObject *prefix##_cancel_msg;                                          \
@@ -495,6 +496,7 @@ future_init(FutureObj *fut, PyObject *loop)
     Py_CLEAR(fut->fut_callbacks);
     Py_CLEAR(fut->fut_result);
     Py_CLEAR(fut->fut_exception);
+    Py_CLEAR(fut->fut_exception_tb);
     Py_CLEAR(fut->fut_source_tb);
     Py_CLEAR(fut->fut_cancel_msg);
     Py_CLEAR(fut->fut_cancelled_exc);
@@ -601,7 +603,9 @@ future_set_exception(FutureObj *fut, PyObject *exc)
     }
 
     assert(!fut->fut_exception);
+    assert(!fut->fut_exception_tb);
     fut->fut_exception = exc_val;
+    fut->fut_exception_tb = PyException_GetTraceback(exc_val);
     fut->fut_state = STATE_FINISHED;
 
     if (future_schedule_callbacks(fut) == -1) {
@@ -656,8 +660,16 @@ future_get_result(FutureObj *fut, PyObject **result)
 
     fut->fut_log_tb = 0;
     if (fut->fut_exception != NULL) {
+        PyObject *tb = fut->fut_exception_tb;
+        if (tb == NULL) {
+            tb = Py_None;
+        }
+        if (PyException_SetTraceback(fut->fut_exception, tb) < 0) {
+            return -1;
+        }
         Py_INCREF(fut->fut_exception);
         *result = fut->fut_exception;
+        Py_CLEAR(fut->fut_exception_tb);
         return 1;
     }
 
@@ -799,6 +811,7 @@ FutureObj_clear(FutureObj *fut)
     Py_CLEAR(fut->fut_callbacks);
     Py_CLEAR(fut->fut_result);
     Py_CLEAR(fut->fut_exception);
+    Py_CLEAR(fut->fut_exception_tb);
     Py_CLEAR(fut->fut_source_tb);
     Py_CLEAR(fut->fut_cancel_msg);
     Py_CLEAR(fut->fut_cancelled_exc);
@@ -815,6 +828,7 @@ FutureObj_traverse(FutureObj *fut, visitproc visit, void *arg)
     Py_VISIT(fut->fut_callbacks);
     Py_VISIT(fut->fut_result);
     Py_VISIT(fut->fut_exception);
+    Py_VISIT(fut->fut_exception_tb);
     Py_VISIT(fut->fut_source_tb);
     Py_VISIT(fut->fut_cancel_msg);
     Py_VISIT(fut->fut_cancelled_exc);