Skip to content

ZJIT: Fix TracePoint return value for methods with rescue #16445

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
k0kubun merged 2 commits into from
Mar 20, 2026
Merged

ZJIT: Fix TracePoint return value for methods with rescue #16445

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
ebac36d
ZJIT: Fix TracePoint return value for methods with rescue
k0kubun Mar 18, 2026
d22c60b
ZJIT: Use counted side exits for send fallbacks while tracing
k0kubun Mar 19, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
19 changes: 19 additions & 0 deletions zjit.c
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -264,6 +264,25 @@ rb_zjit_method_tracing_currently_enabled(void)
return tracing_events & (RUBY_EVENT_C_CALL | RUBY_EVENT_C_RETURN);
}

// Check if any ISEQ trace events are currently enabled.
// Used to prevent ZJIT from compiling while tracing is active, since ZJIT's
// send fallback (rb_vm_opt_send_without_block) uses VM_EXEC which sets
// VM_FRAME_FLAG_FINISH on the callee frame, changing exception handling
// semantics for throw TAG_RETURN (e.g. return from rescue).
bool
rb_zjit_iseq_tracing_currently_enabled(void)
Copy link
Copy Markdown
Contributor

@tekknolagi tekknolagi Mar 18, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why is this different from the yjit approach?

Copy link
Copy Markdown
Member Author

@k0kubun k0kubun Mar 18, 2026 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is a general problem for any fallbacks, like rb_vm_opt_send_without_block. It doesn't reproduce on that particular test case since YJIT specializes Kernel#send, which is called by the tracefunc test, and therefore doesn't use the fallback.

However, if you call something that becomes a fallback for YJIT, it's possible to trigger the same problem with YJIT. It should be addressed in a separate PR though.

def f_raise
  raise
rescue
  return :f_raise_return
end

def call_fwd(...)
  f_raise(...)
end

ary = []
TracePoint.new(:return){|tp|
  ary << [tp.method_id, tp.return_value] if tp.method_id == :f_raise
}.enable{
  call_fwd
}
p ary

On master:

$ .ruby/miniruby -v /tmp/bug.rb
ruby 4.1.0dev (2026-03-18T16:19:25Z master 00a31a46f2) +PRISM [x86_64-linux]
[[:f_raise, :f_raise_return]]

$ .ruby/miniruby -v --yjit-call-threshold=1 /tmp/bug.rb
ruby 4.1.0dev (2026-03-18T16:19:25Z master 00a31a46f2) +YJIT dev +PRISM [x86_64-linux]
[[:f_raise, nil]]

$ .ruby/miniruby -v --zjit-call-threshold=1 /tmp/bug.rb
ruby 4.1.0dev (2026-03-18T16:19:25Z master 00a31a46f2) +ZJIT dev +PRISM [x86_64-linux]
[[:f_raise, nil]]

On this branch:

$ .ruby/miniruby -v --zjit-call-threshold=1 /tmp/yjit_bug.rb
ruby 4.1.0dev (2026-03-18T04:35:56Z zjit-fix-tracefunc ebac36d240) +ZJIT dev +PRISM [x86_64-linux]
[[:f_raise, :f_raise_return]]

Copy link
Copy Markdown
Contributor

@tekknolagi tekknolagi Mar 18, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Sorry, I guess I mean why is our "is tracing enabled" function different than YJIT's?

Copy link
Copy Markdown
Member Author

@k0kubun k0kubun Mar 18, 2026 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm gonna guess you're asking why YJIT has rb_c_method_tracing_currently_enabled but ZJIT doesn't.

For C methods tracing, YJIT uses rb_c_method_tracing_currently_enabled. ZJIT uses rb_zjit_method_tracing_currently_enabled #14575. They look different but seem to be checking essentially the same thing. I'm not sure why we made them different in that PR, but I don't think it's the scope of this PR to make them consistent since they're for C methods tracing.

For Ruby methods, as I said above, YJIT doesn't handle it. ZJIT does now in this PR.

So that's how I see their difference.

Copy link
Copy Markdown
Contributor

@tekknolagi tekknolagi Mar 18, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

oh. gotcha. it is more aggressive. thank you!

Copy link
Copy Markdown
Contributor

@tekknolagi tekknolagi Mar 18, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I was reading the method name as "rb_c_method ... tracing_enabled" rather than "rb_ ... c_method_tracing_enabled"

{
rb_event_flag_t tracing_events;
if (rb_multi_ractor_p()) {
tracing_events = ruby_vm_event_enabled_global_flags;
}
else {
tracing_events = rb_ec_ractor_hooks(GET_EC())->events;
}

return tracing_events & ISEQ_TRACE_EVENTS;
}

bool
rb_zjit_insn_leaf(int insn, const VALUE *opes)
{
Expand Down
1 change: 1 addition & 0 deletions zjit/bindgen/src/main.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -293,6 +293,7 @@ fn main() {
.allowlist_function("rb_set_cfp_(pc|sp)")
.allowlist_function("rb_c_method_tracing_currently_enabled")
.allowlist_function("rb_zjit_method_tracing_currently_enabled")
.allowlist_function("rb_zjit_iseq_tracing_currently_enabled")
.allowlist_function("rb_full_cfunc_return")
.allowlist_function("rb_assert_(iseq|cme)_handle")
.allowlist_function("rb_IMEMO_TYPE_P")
Expand Down
24 changes: 24 additions & 0 deletions zjit/src/codegen_tests.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5107,3 +5107,27 @@ fn test_local_tracepoint() {
called
"), @"true");
}

// Regression test: TracePoint return value for methods with rescue that use `return`.
// ZJIT's send fallback uses rb_vm_opt_send_without_block which calls VM_EXEC,
// setting FLAG_FINISH on the callee frame. This changes how throw TAG_RETURN is
// handled, causing the return value to be nil instead of the actual value.
#[test]
fn test_tracepoint_return_value_with_rescue() {
assert_snapshot!(inspect("
def f_raise
raise
rescue
return :f_raise_return
end

ary = []
TracePoint.new(:return, :b_return){|tp|
ary << [tp.event, tp.method_id, tp.return_value]
}.enable{
send :f_raise
}
ary.pop # last b_return event is not required
ary
"), @"[[:return, :f_raise, :f_raise_return]]");
}
1 change: 1 addition & 0 deletions zjit/src/cruby_bindings.inc.rs
View file
Open in desktop

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

37 changes: 37 additions & 0 deletions zjit/src/hir.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -532,6 +532,7 @@ pub enum SideExitReason {
SplatKwPolymorphic,
SplatKwNotProfiled,
DirectiveInduced,
SendWhileTracing,
}

#[derive(Debug, Clone, Copy)]
Expand Down Expand Up @@ -7517,6 +7518,11 @@ pub fn iseq_to_hir(iseq: *const rb_iseq_t) -> Result<Function, ParseError> {
}
let argc = unsafe { vm_ci_argc((*cd).ci) };

// Side-exit send fallbacks while tracing to avoid FLAG_FINISH breaking throw TAG_RETURN semantics
if unsafe { rb_zjit_iseq_tracing_currently_enabled() } {
fun.push_insn(block, Insn::SideExit { state: exit_id, reason: SideExitReason::SendWhileTracing });
break;
}
let args = state.stack_pop_n(argc as usize)?;
let recv = state.stack_pop()?;
let send = fun.push_insn(block, Insn::Send { recv, cd, blockiseq: None, args, state: exit_id, reason: Uncategorized(opcode) });
Expand Down Expand Up @@ -7646,6 +7652,12 @@ pub fn iseq_to_hir(iseq: *const rb_iseq_t) -> Result<Function, ParseError> {
}
}

// Side-exit send fallbacks while tracing to avoid FLAG_FINISH breaking throw TAG_RETURN semantics
if unsafe { rb_zjit_iseq_tracing_currently_enabled() } {
fun.push_insn(block, Insn::SideExit { state: exit_id, reason: SideExitReason::SendWhileTracing });
break;
}

{
fn new_branch_block(
fun: &mut Function,
Expand Down Expand Up @@ -7732,6 +7744,11 @@ pub fn iseq_to_hir(iseq: *const rb_iseq_t) -> Result<Function, ParseError> {
fun.push_insn(block, Insn::SideExit { state: exit_id, reason: SideExitReason::UnhandledCallType(call_type) });
break; // End the block
}
// Side-exit send fallbacks while tracing to avoid FLAG_FINISH breaking throw TAG_RETURN semantics
if unsafe { rb_zjit_iseq_tracing_currently_enabled() } {
fun.push_insn(block, Insn::SideExit { state: exit_id, reason: SideExitReason::SendWhileTracing });
break;
}
let argc = unsafe { vm_ci_argc((*cd).ci) };
let block_arg = (flags & VM_CALL_ARGS_BLOCKARG) != 0;

Expand Down Expand Up @@ -7766,6 +7783,11 @@ pub fn iseq_to_hir(iseq: *const rb_iseq_t) -> Result<Function, ParseError> {
fun.push_insn(block, Insn::SideExit { state: exit_id, reason: SideExitReason::UnhandledCallType(call_type) });
break; // End the block
}
// Side-exit send fallbacks while tracing to avoid FLAG_FINISH breaking throw TAG_RETURN semantics
if unsafe { rb_zjit_iseq_tracing_currently_enabled() } {
fun.push_insn(block, Insn::SideExit { state: exit_id, reason: SideExitReason::SendWhileTracing });
break;
}
let argc = unsafe { vm_ci_argc((*cd).ci) };

let args = state.stack_pop_n(argc as usize + usize::from(forwarding))?;
Expand Down Expand Up @@ -7795,6 +7817,11 @@ pub fn iseq_to_hir(iseq: *const rb_iseq_t) -> Result<Function, ParseError> {
fun.push_insn(block, Insn::SideExit { state: exit_id, reason: SideExitReason::UnhandledCallType(call_type) });
break; // End the block
}
// Side-exit send fallbacks while tracing to avoid FLAG_FINISH breaking throw TAG_RETURN semantics
if unsafe { rb_zjit_iseq_tracing_currently_enabled() } {
fun.push_insn(block, Insn::SideExit { state: exit_id, reason: SideExitReason::SendWhileTracing });
break;
}
let argc = unsafe { vm_ci_argc((*cd).ci) };
let block_arg = (flags & VM_CALL_ARGS_BLOCKARG) != 0;
let args = state.stack_pop_n(argc as usize + usize::from(block_arg))?;
Expand Down Expand Up @@ -7829,6 +7856,11 @@ pub fn iseq_to_hir(iseq: *const rb_iseq_t) -> Result<Function, ParseError> {
fun.push_insn(block, Insn::SideExit { state: exit_id, reason: SideExitReason::UnhandledCallType(call_type) });
break; // End the block
}
// Side-exit send fallbacks while tracing to avoid FLAG_FINISH breaking throw TAG_RETURN semantics
if unsafe { rb_zjit_iseq_tracing_currently_enabled() } {
fun.push_insn(block, Insn::SideExit { state: exit_id, reason: SideExitReason::SendWhileTracing });
break;
}
let argc = unsafe { vm_ci_argc((*cd).ci) };
let args = state.stack_pop_n(argc as usize + usize::from(forwarding))?;
let recv = state.stack_pop()?;
Expand Down Expand Up @@ -7859,6 +7891,11 @@ pub fn iseq_to_hir(iseq: *const rb_iseq_t) -> Result<Function, ParseError> {
fun.push_insn(block, Insn::SideExit { state: exit_id, reason: SideExitReason::UnhandledCallType(call_type) });
break; // End the block
}
// Side-exit send fallbacks while tracing to avoid FLAG_FINISH breaking throw TAG_RETURN semantics
if unsafe { rb_zjit_iseq_tracing_currently_enabled() } {
fun.push_insn(block, Insn::SideExit { state: exit_id, reason: SideExitReason::SendWhileTracing });
break;
}
let argc = unsafe { vm_ci_argc((*cd).ci) };
let block_arg = (flags & VM_CALL_ARGS_BLOCKARG) != 0;
let args = state.stack_pop_n(argc as usize + usize::from(block_arg))?;
Expand Down
2 changes: 2 additions & 0 deletions zjit/src/stats.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -232,6 +232,7 @@ make_counters! {
exit_splatkw_polymorphic,
exit_splatkw_not_profiled,
exit_directive_induced,
exit_send_while_tracing,
}

// Send fallback counters that are summed as dynamic_send_count
Expand Down Expand Up @@ -616,6 +617,7 @@ pub fn side_exit_counter(reason: crate::hir::SideExitReason) -> Counter {
=> exit_patchpoint_no_singleton_class,
PatchPoint(Invariant::RootBoxOnly)
=> exit_patchpoint_root_box_only,
SendWhileTracing => exit_send_while_tracing,
}
}

Expand Down
Loading