Skip to content

bpo-39606: allow closing async generators that are already closed #18475

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
Feb 13, 2020
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
30 changes: 28 additions & 2 deletions Lib/test/test_asyncgen.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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
Expand All @@ -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()
Original file line number Diff line number Diff line change
@@ -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.
15 changes: 11 additions & 4 deletions Objects/genobject.c
Original file line number Diff line number Diff line change
Expand Up @@ -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;
}

Copy link
Contributor Author

@njsmith njsmith Feb 12, 2020

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

These two hunks are the main change. Originally, this code raised StopIteration if the generator had no frame. The fix for bpo-39386 added the check for the athrow coroutine being in AWAITABLE_STATE_CLOSED to the same if statement, and switched it to raise RuntimeError in both cases. But in fact, the generator having no frame should raise StopIteration; it's only the AWAITABLE_STATE_CLOSED case that should raise RuntimeError. So I split the AWAITABLE_STATE_CLOSED logic out into its own block, and restored the old if (f == NULL || f->f_stacktop) logic.

Everything else is just going through to make sure that we're setting AWAITABLE_STATE_CLOSED at the right times.

if (o->agt_state == AWAITABLE_STATE_INIT) {
if (o->agt_gen->ag_running_async) {
o->agt_state = AWAITABLE_STATE_CLOSED;
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

If the athrow/aclose is being aborted because there's already another task advancing the underlying async generator, then that's a final result, so we should mark the coroutine as complete.

if (o->agt_args == NULL) {
PyErr_SetString(
PyExc_RuntimeError,
Expand Down Expand Up @@ -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;
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

yield_close already sets ag_running_async = 0, so the deleted line was already a no-op.

}
Expand All @@ -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;
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

ag_running_async tracks whether there's some call to asend/athrow in progress. agt_state tracks whether this call to asend/athrow is in progress. So every time we set ag_running_async = 0, we should also set 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
Expand Down Expand Up @@ -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;
Expand Down