ZJIT: Implement SingleRactorMode invalidation

Author: st0012

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

zjit.h

@@ -22,6 +22,7 @@ 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) {}
@@ -31,6 +32,7 @@ static inline void rb_zjit_bop_redefined(int redefined_flag, enum ruby_basic_ope
 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) {}
+static inline void rb_zjit_before_ractor_spawn(void) {}
 #endif // #if USE_YJIT
 
 #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);
             }
-            _ => {
-                debug!("ZJIT: gen_patch_point: unimplemented invariant {invariant:?}");
-                return;
+            Invariant::SingleRactorMode => {
+                let side_exit_ptr = cb.resolve_label(label);
+                track_single_ractor_assumption(code_ptr, side_exit_ptr);
             }
         }
     });

zjit/src/invariants.rs

@@ -1,4 +1,4 @@
-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};
 
@@ -26,6 +26,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
@@ -200,3 +203,38 @@ pub extern "C" fn rb_zjit_constant_state_changed(id: ID) {
         }
     });
 }
+
+/// Track that JIT code for a patch point 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 then 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!(), || {
+        // Clear the set of blocks inside Invariants
+        let jumps = mem::take(&mut ZJITState::get_invariants().single_ractor_patch_points);
+
+        // Invalidate the blocks
+        for jump in &jumps {
+            let cb = ZJITState::get_code_block();
+            cb.with_write_ptr(jump.from, |cb| {
+                let mut asm = Assembler::new();
+                asm_comment!(asm, "Single ractor mode is invalidated");
+                asm.jmp(jump.to.into());
+                asm.compile(cb).expect("can write existing code");
+            });
+        }
+    });
+}