Keep escaped generator frames and wrap asyncgen throw (#8697) · RustPython/RustPython@e18facf · GitHub
Skip to content

Commit e18facf

Browse files
youknowoneyouknowone
authored
Keep escaped generator frames and wrap asyncgen throw (#8697)
* Keep escaped generator frames and wrap asyncgen throw Mark gi_frame/cr_frame/ag_frame as escaped so locals survive deallocation, and let frame.clear() drop those locals after owner finalize. Route throw() through finalize_send_result so StopAsyncIteration is wrapped like send(). Add ag_suspended and drop never-started frame locals on close(). Assisted-by: Grok:grok-4.6 * Store escaped on iframe; skip unused cold alloc Keep the escaped flag on InterpreterFrame so mark_escaped and has_escaped do not allocate FrameColdData. Clear optional cold fields only when they exist. Transfer frame owner in close() the same way send/throw already do. Assisted-by: Grok:grok-4.6 * Replace escaped flag with take_ownership gi_frame returns the frame object. Close uses a uniquely-referenced check to clear locals or take_ownership onto a husk. frame.clear() finalizes a generator-owned frame and only clears FRAME_OBJECT frames. close() leaves the generator suspended when it yields on GeneratorExit. Assisted-by: Grok:grok-4.6 * Drop frame mutex; steal frame_obj atomically take_ownership stores None in an atomic slot instead of locking and swapping a husk. gi_code reads the kept executable. Resume exclusivity stays on the running claim. Assisted-by: Grok:grok-4.6
1 parent f18edb7 commit e18facf

10 files changed

Lines changed: 121 additions & 147 deletions

File tree

‎Lib/test/test_contextlib_async.py‎

Lines changed: 0 additions & 1 deletion

‎Lib/test/test_generators.py‎

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -302,7 +302,6 @@ def __iter__(self):
302302

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

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

722721
# See https://github.com/python/cpython/issues/125723
723722
class GeneratorDeallocTest(unittest.TestCase):
724-
@unittest.expectedFailure # TODO: RUSTPYTHON; frame uses shared Arc, no ownership transfer
725723
def test_frame_outlives_generator(self):
726724
def g1():
727725
a = 42

‎Lib/test/test_inspect/test_inspect.py‎

Lines changed: 0 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -2930,17 +2930,14 @@ def tearDownClass(cls):
29302930
def _asyncgenstate(self):
29312931
return inspect.getasyncgenstate(self.asyncgen)
29322932

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

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

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

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

2958-
@unittest.expectedFailure # TODO: RUSTPYTHON; AttributeError: 'async_generator' object has no attribute 'ag_suspended'
29592954
async def test_running(self):
29602955
async def running_check_asyncgen():
29612956
for number in range(5):

‎crates/vm/src/builtins/asyncgenerator.rs‎

Lines changed: 10 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -132,14 +132,14 @@ impl PyAsyncGen {
132132

133133
#[pygetset]
134134
fn ag_await(&self, _vm: &VirtualMachine) -> Option<PyObjectRef> {
135-
self.inner.frame().yield_from_target()
135+
self.inner.frame_opt().and_then(|f| f.yield_from_target())
136136
}
137137
#[pygetset]
138138
fn ag_frame(&self, _vm: &VirtualMachine) -> Option<FrameObjectRef> {
139139
if self.inner.closed() {
140140
None
141141
} else {
142-
Some(self.inner.frame())
142+
self.inner.frame_opt()
143143
}
144144
}
145145
#[pygetset]
@@ -148,7 +148,11 @@ impl PyAsyncGen {
148148
}
149149
#[pygetset]
150150
fn ag_code(&self, _vm: &VirtualMachine) -> PyRef<PyCode> {
151-
self.inner.frame().iframe().code().to_owned()
151+
self.inner.code()
152+
}
153+
#[pygetset]
154+
fn ag_suspended(&self, _vm: &VirtualMachine) -> bool {
155+
self.inner.suspended()
152156
}
153157

154158
#[pyclassmethod]
@@ -718,8 +722,6 @@ impl PyAnextAwaitable {
718722
if let Some(generator) = wrapped.downcast_ref::<PyGenerator>()
719723
&& generator
720724
.as_coro()
721-
.frame()
722-
.iframe()
723725
.code()
724726
.flags
725727
.contains(crate::bytecode::CodeFlags::ITERABLE_COROUTINE)
@@ -841,7 +843,9 @@ impl Destructor for PyAsyncGen {
841843

842844
impl Drop for PyAsyncGen {
843845
fn drop(&mut self) {
844-
self.inner.frame().clear_generator();
846+
if let Some(frame) = self.inner.frame_opt() {
847+
frame.clear_generator();
848+
}
845849
}
846850
}
847851

‎crates/vm/src/builtins/coroutine.rs‎

Lines changed: 7 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -89,14 +89,14 @@ impl PyCoroutine {
8989

9090
#[pygetset]
9191
fn cr_await(&self, _vm: &VirtualMachine) -> Option<PyObjectRef> {
92-
self.inner.frame().yield_from_target()
92+
self.inner.frame_opt().and_then(|f| f.yield_from_target())
9393
}
9494
#[pygetset]
9595
fn cr_frame(&self, _vm: &VirtualMachine) -> Option<FrameObjectRef> {
9696
if self.inner.closed() {
9797
None
9898
} else {
99-
Some(self.inner.frame())
99+
self.inner.frame_opt()
100100
}
101101
}
102102
#[pygetset]
@@ -105,7 +105,7 @@ impl PyCoroutine {
105105
}
106106
#[pygetset]
107107
fn cr_code(&self, _vm: &VirtualMachine) -> PyRef<PyCode> {
108-
self.inner.frame().iframe().code().to_owned()
108+
self.inner.code()
109109
}
110110
#[pygetset]
111111
fn cr_origin(&self, _vm: &VirtualMachine) -> Option<PyTupleRef> {
@@ -179,7 +179,7 @@ impl Destructor for PyCoroutine {
179179
if zelf.inner.closed() || zelf.inner.running() {
180180
return Ok(());
181181
}
182-
if zelf.inner.frame().lasti() == 0 {
182+
if zelf.inner.frame_opt().is_none_or(|f| f.lasti() == 0) {
183183
crate::warn::warn_unawaited_coroutine(zelf.as_object(), &zelf.inner.qualname(), vm);
184184
zelf.inner.closed.store(true);
185185
return Ok(());
@@ -266,7 +266,9 @@ impl IterNext for PyCoroutineWrapper {
266266

267267
impl Drop for PyCoroutine {
268268
fn drop(&mut self) {
269-
self.inner.frame().clear_generator();
269+
if let Some(frame) = self.inner.frame_opt() {
270+
frame.clear_generator();
271+
}
270272
}
271273
}
272274

‎crates/vm/src/builtins/frame.rs‎

Lines changed: 16 additions & 35 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@
22
33
*/
44

5-
use super::{PyAsyncGen, PyCode, PyCoroutine, PyDictRef, PyIntRef, PyStrRef};
5+
use super::{PyAsyncGen, PyCode, PyCoroutine, PyDictRef, PyGenerator, PyIntRef, PyStrRef};
66
use crate::{
77
Context, Py, PyObjectRef, PyPayload, PyRef, PyResult, VirtualMachine,
88
class::PyClassImpl,
@@ -778,10 +778,8 @@ impl Py<FrameObject> {
778778
);
779779
match owner {
780780
FrameOwner::Generator => {
781-
// Generator frame: check if suspended (lasti > 0 means
782-
// FRAME_SUSPENDED). lasti == 0 means FRAME_CREATED and
783-
// can be cleared. Finalize the owner so a never-started
784-
// coroutine emits its never-awaited warning.
781+
// FRAME_SUSPENDED (lasti > 0) cannot be cleared. FRAME_CREATED
782+
// and finished frames go through the owner finalizer.
785783
if self.lasti() != 0 {
786784
return Err(vm.new_runtime_error("cannot clear a suspended frame"));
787785
}
@@ -790,17 +788,16 @@ impl Py<FrameObject> {
790788
let _ = PyCoroutine::del(coro, vm);
791789
} else if let Some(async_gen) = owner.downcast_ref::<PyAsyncGen>() {
792790
let _ = PyAsyncGen::del(async_gen, vm);
791+
} else if let Some(generator) = owner.downcast_ref::<PyGenerator>() {
792+
let _ = PyGenerator::del(generator, vm);
793793
}
794794
}
795795
return Ok(());
796796
}
797797
FrameOwner::Thread => {
798-
// Thread-owned frame: always executing, cannot clear.
799798
return Err(vm.new_runtime_error("cannot clear an executing frame"));
800799
}
801800
FrameOwner::FrameObject => {
802-
// Check if this materialized frame is backed by a live
803-
// stack-allocated iframe — if so, the frame is executing.
804801
if !self.find_live_source_iframe().is_null() {
805802
return Err(vm.new_runtime_error("cannot clear an executing frame"));
806803
}
@@ -822,27 +819,17 @@ impl Py<FrameObject> {
822819
// Clear the evaluation stack and cell references
823820
self.clear_stack_and_cells();
824821

825-
let cold = self.iframe().cold();
826-
let temporary_refs = {
827-
let mut guard = cold.temporary_refs.lock();
828-
core::mem::take(&mut *guard)
829-
};
830-
let extra_locals = {
831-
let mut guard = cold.f_extra_locals.lock();
832-
guard.take()
833-
};
834-
let locals_cache = {
835-
let mut guard = cold.f_locals_cache.lock();
836-
guard.take()
837-
};
838-
let overwritten = {
839-
let mut guard = cold.f_overwritten_fast_locals.lock();
840-
core::mem::take(&mut *guard)
841-
};
842-
let retained_back = {
843-
let mut guard = cold.retained_back.lock();
844-
guard.take()
845-
};
822+
let (temporary_refs, extra_locals, locals_cache, overwritten, retained_back) =
823+
match self.iframe().cold_opt() {
824+
Some(cold) => (
825+
core::mem::take(&mut *cold.temporary_refs.lock()),
826+
cold.f_extra_locals.lock().take(),
827+
cold.f_locals_cache.lock().take(),
828+
core::mem::take(&mut *cold.f_overwritten_fast_locals.lock()),
829+
cold.retained_back.lock().take(),
830+
),
831+
None => (Vec::new(), None, None, Vec::new(), None),
832+
};
846833
drop((
847834
fastlocals,
848835
temporary_refs,
@@ -858,7 +845,6 @@ impl Py<FrameObject> {
858845
#[pygetset]
859846
fn f_locals(&self, vm: &VirtualMachine) -> PyResult {
860847
if self.uses_locals_proxy(vm)? {
861-
self.mark_escaped();
862848
let proxy = crate::builtins::FrameLocalsProxy::new(self.to_owned());
863849
Ok(proxy.into_ref(&vm.ctx).into())
864850
} else {
@@ -895,7 +881,6 @@ impl Py<FrameObject> {
895881
// Check retained_back for frames whose callers have returned
896882
let retained = self.iframe().cold().retained_back.lock().clone();
897883
if let Some(frame) = retained {
898-
frame.mark_escaped();
899884
return Some(frame);
900885
}
901886
return None;
@@ -911,7 +896,6 @@ impl Py<FrameObject> {
911896
if core::ptr::eq(cur, prev) {
912897
let iframe_ref = unsafe { &*cur };
913898
let fo = iframe_ref.materialize(vm);
914-
fo.mark_escaped();
915899
return Some(fo.to_owned());
916900
}
917901
cur = unsafe { (*cur).previous() };
@@ -921,7 +905,6 @@ impl Py<FrameObject> {
921905
// The caller already returned — check retained_back
922906
let retained = self.iframe().cold().retained_back.lock().clone();
923907
if let Some(frame) = retained {
924-
frame.mark_escaped();
925908
return Some(frame);
926909
}
927910

@@ -936,13 +919,11 @@ impl Py<FrameObject> {
936919
let prev_ref = unsafe { &*prev };
937920
// Fast path: already materialized.
938921
if let Some(fo) = prev_ref.frame_obj() {
939-
fo.mark_escaped();
940922
return Some(fo.to_owned());
941923
}
942924
// Slow path: copy the whole chain, linked through retained_back.
943925
// SAFETY: the world is stopped, so the owning thread is parked.
944926
let fo = unsafe { prev_ref.materialize_detached_chain(vm) };
945-
fo.mark_escaped();
946927
return Some(fo);
947928
}
948929

‎crates/vm/src/builtins/generator.rs‎

Lines changed: 6 additions & 10 deletions

0 commit comments

Comments
 (0)