From 257d291b3a233c18ae4d4fd2aac62dda50090250 Mon Sep 17 00:00:00 2001 From: 1ndahous3 <9205230+1ndahous3@users.noreply.github.com> Date: Mon, 19 Oct 2026 14:27:00 +0300 Subject: [PATCH] Preserve extended opcode arguments under sys.monitoring (#9203) * Preserve extended opcode arguments under sys.monitoring Assisted-by: Codex:GPT-6 * Restart monitored execution after line callback jumps Assisted-by: Codex:GPT-6 * Keep monitoring argument state out of the dispatch loop Assisted-by: Codex:GPT-6 --- crates/vm/src/frame.rs | 46 ++++++++++++++++++++++++------ extra_tests/snippets/stdlib_sys.py | 39 +++++++++++++++++++++++++ 2 files changed, 76 insertions(+), 9 deletions(-) diff --git a/crates/vm/src/frame.rs b/crates/vm/src/frame.rs index 20759fb5f..e23d3a347 100644 --- a/crates/vm/src/frame.rs +++ b/crates/vm/src/frame.rs @@ -8237,20 +8237,30 @@ impl ExecutingFrame<'_> { } // All INSTRUMENTED_* opcodes delegate to a cold function to keep // the hot instruction loop free of monitoring overhead. - _ => self.execute_instrumented(instruction, arg, vm), + _ => { + let (result, do_extend_arg) = self.execute_instrumented(instruction, arg, vm); + *extend_arg = do_extend_arg; + result + } } } /// Handle all INSTRUMENTED_* opcodes. This function is cold — it only /// runs when sys.monitoring has rewritten the bytecode. #[cold] + #[inline(never)] fn execute_instrumented( &mut self, instruction: Instruction, arg: bytecode::OpArg, vm: &VirtualMachine, - ) -> FrameResult { - self.execute_instrumented_from(instruction, arg, false, vm) + ) -> (FrameResult, bool) { + // Keep the mutable flag inside this cold function so its address does + // not escape the dispatch loop and force a stack access for every opcode. + let mut do_extend_arg = false; + let result = + self.execute_instrumented_from(instruction, arg, false, &mut do_extend_arg, vm); + (result, do_extend_arg) } #[cold] @@ -8259,6 +8269,7 @@ impl ExecutingFrame<'_> { instruction: Instruction, arg: bytecode::OpArg, wrapper_dispatch: bool, + do_extend_arg: &mut bool, vm: &VirtualMachine, ) -> FrameResult { debug_assert!( @@ -8625,6 +8636,13 @@ impl ExecutingFrame<'_> { if line != self.prev_line.get() && line > 0 { self.prev_line.set(line); monitoring::fire_line(vm, self.code, offset, line)?; + if self.lasti() != idx as u32 + 1 { + self.apply_pending_lineno_jump(vm); + // Restart at the callback's target without executing the + // old opcode or carrying its EXTENDED_ARG prefix there. + *do_extend_arg = false; + return Ok(None); + } } } @@ -8654,10 +8672,15 @@ 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_from(original_op, arg, wrapper_dispatch, vm) + self.execute_instrumented_from( + original_op, + arg, + wrapper_dispatch, + do_extend_arg, + vm, + ) } else { - let mut do_extend_arg = false; - self.execute_instruction(original_op, arg, &mut do_extend_arg, vm) + self.execute_instruction(original_op, arg, do_extend_arg, vm) }; let orig_caches = original_op.to_base().unwrap_or(original_op).cache_entries(); if orig_caches > 0 && self.lasti() == lasti_before_dispatch { @@ -8694,10 +8717,15 @@ 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_from(original_op, arg, wrapper_dispatch, vm) + self.execute_instrumented_from( + original_op, + arg, + wrapper_dispatch, + do_extend_arg, + vm, + ) } else { - let mut do_extend_arg = false; - self.execute_instruction(original_op, arg, &mut do_extend_arg, vm) + self.execute_instruction(original_op, arg, do_extend_arg, vm) }; let orig_caches = original_op.to_base().unwrap_or(original_op).cache_entries(); if orig_caches > 0 && self.lasti() == lasti_before_dispatch { diff --git a/extra_tests/snippets/stdlib_sys.py b/extra_tests/snippets/stdlib_sys.py index 40daf9530..17253d7e6 100644 --- a/extra_tests/snippets/stdlib_sys.py +++ b/extra_tests/snippets/stdlib_sys.py @@ -155,6 +155,45 @@ finally: monitoring.free_tool_id(2) assert records == ["RAISE", "PY_UNWIND"], records +# A loop body large enough to require EXTENDED_ARG on its jump targets. +namespace = {} +exec( + "def large_loop(skip=False):\n if skip:\n return 'done'\n" + " total = 0\n for x in (1, 2):\n" + " total += x\n" * 100 + " return total\n", + namespace, +) +large_loop = namespace["large_loop"] + + +def jump_out_of_large_loop(code, line): + global loop_hits + if code is large_loop.__code__ and line == 5: + loop_hits += 1 + if loop_hits == 2: + sys._getframe(1).f_lineno = 3 + + +monitoring.use_tool_id(2, "extended-arg") +try: + for event in (monitoring.events.LINE, monitoring.events.INSTRUCTION): + monitoring.register_callback(2, event, lambda *args: None) + for events in ( + monitoring.events.LINE, + monitoring.events.LINE | monitoring.events.INSTRUCTION, + ): + monitoring.set_local_events(2, large_loop.__code__, events) + assert large_loop() == 300 + loop_hits = 0 + monitoring.register_callback(2, monitoring.events.LINE, jump_out_of_large_loop) + monitoring.set_local_events(2, large_loop.__code__, monitoring.events.LINE) + assert large_loop() == "done" + assert loop_hits == 2 +finally: + monitoring.set_local_events(2, large_loop.__code__, 0) + for event in (monitoring.events.LINE, monitoring.events.INSTRUCTION): + monitoring.register_callback(2, event, None) + monitoring.free_tool_id(2) + assert sys.exc_info() == (None, None, None) try: