YJIT: Assert that we have the VM lock while marking

Somewhat important because having the lock is a key part of the
soundness reasoning for the `unsafe` usage here.

4 files changed, 29 insertions, 6 deletions

diff --git a/yjit.c b/yjit.c
index 5bb8a37e53..a9d1d2c38b 100644
--- a/yjit.c
+++ b/yjit.c
@@ -1100,6 +1100,14 @@ object_shape_count(rb_execution_context_t *ec, VALUE self)
     return ULONG2NUM((unsigned long)GET_VM()->next_shape_id);
 }
 
+// Assert that we have the VM lock. Relevant mostly for multi ractor situations.
+// The GC takes the lock before calling us, and this asserts that it indeed happens.
+void
+rb_yjit_assert_holding_vm_lock(void)
+{
+    ASSERT_vm_locking();
+}
+
 // Primitives used by yjit.rb
 VALUE rb_yjit_stats_enabled_p(rb_execution_context_t *ec, VALUE self);
 VALUE rb_yjit_trace_exit_locations_enabled_p(rb_execution_context_t *ec, VALUE self);
diff --git a/yjit/bindgen/src/main.rs b/yjit/bindgen/src/main.rs
index bf762c1401..3dc503f7c3 100644
--- a/yjit/bindgen/src/main.rs
+++ b/yjit/bindgen/src/main.rs
@@ -320,6 +320,8 @@ fn main() {
         .allowlist_function("rb_yjit_exit_locations_dict")
         .allowlist_function("rb_yjit_icache_invalidate")
         .allowlist_function("rb_optimized_call")
+        .allowlist_function("rb_yjit_assert_holding_vm_lock")
+
         // from vm_sync.h
         .allowlist_function("rb_vm_barrier")
 
diff --git a/yjit/src/core.rs b/yjit/src/core.rs
index 4467145f73..5686808b3a 100644
--- a/yjit/src/core.rs
+++ b/yjit/src/core.rs
@@ -863,9 +863,15 @@ pub extern "C" fn rb_yjit_iseq_mark(payload: *mut c_void) {
         // Nothing to mark.
         return;
     } else {
-        // SAFETY: It looks like the GC takes the VM lock while marking
-        // so we should be satisfying aliasing rules here.
-        unsafe { &*(payload as *const IseqPayload) }
+        // SAFETY: The GC takes the VM lock while marking, which
+        // we assert, so we should be synchronized and data race free.
+        //
+        // For aliasing, having the VM lock hopefully also implies that no one
+        // else has an overlapping &mut IseqPayload.
+        unsafe {
+            rb_yjit_assert_holding_vm_lock();
+            &*(payload as *const IseqPayload)
+        }
     };
 
     // For marking VALUEs written into the inline code block.
@@ -915,9 +921,15 @@ pub extern "C" fn rb_yjit_iseq_update_references(payload: *mut c_void) {
         // Nothing to update.
         return;
     } else {
-        // SAFETY: It looks like the GC takes the VM lock while updating references
-        // so we should be satisfying aliasing rules here.
-        unsafe { &*(payload as *const IseqPayload) }
+        // SAFETY: The GC takes the VM lock while marking, which
+        // we assert, so we should be synchronized and data race free.
+        //
+        // For aliasing, having the VM lock hopefully also implies that no one
+        // else has an overlapping &mut IseqPayload.
+        unsafe {
+            rb_yjit_assert_holding_vm_lock();
+            &*(payload as *const IseqPayload)
+        }
     };
 
     // Evict other threads from generated code since we are about to patch them.
diff --git a/yjit/src/cruby_bindings.inc.rs b/yjit/src/cruby_bindings.inc.rs
index 190b163b41..19cea5e682 100644
--- a/yjit/src/cruby_bindings.inc.rs
+++ b/yjit/src/cruby_bindings.inc.rs
@@ -1333,4 +1333,5 @@ extern "C" {
         file: *const ::std::os::raw::c_char,
         line: ::std::os::raw::c_int,
     );
+    pub fn rb_yjit_assert_holding_vm_lock();
 }