feat(env): Add support for virtual thread detection (#9852)

Author: PerfectSlayer

components/environment/build.gradle.kts

@@ -24,7 +24,7 @@ val excludedClassesCoverage by extra {
     "datadog.environment.JavaVirtualMachine.JvmOptionsHolder", // depends on OS and JVM vendor
     "datadog.environment.JvmOptions", // depends on OS and JVM vendor
     "datadog.environment.OperatingSystem**", // depends on OS
-    "datadog.environment.ThreadUtils", // depends on JVM version
+    "datadog.environment.ThreadSupport", // requires Java 21
   )
 }
 val excludedClassesBranchCoverage by extra {

components/environment/src/main/java/datadog/environment/ThreadSupport.java

@@ -0,0 +1,142 @@
+package datadog.environment;
+
+import java.lang.invoke.MethodHandle;
+import java.lang.invoke.MethodHandles;
+import java.lang.invoke.MethodType;
+import java.util.Optional;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+
+/**
+ * Helper class for working with {@link Thread}s.
+ *
+ * <p>Uses feature detection and provides static helpers to work with different versions of Java.
+ *
+ * <p>This class is designed to use {@link MethodHandle}s that constant propagate to minimize the
+ * overhead.
+ */
+public final class ThreadSupport {
+  static final MethodHandle THREAD_ID_MH = findThreadIdMethodHandle();
+  static final MethodHandle IS_VIRTUAL_MH = findIsVirtualMethodHandle();
+  static final MethodHandle NEW_VIRTUAL_THREAD_PER_TASK_EXECUTOR_MH =
+      findNewVirtualThreadPerTaskExecutorMethodHandle();
+
+  private ThreadSupport() {}
+
+  /**
+   * Provides the best identifier available for the current {@link Thread}. Uses {@link
+   * Thread#threadId()} on 19+ or {@link Thread#getId()} on older JVMs.
+   *
+   * @return The best identifier available for the current {@link Thread}.
+   */
+  public static long threadId() {
+    return threadId(Thread.currentThread());
+  }
+
+  /**
+   * Provides the best identifier available for the given {@link Thread}. Uses {@link
+   * Thread#threadId()} on 19+ or {@link Thread#getId()} on older JVMs.
+   *
+   * @return The best identifier available for the given {@link Thread}.
+   */
+  public static long threadId(Thread thread) {
+    if (THREAD_ID_MH != null) {
+      try {
+        return (long) THREAD_ID_MH.invoke(thread);
+      } catch (Throwable ignored) {
+      }
+    }
+    return thread.getId();
+  }
+
+  /**
+   * Checks whether virtual threads are supported on this JVM.
+   *
+   * @return {@code true} if virtual threads are supported, {@code false} otherwise.
+   */
+  public static boolean supportsVirtualThreads() {
+    return (IS_VIRTUAL_MH != null);
+  }
+
+  /**
+   * Checks whether the current thread is a virtual thread.
+   *
+   * @return {@code true} if the current thread is a virtual thread, {@code false} otherwise.
+   */
+  public static boolean isVirtual() {
+    // IS_VIRTUAL_MH will constant propagate -- then dead code eliminate -- and inline as needed
+    return IS_VIRTUAL_MH != null && isVirtual(Thread.currentThread());
+  }
+
+  /**
+   * Checks whether the given thread is a virtual thread.
+   *
+   * @param thread The thread to check.
+   * @return {@code true} if the given thread is virtual, {@code false} otherwise.
+   */
+  public static boolean isVirtual(Thread thread) {
+    // IS_VIRTUAL_MH will constant propagate -- then dead code eliminate -- and inline as needed
+    if (IS_VIRTUAL_MH != null) {
+      try {
+        return (boolean) IS_VIRTUAL_MH.invoke(thread);
+      } catch (Throwable ignored) {
+      }
+    }
+    return false;
+  }
+
+  /**
+   * Returns the virtual thread per task executor if available.
+   *
+   * @return The virtual thread per task executor if available wrapped into an {@link Optional}, or
+   *     {@link Optional#empty()} otherwise.
+   */
+  public static Optional<ExecutorService> newVirtualThreadPerTaskExecutor() {
+    if (NEW_VIRTUAL_THREAD_PER_TASK_EXECUTOR_MH != null) {
+      try {
+        ExecutorService executorService =
+            (ExecutorService) NEW_VIRTUAL_THREAD_PER_TASK_EXECUTOR_MH.invoke();
+        return Optional.of(executorService);
+      } catch (Throwable ignored) {
+      }
+    }
+    return Optional.empty();
+  }
+
+  private static MethodHandle findThreadIdMethodHandle() {
+    if (JavaVirtualMachine.isJavaVersionAtLeast(19)) {
+      try {
+        return MethodHandles.lookup()
+            .findVirtual(Thread.class, "threadId", MethodType.methodType(long.class));
+      } catch (Throwable ignored) {
+        return null;
+      }
+    }
+    return null;
+  }
+
+  private static MethodHandle findIsVirtualMethodHandle() {
+    if (JavaVirtualMachine.isJavaVersionAtLeast(21)) {
+      try {
+        return MethodHandles.lookup()
+            .findVirtual(Thread.class, "isVirtual", MethodType.methodType(boolean.class));
+      } catch (Throwable ignored) {
+      }
+    }
+    return null;
+  }
+
+  private static MethodHandle findNewVirtualThreadPerTaskExecutorMethodHandle() {
+    if (JavaVirtualMachine.isJavaVersionAtLeast(21)) {
+      try {
+        return MethodHandles.lookup()
+            .findStatic(
+                Executors.class,
+                "newVirtualThreadPerTaskExecutor",
+                MethodType.methodType(ExecutorService.class));
+      } catch (Throwable ignored) {
+      }
+    }
+    return null;
+  }
+}

components/environment/src/main/java/datadog/environment/ThreadUtils.java

@@ -0,0 +1,82 @@
+package datadog.environment;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.fail;
+import static org.junit.jupiter.api.condition.JRE.JAVA_21;
+
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.Future;
+import java.util.concurrent.atomic.AtomicLong;
+import org.junit.jupiter.api.AfterAll;
+import org.junit.jupiter.api.BeforeAll;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.condition.EnabledOnJre;
+
+class ThreadSupportTest {
+  private static ExecutorService singleThreadExecutor;
+  private static ExecutorService newVirtualThreadPerTaskExecutor;
+
+  @BeforeAll
+  static void beforeAll() {
+    singleThreadExecutor = Executors.newSingleThreadExecutor();
+    newVirtualThreadPerTaskExecutor = ThreadSupport.newVirtualThreadPerTaskExecutor().orElse(null);
+  }
+
+  @Test
+  public void testThreadId() throws InterruptedException {
+    AtomicLong threadId = new AtomicLong();
+    Thread thread = new Thread(() -> threadId.set(ThreadSupport.threadId()), "foo");
+    thread.start();
+    try {
+      // always works on Thread's where getId isn't overridden by child class
+      assertEquals(thread.getId(), ThreadSupport.threadId(thread));
+    } finally {
+      thread.join();
+    }
+    assertEquals(thread.getId(), threadId.get());
+  }
+
+  @Test
+  void testSupportsVirtualThreads() {
+    assertEquals(
+        JavaVirtualMachine.isJavaVersionAtLeast(21),
+        ThreadSupport.supportsVirtualThreads(),
+        "expected virtual threads support status");
+  }
+
+  @Test
+  void testPlatformThread() {
+    assertVirtualThread(singleThreadExecutor, false);
+  }
+
+  @Test
+  @EnabledOnJre(JAVA_21)
+  void testVirtualThread() {
+    assertVirtualThread(newVirtualThreadPerTaskExecutor, true);
+  }
+
+  static void assertVirtualThread(ExecutorService executorService, boolean expected) {
+    Future<Boolean> futureCurrent = executorService.submit(() -> ThreadSupport.isVirtual());
+    Future<Boolean> futureGiven =
+        executorService.submit(
+            () -> {
+              Thread thread = Thread.currentThread();
+              return ThreadSupport.isVirtual(thread);
+            });
+    try {
+      assertEquals(expected, futureCurrent.get(), "invalid current thread virtual status");
+      assertEquals(expected, futureGiven.get(), "invalid given thread virtual status");
+    } catch (Throwable e) {
+      fail("Can't get thread virtual status", e);
+    }
+  }
+
+  @AfterAll
+  static void afterAll() {
+    singleThreadExecutor.shutdown();
+    if (newVirtualThreadPerTaskExecutor != null) {
+      newVirtualThreadPerTaskExecutor.shutdown();
+    }
+  }
+}