Keep escaped generator frames and wrap asyncgen throw by youknowone · Pull Request #8697 · RustPython/RustPython · GitHub
Skip to content
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
1 change: 0 additions & 1 deletion Lib/test/test_contextlib_async.py
2 changes: 0 additions & 2 deletions Lib/test/test_generators.py
Original file line number Diff line number Diff line change
Expand Up @@ -302,7 +302,6 @@ def __iter__(self):

self.assertEqual([1, 2], list(i for i in C()))

@unittest.expectedFailure # TODO: RUSTPYTHON; AssertionError: False is not true
def test_close_clears_frame(self):
# gh-142766: Test that closing a generator clears its frame
class DetectDelete:
Expand Down Expand Up @@ -721,7 +720,6 @@ def genfn():

# See https://github.com/python/cpython/issues/125723
class GeneratorDeallocTest(unittest.TestCase):
@unittest.expectedFailure # TODO: RUSTPYTHON; frame uses shared Arc, no ownership transfer
def test_frame_outlives_generator(self):
def g1():
a = 42
Expand Down
5 changes: 0 additions & 5 deletions Lib/test/test_inspect/test_inspect.py
Original file line number Diff line number Diff line change
Expand Up @@ -2930,17 +2930,14 @@ def tearDownClass(cls):
def _asyncgenstate(self):
return inspect.getasyncgenstate(self.asyncgen)

@unittest.expectedFailure # TODO: RUSTPYTHON; AttributeError: 'async_generator' object has no attribute 'ag_suspended'
def test_created(self):
self.assertEqual(self._asyncgenstate(), inspect.AGEN_CREATED)

@unittest.expectedFailure # TODO: RUSTPYTHON; AttributeError: 'async_generator' object has no attribute 'ag_suspended'
async def test_suspended(self):
value = await anext(self.asyncgen)
self.assertEqual(self._asyncgenstate(), inspect.AGEN_SUSPENDED)
self.assertEqual(value, 0)

@unittest.expectedFailure # TODO: RUSTPYTHON; AttributeError: 'async_generator' object has no attribute 'ag_suspended'
async def test_closed_after_exhaustion(self):
countdown = 7
with self.assertRaises(StopAsyncIteration):
Expand All @@ -2949,13 +2946,11 @@ async def test_closed_after_exhaustion(self):
self.assertEqual(countdown, 1)
self.assertEqual(self._asyncgenstate(), inspect.AGEN_CLOSED)

@unittest.expectedFailure # TODO: RUSTPYTHON; AttributeError: 'async_generator' object has no attribute 'ag_suspended'
async def test_closed_after_immediate_exception(self):
with self.assertRaises(RuntimeError):
await self.asyncgen.athrow(RuntimeError)
self.assertEqual(self._asyncgenstate(), inspect.AGEN_CLOSED)

@unittest.expectedFailure # TODO: RUSTPYTHON; AttributeError: 'async_generator' object has no attribute 'ag_suspended'
async def test_running(self):
async def running_check_asyncgen():
for number in range(5):
Expand Down
16 changes: 10 additions & 6 deletions crates/vm/src/builtins/asyncgenerator.rs
Original file line number Diff line number Diff line change
Expand Up @@ -132,14 +132,14 @@ impl PyAsyncGen {

#[pygetset]
fn ag_await(&self, _vm: &VirtualMachine) -> Option<PyObjectRef> {
self.inner.frame().yield_from_target()
self.inner.frame_opt().and_then(|f| f.yield_from_target())
}
#[pygetset]
fn ag_frame(&self, _vm: &VirtualMachine) -> Option<FrameObjectRef> {
if self.inner.closed() {
None
} else {
Some(self.inner.frame())
self.inner.frame_opt()
}
}
#[pygetset]
Expand All @@ -148,7 +148,11 @@ impl PyAsyncGen {
}
#[pygetset]
fn ag_code(&self, _vm: &VirtualMachine) -> PyRef<PyCode> {
self.inner.frame().iframe().code().to_owned()
self.inner.code()
}
#[pygetset]
fn ag_suspended(&self, _vm: &VirtualMachine) -> bool {
self.inner.suspended()
Comment on lines +154 to +155

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

P2 Badge Report pending async-generator iterations as running

When an async generator is paused inside an inner await, running_async remains true for the outstanding anext()/asend(), but RunningGuard has already reset inner.running() to false. With this newly exposed property, inspect.getasyncgenstate() consequently sees ag_running == False and ag_suspended == True and reports AGEN_SUSPENDED; CPython reports AGEN_RUNNING until that outstanding asynchronous iteration completes. The public running state needs to account for running_async before enabling this state API.

Useful? React with 👍 / 👎.

Comment on lines +154 to +155

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

P2 Badge Preserve suspension after an ignored async-generator close

When an async generator catches GeneratorExit and yields from await agen.aclose(), PyAsyncGenAThrow::yield_close marks inner.closed even though execution remains paused at that yield. Consequently this new property returns false, inspect.getasyncgenstate() reports AGEN_CLOSED, ag_frame disappears, and a subsequent anext() cannot resume the generator; CPython instead reports AGEN_SUSPENDED and permits the final resume. The ignored-close path must retain the suspended state rather than publishing closure.

Useful? React with 👍 / 👎.

}

#[pyclassmethod]
Expand Down Expand Up @@ -718,8 +722,6 @@ impl PyAnextAwaitable {
if let Some(generator) = wrapped.downcast_ref::<PyGenerator>()
&& generator
.as_coro()
.frame()
.iframe()
.code()
.flags
.contains(crate::bytecode::CodeFlags::ITERABLE_COROUTINE)
Expand Down Expand Up @@ -841,7 +843,9 @@ impl Destructor for PyAsyncGen {

impl Drop for PyAsyncGen {
fn drop(&mut self) {
self.inner.frame().clear_generator();
if let Some(frame) = self.inner.frame_opt() {
frame.clear_generator();
}
}
}

Expand Down
12 changes: 7 additions & 5 deletions crates/vm/src/builtins/coroutine.rs
Original file line number Diff line number Diff line change
Expand Up @@ -89,14 +89,14 @@ impl PyCoroutine {

#[pygetset]
fn cr_await(&self, _vm: &VirtualMachine) -> Option<PyObjectRef> {
self.inner.frame().yield_from_target()
self.inner.frame_opt().and_then(|f| f.yield_from_target())
}
#[pygetset]
fn cr_frame(&self, _vm: &VirtualMachine) -> Option<FrameObjectRef> {
if self.inner.closed() {
None
} else {
Some(self.inner.frame())
self.inner.frame_opt()
}
}
#[pygetset]
Expand All @@ -105,7 +105,7 @@ impl PyCoroutine {
}
#[pygetset]
fn cr_code(&self, _vm: &VirtualMachine) -> PyRef<PyCode> {
self.inner.frame().iframe().code().to_owned()
self.inner.code()
}
#[pygetset]
fn cr_origin(&self, _vm: &VirtualMachine) -> Option<PyTupleRef> {
Expand Down Expand Up @@ -179,7 +179,7 @@ impl Destructor for PyCoroutine {
if zelf.inner.closed() || zelf.inner.running() {
return Ok(());
}
if zelf.inner.frame().lasti() == 0 {
if zelf.inner.frame_opt().is_none_or(|f| f.lasti() == 0) {
crate::warn::warn_unawaited_coroutine(zelf.as_object(), &zelf.inner.qualname(), vm);
zelf.inner.closed.store(true);
return Ok(());
Expand Down Expand Up @@ -266,7 +266,9 @@ impl IterNext for PyCoroutineWrapper {

impl Drop for PyCoroutine {
fn drop(&mut self) {
self.inner.frame().clear_generator();
if let Some(frame) = self.inner.frame_opt() {
frame.clear_generator();
}
}
}

Expand Down
51 changes: 16 additions & 35 deletions crates/vm/src/builtins/frame.rs
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,7 @@

*/

use super::{PyAsyncGen, PyCode, PyCoroutine, PyDictRef, PyIntRef, PyStrRef};
use super::{PyAsyncGen, PyCode, PyCoroutine, PyDictRef, PyGenerator, PyIntRef, PyStrRef};
use crate::{
Context, Py, PyObjectRef, PyPayload, PyRef, PyResult, VirtualMachine,
class::PyClassImpl,
Expand Down Expand Up @@ -778,10 +778,8 @@ impl Py<FrameObject> {
);
match owner {
FrameOwner::Generator => {
// Generator frame: check if suspended (lasti > 0 means
// FRAME_SUSPENDED). lasti == 0 means FRAME_CREATED and
// can be cleared. Finalize the owner so a never-started
// coroutine emits its never-awaited warning.
// FRAME_SUSPENDED (lasti > 0) cannot be cleared. FRAME_CREATED
// and finished frames go through the owner finalizer.
if self.lasti() != 0 {
return Err(vm.new_runtime_error("cannot clear a suspended frame"));
}
Expand All @@ -790,17 +788,16 @@ impl Py<FrameObject> {
let _ = PyCoroutine::del(coro, vm);
} else if let Some(async_gen) = owner.downcast_ref::<PyAsyncGen>() {
let _ = PyAsyncGen::del(async_gen, vm);
} else if let Some(generator) = owner.downcast_ref::<PyGenerator>() {
let _ = PyGenerator::del(generator, vm);
}
Comment on lines 789 to 793

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

P2 Badge Close created async generators before clearing their frames

When ag.ag_frame.clear() is called before the async generator starts, PyAsyncGen::del does not mark the generator closed, but this new fallthrough proceeds to erase its fast locals and stack. The generator therefore still exposes a non-None ag_frame, and a subsequent anext(ag) resumes a cleared frame—typically raising UnboundLocalError for an argument—whereas a cleared created async generator should be closed and raise StopAsyncIteration. Ensure the async-generator owner is synchronously marked closed before clearing its frame.

Useful? React with 👍 / 👎.

}
return Ok(());
}
FrameOwner::Thread => {
// Thread-owned frame: always executing, cannot clear.
return Err(vm.new_runtime_error("cannot clear an executing frame"));
}
FrameOwner::FrameObject => {
// Check if this materialized frame is backed by a live
// stack-allocated iframe — if so, the frame is executing.
if !self.find_live_source_iframe().is_null() {
return Err(vm.new_runtime_error("cannot clear an executing frame"));
}
Expand All @@ -822,27 +819,17 @@ impl Py<FrameObject> {
// Clear the evaluation stack and cell references
self.clear_stack_and_cells();

let cold = self.iframe().cold();
let temporary_refs = {
let mut guard = cold.temporary_refs.lock();
core::mem::take(&mut *guard)
};
let extra_locals = {
let mut guard = cold.f_extra_locals.lock();
guard.take()
};
let locals_cache = {
let mut guard = cold.f_locals_cache.lock();
guard.take()
};
let overwritten = {
let mut guard = cold.f_overwritten_fast_locals.lock();
core::mem::take(&mut *guard)
};
let retained_back = {
let mut guard = cold.retained_back.lock();
guard.take()
};
let (temporary_refs, extra_locals, locals_cache, overwritten, retained_back) =
match self.iframe().cold_opt() {
Some(cold) => (
core::mem::take(&mut *cold.temporary_refs.lock()),
cold.f_extra_locals.lock().take(),
cold.f_locals_cache.lock().take(),
core::mem::take(&mut *cold.f_overwritten_fast_locals.lock()),
cold.retained_back.lock().take(),
),
None => (Vec::new(), None, None, Vec::new(), None),
};
drop((
fastlocals,
temporary_refs,
Expand All @@ -858,7 +845,6 @@ impl Py<FrameObject> {
#[pygetset]
fn f_locals(&self, vm: &VirtualMachine) -> PyResult {
if self.uses_locals_proxy(vm)? {
self.mark_escaped();
let proxy = crate::builtins::FrameLocalsProxy::new(self.to_owned());
Ok(proxy.into_ref(&vm.ctx).into())
} else {
Expand Down Expand Up @@ -895,7 +881,6 @@ impl Py<FrameObject> {
// Check retained_back for frames whose callers have returned
let retained = self.iframe().cold().retained_back.lock().clone();
if let Some(frame) = retained {
frame.mark_escaped();
return Some(frame);
}
return None;
Expand All @@ -911,7 +896,6 @@ impl Py<FrameObject> {
if core::ptr::eq(cur, prev) {
let iframe_ref = unsafe { &*cur };
let fo = iframe_ref.materialize(vm);
fo.mark_escaped();
return Some(fo.to_owned());
}
cur = unsafe { (*cur).previous() };
Expand All @@ -921,7 +905,6 @@ impl Py<FrameObject> {
// The caller already returned — check retained_back
let retained = self.iframe().cold().retained_back.lock().clone();
if let Some(frame) = retained {
frame.mark_escaped();
return Some(frame);
}

Expand All @@ -936,13 +919,11 @@ impl Py<FrameObject> {
let prev_ref = unsafe { &*prev };
// Fast path: already materialized.
if let Some(fo) = prev_ref.frame_obj() {
fo.mark_escaped();
return Some(fo.to_owned());
}
// Slow path: copy the whole chain, linked through retained_back.
// SAFETY: the world is stopped, so the owning thread is parked.
let fo = unsafe { prev_ref.materialize_detached_chain(vm) };
fo.mark_escaped();
return Some(fo);
}

Expand Down
16 changes: 6 additions & 10 deletions crates/vm/src/builtins/generator.rs
Loading
Loading