WIP: Generational code cache roots

Author: wks

.gitignore

@@ -15,4 +15,6 @@ target/
 /*.dylib
 /*.so
 
-/repos/mmtk-core
+/repos/mmtk-core
+
+__pycache__

mmtk/Cargo.lock

@@ -2,7 +2,7 @@
 name = "mmtk_openjdk"
 version = "0.27.0"
 authors = [" <>"]
-rust-version = "1.71.1"
+rust-version = "1.73.0"
 build = "build.rs"
 edition = "2021"
 
@@ -27,14 +27,15 @@ once_cell = "1.10.0"
 atomic = "0.6.0"
 memoffset = "0.9.0"
 cfg-if = "1.0"
+probe = "0.5"
 
 # Be very careful to commit any changes to the following mmtk dependency, as our CI scripts (including mmtk-core CI)
 # rely on matching these lines to modify them: e.g. comment out the git dependency and use the local path.
 # These changes are safe:
 # - change branch
 # - change repo name
 # But other changes including adding/removing whitespaces in commented lines may break the CI.
-mmtk = { git = "https://github.com/mmtk/mmtk-core.git", rev = "45cdf31055b1b6a629bdb8032adaa6dd5a8e32b9" }
+mmtk = { git = "https://github.com/mmtk/mmtk-core.git", rev = "a025c24104d8d456a865aa0122e6e0fb6d77e8f2" }
 # Uncomment the following to build locally
 # mmtk = { path = "../repos/mmtk-core" }
 

mmtk/Cargo.toml

@@ -497,42 +497,31 @@ thread_local! {
 /// Report a list of pointers in nmethod to mmtk.
 #[no_mangle]
 pub extern "C" fn mmtk_add_nmethod_oop(addr: Address) {
-    NMETHOD_SLOTS.with(|x| x.borrow_mut().push(addr))
+    NMETHOD_SLOTS.with_borrow_mut(|x| x.push(addr))
 }
 
 /// Register a nmethod.
 /// The c++ part of the binding should scan the nmethod and report all the pointers to mmtk first, before calling this function.
 /// This function will transfer all the locally cached pointers of this nmethod to the global storage.
 #[no_mangle]
 pub extern "C" fn mmtk_register_nmethod(nm: Address) {
-    let slots = NMETHOD_SLOTS.with(|x| {
-        if x.borrow().len() == 0 {
-            return None;
+    NMETHOD_SLOTS.with_borrow_mut(|slots| {
+        if !slots.is_empty() {
+            let mut roots = crate::NURSERY_CODE_CACHE_ROOTS.lock().unwrap();
+            roots.insert(nm, std::mem::take(slots));
         }
-        Some(x.replace(vec![]))
     });
-    let slots = match slots {
-        Some(slots) => slots,
-        _ => return,
-    };
-    let mut roots = crate::CODE_CACHE_ROOTS.lock().unwrap();
-    // Relaxed add instead of `fetch_add`, since we've already acquired the lock.
-    crate::CODE_CACHE_ROOTS_SIZE.store(
-        crate::CODE_CACHE_ROOTS_SIZE.load(Ordering::Relaxed) + slots.len(),
-        Ordering::Relaxed,
-    );
-    roots.insert(nm, slots);
 }
 
 /// Unregister a nmethod.
 #[no_mangle]
 pub extern "C" fn mmtk_unregister_nmethod(nm: Address) {
-    let mut roots = crate::CODE_CACHE_ROOTS.lock().unwrap();
-    if let Some(slots) = roots.remove(&nm) {
-        // Relaxed sub instead of `fetch_sub`, since we've already acquired the lock.
-        crate::CODE_CACHE_ROOTS_SIZE.store(
-            crate::CODE_CACHE_ROOTS_SIZE.load(Ordering::Relaxed) - slots.len(),
-            Ordering::Relaxed,
-        );
+    {
+        let mut roots = crate::NURSERY_CODE_CACHE_ROOTS.lock().unwrap();
+        roots.remove(&nm);
+    }
+    {
+        let mut roots = crate::MATURE_CODE_CACHE_ROOTS.lock().unwrap();
+        roots.remove(&nm);
     }
 }

mmtk/src/api.rs

@@ -1,10 +1,10 @@
-use std::sync::atomic::Ordering;
-
+use crate::scanning;
 use crate::scanning::to_slots_closure;
 use crate::OpenJDK;
 use crate::OpenJDKSlot;
 use crate::UPCALLS;
 use mmtk::scheduler::*;
+use mmtk::util::Address;
 use mmtk::vm::RootsWorkFactory;
 use mmtk::vm::*;
 use mmtk::MMTK;
@@ -69,16 +69,51 @@ impl<const COMPRESSED: bool, F: RootsWorkFactory<OpenJDKSlot<COMPRESSED>>>
     fn do_work(
         &mut self,
         _worker: &mut GCWorker<OpenJDK<COMPRESSED>>,
-        _mmtk: &'static MMTK<OpenJDK<COMPRESSED>>,
+        mmtk: &'static MMTK<OpenJDK<COMPRESSED>>,
     ) {
-        // Collect all the cached roots
-        let mut slots = Vec::with_capacity(crate::CODE_CACHE_ROOTS_SIZE.load(Ordering::Relaxed));
-        for roots in (*crate::CODE_CACHE_ROOTS.lock().unwrap()).values() {
-            for r in roots {
-                slots.push((*r).into())
+        let is_current_gc_nursery = mmtk
+            .get_plan()
+            .generational()
+            .is_some_and(|gen| gen.is_current_gc_nursery());
+
+            let mut slots = Vec::with_capacity(scanning::WORK_PACKET_CAPACITY);
+
+        let mut nursery_slots = 0;
+        let mut mature_slots = 0;
+
+        let mut add_roots = |roots: &[Address]| {
+            for root in roots {
+                slots.push(OpenJDKSlot::<COMPRESSED>::from(*root));
+                if slots.len() >= scanning::WORK_PACKET_CAPACITY {
+                    self.factory
+                        .create_process_roots_work(std::mem::take(&mut slots));
+                }
+            }
+        };
+
+        {
+            let mut mature = crate::MATURE_CODE_CACHE_ROOTS.lock().unwrap();
+
+            // Only scan mature roots in full-heap collections.
+            if !is_current_gc_nursery {
+                for roots in mature.values() {
+                    mature_slots += roots.len();
+                    add_roots(roots);
+                }
+            }
+
+            {
+                let mut nursery = crate::NURSERY_CODE_CACHE_ROOTS.lock().unwrap();
+                for (key, roots) in nursery.drain() {
+                    nursery_slots += roots.len();
+                    add_roots(&roots);
+                    mature.insert(key, roots);
+                }
             }
         }
-        // Create work packet
+
+        probe!(mmtk_openjdk, code_cache_roots, nursery_slots, mature_slots);
+
         if !slots.is_empty() {
             self.factory.create_process_roots_work(slots);
         }

mmtk/src/gc_work.rs

@@ -1,9 +1,10 @@
 #[macro_use]
 extern crate lazy_static;
+#[macro_use]
+extern crate probe;
 
 use std::collections::HashMap;
 use std::ptr::null_mut;
-use std::sync::atomic::AtomicUsize;
 use std::sync::Mutex;
 
 use libc::{c_char, c_void, uintptr_t};
@@ -202,13 +203,12 @@ pub static MMTK_MARK_COMPACT_HEADER_RESERVED_IN_BYTES: usize =
     mmtk::util::alloc::MarkCompactAllocator::<OpenJDK<false>>::HEADER_RESERVED_IN_BYTES;
 
 lazy_static! {
-    /// A global storage for all the cached CodeCache root pointers
-    static ref CODE_CACHE_ROOTS: Mutex<HashMap<Address, Vec<Address>>> = Mutex::new(HashMap::new());
+    /// A global storage for all the cached CodeCache root pointers added since the last GC.
+    static ref NURSERY_CODE_CACHE_ROOTS: Mutex<HashMap<Address, Vec<Address>>> = Mutex::new(HashMap::new());
+    /// A global storage for all the cached CodeCache root pointers added before the last GC.
+    static ref MATURE_CODE_CACHE_ROOTS: Mutex<HashMap<Address, Vec<Address>>> = Mutex::new(HashMap::new());
 }
 
-/// A counter tracking the total size of the `CODE_CACHE_ROOTS`.
-static CODE_CACHE_ROOTS_SIZE: AtomicUsize = AtomicUsize::new(0);
-
 fn set_compressed_pointer_vm_layout(builder: &mut MMTKBuilder) {
     let max_heap_size = builder.options.gc_trigger.max_heap_size();
     assert!(

mmtk/src/lib.rs

@@ -12,7 +12,7 @@ use mmtk::MutatorContext;
 
 pub struct VMScanning {}
 
-const WORK_PACKET_CAPACITY: usize = 4096;
+pub(crate) const WORK_PACKET_CAPACITY: usize = 4096;
 
 extern "C" fn report_slots_and_renew_buffer<S: Slot, F: RootsWorkFactory<S>>(
     ptr: *mut Address,

mmtk/src/scanning.rs

@@ -0,0 +1,5 @@
+usdt:$MMTK:mmtk_openjdk:code_cache_roots {
+    if (@enable_print) {
+        printf("code_cache_roots,meta,%d,%lu,%lu,%lu\n", tid, nsecs, arg0, arg1);
+    }
+}

tools/tracing/timeline/capture_openjdk.bt

@@ -0,0 +1,13 @@
+#!/usr/bin/env python3
+
+def enrich_meta_extra(log_processor, name, tid, ts, gc, wp, args):
+    if wp is not None:
+        match name:
+            case "code_cache_roots":
+                nursery, mature = int(args[0]), int(args[1])
+                total = nursery + mature
+                wp["args"] |= {
+                    "nursery_slots": nursery,
+                    "mature_slots": mature,
+                    "total_slots": total,
+                }