test/lib/jdk/test/lib/thread/VThreadScheduler.java

*** 21,12 ***
   * questions.
   */
  
  package jdk.test.lib.thread;
  
- import java.lang.reflect.Constructor;
  import java.lang.reflect.Field;
  import java.lang.reflect.InvocationTargetException;
  import java.util.concurrent.Executor;
  import java.util.concurrent.Executors;
  import java.util.concurrent.ThreadFactory;
  
--- 21,12 ---
   * questions.
   */
  
  package jdk.test.lib.thread;
  
  import java.lang.reflect.Field;
+ import java.lang.reflect.Method;
  import java.lang.reflect.InvocationTargetException;
  import java.util.concurrent.Executor;
  import java.util.concurrent.Executors;
  import java.util.concurrent.ThreadFactory;
  

*** 36,21 ***
   * Tests using this class need to open java.base/java.lang.
   */
  public class VThreadScheduler {
      private VThreadScheduler() { }
  
      /**
       * Returns the scheduler for the given virtual thread.
       */
!     public static Executor scheduler(Thread thread) {
          if (!thread.isVirtual())
              throw new IllegalArgumentException("Not a virtual thread");
          try {
              Field scheduler = Class.forName("java.lang.VirtualThread")
                      .getDeclaredField("scheduler");
              scheduler.setAccessible(true);
!             return (Executor) scheduler.get(thread);
          } catch (Exception e) {
              throw new RuntimeException(e);
          }
      }
  
--- 36,41 ---
   * Tests using this class need to open java.base/java.lang.
   */
  public class VThreadScheduler {
      private VThreadScheduler() { }
  
+     /**
+      * Returns the default virtual thread scheduler.
+      */
+     public static Thread.VirtualThreadScheduler defaultScheduler() {
+         try {
+             Method m = Class.forName("java.lang.VirtualThread")
+                     .getDeclaredMethod("defaultScheduler");
+             m.setAccessible(true);
+             return (Thread.VirtualThreadScheduler ) m.invoke(null);
+         } catch (InvocationTargetException e) {
+             Throwable cause = e.getCause();
+             if (cause instanceof RuntimeException re) {
+                 throw re;
+             }
+             throw new RuntimeException(e);
+         } catch (Exception e) {
+             throw new RuntimeException(e);
+         }
+     }
+ 
      /**
       * Returns the scheduler for the given virtual thread.
       */
!     public static Thread.VirtualThreadScheduler scheduler(Thread thread) {
          if (!thread.isVirtual())
              throw new IllegalArgumentException("Not a virtual thread");
          try {
              Field scheduler = Class.forName("java.lang.VirtualThread")
                      .getDeclaredField("scheduler");
              scheduler.setAccessible(true);
!             return (Thread.VirtualThreadScheduler) scheduler.get(thread);
          } catch (Exception e) {
              throw new RuntimeException(e);
          }
      }
  

*** 58,11 ***
       * Return true if custom schedulers are supported.
       */
      public static boolean supportsCustomScheduler() {
          try (var pool = Executors.newCachedThreadPool()) {
              try {
!                 virtualThreadBuilder(pool);
                  return true;
              } catch (UnsupportedOperationException e) {
                  return false;
              }
          }
--- 78,11 ---
       * Return true if custom schedulers are supported.
       */
      public static boolean supportsCustomScheduler() {
          try (var pool = Executors.newCachedThreadPool()) {
              try {
!                 virtualThreadBuilder(pool).unstarted(() -> { });
                  return true;
              } catch (UnsupportedOperationException e) {
                  return false;
              }
          }

*** 70,30 ***
  
      /**
       * Returns a builder to create virtual threads that use the given scheduler.
       * @throws UnsupportedOperationException if custom schedulers are not supported
       */
!     public static Thread.Builder.OfVirtual virtualThreadBuilder(Executor scheduler) {
          try {
              Class<?> clazz = Class.forName("java.lang.ThreadBuilders$VirtualThreadBuilder");
!             Constructor<?> ctor = clazz.getDeclaredConstructor(Executor.class);
!             ctor.setAccessible(true);
!             return (Thread.Builder.OfVirtual) ctor.newInstance(scheduler);
!         } catch (InvocationTargetException e) {
-             Throwable cause = e.getCause();
-             if (cause instanceof RuntimeException re) {
-                 throw re;
-             }
-             throw new RuntimeException(e);
          } catch (Exception e) {
              throw new RuntimeException(e);
          }
      }
  
      /**
       * Returns a ThreadFactory to create virtual threads that use the given scheduler.
       * @throws UnsupportedOperationException if custom schedulers are not supported
       */
!     public static ThreadFactory virtualThreadFactory(Executor scheduler) {
          return virtualThreadBuilder(scheduler).factory();
      }
  }
--- 90,48 ---
  
      /**
       * Returns a builder to create virtual threads that use the given scheduler.
       * @throws UnsupportedOperationException if custom schedulers are not supported
       */
!     public static Thread.Builder.OfVirtual virtualThreadBuilder(Thread.VirtualThreadScheduler scheduler) {
+         Thread.Builder.OfVirtual builder = Thread.ofVirtual();
          try {
              Class<?> clazz = Class.forName("java.lang.ThreadBuilders$VirtualThreadBuilder");
!             Field field = clazz.getDeclaredField("scheduler");
!             field.setAccessible(true);
!             field.set(builder, scheduler);
!             return builder;
          } catch (Exception e) {
              throw new RuntimeException(e);
          }
      }
  
+     /**
+      * Returns a builder to create virtual threads that use the given executor as scheduler.
+      * @throws UnsupportedOperationException if custom schedulers are not supported
+      */
+     public static Thread.Builder.OfVirtual virtualThreadBuilder(Executor executor) {
+         var scheduler = new Thread.VirtualThreadScheduler() {
+             @Override
+             public void onStart(Thread.VirtualThreadTask task) {
+                 executor.execute(task);
+             }
+             @Override
+             public void onContinue(Thread.VirtualThreadTask task) {
+                 executor.execute(task);
+             }
+         };
+         return virtualThreadBuilder(scheduler);
+     }
+ 
      /**
       * Returns a ThreadFactory to create virtual threads that use the given scheduler.
       * @throws UnsupportedOperationException if custom schedulers are not supported
       */
!     public static ThreadFactory virtualThreadFactory(Thread.VirtualThreadScheduler scheduler) {
          return virtualThreadBuilder(scheduler).factory();
      }
+ 
+     public static ThreadFactory virtualThreadFactory(Executor executor) {
+         return virtualThreadBuilder(executor).factory();
+     }
  }