Skip to content

Commit 770484e

Browse files
committed
Restore interrupted with-enter coverage with pending exception support
Deliver VM-owned pending exception classes at bytecode safepoints and clear their roots before terminal owner teardown. Add the native SelfInterruptingContextManager fixture and restore the canonical test_with regression via update_lib. Assisted-by: OpenAI Codex: model version unavailable
1 parent c2bb450 commit 770484e

8 files changed

Lines changed: 337 additions & 6 deletions

File tree

‎Lib/test/test_with.py‎

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@
1010
import unittest
1111
from collections import deque
1212
from contextlib import _GeneratorContextManager, contextmanager, nullcontext
13+
from _testinternalcapi import SelfInterruptingContextManager
1314

1415

1516
def do_with(obj):
@@ -850,5 +851,21 @@ def exit_raises():
850851
expected)
851852

852853

854+
class InterruptDuringEnter(unittest.TestCase):
855+
856+
def test_exit_called_after_interrupt(self):
857+
cm = SelfInterruptingContextManager()
858+
self.assertFalse(cm.within())
859+
try:
860+
with cm:
861+
self.assertTrue(cm.within())
862+
except KeyboardInterrupt:
863+
self.assertFalse(cm.within())
864+
return
865+
except:
866+
self.fail("Wrong exception raised")
867+
self.fail("No exception raised")
868+
869+
853870
if __name__ == '__main__':
854871
unittest.main()

‎crates/vm/src/frame.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3931,7 +3931,7 @@ impl ExecutingFrame<'_> {
39313931
// the allocation that tripped the threshold.
39323932
#[cfg(feature = "threading")]
39333933
vm.run_scheduled_gc();
3934-
Ok(())
3934+
vm.raise_async_exception()
39353935
}
39363936
if vm.eval_breaker_tripped()
39373937
&& let Err(exception) = eval_breaker_work(vm)

‎crates/vm/src/stdlib/_testinternalcapi.rs‎

Lines changed: 43 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -58,6 +58,49 @@ mod _testinternalcapi {
5858
#[pyattr]
5959
const SPECIALIZATION_THRESHOLD: usize = 2;
6060

61+
#[pyattr]
62+
#[pyclass(name = "SelfInterruptingContextManager", module = "_testcapi")]
63+
#[derive(Debug, PyPayload)]
64+
struct SelfInterruptingContextManager {
65+
within: core::sync::atomic::AtomicBool,
66+
}
67+
68+
impl Constructor for SelfInterruptingContextManager {
69+
type Args = FuncArgs;
70+
71+
fn py_new(_cls: &Py<PyType>, _args: FuncArgs, _vm: &VirtualMachine) -> PyResult<Self> {
72+
Ok(Self {
73+
within: core::sync::atomic::AtomicBool::new(false),
74+
})
75+
}
76+
}
77+
78+
#[pyclass(with(Constructor), flags(IMMUTABLETYPE))]
79+
impl SelfInterruptingContextManager {
80+
#[pymethod]
81+
fn __enter__(zelf: PyRef<Self>, vm: &VirtualMachine) -> PyResult<PyRef<Self>> {
82+
zelf.within.store(true, Ordering::Relaxed);
83+
let Some(old) = vm.replace_async_exception(Some(
84+
vm.ctx.exceptions.keyboard_interrupt.to_owned().into(),
85+
)) else {
86+
return Err(vm.new_runtime_error("the current thread state is being cleared"));
87+
};
88+
vm.notify_async_exception();
89+
drop(old);
90+
Ok(zelf)
91+
}
92+
93+
#[pymethod]
94+
fn within(zelf: &Py<Self>) -> bool {
95+
zelf.within.load(Ordering::Relaxed)
96+
}
97+
98+
#[pymethod]
99+
fn __exit__(zelf: &Py<Self>, _args: crate::function::PosArgs<PyObjectRef>) {
100+
zelf.within.store(false, Ordering::Relaxed);
101+
}
102+
}
103+
61104
// ADAPTIVE_COOLDOWN_VALUE + 1
62105
#[pyattr]
63106
const SPECIALIZATION_COOLDOWN: usize = 53;

‎crates/vm/src/stdlib/_thread.rs‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -838,6 +838,7 @@ pub(crate) mod _thread {
838838
// Clean up thread-local storage while VM context is still active
839839
// This ensures __del__ methods are called properly
840840
cleanup_thread_local_data();
841+
vm.clear_async_exception();
841842
// Clean up frame tracking
842843
crate::vm::thread::cleanup_current_thread_frames(vm);
843844
vm.state.thread_count.fetch_sub(1);
@@ -2240,6 +2241,7 @@ pub(crate) mod _thread {
22402241

22412242
// Clean up thread-local data while VM context is still active
22422243
cleanup_thread_local_data();
2244+
vm.clear_async_exception();
22432245

22442246
// Clean up frame tracking
22452247
crate::vm::thread::cleanup_current_thread_frames(vm);
Lines changed: 150 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,150 @@
1+
use super::VirtualMachine;
2+
use crate::{
3+
AsObject, PyAtomicRef, PyObject, PyObjectRef, PyResult,
4+
builtins::{PyBaseException, PyType},
5+
common::wtf8::Wtf8Buf,
6+
};
7+
use core::cell::Cell;
8+
9+
/// Owned by one execution context, like `PyThreadState.async_exc`.
10+
/// This is an external GC root, not an edge owned by a Python object.
11+
/// The VM is not Sync: only its owning execution thread accesses these flags.
12+
/// Cross-thread requests would require a separate synchronized API and wakeup.
13+
pub(super) struct PendingAsyncException {
14+
exception: PyAtomicRef<Option<PyObject>>,
15+
pending: Cell<bool>,
16+
closed: Cell<bool>,
17+
}
18+
19+
impl Default for PendingAsyncException {
20+
fn default() -> Self {
21+
Self {
22+
exception: PyAtomicRef::new_empty(),
23+
pending: Cell::new(false),
24+
closed: Cell::new(false),
25+
}
26+
}
27+
}
28+
29+
impl PendingAsyncException {
30+
pub(super) fn is_pending(&self) -> bool {
31+
self.pending.get()
32+
}
33+
34+
pub(super) fn has_reference(&self) -> bool {
35+
!self.exception.load_ptr().is_null()
36+
}
37+
38+
fn replace(&self, exception: Option<PyObjectRef>) -> Option<Option<PyObjectRef>> {
39+
if self.closed.get() {
40+
return None;
41+
}
42+
Some(self.exception.store(exception))
43+
}
44+
45+
fn notify(&self) {
46+
if !self.closed.get() {
47+
self.pending.set(true);
48+
}
49+
}
50+
51+
fn take(&self) -> Option<PyObjectRef> {
52+
if !self.pending.replace(false) {
53+
return None;
54+
}
55+
self.exception.store(None)
56+
}
57+
58+
fn close(&self) {
59+
// A cleared thread state cannot accept another request. Rejecting
60+
// reentrant requests makes this a single clear, not a finalizer drain.
61+
self.closed.set(true);
62+
self.pending.set(false);
63+
drop(self.exception.store(None));
64+
}
65+
}
66+
67+
impl VirtualMachine {
68+
/// Replace the currently attached thread state's asynchronous exception.
69+
/// The outer `None` means that this VM is inactive or is being cleared.
70+
/// The old owner is returned so callers can release it outside all locks,
71+
/// either before or after notification as required by their API contract.
72+
/// `None` as the input cancels; validation happens only during delivery.
73+
pub(crate) fn replace_async_exception(
74+
&self,
75+
exception: Option<PyObjectRef>,
76+
) -> Option<Option<PyObjectRef>> {
77+
if !super::thread::try_with_current_vm(|vm| core::ptr::eq(vm, self)).unwrap_or(false) {
78+
return None;
79+
}
80+
self.async_exception.replace(exception)
81+
}
82+
83+
pub(crate) fn notify_async_exception(&self) {
84+
self.async_exception.notify();
85+
}
86+
87+
/// Called only at bytecode safepoints, after pending signals and GC.
88+
pub(crate) fn raise_async_exception(&self) -> PyResult<()> {
89+
let Some(exception) = self.async_exception.take() else {
90+
return Ok(());
91+
};
92+
let class = match exception.downcast::<PyType>() {
93+
Ok(class) if class.fast_issubclass(self.ctx.exceptions.base_exception_type) => class,
94+
object => {
95+
let object: PyObjectRef = match object {
96+
Ok(class) => class.into(),
97+
Err(object) => object,
98+
};
99+
let mut message = Wtf8Buf::from("_PyErr_SetObject: exception ");
100+
message.push_wtf8(object.repr(self)?.as_wtf8());
101+
message.push_str(" is not a BaseException subclass");
102+
return Err(self.new_system_error(message));
103+
}
104+
};
105+
106+
// `_PyErr_SetNone` invokes the class through its metaclass, including
107+
// user __call__/__new__/__init__, rather than allocating a bare payload.
108+
let result = class.as_object().call((), self).and_then(|object| {
109+
object.downcast::<PyBaseException>().map_err(|object| {
110+
let mut message = Wtf8Buf::from("calling ");
111+
match class.as_object().repr(self) {
112+
Ok(repr) => message.push_wtf8(repr.as_wtf8()),
113+
Err(error) => return error,
114+
}
115+
message.push_str(" should have returned an instance of BaseException, not ");
116+
message.push_str(&object.class().slot_name());
117+
self.new_type_error(message)
118+
})
119+
});
120+
let exception = match result {
121+
Ok(exception) => exception,
122+
Err(error) => {
123+
// CPython represents the NULL value passed by SetNone as
124+
// '<NULL>'; failure in repr's signal check uses '<unknown>'.
125+
#[cfg(not(target_arch = "wasm32"))]
126+
let args = if crate::signal::check_signals(self).is_ok() {
127+
"<NULL>"
128+
} else {
129+
"<unknown>"
130+
};
131+
#[cfg(target_arch = "wasm32")]
132+
let args = "<NULL>";
133+
let note = self.ctx.new_str(format!(
134+
"Normalization failed: type={} args={args}",
135+
class.slot_name()
136+
));
137+
// A failed attempt to append the note must not replace the
138+
// exception raised by the class constructor.
139+
let _ = error.add_note(note, self);
140+
error
141+
}
142+
};
143+
Err(exception)
144+
}
145+
146+
/// Clear a thread-state root while its VM and frame slot are still usable.
147+
pub(crate) fn clear_async_exception(&self) {
148+
self.async_exception.close();
149+
}
150+
}

‎crates/vm/src/vm/interpreter.rs‎

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -487,6 +487,12 @@ pub struct Interpreter {
487487
vm: VirtualMachine,
488488
}
489489

490+
impl Drop for Interpreter {
491+
fn drop(&mut self) {
492+
thread::clear_async_exception_on_owner_drop(&self.vm);
493+
}
494+
}
495+
490496
impl Interpreter {
491497
/// Create a new interpreter configuration builder.
492498
///
@@ -640,7 +646,7 @@ impl Interpreter {
640646
#[cfg(feature = "threading")]
641647
#[must_use]
642648
pub fn take_main_thread(&mut self) -> Self {
643-
let replacement = self.vm.new_thread().vm;
649+
let replacement = self.vm.new_thread_vm();
644650
let vm = core::mem::replace(&mut self.vm, replacement);
645651
Self {
646652
global_state: self.global_state.clone(),
@@ -845,6 +851,7 @@ impl VirtualMachine {
845851
#[cfg(feature = "threading")]
846852
crate::object::qsbr::QSBR.process();
847853

854+
vm.clear_async_exception();
848855
flush_status
849856
}
850857
}

‎crates/vm/src/vm/mod.rs‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33
//! See also:
44
//! <https://github.com/ProgVal/pythonvm-rust/blob/master/src/processor/mod.rs>
55
6+
mod async_exception;
67
#[cfg(feature = "rustpython-compiler")]
78
mod compile;
89
pub(crate) mod compile_mode;
@@ -158,6 +159,9 @@ pub struct VirtualMachine {
158159
/// UnsafeCell because the VM is per-thread and no reference into the Vec
159160
/// is held across anything that could push to it.
160161
trampoline_stack: core::cell::UnsafeCell<Vec<SuspendedFrame>>,
162+
// Keep the closed state alive through the release of every other VM root.
163+
// A callback during field destruction must not reserve another exception.
164+
async_exception: async_exception::PendingAsyncException,
161165
}
162166

163167
/// Non-owning frame pointer for the non-unix threading frames stack.
@@ -1408,6 +1412,7 @@ impl VirtualMachine {
14081412
pending_tailcall_owner: core::cell::UnsafeCell::new(None),
14091413
pending_gen_resume: core::cell::UnsafeCell::new(None),
14101414
trampoline_stack: core::cell::UnsafeCell::new(Vec::new()),
1415+
async_exception: async_exception::PendingAsyncException::default(),
14111416
};
14121417

14131418
vm.builtins.init_dict(
@@ -3864,6 +3869,9 @@ impl VirtualMachine {
38643869
/// `check_signals` and never park, so `stop_the_world` waits forever.
38653870
#[inline]
38663871
pub(crate) fn eval_breaker_tripped(&self) -> bool {
3872+
if self.async_exception.is_pending() {
3873+
return true;
3874+
}
38673875
#[cfg(feature = "threading")]
38683876
if thread::stop_requested_for_current_thread() || self.state.gc.collection_ready() {
38693877
return true;

0 commit comments

Comments
 (0)