[3.8] bpo-39606: allow closing async generators that are already closed (GH-18475) (GH-18501)

The fix for [bpo-39386](https://bugs.python.org/issue39386) attempted to make it so you couldn't reuse a
agen.aclose() coroutine object. It accidentally also prevented you
from calling aclose() at all on an async generator that was already
closed or exhausted. This commit fixes it so we're only blocking the
actually illegal cases, while allowing the legal cases.

The new tests failed before this patch. Also confirmed that this fixes
the test failures we were seeing in Trio with Python dev builds:
  https://github.com/python-trio/trio/pull/1396


https://bugs.python.org/issue39606
(cherry picked from commit 925dc7fb1d0db85dc137afa4cd14211bf0d67414)


Co-authored-by: Nathaniel J. Smith <njs@pobox.com>


https://bugs.python.org/issue39606



Automerge-Triggered-By: @njsmith

1 files changed, 11 insertions, 4 deletions

diff --git a/Objects/genobject.c b/Objects/genobject.c
index 72aa872..2c06bdc 100644
--- a/Objects/genobject.c
+++ b/Objects/genobject.c
@@ -1812,16 +1812,22 @@ async_gen_athrow_send(PyAsyncGenAThrow *o, PyObject *arg)
     PyFrameObject *f = gen->gi_frame;
     PyObject *retval;
 
-    if (f == NULL || f->f_stacktop == NULL ||
-            o->agt_state == AWAITABLE_STATE_CLOSED) {
+    if (o->agt_state == AWAITABLE_STATE_CLOSED) {
         PyErr_SetString(
             PyExc_RuntimeError,
             "cannot reuse already awaited aclose()/athrow()");
         return NULL;
     }
 
+    if (f == NULL || f->f_stacktop == NULL) {
+        o->agt_state = AWAITABLE_STATE_CLOSED;
+        PyErr_SetNone(PyExc_StopIteration);
+        return NULL;
+    }
+
     if (o->agt_state == AWAITABLE_STATE_INIT) {
         if (o->agt_gen->ag_running_async) {
+            o->agt_state = AWAITABLE_STATE_CLOSED;
             if (o->agt_args == NULL) {
                 PyErr_SetString(
                     PyExc_RuntimeError,
@@ -1893,7 +1899,6 @@ async_gen_athrow_send(PyAsyncGenAThrow *o, PyObject *arg)
         /* aclose() mode */
         if (retval) {
             if (_PyAsyncGenWrappedValue_CheckExact(retval)) {
-                o->agt_gen->ag_running_async = 0;
                 Py_DECREF(retval);
                 goto yield_close;
             }
@@ -1908,16 +1913,17 @@ async_gen_athrow_send(PyAsyncGenAThrow *o, PyObject *arg)
 
 yield_close:
     o->agt_gen->ag_running_async = 0;
+    o->agt_state = AWAITABLE_STATE_CLOSED;
     PyErr_SetString(
         PyExc_RuntimeError, ASYNC_GEN_IGNORED_EXIT_MSG);
     return NULL;
 
 check_error:
     o->agt_gen->ag_running_async = 0;
+    o->agt_state = AWAITABLE_STATE_CLOSED;
     if (PyErr_ExceptionMatches(PyExc_StopAsyncIteration) ||
             PyErr_ExceptionMatches(PyExc_GeneratorExit))
     {
-        o->agt_state = AWAITABLE_STATE_CLOSED;
         if (o->agt_args == NULL) {
             /* when aclose() is called we don't want to propagate
                StopAsyncIteration or GeneratorExit; just raise
@@ -1951,6 +1957,7 @@ async_gen_athrow_throw(PyAsyncGenAThrow *o, PyObject *args)
         /* aclose() mode */
         if (retval && _PyAsyncGenWrappedValue_CheckExact(retval)) {
             o->agt_gen->ag_running_async = 0;
+            o->agt_state = AWAITABLE_STATE_CLOSED;
             Py_DECREF(retval);
             PyErr_SetString(PyExc_RuntimeError, ASYNC_GEN_IGNORED_EXIT_MSG);
             return NULL;