8299414: JVMTI FollowReferences should support references from VirtualThread stack

Reviewed-by: sspitsyn, kevinw

src/hotspot/share/prims/jvmtiTagMap.cpp

@@ -49,6 +49,8 @@
 #include "prims/jvmtiImpl.hpp"
 #include "prims/jvmtiTagMap.hpp"
 #include "prims/jvmtiTagMapTable.hpp"
+#include "prims/jvmtiThreadState.hpp"
+#include "runtime/continuationWrapper.inline.hpp"
 #include "runtime/deoptimization.hpp"
 #include "runtime/frame.inline.hpp"
 #include "runtime/handles.inline.hpp"
@@ -2205,6 +2207,157 @@ class JNILocalRootsClosure : public OopClosure {
   virtual void do_oop(narrowOop* obj_p) { ShouldNotReachHere(); }
 };
 
+// Helper class to collect/report stack references.
+class StackRefCollector {
+private:
+  JvmtiTagMap* _tag_map;
+  JNILocalRootsClosure* _blk;
+  // java_thread is needed only to report JNI local on top native frame;
+  // I.e. it's required only for platform/carrier threads or mounted virtual threads.
+  JavaThread* _java_thread;
+
+  oop _threadObj;
+  jlong _thread_tag;
+  jlong _tid;
+
+  bool _is_top_frame;
+  int _depth;
+  frame* _last_entry_frame;
+
+  bool report_java_stack_refs(StackValueCollection* values, jmethodID method, jlocation bci, jint slot_offset);
+  bool report_native_stack_refs(jmethodID method);
+
+public:
+  StackRefCollector(JvmtiTagMap* tag_map, JNILocalRootsClosure* blk, JavaThread* java_thread)
+    : _tag_map(tag_map), _blk(blk), _java_thread(java_thread),
+      _threadObj(nullptr), _thread_tag(0), _tid(0),
+      _is_top_frame(true), _depth(0), _last_entry_frame(nullptr)
+  {
+  }
+
+  bool set_thread(oop o);
+  // Sets the thread and reports the reference to it with the specified kind.
+  bool set_thread(jvmtiHeapReferenceKind kind, oop o);
+
+  bool do_frame(vframe* vf);
+  // Handles frames until vf->sender() is null.
+  bool process_frames(vframe* vf);
+};
+
+bool StackRefCollector::set_thread(oop o) {
+  _threadObj = o;
+  _thread_tag = tag_for(_tag_map, _threadObj);
+  _tid = java_lang_Thread::thread_id(_threadObj);
+
+  _is_top_frame = true;
+  _depth = 0;
+  _last_entry_frame = nullptr;
+
+  return true;
+}
+
+bool StackRefCollector::set_thread(jvmtiHeapReferenceKind kind, oop o) {
+  return set_thread(o)
+         && CallbackInvoker::report_simple_root(kind, _threadObj);
+}
+
+bool StackRefCollector::report_java_stack_refs(StackValueCollection* values, jmethodID method, jlocation bci, jint slot_offset) {
+  for (int index = 0; index < values->size(); index++) {
+    if (values->at(index)->type() == T_OBJECT) {
+      oop obj = values->obj_at(index)();
+      if (obj == nullptr) {
+        continue;
+      }
+      // stack reference
+      if (!CallbackInvoker::report_stack_ref_root(_thread_tag, _tid, _depth, method,
+                                                  bci, slot_offset + index, obj)) {
+        return false;
+      }
+    }
+  }
+  return true;
+}
+
+bool StackRefCollector::report_native_stack_refs(jmethodID method) {
+  _blk->set_context(_thread_tag, _tid, _depth, method);
+  if (_is_top_frame) {
+    // JNI locals for the top frame.
+    assert(_java_thread != nullptr, "sanity");
+    _java_thread->active_handles()->oops_do(_blk);
+    if (_blk->stopped()) {
+      return false;
+    }
+  } else {
+    if (_last_entry_frame != nullptr) {
+      // JNI locals for the entry frame.
+      assert(_last_entry_frame->is_entry_frame(), "checking");
+      _last_entry_frame->entry_frame_call_wrapper()->handles()->oops_do(_blk);
+      if (_blk->stopped()) {
+        return false;
+      }
+    }
+  }
+  return true;
+}
+
+bool StackRefCollector::do_frame(vframe* vf) {
+  if (vf->is_java_frame()) {
+    // java frame (interpreted, compiled, ...)
+    javaVFrame* jvf = javaVFrame::cast(vf);
+
+    jmethodID method = jvf->method()->jmethod_id();
+
+    if (!(jvf->method()->is_native())) {
+      jlocation bci = (jlocation)jvf->bci();
+      StackValueCollection* locals = jvf->locals();
+      if (!report_java_stack_refs(locals, method, bci, 0)) {
+        return false;
+      }
+      if (!report_java_stack_refs(jvf->expressions(), method, bci, locals->size())) {
+        return false;
+      }
+
+      // Follow oops from compiled nmethod.
+      if (jvf->cb() != nullptr && jvf->cb()->is_nmethod()) {
+        _blk->set_context(_thread_tag, _tid, _depth, method);
+        jvf->cb()->as_nmethod()->oops_do(_blk);
+        if (_blk->stopped()) {
+          return false;
+        }
+      }
+    } else {
+      // native frame
+      if (!report_native_stack_refs(method)) {
+        return false;
+      }
+    }
+    _last_entry_frame = nullptr;
+    _depth++;
+  } else {
+    // externalVFrame - for an entry frame then we report the JNI locals
+    // when we find the corresponding javaVFrame
+    frame* fr = vf->frame_pointer();
+    assert(fr != nullptr, "sanity check");
+    if (fr->is_entry_frame()) {
+      _last_entry_frame = fr;
+    }
+  }
+
+  _is_top_frame = false;
+
+  return true;
+}
+
+bool StackRefCollector::process_frames(vframe* vf) {
+  while (vf != nullptr) {
+    if (!do_frame(vf)) {
+      return false;
+    }
+    vf = vf->sender();
+  }
+  return true;
+}
+
 
 // A VM operation to iterate over objects that are reachable from
 // a set of roots or an initial object.
@@ -2268,7 +2421,8 @@ class VM_HeapWalkOperation: public VM_Operation {
   // root collection
   inline bool collect_simple_roots();
   inline bool collect_stack_roots();
-  inline bool collect_stack_roots(JavaThread* java_thread, JNILocalRootsClosure* blk);
+  inline bool collect_stack_refs(JavaThread* java_thread, JNILocalRootsClosure* blk);
+  inline bool collect_vthread_stack_refs(oop vt);
 
   // visit an object
   inline bool visit(oop o);
@@ -2617,121 +2771,66 @@ inline bool VM_HeapWalkOperation::collect_simple_roots() {
   return true;
 }
 
-// Walk the stack of a given thread and find all references (locals
-// and JNI calls) and report these as stack references
-inline bool VM_HeapWalkOperation::collect_stack_roots(JavaThread* java_thread,
-                                                      JNILocalRootsClosure* blk)
+// Reports the thread as JVMTI_HEAP_REFERENCE_THREAD,
+// walks the stack of the thread, finds all references (locals
+// and JNI calls) and reports these as stack references.
+inline bool VM_HeapWalkOperation::collect_stack_refs(JavaThread* java_thread,
+                                                     JNILocalRootsClosure* blk)
 {
   oop threadObj = java_thread->threadObj();
+  oop mounted_vt = java_thread->is_vthread_mounted() ? java_thread->vthread() : nullptr;
+  if (mounted_vt != nullptr && !JvmtiEnvBase::is_vthread_alive(mounted_vt)) {
+    mounted_vt = nullptr;
+  }
   assert(threadObj != nullptr, "sanity check");
 
-  // only need to get the thread's tag once per thread
-  jlong thread_tag = tag_for(_tag_map, threadObj);
+  StackRefCollector stack_collector(tag_map(), blk, java_thread);
 
-  // also need the thread id
-  jlong tid = java_lang_Thread::thread_id(threadObj);
+  if (!java_thread->has_last_Java_frame()) {
+    if (!stack_collector.set_thread(JVMTI_HEAP_REFERENCE_THREAD, threadObj)) {
+      return false;
+    }
+    // no last java frame but there may be JNI locals
+    blk->set_context(tag_for(_tag_map, threadObj), java_lang_Thread::thread_id(threadObj), 0, (jmethodID)nullptr);
+    java_thread->active_handles()->oops_do(blk);
+    return !blk->stopped();
+  }
+  // vframes are resource allocated
+  Thread* current_thread = Thread::current();
+  ResourceMark rm(current_thread);
+  HandleMark hm(current_thread);
 
+  RegisterMap reg_map(java_thread,
+                      RegisterMap::UpdateMap::include,
+                      RegisterMap::ProcessFrames::include,
+                      RegisterMap::WalkContinuation::include);
 
-  if (java_thread->has_last_Java_frame()) {
-
-    // vframes are resource allocated
-    Thread* current_thread = Thread::current();
-    ResourceMark rm(current_thread);
-    HandleMark hm(current_thread);
-
-    RegisterMap reg_map(java_thread,
-                        RegisterMap::UpdateMap::include,
-                        RegisterMap::ProcessFrames::include,
-                        RegisterMap::WalkContinuation::skip);
+  // first handle mounted vthread (if any)
+  if (mounted_vt != nullptr) {
     frame f = java_thread->last_frame();
     vframe* vf = vframe::new_vframe(&f, &reg_map, java_thread);
-
-    bool is_top_frame = true;
-    int depth = 0;
-    frame* last_entry_frame = nullptr;
-
-    while (vf != nullptr) {
-      if (vf->is_java_frame()) {
-
-        // java frame (interpreted, compiled, ...)
-        javaVFrame *jvf = javaVFrame::cast(vf);
-
-        // the jmethodID
-        jmethodID method = jvf->method()->jmethod_id();
-
-        if (!(jvf->method()->is_native())) {
-          jlocation bci = (jlocation)jvf->bci();
-          StackValueCollection* locals = jvf->locals();
-          for (int slot=0; slot<locals->size(); slot++) {
-            if (locals->at(slot)->type() == T_OBJECT) {
-              oop o = locals->obj_at(slot)();
-              if (o == nullptr) {
-                continue;
-              }
-
-              // stack reference
-              if (!CallbackInvoker::report_stack_ref_root(thread_tag, tid, depth, method,
-                                                   bci, slot, o)) {
-                return false;
-              }
-            }
-          }
-
-          StackValueCollection* exprs = jvf->expressions();
-          for (int index=0; index < exprs->size(); index++) {
-            if (exprs->at(index)->type() == T_OBJECT) {
-              oop o = exprs->obj_at(index)();
-              if (o == nullptr) {
-                continue;
-              }
-
-              // stack reference
-              if (!CallbackInvoker::report_stack_ref_root(thread_tag, tid, depth, method,
-                                                   bci, locals->size() + index, o)) {
-                return false;
-              }
-            }
-          }
-
-          // Follow oops from compiled nmethod
-          if (jvf->cb() != nullptr && jvf->cb()->is_nmethod()) {
-            blk->set_context(thread_tag, tid, depth, method);
-            jvf->cb()->as_nmethod()->oops_do(blk);
-          }
-        } else {
-          blk->set_context(thread_tag, tid, depth, method);
-          if (is_top_frame) {
-            // JNI locals for the top frame.
-            java_thread->active_handles()->oops_do(blk);
-          } else {
-            if (last_entry_frame != nullptr) {
-              // JNI locals for the entry frame
-              assert(last_entry_frame->is_entry_frame(), "checking");
-              last_entry_frame->entry_frame_call_wrapper()->handles()->oops_do(blk);
-            }
-          }
-        }
-        last_entry_frame = nullptr;
-        depth++;
-      } else {
-        // externalVFrame - for an entry frame then we report the JNI locals
-        // when we find the corresponding javaVFrame
-        frame* fr = vf->frame_pointer();
-        assert(fr != nullptr, "sanity check");
-        if (fr->is_entry_frame()) {
-          last_entry_frame = fr;
-        }
-      }
-
-      vf = vf->sender();
-      is_top_frame = false;
+    // report virtual thread as JVMTI_HEAP_REFERENCE_OTHER
+    if (!stack_collector.set_thread(JVMTI_HEAP_REFERENCE_OTHER, mounted_vt)) {
+      return false;
+    }
+    // split virtual thread and carrier thread stacks by vthread entry ("enterSpecial") frame,
+    // consider vthread entry frame as the last vthread stack frame
+    while (vf != nullptr) {
+      if (!stack_collector.do_frame(vf)) {
+        return false;
+      }
+      if (vf->is_vthread_entry()) {
+        break;
+      }
+      vf = vf->sender();
     }
-  } else {
-    // no last java frame but there may be JNI locals
-    blk->set_context(thread_tag, tid, 0, (jmethodID)nullptr);
-    java_thread->active_handles()->oops_do(blk);
   }
-  return true;
+  // Platform or carrier thread.
+  vframe* vf = JvmtiEnvBase::get_cthread_last_java_vframe(java_thread, &reg_map);
+  if (!stack_collector.set_thread(JVMTI_HEAP_REFERENCE_THREAD, threadObj)) {
+    return false;
+  }
+  return stack_collector.process_frames(vf);
 }
 
 
@@ -2743,13 +2842,7 @@ inline bool VM_HeapWalkOperation::collect_stack_roots() {
   for (JavaThreadIteratorWithHandle jtiwh; JavaThread *thread = jtiwh.next(); ) {
     oop threadObj = thread->threadObj();
     if (threadObj != nullptr && !thread->is_exiting() && !thread->is_hidden_from_external_view()) {
-      // Collect the simple root for this thread before we
-      // collect its stack roots
-      if (!CallbackInvoker::report_simple_root(JVMTI_HEAP_REFERENCE_THREAD,
-                                               threadObj)) {
-        return false;
-      }
-      if (!collect_stack_roots(thread, &blk)) {
+      if (!collect_stack_refs(thread, &blk)) {
         return false;
       }
     }
@@ -2757,6 +2850,42 @@ inline bool VM_HeapWalkOperation::collect_stack_roots() {
   return true;
 }
 
+// Reports stack references for the unmounted virtual thread.
+inline bool VM_HeapWalkOperation::collect_vthread_stack_refs(oop vt) {
+  if (!JvmtiEnvBase::is_vthread_alive(vt)) {
+    return true;
+  }
+  ContinuationWrapper cont(java_lang_VirtualThread::continuation(vt));
+  if (cont.is_empty()) {
+    return true;
+  }
+  assert(!cont.is_mounted(), "sanity check");
+
+  stackChunkOop chunk = cont.last_nonempty_chunk();
+  if (chunk == nullptr || chunk->is_empty()) {
+    return true;
+  }
+
+  // vframes are resource allocated
+  Thread* current_thread = Thread::current();
+  ResourceMark rm(current_thread);
+  HandleMark hm(current_thread);
+
+  RegisterMap reg_map(cont.continuation(), RegisterMap::UpdateMap::include);
+
+  JNILocalRootsClosure blk;
+  // JavaThread is not required for unmounted virtual threads
+  StackRefCollector stack_collector(tag_map(), &blk, nullptr);
+  // reference to the vthread is already reported
+  if (!stack_collector.set_thread(vt)) {
+    return false;
+  }
+
+  frame fr = chunk->top_frame(&reg_map);
+  vframe* vf = vframe::new_vframe(&fr, &reg_map, nullptr);
+  return stack_collector.process_frames(vf);
+}
+
 // visit an object
 // first mark the object as visited
 // second get all the outbound references from this object (in other words, all
@@ -2775,6 +2904,13 @@ bool VM_HeapWalkOperation::visit(oop o) {
         return iterate_over_class(o);
       }
     } else {
+      // we report stack references only when initial object is not specified
+      // (in the case we start from heap roots which include platform thread stack references)
+      if (initial_object().is_null() && java_lang_VirtualThread::is_subclass(o->klass())) {
+        if (!collect_vthread_stack_refs(o)) {
+          return false;
+        }
+      }
       return iterate_over_object(o);
     }
   }
@@ -2837,6 +2973,9 @@ void JvmtiTagMap::iterate_over_reachable_objects(jvmtiHeapRootCallback heap_root
   eb.deoptimize_objects_all_threads();
   Arena dead_object_arena(mtServiceability);
   GrowableArray<jlong> dead_objects(&dead_object_arena, 10, 0, 0);
+
+  JvmtiVTMSTransitionDisabler disabler;
+
   {
     MutexLocker ml(Heap_lock);
     BasicHeapWalkContext context(heap_root_callback, stack_ref_callback, object_ref_callback);
@@ -2856,6 +2995,9 @@ void JvmtiTagMap::iterate_over_objects_reachable_from_object(jobject object,
 
   Arena dead_object_arena(mtServiceability);
   GrowableArray<jlong> dead_objects(&dead_object_arena, 10, 0, 0);
+
+  JvmtiVTMSTransitionDisabler disabler;
+
   {
     MutexLocker ml(Heap_lock);
     BasicHeapWalkContext context(nullptr, nullptr, object_ref_callback);
@@ -2884,6 +3026,9 @@ void JvmtiTagMap::follow_references(jint heap_filter,
 
   Arena dead_object_arena(mtServiceability);
   GrowableArray<jlong> dead_objects(&dead_object_arena, 10, 0, 0);
+
+  JvmtiVTMSTransitionDisabler disabler;
+
   {
     MutexLocker ml(Heap_lock);
     AdvancedHeapWalkContext context(heap_filter, klass, callbacks);

test/hotspot/jtreg/serviceability/jvmti/vthread/FollowReferences/VThreadStackRefTest.java

@@ -0,0 +1,246 @@
+/*
+ * Copyright (c) 2023, Oracle and/or its affiliates. All rights reserved.
+ * DO NOT ALTER OR REMOVE COPYRIGHT NOTICES OR THIS FILE HEADER.
+ *
+ * This code is free software; you can redistribute it and/or modify it
+ * under the terms of the GNU General Public License version 2 only, as
+ * published by the Free Software Foundation.
+ *
+ * This code is distributed in the hope that it will be useful, but WITHOUT
+ * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
+ * FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public License
+ * version 2 for more details (a copy is included in the LICENSE file that
+ * accompanied this code).
+ *
+ * You should have received a copy of the GNU General Public License version
+ * 2 along with this work; if not, write to the Free Software Foundation,
+ * Inc., 51 Franklin St, Fifth Floor, Boston, MA 02110-1301 USA.
+ *
+ * Please contact Oracle, 500 Oracle Parkway, Redwood Shores, CA 94065 USA
+ * or visit www.oracle.com if you need additional information or have any
+ * questions.
+ */
+
+/**
+ * @test id=default
+ * @requires vm.jvmti
+ * @requires vm.continuations
+ * @run main/othervm/native
+ *      -Djdk.virtualThreadScheduler.parallelism=1
+ *      -agentlib:VThreadStackRefTest
+ *      VThreadStackRefTest
+ */
+
+/**
+ * @test id=no-vmcontinuations
+ * @requires vm.jvmti
+ * @run main/othervm/native
+ *      -XX:+UnlockExperimentalVMOptions -XX:-VMContinuations
+ *      -agentlib:VThreadStackRefTest
+ *      VThreadStackRefTest NoMountCheck
+ */
+
+import java.lang.ref.Reference;
+import java.util.stream.Stream;
+import java.util.concurrent.CountDownLatch;
+
+/*
+ * The test verifies JVMTI FollowReferences function reports references from
+ * mounted and unmounted virtual threads and reports correct thread id
+ * (for mounted vthread it should be vthread id, and not carrier thread id).
+ * Additionally tests that references from platform threads are reported correctly
+ * and that references from terminated vthread are not reported.
+ * To get both mounted and unmounted vthreads the test:
+ * - limits the number of carrier threads to 1;
+ * - starts vthread that creates a stack local and JNI local
+ *   and then waits in CountDownLatch.await();
+ * - starts another vthread that create stack local and JNI local (on top frame)
+ *   and waits in native to avoid unmounting.
+ */
+public class VThreadStackRefTest {
+
+    // Currently we cannot test JNI locals for unmounted threads
+    // as native calls pin virtual thread.
+    // TODO: revise if this changes.
+    static final boolean testUnmountedJNILocals = false;
+
+    // The flag is set by createObjAndWait method.
+    static volatile boolean mountedVthreadReady = false;
+
+    public static void main(String[] args) throws InterruptedException {
+        boolean noMountCheck = args.length > 0
+                               && args[0].equalsIgnoreCase("NoMountCheck");
+        CountDownLatch dumpedLatch = new CountDownLatch(1);
+
+        CountDownLatch unmountedThreadReady = new CountDownLatch(1);
+        // Unmounted virtual thread with stack local.
+        Thread vthreadUnmounted = Thread.ofVirtual().start(() -> {
+            Object referenced = new VThreadUnmountedReferenced();
+            System.out.println("created " + referenced.getClass());
+            if (testUnmountedJNILocals) {
+                createObjAndCallback(VThreadUnmountedJNIReferenced.class,
+                    new Runnable() {
+                        public void run() {
+                            unmountedThreadReady.countDown();
+                            await(dumpedLatch);
+                        }
+                    });
+            } else {
+                unmountedThreadReady.countDown();
+                await(dumpedLatch);
+            }
+            Reference.reachabilityFence(referenced);
+        });
+        // Wait until unmounted thread is ready.
+        unmountedThreadReady.await();
+
+        // Ended virtual thread with stack local - should not be reported.
+        Thread vthreadEnded = Thread.ofVirtual().start(() -> {
+            Object referenced = new VThreadUnmountedEnded();
+            System.out.println("created " + referenced.getClass());
+            Reference.reachabilityFence(referenced);
+        });
+        // Make sure this vthread has exited so we can test
+        // that it no longer holds any stack references.
+        vthreadEnded.join();
+
+        // Mounted virtual thread with stack local and JNI local on top frame.
+        Thread vthreadMounted = Thread.ofVirtual().start(() -> {
+            Object referenced = new VThreadMountedReferenced();
+            System.out.println("created " + referenced.getClass());
+            createObjAndWait(VThreadMountedJNIReferenced.class);
+            Reference.reachabilityFence(referenced);
+        });
+        // Wait until mounted vthread is ready.
+        while (!mountedVthreadReady) {
+            Thread.sleep(10);
+        }
+
+        CountDownLatch pThreadReady = new CountDownLatch(1);
+        // Sanity check - reference from platform thread stack.
+        Thread pthread = Thread.ofPlatform().start(() -> {
+            Object referenced = new PThreadReferenced();
+            System.out.println("created " + referenced.getClass());
+            pThreadReady.countDown();
+            await(dumpedLatch);
+            Reference.reachabilityFence(referenced);
+        });
+        // Wait until platform thread is ready.
+        pThreadReady.await();
+
+        System.out.println("threads:");
+        System.out.println("  - vthreadUnmounted: " + vthreadUnmounted);
+        System.out.println("  - vthreadEnded: " + vthreadEnded);
+        System.out.println("  - vthreadMounted: " + vthreadMounted);
+        System.out.println("  - pthread: " + pthread);
+
+        TestCase[] testCases = new TestCase[] {
+            new TestCase(VThreadUnmountedReferenced.class, 1, vthreadUnmounted.getId()),
+            new TestCase(VThreadUnmountedJNIReferenced.class,
+                         testUnmountedJNILocals ? 1 : 0,
+                         testUnmountedJNILocals ? vthreadUnmounted.getId() : 0),
+            new TestCase(VThreadMountedReferenced.class, 1, vthreadMounted.getId()),
+            new TestCase(VThreadMountedJNIReferenced.class, 1, vthreadMounted.getId()),
+            new TestCase(PThreadReferenced.class, 1, pthread.getId()),
+            // expected to be unreported as stack local
+            new TestCase(VThreadUnmountedEnded.class, 0, 0)
+        };
+
+        Class[] testClasses = Stream.of(testCases).map(c -> c.cls()).toArray(Class[]::new);
+        System.out.println("test classes:");
+        for (int i = 0; i < testClasses.length; i++) {
+            System.out.println("  (" + i + ") " + testClasses[i]);
+        }
+
+        try {
+            if (noMountCheck) {
+                System.out.println("INFO: No mount/unmount checks");
+            } else {
+                verifyVthreadMounted(vthreadUnmounted, false);
+                verifyVthreadMounted(vthreadMounted, true);
+            }
+
+            test(testClasses);
+        } finally {
+            // Finish all threads
+            endWait();               // signal mounted vthread to exit
+            dumpedLatch.countDown(); // signal unmounted vthread and platform thread to exit
+        }
+
+        vthreadMounted.join();
+        vthreadUnmounted.join();
+        pthread.join();
+
+        boolean failed = false;
+        for (int i = 0; i < testCases.length; i++) {
+            int refCount = getRefCount(i);
+            long threadId = getRefThreadID(i);
+            String status = "OK";
+            if (refCount != testCases[i].expectedCount()
+                    || threadId != testCases[i].expectedThreadId()) {
+                failed = true;
+                status = "ERROR";
+            }
+            System.out.println("  (" + i + ") " + status
+                               + " " + testCases[i].cls()
+                               + ": ref count = " + refCount
+                               + " (expected " + testCases[i].expectedCount() + ")"
+                               + ", thread id = " + threadId
+                               + " (expected " + testCases[i].expectedThreadId() + ")");
+        }
+        if (failed) {
+            throw new RuntimeException("Test failed");
+        }
+    }
+
+    private static void await(CountDownLatch dumpedLatch) {
+        try {
+            dumpedLatch.await();
+        } catch (InterruptedException e) {
+            throw new RuntimeException(e);
+        }
+    }
+
+    private static void verifyVthreadMounted(Thread t, boolean expectedMounted) {
+        // Hacky, but simple.
+        // If virtual thread is mounted, its toString() contains
+        // info about carrier thread, something like
+        // VirtualThread[#27]/runnable@ForkJoinPool-1-worker-1
+        String s = t.toString();
+        boolean mounted = t.isVirtual() && s.contains("/runnable@");
+        System.out.println("Thread " + t + ": " + (mounted ? "mounted" : "unmounted"));
+        if (mounted != expectedMounted) {
+            throw new RuntimeException("Thread " + t + " has unexpected mount state");
+        }
+    }
+
+    private static native void test(Class<?>... classes);
+    private static native int getRefCount(int index);
+    private static native long getRefThreadID(int index);
+
+    // Creates object of the the specified class (local JNI)
+    // and calls the provided callback.
+    private static native void createObjAndCallback(Class cls, Runnable callback);
+    // Creates object of the the specified class (local JNI),
+    // sets mountedVthreadReady static field,
+    // and then waits until endWait() method is called.
+    private static native void createObjAndWait(Class cls);
+    // Signals createObjAndWait() to exit.
+    private static native void endWait();
+
+    private record TestCase(Class cls, int expectedCount, long expectedThreadId) {
+    }
+
+    public static class VThreadUnmountedReferenced {
+    }
+    public static class VThreadUnmountedJNIReferenced {
+    }
+    public static class VThreadUnmountedEnded {
+    }
+    public static class VThreadMountedReferenced {
+    }
+    public static class VThreadMountedJNIReferenced {
+    }
+    public static class PThreadReferenced {
+    }
+}

test/hotspot/jtreg/serviceability/jvmti/vthread/FollowReferences/libVThreadStackRefTest.cpp

@@ -0,0 +1,208 @@
+/*
+ * Copyright (c) 2023, Oracle and/or its affiliates. All rights reserved.
+ * DO NOT ALTER OR REMOVE COPYRIGHT NOTICES OR THIS FILE HEADER.
+ *
+ * This code is free software; you can redistribute it and/or modify it
+ * under the terms of the GNU General Public License version 2 only, as
+ * published by the Free Software Foundation.
+ *
+ * This code is distributed in the hope that it will be useful, but WITHOUT
+ * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
+ * FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public License
+ * version 2 for more details (a copy is included in the LICENSE file that
+ * accompanied this code).
+ *
+ * You should have received a copy of the GNU General Public License version
+ * 2 along with this work; if not, write to the Free Software Foundation,
+ * Inc., 51 Franklin St, Fifth Floor, Boston, MA 02110-1301 USA.
+ *
+ * Please contact Oracle, 500 Oracle Parkway, Redwood Shores, CA 94065 USA
+ * or visit www.oracle.com if you need additional information or have any
+ * questions.
+ */
+
+#include <jni.h>
+#include <jvmti.h>
+#include <jvmti_common.h>
+#include <atomic>
+#include <string.h>
+
+namespace {
+
+jvmtiEnv *jvmti = nullptr;
+
+const int TAG_START = 100;
+
+struct RefCounters {
+  jint test_class_count;
+  jint *count;
+  jlong *thread_id;
+
+  RefCounters(): test_class_count(0), count(nullptr) {}
+
+  void* alloc(JNIEnv* env, jlong size) {
+    unsigned char* ptr;
+    jvmtiError err = jvmti->Allocate(size, &ptr);
+    if (err != JVMTI_ERROR_NONE) {
+      env->FatalError("jvmti->Allocate failed");
+    }
+    memset(ptr, 0, size);
+    return ptr;
+  }
+
+  void init(JNIEnv* env, jint test_class_count) {
+    this->test_class_count = test_class_count;
+    count = (jint*)alloc(env, sizeof(count[0]) *  test_class_count);
+    thread_id = (jlong*)alloc(env, sizeof(thread_id[0]) *  test_class_count);
+  }
+} refCounters;
+
+}
+
+/////////////////////////////////////////
+// Agent functions
+/////////////////////////////////////////
+jint JNICALL
+HeapReferenceCallback(jvmtiHeapReferenceKind reference_kind,
+                      const jvmtiHeapReferenceInfo* reference_info,
+                      jlong class_tag, jlong referrer_class_tag, jlong size,
+                      jlong* tag_ptr, jlong* referrer_tag_ptr, jint length, void* user_data) {
+  if (class_tag >= TAG_START) {
+    jlong index = class_tag - TAG_START;
+    switch (reference_kind) {
+    case JVMTI_HEAP_REFERENCE_STACK_LOCAL: {
+      jvmtiHeapReferenceInfoStackLocal *stackInfo = (jvmtiHeapReferenceInfoStackLocal *)reference_info;
+      refCounters.count[index]++;
+      refCounters.thread_id[index] = stackInfo->thread_id;
+      LOG("Stack local: index = %d, thread_id = %d\n",
+          (int)index, (int)stackInfo->thread_id);
+      if (refCounters.count[index] > 1) {
+        LOG("ERROR: count > 1: %d\n", (int)refCounters.count[index]);
+      }
+    }
+    break;
+    case JVMTI_HEAP_REFERENCE_JNI_LOCAL: {
+      jvmtiHeapReferenceInfoJniLocal *jniInfo = (jvmtiHeapReferenceInfoJniLocal *)reference_info;
+      refCounters.count[index]++;
+      refCounters.thread_id[index] = jniInfo->thread_id;
+      LOG("JNI local: index = %d, thread_id = %d\n",
+          (int)index, (int)jniInfo->thread_id);
+      if (refCounters.count[index] > 1) {
+        LOG("ERROR: count > 1: %d\n", (int)refCounters.count[index]);
+      }
+    }
+    break;
+    default:
+      // unexpected ref.kind
+      LOG("ERROR: unexpected ref_kind for class %d: %d\n",
+          (int)index, (int)reference_kind);
+    }
+  }
+  return JVMTI_VISIT_OBJECTS;
+}
+
+extern "C" JNIEXPORT jint JNICALL
+Agent_OnLoad(JavaVM *vm, char *options, void *reserved) {
+  if (vm->GetEnv((void **)&jvmti, JVMTI_VERSION) != JNI_OK) {
+    LOG("Could not initialize JVMTI\n");
+    return JNI_ERR;
+  }
+  jvmtiCapabilities capabilities;
+  memset(&capabilities, 0, sizeof(capabilities));
+  capabilities.can_tag_objects = 1;
+  jvmtiError err = jvmti->AddCapabilities(&capabilities);
+  if (err != JVMTI_ERROR_NONE) {
+    LOG("JVMTI AddCapabilities error: %d\n", err);
+    return JNI_ERR;
+  }
+
+  return JNI_OK;
+}
+
+
+/////////////////////////////////////////
+// Test native methods
+/////////////////////////////////////////
+extern "C" JNIEXPORT void JNICALL
+Java_VThreadStackRefTest_test(JNIEnv* env, jclass clazz, jobjectArray classes) {
+  jsize classes_count = env->GetArrayLength(classes);
+  for (int i = 0; i < classes_count; i++) {
+    jvmti->SetTag(env->GetObjectArrayElement(classes, i), TAG_START + i);
+  }
+  refCounters.init(env, classes_count);
+  jvmtiHeapCallbacks callbacks;
+  memset(&callbacks, 0, sizeof(jvmtiHeapCallbacks));
+  callbacks.heap_reference_callback = HeapReferenceCallback;
+  jvmtiError err = jvmti->FollowReferences(0, nullptr, nullptr, &callbacks, nullptr);
+  if (err != JVMTI_ERROR_NONE) {
+    LOG("JVMTI FollowReferences error: %d\n", err);
+    env->FatalError("FollowReferences failed");
+  }
+}
+
+extern "C" JNIEXPORT jint JNICALL
+Java_VThreadStackRefTest_getRefCount(JNIEnv* env, jclass clazz, jint index) {
+  return refCounters.count[index];
+}
+
+extern "C" JNIEXPORT jlong JNICALL
+Java_VThreadStackRefTest_getRefThreadID(JNIEnv* env, jclass clazz, jint index) {
+  return refCounters.thread_id[index];
+}
+
+static void print_created_class(JNIEnv* env, jclass cls) {
+  jmethodID mid = env->GetMethodID(cls, "toString", "()Ljava/lang/String;");
+  if (mid == nullptr) {
+    env->FatalError("failed to get toString method");
+    return;
+  }
+  jstring jstr = (jstring)env->CallObjectMethod(cls, mid);
+  const char* str = env->GetStringUTFChars(jstr, 0);
+  LOG("created %s\n", str);
+  env->ReleaseStringUTFChars(jstr, str);
+}
+
+// Creates object of the the specified class (local JNI)
+// and calls the provided callback.
+extern "C" JNIEXPORT void JNICALL
+Java_VThreadStackRefTest_createObjAndCallback(JNIEnv* env, jclass clazz, jclass cls, jobject callback) {
+  jobject jobj = env->AllocObject(cls);
+  print_created_class(env, cls);
+
+  jclass callbackClass = env->GetObjectClass(callback);
+  jmethodID mid = env->GetMethodID(callbackClass, "run", "()V");
+  if (mid == nullptr) {
+    env->FatalError("cannot get run method");
+    return;
+  }
+  env->CallVoidMethod(callback, mid);
+}
+
+static std::atomic<bool> time_to_exit(false);
+
+// Creates object of the the specified class (local JNI),
+// sets mountedVthreadReady static field,
+// and then waits until endWait() method is called.
+extern "C" JNIEXPORT void JNICALL
+Java_VThreadStackRefTest_createObjAndWait(JNIEnv* env, jclass clazz, jclass cls) {
+  jobject jobj = env->AllocObject(cls);
+  print_created_class(env, cls);
+
+  // Notify main thread that we are ready
+  jfieldID fid = env->GetStaticFieldID(clazz, "mountedVthreadReady", "Z");
+  if (fid == nullptr) {
+    env->FatalError("cannot get mountedVthreadReady field");
+    return;
+  }
+  env->SetStaticBooleanField(clazz, fid, JNI_TRUE);
+
+  while (!time_to_exit) {
+    sleep_ms(100);
+  }
+}
+
+// Signals createObjAndWait() to exit.
+extern "C" JNIEXPORT void JNICALL
+Java_VThreadStackRefTest_endWait(JNIEnv* env, jclass clazz) {
+  time_to_exit = true;
+}