Fix memory leak due to incorrect root scanning

mmtk/src/api.rs

@@ -274,3 +274,22 @@ pub extern "C" fn get_finalized_object() -> ObjectReference {
         None => unsafe { Address::ZERO.to_object_reference() },
     }
 }
+
+/// Test if an object is live at the end of a GC.
+/// Note: only call this method after the liveness tracing and before gc release.
+#[no_mangle]
+pub extern "C" fn mmtk_is_live(object: ObjectReference) -> usize {
+    if object.is_null() {
+        return 0;
+    }
+    object.is_live() as _
+}
+
+/// If the object is non-null and forwarded, return the forwarded pointer. Otherwise, return the original pointer.
+#[no_mangle]
+pub extern "C" fn mmtk_get_forwarded_ref(object: ObjectReference) -> ObjectReference {
+    if object.is_null() {
+        return object;
+    }
+    object.get_forwarded_object().unwrap_or(object)
+}

mmtk/src/collection.rs

@@ -1,4 +1,4 @@
-use mmtk::scheduler::WorkBucketStage;
+use mmtk::scheduler::{WorkBucketStage, GCWorker};
 use mmtk::scheduler::{ProcessEdgesWork, ScanStackRoot};
 use mmtk::util::alloc::AllocationError;
 use mmtk::util::opaque_pointer::*;
@@ -87,4 +87,10 @@ impl Collection<OpenJDK> for VMCollection {
             ((*UPCALLS).schedule_finalizer)();
         }
     }
+
+    fn process_weak_refs<E: ProcessEdgesWork<VM = OpenJDK>>(_worker: &mut GCWorker<OpenJDK>) {
+        unsafe {
+            ((*UPCALLS).process_weak_refs)();
+        }
+    }
 }

mmtk/src/gc_work.rs

@@ -116,38 +116,6 @@ impl<E: ProcessEdgesWork<VM = OpenJDK>> GCWork<OpenJDK> for ScanSystemDictionary
     }
 }
 
-pub struct ScanCodeCacheRoots<E: ProcessEdgesWork<VM = OpenJDK>>(PhantomData<E>);
-
-impl<E: ProcessEdgesWork<VM = OpenJDK>> ScanCodeCacheRoots<E> {
-    pub fn new() -> Self {
-        Self(PhantomData)
-    }
-}
-
-impl<E: ProcessEdgesWork<VM = OpenJDK>> GCWork<OpenJDK> for ScanCodeCacheRoots<E> {
-    fn do_work(&mut self, _worker: &mut GCWorker<OpenJDK>, _mmtk: &'static MMTK<OpenJDK>) {
-        unsafe {
-            ((*UPCALLS).scan_code_cache_roots)(create_process_edges_work::<E> as _);
-        }
-    }
-}
-
-pub struct ScanStringTableRoots<E: ProcessEdgesWork<VM = OpenJDK>>(PhantomData<E>);
-
-impl<E: ProcessEdgesWork<VM = OpenJDK>> ScanStringTableRoots<E> {
-    pub fn new() -> Self {
-        Self(PhantomData)
-    }
-}
-
-impl<E: ProcessEdgesWork<VM = OpenJDK>> GCWork<OpenJDK> for ScanStringTableRoots<E> {
-    fn do_work(&mut self, _worker: &mut GCWorker<OpenJDK>, _mmtk: &'static MMTK<OpenJDK>) {
-        unsafe {
-            ((*UPCALLS).scan_string_table_roots)(create_process_edges_work::<E> as _);
-        }
-    }
-}
-
 pub struct ScanClassLoaderDataGraphRoots<E: ProcessEdgesWork<VM = OpenJDK>>(PhantomData<E>);
 
 impl<E: ProcessEdgesWork<VM = OpenJDK>> ScanClassLoaderDataGraphRoots<E> {
@@ -164,22 +132,6 @@ impl<E: ProcessEdgesWork<VM = OpenJDK>> GCWork<OpenJDK> for ScanClassLoaderDataG
     }
 }
 
-pub struct ScanWeakProcessorRoots<E: ProcessEdgesWork<VM = OpenJDK>>(PhantomData<E>);
-
-impl<E: ProcessEdgesWork<VM = OpenJDK>> ScanWeakProcessorRoots<E> {
-    pub fn new() -> Self {
-        Self(PhantomData)
-    }
-}
-
-impl<E: ProcessEdgesWork<VM = OpenJDK>> GCWork<OpenJDK> for ScanWeakProcessorRoots<E> {
-    fn do_work(&mut self, _worker: &mut GCWorker<OpenJDK>, _mmtk: &'static MMTK<OpenJDK>) {
-        unsafe {
-            ((*UPCALLS).scan_weak_processor_roots)(create_process_edges_work::<E> as _);
-        }
-    }
-}
-
 pub struct ScanVMThreadRoots<E: ProcessEdgesWork<VM = OpenJDK>>(PhantomData<E>);
 
 impl<E: ProcessEdgesWork<VM = OpenJDK>> ScanVMThreadRoots<E> {

mmtk/src/lib.rs

@@ -75,15 +75,13 @@ pub struct OpenJDK_Upcalls {
     pub scan_jvmti_export_roots: extern "C" fn(process_edges: ProcessEdgesFn),
     pub scan_aot_loader_roots: extern "C" fn(process_edges: ProcessEdgesFn),
     pub scan_system_dictionary_roots: extern "C" fn(process_edges: ProcessEdgesFn),
-    pub scan_code_cache_roots: extern "C" fn(process_edges: ProcessEdgesFn),
-    pub scan_string_table_roots: extern "C" fn(process_edges: ProcessEdgesFn),
     pub scan_class_loader_data_graph_roots: extern "C" fn(process_edges: ProcessEdgesFn),
-    pub scan_weak_processor_roots: extern "C" fn(process_edges: ProcessEdgesFn),
     pub scan_vm_thread_roots: extern "C" fn(process_edges: ProcessEdgesFn),
     pub number_of_mutators: extern "C" fn() -> usize,
     pub schedule_finalizer: extern "C" fn(),
     pub prepare_for_roots_re_scanning: extern "C" fn(),
     pub object_alignment: extern "C" fn() -> i32,
+    pub process_weak_refs: extern "C" fn(),
 }
 
 pub static mut UPCALLS: *const OpenJDK_Upcalls = null_mut();

mmtk/src/scanning.rs

@@ -83,10 +83,7 @@ impl Scanning<OpenJDK> for VMScanning {
                 box ScanJvmtiExportRoots::<W>::new(),
                 box ScanAOTLoaderRoots::<W>::new(),
                 box ScanSystemDictionaryRoots::<W>::new(),
-                box ScanCodeCacheRoots::<W>::new(),
-                box ScanStringTableRoots::<W>::new(),
                 box ScanClassLoaderDataGraphRoots::<W>::new(),
-                box ScanWeakProcessorRoots::<W>::new(),
             ],
         );
         if !(Self::SCAN_MUTATORS_IN_SAFEPOINT && Self::SINGLE_THREAD_MUTATOR_SCANNING) {

openjdk/mmtk.h

@@ -87,6 +87,9 @@ extern void handle_user_collection_request(void *tls);
 extern void start_control_collector(void *tls, void *context);
 extern void start_worker(void *tls, void* worker);
 
+extern size_t mmtk_is_live(void* object);
+extern void* mmtk_get_forwarded_ref(void* object);
+
 /**
  * VM Accounting
  */
@@ -136,14 +139,13 @@ typedef struct {
     void (*scan_jvmti_export_roots) (ProcessEdgesFn process_edges);
     void (*scan_aot_loader_roots) (ProcessEdgesFn process_edges);
     void (*scan_system_dictionary_roots) (ProcessEdgesFn process_edges);
-    void (*scan_code_cache_roots) (ProcessEdgesFn process_edges);
-    void (*scan_string_table_roots) (ProcessEdgesFn process_edges);
     void (*scan_class_loader_data_graph_roots) (ProcessEdgesFn process_edges);
-    void (*scan_weak_processor_roots) (ProcessEdgesFn process_edges);
     void (*scan_vm_thread_roots) (ProcessEdgesFn process_edges);
     size_t (*number_of_mutators)();
     void (*schedule_finalizer)();
     void (*prepare_for_roots_re_scanning)();
+    int32_t (*object_alignment)();
+    void (*process_weak_ref)();
 } OpenJDK_Upcalls;
 
 extern void openjdk_gc_init(OpenJDK_Upcalls *calls, size_t heap_size);

openjdk/mmtkHeap.cpp

@@ -343,98 +343,50 @@ bool MMTkHeap::is_scavengable(oop obj) {return false;}
 // Heap verification
 void MMTkHeap::verify(VerifyOption option) {}
 
-template<int MAX_TASKS = 12>
-struct MMTkRootScanWorkScope {
-  int* _num_root_scan_tasks;
-  int _current_task_ordinal;
-  MMTkRootScanWorkScope(int* num_root_scan_tasks): _num_root_scan_tasks(num_root_scan_tasks), _current_task_ordinal(0) {
-    _current_task_ordinal = Atomic::add(1, _num_root_scan_tasks);
-    if (_current_task_ordinal == 1) {
-      nmethod::oops_do_marking_prologue();
-    }
-  }
-  ~MMTkRootScanWorkScope() {
-    if (_current_task_ordinal == MAX_TASKS) {
-      _current_task_ordinal = 0;
-      nmethod::oops_do_marking_epilogue();
-    }
-  }
-};
-
 void MMTkHeap::scan_static_roots(OopClosure& cl) {
 }
 
 
 void MMTkHeap::scan_universe_roots(OopClosure& cl) {
   ResourceMark rm;
-  MMTkRootScanWorkScope<> root_scan_work(&_num_root_scan_tasks);
   Universe::oops_do(&cl);
 }
 void MMTkHeap::scan_jni_handle_roots(OopClosure& cl) {
   ResourceMark rm;
-  MMTkRootScanWorkScope<> root_scan_work(&_num_root_scan_tasks);
   JNIHandles::oops_do(&cl);
 }
 void MMTkHeap::scan_object_synchronizer_roots(OopClosure& cl) {
   ResourceMark rm;
-  MMTkRootScanWorkScope<> root_scan_work(&_num_root_scan_tasks);
   ObjectSynchronizer::oops_do(&cl);
 }
 void MMTkHeap::scan_management_roots(OopClosure& cl) {
   ResourceMark rm;
-  MMTkRootScanWorkScope<> root_scan_work(&_num_root_scan_tasks);
   Management::oops_do(&cl);
 }
 void MMTkHeap::scan_jvmti_export_roots(OopClosure& cl) {
   ResourceMark rm;
-  MMTkRootScanWorkScope<> root_scan_work(&_num_root_scan_tasks);
   JvmtiExport::oops_do(&cl);
 }
 void MMTkHeap::scan_aot_loader_roots(OopClosure& cl) {
   ResourceMark rm;
-  MMTkRootScanWorkScope<> root_scan_work(&_num_root_scan_tasks);
   AOTLoader::oops_do(&cl);
 }
 void MMTkHeap::scan_system_dictionary_roots(OopClosure& cl) {
   ResourceMark rm;
-  MMTkRootScanWorkScope<> root_scan_work(&_num_root_scan_tasks);
   SystemDictionary::oops_do(&cl);
 }
-void MMTkHeap::scan_code_cache_roots(OopClosure& cl) {
-  ResourceMark rm;
-  MMTkRootScanWorkScope<> root_scan_work(&_num_root_scan_tasks);
-  CodeBlobToOopClosure cb_cl(&cl, true);
-  {
-    MutexLockerEx lock(CodeCache_lock, Mutex::_no_safepoint_check_flag);
-    CodeCache::scavenge_root_nmethods_do(&cb_cl);
-    CodeCache::blobs_do(&cb_cl);
-  }
-}
-void MMTkHeap::scan_string_table_roots(OopClosure& cl) {
-  ResourceMark rm;
-  MMTkRootScanWorkScope<> root_scan_work(&_num_root_scan_tasks);
-  StringTable::oops_do(&cl);
-}
 void MMTkHeap::scan_class_loader_data_graph_roots(OopClosure& cl) {
   ResourceMark rm;
-  MMTkRootScanWorkScope<> root_scan_work(&_num_root_scan_tasks);
-  CLDToOopClosure cld_cl(&cl, false);
-  ClassLoaderDataGraph::cld_do(&cld_cl);
-}
-void MMTkHeap::scan_weak_processor_roots(OopClosure& cl) {
-  ResourceMark rm;
-  MMTkRootScanWorkScope<> root_scan_work(&_num_root_scan_tasks);
-  WeakProcessor::oops_do(&cl); // (really needed???)
+  CLDToOopClosure cld_cl(&cl);
+  ClassLoaderDataGraph::roots_cld_do(&cld_cl, &cld_cl);
 }
 void MMTkHeap::scan_vm_thread_roots(OopClosure& cl) {
   ResourceMark rm;
-  MMTkRootScanWorkScope<> root_scan_work(&_num_root_scan_tasks);
   VMThread::vm_thread()->oops_do(&cl, NULL);
 }
 
 void MMTkHeap::scan_global_roots(OopClosure& cl) {
   ResourceMark rm;
-  MMTkRootScanWorkScope<> root_scan_work(&_num_root_scan_tasks);
 
   CodeBlobToOopClosure cb_cl(&cl, true);
   CLDToOopClosure cld_cl(&cl, false);
@@ -462,7 +414,6 @@ void MMTkHeap::scan_global_roots(OopClosure& cl) {
 
 void MMTkHeap::scan_thread_roots(OopClosure& cl) {
   ResourceMark rm;
-  MMTkRootScanWorkScope<> root_scan_work(&_num_root_scan_tasks);
   Threads::possibly_parallel_oops_do(false, &cl, NULL);
 }
 
@@ -516,6 +467,15 @@ void MMTkHeap::report_java_thread_yield(JavaThread* thread) {
   if (_create_stack_scan_work != NULL) _create_stack_scan_work((void*) &thread->third_party_heap_mutator);
 }
 
+void MMTkHeap::register_nmethod(nmethod* nm) {
+  CodeCache::register_scavenge_root_nmethod(nm);
+}
+
+void MMTkHeap::verify_nmethod(nmethod* nm) {
+  CodeCache::verify_scavenge_root_nmethod(nm);
+}
+
+
 /*
  * files with prints currently:
  * collectedHeap.inline.hpp, mmtkHeap.cpp,

openjdk/mmtkHeap.hpp

@@ -156,6 +156,8 @@ class MMTkHeap : public CollectedHeap {
 
   void prepare_for_verify() ;
 
+  virtual void register_nmethod(nmethod* nm);
+  virtual void verify_nmethod(nmethod* nm);
 
 private:
 

openjdk/mmtkUpcalls.cpp

@@ -40,6 +40,7 @@
 #include "runtime/threadSMR.hpp"
 #include "runtime/vmThread.hpp"
 #include "utilities/debug.hpp"
+#include "gc/shared/weakProcessor.hpp"
 
 static size_t mmtk_start_the_world_count = 0;
 
@@ -63,10 +64,13 @@ static void mmtk_stop_all_mutators(void *tls, void (*create_stack_scan_work)(voi
     }
   }
   log_debug(gc)("Finished enumerating threads.");
+  BiasedLocking::preserve_marks();
+  nmethod::oops_do_marking_prologue();
 }
 
 static void mmtk_resume_mutators(void *tls) {
   ClassLoaderDataGraph::purge();
+  BiasedLocking::restore_marks();
   CodeCache::gc_epilogue();
   JvmtiExport::gc_epilogue();
 #if COMPILER2_OR_JVMCI
@@ -302,10 +306,7 @@ static void mmtk_scan_management_roots(ProcessEdgesFn process_edges) { MMTkRoots
 static void mmtk_scan_jvmti_export_roots(ProcessEdgesFn process_edges) { MMTkRootsClosure2 cl(process_edges); MMTkHeap::heap()->scan_jvmti_export_roots(cl); }
 static void mmtk_scan_aot_loader_roots(ProcessEdgesFn process_edges) { MMTkRootsClosure2 cl(process_edges); MMTkHeap::heap()->scan_aot_loader_roots(cl); }
 static void mmtk_scan_system_dictionary_roots(ProcessEdgesFn process_edges) { MMTkRootsClosure2 cl(process_edges); MMTkHeap::heap()->scan_system_dictionary_roots(cl); }
-static void mmtk_scan_code_cache_roots(ProcessEdgesFn process_edges) { MMTkRootsClosure2 cl(process_edges); MMTkHeap::heap()->scan_code_cache_roots(cl); }
-static void mmtk_scan_string_table_roots(ProcessEdgesFn process_edges) { MMTkRootsClosure2 cl(process_edges); MMTkHeap::heap()->scan_string_table_roots(cl); }
 static void mmtk_scan_class_loader_data_graph_roots(ProcessEdgesFn process_edges) { MMTkRootsClosure2 cl(process_edges); MMTkHeap::heap()->scan_class_loader_data_graph_roots(cl); }
-static void mmtk_scan_weak_processor_roots(ProcessEdgesFn process_edges) { MMTkRootsClosure2 cl(process_edges); MMTkHeap::heap()->scan_weak_processor_roots(cl); }
 static void mmtk_scan_vm_thread_roots(ProcessEdgesFn process_edges) { MMTkRootsClosure2 cl(process_edges); MMTkHeap::heap()->scan_vm_thread_roots(cl); }
 
 static size_t mmtk_number_of_mutators() {
@@ -319,6 +320,59 @@ static void mmtk_prepare_for_roots_re_scanning() {
 #endif
 }
 
+static int32_t mmtk_object_alignment() {
+  ShouldNotReachHere();
+  return 0;
+}
+
+class MMTkIsAliveClosure : public BoolObjectClosure {
+public:
+  virtual bool do_object_b(oop p) {
+    return mmtk_is_live((void*) p) != 0;
+  }
+};
+
+class MMTkForwardClosure : public OopClosure {
+ public:
+  virtual void do_oop(oop* o) {
+    *o = (oop) mmtk_get_forwarded_ref((void*) *o);
+  }
+  virtual void do_oop(narrowOop* o) {}
+};
+
+/// Clean up the weak-ref storage and update pointers.
+static void mmtk_process_weak_ref() {
+
+  HandleMark hm;
+
+  MMTkIsAliveClosure is_alive;
+  MMTkForwardClosure forward;
+
+  nmethod::oops_do_marking_epilogue();
+
+  WeakProcessor::weak_oops_do(&is_alive, &forward);
+
+  // TODO: uncomment below to support class unloading
+  // bool purged_class = SystemDictionary::do_unloading(NULL);
+  bool purged_class = false;
+
+  CodeBlobToOopClosure adjust_from_blobs(&forward, CodeBlobToOopClosure::FixRelocations);
+  CodeCache::blobs_do(&adjust_from_blobs);
+  CodeCache::do_unloading(&is_alive, purged_class);
+  Klass::clean_weak_klass_links(purged_class);
+
+  StringTable::oops_do(&forward);
+  StringTable::unlink(&is_alive);
+
+  SymbolTable::unlink();
+
+  // TODO: Enable below to support class unloading
+  // ClassLoaderDataGraph::clear_claimed_marks();
+  // CLDToOopClosure adjust_cld_closure(&forward);
+  // ClassLoaderDataGraph::cld_do(&adjust_cld_closure);
+  // ClassLoaderDataGraph::oops_do(&forward, true);
+}
+
 OpenJDK_Upcalls mmtk_upcalls = {
   mmtk_stop_all_mutators,
   mmtk_resume_mutators,
@@ -352,12 +406,11 @@ OpenJDK_Upcalls mmtk_upcalls = {
   mmtk_scan_jvmti_export_roots,
   mmtk_scan_aot_loader_roots,
   mmtk_scan_system_dictionary_roots,
-  mmtk_scan_code_cache_roots,
-  mmtk_scan_string_table_roots,
   mmtk_scan_class_loader_data_graph_roots,
-  mmtk_scan_weak_processor_roots,
   mmtk_scan_vm_thread_roots,
   mmtk_number_of_mutators,
   mmtk_schedule_finalizer,
   mmtk_prepare_for_roots_re_scanning,
+  mmtk_object_alignment,
+  mmtk_process_weak_ref,
 };