8315373: Change VirtualThread to unmount after freezing, re-mount before thawing

Reviewed-by: pchilanomate, mgronlun, sspitsyn

src/hotspot/share/jfr/recorder/stacktrace/jfrStackTrace.cpp

@@ -138,9 +138,9 @@ void JfrStackFrame::write(JfrCheckpointWriter& cpw) const {
 
 class JfrVframeStream : public vframeStreamCommon {
  private:
+  bool _vthread;
   const ContinuationEntry* _cont_entry;
   bool _async_mode;
-  bool _vthread;
   bool step_to_sender();
   void next_frame();
  public:
@@ -165,8 +165,9 @@ JfrVframeStream::JfrVframeStream(JavaThread* jt, const frame& fr, bool stop_at_j
                                  RegisterMap::UpdateMap::skip,
                                  RegisterMap::ProcessFrames::skip,
                                  walk_continuation(jt))),
-    _cont_entry(JfrThreadLocal::is_vthread(jt) ? jt->last_continuation() : nullptr),
-    _async_mode(async_mode), _vthread(JfrThreadLocal::is_vthread(jt)) {
+    _vthread(JfrThreadLocal::is_vthread(jt)),
+    _cont_entry(_vthread ? jt->last_continuation() : nullptr),
+    _async_mode(async_mode) {
   assert(!_vthread || _cont_entry != nullptr, "invariant");
   _reg_map.set_async(async_mode);
   _frame = fr;

src/hotspot/share/jfr/support/jfrThreadLocal.cpp

@@ -395,11 +395,14 @@ traceid JfrThreadLocal::thread_id(const Thread* t) {
     return t->jfr_thread_local()->_thread_id_alias;
   }
   JfrThreadLocal* const tl = t->jfr_thread_local();
-  if (!t->is_Java_thread() || !Atomic::load_acquire(&tl->_vthread)) {
+  if (!t->is_Java_thread()) {
+    return jvm_thread_id(t, tl);
+  }
+  const JavaThread* jt = JavaThread::cast(t);
+  if (!is_vthread(jt)) {
     return jvm_thread_id(t, tl);
   }
   // virtual thread
-  const JavaThread* jt = JavaThread::cast(t);
   const traceid tid = vthread_id(jt);
   assert(tid != 0, "invariant");
   if (!tl->is_vthread_excluded()) {
@@ -456,7 +459,7 @@ traceid JfrThreadLocal::jvm_thread_id(const Thread* t) {
 
 bool JfrThreadLocal::is_vthread(const JavaThread* jt) {
   assert(jt != nullptr, "invariant");
-  return Atomic::load_acquire(&jt->jfr_thread_local()->_vthread);
+  return Atomic::load_acquire(&jt->jfr_thread_local()->_vthread) && jt->last_continuation() != nullptr;
 }
 
 inline bool is_virtual(const JavaThread* jt, oop thread) {

src/java.base/share/classes/java/lang/VirtualThread.java

@@ -194,8 +194,11 @@ final class VirtualThread extends BaseVirtualThread {
     }
 
     /**
-     * Runs or continues execution of the continuation on the current thread.
+     * Runs or continues execution on the current thread. The virtual thread is mounted
+     * on the current thread before the task runs or continues. It unmounts when the
+     * task completes or yields.
      */
+    @ChangesCurrentThread
     private void runContinuation() {
         // the carrier must be a platform thread
         if (Thread.currentThread().isVirtual()) {
@@ -217,11 +220,13 @@ final class VirtualThread extends BaseVirtualThread {
         // notify JVMTI before mount
         notifyJvmtiMount(/*hide*/true);
 
+        mount();
         try {
             cont.run();
         } finally {
+            unmount();
             if (cont.isDone()) {
-                afterTerminate();
+                afterDone();
             } else {
                 afterYield();
             }
@@ -285,16 +290,13 @@ final class VirtualThread extends BaseVirtualThread {
     }
 
     /**
-     * Runs a task in the context of this virtual thread. The virtual thread is
-     * mounted on the current (carrier) thread before the task runs. It unmounts
-     * from its carrier thread when the task completes.
+     * Runs a task in the context of this virtual thread.
      */
     @ChangesCurrentThread
     private void run(Runnable task) {
-        assert state == RUNNING;
+        assert Thread.currentThread() == this && state == RUNNING;
 
-        // first mount
-        mount();
+        // notify JVMTI, may post VirtualThreadStart event
         notifyJvmtiStart();
 
         // emit JFR event if enabled
@@ -322,12 +324,8 @@ final class VirtualThread extends BaseVirtualThread {
                 }
 
             } finally {
-                // last unmount
+                // notify JVMTI, may post VirtualThreadEnd event
                 notifyJvmtiEnd();
-                unmount();
-
-                // final state
-                setState(TERMINATED);
             }
         }
     }
@@ -417,21 +415,15 @@ final class VirtualThread extends BaseVirtualThread {
      }
 
     /**
-     * Unmounts this virtual thread, invokes Continuation.yield, and re-mounts the
-     * thread when continued. When enabled, JVMTI must be notified from this method.
-     * @return true if the yield was successful
+     * Invokes Continuation.yield, notifying JVMTI (if enabled) to hide frames until
+     * the continuation continues.
      */
     @Hidden
-    @ChangesCurrentThread
     private boolean yieldContinuation() {
-        // unmount
         notifyJvmtiUnmount(/*hide*/true);
-        unmount();
         try {
             return Continuation.yield(VTHREAD_SCOPE);
         } finally {
-            // re-mount
-            mount();
             notifyJvmtiMount(/*hide*/false);
         }
     }
@@ -477,22 +469,22 @@ final class VirtualThread extends BaseVirtualThread {
     }
 
     /**
-     * Invoked after the thread terminates execution. It notifies anyone
-     * waiting for the thread to terminate.
+     * Invoked after the continuation completes.
      */
-    private void afterTerminate() {
-        afterTerminate(true, true);
+    private void afterDone() {
+        afterDone(true, true);
     }
 
     /**
-     * Invoked after the thread terminates (or start failed). This method
-     * notifies anyone waiting for the thread to terminate.
+     * Invoked after the continuation completes (or start failed). Sets the thread
+     * state to TERMINATED and notifies anyone waiting for the thread to terminate.
      *
      * @param notifyContainer true if its container should be notified
      * @param executed true if the thread executed, false if it failed to start
      */
-    private void afterTerminate(boolean notifyContainer, boolean executed) {
-        assert (state() == TERMINATED) && (carrierThread == null);
+    private void afterDone(boolean notifyContainer, boolean executed) {
+        assert carrierThread == null;
+        setState(TERMINATED);
 
         if (executed) {
             notifyJvmtiUnmount(/*hide*/false);
@@ -546,8 +538,7 @@ final class VirtualThread extends BaseVirtualThread {
             started = true;
         } finally {
             if (!started) {
-                setState(TERMINATED);
-                afterTerminate(addedToContainer, /*executed*/false);
+                afterDone(addedToContainer, /*executed*/false);
             }
         }
     }