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.Field;
27 import java.lang.reflect.Method;
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 default virtual thread scheduler.
43 */
44 public static Thread.VirtualThreadScheduler defaultScheduler() {
45 try {
46 Method m = Class.forName("java.lang.VirtualThread")
47 .getDeclaredMethod("defaultScheduler");
48 m.setAccessible(true);
49 return (Thread.VirtualThreadScheduler ) m.invoke(null);
50 } catch (InvocationTargetException e) {
51 Throwable cause = e.getCause();
52 if (cause instanceof RuntimeException re) {
53 throw re;
54 }
55 throw new RuntimeException(e);
56 } catch (Exception e) {
57 throw new RuntimeException(e);
58 }
59 }
60
61 /**
62 * Returns the scheduler for the given virtual thread.
63 */
64 public static Thread.VirtualThreadScheduler scheduler(Thread thread) {
65 if (!thread.isVirtual())
66 throw new IllegalArgumentException("Not a virtual thread");
67 try {
68 Field scheduler = Class.forName("java.lang.VirtualThread")
69 .getDeclaredField("scheduler");
70 scheduler.setAccessible(true);
71 return (Thread.VirtualThreadScheduler) scheduler.get(thread);
72 } catch (Exception e) {
73 throw new RuntimeException(e);
74 }
75 }
76
77 /**
78 * Return true if custom schedulers are supported.
79 */
80 public static boolean supportsCustomScheduler() {
81 try (var pool = Executors.newCachedThreadPool()) {
82 try {
83 virtualThreadBuilder(pool).unstarted(() -> { });
84 return true;
85 } catch (UnsupportedOperationException e) {
86 return false;
87 }
88 }
89 }
90
91 /**
92 * Returns a builder to create virtual threads that use the given scheduler.
93 * @throws UnsupportedOperationException if custom schedulers are not supported
94 */
95 public static Thread.Builder.OfVirtual virtualThreadBuilder(Thread.VirtualThreadScheduler scheduler) {
96 Thread.Builder.OfVirtual builder = Thread.ofVirtual();
97 try {
98 Class<?> clazz = Class.forName("java.lang.ThreadBuilders$VirtualThreadBuilder");
99 Field field = clazz.getDeclaredField("scheduler");
100 field.setAccessible(true);
101 field.set(builder, scheduler);
102 return builder;
103 } catch (Exception e) {
104 throw new RuntimeException(e);
105 }
106 }
107
108 /**
109 * Returns a builder to create virtual threads that use the given executor as scheduler.
110 * @throws UnsupportedOperationException if custom schedulers are not supported
111 */
112 public static Thread.Builder.OfVirtual virtualThreadBuilder(Executor executor) {
113 var scheduler = new Thread.VirtualThreadScheduler() {
114 @Override
115 public void onStart(Thread.VirtualThreadTask task) {
116 executor.execute(task);
117 }
118 @Override
119 public void onContinue(Thread.VirtualThreadTask task) {
120 executor.execute(task);
121 }
122 };
123 return virtualThreadBuilder(scheduler);
124 }
125
126 /**
127 * Returns a ThreadFactory to create virtual threads that use the given scheduler.
128 * @throws UnsupportedOperationException if custom schedulers are not supported
129 */
130 public static ThreadFactory virtualThreadFactory(Thread.VirtualThreadScheduler scheduler) {
131 return virtualThreadBuilder(scheduler).factory();
132 }
133
134 public static ThreadFactory virtualThreadFactory(Executor executor) {
135 return virtualThreadBuilder(executor).factory();
136 }
137 }