8254980: ZGC: ZHeapIterator visits armed nmethods with -XX:-ClassUnloading

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

@@ -126,6 +126,17 @@ class ZHeapIteratorRootOopClosure : public ZRootsIteratorClosure {
     CodeBlobToOopClosure code_cl(this, false /* fix_oop_relocations */);
     thread->oops_do(this, &code_cl);
   }
+
+  virtual ZNMethodEntry nmethod_entry() const {
+    if (ClassUnloading) {
+      // All encountered nmethods should have been "entered" during stack walking
+      return ZNMethodEntry::VerifyDisarmed;
+    } else {
+      // All nmethods are considered roots and will be visited.
+      // Make sure that the unvisited gets fixed and disarmed before proceeding.
+      return ZNMethodEntry::PreBarrier;
+    }
+  }
 };
 
 template <bool VisitReferents>

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

@@ -582,8 +582,9 @@ class ZMarkConcurrentRootsIteratorClosure : public ZRootsIteratorClosure {
     ZThreadLocalAllocBuffer::publish_statistics();
   }
 
-  virtual bool should_disarm_nmethods() const {
-    return true;
+  virtual ZNMethodEntry nmethod_entry() const {
+    // Only apply closure to armed nmethods, and then disarm them.
+    return ZNMethodEntry::Disarm;
   }
 
   virtual void do_thread(Thread* thread) {

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

@@ -236,28 +236,42 @@ void ZNMethod::nmethod_oops_do(nmethod* nm, OopClosure* cl) {
 
 class ZNMethodToOopsDoClosure : public NMethodClosure {
 private:
-  OopClosure* const _cl;
-  const bool        _should_disarm_nmethods;
+  OopClosure* const        _cl;
+  const ZNMethodEntry      _entry;
+  BarrierSetNMethod* const _bs_nm;
 
 public:
-  ZNMethodToOopsDoClosure(OopClosure* cl, bool should_disarm_nmethods) :
+  ZNMethodToOopsDoClosure(OopClosure* cl, ZNMethodEntry entry) :
       _cl(cl),
-      _should_disarm_nmethods(should_disarm_nmethods) {}
+      _entry(entry),
+      _bs_nm(BarrierSet::barrier_set()->barrier_set_nmethod()) {}
 
   virtual void do_nmethod(nmethod* nm) {
+    if (_entry == ZNMethodEntry::PreBarrier) {
+      // Apply entry barrier before proceeding with closure
+      _bs_nm->nmethod_entry_barrier(nm);
+    }
+
     ZLocker<ZReentrantLock> locker(ZNMethod::lock_for_nmethod(nm));
     if (!nm->is_alive()) {
       return;
     }
 
-    if (_should_disarm_nmethods) {
+    if (_entry == ZNMethodEntry::Disarm) {
+      // Apply closure and disarm only armed nmethods
       if (ZNMethod::is_armed(nm)) {
         ZNMethod::nmethod_oops_do(nm, _cl);
         ZNMethod::disarm(nm);
       }
-    } else {
-      ZNMethod::nmethod_oops_do(nm, _cl);
+      return;
     }
+
+    if (_entry == ZNMethodEntry::VerifyDisarmed) {
+      // Only verify
+      assert(!ZNMethod::is_armed(nm), "Must be disarmed");
+    }
+
+    ZNMethod::nmethod_oops_do(nm, _cl);
   }
 };
 
@@ -269,8 +283,8 @@ void ZNMethod::oops_do_end() {
   ZNMethodTable::nmethods_do_end();
 }
 
-void ZNMethod::oops_do(OopClosure* cl, bool should_disarm_nmethods) {
-  ZNMethodToOopsDoClosure nmethod_cl(cl, should_disarm_nmethods);
+void ZNMethod::oops_do(OopClosure* cl, ZNMethodEntry entry) {
+  ZNMethodToOopsDoClosure nmethod_cl(cl, entry);
   ZNMethodTable::nmethods_do(&nmethod_cl);
 }
 

src/hotspot/share/gc/z/zNMethod.hpp

@@ -31,6 +31,13 @@ class OopClosure;
 class ZReentrantLock;
 class ZWorkers;
 
+enum class ZNMethodEntry {
+  PreBarrier,
+  Disarm,
+  VerifyDisarmed,
+  None
+};
+
 class ZNMethod : public AllStatic {
 private:
   static void attach_gc_data(nmethod* nm);
@@ -52,7 +59,7 @@ class ZNMethod : public AllStatic {
 
   static void oops_do_begin();
   static void oops_do_end();
-  static void oops_do(OopClosure* cl, bool should_disarm_nmethods);
+  static void oops_do(OopClosure* cl, ZNMethodEntry entry);
 
   static ZReentrantLock* lock_for_nmethod(nmethod* nm);
 

src/hotspot/share/gc/z/zOopClosures.hpp

@@ -57,12 +57,16 @@ class ZPhantomKeepAliveOopClosure : public ZRootsIteratorClosure {
 public:
   virtual void do_oop(oop* p);
   virtual void do_oop(narrowOop* p);
+
+  virtual ZNMethodEntry nmethod_entry() const;
 };
 
 class ZPhantomCleanOopClosure : public ZRootsIteratorClosure {
 public:
   virtual void do_oop(oop* p);
   virtual void do_oop(narrowOop* p);
+
+  virtual ZNMethodEntry nmethod_entry() const;
 };
 
 #endif // SHARE_GC_Z_ZOOPCLOSURES_HPP

src/hotspot/share/gc/z/zOopClosures.inline.hpp

@@ -80,6 +80,11 @@ inline void ZPhantomKeepAliveOopClosure::do_oop(oop* p) {
   ZBarrier::keep_alive_barrier_on_phantom_oop_field(p);
 }
 
+inline ZNMethodEntry ZPhantomKeepAliveOopClosure::nmethod_entry() const {
+  ShouldNotReachHere();
+  return ZNMethodEntry::None;
+}
+
 inline void ZPhantomKeepAliveOopClosure::do_oop(narrowOop* p) {
   ShouldNotReachHere();
 }
@@ -104,4 +109,9 @@ inline void ZPhantomCleanOopClosure::do_oop(narrowOop* p) {
   ShouldNotReachHere();
 }
 
+inline ZNMethodEntry ZPhantomCleanOopClosure::nmethod_entry() const {
+  ShouldNotReachHere();
+  return ZNMethodEntry::None;
+}
+
 #endif // SHARE_GC_Z_ZOOPCLOSURES_INLINE_HPP

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

@@ -122,7 +122,7 @@ void ZConcurrentRootsIterator::do_class_loader_data_graph(ZRootsIteratorClosure*
 
 void ZConcurrentRootsIterator::do_code_cache(ZRootsIteratorClosure* cl) {
   ZStatTimer timer(ZSubPhaseConcurrentRootsCodeCache);
-  ZNMethod::oops_do(cl, cl->should_disarm_nmethods());
+  ZNMethod::oops_do(cl, cl->nmethod_entry());
 }
 
 class ZConcurrentRootsIteratorThreadClosure : public ThreadClosure {

src/hotspot/share/gc/z/zRootsIterator.hpp

@@ -26,6 +26,7 @@
 
 #include "classfile/classLoaderDataGraph.hpp"
 #include "gc/shared/oopStorageSetParState.hpp"
+#include "gc/z/zNMethod.hpp"
 #include "memory/allocation.hpp"
 #include "memory/iterator.hpp"
 #include "runtime/threadSMR.hpp"
@@ -61,9 +62,7 @@ class ZRootsIteratorClosure : public OopClosure {
 public:
   virtual void do_thread(Thread* thread) {}
 
-  virtual bool should_disarm_nmethods() const {
-    return false;
-  }
+  virtual ZNMethodEntry nmethod_entry() const = 0;
 };
 
 class ZJavaThreadsIterator {

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

@@ -94,6 +94,11 @@ class ZVerifyRootClosure : public ZRootsIteratorClosure {
   bool verify_fixed() const {
     return _verify_fixed;
   }
+
+  virtual ZNMethodEntry nmethod_entry() const {
+    // Verification performs its own verification
+    return ZNMethodEntry::None;
+  }
 };
 
 class ZVerifyCodeBlobClosure : public CodeBlobToOopClosure {
@@ -188,7 +193,7 @@ void ZVerifyRootClosure::do_thread(Thread* thread) {
   verify_stack.verify_frames();
 }
 
-class ZVerifyOopClosure : public ClaimMetadataVisitingOopIterateClosure, public ZRootsIteratorClosure  {
+class ZVerifyOopClosure : public ClaimMetadataVisitingOopIterateClosure {
 private:
   const bool _verify_weaks;