From efa8c2692919e91eef871710e456ce8401dbc7e4 Mon Sep 17 00:00:00 2001 From: youknowdot Date: Sat, 17 Oct 2026 18:57:23 -0700 Subject: [PATCH] Align code introspection and monitoring with Python 3.15 (#9178) * Align code metadata and monitoring event semantics with Python 3.15 Preserve canonical qualified-name interning and subclass values, normalize valid debug-range tables, and correct event-specific monitoring disable and callback ownership boundaries. Assisted-by: OpenAI Codex:model version unavailable * Import canonical monitoring shutdown helper Copy the missing test_monitoring dependency from pinned CPython 3.15.0 through scripts/update_lib, retaining existing suite guards. Assisted-by: OpenAI Codex:model version unavailable * Match observed monitoring source layers and tool release ordering Model active source instrumentation across all subscribed tools instead of unconditional legacy edge pairing. Preserve tool names during unlocked callback finalization before releasing them. Assisted-by: OpenAI Codex:model version unavailable * Preserve code qualnames and monitoring event-site identity Retain decoded WTF8 qualnames, use exact-str marshal dispatch, and model wrapped branch dispatch with physical DISABLE sites. Retire the proven monitoring shutdown guard through update_lib. Assisted-by: OpenAI Codex:model version unavailable --- Lib/test/_test_monitoring_shutdown.py | 30 ++ Lib/test/test_monitoring.py | 1 - crates/compiler-core/src/marshal.rs | 7 +- crates/vm/src/builtins/code.rs | 115 +++++++- crates/vm/src/builtins/code/const_cache.rs | 1 + crates/vm/src/builtins/function.rs | 2 +- crates/vm/src/frame.rs | 71 ++++- crates/vm/src/stdlib/marshal.rs | 10 +- crates/vm/src/stdlib/sys/monitoring.rs | 318 +++++++++++++++------ 9 files changed, 449 insertions(+), 106 deletions(-) create mode 100644 Lib/test/_test_monitoring_shutdown.py diff --git a/Lib/test/_test_monitoring_shutdown.py b/Lib/test/_test_monitoring_shutdown.py new file mode 100644 index 000000000..3d0fbec02 --- /dev/null +++ b/Lib/test/_test_monitoring_shutdown.py @@ -0,0 +1,30 @@ +#!/usr/bin/env python3 + +# gh-115832: An object destructor running during the final GC of interpreter +# shutdown triggered an infinite loop in the instrumentation code. + +import sys + +class CallableCycle: + def __init__(self): + self._cycle = self + + def __del__(self): + pass + + def __call__(self, code, instruction_offset): + pass + +def tracefunc(frame, event, arg): + pass + +def main(): + tool_id = sys.monitoring.PROFILER_ID + event_id = sys.monitoring.events.PY_START + + sys.monitoring.use_tool_id(tool_id, "test profiler") + sys.monitoring.set_events(tool_id, event_id) + sys.monitoring.register_callback(tool_id, event_id, CallableCycle()) + +if __name__ == "__main__": + sys.exit(main()) diff --git a/Lib/test/test_monitoring.py b/Lib/test/test_monitoring.py index 006480abe..ad9b61a65 100644 --- a/Lib/test/test_monitoring.py +++ b/Lib/test/test_monitoring.py @@ -2576,7 +2576,6 @@ class TestTier2Optimizer(CheckEvents): class TestMonitoringAtShutdown(unittest.TestCase): - @unittest.expectedFailure # TODO: RUSTPYTHON; - requires subprocess support def test_monitoring_live_at_shutdown(self): # gh-115832: An object destructor running during the final GC of # interpreter shutdown triggered an infinite loop in the diff --git a/crates/compiler-core/src/marshal.rs b/crates/compiler-core/src/marshal.rs index 5587c6314..8f0c7253d 100644 --- a/crates/compiler-core/src/marshal.rs +++ b/crates/compiler-core/src/marshal.rs @@ -559,6 +559,7 @@ fn read_const_value( /// Object-valued code fields as read from the stream (`r_object` TYPE_CODE). pub struct CodeFields { pub consts: V, + pub qualname: V, pub names: V, pub localsplusnames: V, pub localspluskinds: V, @@ -977,7 +978,10 @@ fn deserialize_code_value_inner( }; let source_path_raw = read_string(rdr, refs)?; let obj_name_raw = read_string(rdr, refs)?; - let qualname_raw = read_string(rdr, refs)?; + let qualname_value = deserialize_value_depth(rdr, bag, child_depth, refs)?; + let qualname_raw = bag + .str_from_value(&qualname_value) + .ok_or(MarshalError::BadType)?; let first_line_raw = rdr.read_u32()? as i32; let first_line_number = if first_line_raw > 0 { @@ -1049,6 +1053,7 @@ fn deserialize_code_value_inner( code, CodeFields { consts: consts_value, + qualname: qualname_value, names: names_value, localsplusnames: localsplusnames_value, localspluskinds: kinds_value, diff --git a/crates/vm/src/builtins/code.rs b/crates/vm/src/builtins/code.rs index 89f760ad2..be6807efe 100644 --- a/crates/vm/src/builtins/code.rs +++ b/crates/vm/src/builtins/code.rs @@ -391,7 +391,9 @@ pub trait IntoCodeObject { where Self: Sized, { - PyCode::new(self.into_code_object(vm), &vm.ctx) + let mut code = PyCode::new(self.into_code_object(vm), &vm.ctx); + code.apply_debug_ranges(vm); + code } } @@ -444,7 +446,6 @@ pub struct PyCode { #[pymember(name = "co_posonlyargcount", path = "posonlyarg_count")] #[pymember(name = "co_kwonlyargcount", path = "kwonlyarg_count")] #[pymember(name = "co_stacksize", path = "max_stackdepth")] - #[pymember(name = "co_qualname", path = "qualname")] #[pymember(name = "co_flags", path = "flags")] pub code: CodeObject, /// Slot-indexed names, equivalent to CPython's `co_localsplusnames`. @@ -452,6 +453,8 @@ pub struct PyCode { /// merged cell variables. localsplus_names: Box<[&'static PyStrInterned]>, pub(crate) co_name: PyStrRef, + #[pymember(name = "co_qualname")] + qualified_name: PyStrRef, pub(crate) co_names: PyTupleRef, pub(crate) co_localsplusnames: PyTupleRef, pub(crate) co_localspluskinds: PyBytesRef, @@ -625,6 +628,7 @@ impl PyCode { Self { first_traceable, co_name: code.obj_name.to_owned(), + qualified_name: code.qualname.to_owned(), co_names: names_tuple(&code.names), co_localsplusnames: names_tuple(&localsplus_names), co_localspluskinds: ctx.new_bytes(code.localspluskinds.to_vec()), @@ -654,6 +658,10 @@ impl PyCode { obj.downcast().map_err(|_| MarshalError::BadType) } code.constants = bytecode::Constants::new(LiteralTuple(downcast(fields.consts)?)); + // The compiler's string conversion is lossy for surrogates; retain the + // decoded value and apply code construction's exact-string interning. + let qualname = downcast::(fields.qualname)?; + code.qualname = ctx.intern_str(qualname.as_wtf8()); let mut code = Self::new(code, ctx); code.set_names(downcast(fields.names)?); code.set_localsplusnames(downcast(fields.localsplusnames)?); @@ -678,6 +686,19 @@ impl PyCode { } } + pub(crate) fn apply_debug_ranges(&mut self, vm: &VirtualMachine) { + if !vm.state.config.settings.code_debug_ranges + && let Some(linetable) = remove_column_info(&self.code.linetable) + { + self.code.linetable = linetable; + self.co_linetable = vm.ctx.new_bytes(self.code.linetable.to_vec()); + } + } + + pub(crate) fn co_qualname(&self) -> PyStrRef { + self.qualified_name.clone() + } + #[inline(always)] pub(crate) fn localsplus_name(&self, index: usize) -> &'static PyStrInterned { // Bytecode operands and frame-proxy indices are validated against @@ -734,11 +755,9 @@ impl PyCode { } pub fn new_ref_with_bag(vm: &VirtualMachine, code: CodeObject) -> PyRef { - PyRef::new_ref( - Self::new(code, &vm.ctx), - vm.ctx.types.code_type.to_owned(), - None, - ) + let mut code = Self::new(code, &vm.ctx); + code.apply_debug_ranges(vm); + PyRef::new_ref(code, vm.ctx.types.code_type.to_owned(), None) } pub fn new_ref_from_bytecode(vm: &VirtualMachine, code: bytecode::CodeObject) -> PyRef { @@ -1170,9 +1189,13 @@ impl Constructor for PyCode { let mut code = Self::new(code, &vm.ctx); code.co_name = args.name; + if !args.qualname.class().is(vm.ctx.types.str_type) { + code.qualified_name = args.qualname; + } code.set_names(args.names); code.co_linetable = args.linetable; code.co_exceptiontable = args.exceptiontable; + code.apply_debug_ranges(vm); Ok(code) } } @@ -1633,7 +1656,7 @@ impl Py { let qualname = match co_qualname { OptionalArg::Present(qualname) => qualname, - OptionalArg::Missing => self.code.qualname.to_owned(), + OptionalArg::Missing => self.qualified_name.clone(), }; // Room for one value is always reserved, even where nothing is pushed. @@ -1741,9 +1764,13 @@ impl Py { let mut code = PyCode::new(new_code, &vm.ctx); code.co_name = obj_name; + if !qualname.class().is(vm.ctx.types.str_type) { + code.qualified_name = qualname; + } code.set_names(names); code.co_linetable = linetable; code.co_exceptiontable = exceptiontable; + code.apply_debug_ranges(vm); Ok(code) } @@ -1800,6 +1827,78 @@ impl ToPyObject for bytecode::CodeObject { } } +// Preserve entry lengths and line deltas while dropping columns, as code +// construction does under -X no_debug_ranges in CPython. Decode and re-encode +// signed deltas to normalize redundant encodings, including negative zero. +// Leave unconvertible tables unchanged rather than inventing line deltas or +// introducing new constructor errors. This is not a general table validator. +fn remove_column_info(linetable: &[u8]) -> Option> { + let mut result = Vec::with_capacity(linetable.len()); + let mut reader = LineTableReader::new(linetable); + while let Some(header) = reader.read_byte() { + if header & 0x80 == 0 { + return None; + } + let kind = (header >> 3) & 15; + if kind == PyCodeLocationInfoKind::None as u8 { + result.push(header); + } else { + result.push(0x80 | ((PyCodeLocationInfoKind::NoColumns as u8) << 3) | (header & 7)); + let delta = match PyCodeLocationInfoKind::from_code(kind) { + Some(PyCodeLocationInfoKind::Long | PyCodeLocationInfoKind::NoColumns) => { + read_location_delta(&mut reader)? + } + Some(PyCodeLocationInfoKind::OneLine1) => 1, + Some(PyCodeLocationInfoKind::OneLine2) => 2, + _ => 0, + }; + let mut encoded = (delta.unsigned_abs() << 1) | u32::from(delta < 0); + while encoded >= 64 { + result.push(0x40 | (encoded & 0x3f) as u8); + encoded >>= 6; + } + result.push(encoded as u8); + } + while reader.peek_byte().is_some_and(|byte| byte & 0x80 == 0) { + reader.read_byte(); + } + } + Some(result.into_boxed_slice()) +} + +// Location deltas use an unsigned 32-bit sign/magnitude representation. Allow +// redundant zero groups, but reject overflow and truncation without unchecked +// shifts or consuming the next header. This does not validate absolute lines. +fn read_location_delta(reader: &mut LineTableReader<'_>) -> Option { + let mut value = 0u32; + let mut shift = 0; + loop { + let byte = reader.peek_byte()?; + if byte & 0x80 != 0 { + return None; + } + reader.read_byte(); + let digit = u32::from(byte & 0x3f); + if shift < u32::BITS { + if digit > (u32::MAX >> shift) { + return None; + } + value |= digit << shift; + } else if digit != 0 { + return None; + } + if byte & 0x40 == 0 { + let magnitude = (value >> 1) as i32; + return Some(if value & 1 != 0 { + -magnitude + } else { + magnitude + }); + } + shift = (shift + 6).min(u32::BITS); + } +} + // Helper struct for reading linetable struct LineTableReader<'a> { data: &'a [u8], diff --git a/crates/vm/src/builtins/code/const_cache.rs b/crates/vm/src/builtins/code/const_cache.rs index 026dfda13..6515219e7 100644 --- a/crates/vm/src/builtins/code/const_cache.rs +++ b/crates/vm/src/builtins/code/const_cache.rs @@ -58,6 +58,7 @@ impl CompileBag<'_> { /// merges them, in the same order. fn new_code(&self, code: CodeObject) -> PyCode { let mut code = PyCode::new(code, &self.vm.ctx); + code.apply_debug_ranges(self.vm); code.co_exceptiontable = self.merge_one(code.co_exceptiontable); code.co_linetable = self.merge_one(code.co_linetable); code.co_names = self.merge_one(code.co_names); diff --git a/crates/vm/src/builtins/function.rs b/crates/vm/src/builtins/function.rs index b7629cf20..5f59675b9 100644 --- a/crates/vm/src/builtins/function.rs +++ b/crates/vm/src/builtins/function.rs @@ -207,7 +207,7 @@ impl PyFunction { vm.ctx.none() }; - let qualname = vm.ctx.new_str(code.qualname.as_str()); + let qualname = code.co_qualname(); let func = Self { code: PyAtomicRef::from(code), globals, diff --git a/crates/vm/src/frame.rs b/crates/vm/src/frame.rs index 9520ee3ef..95a382ce7 100644 --- a/crates/vm/src/frame.rs +++ b/crates/vm/src/frame.rs @@ -3820,7 +3820,9 @@ impl ExecutingFrame<'_> { && self.trace_lines_is_set() && !matches!( self.code.instructions.read_op(idx), - Instruction::Resume { .. } | Instruction::InstrumentedResume + Instruction::Resume { .. } + | Instruction::ResumeCheck + | Instruction::InstrumentedResume ) { let line = self.code.addr2line(idx as i32 * 2); @@ -3906,7 +3908,10 @@ impl ExecutingFrame<'_> { && self.trace_opcodes_is_set() && !matches!( op.into(), - Opcode::Resume | Opcode::InstrumentedResume | Opcode::ExtendedArg + Opcode::Resume + | Opcode::ResumeCheck + | Opcode::InstrumentedResume + | Opcode::ExtendedArg ) { vm.trace_event(crate::protocol::TraceEvent::Opcode, None)?; @@ -8230,6 +8235,17 @@ impl ExecutingFrame<'_> { instruction: Instruction, arg: bytecode::OpArg, vm: &VirtualMachine, + ) -> FrameResult { + self.execute_instrumented_from(instruction, arg, false, vm) + } + + #[cold] + fn execute_instrumented_from( + &mut self, + instruction: Instruction, + arg: bytecode::OpArg, + wrapper_dispatch: bool, + vm: &VirtualMachine, ) -> FrameResult { debug_assert!( instruction.is_instrumented(), @@ -8405,7 +8421,13 @@ impl ExecutingFrame<'_> { if continued { if self.monitoring_mask & MonitoringEvent::BranchLeft.mask() != 0 { let dest_offset = (self.lasti() + 1) * 2; // after caches - monitoring::fire_branch_left(vm, self.code, src_offset, dest_offset)?; + monitoring::fire_branch_left( + vm, + self.code, + src_offset, + src_offset, + dest_offset, + )?; } } else if self.monitoring_mask & MonitoringEvent::BranchRight.mask() != 0 { // INSTRUMENTED_POP_ITER: dest is the instruction after @@ -8512,9 +8534,14 @@ impl ExecutingFrame<'_> { { branch_idx -= 1; } - let src_offset = (branch_idx as u32) * 2; + let site = not_taken_idx as u32 * 2; + let src_offset = if wrapper_dispatch { + site + } else { + (branch_idx as u32) * 2 + }; let dest_offset = self.lasti() * 2; - monitoring::fire_branch_left(vm, self.code, src_offset, dest_offset)?; + monitoring::fire_branch_left(vm, self.code, site, src_offset, dest_offset)?; } Ok(None) } @@ -8567,6 +8594,17 @@ impl ExecutingFrame<'_> { "INSTRUMENTED_LINE at {idx} without stored opcode" ); + let line_active = + monitoring::is_event_active_at(vm, self.code, offset, MonitoringEvent::Line); + let instruction_active = also_instruction + && monitoring::is_event_active_at( + vm, + self.code, + offset, + MonitoringEvent::Instruction, + ); + let wrapper_dispatch = line_active || instruction_active; + // Fire LINE event only if line changed if let Some((loc, _)) = self.code.locations.get(idx) { let line = loc.line.get() as u32; @@ -8576,8 +8614,16 @@ impl ExecutingFrame<'_> { } } - // If the LINE position also had INSTRUCTION, fire that event too - if also_instruction { + // Removing the last LINE layer changes its saved opcode in CPython. + // That redispatch skips nested INSTRUCTION once, not on later visits. + let line_removed = line_active + && !monitoring::is_event_active_at( + vm, + self.code, + offset, + MonitoringEvent::Line, + ); + if also_instruction && !line_removed { monitoring::fire_instruction(vm, self.code, offset)?; } @@ -8594,7 +8640,7 @@ impl ExecutingFrame<'_> { .expect("invalid opcode in side-table chain"); let lasti_before_dispatch = self.lasti(); let result = if original_op.to_base().is_some() { - self.execute_instrumented(original_op, arg, vm) + self.execute_instrumented_from(original_op, arg, wrapper_dispatch, vm) } else { let mut do_extend_arg = false; self.execute_instruction(original_op, arg, &mut do_extend_arg, vm) @@ -8619,6 +8665,13 @@ impl ExecutingFrame<'_> { "INSTRUMENTED_INSTRUCTION at {idx} without stored opcode" ); + let wrapper_dispatch = monitoring::is_event_active_at( + vm, + self.code, + offset, + MonitoringEvent::Instruction, + ); + // Fire INSTRUCTION event monitoring::fire_instruction(vm, self.code, offset)?; @@ -8627,7 +8680,7 @@ impl ExecutingFrame<'_> { .expect("invalid opcode in instruction side-table"); let lasti_before_dispatch = self.lasti(); let result = if original_op.to_base().is_some() { - self.execute_instrumented(original_op, arg, vm) + self.execute_instrumented_from(original_op, arg, wrapper_dispatch, vm) } else { let mut do_extend_arg = false; self.execute_instruction(original_op, arg, &mut do_extend_arg, vm) diff --git a/crates/vm/src/stdlib/marshal.rs b/crates/vm/src/stdlib/marshal.rs index 7e6d17000..62bed8d7a 100644 --- a/crates/vm/src/stdlib/marshal.rs +++ b/crates/vm/src/stdlib/marshal.rs @@ -344,7 +344,9 @@ mod decl { write_float_str(buf, cv.re); write_float_str(buf, cv.im); } - } else if let Some(s) = obj.downcast_ref::() { + } else if let Some(s) = obj.downcast_ref::() + && obj.class().is(PyStr::static_type()) + { let bytes = s.as_wtf8().as_bytes(); // Only the formats that carry back-references tell an interned // string apart, and only those from 4 on have the ascii-only forms. @@ -428,6 +430,7 @@ mod decl { buf.write_u32(value); } let co_code = co.get_code(vm); + let co_qualname = co.co_qualname(); let mut write_field = |buf: &mut Vec, field: &PyObject| { write_object_depth(buf, field, refs, version, allow_code, vm, depth - 1) }; @@ -439,7 +442,7 @@ mod decl { co.co_localspluskinds.as_object(), co.source_path().as_object(), co.co_name.as_object(), - code.qualname.as_object(), + co_qualname.as_object(), ] { write_field(buf, field)?; } @@ -652,7 +655,8 @@ mod decl { .new_value_error("unmarshalling code objects is disallowed"), )); } - let code = PyCode::from_marshal(code, fields, &self.vm.ctx)?; + let mut code = PyCode::from_marshal(code, fields, &self.vm.ctx)?; + code.apply_debug_ranges(self.vm); Ok(code.into_ref(&self.vm.ctx).into()) } fn make_stop_iter(&self) -> Result { diff --git a/crates/vm/src/stdlib/sys/monitoring.rs b/crates/vm/src/stdlib/sys/monitoring.rs index b6d265f94..cf2b13c78 100644 --- a/crates/vm/src/stdlib/sys/monitoring.rs +++ b/crates/vm/src/stdlib/sys/monitoring.rs @@ -98,8 +98,11 @@ pub struct MonitoringState { pub global_events: [u32; TOOL_LIMIT], pub local_events: HashMap<(usize, usize), u32>, pub callbacks: HashMap<(usize, usize), PyObjectRef>, - /// Per-instruction disabled tools: (code_id, offset, tool) - pub disabled: HashSet<(usize, usize, usize)>, + // Legacy BRANCH callbacks also check whether the source instruction was + // disabled. Direct BRANCH_LEFT/RIGHT registrations bypass that check. + legacy_branch_callbacks: HashSet<(usize, usize)>, + /// Per-instruction disabled tools: (code_id, offset, event, tool) + pub disabled: HashSet<(usize, usize, usize, usize)>, /// Cached MISSING sentinel singleton pub missing: Option, /// Cached DISABLE sentinel singleton @@ -113,6 +116,7 @@ impl Default for MonitoringState { global_events: [0; TOOL_LIMIT], local_events: HashMap::new(), callbacks: HashMap::new(), + legacy_branch_callbacks: HashSet::new(), disabled: HashSet::new(), missing: None, disable: None, @@ -635,33 +639,44 @@ fn use_tool_id(tool_id: i32, name: &str, vm: &VirtualMachine) -> PyResult<()> { } fn clear_tool_id(tool_id: i32, vm: &VirtualMachine) -> PyResult<()> { - let tool = check_valid_tool(tool_id, vm)?; - let mut state = vm.state.monitoring.lock(); - if state.tool_names[tool].is_some() { - state.global_events[tool] = 0; - state - .local_events - .retain(|(local_tool, _), _| *local_tool != tool); - state.callbacks.retain(|(cb_tool, _), _| *cb_tool != tool); - state.disabled.retain(|&(_, _, t)| t != tool); - } - update_events_mask(vm, &state); - Ok(()) + release_tool_callbacks(tool_id, false, vm) } fn free_tool_id(tool_id: i32, vm: &VirtualMachine) -> PyResult<()> { + release_tool_callbacks(tool_id, true, vm) +} + +fn release_tool_callbacks(tool_id: i32, free: bool, vm: &VirtualMachine) -> PyResult<()> { let tool = check_valid_tool(tool_id, vm)?; - let mut state = vm.state.monitoring.lock(); - if state.tool_names[tool].is_some() { - state.global_events[tool] = 0; - state - .local_events - .retain(|(local_tool, _), _| *local_tool != tool); - state.callbacks.retain(|(cb_tool, _), _| *cb_tool != tool); - state.disabled.retain(|&(_, _, t)| t != tool); - state.tool_names[tool] = None; + let callbacks = { + let mut state = vm.state.monitoring.lock(); + let callbacks = if state.tool_names[tool].is_some() { + state.global_events[tool] = 0; + state + .local_events + .retain(|(local_tool, _), _| *local_tool != tool); + let callbacks: Vec<_> = state + .callbacks + .extract_if(|(cb_tool, _), _| *cb_tool == tool) + .map(|(_, callback)| callback) + .collect(); + state.disabled.retain(|&(_, _, _, t)| t != tool); + state + .legacy_branch_callbacks + .retain(|(cb_tool, _)| *cb_tool != tool); + callbacks + } else { + Vec::new() + }; + update_events_mask(vm, &state); + callbacks + }; + // Callback finalizers may reenter monitoring; release owners after unlocking. + drop(callbacks); + // Keep the name visible during callback finalization, as CPython does. + if free { + vm.state.monitoring.lock().tool_names[tool] = None; } - update_events_mask(vm, &state); Ok(()) } @@ -682,28 +697,35 @@ fn register_callback( vm.audit("sys.monitoring.register_callback", || (func.clone(),))?; - let mut state = vm.state.monitoring.lock(); - let prev = state - .callbacks - .remove(&(tool, event_id)) - .unwrap_or_else(|| vm.ctx.none()); - let branch_id = MonitoringEvent::Branch as usize; - let branch_left_id = MonitoringEvent::BranchLeft as usize; - let branch_right_id = MonitoringEvent::BranchRight as usize; - if !vm.is_none(&func) { - state.callbacks.insert((tool, event_id), func.clone()); - // BRANCH is a composite event: also register for BRANCH_LEFT/RIGHT - if event_id == branch_id { - state.callbacks.insert((tool, branch_left_id), func.clone()); - state.callbacks.insert((tool, branch_right_id), func); + let (prev, displaced) = { + let mut state = vm.state.monitoring.lock(); + let prev = state + .callbacks + .remove(&(tool, event_id)) + .unwrap_or_else(|| vm.ctx.none()); + let mut displaced = Vec::new(); + state.legacy_branch_callbacks.remove(&(tool, event_id)); + if !vm.is_none(&func) { + state.callbacks.insert((tool, event_id), func.clone()); } - } else { - // Also clear BRANCH_LEFT/RIGHT when clearing BRANCH - if event_id == branch_id { - state.callbacks.remove(&(tool, branch_left_id)); - state.callbacks.remove(&(tool, branch_right_id)); + // BRANCH replaces both edge registrations, including independently + // registered callbacks. Keep their owners alive until after unlocking. + if event_id == MonitoringEvent::Branch as usize { + for edge in [MonitoringEvent::BranchLeft, MonitoringEvent::BranchRight] { + let key = (tool, edge as usize); + let old = if vm.is_none(&func) { + state.legacy_branch_callbacks.remove(&key); + state.callbacks.remove(&key) + } else { + state.legacy_branch_callbacks.insert(key); + state.callbacks.insert(key, func.clone()) + }; + displaced.extend(old); + } } - } + (prev, displaced) + }; + drop(displaced); Ok(prev) } @@ -819,6 +841,19 @@ fn fire( code: &Py, offset: u32, cb_extra: &[PyObjectRef], +) -> PyResult<()> { + fire_at(vm, event, code, offset, offset, cb_extra) +} + +// The physical instruction owns DISABLE state; branch callback source offsets +// can differ from that instruction when NOT_TAKEN follows a branch or wrapper. +fn fire_at( + vm: &VirtualMachine, + event: MonitoringEvent, + code: &Py, + offset: u32, + source_offset: u32, + cb_extra: &[PyObjectRef], ) -> PyResult<()> { // Prevent recursive event firing if vm.tracing_is_suppressed() || FIRING.with(|f| f.get()) { @@ -828,6 +863,11 @@ fn fire( let event_id = event as usize; let event_mask = event.mask(); let code_id = code.get_id(); + let disabled_event = if matches!(event, MonitoringEvent::CReturn | MonitoringEvent::CRaise) { + MonitoringEvent::Call as usize + } else { + event_id + }; // C_RETURN and C_RAISE are implicitly enabled when CALL is set. let check_bit = if matches!(event, MonitoringEvent::CReturn | MonitoringEvent::CRaise) { @@ -836,9 +876,18 @@ fn fire( event_mask }; - // Collect callbacks and snapshot the DISABLE sentinel under a single lock. - let (callbacks, disable_sentinel): (Vec<(usize, PyObjectRef)>, Option) = { - let state = vm.state.monitoring.lock(); + let branch_source = if matches!( + event, + MonitoringEvent::BranchLeft | MonitoringEvent::BranchRight + ) { + Some(branch_source_events(code, source_offset)) + } else { + None + }; + + // Collect callbacks without holding the monitoring lock across Python code. + let callbacks = { + let mut state = vm.state.monitoring.lock(); let mut cbs = Vec::new(); for tool in 0..TOOL_LIMIT { let global = state.global_events[tool]; @@ -850,14 +899,32 @@ fn fire( if ((global | local) & check_bit) == 0 { continue; } - if state.disabled.contains(&(code_id, offset as usize, tool)) { + if state + .disabled + .contains(&(code_id, offset as usize, disabled_event, tool)) + { continue; } - if let Some(cb) = state.callbacks.get(&(tool, event_id)) { - cbs.push((tool, cb.clone())); + if state.callbacks.contains_key(&(tool, event_id)) { + let legacy_branch = state.legacy_branch_callbacks.contains(&(tool, event_id)); + if legacy_branch + && branch_source.as_ref().is_some_and(|events| { + !events.iter().any(|source_event| { + event_active_at(&state, code_id, source_offset, *source_event) + }) + }) + { + // The legacy wrapper returns DISABLE without calling the + // user when the source instruction is no longer instrumented. + state + .disabled + .insert((code_id, offset as usize, event_id, tool)); + continue; + } + cbs.push((tool, state.callbacks[&(tool, event_id)].clone())); } } - (cbs, state.disable.clone()) + cbs }; if callbacks.is_empty() { @@ -869,43 +936,126 @@ fn fire( args_vec.extend_from_slice(cb_extra); let args = FuncArgs::from(args_vec); + for (tool, cb) in callbacks { + if call_instrument(vm, event, &cb, args.clone())? { + if is_instrumented_event(event_id) { + let mut state = vm.state.monitoring.lock(); + state + .disabled + .insert((code_id, offset as usize, event_id, tool)); + } else if event_id < UNGROUPED_EVENTS_COUNT { + // Non-instrumented events are disabled for this whole code + // object by removing only the local subscription. Global + // subscriptions continue to fire, including for this code. + let mut state = vm.state.monitoring.lock(); + if let Some(local) = state.local_events.get_mut(&(tool, code_id)) { + *local &= !event_mask; + if *local == 0 { + state.local_events.remove(&(tool, code_id)); + } + update_events_mask(vm, &state); + } + } else { + return Err(cannot_disable(vm, event, tool)); + } + } + } + Ok(()) +} + +fn event_active_at( + state: &MonitoringState, + code_id: usize, + offset: u32, + event: MonitoringEvent, +) -> bool { + (0..TOOL_LIMIT).any(|tool| { + let active = state.global_events[tool] + | state + .local_events + .get(&(tool, code_id)) + .copied() + .unwrap_or(0); + state.tool_names[tool].is_some() + && active & event.mask() != 0 + && !state + .disabled + .contains(&(code_id, offset as usize, event as usize, tool)) + }) +} + +pub(crate) fn is_event_active_at( + vm: &VirtualMachine, + code: &Py, + offset: u32, + event: MonitoringEvent, +) -> bool { + event_active_at(&vm.state.monitoring.lock(), code.get_id(), offset, event) +} + +// Model the source's instrumentation layers without changing its bytecode. +// A layer stays instrumented until its last subscribed tool disables it. +fn branch_source_events(code: &Py, offset: u32) -> Vec { + use rustpython_compiler_core::bytecode::Instruction; + let index = offset as usize / 2; + let data = code.monitoring_data.lock(); + let mut events = Vec::with_capacity(3); + let mut op = code.code.instructions.read_op(index); + if let Some(data) = data.as_ref() { + if matches!(op, Instruction::InstrumentedLine) { + events.push(MonitoringEvent::Line); + op = Instruction::try_from(data.line_opcodes[index]) + .expect("INSTRUMENTED_LINE without original opcode"); + } + if matches!(op, Instruction::InstrumentedInstruction) { + events.push(MonitoringEvent::Instruction); + op = Instruction::try_from(data.per_instruction_opcodes[index]) + .expect("INSTRUMENTED_INSTRUCTION without original opcode"); + } + } + if let Some(event) = event_for_opcode(op, code.code.instructions.read_arg(index).as_u8()) { + events.push(event); + } + events +} + +// Restore tracing state before processing a callback result or exception. +fn call_instrument( + vm: &VirtualMachine, + event: MonitoringEvent, + callback: &PyObject, + args: FuncArgs, +) -> PyResult { + if vm.tracing_is_suppressed() || FIRING.with(|f| f.get()) { + return Ok(false); + } + let disable_sentinel = vm.state.monitoring.lock().disable.clone(); FIRING.with(|f| f.set(true)); vm.enter_tracing(); let old_what = vm.what_event.replace(Some(event)); - let result = (|| { - for (tool, cb) in callbacks { - let result = cb.call(args.clone(), vm)?; - if disable_sentinel.as_ref().is_some_and(|d| result.is(d)) { - if is_instrumented_event(event_id) { - let mut state = vm.state.monitoring.lock(); - state.disabled.insert((code_id, offset as usize, tool)); - } else if event_id < UNGROUPED_EVENTS_COUNT { - // Other (non-instrumented) event: disable for this code object. - let mut state = vm.state.monitoring.lock(); - if let Some(events) = state.local_events.get_mut(&(tool, code_id)) { - *events &= !event_mask; - if *events == 0 { - state.local_events.remove(&(tool, code_id)); - } - } - update_events_mask(vm, &state); - } else { - // Remove the callback. - let mut state = vm.state.monitoring.lock(); - state.callbacks.remove(&(tool, event_id)); - return Err(vm.new_value_error(format!( - "Cannot disable {} events. Callback removed.", - EVENT_NAMES[event_id] - ))); - } - } - } - Ok(()) - })(); + let result = callback.call(args, vm); vm.what_event.set(old_what); vm.leave_tracing(); FIRING.with(|f| f.set(false)); - result + let result = result?; + let disable = disable_sentinel.as_ref().is_some_and(|d| result.is(d)); + Ok(disable) +} + +fn cannot_disable( + vm: &VirtualMachine, + event: MonitoringEvent, + tool: usize, +) -> crate::builtins::PyBaseExceptionRef { + let callback = { + let mut state = vm.state.monitoring.lock(); + state.callbacks.remove(&(tool, event as usize)) + }; + drop(callback); + vm.new_value_error(format!( + "Cannot disable {} events. Callback removed.", + EVENT_NAMES[event as usize] + )) } // Public dispatch functions (called from frame.rs) @@ -1173,13 +1323,15 @@ pub(crate) fn fire_jump( pub(crate) fn fire_branch_left( vm: &VirtualMachine, code: &Py, + site: u32, offset: u32, destination: u32, ) -> PyResult<()> { - fire( + fire_at( vm, MonitoringEvent::BranchLeft, code, + site, offset, &[ vm.ctx.new_int(offset).into(),