Skip to content

Commit 8418dcc

Browse files
committed
GH-117881: fix athrow().throw()/asend().throw() concurrent access
1 parent 7bcc257 commit 8418dcc

File tree

2 files changed

+97
-9
lines changed

2 files changed

+97
-9
lines changed

Lib/test/test_asyncgen.py

Lines changed: 61 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -393,6 +393,64 @@ async def gen():
393393
r'anext\(\): asynchronous generator is already running'):
394394
an.__next__()
395395

396+
def test_async_gen_asend_throw_concurrent(self):
397+
import types
398+
399+
@types.coroutine
400+
def _async_yield(v):
401+
return (yield v)
402+
403+
class MyExc(Exception):
404+
pass
405+
406+
async def agenfn():
407+
while True:
408+
try:
409+
await _async_yield(None)
410+
except MyExc:
411+
pass
412+
return
413+
yield
414+
415+
416+
agen = agenfn()
417+
gen = agen.asend(None)
418+
gen.send(None)
419+
gen2 = agen.asend(None)
420+
421+
with self.assertRaisesRegex(RuntimeError,
422+
r'anext\(\): asynchronous generator is already running'):
423+
gen2.throw(MyExc)
424+
425+
def test_async_gen_athrow_throw_concurrent(self):
426+
import types
427+
428+
@types.coroutine
429+
def _async_yield(v):
430+
return (yield v)
431+
432+
class MyExc(Exception):
433+
pass
434+
435+
async def agenfn():
436+
while True:
437+
try:
438+
await _async_yield(None)
439+
except MyExc:
440+
pass
441+
return
442+
yield
443+
444+
445+
agen = agenfn()
446+
gen = agen.asend(None)
447+
gen.send(None)
448+
gen2 = agen.athrow(MyExc)
449+
450+
with self.assertRaisesRegex(RuntimeError,
451+
r'athrow\(\): asynchronous generator is already running'):
452+
gen2.throw(MyExc)
453+
396454
def test_async_gen_3_arg_deprecation_warning(self):
397455
async def gen():
398456
yield 123
@@ -1572,11 +1630,8 @@ async def main():
15721630
self.assertIsInstance(message['exception'], ZeroDivisionError)
15731631
self.assertIn('unhandled exception during asyncio.run() shutdown',
15741632
message['message'])
1575-
with self.assertWarnsRegex(RuntimeWarning,
1576-
f"coroutine method 'aclose' of '{async_iterate.__qualname__}' "
1577-
f"was never awaited"):
1578-
del message, messages
1579-
gc_collect()
1633+
del message, messages
1634+
gc_collect()
15801635

15811636
def test_async_gen_expression_01(self):
15821637
async def arange(n):
@@ -1630,10 +1685,7 @@ async def main():
16301685
asyncio.run(main())
16311686

16321687
self.assertEqual([], messages)
1633-
with self.assertWarnsRegex(RuntimeWarning,
1634-
f"coroutine method 'aclose' of '{async_iterate.__qualname__}' "
1635-
f"was never awaited"):
1636-
gc_collect()
1688+
gc_collect()
16371689

16381690
def test_async_gen_await_same_anext_coro_twice(self):
16391691
async def async_iterate():

Objects/genobject.c

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1771,6 +1771,7 @@ async_gen_asend_send(PyAsyncGenASend *o, PyObject *arg)
17711771

17721772
if (o->ags_state == AWAITABLE_STATE_INIT) {
17731773
if (o->ags_gen->ag_running_async) {
1774+
o->ags_state = AWAITABLE_STATE_CLOSED;
17741775
PyErr_SetString(
17751776
PyExc_RuntimeError,
17761777
"anext(): asynchronous generator is already running");
@@ -1814,10 +1815,23 @@ async_gen_asend_throw(PyAsyncGenASend *o, PyObject *const *args, Py_ssize_t narg
18141815
return NULL;
18151816
}
18161817

1818+
if (o->ags_state == AWAITABLE_STATE_INIT) {
1819+
if (o->ags_gen->ag_running_async) {
1820+
o->ags_state = AWAITABLE_STATE_CLOSED;
1821+
PyErr_SetString(
1822+
PyExc_RuntimeError,
1823+
"anext(): asynchronous generator is already running");
1824+
return NULL;
1825+
}
1826+
1827+
o->ags_state = AWAITABLE_STATE_ITER;
1828+
}
1829+
18171830
result = gen_throw((PyGenObject*)o->ags_gen, args, nargs);
18181831
result = async_gen_unwrap_value(o->ags_gen, result);
18191832

18201833
if (result == NULL) {
1834+
o->ags_gen->ag_running_async = 0;
18211835
o->ags_state = AWAITABLE_STATE_CLOSED;
18221836
}
18231837

@@ -2206,6 +2220,25 @@ async_gen_athrow_throw(PyAsyncGenAThrow *o, PyObject *const *args, Py_ssize_t na
22062220
return NULL;
22072221
}
22082222

2223+
if (o->agt_state == AWAITABLE_STATE_INIT) {
2224+
if (o->agt_gen->ag_running_async) {
2225+
o->agt_state = AWAITABLE_STATE_CLOSED;
2226+
if (o->agt_args == NULL) {
2227+
PyErr_SetString(
2228+
PyExc_RuntimeError,
2229+
"aclose(): asynchronous generator is already running");
2230+
}
2231+
else {
2232+
PyErr_SetString(
2233+
PyExc_RuntimeError,
2234+
"athrow(): asynchronous generator is already running");
2235+
}
2236+
return NULL;
2237+
}
2238+
2239+
o->agt_state = AWAITABLE_STATE_ITER;
2240+
}
2241+
22092242
retval = gen_throw((PyGenObject*)o->agt_gen, args, nargs);
22102243
if (o->agt_args) {
22112244
return async_gen_unwrap_value(o->agt_gen, retval);
@@ -2218,6 +2251,9 @@ async_gen_athrow_throw(PyAsyncGenAThrow *o, PyObject *const *args, Py_ssize_t na
22182251
PyErr_SetString(PyExc_RuntimeError, ASYNC_GEN_IGNORED_EXIT_MSG);
22192252
return NULL;
22202253
}
2254+
if (retval == NULL) {
2255+
o->agt_gen->ag_running_async = 0;
2256+
}
22212257
if (PyErr_ExceptionMatches(PyExc_StopAsyncIteration) ||
22222258
PyErr_ExceptionMatches(PyExc_GeneratorExit))
22232259
{

0 commit comments

Comments
 (0)