YJIT: Cancel on-stack jit_return on invalidation

cont.c

@@ -1290,18 +1290,37 @@ rb_jit_cont_each_iseq(rb_iseq_callback callback, void *data)
         if (cont->ec->vm_stack == NULL)
             continue;
 
-        const rb_control_frame_t *cfp;
-        for (cfp = RUBY_VM_END_CONTROL_FRAME(cont->ec) - 1; ; cfp = RUBY_VM_NEXT_CONTROL_FRAME(cfp)) {
-            const rb_iseq_t *iseq;
-            if (cfp->pc && (iseq = cfp->iseq) != NULL && imemo_type((VALUE)iseq) == imemo_iseq) {
-                callback(iseq, data);
+        const rb_control_frame_t *cfp = cont->ec->cfp;
+        while (!RUBY_VM_CONTROL_FRAME_STACK_OVERFLOW_P(cont->ec, cfp)) {
+            if (cfp->pc && cfp->iseq && imemo_type((VALUE)cfp->iseq) == imemo_iseq) {
+                callback(cfp->iseq, data);
             }
+            cfp = RUBY_VM_PREVIOUS_CONTROL_FRAME(cfp);
+        }
+    }
+}
+
+#if USE_YJIT
+// Update the jit_return of all CFPs to leave_exit unless it's leave_exception or not set.
+// This prevents jit_exec_exception from jumping to the caller after invalidation.
+void
+rb_yjit_cancel_jit_return(void *leave_exit, void *leave_exception)
+{
+    struct rb_jit_cont *cont;
+    for (cont = first_jit_cont; cont != NULL; cont = cont->next) {
+        if (cont->ec->vm_stack == NULL)
+            continue;
 
-            if (cfp == cont->ec->cfp)
-                break; // reached the most recent cfp
+        const rb_control_frame_t *cfp = cont->ec->cfp;
+        while (!RUBY_VM_CONTROL_FRAME_STACK_OVERFLOW_P(cont->ec, cfp)) {
+            if (cfp->jit_return && cfp->jit_return != leave_exception) {
+                ((rb_control_frame_t *)cfp)->jit_return = leave_exit;
+            }
+            cfp = RUBY_VM_PREVIOUS_CONTROL_FRAME(cfp);
         }
     }
 }
+#endif
 
 // Finish working with jit_cont.
 void

test/ruby/test_yjit.rb

@@ -1324,6 +1324,36 @@ def freeze = :ugokanai
     RUBY
   end
 
+  def test_return_to_invalidated_frame
+    assert_compiles(code_gc_helpers + <<~RUBY, exits: :any, result: :ok)
+      def jump
+        [] # something not inlined
+      end
+
+      def entry(code_gc)
+        jit_exception(code_gc)
+        jump # faulty jump after code GC. #jit_exception should not come back.
+      end
+
+      def jit_exception(code_gc)
+        if code_gc
+          tap do
+            RubyVM::YJIT.code_gc
+            break # jit_exec_exception catches TAG_BREAK and re-enters JIT code
+          end
+        end
+      end
+
+      add_pages(100)
+      jump           # Compile #jump in a non-first page
+      add_pages(100)
+      entry(false)   # Compile #entry and its call to #jump in another page
+      entry(true)    # Free #jump but not #entry
+
+      :ok
+    RUBY
+  end
+
   def test_setivar_on_class
     # Bug in https://github.com/ruby/ruby/pull/8152
     assert_compiles(<<~RUBY, result: :ok)

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();