Merge pull request #3870 from armanbilge/issue/3869

Block in-place if running on a virtual thread

Author: djspiewak

core/js-native/src/main/scala/cats/effect/IOFiberConstants.scala

@@ -16,6 +16,8 @@
 
 package cats.effect
 
+import org.typelevel.scalaccompat.annotation._
+
 // defined in Java for the JVM, Scala for ScalaJS (where object field access is faster)
 private object IOFiberConstants {
 
@@ -43,4 +45,7 @@ private object IOFiberConstants {
   final val CedeR = 6
   final val AutoCedeR = 7
   final val DoneR = 8
+
+  @nowarn212
+  @inline def isVirtualThread(t: Thread): Boolean = false
 }

core/jvm/src/main/java/cats/effect/IOFiberConstants.java

@@ -16,6 +16,10 @@
 
 package cats.effect;
 
+import java.lang.invoke.MethodHandle;
+import java.lang.invoke.MethodHandles;
+import java.lang.invoke.MethodType;
+
 // defined in Java since Scala doesn't let us define static fields
 final class IOFiberConstants {
 
@@ -43,4 +47,28 @@ final class IOFiberConstants {
   static final byte CedeR = 6;
   static final byte AutoCedeR = 7;
   static final byte DoneR = 8;
+
+  static boolean isVirtualThread(final Thread thread) {
+    try {
+      return (boolean) THREAD_IS_VIRTUAL_HANDLE.invokeExact(thread);
+    } catch (Throwable t) {
+      return false;
+    }
+  }
+
+  private static final MethodHandle THREAD_IS_VIRTUAL_HANDLE;
+
+  static {
+    final MethodHandles.Lookup lookup = MethodHandles.publicLookup();
+    final MethodType mt = MethodType.methodType(boolean.class);
+    MethodHandle mh;
+    try {
+      mh = lookup.findVirtual(Thread.class, "isVirtual", mt);
+    } catch (Throwable t) {
+      mh =
+          MethodHandles.dropArguments(
+              MethodHandles.constant(boolean.class, false), 0, Thread.class);
+    }
+    THREAD_IS_VIRTUAL_HANDLE = mh;
+  }
 }

core/shared/src/main/scala/cats/effect/IOFiber.scala

@@ -996,6 +996,20 @@ private final class IOFiber[A](
               } else {
                 blockingFallback(cur)
               }
+            } else if (isVirtualThread(Thread.currentThread())) {
+              var error: Throwable = null
+              val r =
+                try {
+                  cur.thunk()
+                } catch {
+                  case t if NonFatal(t) =>
+                    error = t
+                  case t: Throwable =>
+                    onFatalFailure(t)
+                }
+
+              val next = if (error eq null) succeeded(r, 0) else failed(error, 0)
+              runLoop(next, nextCancelation, nextAutoCede)
             } else {
               blockingFallback(cur)
             }

tests/jvm/src/test/scala/cats/effect/DetectPlatform.scala

@@ -21,4 +21,7 @@ trait DetectPlatform {
   def isJS: Boolean = false
   def isJVM: Boolean = true
   def isNative: Boolean = false
+
+  def javaMajorVersion: Int =
+    System.getProperty("java.version").stripPrefix("1.").takeWhile(_.isDigit).toInt
 }

tests/jvm/src/test/scala/cats/effect/IOPlatformSpecification.scala

@@ -36,12 +36,13 @@ import java.util.concurrent.{
   CancellationException,
   CompletableFuture,
   CountDownLatch,
+  ExecutorService,
   Executors,
   ThreadLocalRandom
 }
 import java.util.concurrent.atomic.{AtomicBoolean, AtomicInteger, AtomicLong, AtomicReference}
 
-trait IOPlatformSpecification { self: BaseSpec with ScalaCheck =>
+trait IOPlatformSpecification extends DetectPlatform { self: BaseSpec with ScalaCheck =>
 
   def platformSpecs = {
     "platform" should {
@@ -531,6 +532,25 @@ trait IOPlatformSpecification { self: BaseSpec with ScalaCheck =>
           runtime.shutdown()
         }
       }
+
+      if (javaMajorVersion >= 21)
+        "block in-place on virtual threads" in real {
+          val loomExec = classOf[Executors]
+            .getDeclaredMethod("newVirtualThreadPerTaskExecutor")
+            .invoke(null)
+            .asInstanceOf[ExecutorService]
+
+          val loomEc = ExecutionContext.fromExecutor(loomExec)
+
+          IO.blocking {
+            classOf[Thread]
+              .getDeclaredMethod("isVirtual")
+              .invoke(Thread.currentThread())
+              .asInstanceOf[Boolean]
+          }.evalOn(loomEc)
+        }
+      else
+        "block in-place on virtual threads" in skipped("virtual threads not supported")
     }
   }
 }