8350765: Need to pin when accessing thread container from virtual thread

Reviewed-by: vklang, jpai

Author: Alan Bateman

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

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 1994, 2024, Oracle and/or its affiliates. All rights reserved.
+ * Copyright (c) 1994, 2025, 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
@@ -1409,7 +1409,7 @@ void start(ThreadContainer container) {
 
             // start thread
             boolean started = false;
-            container.onStart(this);  // may throw
+            container.add(this);  // may throw
             try {
                 // scoped values may be inherited
                 inheritScopedValueBindings(container);
@@ -1418,7 +1418,7 @@ void start(ThreadContainer container) {
                 started = true;
             } finally {
                 if (!started) {
-                    container.onExit(this);
+                    container.remove(this);
                 }
             }
         }
@@ -1487,7 +1487,7 @@ private void exit() {
             // notify container that thread is exiting
             ThreadContainer container = threadContainer();
             if (container != null) {
-                container.onExit(this);
+                container.remove(this);
             }
         }
 

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

@@ -664,7 +664,7 @@ private void afterDone(boolean notifyContainer) {
 
         // notify container
         if (notifyContainer) {
-            threadContainer().onExit(this);
+            threadContainer().remove(this);
         }
 
         // clear references to thread locals
@@ -692,7 +692,7 @@ void start(ThreadContainer container) {
         boolean addedToContainer = false;
         boolean started = false;
         try {
-            container.onStart(this);  // may throw
+            container.add(this);  // may throw
             addedToContainer = true;
 
             // scoped values may be inherited

src/java.base/share/classes/java/lang/ref/ReferenceQueue.java

@@ -27,7 +27,6 @@
 
 import java.util.function.Consumer;
 import jdk.internal.misc.VM;
-import jdk.internal.vm.Continuation;
 import jdk.internal.vm.ContinuationSupport;
 
 /**
@@ -145,19 +144,6 @@ boolean enqueue(Reference<? extends T> r) { /* Called only by Reference class */
         }
     }
 
-    private boolean tryDisablePreempt() {
-        if (Thread.currentThread().isVirtual() && ContinuationSupport.isSupported()) {
-            Continuation.pin();
-            return true;
-        } else {
-            return false;
-        }
-    }
-
-    private void enablePreempt() {
-        Continuation.unpin();
-    }
-
     /**
      * Polls this queue to see if a reference object is available.  If one is
      * available without further delay then it is removed from the queue and
@@ -173,13 +159,13 @@ public Reference<? extends T> poll() {
 
         // Prevent a virtual thread from being preempted as this could potentially
         // deadlock with a carrier that is polling the same reference queue.
-        boolean disabled = tryDisablePreempt();
+        ContinuationSupport.pinIfSupported();
         try {
             synchronized (lock) {
                 return poll0();
             }
         } finally {
-            if (disabled) enablePreempt();
+            ContinuationSupport.unpinIfSupported();
         }
     }
 

src/java.base/share/classes/jdk/internal/vm/ContinuationSupport.java

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2022, Oracle and/or its affiliates. All rights reserved.
+ * Copyright (c) 2022, 2025, 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
@@ -50,5 +50,23 @@ public static void ensureSupported() {
         }
     }
 
+    /**
+     * Pins the current continuation if the VM has continuations support.
+     */
+    public static void pinIfSupported() {
+        if (isSupported()) {
+            Continuation.pin();
+        }
+    }
+
+    /**
+     * Unpins the current continuation if the VM has continuations support.
+     */
+    public static void unpinIfSupported() {
+        if (isSupported()) {
+            Continuation.unpin();
+        }
+    }
+
     private static native boolean isSupported0();
 }

src/java.base/share/classes/jdk/internal/vm/ThreadContainer.java

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2021, 2023, Oracle and/or its affiliates. All rights reserved.
+ * Copyright (c) 2021, 2025, 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
@@ -76,24 +76,49 @@ public long threadCount() {
     public abstract Stream<Thread> threads();
 
     /**
-     * Invoked by Thread::start before the given Thread is started.
+     * Invoked by {@code add} to add a thread to this container before it starts.
      */
-    public void onStart(Thread thread) {
-        // do nothing
+    protected void onStart(Thread thread) {
     }
 
     /**
-     * Invoked when a Thread terminates or starting it fails.
-     *
-     * For a platform thread, this method is invoked by the thread itself when it
-     * terminates. For a virtual thread, this method is invoked on its carrier
-     * after the virtual thread has terminated.
-     *
-     * If starting the Thread failed then this method is invoked on the thread
-     * that invoked onStart.
+     * Invoked by {@code remove} to remove a thread from this container when it
+     * terminates (or failed to start).
      */
-    public void onExit(Thread thread) {
-        // do nothing
+    protected void onExit(Thread thread) {
+    }
+
+    /**
+     * Adds a thread to this container. This method should be invoked before the
+     * thread executes.
+     */
+    public final void add(Thread thread) {
+        // Prevent a virtual thread from being preempted as this could potentially
+        // deadlock when scheduled to continue and all carriers are blocked adding
+        // or removing virtual threads.
+        ContinuationSupport.pinIfSupported();
+        try {
+            onStart(thread);
+        } finally {
+            ContinuationSupport.unpinIfSupported();
+        }
+    }
+
+    /**
+     * Remove a thread from this container. This method can be invoked by the thread
+     * itself as it terminates, or it can be invoked by another thread after the given
+     * thread has terminated (or failed to start).
+     */
+    public final void remove(Thread thread) {
+        // Prevent a virtual thread from being preempted as this could potentially
+        // deadlock when scheduled to continue and all carriers are blocked adding
+        // or removing virtual threads.
+        ContinuationSupport.pinIfSupported();
+        try {
+            onExit(thread);
+        } finally {
+            ContinuationSupport.unpinIfSupported();
+        }
     }
 
     /**

src/java.base/windows/classes/sun/nio/fs/WindowsSecurity.java

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2008, 2024, Oracle and/or its affiliates. All rights reserved.
+ * Copyright (c) 2008, 2025, 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
@@ -25,7 +25,6 @@
 
 package sun.nio.fs;
 
-import jdk.internal.vm.Continuation;
 import jdk.internal.vm.ContinuationSupport;
 
 import static sun.nio.fs.WindowsNativeDispatcher.*;
@@ -106,9 +105,7 @@ static Privilege enablePrivilege(String priv) {
         final boolean needToRevert = elevated;
 
         // prevent yielding with privileges
-        if (ContinuationSupport.isSupported())
-            Continuation.pin();
-
+        ContinuationSupport.pinIfSupported();
         return () -> {
             try {
                 if (token != 0L) {
@@ -126,8 +123,7 @@ else if (needToRevert)
                 }
             } finally {
                 LocalFree(pLuid);
-                if (ContinuationSupport.isSupported())
-                    Continuation.unpin();
+                ContinuationSupport.unpinIfSupported();
             }
         };
     }