From 38a68f24e7fa5c3d8aa8b264d0bc6ef76cd4deeb Mon Sep 17 00:00:00 2001
From: "Nathaniel J. Smith" <njs@pobox.com>
Date: Wed, 12 Feb 2020 00:00:08 -0800
Subject: [PATCH] bpo-39606: allow closing async generators that are already
 closed

The fix for bpo-39386 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
---
 Lib/test/test_asyncgen.py                     | 30 +++++++++++++++++--
 .../2020-02-11-23-59-07.bpo-39606.a72Sxc.rst  |  2 ++
 Objects/genobject.c                           | 15 +++++++---
 3 files changed, 41 insertions(+), 6 deletions(-)
 create mode 100644 Misc/NEWS.d/next/Core and Builtins/2020-02-11-23-59-07.bpo-39606.a72Sxc.rst

diff --git a/Lib/test/test_asyncgen.py b/Lib/test/test_asyncgen.py
index 24b20bec2b2d1e..fb6321d2264f31 100644
--- a/Lib/test/test_asyncgen.py
+++ b/Lib/test/test_asyncgen.py
@@ -1128,7 +1128,7 @@ async def main():
 
         self.assertEqual([], messages)
 
-    def test_async_gen_await_anext_twice(self):
+    def test_async_gen_await_same_anext_coro_twice(self):
         async def async_iterate():
             yield 1
             yield 2
@@ -1147,7 +1147,7 @@ async def run():
 
         self.loop.run_until_complete(run())
 
-    def test_async_gen_await_aclose_twice(self):
+    def test_async_gen_await_same_aclose_coro_twice(self):
         async def async_iterate():
             yield 1
             yield 2
@@ -1164,6 +1164,32 @@ async def run():
 
         self.loop.run_until_complete(run())
 
+    def test_async_gen_aclose_twice_with_different_coros(self):
+        # Regression test for https://bugs.python.org/issue39606
+        async def async_iterate():
+            yield 1
+            yield 2
+
+        async def run():
+            it = async_iterate()
+            await it.aclose()
+            await it.aclose()
+
+        self.loop.run_until_complete(run())
+
+    def test_async_gen_aclose_after_exhaustion(self):
+        # Regression test for https://bugs.python.org/issue39606
+        async def async_iterate():
+            yield 1
+            yield 2
+
+        async def run():
+            it = async_iterate()
+            async for _ in it:
+                pass
+            await it.aclose()
+
+        self.loop.run_until_complete(run())
 
 if __name__ == "__main__":
     unittest.main()
diff --git a/Misc/NEWS.d/next/Core and Builtins/2020-02-11-23-59-07.bpo-39606.a72Sxc.rst b/Misc/NEWS.d/next/Core and Builtins/2020-02-11-23-59-07.bpo-39606.a72Sxc.rst
new file mode 100644
index 00000000000000..b7cbe4e91f59c9
--- /dev/null
+++ b/Misc/NEWS.d/next/Core and Builtins/2020-02-11-23-59-07.bpo-39606.a72Sxc.rst	
@@ -0,0 +1,2 @@
+Fix regression caused by fix for bpo-39386, that prevented calling
+``aclose`` on an async generator that had already been closed or exhausted.
diff --git a/Objects/genobject.c b/Objects/genobject.c
index 576d6856c7f305..0efd57de7a5a34 100644
--- a/Objects/genobject.c
+++ b/Objects/genobject.c
@@ -1797,16 +1797,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,
@@ -1878,7 +1884,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;
             }
@@ -1893,16 +1898,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
@@ -1936,6 +1942,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;