8321270: Virtual Thread.yield consumes parking permit

Reviewed-by: sspitsyn

Author: Alan Bateman

src/hotspot/share/classfile/javaClasses.cpp

@@ -1986,26 +1986,27 @@ int java_lang_VirtualThread::state(oop vthread) {
 JavaThreadStatus java_lang_VirtualThread::map_state_to_thread_status(int state) {
   JavaThreadStatus status = JavaThreadStatus::NEW;
   switch (state & ~SUSPENDED) {
-    case NEW :
+    case NEW:
       status = JavaThreadStatus::NEW;
       break;
-    case STARTED :
-    case RUNNABLE :
-    case RUNNING :
-    case PARKING :
+    case STARTED:
+    case RUNNING:
+    case PARKING:
     case TIMED_PARKING:
-    case YIELDING :
+    case UNPARKED:
+    case YIELDING:
+    case YIELDED:
       status = JavaThreadStatus::RUNNABLE;
       break;
-    case PARKED :
-    case PINNED :
+    case PARKED:
+    case PINNED:
       status = JavaThreadStatus::PARKED;
       break;
     case TIMED_PARKED:
     case TIMED_PINNED:
       status = JavaThreadStatus::PARKED_TIMED;
       break;
-    case TERMINATED :
+    case TERMINATED:
       status = JavaThreadStatus::TERMINATED;
       break;
     default:

src/hotspot/share/classfile/javaClasses.hpp

@@ -523,15 +523,16 @@ class java_lang_VirtualThread : AllStatic {
   enum {
     NEW           = 0,
     STARTED       = 1,
-    RUNNABLE      = 2,
-    RUNNING       = 3,
-    PARKING       = 4,
-    PARKED        = 5,
-    PINNED        = 6,
-    TIMED_PARKING = 7,
-    TIMED_PARKED  = 8,
-    TIMED_PINNED  = 9,
+    RUNNING       = 2,
+    PARKING       = 3,
+    PARKED        = 4,
+    PINNED        = 5,
+    TIMED_PARKING = 6,
+    TIMED_PARKED  = 7,
+    TIMED_PINNED  = 8,
+    UNPARKED      = 9,
     YIELDING      = 10,
+    YIELDED       = 11,
     TERMINATED    = 99,
 
     // additional state bits

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

@@ -96,37 +96,37 @@ final class VirtualThread extends BaseVirtualThread {
      *  RUNNING -> PARKING         // Thread parking with LockSupport.park
      *  PARKING -> PARKED          // cont.yield successful, parked indefinitely
      *  PARKING -> PINNED          // cont.yield failed, parked indefinitely on carrier
-     *   PARKED -> RUNNABLE        // unparked, schedule to continue
+     *   PARKED -> UNPARKED        // unparked, may be scheduled to continue
      *   PINNED -> RUNNING         // unparked, continue execution on same carrier
+     * UNPARKED -> RUNNING         // continue execution after park
      *
      *       RUNNING -> TIMED_PARKING   // Thread parking with LockSupport.parkNanos
      * TIMED_PARKING -> TIMED_PARKED    // cont.yield successful, timed-parked
      * TIMED_PARKING -> TIMED_PINNED    // cont.yield failed, timed-parked on carrier
-     *  TIMED_PARKED -> RUNNABLE        // unparked, schedule to continue
+     *  TIMED_PARKED -> UNPARKED        // unparked, may be scheduled to continue
      *  TIMED_PINNED -> RUNNING         // unparked, continue execution on same carrier
      *
-     * RUNNABLE -> RUNNING         // continue execution
-     *
      *  RUNNING -> YIELDING        // Thread.yield
-     * YIELDING -> RUNNABLE        // yield successful
-     * YIELDING -> RUNNING         // yield failed
+     * YIELDING -> YIELDED         // cont.yield successful, may be scheduled to continue
+     * YIELDING -> RUNNING         // cont.yield failed
+     *  YIELDED -> RUNNING         // continue execution after Thread.yield
      */
     private static final int NEW      = 0;
     private static final int STARTED  = 1;
-    private static final int RUNNABLE = 2;     // runnable-unmounted
-    private static final int RUNNING  = 3;     // runnable-mounted
-
-    // untimed parking
-    private static final int PARKING  = 4;
-    private static final int PARKED   = 5;     // unmounted
-    private static final int PINNED   = 6;     // mounted
+    private static final int RUNNING  = 2;     // runnable-mounted
 
-    // timed parking
-    private static final int TIMED_PARKING = 7;
-    private static final int TIMED_PARKED  = 8;
-    private static final int TIMED_PINNED  = 9;
+    // untimed and timed parking
+    private static final int PARKING       = 3;
+    private static final int PARKED        = 4;     // unmounted
+    private static final int PINNED        = 5;     // mounted
+    private static final int TIMED_PARKING = 6;
+    private static final int TIMED_PARKED  = 7;     // unmounted
+    private static final int TIMED_PINNED  = 8;     // mounted
+    private static final int UNPARKED      = 9;     // unmounted but runnable
 
-    private static final int YIELDING   = 10;  // Thread.yield
+    // Thread.yield
+    private static final int YIELDING = 10;
+    private static final int YIELDED  = 11;         // unmounted but runnable
 
     private static final int TERMINATED = 99;  // final state
 
@@ -218,11 +218,15 @@ private void runContinuation() {
 
         // set state to RUNNING
         int initialState = state();
-        if (initialState == STARTED && compareAndSetState(STARTED, RUNNING)) {
-            // first run
-        } else if (initialState == RUNNABLE && compareAndSetState(RUNNABLE, RUNNING)) {
-            // consume parking permit
-            setParkPermit(false);
+        if (initialState == STARTED || initialState == UNPARKED || initialState == YIELDED) {
+            // newly started or continue after parking/blocking/Thread.yield
+            if (!compareAndSetState(initialState, RUNNING)) {
+                return;
+            }
+            // consume parking permit when continuing after parking
+            if (initialState == UNPARKED) {
+                setParkPermit(false);
+            }
         } else {
             // not runnable
             return;
@@ -244,8 +248,7 @@ private void runContinuation() {
     /**
      * Submits the runContinuation task to the scheduler. For the default scheduler,
      * and calling it on a worker thread, the task will be pushed to the local queue,
-     * otherwise it will be pushed to a submission queue.
-     *
+     * otherwise it will be pushed to an external submission queue.
      * @throws RejectedExecutionException
      */
     private void submitRunContinuation() {
@@ -258,7 +261,7 @@ private void submitRunContinuation() {
     }
 
     /**
-     * Submits the runContinuation task to the scheduler with a lazy submit.
+     * Submits the runContinuation task to given scheduler with a lazy submit.
      * @throws RejectedExecutionException
      * @see ForkJoinPool#lazySubmit(ForkJoinTask)
      */
@@ -272,7 +275,7 @@ private void lazySubmitRunContinuation(ForkJoinPool pool) {
     }
 
     /**
-     * Submits the runContinuation task to the scheduler as an external submit.
+     * Submits the runContinuation task to the given scheduler as an external submit.
      * @throws RejectedExecutionException
      * @see ForkJoinPool#externalSubmit(ForkJoinTask)
      */
@@ -457,7 +460,7 @@ private void afterYield() {
             setState(newState);
 
             // may have been unparked while parking
-            if (parkPermit && compareAndSetState(newState, RUNNABLE)) {
+            if (parkPermit && compareAndSetState(newState, UNPARKED)) {
                 // lazy submit to continue on the current thread as carrier if possible
                 if (currentThread() instanceof CarrierThread ct) {
                     lazySubmitRunContinuation(ct.getPool());
@@ -471,7 +474,7 @@ private void afterYield() {
 
         // Thread.yield
         if (s == YIELDING) {
-            setState(RUNNABLE);
+            setState(YIELDED);
 
             // external submit if there are no tasks in the local task queue
             if (currentThread() instanceof CarrierThread ct && ct.getQueuedTaskCount() == 0) {
@@ -618,7 +621,7 @@ void parkNanos(long nanos) {
             long startTime = System.nanoTime();
 
             boolean yielded = false;
-            Future<?> unparker = scheduleUnpark(this::unpark, nanos);
+            Future<?> unparker = scheduleUnpark(nanos);  // may throw OOME
             setState(TIMED_PARKING);
             try {
                 yielded = yieldContinuation();  // may throw
@@ -683,14 +686,15 @@ private void parkOnCarrierThread(boolean timed, long nanos) {
     }
 
     /**
-     * Schedule an unpark task to run after a given delay.
+     * Schedule this virtual thread to be unparked after a given delay.
      */
     @ChangesCurrentThread
-    private Future<?> scheduleUnpark(Runnable unparker, long nanos) {
+    private Future<?> scheduleUnpark(long nanos) {
+        assert Thread.currentThread() == this;
         // need to switch to current carrier thread to avoid nested parking
         switchToCarrierThread();
         try {
-            return UNPARKER.schedule(unparker, nanos, NANOSECONDS);
+            return UNPARKER.schedule(this::unpark, nanos, NANOSECONDS);
         } finally {
             switchToVirtualThread(this);
         }
@@ -726,7 +730,7 @@ void unpark() {
         if (!getAndSetParkPermit(true) && currentThread != this) {
             int s = state();
             boolean parked = (s == PARKED) || (s == TIMED_PARKED);
-            if (parked && compareAndSetState(s, RUNNABLE)) {
+            if (parked && compareAndSetState(s, UNPARKED)) {
                 if (currentThread instanceof VirtualThread vthread) {
                     vthread.switchToCarrierThread();
                     try {
@@ -738,7 +742,7 @@ void unpark() {
                     submitRunContinuation();
                 }
             } else if ((s == PINNED) || (s == TIMED_PINNED)) {
-                // unpark carrier thread when pinned.
+                // unpark carrier thread when pinned
                 synchronized (carrierThreadAccessLock()) {
                     Thread carrier = carrierThread;
                     if (carrier != null && ((s = state()) == PINNED || s == TIMED_PINNED)) {
@@ -889,7 +893,8 @@ Thread.State threadState() {
                 } else {
                     return Thread.State.RUNNABLE;
                 }
-            case RUNNABLE:
+            case UNPARKED:
+            case YIELDED:
                 // runnable, not mounted
                 return Thread.State.RUNNABLE;
             case RUNNING:
@@ -905,7 +910,7 @@ Thread.State threadState() {
             case PARKING:
             case TIMED_PARKING:
             case YIELDING:
-                // runnable, mounted, not yet waiting
+                // runnable, in transition
                 return Thread.State.RUNNABLE;
             case PARKED:
             case PINNED:
@@ -947,35 +952,58 @@ StackTraceElement[] asyncGetStackTrace() {
 
     /**
      * Returns the stack trace for this virtual thread if it is unmounted.
-     * Returns null if the thread is in another state.
+     * Returns null if the thread is mounted or in transition.
      */
     private StackTraceElement[] tryGetStackTrace() {
         int initialState = state();
-        return switch (initialState) {
-            case RUNNABLE, PARKED, TIMED_PARKED -> {
-                int suspendedState = initialState | SUSPENDED;
-                if (compareAndSetState(initialState, suspendedState)) {
-                    try {
-                        yield cont.getStackTrace();
-                    } finally {
-                        assert state == suspendedState;
-                        setState(initialState);
-
-                        // re-submit if runnable
-                        // re-submit if unparked while suspended
-                        if (initialState == RUNNABLE
-                            || (parkPermit && compareAndSetState(initialState, RUNNABLE))) {
-                            try {
-                                submitRunContinuation();
-                            } catch (RejectedExecutionException ignore) { }
-                        }
-                    }
-                }
-                yield null;
+        switch (initialState) {
+            case NEW, STARTED, TERMINATED -> {
+                return new StackTraceElement[0];  // unmounted, empty stack
+            }
+            case RUNNING, PINNED -> {
+                return null;   // mounted
+            }
+            case PARKED, TIMED_PARKED -> {
+                // unmounted, not runnable
+            }
+            case UNPARKED, YIELDED -> {
+                // unmounted, runnable
             }
-            case NEW, STARTED, TERMINATED ->  new StackTraceElement[0];  // empty stack
-            default -> null;
+            case PARKING, TIMED_PARKING, YIELDING -> {
+                return null;  // in transition
+            }
+            default -> throw new InternalError();
+        }
+
+        // thread is unmounted, prevent it from continuing
+        int suspendedState = initialState | SUSPENDED;
+        if (!compareAndSetState(initialState, suspendedState)) {
+            return null;
+        }
+
+        // get stack trace and restore state
+        StackTraceElement[] stack;
+        try {
+            stack = cont.getStackTrace();
+        } finally {
+            assert state == suspendedState;
+            setState(initialState);
+        }
+        boolean resubmit = switch (initialState) {
+            case UNPARKED, YIELDED -> {
+                // resubmit as task may have run while suspended
+                yield true;
+            }
+            case PARKED, TIMED_PARKED -> {
+                // resubmit if unparked while suspended
+                yield parkPermit && compareAndSetState(initialState, UNPARKED);
+            }
+            default -> throw new InternalError();
         };
+        if (resubmit) {
+            submitRunContinuation();
+        }
+        return stack;
     }
 
     @Override

test/jdk/java/lang/Thread/virtual/ThreadAPI.java

@@ -23,7 +23,7 @@
 
 /*
  * @test id=default
- * @bug 8284161 8286788
+ * @bug 8284161 8286788 8321270
  * @summary Test Thread API with virtual threads
  * @modules java.base/java.lang:+open
  * @library /test/lib
@@ -1191,6 +1191,36 @@ void testYield2() throws Exception {
         assertEquals(List.of("A", "A", "B"), list);
     }
 
+    /**
+     * Test that Thread.yield does not consume the thread's parking permit.
+     */
+    @Test
+    void testYield3() throws Exception {
+        var thread = Thread.ofVirtual().start(() -> {
+            LockSupport.unpark(Thread.currentThread());
+            Thread.yield();
+            LockSupport.park();  // should not park
+        });
+        thread.join();
+    }
+
+    /**
+     * Test that Thread.yield does not make available the thread's parking permit.
+     */
+    @Test
+    void testYield4() throws Exception {
+        var thread = Thread.ofVirtual().start(() -> {
+            Thread.yield();
+            LockSupport.park();  // should park
+        });
+        try {
+            await(thread, Thread.State.WAITING);
+        } finally {
+            LockSupport.unpark(thread);
+            thread.join();
+        }
+    }
+
     /**
      * Test Thread.onSpinWait.
      */