ZJIT: Implement SingleRactorMode invalidation (#14121)

* ZJIT: Implement SingleRactorMode invalidation

* ZJIT: Add macro for compiling jumps

* ZJIT: Fix typo in comment

* YJIT: Fix typo in comment

* ZJIT: Avoid using unexported types in zjit.h

`enum ruby_vminsn_type` is declared in `insns.inc` and is not exported.
Using it in `zjit.h` would cause build errors when the file including it
doesn't include `insns.inc`.

depend

@@ -12702,6 +12702,7 @@ ractor.$(OBJEXT): {$(VPATH)}vm_debug.h
 ractor.$(OBJEXT): {$(VPATH)}vm_opts.h
 ractor.$(OBJEXT): {$(VPATH)}vm_sync.h
 ractor.$(OBJEXT): {$(VPATH)}yjit.h
+ractor.$(OBJEXT): {$(VPATH)}zjit.h
 random.$(OBJEXT): $(CCAN_DIR)/check_type/check_type.h
 random.$(OBJEXT): $(CCAN_DIR)/container_of/container_of.h
 random.$(OBJEXT): $(CCAN_DIR)/list/list.h

ractor.c

@@ -19,6 +19,7 @@
 #include "internal/thread.h"
 #include "variable.h"
 #include "yjit.h"
+#include "zjit.h"
 
 VALUE rb_cRactor;
 static VALUE rb_cRactorSelector;
@@ -511,6 +512,7 @@ ractor_create(rb_execution_context_t *ec, VALUE self, VALUE loc, VALUE name, VAL
     r->debug = cr->debug;
 
     rb_yjit_before_ractor_spawn();
+    rb_zjit_before_ractor_spawn();
     rb_thread_create_ractor(r, args, block);
 
     RB_GC_GUARD(rv);

test/ruby/test_zjit.rb

@@ -950,6 +950,26 @@ class TestZJIT < Test::Unit::TestCase
     RUBY
   end
 
+  def test_single_ractor_mode_invalidation
+    # Without invalidating the single-ractor mode, the test would crash
+    assert_compiles '"errored but not crashed"', <<~RUBY, call_threshold: 2, insns: [:opt_getconstant_path]
+      C = Object.new
+
+      def test
+        C
+      rescue Ractor::IsolationError
+        "errored but not crashed"
+      end
+
+      test
+      test
+
+      Ractor.new {
+        test
+      }.value
+    RUBY
+  end
+
   def test_dupn
     assert_compiles '[[1], [1, 1], :rhs, [nil, :rhs]]', <<~RUBY, insns: [:dupn]
       def test(array) = (array[1, 2] ||= :rhs)

yjit/src/invariants.rs

@@ -303,7 +303,7 @@ pub extern "C" fn rb_yjit_cme_invalidate(callee_cme: *const rb_callable_method_e
     });
 }
 
-/// Callback for then Ruby is about to spawn a ractor. In that case we need to
+/// Callback for when Ruby is about to spawn a ractor. In that case we need to
 /// invalidate every block that is assuming single ractor mode.
 #[no_mangle]
 pub extern "C" fn rb_yjit_before_ractor_spawn() {

zjit.h

@@ -14,7 +14,7 @@ extern bool rb_zjit_enabled_p;
 extern uint64_t rb_zjit_call_threshold;
 extern uint64_t rb_zjit_profile_threshold;
 void rb_zjit_compile_iseq(const rb_iseq_t *iseq, rb_execution_context_t *ec, bool jit_exception);
-void rb_zjit_profile_insn(enum ruby_vminsn_type insn, rb_execution_context_t *ec);
+void rb_zjit_profile_insn(uint32_t insn, rb_execution_context_t *ec);
 void rb_zjit_profile_enable(const rb_iseq_t *iseq);
 void rb_zjit_bop_redefined(int redefined_flag, enum ruby_basic_operators bop);
 void rb_zjit_cme_invalidate(const rb_callable_method_entry_t *cme);
@@ -22,15 +22,17 @@ void rb_zjit_invalidate_ep_is_bp(const rb_iseq_t *iseq);
 void rb_zjit_constant_state_changed(ID id);
 void rb_zjit_iseq_mark(void *payload);
 void rb_zjit_iseq_update_references(void *payload);
+void rb_zjit_before_ractor_spawn(void);
 #else
 #define rb_zjit_enabled_p false
 static inline void rb_zjit_compile_iseq(const rb_iseq_t *iseq, rb_execution_context_t *ec, bool jit_exception) {}
-static inline void rb_zjit_profile_insn(enum ruby_vminsn_type insn, rb_execution_context_t *ec) {}
+static inline void rb_zjit_profile_insn(uint32_t insn, rb_execution_context_t *ec) {}
 static inline void rb_zjit_profile_enable(const rb_iseq_t *iseq) {}
 static inline void rb_zjit_bop_redefined(int redefined_flag, enum ruby_basic_operators bop) {}
 static inline void rb_zjit_cme_invalidate(const rb_callable_method_entry_t *cme) {}
 static inline void rb_zjit_invalidate_ep_is_bp(const rb_iseq_t *iseq) {}
 static inline void rb_zjit_constant_state_changed(ID id) {}
-#endif // #if USE_YJIT
+static inline void rb_zjit_before_ractor_spawn(void) {}
+#endif // #if USE_ZJIT
 
 #endif // #ifndef ZJIT_H

zjit/src/codegen.rs

@@ -4,7 +4,7 @@ use std::ffi::{c_int, c_void};
 
 use crate::asm::Label;
 use crate::backend::current::{Reg, ALLOC_REGS};
-use crate::invariants::{track_bop_assumption, track_cme_assumption, track_stable_constant_names_assumption};
+use crate::invariants::{track_bop_assumption, track_cme_assumption, track_single_ractor_assumption, track_stable_constant_names_assumption};
 use crate::gc::{get_or_create_iseq_payload, append_gc_offsets};
 use crate::state::ZJITState;
 use crate::stats::{counter_ptr, Counter};
@@ -542,9 +542,9 @@ fn gen_patch_point(jit: &mut JITState, asm: &mut Assembler, invariant: &Invarian
                 let side_exit_ptr = cb.resolve_label(label);
                 track_stable_constant_names_assumption(idlist, code_ptr, side_exit_ptr);
             }
-            _ => {
+            Invariant::SingleRactorMode => {
-                debug!("ZJIT: gen_patch_point: unimplemented invariant {invariant:?}");
+                let side_exit_ptr = cb.resolve_label(label);
-                return;
+                track_single_ractor_assumption(code_ptr, side_exit_ptr);
             }
         }
     });

zjit/src/invariants.rs

@@ -1,7 +1,20 @@
-use std::{collections::{HashMap, HashSet}};
+use std::{collections::{HashMap, HashSet}, mem};
 
 use crate::{backend::lir::{asm_comment, Assembler}, cruby::{rb_callable_method_entry_t, ruby_basic_operators, src_loc, with_vm_lock, IseqPtr, RedefinitionFlag, ID}, hir::Invariant, options::debug, state::{zjit_enabled_p, ZJITState}, virtualmem::CodePtr};
 
+macro_rules! compile_jumps {
+    ($cb:expr, $jumps:expr, $($comment_args:tt)*) => {
+        for jump in $jumps {
+            $cb.with_write_ptr(jump.from, |cb| {
+                let mut asm = Assembler::new();
+                asm_comment!(asm, $($comment_args)*);
+                asm.jmp(jump.to.into());
+                asm.compile(cb).expect("can write existing code");
+            });
+        }
+    };
+}
+
 #[derive(Debug, Eq, Hash, PartialEq)]
 struct Jump {
     from: CodePtr,
@@ -26,6 +39,9 @@ pub struct Invariants {
 
     /// Map from constant ID to patch points that assume the constant hasn't been redefined
     constant_state_patch_points: HashMap<ID, HashSet<Jump>>,
+
+    /// Set of patch points that assume that the interpreter is running with only one ractor
+    single_ractor_patch_points: HashSet<Jump>,
 }
 
 /// Called when a basic operator is redefined. Note that all the blocks assuming
@@ -46,14 +62,7 @@ pub extern "C" fn rb_zjit_bop_redefined(klass: RedefinitionFlag, bop: ruby_basic
             debug!("BOP is redefined: {}", bop);
 
             // Invalidate all patch points for this BOP
-            for jump in jumps {
+            compile_jumps!(cb, jumps, "BOP is redefined: {}", bop);
-                cb.with_write_ptr(jump.from, |cb| {
-                    let mut asm = Assembler::new();
-                    asm_comment!(asm, "BOP is redefined: {}", bop);
-                    asm.jmp(jump.to.into());
-                    asm.compile(cb).expect("can write existing code");
-                });
-            }
 
             cb.mark_all_executable();
         }
@@ -159,14 +168,8 @@ pub extern "C" fn rb_zjit_cme_invalidate(cme: *const rb_callable_method_entry_t)
             debug!("CME is invalidated: {:?}", cme);
 
             // Invalidate all patch points for this CME
-            for jump in jumps {
+            compile_jumps!(cb, jumps, "CME is invalidated: {:?}", cme);
-                cb.with_write_ptr(jump.from, |cb| {
+
-                    let mut asm = Assembler::new();
-                    asm_comment!(asm, "CME is invalidated: {:?}", cme);
-                    asm.jmp(jump.to.into());
-                    asm.compile(cb).expect("can write existing code");
-                });
-            }
             cb.mark_all_executable();
         }
     });
@@ -187,16 +190,38 @@ pub extern "C" fn rb_zjit_constant_state_changed(id: ID) {
             debug!("Constant state changed: {:?}", id);
 
             // Invalidate all patch points for this constant ID
-            for jump in jumps {
+            compile_jumps!(cb, jumps, "Constant state changed: {:?}", id);
-                cb.with_write_ptr(jump.from, |cb| {
-                    let mut asm = Assembler::new();
-                    asm_comment!(asm, "Constant state changed: {:?}", id);
-                    asm.jmp(jump.to.into());
-                    asm.compile(cb).expect("can write existing code");
-                });
-            }
 
             cb.mark_all_executable();
         }
     });
 }
+
+/// Track the JIT code that assumes that the interpreter is running with only one ractor
+pub fn track_single_ractor_assumption(patch_point_ptr: CodePtr, side_exit_ptr: CodePtr) {
+    let invariants = ZJITState::get_invariants();
+    invariants.single_ractor_patch_points.insert(Jump {
+        from: patch_point_ptr,
+        to: side_exit_ptr,
+    });
+}
+
+/// Callback for when Ruby is about to spawn a ractor. In that case we need to
+/// invalidate every block that is assuming single ractor mode.
+#[unsafe(no_mangle)]
+pub extern "C" fn rb_zjit_before_ractor_spawn() {
+    // If ZJIT isn't enabled, do nothing
+    if !zjit_enabled_p() {
+        return;
+    }
+
+    with_vm_lock(src_loc!(), || {
+        let cb = ZJITState::get_code_block();
+        let jumps = mem::take(&mut ZJITState::get_invariants().single_ractor_patch_points);
+
+        // Invalidate all patch points for single ractor mode
+        compile_jumps!(cb, jumps, "Another ractor spawned, invalidating single ractor mode assumption");
+
+        cb.mark_all_executable();
+    });
+}

zjit/src/profile.rs

@@ -39,10 +39,10 @@ impl Profiler {
 
 /// API called from zjit_* instruction. opcode is the bare (non-zjit_*) instruction.
 #[unsafe(no_mangle)]
-pub extern "C" fn rb_zjit_profile_insn(bare_opcode: ruby_vminsn_type, ec: EcPtr) {
+pub extern "C" fn rb_zjit_profile_insn(bare_opcode: u32, ec: EcPtr) {
     with_vm_lock(src_loc!(), || {
         let mut profiler = Profiler::new(ec);
-        profile_insn(&mut profiler, bare_opcode);
+        profile_insn(&mut profiler, bare_opcode as ruby_vminsn_type);
     });
 }