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

@@ -21,12 +21,12 @@
   * questions.
   */
  
  package jdk.test.lib.thread;
  
- import java.lang.reflect.Constructor;
  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 +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 Executor scheduler(Thread 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 (Executor) scheduler.get(thread);
+             return (Thread.VirtualThreadScheduler) scheduler.get(thread);
          } catch (Exception e) {
              throw new RuntimeException(e);
          }
      }
  

@@ -70,30 +90,36 @@
  
      /**
       * 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;
+     @SuppressWarnings("restricted")
+     public static Thread.Builder.OfVirtual virtualThreadBuilder(Thread.VirtualThreadScheduler scheduler) {
+         return Thread.ofVirtual().scheduler(scheduler);
+     }
+ 
+     public static Thread.Builder.OfVirtual virtualThreadBuilder(Executor executor) {
+         var scheduler = new Thread.VirtualThreadScheduler() {
+             @Override
+             public void onStart(Thread.VirtualThreadTask task) {
+                 executor.execute(task);
              }
-             throw new RuntimeException(e);
-         } catch (Exception e) {
-             throw new RuntimeException(e);
-         }
+             @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(Executor scheduler) {
+     public static ThreadFactory virtualThreadFactory(Thread.VirtualThreadScheduler scheduler) {
          return virtualThreadBuilder(scheduler).factory();
      }
+ 
+     public static ThreadFactory virtualThreadFactory(Executor executor) {
+         return virtualThreadBuilder(executor).factory();
+     }
  }