X Tutup
Skip to content

Commit e320cfe

Browse files
youknowoneyouknowone
committed
Optimize coroutine exception handling and fix gen_throw traceback
- Replace ExceptionStack linked list with Vec for O(1) push/pop - Introduce FramePtr (NonNull<Py<Frame>>) to avoid Arc clone in frame stack - Use FramePtr in ThreadSlot for lock-protected cross-thread frame access - Add resume_gen_frame for lightweight generator/coroutine resume - Replace Coro.exception PyMutex with PyAtomicRef for lock-free swap - Add with_frame_exc to support initial exception state for generators - Use scopeguard for panic-safe cleanup in with_frame_exc/resume_gen_frame - Add traceback entries in gen_throw delegate/close error paths - Treat EndAsyncFor and CleanupThrow as reraise in handle_exception - Update current_frame()/current_globals() to return owned values - Add ThreadSlot with atomic exception field for sys._current_exceptions() - Push/pop thread frames in resume_gen_frame for sys._current_frames()
1 parent e3c533a commit e320cfe

File tree

11 files changed

+359
-165
lines changed

11 files changed

+359
-165
lines changed

‎crates/stdlib/src/faulthandler.rs‎

Lines changed: 17 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -96,9 +96,8 @@ mod decl {
9696
all_threads: AtomicBool::new(true),
9797
};
9898

99-
/// Arc<Mutex<Vec<FrameRef>>> - shared frame slot for a thread
10099
#[cfg(feature = "threading")]
101-
type ThreadFrameSlot = Arc<parking_lot::Mutex<Vec<crate::vm::frame::FrameRef>>>;
100+
type ThreadFrameSlot = Arc<rustpython_vm::vm::thread::ThreadSlot>;
102101

103102
// Watchdog thread state for dump_traceback_later
104103
struct WatchdogState {
@@ -326,7 +325,7 @@ mod decl {
326325

327326
/// Write a frame's info to an fd using signal-safe I/O.
328327
#[cfg(any(unix, windows))]
329-
fn dump_frame_from_ref(fd: i32, frame: &crate::vm::PyRef<Frame>) {
328+
fn dump_frame_from_ref(fd: i32, frame: &crate::vm::Py<Frame>) {
330329
let funcname = frame.code.obj_name.as_str();
331330
let filename = frame.code.source_path().as_str();
332331
let lineno = if frame.lasti() == 0 {
@@ -345,20 +344,23 @@ mod decl {
345344
}
346345

347346
/// Dump traceback for a thread given its frame stack (for cross-thread dumping).
347+
/// # Safety
348+
/// Each `FramePtr` must point to a live frame (caller holds the Mutex).
348349
#[cfg(all(any(unix, windows), feature = "threading"))]
349350
fn dump_traceback_thread_frames(
350351
fd: i32,
351352
thread_id: u64,
352353
is_current: bool,
353-
frames: &[crate::vm::frame::FrameRef],
354+
frames: &[rustpython_vm::vm::FramePtr],
354355
) {
355356
write_thread_id(fd, thread_id, is_current);
356357

357358
if frames.is_empty() {
358359
puts(fd, " <no Python frame>\n");
359360
} else {
360-
for frame in frames.iter().rev() {
361-
dump_frame_from_ref(fd, frame);
361+
for fp in frames.iter().rev() {
362+
// SAFETY: caller holds the Mutex, so the owning thread can't pop.
363+
dump_frame_from_ref(fd, unsafe { fp.as_ref() });
362364
}
363365
}
364366
}
@@ -382,8 +384,9 @@ mod decl {
382384
} else {
383385
puts(fd, "Stack (most recent call first):\n");
384386
let frames = vm.frames.borrow();
385-
for frame in frames.iter().rev() {
386-
dump_frame_from_ref(fd, frame);
387+
for fp in frames.iter().rev() {
388+
// SAFETY: the frame is alive while it's in the Vec
389+
dump_frame_from_ref(fd, unsafe { fp.as_ref() });
387390
}
388391
}
389392
}
@@ -410,7 +413,7 @@ mod decl {
410413
if tid == current_tid {
411414
continue;
412415
}
413-
let frames_guard = slot.lock();
416+
let frames_guard = slot.frames.lock();
414417
dump_traceback_thread_frames(fd, tid, false, &frames_guard);
415418
puts(fd, "\n");
416419
}
@@ -421,8 +424,8 @@ mod decl {
421424
if frames.is_empty() {
422425
puts(fd, " <no Python frame>\n");
423426
} else {
424-
for frame in frames.iter().rev() {
425-
dump_frame_from_ref(fd, frame);
427+
for fp in frames.iter().rev() {
428+
dump_frame_from_ref(fd, unsafe { fp.as_ref() });
426429
}
427430
}
428431
}
@@ -431,8 +434,8 @@ mod decl {
431434
{
432435
write_thread_id(fd, current_thread_id(), true);
433436
let frames = vm.frames.borrow();
434-
for frame in frames.iter().rev() {
435-
dump_frame_from_ref(fd, frame);
437+
for fp in frames.iter().rev() {
438+
dump_frame_from_ref(fd, unsafe { fp.as_ref() });
436439
}
437440
}
438441
}
@@ -870,7 +873,7 @@ mod decl {
870873
#[cfg(feature = "threading")]
871874
{
872875
for (tid, slot) in &thread_frame_slots {
873-
let frames = slot.lock();
876+
let frames = slot.frames.lock();
874877
dump_traceback_thread_frames(fd, *tid, false, &frames);
875878
}
876879
}

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

Lines changed: 36 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@
44

55
use super::{PyCode, PyDictRef, PyIntRef, PyStrRef};
66
use crate::{
7-
AsObject, Context, Py, PyObjectRef, PyRef, PyResult, VirtualMachine,
7+
Context, Py, PyObjectRef, PyRef, PyResult, VirtualMachine,
88
class::PyClassImpl,
99
frame::{Frame, FrameOwner, FrameRef},
1010
function::PySetterValue,
@@ -195,16 +195,43 @@ impl Py<Frame> {
195195

196196
#[pygetset]
197197
pub fn f_back(&self, vm: &VirtualMachine) -> Option<PyRef<Frame>> {
198-
// TODO: actually store f_back inside Frame struct
198+
let previous = self.previous_frame();
199+
if previous.is_null() {
200+
return None;
201+
}
199202

200-
// get the frame in the frame stack that appears before this one.
201-
// won't work if this frame isn't in the frame stack, hence the todo above
202-
vm.frames
203+
if let Some(frame) = vm
204+
.frames
203205
.borrow()
204206
.iter()
205-
.rev()
206-
.skip_while(|p| !p.is(self.as_object()))
207-
.nth(1)
208-
.cloned()
207+
.find(|fp| {
208+
// SAFETY: the caller keeps the FrameRef alive while it's in the Vec
209+
let py: &crate::Py<Frame> = unsafe { fp.as_ref() };
210+
let ptr: *const Frame = &**py;
211+
core::ptr::eq(ptr, previous)
212+
})
213+
.map(|fp| unsafe { fp.as_ref() }.to_owned())
214+
{
215+
return Some(frame);
216+
}
217+
218+
#[cfg(feature = "threading")]
219+
{
220+
let registry = vm.state.thread_frames.lock();
221+
for slot in registry.values() {
222+
let frames = slot.frames.lock();
223+
// SAFETY: the owning thread can't pop while we hold the Mutex,
224+
// so FramePtr is valid for the duration of the lock.
225+
if let Some(frame) = frames.iter().find_map(|fp| {
226+
let f = unsafe { fp.as_ref() };
227+
let ptr: *const Frame = &**f;
228+
core::ptr::eq(ptr, previous).then(|| f.to_owned())
229+
}) {
230+
return Some(frame);
231+
}
232+
}
233+
}
234+
235+
None
209236
}
210237
}

‎crates/vm/src/coroutine.rs‎

Lines changed: 15 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,11 @@
11
use crate::{
22
AsObject, Py, PyObject, PyObjectRef, PyResult, TryFromObject, VirtualMachine,
3-
builtins::{PyBaseExceptionRef, PyStrRef},
3+
builtins::PyStrRef,
44
common::lock::PyMutex,
55
exceptions::types::PyBaseException,
6-
frame::{ExecutionResult, FrameOwner, FrameRef},
6+
frame::{ExecutionResult, Frame, FrameOwner, FrameRef},
77
function::OptionalArg,
8-
object::{Traverse, TraverseFn},
8+
object::{PyAtomicRef, Traverse, TraverseFn},
99
protocol::PyIterReturn,
1010
};
1111
use crossbeam_utils::atomic::AtomicCell;
@@ -36,15 +36,17 @@ pub struct Coro {
3636
// _weakreflist
3737
name: PyMutex<PyStrRef>,
3838
qualname: PyMutex<PyStrRef>,
39-
exception: PyMutex<Option<PyBaseExceptionRef>>, // exc_state
39+
exception: PyAtomicRef<Option<PyBaseException>>, // exc_state
4040
}
4141

4242
unsafe impl Traverse for Coro {
4343
fn traverse(&self, tracer_fn: &mut TraverseFn<'_>) {
4444
self.frame.traverse(tracer_fn);
4545
self.name.traverse(tracer_fn);
4646
self.qualname.traverse(tracer_fn);
47-
self.exception.traverse(tracer_fn);
47+
if let Some(exc) = self.exception.deref() {
48+
exc.traverse(tracer_fn);
49+
}
4850
}
4951
}
5052

@@ -65,7 +67,7 @@ impl Coro {
6567
frame,
6668
closed: AtomicCell::new(false),
6769
running: AtomicCell::new(false),
68-
exception: PyMutex::default(),
70+
exception: PyAtomicRef::from(None),
6971
name: PyMutex::new(name),
7072
qualname: PyMutex::new(qualname),
7173
}
@@ -92,33 +94,20 @@ impl Coro {
9294
func: F,
9395
) -> PyResult<ExecutionResult>
9496
where
95-
F: FnOnce(FrameRef) -> PyResult<ExecutionResult>,
97+
F: FnOnce(&Py<Frame>) -> PyResult<ExecutionResult>,
9698
{
9799
if self.running.compare_exchange(false, true).is_err() {
98100
return Err(vm.new_value_error(format!("{} already executing", gen_name(jen, vm))));
99101
}
100102

101-
// swap exception state
102-
// Get generator's saved exception state from last yield
103-
let gen_exc = self.exception.lock().take();
104-
105-
// Use a slot to capture generator's exception state before with_frame pops
106-
let exception_slot = &self.exception;
103+
// SAFETY: running.compare_exchange guarantees exclusive access
104+
let gen_exc = unsafe { self.exception.swap(None) };
105+
let exception_ptr = &self.exception as *const PyAtomicRef<Option<PyBaseException>>;
107106

108-
// Run the generator frame
109-
// with_frame does push_exception(None) which creates a new exception context
110-
// The caller's exception remains in the chain via prev, so topmost_exception()
111-
// will find it if generator's exception is None
112-
let result = vm.with_frame(self.frame.clone(), |f| {
113-
// with_frame pushed None, creating: { exc: None, prev: caller's exc_info }
114-
// Pop None and push generator's exception instead
115-
// This maintains the chain: { exc: gen_exc, prev: caller's exc_info }
116-
vm.pop_exception();
117-
vm.push_exception(gen_exc);
107+
let result = vm.resume_gen_frame(&self.frame, gen_exc, |f| {
118108
let result = func(f);
119-
// Save generator's exception state BEFORE with_frame pops
120-
// This is the generator's current exception context
121-
*exception_slot.lock() = vm.current_exception();
109+
// SAFETY: exclusive access guaranteed by running flag
110+
let _old = unsafe { (*exception_ptr).swap(vm.current_exception()) };
122111
result
123112
});
124113

‎crates/vm/src/frame.rs‎

Lines changed: 38 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -463,7 +463,7 @@ impl ExecutingFrame<'_> {
463463
// Execute until return or exception:
464464
let instructions = &self.code.instructions;
465465
let mut arg_state = bytecode::OpArgState::default();
466-
let mut prev_line: usize = 0;
466+
let mut prev_line: u32 = 0;
467467
loop {
468468
let idx = self.lasti() as usize;
469469
// Fire 'line' trace event when line number changes.
@@ -472,9 +472,9 @@ impl ExecutingFrame<'_> {
472472
if vm.use_tracing.get()
473473
&& !vm.is_none(&self.object.trace.lock())
474474
&& let Some((loc, _)) = self.code.locations.get(idx)
475-
&& loc.line.get() != prev_line
475+
&& loc.line.get() as u32 != prev_line
476476
{
477-
prev_line = loc.line.get();
477+
prev_line = loc.line.get() as u32;
478478
vm.trace_event(crate::protocol::TraceEvent::Line, None)?;
479479
}
480480
self.update_lasti(|i| *i += 1);
@@ -543,13 +543,16 @@ impl ExecutingFrame<'_> {
543543
// Check if this is a RERAISE instruction
544544
// Both AnyInstruction::Raise { kind: Reraise/ReraiseFromStack } and
545545
// AnyInstruction::Reraise are reraise operations that should not add
546-
// new traceback entries
546+
// new traceback entries.
547+
// EndAsyncFor and CleanupThrow also re-raise non-matching exceptions.
547548
let is_reraise = match op {
548549
Instruction::RaiseVarargs { kind } => matches!(
549550
kind.get(arg),
550551
bytecode::RaiseKind::BareRaise | bytecode::RaiseKind::ReraiseFromStack
551552
),
552-
Instruction::Reraise { .. } => true,
553+
Instruction::Reraise { .. }
554+
| Instruction::EndAsyncFor
555+
| Instruction::CleanupThrow => true,
553556
_ => false,
554557
};
555558

@@ -653,6 +656,19 @@ impl ExecutingFrame<'_> {
653656
Ok(())
654657
};
655658
if let Err(err) = close_result {
659+
let idx = self.lasti().saturating_sub(1) as usize;
660+
if idx < self.code.locations.len() {
661+
let (loc, _end_loc) = self.code.locations[idx];
662+
let next = err.__traceback__();
663+
let new_traceback = PyTraceback::new(
664+
next,
665+
self.object.to_owned(),
666+
idx as u32 * 2,
667+
loc.line,
668+
);
669+
err.set_traceback_typed(Some(new_traceback.into_ref(&vm.ctx)));
670+
}
671+
656672
self.push_value(vm.ctx.none());
657673
vm.contextualize_exception(&err);
658674
return match self.unwind_blocks(vm, UnwindReason::Raising { exception: err }) {
@@ -678,6 +694,23 @@ impl ExecutingFrame<'_> {
678694
Either::B(meth) => meth.call((exc_type, exc_val, exc_tb), vm),
679695
};
680696
return ret.map(ExecutionResult::Yield).or_else(|err| {
697+
// Add traceback entry for the yield-from/await point.
698+
// gen_send_ex2 resumes the frame with a pending exception,
699+
// which goes through error: → PyTraceBack_Here. We add the
700+
// entry here before calling unwind_blocks.
701+
let idx = self.lasti().saturating_sub(1) as usize;
702+
if idx < self.code.locations.len() {
703+
let (loc, _end_loc) = self.code.locations[idx];
704+
let next = err.__traceback__();
705+
let new_traceback = PyTraceback::new(
706+
next,
707+
self.object.to_owned(),
708+
idx as u32 * 2,
709+
loc.line,
710+
);
711+
err.set_traceback_typed(Some(new_traceback.into_ref(&vm.ctx)));
712+
}
713+
681714
self.push_value(vm.ctx.none());
682715
vm.contextualize_exception(&err);
683716
match self.unwind_blocks(vm, UnwindReason::Raising { exception: err }) {

‎crates/vm/src/protocol/callable.rs‎

Lines changed: 4 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@ use crate::{
22
builtins::{PyBoundMethod, PyFunction},
33
function::{FuncArgs, IntoFuncArgs},
44
types::GenericMethod,
5-
{AsObject, PyObject, PyObjectRef, PyResult, VirtualMachine},
5+
{PyObject, PyObjectRef, PyResult, VirtualMachine},
66
};
77

88
impl PyObject {
@@ -111,12 +111,11 @@ impl VirtualMachine {
111111
return Ok(());
112112
}
113113

114-
let frame_ref = self.current_frame();
115-
if frame_ref.is_none() {
114+
let Some(frame_ref) = self.current_frame() else {
116115
return Ok(());
117-
}
116+
};
118117

119-
let frame = frame_ref.unwrap().as_object().to_owned();
118+
let frame: PyObjectRef = frame_ref.into();
120119
let event = self.ctx.new_str(event.to_string()).into();
121120
let args = vec![frame, event, arg.unwrap_or_else(|| self.ctx.none())];
122121

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -384,7 +384,7 @@ mod builtins {
384384
)
385385
}
386386
None => (
387-
vm.current_globals().clone(),
387+
vm.current_globals(),
388388
if let Some(locals) = self.locals {
389389
locals
390390
} else {
@@ -503,7 +503,7 @@ mod builtins {
503503

504504
#[pyfunction]
505505
fn globals(vm: &VirtualMachine) -> PyDictRef {
506-
vm.current_globals().clone()
506+
vm.current_globals()
507507
}
508508

509509
#[pyfunction]

0 commit comments

Comments
 (0)
X Tutup