Unprotect and protect code memory during GC

mmtk/Cargo.toml

@@ -12,7 +12,7 @@ edition = "2021"
 # Metadata for the Ruby repository
 [package.metadata.ci-repos.ruby]
 repo = "mmtk/ruby" # This is used by actions/checkout, so the format is "owner/repo", not URL.
-rev = "f3f893c6fb8cbe08c3bc4c68d4cd0deceae42461"
+rev = "2bfc164e05d672cac6dfb5dad303636c6195828a"
 
 [lib]
 name = "mmtk_ruby"
@@ -37,7 +37,7 @@ features = ["is_mmtk_object", "object_pinning", "sticky_immix_non_moving_nursery
 
 # Uncomment the following lines to use mmtk-core from the official repository.
 git = "https://github.com/mmtk/mmtk-core.git"
-rev = "6ac4f73e21e0b08b2de1d456175ae7caaaf886fd"
+rev = "31a78a41f02fc7228780b501c4944ba750e32ee4"
 
 # Uncomment the following line to use mmtk-core from a local repository.
 #path = "../../mmtk-core"

mmtk/src/abi.rs

@@ -390,6 +390,9 @@ pub struct RubyUpcalls {
         end: usize,
         stats: *mut ConcurrentSetStats,
     ),
+    // Memory protection for code memory
+    pub before_updating_jit_code: extern "C" fn(),
+    pub after_updating_jit_code: extern "C" fn(),
 }
 
 unsafe impl Sync for RubyUpcalls {}

mmtk/src/api.rs

@@ -408,3 +408,16 @@ pub unsafe extern "C" fn mmtk_hidden_header_is_sane(hidden_header: *const Hidden
     let hidden_header = unsafe { &*hidden_header };
     hidden_header.is_sane()
 }
+
+#[no_mangle]
+pub extern "C" fn mmtk_current_gc_may_move_object() -> bool {
+    crate::mmtk().get_plan().current_gc_may_move_object()
+}
+
+#[no_mangle]
+pub extern "C" fn mmtk_current_gc_is_nursery() -> bool {
+    crate::mmtk()
+        .get_plan()
+        .generational()
+        .is_some_and(|gen| gen.is_current_gc_nursery())
+}

mmtk/src/collection.rs

@@ -26,6 +26,7 @@ impl Collection<Ruby> for VMCollection {
             Self::notify_mutator_ready::<F>,
             &mut mutator_visitor as *mut F as *mut _,
         );
+        crate::yjit_support::schedule_jit_code_protection_work_packets(tls);
     }
 
     fn resume_mutators(tls: VMWorkerThread) {

mmtk/src/lib.rs

@@ -31,6 +31,7 @@ pub mod reference_glue;
 pub mod scanning;
 pub mod utils;
 pub mod weak_proc;
+pub mod yjit_support;
 
 #[derive(Default)]
 pub struct Ruby;

mmtk/src/yjit_support.rs

@@ -0,0 +1,39 @@
+use mmtk::{
+    scheduler::{GCWork, WorkBucketStage},
+    util::VMWorkerThread,
+};
+
+use crate::{abi::GCThreadTLS, upcalls, Ruby};
+
+struct BeforeUpdatingJitCode;
+
+impl GCWork<Ruby> for BeforeUpdatingJitCode {
+    fn do_work(
+        &mut self,
+        _worker: &mut mmtk::scheduler::GCWorker<Ruby>,
+        _mmtk: &'static mmtk::MMTK<Ruby>,
+    ) {
+        (upcalls().before_updating_jit_code)();
+    }
+}
+
+struct AfterUpdatingJitCode;
+
+impl GCWork<Ruby> for AfterUpdatingJitCode {
+    fn do_work(
+        &mut self,
+        _worker: &mut mmtk::scheduler::GCWorker<Ruby>,
+        _mmtk: &'static mmtk::MMTK<Ruby>,
+    ) {
+        (upcalls().after_updating_jit_code)();
+    }
+}
+
+pub fn schedule_jit_code_protection_work_packets(tls: VMWorkerThread) {
+    let gc_tls: &'static mut GCThreadTLS = unsafe { GCThreadTLS::from_vwt_check(tls) };
+    let worker = gc_tls.worker();
+    if crate::mmtk().get_plan().current_gc_may_move_object() {
+        worker.scheduler().work_buckets[WorkBucketStage::Prepare].add(BeforeUpdatingJitCode);
+        worker.scheduler().work_buckets[WorkBucketStage::VMRefClosure].add(AfterUpdatingJitCode);
+    }
+}