1 /*
2 * Copyright (c) 2020, 2025, 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 /**
25 * @test
26 * @summary Test virtual threads using a custom scheduler
27 * @requires vm.continuations
28 * @modules java.base/java.lang:+open
29 * @library /test/lib
30 * @run junit CustomScheduler
31 */
32
33 import java.lang.reflect.Field;
34 import java.time.Duration;
35 import java.util.ArrayList;
36 import java.util.List;
37 import java.util.concurrent.Executor;
38 import java.util.concurrent.Executors;
39 import java.util.concurrent.ExecutorService;
40 import java.util.concurrent.atomic.AtomicBoolean;
41 import java.util.concurrent.atomic.AtomicInteger;
42 import java.util.concurrent.atomic.AtomicReference;
43 import java.util.concurrent.locks.LockSupport;
44
45 import jdk.test.lib.thread.VThreadRunner;
46 import jdk.test.lib.thread.VThreadScheduler;
47 import org.junit.jupiter.api.Test;
48 import org.junit.jupiter.api.BeforeAll;
49 import org.junit.jupiter.api.AfterAll;
50 import static org.junit.jupiter.api.Assertions.*;
51 import static org.junit.jupiter.api.Assumptions.*;
52
53 class CustomScheduler {
54 private static Thread.VirtualThreadScheduler defaultScheduler;
55 private static ExecutorService threadPool1, threadPool2;
56 private static Thread.VirtualThreadScheduler scheduler1, scheduler2;
57
58 @BeforeAll
59 static void setup() throws Exception {
60 var ref = new AtomicReference<Thread.VirtualThreadScheduler>();
61 Thread thread = Thread.startVirtualThread(() -> {
62 ref.set(currentScheduler());
63 });
64 thread.join();
65 defaultScheduler = ref.get();
66
67 threadPool1 = Executors.newFixedThreadPool(1);
68 threadPool2 = Executors.newFixedThreadPool(1);
69 scheduler1 = adapt(threadPool1);
70 scheduler2 = adapt(threadPool2);
71 }
72
73 @AfterAll
74 static void shutdown() {
75 threadPool1.shutdown();
76 threadPool2.shutdown();
77 }
78
79 /**
80 * Test platform thread creating a virtual thread that uses a custom scheduler.
81 */
82 @Test
83 void testCustomScheduler1() throws Exception {
84 var ref = new AtomicReference<Thread.VirtualThreadScheduler>();
85 Thread thread = VThreadScheduler.virtualThreadBuilder(scheduler1).start(() -> {
86 ref.set(currentScheduler());
87 });
88 thread.join();
89 assertTrue(ref.get() == scheduler1);
90 }
91
92 /**
93 * Test virtual thread creating a virtual thread that uses a custom scheduler.
94 */
95 @Test
96 void testCustomScheduler2() throws Exception {
97 VThreadRunner.run(this::testCustomScheduler1);
98 }
99
100 /**
101 * Test virtual thread using custom scheduler creating a virtual thread that uses
102 * the default scheduler.
103 */
104 @Test
105 void testCustomScheduler3() throws Exception {
106 var ref = new AtomicReference<Thread.VirtualThreadScheduler>();
107 Thread thread = VThreadScheduler.virtualThreadBuilder(scheduler1).start(() -> {
108 try {
109 Thread.ofVirtual().start(() -> {
110 ref.set(currentScheduler());
111 }).join();
112 } catch (Exception e) {
113 e.printStackTrace();
114 }
115 });
116 thread.join();
117 assertTrue(ref.get() == defaultScheduler);
118 }
119
120 /**
121 * Test virtual thread using custom scheduler creating a virtual thread
122 * that uses a different custom scheduler.
123 */
124 @Test
125 void testCustomScheduler4() throws Exception {
126 var ref = new AtomicReference<Thread.VirtualThreadScheduler>();
127 Thread thread1 = VThreadScheduler.virtualThreadBuilder(scheduler1).start(() -> {
128 try {
129 Thread thread2 = VThreadScheduler.virtualThreadBuilder(scheduler2).start(() -> {
130 ref.set(currentScheduler());
131 });
132 thread2.join();
133 } catch (Exception e) {
134 e.printStackTrace();
135 }
136 });
137 thread1.join();
138 assertTrue(ref.get() == scheduler2);
139 }
140
141 /**
142 * Test running task on a virtual thread, should thrown WrongThreadException.
143 */
144 @Test
145 void testBadCarrier() {
146 Executor executor = (task) -> {
147 var exc = new AtomicReference<Throwable>();
148 try {
149 Thread.ofVirtual().start(() -> {
150 try {
151 task.run();
152 fail();
153 } catch (Throwable e) {
154 exc.set(e);
155 }
156 }).join();
157 } catch (InterruptedException e) {
158 fail();
159 }
160 assertTrue(exc.get() instanceof WrongThreadException);
161 };
162 var scheduler = adapt(executor);
163 VThreadScheduler.virtualThreadBuilder(scheduler).start(LockSupport::park);
164 }
165
166 /**
167 * Test parking with the virtual thread interrupt set, should not leak to the
168 * carrier thread when the task completes.
169 */
170 @Test
171 void testParkWithInterruptSet() {
172 Thread carrier = Thread.currentThread();
173 assumeFalse(carrier.isVirtual(), "Main thread is a virtual thread");
174 try {
175 var scheduler = adapt(Runnable::run);
176 Thread vthread = VThreadScheduler.virtualThreadBuilder(scheduler).start(() -> {
177 Thread.currentThread().interrupt();
178 Thread.yield();
179 });
180 assertTrue(vthread.isInterrupted());
181 assertFalse(carrier.isInterrupted());
182 } finally {
183 Thread.interrupted();
184 }
185 }
186
187 /**
188 * Test terminating with the virtual thread interrupt set, should not leak to
189 * the carrier thread when the task completes.
190 */
191 @Test
192 void testTerminateWithInterruptSet() {
193 Thread carrier = Thread.currentThread();
194 assumeFalse(carrier.isVirtual(), "Main thread is a virtual thread");
195 try {
196 var scheduler = adapt(Runnable::run);
197 Thread vthread = VThreadScheduler.virtualThreadBuilder(scheduler).start(() -> {
198 Thread.currentThread().interrupt();
199 });
200 assertTrue(vthread.isInterrupted());
201 assertFalse(carrier.isInterrupted());
202 } finally {
203 Thread.interrupted();
204 }
205 }
206
207 /**
208 * Test running task with the carrier's interrupted status set.
209 */
210 @Test
211 void testRunWithInterruptSet() throws Exception {
212 assumeFalse(Thread.currentThread().isVirtual(), "Main thread is a virtual thread");
213 var scheduler = adapt(task -> {
214 Thread.currentThread().interrupt();
215 task.run();
216 });
217 try {
218 AtomicBoolean interrupted = new AtomicBoolean();
219 Thread vthread = VThreadScheduler.virtualThreadBuilder(scheduler).start(() -> {
220 interrupted.set(Thread.currentThread().isInterrupted());
221 });
222 assertFalse(vthread.isInterrupted());
223 } finally {
224 Thread.interrupted();
225 }
226 }
227
228 /**
229 * Test custom scheduler throwing OOME when starting a thread.
230 */
231 @Test
232 void testThreadStartOOME() throws Exception {
233 var scheduler = adapt(task -> {
234 System.err.println("OutOfMemoryError");
235 throw new OutOfMemoryError();
236 });
237 Thread thread = VThreadScheduler.virtualThreadBuilder(scheduler).unstarted(() -> { });
238 assertThrows(OutOfMemoryError.class, thread::start);
239 }
240
241 /**
242 * Test custom scheduler throwing OOME when unparking a thread.
243 */
244 @Test
245 void testThreadUnparkOOME() throws Exception {
246 try (ExecutorService executor = Executors.newFixedThreadPool(1)) {
247 AtomicInteger counter = new AtomicInteger();
248 var scheduler = adapt(task -> {
249 switch (counter.getAndIncrement()) {
250 case 0 -> executor.execute(task); // Thread.start
251 case 1, 2 -> { // unpark attempt 1+2
252 System.err.println("OutOfMemoryError");
253 throw new OutOfMemoryError();
254 }
255 default -> executor.execute(task);
256 }
257 executor.execute(task);
258 });
259
260 // start thread and wait for it to park
261 var thread = VThreadScheduler.virtualThreadBuilder(scheduler).start(LockSupport::park);
262 await(thread, Thread.State.WAITING);
263
264 // unpark thread, this should retry until OOME is not thrown
265 LockSupport.unpark(thread);
266 thread.join();
267 }
268 }
269
270 /**
271 * Returns the scheduler for the current virtual thread.
272 */
273 private static Thread.VirtualThreadScheduler currentScheduler() {
274 return VThreadScheduler.scheduler(Thread.currentThread());
275 }
276
277 private static Thread.VirtualThreadScheduler adapt(Executor executor) {
278 return new Thread.VirtualThreadScheduler() {
279 @Override
280 public void onStart(Thread.VirtualThreadTask task) {
281 executor.execute(task);
282 }
283 @Override
284 public void onContinue(Thread.VirtualThreadTask task) {
285 executor.execute(task);
286 }
287 };
288 }
289
290 /**
291 * Waits for the given thread to reach a given state.
292 */
293 private static void await(Thread thread, Thread.State expectedState) throws InterruptedException {
294 Thread.State state = thread.getState();
295 while (state != expectedState) {
296 assertTrue(state != Thread.State.TERMINATED, "Thread has terminated");
297 Thread.sleep(10);
298 state = thread.getState();
299 }
300 }
301 }