8255662: ZGC: Unify nmethod closures in the heap iterator

Reviewed-by: eosterlund, pliden

src/hotspot/share/gc/z/zHeapIterator.cpp

@@ -250,9 +250,7 @@ public:
       _bs_nm(BarrierSet::barrier_set()->barrier_set_nmethod()) {}
 
   virtual void do_nmethod(nmethod* nm) {
-    assert(!ClassUnloading, "Only used if class unloading is turned off");
-
-    // ClassUnloading is turned off, all nmethods are considered strong,
+    // If ClassUnloading is turned off, all nmethods are considered strong,
     // not only those on the call stacks. The heap iteration might happen
     // before the concurrent processign of the code cache, make sure that
     // all nmethods have been processed before visiting the oops.
@@ -264,27 +262,16 @@ public:
 
 class ZHeapIteratorThreadClosure : public ThreadClosure {
 private:
-  OopClosure* const _cl;
-
-  class NMethodVisitor : public CodeBlobToOopClosure {
-  public:
-    NMethodVisitor(OopClosure* cl) :
-        CodeBlobToOopClosure(cl, false /* fix_oop_relocations */) {}
-
-    void do_code_blob(CodeBlob* cb) {
-      assert(!cb->is_nmethod() || !ZNMethod::is_armed(cb->as_nmethod()),
-          "NMethods on stack should have been fixed and disarmed");
-
-      CodeBlobToOopClosure::do_code_blob(cb);
-    }
-  };
+  OopClosure* const        _cl;
+  CodeBlobToNMethodClosure _cb_cl;
 
 public:
-  ZHeapIteratorThreadClosure(OopClosure* cl) : _cl(cl) {}
+  ZHeapIteratorThreadClosure(OopClosure* cl, NMethodClosure* nm_cl) :
+      _cl(cl),
+      _cb_cl(nm_cl) {}
 
   void do_thread(Thread* thread) {
-    NMethodVisitor code_cl(_cl);
-    thread->oops_do(_cl, &code_cl);
+    thread->oops_do(_cl, &_cb_cl);
   }
 };
 
@@ -292,7 +279,7 @@ void ZHeapIterator::push_strong_roots(const ZHeapIteratorContext& context) {
   ZHeapIteratorRootOopClosure<false /* Weak */> cl(context);
   ZHeapIteratorCLDCLosure cld_cl(&cl);
   ZHeapIteratorNMethodClosure nm_cl(&cl);
-  ZHeapIteratorThreadClosure thread_cl(&cl);
+  ZHeapIteratorThreadClosure thread_cl(&cl, &nm_cl);
 
   _concurrent_roots.apply(&cl,
                           &cld_cl,

src/hotspot/share/memory/iterator.cpp

@@ -64,3 +64,10 @@ void MarkingCodeBlobClosure::do_code_blob(CodeBlob* cb) {
     do_nmethod(nm);
   }
 }
+
+void CodeBlobToNMethodClosure::do_code_blob(CodeBlob* cb) {
+  nmethod* nm = cb->as_nmethod_or_null();
+  if (nm != NULL) {
+    _nm_cl->do_nmethod(nm);
+  }
+}

src/hotspot/share/memory/iterator.hpp

@@ -261,6 +261,15 @@ class NMethodClosure : public Closure {
   virtual void do_nmethod(nmethod* n) = 0;
 };
 
+class CodeBlobToNMethodClosure : public CodeBlobClosure {
+  NMethodClosure* const _nm_cl;
+
+ public:
+  CodeBlobToNMethodClosure(NMethodClosure* nm_cl) : _nm_cl(nm_cl) {}
+
+  virtual void do_code_blob(CodeBlob* cb);
+};
+
 // MonitorClosure is used for iterating over monitors in the monitors cache
 
 class ObjectMonitor;