Merge pull request ruby#227

Move branch_stub_hit() to Rust

yjit.c

@@ -658,34 +658,15 @@ rb_RSTRUCT_SET(VALUE st, int k, VALUE v)
     RSTRUCT_SET(st, k, v);
 }
 
-const struct rb_callinfo*
-rb_get_call_data_ci(struct rb_call_data* cd) {
-    return cd->ci;
-}
-
-const uint8_t *
-rb_yjit_branch_stub_hit(void *branch_ptr, uint32_t target_idx, rb_execution_context_t *ec)
+const struct rb_callinfo *
+rb_get_call_data_ci(struct rb_call_data *cd)
 {
-    const uint8_t *ret;
-    // Acquire the VM lock and then signal all other Ruby threads (ractors) to
-    // contend for the VM lock, putting them to sleep. YJIT uses this to evict
-    // threads running inside generated code so among other things, it can
-    // safely change memory protection of regions housing generated code.
-    RB_VM_LOCK_ENTER();
-    rb_vm_barrier();
-
-    const uint8_t *rb_yjit_rust_branch_stub_hit(void *branch_ptr, uint32_t target_idx, rb_execution_context_t *ec);
-    ret = rb_yjit_rust_branch_stub_hit(branch_ptr, target_idx, ec);
-
-    // Release the VM lock. Note, watch out for Ruby exceptions and Rust panics
-    // to ensure that the lock is properly released in exceptional situations.
-    RB_VM_LOCK_LEAVE();
-
-    return ret;
+    return cd->ci;
 }
 
 bool
-rb_BASIC_OP_UNREDEFINED_P(enum ruby_basic_operators bop, uint32_t klass) {
+rb_BASIC_OP_UNREDEFINED_P(enum ruby_basic_operators bop, uint32_t klass)
+{
     return BASIC_OP_UNREDEFINED_P(bop, klass);
 }
 
@@ -700,6 +681,27 @@ rb_yjit_multi_ractor_p(void) {
     return rb_multi_ractor_p();
 }
 
+// Acquire the VM lock and then signal all other Ruby threads (ractors) to
+// contend for the VM lock, putting them to sleep. YJIT uses this to evict
+// threads running inside generated code so among other things, it can
+// safely change memory protection of regions housing generated code.
+//
+// Using rb_vm_lock() rather than RB_VM_LOCK_ENTER() because
+//   - Simpler signature for FFI
+//   - The lev parameter appears to be for debugging purposes only
+void
+rb_yjit_vm_lock_then_barrier(const char *file, int line)
+{
+    rb_vm_lock(file, line);
+    rb_vm_barrier();
+}
+
+void
+rb_yjit_vm_unlock(const char *file, int line)
+{
+    rb_vm_unlock(file, line);
+}
+
 #include "yjit_iface.c"
 
 #endif // if JIT_ENABLED && PLATFORM_SUPPORTED_P

yjit/bindgen/src/main.rs

@@ -202,6 +202,8 @@ fn main() {
         .allowlist_function("rb_yjit_multi_ractor_p")
         .allowlist_function("rb_c_method_tracing_currently_enabled")
         .allowlist_function("rb_full_cfunc_return")
+        .allowlist_function("rb_yjit_vm_lock_then_barrier")
+        .allowlist_function("rb_yjit_vm_unlock")
 
         // Not sure why it's picking these up, but don't.
         .blocklist_type("FILE")

yjit/src/core.rs

@@ -1234,10 +1234,20 @@ fn make_branch_entry(block: BlockRef, src_ctx: &Context, gen_fn: BranchGenFn) ->
     return branchref;
 }
 
-// Called by the generated code when a branch stub is executed
-// Triggers compilation of branches and code patching
-#[no_mangle]
-pub extern "C" fn rb_yjit_rust_branch_stub_hit(branch_ptr: *const c_void, target_idx: u32, ec: EcPtr) -> *const u8
+/// Generated code calls this function with the SysV calling convention.
+/// See [get_branch_target].
+extern "sysv64" fn branch_stub_hit(
+    branch_ptr: *const c_void,
+    target_idx: u32,
+    ec: EcPtr
+) -> *const u8
+{
+    with_vm_lock(src_loc!(), || branch_stub_hit_body(branch_ptr, target_idx, ec))
+}
+
+/// Called by the generated code when a branch stub is executed
+/// Triggers compilation of branches and code patching
+fn branch_stub_hit_body(branch_ptr: *const c_void, target_idx: u32, ec: EcPtr) -> *const u8
 {
     assert!(!branch_ptr.is_null());
 
@@ -1409,7 +1419,7 @@ fn get_branch_target(
     mov(ocb, C_ARG_REGS[2], REG_EC);
     mov(ocb, C_ARG_REGS[1], uimm_opnd(target_idx as u64));
     mov(ocb, C_ARG_REGS[0], const_ptr_opnd(branch_ptr as *const u8));
-    call_ptr(ocb, REG0, rb_yjit_branch_stub_hit as *mut u8);
+    call_ptr(ocb, REG0, branch_stub_hit as *mut u8);
 
     // Jump to the address returned by the
     // branch_stub_hit call

yjit/src/cruby.rs

@@ -85,6 +85,7 @@
 use std::ffi::CString;
 use std::convert::From;
 use std::os::raw::{c_int, c_uint, c_long, c_char, c_void};
+use std::panic::{UnwindSafe, catch_unwind};
 
 // We check that we can do this with the configure script and a couple of
 // static asserts. u64 and not usize to play nice with lowering to x86.
@@ -567,6 +568,61 @@ pub fn rust_str_to_sym(str: &str) -> VALUE
     }
 }
 
+/// A location in Rust code for integrating with debugging facilities defined in C.
+/// Use the [src_loc!] macro to crate an instance.
+pub struct SourceLocation {
+    pub file: CString,
+    pub line: c_int,
+}
+
+/// Make a [SourceLocation] at the current spot.
+macro_rules! src_loc {
+    () => {
+        crate::cruby::SourceLocation{
+            file: std::ffi::CString::new(file!()).unwrap(), // ASCII source file paths
+            line: line!().try_into().unwrap(), // not that many lines
+        }
+    }
+}
+
+pub(crate) use src_loc;
+
+/// Acquire the VM lock, make sure all other Ruby threads are asleep then run
+/// some code while holding the lock. Returns whatever `func` returns.
+/// Use with [src_loc!].
+///
+/// Required for code patching in the presence of ractors.
+pub fn with_vm_lock<F, R>(loc: SourceLocation, func: F) -> R
+    where F: FnOnce() -> R + UnwindSafe
+{
+    let file = loc.file.as_ptr();
+    let line = loc.line;
+
+    unsafe { rb_yjit_vm_lock_then_barrier(file, line) };
+
+    let ret = match catch_unwind(func) {
+        Ok(result) => result,
+        Err(_) => {
+            // Theoretically we can recover from some of these panics,
+            // but it's too late if the unwind reaches here.
+            use std::{io, process, str};
+
+            let _ = catch_unwind(|| { // IO functions can panic too.
+                eprintln!(
+                    "YJIT panicked while holding VM lock acquired at {}:{}. Aborting...",
+                    str::from_utf8(loc.file.as_bytes()).unwrap_or("<not utf8>"),
+                    line,
+                );
+            });
+            process::abort();
+        }
+    };
+
+    unsafe { rb_yjit_vm_unlock(file, line) };
+
+    ret
+}
+
 // Non-idiomatic capitalization for consistency with CRuby code
 #[allow(non_upper_case_globals)]
 pub const Qfalse: VALUE = VALUE(0);

yjit/src/cruby_bindings.inc.rs

@@ -684,3 +684,12 @@ extern "C" {
 extern "C" {
     pub fn rb_yjit_multi_ractor_p() -> bool;
 }
+extern "C" {
+    pub fn rb_yjit_vm_lock_then_barrier(
+        file: *const ::std::os::raw::c_char,
+        line: ::std::os::raw::c_int,
+    );
+}
+extern "C" {
+    pub fn rb_yjit_vm_unlock(file: *const ::std::os::raw::c_char, line: ::std::os::raw::c_int);
+}