YJIT: Set code mem permissions in bulk

Some GC modules, notably MMTk, support parallel GC, i.e. multiple GC
threads work in parallel during a GC.  Currently, when two GC threads
scan two iseq objects simultaneously when YJIT is enabled, both threads
will attempt to borrow `CodeBlock::mem_block`, which will result in
panic.

This commit makes one part of the change.

We now set the YJIT code memory to writable in bulk before the
reference-updating phase, and reset it to executable in bulk after the
reference-updating phase.  Previously, YJIT lazily sets memory pages
writable while updating object references embedded in JIT-compiled
machine code, and sets the memory back to executable by calling
`mark_all_executable`.  This approach is inherently unfriendly to
parallel GC because (1) it borrows `CodeBlock::mem_block`, and (2) it
sets the whole `CodeBlock` as executable which races with other GC
threads that are updating other iseq objects.  It also has performance
overhead due to the frequent invocation of system calls.  We now set the
permission of all the code memory in bulk before and after the reference
updating phase.  Multiple GC threads can now perform raw memory writes
in parallel.  We should also see performance improvement during moving
GC because of the reduced number of `mprotect` system calls.

yjit/src/asm/mod.rs

@@ -590,6 +590,10 @@ impl CodeBlock {
         self.label_refs = state.label_refs;
     }
 
+    pub fn mark_all_writeable(&mut self) {
+        self.mem_block.borrow_mut().mark_all_writeable();
+    }
+
     pub fn mark_all_executable(&mut self) {
         self.mem_block.borrow_mut().mark_all_executable();
     }

yjit/src/core.rs

@@ -2035,13 +2035,6 @@ pub extern "C" fn rb_yjit_iseq_update_references(iseq: IseqPtr) {
         block_update_references(block, cb, true);
     }
 
-    // Note that we would have returned already if YJIT is off.
-    cb.mark_all_executable();
-
-    CodegenGlobals::get_outlined_cb()
-        .unwrap()
-        .mark_all_executable();
-
     return;
 
     fn block_update_references(block: &Block, cb: &mut CodeBlock, dead: bool) {
@@ -2110,6 +2103,34 @@ pub extern "C" fn rb_yjit_iseq_update_references(iseq: IseqPtr) {
     }
 }
 
+/// Mark all code memory as writable.
+/// This function is useful for garbage collectors that update references in JIT-compiled code in
+/// bulk.
+#[no_mangle]
+pub extern "C" fn rb_yjit_mark_all_writeable() {
+    if CodegenGlobals::has_instance() {
+        CodegenGlobals::get_inline_cb().mark_all_writeable();
+
+        CodegenGlobals::get_outlined_cb()
+            .unwrap()
+            .mark_all_writeable();
+    }
+}
+
+/// Mark all code memory as executable.
+/// This function is useful for garbage collectors that update references in JIT-compiled code in
+/// bulk.
+#[no_mangle]
+pub extern "C" fn rb_yjit_mark_all_executable() {
+    if CodegenGlobals::has_instance() {
+        CodegenGlobals::get_inline_cb().mark_all_executable();
+
+        CodegenGlobals::get_outlined_cb()
+            .unwrap()
+            .mark_all_executable();
+    }
+}
+
 /// Get all blocks for a particular place in an iseq.
 fn get_version_list(blockid: BlockId) -> Option<&'static mut VersionList> {
     let insn_idx = blockid.idx.as_usize();

yjit/src/virtualmem.rs

@@ -231,6 +231,23 @@ impl<A: Allocator> VirtualMemory<A> {
         Ok(())
     }
 
+    /// Make all the code in the region writeable.
+    /// Call this during GC before the phase of updating reference fields.
+    pub fn mark_all_writeable(&mut self) {
+        self.current_write_page = None;
+
+        let region_start = self.region_start;
+        let mapped_region_bytes: u32 = self.mapped_region_bytes.try_into().unwrap();
+
+        // Make mapped region executable
+        if !self.allocator.mark_writable(region_start.as_ptr(), mapped_region_bytes) {
+            panic!("Cannot make memory region writable: {:?}-{:?}",
+                region_start.as_ptr(),
+                unsafe { region_start.as_ptr().add(mapped_region_bytes as usize)}
+            );
+        }
+    }
+
     /// Make all the code in the region executable. Call this at the end of a write session.
     /// See [Self] for usual usage flow.
     pub fn mark_all_executable(&mut self) {