diff --git a/crates/compiler-core/src/bytecode.rs b/crates/compiler-core/src/bytecode.rs index a6d197950..522bc7c32 100644 --- a/crates/compiler-core/src/bytecode.rs +++ b/crates/compiler-core/src/bytecode.rs @@ -99,6 +99,34 @@ pub fn find_exception_handler(table: &[u8], offset: u32) -> Option Vec { + let mut entries = Vec::new(); + let mut pos = 0; + while pos < table.len() { + let Some(start) = read_varint_with_start(table, &mut pos) else { + break; + }; + let Some(size) = read_varint(table, &mut pos) else { + break; + }; + let Some(target) = read_varint(table, &mut pos) else { + break; + }; + let Some(depth_lasti) = read_varint(table, &mut pos) else { + break; + }; + entries.push(ExceptionTableEntry { + start, + end: start + size, + target, + depth: (depth_lasti >> 1) as u16, + push_lasti: (depth_lasti & 1) != 0, + }); + } + entries +} + /// CPython 3.11+ linetable location info codes #[derive(Copy, Clone, Debug, PartialEq, Eq)] #[repr(u8)] @@ -370,6 +398,23 @@ impl Deref for CodeUnits { } } +impl CodeUnits { + /// Replace the opcode at `index` in-place without changing the arg byte. + /// + /// # Safety + /// Caller must ensure `index` is in bounds and `new_op` has the same + /// arg semantics as the original opcode. + pub unsafe fn replace_op(&self, index: usize, new_op: Instruction) { + unsafe { + let ptr = self.0.as_ptr() as *mut CodeUnit; + let unit_ptr = ptr.add(index); + // Write only the opcode byte (first byte of CodeUnit due to #[repr(C)]) + let op_ptr = unit_ptr as *mut u8; + core::ptr::write_volatile(op_ptr, new_op.into()); + } + } +} + /// A Constant (which usually encapsulates data within it) /// /// # Examples diff --git a/crates/compiler-core/src/bytecode/instruction.rs b/crates/compiler-core/src/bytecode/instruction.rs index 6f43e79c1..fab01c29f 100644 --- a/crates/compiler-core/src/bytecode/instruction.rs +++ b/crates/compiler-core/src/bytecode/instruction.rs @@ -419,6 +419,100 @@ impl TryFrom for Instruction { } } +impl Instruction { + /// Returns `true` if this is any instrumented opcode + /// (regular INSTRUMENTED_*, INSTRUMENTED_LINE, or INSTRUMENTED_INSTRUCTION). + pub fn is_instrumented(self) -> bool { + self.to_base().is_some() + || matches!(self, Self::InstrumentedLine | Self::InstrumentedInstruction) + } + + /// Map a base opcode to its INSTRUMENTED_* variant. + /// Returns `None` if this opcode has no instrumented counterpart. + /// + /// # Panics (debug) + /// Panics if called on an already-instrumented opcode. + pub fn to_instrumented(self) -> Option { + debug_assert!( + self.to_base().is_none(), + "to_instrumented called on already-instrumented opcode {self:?}" + ); + Some(match self { + Self::Resume { .. } => Self::InstrumentedResume, + Self::ReturnValue => Self::InstrumentedReturnValue, + Self::YieldValue { .. } => Self::InstrumentedYieldValue, + Self::Call { .. } => Self::InstrumentedCall, + Self::CallKw { .. } => Self::InstrumentedCallKw, + Self::CallFunctionEx => Self::InstrumentedCallFunctionEx, + Self::LoadSuperAttr { .. } => Self::InstrumentedLoadSuperAttr, + Self::JumpForward { .. } => Self::InstrumentedJumpForward, + Self::JumpBackward { .. } => Self::InstrumentedJumpBackward, + Self::ForIter { .. } => Self::InstrumentedForIter, + Self::EndFor => Self::InstrumentedEndFor, + Self::EndSend => Self::InstrumentedEndSend, + Self::PopJumpIfTrue { .. } => Self::InstrumentedPopJumpIfTrue, + Self::PopJumpIfFalse { .. } => Self::InstrumentedPopJumpIfFalse, + Self::PopJumpIfNone { .. } => Self::InstrumentedPopJumpIfNone, + Self::PopJumpIfNotNone { .. } => Self::InstrumentedPopJumpIfNotNone, + Self::NotTaken => Self::InstrumentedNotTaken, + Self::PopIter => Self::InstrumentedPopIter, + Self::EndAsyncFor => Self::InstrumentedEndAsyncFor, + _ => return None, + }) + } + + /// Map an INSTRUMENTED_* opcode back to its base variant. + /// Returns `None` if this is not an instrumented opcode. + /// + /// The returned base opcode uses `Arg::marker()` for typed fields — + /// only the opcode byte matters since `replace_op` preserves the arg byte. + /// + /// # Panics (debug) + /// Panics if called on a base opcode that has an instrumented counterpart. + pub fn to_base(self) -> Option { + Some(match self { + Self::InstrumentedResume => Self::Resume { arg: Arg::marker() }, + Self::InstrumentedReturnValue => Self::ReturnValue, + Self::InstrumentedYieldValue => Self::YieldValue { arg: Arg::marker() }, + Self::InstrumentedCall => Self::Call { + nargs: Arg::marker(), + }, + Self::InstrumentedCallKw => Self::CallKw { + nargs: Arg::marker(), + }, + Self::InstrumentedCallFunctionEx => Self::CallFunctionEx, + Self::InstrumentedLoadSuperAttr => Self::LoadSuperAttr { arg: Arg::marker() }, + Self::InstrumentedJumpForward => Self::JumpForward { + target: Arg::marker(), + }, + Self::InstrumentedJumpBackward => Self::JumpBackward { + target: Arg::marker(), + }, + Self::InstrumentedForIter => Self::ForIter { + target: Arg::marker(), + }, + Self::InstrumentedEndFor => Self::EndFor, + Self::InstrumentedEndSend => Self::EndSend, + Self::InstrumentedPopJumpIfTrue => Self::PopJumpIfTrue { + target: Arg::marker(), + }, + Self::InstrumentedPopJumpIfFalse => Self::PopJumpIfFalse { + target: Arg::marker(), + }, + Self::InstrumentedPopJumpIfNone => Self::PopJumpIfNone { + target: Arg::marker(), + }, + Self::InstrumentedPopJumpIfNotNone => Self::PopJumpIfNotNone { + target: Arg::marker(), + }, + Self::InstrumentedNotTaken => Self::NotTaken, + Self::InstrumentedPopIter => Self::PopIter, + Self::InstrumentedEndAsyncFor => Self::EndAsyncFor, + _ => return None, + }) + } +} + impl InstructionMetadata for Instruction { #[inline] fn label_arg(&self) -> Option> { diff --git a/crates/vm/src/builtins/code.rs b/crates/vm/src/builtins/code.rs index fcaff6879..5bb92db9d 100644 --- a/crates/vm/src/builtins/code.rs +++ b/crates/vm/src/builtins/code.rs @@ -1,6 +1,7 @@ //! Infamous code object. The python class `code` use super::{PyBytesRef, PyStrRef, PyTupleRef, PyType}; +use crate::common::lock::PyMutex; use crate::{ AsObject, Context, Py, PyObject, PyObjectRef, PyPayload, PyRef, PyResult, VirtualMachine, builtins::PyStrInterned, @@ -15,7 +16,7 @@ use alloc::fmt; use core::{ borrow::Borrow, ops::Deref, - sync::atomic::{AtomicPtr, Ordering}, + sync::atomic::{AtomicPtr, AtomicU64, Ordering}, }; use malachite_bigint::BigInt; use num_traits::Zero; @@ -324,10 +325,27 @@ impl> IntoCodeObject for frozen::FrozenCodeObject { } } +/// Per-code-object monitoring data (_PyCoMonitoringData). +/// Stores original opcodes displaced by INSTRUMENTED_LINE / INSTRUMENTED_INSTRUCTION. +pub struct CoMonitoringData { + /// Original opcodes at positions with INSTRUMENTED_LINE. + /// Indexed by instruction index. 0 = not instrumented for LINE. + pub line_opcodes: Vec, + + /// Original opcodes at positions with INSTRUMENTED_INSTRUCTION. + /// Indexed by instruction index. 0 = not instrumented for INSTRUCTION. + pub per_instruction_opcodes: Vec, +} + #[pyclass(module = false, name = "code")] pub struct PyCode { pub code: CodeObject, source_path: AtomicPtr, + /// Version counter for lazy re-instrumentation. + /// Compared against `PyGlobalState::instrumentation_version` at RESUME. + pub instrumentation_version: AtomicU64, + /// Side-table for INSTRUMENTED_LINE / INSTRUMENTED_INSTRUCTION. + pub monitoring_data: PyMutex>, } impl Deref for PyCode { @@ -343,6 +361,8 @@ impl PyCode { Self { code, source_path: AtomicPtr::new(sp), + instrumentation_version: AtomicU64::new(0), + monitoring_data: PyMutex::new(None), } } diff --git a/crates/vm/src/frame.rs b/crates/vm/src/frame.rs index 9425a533a..7f9d335cf 100644 --- a/crates/vm/src/frame.rs +++ b/crates/vm/src/frame.rs @@ -370,6 +370,7 @@ impl Py { object: self, state: &mut state, monitoring_mask: 0, + prev_line: 0, }; f(exec) } @@ -416,6 +417,7 @@ impl Py { object: self, state: &mut state, monitoring_mask: 0, + prev_line: 0, }; exec.yield_from_target().map(PyObject::to_owned) } @@ -452,6 +454,8 @@ struct ExecutingFrame<'a> { state: &'a mut FrameState, /// Cached monitoring events mask. Reloaded at Resume instruction only, monitoring_mask: u32, + /// Previous line number for LINE event suppression. + prev_line: u32, } impl fmt::Debug for ExecutingFrame<'_> { @@ -501,8 +505,6 @@ impl ExecutingFrame<'_> { // Execute until return or exception: let instructions = &self.code.instructions; let mut arg_state = bytecode::OpArgState::default(); - let mut prev_line: u32 = 0; - self.monitoring_mask = vm.state.monitoring_events.load(); loop { let idx = self.lasti() as usize; // Fire 'line' trace event when line number changes. @@ -511,9 +513,9 @@ impl ExecutingFrame<'_> { if vm.use_tracing.get() && !vm.is_none(&self.object.trace.lock()) && let Some((loc, _)) = self.code.locations.get(idx) - && loc.line.get() as u32 != prev_line + && loc.line.get() as u32 != self.prev_line { - prev_line = loc.line.get() as u32; + self.prev_line = loc.line.get() as u32; vm.trace_event(crate::protocol::TraceEvent::Line, None)?; } self.update_lasti(|i| *i += 1); @@ -521,22 +523,15 @@ impl ExecutingFrame<'_> { let arg = arg_state.extend(arg); let mut do_extend_arg = false; - // Always track current line for LINE monitoring, even when - // monitoring is off. This prevents spurious LINE events when - // monitoring is enabled mid-function. - if !matches!(op, Instruction::Resume { .. } | Instruction::ExtendedArg) { - let line = self.code.locations[idx].0.line.get() as u32; - if self.monitoring_mask != 0 { - use crate::stdlib::sys::monitoring; - let offset = idx as u32 * 2; - if self.monitoring_mask & monitoring::EVENT_LINE != 0 && line != prev_line { - monitoring::fire_line(vm, self.code, offset, line)?; - } - if self.monitoring_mask & monitoring::EVENT_INSTRUCTION != 0 { - monitoring::fire_instruction(vm, self.code, offset)?; - } - } - prev_line = line; + // Track current line for settrace LINE event and monitoring. + if !matches!( + op, + Instruction::Resume { .. } + | Instruction::ExtendedArg + | Instruction::InstrumentedLine + ) && let Some((loc, _)) = self.code.locations.get(idx) + { + self.prev_line = loc.line.get() as u32; } let result = self.execute_instruction(op, arg, &mut do_extend_arg, vm); @@ -632,8 +627,9 @@ impl ExecutingFrame<'_> { // replace the original exception. let exception = { use crate::stdlib::sys::monitoring; + let mon_events = vm.state.monitoring_events.load(); if is_reraise { - if self.monitoring_mask & monitoring::EVENT_RERAISE != 0 { + if mon_events & monitoring::EVENT_RERAISE != 0 { let offset = idx as u32 * 2; let exc_obj: PyObjectRef = exception.clone().into(); match monitoring::fire_reraise(vm, self.code, offset, &exc_obj) { @@ -643,7 +639,7 @@ impl ExecutingFrame<'_> { } else { exception } - } else if self.monitoring_mask & monitoring::EVENT_RAISE != 0 { + } else if mon_events & monitoring::EVENT_RAISE != 0 { let offset = idx as u32 * 2; let exc_obj: PyObjectRef = exception.clone().into(); match monitoring::fire_raise(vm, self.code, offset, &exc_obj) { @@ -660,7 +656,7 @@ impl ExecutingFrame<'_> { Ok(Some(result)) => break Ok(result), Err(exception) => { // Fire PY_UNWIND: exception escapes this frame - let exception = if self.monitoring_mask + let exception = if vm.state.monitoring_events.load() & crate::stdlib::sys::monitoring::EVENT_PY_UNWIND != 0 { @@ -858,7 +854,8 @@ impl ExecutingFrame<'_> { // when a callback fails, so we mirror that here. let exception = { use crate::stdlib::sys::monitoring; - let exception = if self.monitoring_mask & monitoring::EVENT_PY_THROW != 0 { + let mon_events = vm.state.monitoring_events.load(); + let exception = if mon_events & monitoring::EVENT_PY_THROW != 0 { let offset = idx as u32 * 2; let exc_obj: PyObjectRef = exception.clone().into(); match monitoring::fire_py_throw(vm, self.code, offset, &exc_obj) { @@ -868,7 +865,7 @@ impl ExecutingFrame<'_> { } else { exception }; - if self.monitoring_mask & monitoring::EVENT_RAISE != 0 { + if mon_events & monitoring::EVENT_RAISE != 0 { let offset = idx as u32 * 2; let exc_obj: PyObjectRef = exception.clone().into(); match monitoring::fire_raise(vm, self.code, offset, &exc_obj) { @@ -894,7 +891,7 @@ impl ExecutingFrame<'_> { Err(exception) => { // Fire PY_UNWIND: exception escapes the generator frame. // do_monitor_exc replaces the exception on callback failure. - let exception = if self.monitoring_mask + let exception = if vm.state.monitoring_events.load() & crate::stdlib::sys::monitoring::EVENT_PY_UNWIND != 0 { @@ -1313,7 +1310,10 @@ impl ExecutingFrame<'_> { *extend_arg = true; Ok(None) } - Instruction::ForIter { target } => self.execute_for_iter(vm, target.get(arg)), + Instruction::ForIter { target } => { + self.execute_for_iter(vm, target.get(arg))?; + Ok(None) + } Instruction::FormatSimple => { let value = self.pop_value(); let formatted = vm.format(&value, vm.ctx.new_str(""))?; @@ -1489,18 +1489,7 @@ impl ExecutingFrame<'_> { Ok(None) } Instruction::JumpBackward { target } => { - let src_offset = (self.lasti() - 1) * 2; - let dest = target.get(arg); - self.jump(dest); - // JUMP events fire only for backward jumps (loop iterations) - if self.monitoring_mask & crate::stdlib::sys::monitoring::EVENT_JUMP != 0 { - crate::stdlib::sys::monitoring::fire_jump( - vm, - self.code, - src_offset, - dest.0 * 2, - )?; - } + self.jump(target.get(arg)); Ok(None) } Instruction::JumpBackwardNoInterrupt { target } => { @@ -2070,8 +2059,6 @@ impl ExecutingFrame<'_> { Instruction::Nop => Ok(None), // NOT_TAKEN is a branch prediction hint - functionally a NOP Instruction::NotTaken => Ok(None), - // Instrumented version of NOT_TAKEN - NOP without monitoring - Instruction::InstrumentedNotTaken => Ok(None), // CACHE is used by adaptive interpreter for inline caching - NOP for us Instruction::Cache => Ok(None), Instruction::ReturnGenerator => { @@ -2105,24 +2092,16 @@ impl ExecutingFrame<'_> { Instruction::PopJumpIfTrue { target } => self.pop_jump_if(vm, target.get(arg), true), Instruction::PopJumpIfNone { target } => { let value = self.pop_value(); - let branch_taken = vm.is_none(&value); - let target = target.get(arg); - let src_offset = (self.lasti() - 1) * 2; - if branch_taken { - self.jump(target); + if vm.is_none(&value) { + self.jump(target.get(arg)); } - self.fire_branch_event(vm, src_offset, branch_taken, target)?; Ok(None) } Instruction::PopJumpIfNotNone { target } => { let value = self.pop_value(); - let branch_taken = !vm.is_none(&value); - let target = target.get(arg); - let src_offset = (self.lasti() - 1) * 2; - if branch_taken { - self.jump(target); + if !vm.is_none(&value) { + self.jump(target.get(arg)); } - self.fire_branch_event(vm, src_offset, branch_taken, target)?; Ok(None) } Instruction::PopTop => { @@ -2146,26 +2125,29 @@ impl ExecutingFrame<'_> { Ok(None) } Instruction::RaiseVarargs { kind } => self.execute_raise(vm, kind.get(arg)), - Instruction::Resume { arg: resume_arg } => { - use crate::stdlib::sys::monitoring; - let resume_type = resume_arg.get(arg); - self.monitoring_mask = vm.state.monitoring_events.load(); - let offset = (self.lasti() - 1) * 2; - if resume_type == 0 { - if self.monitoring_mask & monitoring::EVENT_PY_START != 0 { - monitoring::fire_py_start(vm, self.code, offset)?; - } - } else if self.monitoring_mask & monitoring::EVENT_PY_RESUME != 0 { - monitoring::fire_py_resume(vm, self.code, offset)?; + Instruction::Resume { .. } => { + // Check if bytecode needs re-instrumentation + let global_ver = vm + .state + .instrumentation_version + .load(atomic::Ordering::Acquire); + let code_ver = self + .code + .instrumentation_version + .load(atomic::Ordering::Acquire); + if code_ver != global_ver { + let events = vm.state.monitoring_events.load(); + crate::stdlib::sys::monitoring::instrument_code(self.code, events); + self.code + .instrumentation_version + .store(global_ver, atomic::Ordering::Release); + // Re-execute this instruction (it may now be INSTRUMENTED_RESUME) + self.update_lasti(|i| *i -= 1); } Ok(None) } Instruction::ReturnValue => { let value = self.pop_value(); - if self.monitoring_mask & crate::stdlib::sys::monitoring::EVENT_PY_RETURN != 0 { - let offset = (self.lasti() - 1) * 2; - crate::stdlib::sys::monitoring::fire_py_return(vm, self.code, offset, &value)?; - } self.unwind_blocks(vm, UnwindReason::Returning { value }) } Instruction::SetAdd { i } => { @@ -2388,10 +2370,6 @@ impl ExecutingFrame<'_> { } Instruction::YieldValue { arg: oparg } => { let value = self.pop_value(); - if self.monitoring_mask & crate::stdlib::sys::monitoring::EVENT_PY_YIELD != 0 { - let offset = (self.lasti() - 1) * 2; - crate::stdlib::sys::monitoring::fire_py_yield(vm, self.code, offset, &value)?; - } // arg=0: direct yield (wrapped for async generators) // arg=1: yield from await/yield-from (NOT wrapped) let wrap = oparg.get(arg) == 0; @@ -2499,13 +2477,52 @@ impl ExecutingFrame<'_> { self.push_value(vm.ctx.new_bool(!value).into()); Ok(None) } - // --- Instrumented opcode handlers --- - // Each instrumented variant delegates to the same logic as its base opcode. - // Monitoring events are already fired from the base opcode implementations. + // All INSTRUMENTED_* opcodes delegate to a cold function to keep + // the hot instruction loop free of monitoring overhead. + _ => self.execute_instrumented(instruction, arg, vm), + } + } + + /// Handle all INSTRUMENTED_* opcodes. This function is cold — it only + /// runs when sys.monitoring has rewritten the bytecode. + #[cold] + fn execute_instrumented( + &mut self, + instruction: Instruction, + arg: bytecode::OpArg, + vm: &VirtualMachine, + ) -> FrameResult { + debug_assert!( + instruction.is_instrumented(), + "execute_instrumented called with non-instrumented opcode {instruction:?}" + ); + // Refresh monitoring mask from global state on every instrumented opcode + // execution. This ensures frames already past RESUME pick up events that + // were enabled by a set_events() call while the frame was executing. + self.monitoring_mask = vm.state.monitoring_events.load(); + use crate::stdlib::sys::monitoring; + match instruction { Instruction::InstrumentedResume => { - use crate::stdlib::sys::monitoring; + // Version check: re-instrument if stale + let global_ver = vm + .state + .instrumentation_version + .load(atomic::Ordering::Acquire); + let code_ver = self + .code + .instrumentation_version + .load(atomic::Ordering::Acquire); + if code_ver != global_ver { + let events = vm.state.monitoring_events.load(); + monitoring::instrument_code(self.code, events); + self.code + .instrumentation_version + .store(global_ver, atomic::Ordering::Release); + // Re-execute (may have been de-instrumented to base Resume) + self.update_lasti(|i| *i -= 1); + return Ok(None); + } let resume_type = u32::from(arg); - self.monitoring_mask = vm.state.monitoring_events.load(); let offset = (self.lasti() - 1) * 2; if resume_type == 0 { if self.monitoring_mask & monitoring::EVENT_PY_START != 0 { @@ -2518,17 +2535,17 @@ impl ExecutingFrame<'_> { } Instruction::InstrumentedReturnValue => { let value = self.pop_value(); - if self.monitoring_mask & crate::stdlib::sys::monitoring::EVENT_PY_RETURN != 0 { + if self.monitoring_mask & monitoring::EVENT_PY_RETURN != 0 { let offset = (self.lasti() - 1) * 2; - crate::stdlib::sys::monitoring::fire_py_return(vm, self.code, offset, &value)?; + monitoring::fire_py_return(vm, self.code, offset, &value)?; } self.unwind_blocks(vm, UnwindReason::Returning { value }) } Instruction::InstrumentedYieldValue => { let value = self.pop_value(); - if self.monitoring_mask & crate::stdlib::sys::monitoring::EVENT_PY_YIELD != 0 { + if self.monitoring_mask & monitoring::EVENT_PY_YIELD != 0 { let offset = (self.lasti() - 1) * 2; - crate::stdlib::sys::monitoring::fire_py_yield(vm, self.code, offset, &value)?; + monitoring::fire_py_yield(vm, self.code, offset, &value)?; } let oparg = u32::from(arg); let wrap = oparg == 0; @@ -2541,90 +2558,201 @@ impl ExecutingFrame<'_> { } Instruction::InstrumentedCall => { let args = self.collect_positional_args(u32::from(arg)); - self.execute_call(args, vm) + self.execute_call_instrumented(args, vm) } Instruction::InstrumentedCallKw => { let args = self.collect_keyword_args(u32::from(arg)); - self.execute_call(args, vm) + self.execute_call_instrumented(args, vm) } Instruction::InstrumentedCallFunctionEx => { let args = self.collect_ex_args(vm)?; - self.execute_call(args, vm) + self.execute_call_instrumented(args, vm) } Instruction::InstrumentedLoadSuperAttr => { - let oparg = bytecode::LoadSuperAttr::try_from(u32::from(arg)) - .map_err(|_| vm.new_value_error("invalid oparg".to_owned()))?; - self.load_super_attr(vm, oparg) + let oparg = bytecode::LoadSuperAttr::from(u32::from(arg)); + let offset = (self.lasti() - 1) * 2; + // Fire CALL event before super() call + let call_args = if self.monitoring_mask & monitoring::EVENT_CALL != 0 { + let global_super: PyObjectRef = self.nth_value(2).to_owned(); + let arg0 = if oparg.has_class() { + self.nth_value(1).to_owned() + } else { + monitoring::get_missing(vm) + }; + monitoring::fire_call(vm, self.code, offset, &global_super, arg0.clone())?; + Some((global_super, arg0)) + } else { + None + }; + match self.load_super_attr(vm, oparg) { + Ok(result) => { + // Fire C_RETURN on success + if let Some((global_super, arg0)) = call_args { + monitoring::fire_c_return(vm, self.code, offset, &global_super, arg0)?; + } + Ok(result) + } + Err(exc) => { + // Fire C_RAISE on failure + if let Some((global_super, arg0)) = call_args { + let _ = monitoring::fire_c_raise( + vm, + self.code, + offset, + &global_super, + arg0, + ); + } + Err(exc) + } + } } Instruction::InstrumentedJumpForward => { + let src_offset = (self.lasti() - 1) * 2; let target = bytecode::Label::from(u32::from(arg)); self.jump(target); + if self.monitoring_mask & monitoring::EVENT_JUMP != 0 { + monitoring::fire_jump(vm, self.code, src_offset, target.0 * 2)?; + } Ok(None) } Instruction::InstrumentedJumpBackward => { let src_offset = (self.lasti() - 1) * 2; let dest = bytecode::Label::from(u32::from(arg)); self.jump(dest); - if self.monitoring_mask & crate::stdlib::sys::monitoring::EVENT_JUMP != 0 { - crate::stdlib::sys::monitoring::fire_jump( - vm, - self.code, - src_offset, - dest.0 * 2, - )?; + if self.monitoring_mask & monitoring::EVENT_JUMP != 0 { + monitoring::fire_jump(vm, self.code, src_offset, dest.0 * 2)?; } Ok(None) } Instruction::InstrumentedForIter => { + let src_offset = (self.lasti() - 1) * 2; let target = bytecode::Label::from(u32::from(arg)); - self.execute_for_iter(vm, target) + let continued = self.execute_for_iter(vm, target)?; + if continued { + if self.monitoring_mask & monitoring::EVENT_BRANCH_LEFT != 0 { + let dest_offset = self.lasti() * 2; + monitoring::fire_branch_left(vm, self.code, src_offset, dest_offset)?; + } + } else if self.monitoring_mask & monitoring::EVENT_BRANCH_RIGHT != 0 { + let dest_offset = self.lasti() * 2; + monitoring::fire_branch_right(vm, self.code, src_offset, dest_offset)?; + } + Ok(None) } Instruction::InstrumentedEndFor => { - self.pop_value(); + // Stack: [value, receiver(iter), ...] + // PyGen_Check: only fire STOP_ITERATION for generators + let is_gen = self + .nth_value(1) + .downcast_ref::() + .is_some(); + let value = self.pop_value(); + if is_gen && self.monitoring_mask & monitoring::EVENT_STOP_ITERATION != 0 { + let offset = (self.lasti() - 1) * 2; + monitoring::fire_stop_iteration(vm, self.code, offset, &value)?; + } Ok(None) } Instruction::InstrumentedEndSend => { let value = self.pop_value(); - self.pop_value(); // discard receiver + let receiver = self.pop_value(); + // PyGen_Check || PyCoro_CheckExact + let is_gen_or_coro = receiver + .downcast_ref::() + .is_some() + || receiver + .downcast_ref::() + .is_some(); + if is_gen_or_coro && self.monitoring_mask & monitoring::EVENT_STOP_ITERATION != 0 { + let offset = (self.lasti() - 1) * 2; + monitoring::fire_stop_iteration(vm, self.code, offset, &value)?; + } self.push_value(value); Ok(None) } Instruction::InstrumentedPopJumpIfTrue => { let target = bytecode::Label::from(u32::from(arg)); - self.pop_jump_if(vm, target, true) + let obj = self.pop_value(); + let value = obj.try_to_bool(vm)?; + if value { + self.jump(target); + // Branch taken → fire BRANCH_RIGHT + if self.monitoring_mask & monitoring::EVENT_BRANCH_RIGHT != 0 { + let src_offset = (self.lasti() - 1) * 2; + monitoring::fire_branch_right(vm, self.code, src_offset, target.0 * 2)?; + } + } + // Branch not taken → InstrumentedNotTaken fires BRANCH_LEFT + Ok(None) } Instruction::InstrumentedPopJumpIfFalse => { let target = bytecode::Label::from(u32::from(arg)); - self.pop_jump_if(vm, target, false) + let obj = self.pop_value(); + let value = obj.try_to_bool(vm)?; + if !value { + self.jump(target); + // Branch taken → fire BRANCH_RIGHT + if self.monitoring_mask & monitoring::EVENT_BRANCH_RIGHT != 0 { + let src_offset = (self.lasti() - 1) * 2; + monitoring::fire_branch_right(vm, self.code, src_offset, target.0 * 2)?; + } + } + Ok(None) } Instruction::InstrumentedPopJumpIfNone => { let value = self.pop_value(); - let branch_taken = vm.is_none(&value); let target = bytecode::Label::from(u32::from(arg)); - let src_offset = (self.lasti() - 1) * 2; - if branch_taken { + if vm.is_none(&value) { self.jump(target); + // Branch taken → fire BRANCH_RIGHT + if self.monitoring_mask & monitoring::EVENT_BRANCH_RIGHT != 0 { + let src_offset = (self.lasti() - 1) * 2; + monitoring::fire_branch_right(vm, self.code, src_offset, target.0 * 2)?; + } } - self.fire_branch_event(vm, src_offset, branch_taken, target)?; Ok(None) } Instruction::InstrumentedPopJumpIfNotNone => { let value = self.pop_value(); - let branch_taken = !vm.is_none(&value); let target = bytecode::Label::from(u32::from(arg)); - let src_offset = (self.lasti() - 1) * 2; - if branch_taken { + if !vm.is_none(&value) { self.jump(target); + // Branch taken → fire BRANCH_RIGHT + if self.monitoring_mask & monitoring::EVENT_BRANCH_RIGHT != 0 { + let src_offset = (self.lasti() - 1) * 2; + monitoring::fire_branch_right(vm, self.code, src_offset, target.0 * 2)?; + } + } + Ok(None) + } + Instruction::InstrumentedNotTaken => { + if self.monitoring_mask & monitoring::EVENT_BRANCH_LEFT != 0 { + let offset = (self.lasti() - 1) * 2; + let dest_offset = self.lasti() * 2; + monitoring::fire_branch_left( + vm, + self.code, + offset.saturating_sub(2), + dest_offset, + )?; } - self.fire_branch_event(vm, src_offset, branch_taken, target)?; Ok(None) } - // InstrumentedNotTaken is already handled above (line 2032) Instruction::InstrumentedPopIter => { + // BRANCH_RIGHT is fired by InstrumentedForIter, not here. self.pop_value(); Ok(None) } Instruction::InstrumentedEndAsyncFor => { + if self.monitoring_mask & monitoring::EVENT_BRANCH_RIGHT != 0 { + let oparg_val = u32::from(arg); + // src = next_instr - oparg (END_SEND position) + let src_offset = (self.lasti() - oparg_val) * 2; + // dest = this_instr + 1 + let dest_offset = self.lasti() * 2; + monitoring::fire_branch_right(vm, self.code, src_offset, dest_offset)?; + } let exc = self.pop_value(); let _awaitable = self.pop_value(); let exc = exc @@ -2637,10 +2765,84 @@ impl ExecutingFrame<'_> { Err(exc) } } - // INSTRUMENTED_LINE and INSTRUMENTED_INSTRUCTION are handled in the - // execution loop (line/instruction monitoring), so they are NOPs here. - Instruction::InstrumentedLine => Ok(None), - Instruction::InstrumentedInstruction => Ok(None), + Instruction::InstrumentedLine => { + let idx = self.lasti() as usize - 1; + let offset = idx as u32 * 2; + + // Read the full side-table chain before firing any events, + // because a callback may de-instrument and clear the tables. + let (real_op_byte, also_instruction) = { + let data = self.code.monitoring_data.lock(); + let line_op = data.as_ref().map(|d| d.line_opcodes[idx]).unwrap_or(0); + if line_op == u8::from(Instruction::InstrumentedInstruction) { + // LINE wraps INSTRUCTION: resolve the INSTRUCTION side-table too + let inst_op = data + .as_ref() + .map(|d| d.per_instruction_opcodes[idx]) + .unwrap_or(0); + (inst_op, true) + } else { + (line_op, false) + } + }; + debug_assert!( + real_op_byte != 0, + "INSTRUMENTED_LINE at {idx} without stored opcode" + ); + + // Fire LINE event only if line changed + if let Some((loc, _)) = self.code.locations.get(idx) { + let line = loc.line.get() as u32; + if line != self.prev_line && line > 0 { + self.prev_line = line; + monitoring::fire_line(vm, self.code, offset, line)?; + } + } + + // If the LINE position also had INSTRUCTION, fire that event too + if also_instruction { + monitoring::fire_instruction(vm, self.code, offset)?; + } + + // Re-dispatch to the real original opcode + let original_op = Instruction::try_from(real_op_byte) + .expect("invalid opcode in side-table chain"); + if original_op.to_base().is_some() { + self.execute_instrumented(original_op, arg, vm) + } else { + let mut do_extend_arg = false; + self.execute_instruction(original_op, arg, &mut do_extend_arg, vm) + } + } + Instruction::InstrumentedInstruction => { + let idx = self.lasti() as usize - 1; + let offset = idx as u32 * 2; + + // Get original opcode from side-table + let original_op_byte = { + let data = self.code.monitoring_data.lock(); + data.as_ref() + .map(|d| d.per_instruction_opcodes[idx]) + .unwrap_or(0) + }; + debug_assert!( + original_op_byte != 0, + "INSTRUMENTED_INSTRUCTION at {idx} without stored opcode" + ); + + // Fire INSTRUCTION event + monitoring::fire_instruction(vm, self.code, offset)?; + + // Re-dispatch to original opcode + let original_op = Instruction::try_from(original_op_byte) + .expect("invalid opcode in instruction side-table"); + if original_op.to_base().is_some() { + self.execute_instrumented(original_op, arg, vm) + } else { + let mut do_extend_arg = false; + self.execute_instruction(original_op, arg, &mut do_extend_arg, vm) + } + } _ => { unreachable!("{instruction:?} instruction should not be executed") } @@ -2865,7 +3067,7 @@ impl ExecutingFrame<'_> { // Fire EXCEPTION_HANDLED before setting up handler. // If the callback raises, the handler is NOT set up and the // new exception propagates instead. - if self.monitoring_mask + if vm.state.monitoring_events.load() & crate::stdlib::sys::monitoring::EVENT_EXCEPTION_HANDLED != 0 { @@ -3082,8 +3284,6 @@ impl ExecutingFrame<'_> { #[inline] fn execute_call(&mut self, args: FuncArgs, vm: &VirtualMachine) -> FrameResult { - use crate::stdlib::sys::monitoring; - // Stack: [callable, self_or_null, ...] let self_or_null = self.pop_value_opt(); // Option let callable = self.pop_value(); @@ -3102,9 +3302,38 @@ impl ExecutingFrame<'_> { args }; + match callable.call(final_args, vm) { + Ok(value) => { + self.push_value(value); + Ok(None) + } + Err(exc) => Err(exc), + } + } + + /// Instrumented version of execute_call: fires CALL, C_RETURN, and C_RAISE events. + fn execute_call_instrumented(&mut self, args: FuncArgs, vm: &VirtualMachine) -> FrameResult { + use crate::stdlib::sys::monitoring; + + // Stack: [callable, self_or_null, ...] + let self_or_null = self.pop_value_opt(); // Option + let callable = self.pop_value(); + + // If self_or_null is Some (not NULL), prepend it to args + let final_args = if let Some(self_val) = self_or_null { + let mut all_args = vec![self_val]; + all_args.extend(args.args); + FuncArgs { + args: all_args, + kwargs: args.kwargs, + } + } else { + args + }; + let is_python_call = callable.downcast_ref::().is_some(); - // Compute arg0 once for CALL, C_RETURN, and C_RAISE events + // Fire CALL event let call_arg0 = if self.monitoring_mask & monitoring::EVENT_CALL != 0 { let arg0 = final_args .args @@ -3130,8 +3359,6 @@ impl ExecutingFrame<'_> { Ok(None) } Err(exc) => { - // call_instrumentation_vector_protected replaces the active - // exception when the callback itself raises. let exc = if let Some(arg0) = call_arg0 && !is_python_call { @@ -3274,43 +3501,6 @@ impl ExecutingFrame<'_> { self.update_lasti(|i| *i = target_pc); } - #[inline] - fn fire_branch_event( - &self, - vm: &VirtualMachine, - src_offset: u32, - branch_taken: bool, - target: bytecode::Label, - ) -> PyResult<()> { - if self.monitoring_mask - & (crate::stdlib::sys::monitoring::EVENT_BRANCH_LEFT - | crate::stdlib::sys::monitoring::EVENT_BRANCH_RIGHT) - != 0 - { - let dest_offset = if branch_taken { - target.0 * 2 - } else { - self.lasti() * 2 - }; - if branch_taken { - crate::stdlib::sys::monitoring::fire_branch_right( - vm, - self.code, - src_offset, - dest_offset, - )?; - } else { - crate::stdlib::sys::monitoring::fire_branch_left( - vm, - self.code, - src_offset, - dest_offset, - )?; - } - } - Ok(()) - } - #[inline] fn pop_jump_if( &mut self, @@ -3320,54 +3510,47 @@ impl ExecutingFrame<'_> { ) -> FrameResult { let obj = self.pop_value(); let value = obj.try_to_bool(vm)?; - let branch_taken = value == flag; - let src_offset = (self.lasti() - 1) * 2; - if branch_taken { + if value == flag { self.jump(target); } - self.fire_branch_event(vm, src_offset, branch_taken, target)?; Ok(None) } - /// The top of stack contains the iterator, lets push it forward - fn execute_for_iter(&mut self, vm: &VirtualMachine, target: bytecode::Label) -> FrameResult { + /// Advance the iterator on top of stack. + /// Returns `true` if iteration continued (item pushed), `false` if exhausted (jumped). + fn execute_for_iter( + &mut self, + vm: &VirtualMachine, + target: bytecode::Label, + ) -> Result { let top_of_stack = PyIter::new(self.top_value()); let next_obj = top_of_stack.next(vm); - let src_offset = (self.lasti() - 1) * 2; - // Check the next object: match next_obj { Ok(PyIterReturn::Return(value)) => { self.push_value(value); - self.fire_branch_event(vm, src_offset, false, target)?; - Ok(None) + Ok(true) } Ok(PyIterReturn::StopIteration(_)) => { - // Check if target instruction is END_FOR (CPython 3.14 pattern) - // If so, skip it and jump to target + 1 instruction (POP_ITER) + // Skip END_FOR if followed by POP_ITER (both base and instrumented). + // PopIter may be further wrapped by InstrumentedInstruction / InstrumentedLine. let target_idx = target.0 as usize; let jump_target = if let Some(unit) = self.code.instructions.get(target_idx) { - if matches!(unit.op, bytecode::Instruction::EndFor) - && matches!( - self.code.instructions.get(target_idx + 1).map(|u| &u.op), - Some(bytecode::Instruction::PopIter) - ) - { - // Skip END_FOR, jump to POP_ITER + if matches!( + unit.op, + bytecode::Instruction::EndFor | bytecode::Instruction::InstrumentedEndFor + ) { bytecode::Label(target.0 + 1) } else { - // Legacy pattern: jump directly to target (POP_TOP/POP_ITER) target } } else { target }; self.jump(jump_target); - self.fire_branch_event(vm, src_offset, true, jump_target)?; - Ok(None) + Ok(false) } Err(next_error) => { - // On error, pop iterator and propagate self.pop_value(); Err(next_error) } diff --git a/crates/vm/src/stdlib/sys/monitoring.rs b/crates/vm/src/stdlib/sys/monitoring.rs index f024e9dad..f64c29986 100644 --- a/crates/vm/src/stdlib/sys/monitoring.rs +++ b/crates/vm/src/stdlib/sys/monitoring.rs @@ -1,8 +1,9 @@ use crate::{ AsObject, PyObjectRef, PyPayload, PyRef, PyResult, VirtualMachine, - builtins::{PyCode, PyDictRef, PyNamespace, PyUtf8StrRef}, + builtins::{PyCode, PyDictRef, PyNamespace, PyUtf8StrRef, code::CoMonitoringData}, function::FuncArgs, }; +use core::sync::atomic::Ordering; use crossbeam_utils::atomic::AtomicCell; use std::collections::{HashMap, HashSet}; @@ -53,7 +54,7 @@ pub const EVENT_EXCEPTION_HANDLED: u32 = MonitoringEvents::EXCEPTION_HANDLED.bit pub const EVENT_PY_UNWIND: u32 = MonitoringEvents::PY_UNWIND.bits(); pub const EVENT_C_RETURN: u32 = MonitoringEvents::C_RETURN.bits(); const EVENT_C_RAISE: u32 = MonitoringEvents::C_RAISE.bits(); -const EVENT_STOP_ITERATION: u32 = MonitoringEvents::STOP_ITERATION.bits(); +pub const EVENT_STOP_ITERATION: u32 = MonitoringEvents::STOP_ITERATION.bits(); pub const EVENT_PY_THROW: u32 = MonitoringEvents::PY_THROW.bits(); const EVENT_BRANCH: u32 = MonitoringEvents::BRANCH.bits(); pub const EVENT_RERAISE: u32 = MonitoringEvents::RERAISE.bits(); @@ -216,9 +217,269 @@ fn normalize_event_set(event_set: i32, local: bool, vm: &VirtualMachine) -> PyRe Ok(event_set) } +/// Rewrite a code object's bytecode in-place with layered instrumentation. +/// +/// Three layers (outermost first): +/// 1. INSTRUMENTED_LINE — wraps line-start instructions (stores original in side-table) +/// 2. INSTRUMENTED_INSTRUCTION — wraps all traceable instructions (stores original in side-table) +/// 3. Regular INSTRUMENTED_* — direct 1:1 opcode swap (no side-table needed) +/// +/// De-instrumentation peels layers in reverse order. +pub fn instrument_code(code: &PyCode, events: u32) { + use rustpython_compiler_core::bytecode::{self, Instruction}; + + let len = code.code.instructions.len(); + let mut monitoring_data = code.monitoring_data.lock(); + + // === Phase 1-3: De-instrument all layers (outermost first) === + + // Phase 1: Remove INSTRUMENTED_LINE → restore from side-table + if let Some(data) = monitoring_data.as_mut() { + for i in 0..len { + if data.line_opcodes[i] != 0 { + let original = Instruction::try_from(data.line_opcodes[i]) + .expect("invalid opcode in line side-table"); + unsafe { + code.code.instructions.replace_op(i, original); + } + data.line_opcodes[i] = 0; + } + } + } + + // Phase 2: Remove INSTRUMENTED_INSTRUCTION → restore from side-table + if let Some(data) = monitoring_data.as_mut() { + for i in 0..len { + if data.per_instruction_opcodes[i] != 0 { + let original = Instruction::try_from(data.per_instruction_opcodes[i]) + .expect("invalid opcode in instruction side-table"); + unsafe { + code.code.instructions.replace_op(i, original); + } + data.per_instruction_opcodes[i] = 0; + } + } + } + + // Phase 3: Remove regular INSTRUMENTED_* → restore base opcodes + for i in 0..len { + let op = code.code.instructions[i].op; + if let Some(base) = op.to_base() { + unsafe { + code.code.instructions.replace_op(i, base); + } + } + } + + // All opcodes are now base opcodes. + + if events == 0 { + *monitoring_data = None; + return; + } + + // === Phase 4-6: Re-instrument (innermost first) === + + // Ensure monitoring data exists + if monitoring_data.is_none() { + *monitoring_data = Some(CoMonitoringData { + line_opcodes: vec![0u8; len], + per_instruction_opcodes: vec![0u8; len], + }); + } + let data = monitoring_data.as_mut().unwrap(); + // Resize if code length changed (shouldn't happen, but be safe) + data.line_opcodes.resize(len, 0); + data.per_instruction_opcodes.resize(len, 0); + + // Find _co_firsttraceable: index of first RESUME instruction + let first_traceable = code + .code + .instructions + .iter() + .position(|u| matches!(u.op, Instruction::Resume { .. })) + .unwrap_or(0); + + // Phase 4: Place regular INSTRUMENTED_* opcodes + for i in 0..len { + let op = code.code.instructions[i].op; + if let Some(instrumented) = op.to_instrumented() { + unsafe { + code.code.instructions.replace_op(i, instrumented); + } + } + } + + // Phase 5: Place INSTRUMENTED_INSTRUCTION (if EVENT_INSTRUCTION is active) + if events & EVENT_INSTRUCTION != 0 { + for i in first_traceable..len { + let op = code.code.instructions[i].op; + // Skip ExtendedArg + if matches!(op, Instruction::ExtendedArg) { + continue; + } + // Excluded: RESUME and END_FOR (and their instrumented variants) + let base = op.to_base().map_or(op, |b| b); + if matches!(base, Instruction::Resume { .. } | Instruction::EndFor) { + continue; + } + // Store current opcode (may already be INSTRUMENTED_*) and replace + data.per_instruction_opcodes[i] = u8::from(op); + unsafe { + code.code + .instructions + .replace_op(i, Instruction::InstrumentedInstruction); + } + } + } + + // Phase 6: Place INSTRUMENTED_LINE (if EVENT_LINE is active) + // Mirrors CPython's initialize_lines: first determine which positions + // are line starts, then mark branch/jump targets, then place opcodes. + if events & EVENT_LINE != 0 { + // is_line_start[i] = true if position i should have INSTRUMENTED_LINE + let mut is_line_start = vec![false; len]; + + // First pass: mark positions where the source line changes + let mut prev_line: Option = None; + for (i, unit) in code + .code + .instructions + .iter() + .enumerate() + .take(len) + .skip(first_traceable) + { + let op = unit.op; + let base = op.to_base().map_or(op, |b| b); + if matches!(base, Instruction::ExtendedArg) { + continue; + } + // Excluded opcodes + if matches!( + base, + Instruction::Resume { .. } + | Instruction::EndFor + | Instruction::EndSend + | Instruction::PopIter + | Instruction::EndAsyncFor + ) { + continue; + } + if let Some((loc, _)) = code.code.locations.get(i) { + let line = loc.line.get() as u32; + let is_new = prev_line != Some(line); + prev_line = Some(line); + if is_new && line > 0 { + is_line_start[i] = true; + } + } + } + + // Second pass: mark branch/jump targets as line starts. + // Every jump/branch target must be a line start, even if on the + // same source line as the preceding instruction. Critical for loops + // (JUMP_BACKWARD → FOR_ITER). + let mut arg_state = bytecode::OpArgState::default(); + for unit in code.code.instructions[first_traceable..len].iter().copied() { + let (op, arg) = arg_state.get(unit); + let base = op.to_base().map_or(op, |b| b); + + if matches!(base, Instruction::ExtendedArg) { + continue; + } + + let target: Option = match base { + Instruction::PopJumpIfFalse { .. } + | Instruction::PopJumpIfTrue { .. } + | Instruction::PopJumpIfNone { .. } + | Instruction::PopJumpIfNotNone { .. } + | Instruction::JumpForward { .. } + | Instruction::JumpBackward { .. } + | Instruction::JumpBackwardNoInterrupt { .. } => Some(u32::from(arg) as usize), + Instruction::ForIter { .. } | Instruction::Send { .. } => { + // Skip over END_FOR/END_SEND + Some(u32::from(arg) as usize + 1) + } + _ => None, + }; + + if let Some(target_idx) = target + && target_idx < len + && !is_line_start[target_idx] + { + let target_op = code.code.instructions[target_idx].op; + let target_base = target_op.to_base().map_or(target_op, |b| b); + // Skip POP_ITER targets + if matches!(target_base, Instruction::PopIter) { + continue; + } + if let Some((loc, _)) = code.code.locations.get(target_idx) + && loc.line.get() > 0 + { + is_line_start[target_idx] = true; + } + } + } + + // Third pass: mark exception handler targets as line starts. + for entry in bytecode::decode_exception_table(&code.code.exceptiontable) { + let target_idx = entry.target as usize; + if target_idx < len && !is_line_start[target_idx] { + let target_op = code.code.instructions[target_idx].op; + let target_base = target_op.to_base().map_or(target_op, |b| b); + if !matches!(target_base, Instruction::PopIter) + && let Some((loc, _)) = code.code.locations.get(target_idx) + && loc.line.get() > 0 + { + is_line_start[target_idx] = true; + } + } + } + + // Fourth pass: actually place INSTRUMENTED_LINE at all marked positions + for (i, marked) in is_line_start + .iter() + .copied() + .enumerate() + .take(len) + .skip(first_traceable) + { + if marked { + let op = code.code.instructions[i].op; + data.line_opcodes[i] = u8::from(op); + unsafe { + code.code + .instructions + .replace_op(i, Instruction::InstrumentedLine); + } + } + } + } +} + /// Update the global monitoring_events atomic mask from current state. fn update_events_mask(vm: &VirtualMachine, state: &MonitoringState) { - vm.state.monitoring_events.store(state.combined_events()); + let events = state.combined_events(); + vm.state.monitoring_events.store(events); + let new_ver = vm + .state + .instrumentation_version + .fetch_add(1, Ordering::Release) + + 1; + // Eagerly re-instrument all frames on the current thread's stack so that + // code objects already past their RESUME pick up the new event set. + for fp in vm.frames.borrow().iter() { + // SAFETY: frames in the Vec are alive while their FrameRef is on the call stack. + let frame = unsafe { fp.as_ref() }; + let code = &frame.code; + let code_ver = code.instrumentation_version.load(Ordering::Acquire); + if code_ver != new_ver { + instrument_code(code, events); + code.instrumentation_version + .store(new_ver, Ordering::Release); + } + } } fn use_tool_id(tool_id: i32, name: &str, vm: &VirtualMachine) -> PyResult<()> { @@ -472,8 +733,11 @@ fn fire( // Non-local events (RAISE, EXCEPTION_HANDLED, PY_UNWIND, etc.) // cannot be disabled per code object. if event_id >= LOCAL_EVENTS_COUNT { + // Remove the callback, matching CPython behavior. + let mut state = vm.state.monitoring.lock(); + state.callbacks.remove(&(tool, event_id)); return Err(vm.new_value_error(format!( - "cannot disable {} events", + "Cannot disable {} events. Callback removed.", EVENT_NAMES[event_id] ))); } diff --git a/crates/vm/src/vm/interpreter.rs b/crates/vm/src/vm/interpreter.rs index 9f6f4de3d..2a664e39f 100644 --- a/crates/vm/src/vm/interpreter.rs +++ b/crates/vm/src/vm/interpreter.rs @@ -46,7 +46,7 @@ where use crate::common::hash::HashSecret; use crate::common::lock::PyMutex; use crate::warn::WarningsState; - use core::sync::atomic::AtomicBool; + use core::sync::atomic::{AtomicBool, AtomicU64}; use crossbeam_utils::atomic::AtomicCell; let paths = getpath::init_path_config(&settings); @@ -126,6 +126,7 @@ where shutdown_handles: parking_lot::Mutex::new(Vec::new()), monitoring: PyMutex::default(), monitoring_events: AtomicCell::new(0), + instrumentation_version: AtomicU64::new(0), }); // Create VM with the global state diff --git a/crates/vm/src/vm/mod.rs b/crates/vm/src/vm/mod.rs index 3d1fa467d..ead411702 100644 --- a/crates/vm/src/vm/mod.rs +++ b/crates/vm/src/vm/mod.rs @@ -43,7 +43,7 @@ use alloc::{borrow::Cow, collections::BTreeMap}; use core::{ cell::{Cell, OnceCell, RefCell}, ptr::NonNull, - sync::atomic::{AtomicBool, Ordering}, + sync::atomic::{AtomicBool, AtomicU64, Ordering}, }; use crossbeam_utils::atomic::AtomicCell; #[cfg(unix)] @@ -159,6 +159,9 @@ pub struct PyGlobalState { pub monitoring: PyMutex, /// Fast-path mask: OR of all tools' events. 0 means no monitoring overhead. pub monitoring_events: stdlib::sys::monitoring::MonitoringEventsMask, + /// Incremented on every monitoring state change. Code objects compare their + /// local version against this to decide whether re-instrumentation is needed. + pub instrumentation_version: AtomicU64, } pub fn process_hash_secret_seed() -> u32 {