Generational code cache roots (#282)

This PR ports the generational code cache roots handling code from the
LXR branch. See
https://github.com/wenyuzhao/mmtk-openjdk/blob/lxr/mmtk/src/gc_work.rs#L242-L262

We split code cache roots into two hash tables, one for nursery roots
and one for mature roots. Nursery roots are those added since the
previous GC. In a nursery GC, we only scan the nursery code cache roots,
assuming mature code cache roots all point to mature objects. In a
full-heap GC, we scan both nursery and mature code cache roots.

Because code cache roots are added in bulk in the beginning of
execution, and are seldom added later, we expect the number of code
cache roots to reduce to zero for most nursery GCs, and greatly reduce
the root-scanning time.

This PR also adds USDT tracepoints for code cache roots, and make use of
the extensible eBPF timeline tools in mmtk-core introduced in
mmtk/mmtk-core#1162.

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

@@ -490,49 +490,41 @@ pub extern "C" fn get_finalized_object() -> NullableObjectReference {
 }
 
 thread_local! {
-    /// Cache all the pointers reported by the current thread.
+    /// Cache reference slots of an nmethod while the current thread is executing
+    /// `MMTkRegisterNMethodOopClosure`.
     static NMETHOD_SLOTS: RefCell<Vec<Address>> = const { RefCell::new(vec![]) };
 }
 
-/// Report a list of pointers in nmethod to mmtk.
+/// Report one reference slot in an 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.
+/// Register an nmethod.
+///
+/// The C++ part of the binding should have scanned the nmethod and reported all the reference slots
+/// using `mmtk_add_nmethod_oop` before calling this function. This function will transfer all the
+/// locally cached slots 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.
+/// Unregister an 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 roots 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 roots 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,
+                }