8212181: ZGC: Fix incorrect root iteration in ZHeapIterator

Reviewed-by: eosterlund

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

@@ -29,7 +29,6 @@
 #include "gc/z/zOop.inline.hpp"
 #include "gc/z/zRootsIterator.hpp"
 #include "memory/iterator.inline.hpp"
-#include "oops/oop.inline.hpp"
 #include "utilities/bitMap.inline.hpp"
 #include "utilities/stack.inline.hpp"
 
@@ -54,19 +53,18 @@ public:
 class ZHeapIteratorRootOopClosure : public ZRootsIteratorClosure {
 private:
   ZHeapIterator* const _iter;
-  ObjectClosure* const _cl;
 
 public:
-  ZHeapIteratorRootOopClosure(ZHeapIterator* iter, ObjectClosure* cl) :
-      _iter(iter),
-      _cl(cl) {}
+  ZHeapIteratorRootOopClosure(ZHeapIterator* iter) :
+      _iter(iter) {}
 
   virtual void do_oop(oop* p) {
     // Load barrier needed here for the same reason we
-    // need fixup_partial_loads() in ZHeap::mark_end()
+    // need fixup_partial_loads() in ZHeap::mark_end().
+    // This barrier is also needed here in case we're
+    // treating the JVMTI weak tag map as strong roots.
     const oop obj = ZBarrier::load_barrier_on_oop_field(p);
     _iter->push(obj);
-    _iter->drain(_cl);
   }
 
   virtual void do_oop(narrowOop* p) {
@@ -74,19 +72,13 @@ public:
   }
 };
 
-class ZHeapIteratorPushOopClosure : public BasicOopIterateClosure {
+class ZHeapIteratorOopClosure : public BasicOopIterateClosure {
 private:
   ZHeapIterator* const _iter;
   const oop            _base;
   const bool           _visit_referents;
 
-public:
-  ZHeapIteratorPushOopClosure(ZHeapIterator* iter, oop base) :
-      _iter(iter),
-      _base(base),
-      _visit_referents(iter->visit_referents()) {}
-
-  oop load_oop(oop* p) {
+  oop load_oop(oop* p) const {
     if (_visit_referents) {
       return HeapAccess<ON_UNKNOWN_OOP_REF>::oop_load_at(_base, _base->field_offset(p));
     } else {
@@ -94,6 +86,12 @@ public:
     }
   }
 
+public:
+  ZHeapIteratorOopClosure(ZHeapIterator* iter, oop base, bool visit_referents) :
+      _iter(iter),
+      _base(base),
+      _visit_referents(visit_referents) {}
+
   virtual ReferenceIterationMode reference_iteration_mode() {
     return _visit_referents ? DO_FIELDS : DO_FIELDS_EXCEPT_REFERENT;
   }
@@ -126,11 +124,11 @@ ZHeapIterator::~ZHeapIterator() {
   }
 }
 
-size_t ZHeapIterator::object_index_max() const {
+static size_t object_index_max() {
   return ZPageSizeMin >> ZObjectAlignmentSmallShift;
 }
 
-size_t ZHeapIterator::object_index(oop obj) const {
+static size_t object_index(oop obj) {
   const uintptr_t addr = ZOop::to_address(obj);
   const uintptr_t offset = ZAddress::offset(addr);
   const uintptr_t mask = (1 << ZPageSizeMinShift) - 1;
@@ -165,7 +163,20 @@ void ZHeapIterator::push(oop obj) {
   _visit_stack.push(obj);
 }
 
-void ZHeapIterator::drain(ObjectClosure* cl) {
+void ZHeapIterator::objects_do(ObjectClosure* cl) {
+  // Push roots onto stack
+  {
+    // Note that we also visit the JVMTI weak tag map as if they were
+    // strong roots to make sure we visit all tagged objects, even
+    // those that might now have become unreachable. If we didn't do
+    // this the user would have expected to see ObjectFree events for
+    // unreachable objects in the tag map.
+    ZRootsIterator roots;
+    ZHeapIteratorRootOopClosure root_cl(this);
+    roots.oops_do(&root_cl, true /* visit_jvmti_weak_export */);
+  }
+
+  // Drain stack
   while (!_visit_stack.is_empty()) {
     const oop obj = _visit_stack.pop();
 
@@ -173,23 +184,7 @@ void ZHeapIterator::drain(ObjectClosure* cl) {
     cl->do_object(obj);
 
     // Push members to visit
-    ZHeapIteratorPushOopClosure push_cl(this, obj);
+    ZHeapIteratorOopClosure push_cl(this, obj, _visit_referents);
     obj->oop_iterate(&push_cl);
   }
 }
-
-bool ZHeapIterator::visit_referents() const {
-  return _visit_referents;
-}
-
-void ZHeapIterator::objects_do(ObjectClosure* cl) {
-  ZHeapIteratorRootOopClosure root_cl(this, cl);
-  ZRootsIterator roots;
-
-  // Follow roots. Note that we also visit the JVMTI weak tag map
-  // as if they were strong roots to make sure we visit all tagged
-  // objects, even those that might now have become unreachable.
-  // If we didn't do this the user would have expected to see
-  // ObjectFree events for unreachable objects in the tag map.
-  roots.oops_do(&root_cl, true /* visit_jvmti_weak_export */);
-}