1 /*
 2  * Copyright (c) 2019, 2024, Oracle and/or its affiliates. All rights reserved.
 3  * DO NOT ALTER OR REMOVE COPYRIGHT NOTICES OR THIS FILE HEADER.
 4  *
 5  * This code is free software; you can redistribute it and/or modify it
 6  * under the terms of the GNU General Public License version 2 only, as
 7  * published by the Free Software Foundation.
 8  *
 9  * This code is distributed in the hope that it will be useful, but WITHOUT
10  * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
11  * FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public License
12  * version 2 for more details (a copy is included in the LICENSE file that
13  * accompanied this code).
14  *
15  * You should have received a copy of the GNU General Public License version
16  * 2 along with this work; if not, write to the Free Software Foundation,
17  * Inc., 51 Franklin St, Fifth Floor, Boston, MA 02110-1301 USA.
18  *
19  * Please contact Oracle, 500 Oracle Parkway, Redwood Shores, CA 94065 USA
20  * or visit www.oracle.com if you need additional information or have any
21  * questions.
22  */
23 
24 package jdk.test.lib.thread;
25 
26 import java.lang.reflect.Constructor;
27 import java.lang.reflect.Field;
28 import java.lang.reflect.InvocationTargetException;
29 import java.util.concurrent.Executor;
30 import java.util.concurrent.Executors;
31 import java.util.concurrent.ThreadFactory;
32 
33 /**
34  * Helper class to allow tests run virtual threads with a custom scheduler.
35  *
36  * Tests using this class need to open java.base/java.lang.
37  */
38 public class VThreadScheduler {
39     private VThreadScheduler() { }
40 
41     /**
42      * Returns the scheduler for the given virtual thread.
43      */
44     public static Executor scheduler(Thread thread) {
45         if (!thread.isVirtual())
46             throw new IllegalArgumentException("Not a virtual thread");
47         try {
48             Field scheduler = Class.forName("java.lang.VirtualThread")
49                     .getDeclaredField("scheduler");
50             scheduler.setAccessible(true);
51             return (Executor) scheduler.get(thread);
52         } catch (Exception e) {
53             throw new RuntimeException(e);
54         }
55     }
56 
57     /**
58      * Return true if custom schedulers are supported.
59      */
60     public static boolean supportsCustomScheduler() {
61         try (var pool = Executors.newCachedThreadPool()) {
62             try {
63                 virtualThreadBuilder(pool);
64                 return true;
65             } catch (UnsupportedOperationException e) {
66                 return false;
67             }
68         }
69     }
70 
71     /**
72      * Returns a builder to create virtual threads that use the given scheduler.
73      * @throws UnsupportedOperationException if custom schedulers are not supported
74      */
75     public static Thread.Builder.OfVirtual virtualThreadBuilder(Executor scheduler) {
76         try {
77             Class<?> clazz = Class.forName("java.lang.ThreadBuilders$VirtualThreadBuilder");
78             Constructor<?> ctor = clazz.getDeclaredConstructor(Executor.class);
79             ctor.setAccessible(true);
80             return (Thread.Builder.OfVirtual) ctor.newInstance(scheduler);
81         } catch (InvocationTargetException e) {
82             Throwable cause = e.getCause();
83             if (cause instanceof RuntimeException re) {
84                 throw re;
85             }
86             throw new RuntimeException(e);
87         } catch (Exception e) {
88             throw new RuntimeException(e);
89         }
90     }
91 
92     /**
93      * Returns a ThreadFactory to create virtual threads that use the given scheduler.
94      * @throws UnsupportedOperationException if custom schedulers are not supported
95      */
96     public static ThreadFactory virtualThreadFactory(Executor scheduler) {
97         return virtualThreadBuilder(scheduler).factory();
98     }
99 }