Scan registers in task and thread

julia/mmtk_julia.c

@@ -210,6 +210,8 @@ JL_DLLEXPORT void jl_gc_prepare_to_collect(void)
     gc_num.total_time_to_safepoint += duration;
 
     if (!jl_atomic_load_acquire(&jl_gc_disable_counter)) {
+        jl_save_context_for_conservative_scanning(ptls, NULL);
+
         JL_LOCK_NOGC(&finalizers_lock); // all the other threads are stopped, so this does not make sense, right? otherwise, failing that, this seems like plausibly a deadlock
         combine_thread_gc_counts(&gc_num);
         mmtk_block_thread_for_gc(gc_n_threads);

mmtk/src/julia_scanning.rs

@@ -487,7 +487,7 @@ pub fn unpin_conservative_roots() {
     );
 }
 
-pub unsafe fn mmtk_conservative_scan_native_stack(ta: *const mmtk_jl_task_t) {
+pub fn mmtk_conservative_scan_native_stack(ta: *const mmtk_jl_task_t) {
     let mut size: u64 = 0;
     let mut ptid: i32 = 0;
     log::info!("mmtk_conservative_scan_native_stack begin ta = {:?}", ta);
@@ -515,6 +515,24 @@ pub unsafe fn mmtk_conservative_scan_native_stack(ta: *const mmtk_jl_task_t) {
     }
 }
 
+pub fn mmtk_conservative_scan_task_registers(ta: *const mmtk_jl_task_t) {
+    let (lo, hi) = get_range(&unsafe { &*ta }.ctx);
+    conservative_scan_range(lo, hi);
+}
+
+pub fn mmtk_conservative_scan_ptls_registers(ptls: &mut mmtk_jl_tls_states_t) {
+    let (lo, hi) = get_range(&ptls.ctx_at_the_time_gc_started);
+    conservative_scan_range(lo, hi);
+}
+
+// TODO: This scans the entire context type, which is slower.
+// We actually only need to scan registers.
+pub fn get_range<T>(ctx: &T) -> (Address, Address) {
+    let start = Address::from_ptr(ctx);
+    let ty_size = std::mem::size_of::<T>();
+    (start, start + ty_size)
+}
+
 pub fn conservative_scan_range(lo: Address, hi: Address) {
     // The high address is exclusive
     let hi = if hi.is_aligned_to(BYTES_IN_ADDRESS) {

mmtk/src/scanning.rs

@@ -55,6 +55,9 @@ impl Scanning<JuliaVM> for VMScanning {
         let mut pinning_slot_buffer = SlotBuffer { buffer: vec![] }; // roots from the shadow stack that we know that do not need to be transitively pinned
         let mut node_buffer = vec![];
 
+        // Conservatively scan registers saved with the thread
+        mmtk_conservative_scan_ptls_registers(ptls);
+
         // Scan thread local from ptls: See gc_queue_thread_local in gc.c
         let mut root_scan_task = |task: *const mmtk__jl_task_t, task_is_root: bool| {
             if !task.is_null() {
@@ -76,9 +79,10 @@ impl Scanning<JuliaVM> for VMScanning {
                     mutator.mutator_tls,
                     pthread
                 );
-                unsafe {
-                    mmtk_conservative_scan_native_stack(task);
-                }
+                // Conservative scan stack and registers
+                mmtk_conservative_scan_native_stack(task);
+                mmtk_conservative_scan_task_registers(task);
+
                 if task_is_root {
                     // captures wrong root nodes before creating the work
                     debug_assert!(