Skip to content

Commit 1b61f5d

Browse files
youknowoneyouknowone
committed
Extract datastack_frame_size_bytes_for_code, skip monitoring for init_cleanup frames, guard trace dispatch
- Extract datastack_frame_size_bytes_for_code as free function, use it to compute init_cleanup stack bytes instead of hardcoded constant - Add monitoring_disabled_for_code to skip instrumentation for synthetic init_cleanup code object in RESUME and execute_instrumented - Add is_trace_event guard so profile-only events skip trace_func dispatch - Reformat core.rs (rustfmt)
1 parent e84ac3c commit 1b61f5d

File tree

4 files changed

+78
-48
lines changed

4 files changed

+78
-48
lines changed

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

Lines changed: 17 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -673,20 +673,7 @@ impl Py<PyFunction> {
673673
/// Returns `None` for generator/coroutine code paths that do not push a
674674
/// regular datastack-backed frame in the fast call path.
675675
pub(crate) fn datastack_frame_size_bytes(&self) -> Option<usize> {
676-
let code: &Py<PyCode> = &self.code;
677-
if code
678-
.flags
679-
.intersects(bytecode::CodeFlags::GENERATOR | bytecode::CodeFlags::COROUTINE)
680-
{
681-
return None;
682-
}
683-
let nlocalsplus = code
684-
.varnames
685-
.len()
686-
.checked_add(code.cellvars.len())?
687-
.checked_add(code.freevars.len())?;
688-
let capacity = nlocalsplus.checked_add(code.max_stackdepth as usize)?;
689-
capacity.checked_mul(core::mem::size_of::<usize>())
676+
datastack_frame_size_bytes_for_code(&self.code)
690677
}
691678

692679
pub(crate) fn prepare_exact_args_frame(
@@ -782,6 +769,22 @@ impl Py<PyFunction> {
782769
}
783770
}
784771

772+
pub(crate) fn datastack_frame_size_bytes_for_code(code: &Py<PyCode>) -> Option<usize> {
773+
if code
774+
.flags
775+
.intersects(bytecode::CodeFlags::GENERATOR | bytecode::CodeFlags::COROUTINE)
776+
{
777+
return None;
778+
}
779+
let nlocalsplus = code
780+
.varnames
781+
.len()
782+
.checked_add(code.cellvars.len())?
783+
.checked_add(code.freevars.len())?;
784+
let capacity = nlocalsplus.checked_add(code.max_stackdepth as usize)?;
785+
capacity.checked_mul(core::mem::size_of::<usize>())
786+
}
787+
785788
impl PyPayload for PyFunction {
786789
#[inline]
787790
fn class(ctx: &Context) -> &'static Py<PyType> {

‎crates/vm/src/frame.rs‎

Lines changed: 39 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,10 @@ use crate::{
1414
builtin_func::PyNativeFunction,
1515
descriptor::{MemberGetter, PyMemberDescriptor, PyMethodDescriptor},
1616
frame::stack_analysis,
17-
function::{PyBoundMethod, PyCell, PyCellRef, PyFunction, vectorcall_function},
17+
function::{
18+
PyBoundMethod, PyCell, PyCellRef, PyFunction, datastack_frame_size_bytes_for_code,
19+
vectorcall_function,
20+
},
1821
list::PyListIterator,
1922
range::PyRangeIterator,
2023
tuple::{PyTuple, PyTupleIterator, PyTupleRef},
@@ -1259,6 +1262,11 @@ impl fmt::Debug for ExecutingFrame<'_> {
12591262
}
12601263

12611264
impl ExecutingFrame<'_> {
1265+
#[inline]
1266+
fn monitoring_disabled_for_code(&self, vm: &VirtualMachine) -> bool {
1267+
self.code.is(&vm.ctx.init_cleanup_code)
1268+
}
1269+
12621270
fn specialization_new_init_cleanup_frame(&self, vm: &VirtualMachine) -> FrameRef {
12631271
Frame::new(
12641272
vm.ctx.init_cleanup_code.clone(),
@@ -3146,6 +3154,17 @@ impl ExecutingFrame<'_> {
31463154
self.code.instructions.quicken();
31473155
atomic::fence(atomic::Ordering::Release);
31483156
}
3157+
if self.monitoring_disabled_for_code(vm) {
3158+
let global_ver = vm
3159+
.state
3160+
.instrumentation_version
3161+
.load(atomic::Ordering::Acquire);
3162+
monitoring::instrument_code(self.code, 0);
3163+
self.code
3164+
.instrumentation_version
3165+
.store(global_ver, atomic::Ordering::Release);
3166+
return Ok(None);
3167+
}
31493168
// Check if bytecode needs re-instrumentation
31503169
let global_ver = vm
31513170
.state
@@ -4757,22 +4776,16 @@ impl ExecutingFrame<'_> {
47574776
&& let Some(init_func) = cls.get_cached_init_for_specialization(cached_version)
47584777
&& let Some(cls_alloc) = cls.slots.alloc.load()
47594778
{
4760-
// CPython guard is:
4761-
// code->co_framesize + _Py_InitCleanup.co_framesize
4762-
// and _Py_InitCleanup.co_framesize is defined as:
4763-
// 2 + FRAME_SPECIALS_SIZE
4764-
// (see cpython/Python/specialize.c).
4765-
//
4766-
// RustPython datastack stores localsplus payload only, but we
4767-
// keep the guard conservative and CPython-shaped by accounting
4768-
// for the full cleanup frame size.
4769-
const CPYTHON_INIT_CLEANUP_CO_FRAMESIZE_SLOTS: usize = 11;
4770-
const INIT_CLEANUP_STACK_BYTES: usize =
4771-
CPYTHON_INIT_CLEANUP_CO_FRAMESIZE_SLOTS * core::mem::size_of::<usize>();
4779+
// Match CPython's `code->co_framesize + _Py_InitCleanup.co_framesize`
4780+
// shape, using RustPython's datastack-backed frame size
4781+
// equivalent for the extra shim frame.
4782+
let init_cleanup_stack_bytes =
4783+
datastack_frame_size_bytes_for_code(&vm.ctx.init_cleanup_code)
4784+
.expect("_Py_InitCleanup shim is not a generator/coroutine");
47724785
if !self.specialization_has_datastack_space_for_func_with_extra(
47734786
vm,
47744787
&init_func,
4775-
INIT_CLEANUP_STACK_BYTES,
4788+
init_cleanup_stack_bytes,
47764789
) {
47774790
return self.execute_call_vectorcall(nargs, vm);
47784791
}
@@ -5586,6 +5599,18 @@ impl ExecutingFrame<'_> {
55865599
instruction.is_instrumented(),
55875600
"execute_instrumented called with non-instrumented opcode {instruction:?}"
55885601
);
5602+
if self.monitoring_disabled_for_code(vm) {
5603+
let global_ver = vm
5604+
.state
5605+
.instrumentation_version
5606+
.load(atomic::Ordering::Acquire);
5607+
monitoring::instrument_code(self.code, 0);
5608+
self.code
5609+
.instrumentation_version
5610+
.store(global_ver, atomic::Ordering::Release);
5611+
self.update_lasti(|i| *i -= 1);
5612+
return Ok(None);
5613+
}
55895614
self.monitoring_mask = vm.state.monitoring_events.load();
55905615
match instruction {
55915616
Instruction::InstrumentedResume => {

‎crates/vm/src/object/core.rs‎

Lines changed: 12 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -331,8 +331,7 @@ const _: () = assert!(core::mem::align_of::<ObjExt>() >= core::mem::align_of::<P
331331
const _: () = assert!(
332332
core::mem::size_of::<WeakRefList>().is_multiple_of(core::mem::align_of::<WeakRefList>())
333333
);
334-
const _: () =
335-
assert!(core::mem::align_of::<WeakRefList>() >= core::mem::align_of::<PyInner<()>>());
334+
const _: () = assert!(core::mem::align_of::<WeakRefList>() >= core::mem::align_of::<PyInner<()>>());
336335

337336
/// This is an actual python object. It consists of a `typ` which is the
338337
/// python class, and carries some rust payload optionally. This rust
@@ -375,8 +374,7 @@ impl<T> PyInner<T> {
375374
#[inline(always)]
376375
pub(super) fn ext_ref(&self) -> Option<&ObjExt> {
377376
let (flags, member_count) = self.read_type_flags();
378-
let has_ext =
379-
flags.has_feature(crate::types::PyTypeFlags::HAS_DICT) || member_count > 0;
377+
let has_ext = flags.has_feature(crate::types::PyTypeFlags::HAS_DICT) || member_count > 0;
380378
if !has_ext {
381379
return None;
382380
}
@@ -387,8 +385,7 @@ impl<T> PyInner<T> {
387385
EXT_OFFSET
388386
};
389387
let self_addr = (self as *const Self as *const u8).addr();
390-
let ext_ptr =
391-
core::ptr::with_exposed_provenance::<ObjExt>(self_addr.wrapping_sub(offset));
388+
let ext_ptr = core::ptr::with_exposed_provenance::<ObjExt>(self_addr.wrapping_sub(offset));
392389
Some(unsafe { &*ext_ptr })
393390
}
394391

@@ -957,8 +954,8 @@ impl<T: PyPayload> PyInner<T> {
957954
unsafe fn dealloc(ptr: *mut Self) {
958955
unsafe {
959956
let (flags, member_count) = (*ptr).read_type_flags();
960-
let has_ext = flags.has_feature(crate::types::PyTypeFlags::HAS_DICT)
961-
|| member_count > 0;
957+
let has_ext =
958+
flags.has_feature(crate::types::PyTypeFlags::HAS_DICT) || member_count > 0;
962959
let has_weakref = flags.has_feature(crate::types::PyTypeFlags::HAS_WEAKREF);
963960

964961
if has_ext || has_weakref {
@@ -977,9 +974,8 @@ impl<T: PyPayload> PyInner<T> {
977974
.unwrap()
978975
.0;
979976
}
980-
let (combined, inner_offset) = layout
981-
.extend(core::alloc::Layout::new::<Self>())
982-
.unwrap();
977+
let (combined, inner_offset) =
978+
layout.extend(core::alloc::Layout::new::<Self>()).unwrap();
983979
let combined = combined.pad_to_align();
984980

985981
let alloc_ptr = (ptr as *mut u8).sub(inner_offset);
@@ -1045,9 +1041,8 @@ impl<T: PyPayload + core::fmt::Debug> PyInner<T> {
10451041
None
10461042
};
10471043

1048-
let (combined, inner_offset) = layout
1049-
.extend(core::alloc::Layout::new::<Self>())
1050-
.unwrap();
1044+
let (combined, inner_offset) =
1045+
layout.extend(core::alloc::Layout::new::<Self>()).unwrap();
10511046
let combined = combined.pad_to_align();
10521047

10531048
let alloc_ptr = unsafe { alloc::alloc::alloc(combined) };
@@ -1757,8 +1752,7 @@ impl PyObject {
17571752
// the pointer without clearing dict contents. The dict may still be
17581753
// referenced by other live objects (e.g. function.__globals__).
17591754
let (flags, member_count) = obj.0.read_type_flags();
1760-
let has_ext = flags.has_feature(crate::types::PyTypeFlags::HAS_DICT)
1761-
|| member_count > 0;
1755+
let has_ext = flags.has_feature(crate::types::PyTypeFlags::HAS_DICT) || member_count > 0;
17621756
if has_ext {
17631757
let has_weakref = flags.has_feature(crate::types::PyTypeFlags::HAS_WEAKREF);
17641758
let offset = if has_weakref {
@@ -1767,9 +1761,8 @@ impl PyObject {
17671761
EXT_OFFSET
17681762
};
17691763
let self_addr = (ptr as *const u8).addr();
1770-
let ext_ptr = core::ptr::with_exposed_provenance_mut::<ObjExt>(
1771-
self_addr.wrapping_sub(offset),
1772-
);
1764+
let ext_ptr =
1765+
core::ptr::with_exposed_provenance_mut::<ObjExt>(self_addr.wrapping_sub(offset));
17731766
let ext = unsafe { &mut *ext_ptr };
17741767
if let Some(old_dict) = ext.dict.take() {
17751768
// Get the dict ref before dropping InstanceDict

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

Lines changed: 10 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -146,6 +146,14 @@ pub(crate) enum TraceEvent {
146146
}
147147

148148
impl TraceEvent {
149+
/// Whether sys.settrace receives this event.
150+
fn is_trace_event(&self) -> bool {
151+
matches!(
152+
self,
153+
Self::Call | Self::Return | Self::Exception | Self::Line | Self::Opcode
154+
)
155+
}
156+
149157
/// Whether sys.setprofile receives this event.
150158
/// In legacy_tracing.c, profile callbacks are only registered for
151159
/// PY_RETURN, PY_UNWIND, C_CALL, C_RETURN, C_RAISE.
@@ -211,6 +219,7 @@ impl VirtualMachine {
211219
return Ok(None);
212220
}
213221

222+
let is_trace_event = event.is_trace_event();
214223
let is_profile_event = event.is_profile_event();
215224
let is_opcode_event = event.is_opcode_event();
216225

@@ -231,7 +240,7 @@ impl VirtualMachine {
231240

232241
// temporarily disable tracing, during the call to the
233242
// tracing function itself.
234-
if !self.is_none(&trace_func) {
243+
if is_trace_event && !self.is_none(&trace_func) {
235244
self.use_tracing.set(false);
236245
let res = trace_func.call(args.clone(), self);
237246
self.use_tracing.set(true);

0 commit comments

Comments
 (0)