YJIT: Cancel on-stack jit_return on invalidation (#9086)

* YJIT: Cancel on-stack jit_return on invalidation

Co-authored-by: Alan Wu <alansi.xingwu@shopify.com>

* Use RUBY_VM_CONTROL_FRAME_STACK_OVERFLOW_P

---------

Co-authored-by: Alan Wu <alansi.xingwu@shopify.com>

1 files changed, 12 insertions, 0 deletions

diff --git a/yjit/src/invariants.rs b/yjit/src/invariants.rs
index 13995ce5fb..ad684cd54e 100644
--- a/yjit/src/invariants.rs
+++ b/yjit/src/invariants.rs
@@ -11,6 +11,7 @@ use crate::utils::IntoUsize;
 use crate::yjit::yjit_enabled_p;
 
 use std::collections::{HashMap, HashSet};
+use std::os::raw::c_void;
 use std::mem;
 
 // Invariants to track:
@@ -517,6 +518,17 @@ pub extern "C" fn rb_yjit_tracing_invalidate_all() {
 
         let cb = CodegenGlobals::get_inline_cb();
 
+        // Prevent on-stack frames from jumping to the caller on jit_exec_exception
+        extern "C" {
+            fn rb_yjit_cancel_jit_return(leave_exit: *mut c_void, leave_exception: *mut c_void) -> VALUE;
+        }
+        unsafe {
+            rb_yjit_cancel_jit_return(
+                CodegenGlobals::get_leave_exit_code().raw_ptr(cb) as _,
+                CodegenGlobals::get_leave_exception_code().raw_ptr(cb) as _,
+            );
+        }
+
         // Apply patches
         let old_pos = cb.get_write_pos();
         let old_dropped_bytes = cb.has_dropped_bytes();