bpo-40696: Fix a hang that can arise after gen.throw() (GH-20287)

This updates _PyErr_ChainStackItem() to use _PyErr_SetObject()
instead of _PyErr_ChainExceptions(). This prevents a hang in
certain circumstances because _PyErr_SetObject() performs checks
to prevent cycles in the exception context chain while
_PyErr_ChainExceptions() doesn't.

7 files changed, 130 insertions, 23 deletions

diff --git a/Include/internal/pycore_pyerrors.h b/Include/internal/pycore_pyerrors.h
index 3290a37..2cf1160 100644
--- a/Include/internal/pycore_pyerrors.h
+++ b/Include/internal/pycore_pyerrors.h
@@ -51,7 +51,7 @@ PyAPI_FUNC(void) _PyErr_SetObject(
     PyObject *value);
 
 PyAPI_FUNC(void) _PyErr_ChainStackItem(
-    _PyErr_StackItem *exc_state);
+    _PyErr_StackItem *exc_info);
 
 PyAPI_FUNC(void) _PyErr_Clear(PyThreadState *tstate);
 
diff --git a/Lib/test/test_asyncio/test_tasks.py b/Lib/test/test_asyncio/test_tasks.py
index 63968e2..3734013 100644
--- a/Lib/test/test_asyncio/test_tasks.py
+++ b/Lib/test/test_asyncio/test_tasks.py
@@ -536,9 +536,42 @@ class BaseTaskTests:
                     self.assertEqual((type(chained), chained.args),
                         (KeyError, (3,)))
 
-        task = self.new_task(loop, run())
-        loop.run_until_complete(task)
-        loop.close()
+        try:
+            task = self.new_task(loop, run())
+            loop.run_until_complete(task)
+        finally:
+            loop.close()
+
+    def test_exception_chaining_after_await_with_context_cycle(self):
+        # Check trying to create an exception context cycle:
+        # https://bugs.python.org/issue40696
+        has_cycle = None
+        loop = asyncio.new_event_loop()
+        self.set_event_loop(loop)
+
+        async def process_exc(exc):
+            raise exc
+
+        async def run():
+            nonlocal has_cycle
+            try:
+                raise KeyError('a')
+            except Exception as exc:
+                task = self.new_task(loop, process_exc(exc))
+                try:
+                    await task
+                except BaseException as exc:
+                    has_cycle = (exc is exc.__context__)
+                    # Prevent a hang if has_cycle is True.
+                    exc.__context__ = None
+
+        try:
+            task = self.new_task(loop, run())
+            loop.run_until_complete(task)
+        finally:
+            loop.close()
+        # This also distinguishes from the initial has_cycle=None.
+        self.assertEqual(has_cycle, False)
 
     def test_cancel(self):
 
diff --git a/Lib/test/test_generators.py b/Lib/test/test_generators.py
index 87cc2df..bf48221 100644
--- a/Lib/test/test_generators.py
+++ b/Lib/test/test_generators.py
@@ -371,6 +371,32 @@ class GeneratorThrowTest(unittest.TestCase):
         context = cm.exception.__context__
         self.assertEqual((type(context), context.args), (KeyError, ('a',)))
 
+    def test_exception_context_with_yield_from_with_context_cycle(self):
+        # Check trying to create an exception context cycle:
+        # https://bugs.python.org/issue40696
+        has_cycle = None
+
+        def f():
+            yield
+
+        def g(exc):
+            nonlocal has_cycle
+            try:
+                raise exc
+            except Exception:
+                try:
+                    yield from f()
+                except Exception as exc:
+                    has_cycle = (exc is exc.__context__)
+            yield
+
+        exc = KeyError('a')
+        gen = g(exc)
+        gen.send(None)
+        gen.throw(exc)
+        # This also distinguishes from the initial has_cycle=None.
+        self.assertEqual(has_cycle, False)
+
     def test_throw_after_none_exc_type(self):
         def g():
             try:
diff --git a/Misc/NEWS.d/next/Core and Builtins/2020-05-21-01-54-00.bpo-40696.u3n8Wx.rst b/Misc/NEWS.d/next/Core and Builtins/2020-05-21-01-54-00.bpo-40696.u3n8Wx.rst
new file mode 100644
index 0000000..f99bdea
--- /dev/null
+++ b/Misc/NEWS.d/next/Core and Builtins/2020-05-21-01-54-00.bpo-40696.u3n8Wx.rst
@@ -0,0 +1,2 @@
+Fix a hang that can arise after :meth:`generator.throw` due to a cycle
+in the exception context chain.
diff --git a/Modules/_asynciomodule.c b/Modules/_asynciomodule.c
index 1b6a579..0608c40 100644
--- a/Modules/_asynciomodule.c
+++ b/Modules/_asynciomodule.c
@@ -612,19 +612,20 @@ create_cancelled_error(PyObject *msg)
 }
 
 static void
-set_cancelled_error(PyObject *msg)
+future_set_cancelled_error(FutureObj *fut)
 {
-    PyObject *exc = create_cancelled_error(msg);
+    PyObject *exc = create_cancelled_error(fut->fut_cancel_msg);
     PyErr_SetObject(asyncio_CancelledError, exc);
     Py_DECREF(exc);
+
+    _PyErr_ChainStackItem(&fut->fut_cancelled_exc_state);
 }
 
 static int
 future_get_result(FutureObj *fut, PyObject **result)
 {
     if (fut->fut_state == STATE_CANCELLED) {
-        set_cancelled_error(fut->fut_cancel_msg);
-        _PyErr_ChainStackItem(&fut->fut_cancelled_exc_state);
+        future_set_cancelled_error(fut);
         return -1;
     }
 
@@ -866,8 +867,7 @@ _asyncio_Future_exception_impl(FutureObj *self)
     }
 
     if (self->fut_state == STATE_CANCELLED) {
-        set_cancelled_error(self->fut_cancel_msg);
-        _PyErr_ChainStackItem(&self->fut_cancelled_exc_state);
+        future_set_cancelled_error(self);
         return NULL;
     }
 
diff --git a/Objects/genobject.c b/Objects/genobject.c
index 271720b..09efbab 100644
--- a/Objects/genobject.c
+++ b/Objects/genobject.c
@@ -203,13 +203,15 @@ gen_send_ex(PyGenObject *gen, PyObject *arg, int exc, int closing)
     assert(f->f_back == NULL);
     f->f_back = tstate->frame;
 
-    if (exc) {
-        _PyErr_ChainStackItem(&gen->gi_exc_state);
-    }
-
     gen->gi_running = 1;
     gen->gi_exc_state.previous_item = tstate->exc_info;
     tstate->exc_info = &gen->gi_exc_state;
+
+    if (exc) {
+        assert(_PyErr_Occurred(tstate));
+        _PyErr_ChainStackItem(NULL);
+    }
+
     result = _PyEval_EvalFrame(tstate, f, exc);
     tstate->exc_info = gen->gi_exc_state.previous_item;
     gen->gi_exc_state.previous_item = NULL;
diff --git a/Python/errors.c b/Python/errors.c
index 3b42c11..70365aa 100644
--- a/Python/errors.c
+++ b/Python/errors.c
@@ -477,7 +477,9 @@ PyErr_SetExcInfo(PyObject *p_type, PyObject *p_value, PyObject *p_traceback)
 
 /* Like PyErr_Restore(), but if an exception is already set,
    set the context associated with it.
- */
+
+   The caller is responsible for ensuring that this call won't create
+   any cycles in the exception context chain. */
 void
 _PyErr_ChainExceptions(PyObject *exc, PyObject *val, PyObject *tb)
 {
@@ -512,18 +514,60 @@ _PyErr_ChainExceptions(PyObject *exc, PyObject *val, PyObject *tb)
     }
 }
 
+/* Set the currently set exception's context to the given exception.
+
+   If the provided exc_info is NULL, then the current Python thread state's
+   exc_info will be used for the context instead.
+
+   This function can only be called when _PyErr_Occurred() is true.
+   Also, this function won't create any cycles in the exception context
+   chain to the extent that _PyErr_SetObject ensures this. */
 void
-_PyErr_ChainStackItem(_PyErr_StackItem *exc_state)
+_PyErr_ChainStackItem(_PyErr_StackItem *exc_info)
 {
-    if (exc_state->exc_type == NULL || exc_state->exc_type == Py_None) {
+    PyThreadState *tstate = _PyThreadState_GET();
+    assert(_PyErr_Occurred(tstate));
+
+    int exc_info_given;
+    if (exc_info == NULL) {
+        exc_info_given = 0;
+        exc_info = tstate->exc_info;
+    } else {
+        exc_info_given = 1;
+    }
+    if (exc_info->exc_type == NULL || exc_info->exc_type == Py_None) {
         return;
     }
-    Py_INCREF(exc_state->exc_type);
-    Py_XINCREF(exc_state->exc_value);
-    Py_XINCREF(exc_state->exc_traceback);
-    _PyErr_ChainExceptions(exc_state->exc_type,
-                           exc_state->exc_value,
-                           exc_state->exc_traceback);
+
+    _PyErr_StackItem *saved_exc_info;
+    if (exc_info_given) {
+        /* Temporarily set the thread state's exc_info since this is what
+           _PyErr_SetObject uses for implicit exception chaining. */
+        saved_exc_info = tstate->exc_info;
+        tstate->exc_info = exc_info;
+    }
+
+    PyObject *exc, *val, *tb;
+    _PyErr_Fetch(tstate, &exc, &val, &tb);
+
+    PyObject *exc2, *val2, *tb2;
+    exc2 = exc_info->exc_type;
+    val2 = exc_info->exc_value;
+    tb2 = exc_info->exc_traceback;
+    _PyErr_NormalizeException(tstate, &exc2, &val2, &tb2);
+    if (tb2 != NULL) {
+        PyException_SetTraceback(val2, tb2);
+    }
+
+    /* _PyErr_SetObject sets the context from PyThreadState. */
+    _PyErr_SetObject(tstate, exc, val);
+    Py_DECREF(exc);  // since _PyErr_Occurred was true
+    Py_XDECREF(val);
+    Py_XDECREF(tb);
+
+    if (exc_info_given) {
+        tstate->exc_info = saved_exc_info;
+    }
 }
 
 static PyObject *