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
38 import java.util.concurrent.Executor;
39 import java.util.concurrent.Executors;
40 import java.util.concurrent.ExecutorService;
41 import java.util.concurrent.ThreadFactory;
42 import java.util.concurrent.atomic.AtomicBoolean;
43 import java.util.concurrent.atomic.AtomicInteger;
44 import java.util.concurrent.atomic.AtomicReference;
45 import java.util.concurrent.locks.LockSupport;
46 import java.util.concurrent.CountDownLatch;
47
48 import jdk.test.lib.thread.VThreadRunner;
49 import jdk.test.lib.thread.VThreadScheduler;
50 import org.junit.jupiter.api.Test;
51 import org.junit.jupiter.api.BeforeAll;
52 import org.junit.jupiter.api.AfterAll;
53 import static org.junit.jupiter.api.Assertions.*;
54 import static org.junit.jupiter.api.Assumptions.*;
55
56 class CustomScheduler {
57 private static Thread.VirtualThreadScheduler defaultScheduler;
58 private static ExecutorService threadPool1, threadPool2;
59 private static Thread.VirtualThreadScheduler scheduler1, scheduler2;
60
61 @BeforeAll
62 static void setup() throws Exception {
63 var ref = new AtomicReference<Thread.VirtualThreadScheduler>();
64 Thread thread = Thread.startVirtualThread(() -> {
65 ref.set(currentScheduler());
66 });
67 thread.join();
68 defaultScheduler = ref.get();
69
70 threadPool1 = Executors.newFixedThreadPool(1);
71 threadPool2 = Executors.newFixedThreadPool(1);
72 scheduler1 = adapt(threadPool1);
73 scheduler2 = adapt(threadPool2);
74 }
75
76 @AfterAll
77 static void shutdown() {
78 threadPool1.shutdown();
79 threadPool2.shutdown();
80 }
81
82 /**
83 * Test platform thread creating a virtual thread that uses a custom scheduler.
84 */
85 @Test
86 void testCustomScheduler1() throws Exception {
87 var ref = new AtomicReference<Thread.VirtualThreadScheduler>();
88 Thread thread = VThreadScheduler.virtualThreadBuilder(scheduler1).start(() -> {
89 ref.set(currentScheduler());
90 });
91 thread.join();
92 assertTrue(ref.get() == scheduler1);
93 }
94
95 /**
96 * Test virtual thread creating a virtual thread that uses a custom scheduler.
97 */
98 @Test
99 void testCustomScheduler2() throws Exception {
100 VThreadRunner.run(this::testCustomScheduler1);
101 }
102
103 /**
104 * Test virtual thread using custom scheduler creating a virtual thread that uses
105 * the default scheduler.
106 */
107 @Test
108 void testCustomScheduler3() throws Exception {
109 var ref = new AtomicReference<Thread.VirtualThreadScheduler>();
110 Thread thread = VThreadScheduler.virtualThreadBuilder(scheduler1).start(() -> {
111 try {
112 Thread.ofVirtual().start(() -> {
113 ref.set(currentScheduler());
114 }).join();
115 } catch (Exception e) {
116 e.printStackTrace();
117 }
118 });
119 thread.join();
120 assertTrue(ref.get() == defaultScheduler);
121 }
122
123 /**
124 * Test virtual thread using custom scheduler creating a virtual thread
125 * that uses a different custom scheduler.
126 */
127 @Test
128 void testCustomScheduler4() throws Exception {
129 var ref = new AtomicReference<Thread.VirtualThreadScheduler>();
130 Thread thread1 = VThreadScheduler.virtualThreadBuilder(scheduler1).start(() -> {
131 try {
132 Thread thread2 = VThreadScheduler.virtualThreadBuilder(scheduler2).start(() -> {
133 ref.set(currentScheduler());
134 });
135 thread2.join();
136 } catch (Exception e) {
137 e.printStackTrace();
138 }
139 });
140 thread1.join();
141 assertTrue(ref.get() == scheduler2);
142 }
143
144 /**
145 * Test running task on a virtual thread, should thrown WrongThreadException.
146 */
147 @Test
148 void testBadCarrier() {
149 Executor executor = (task) -> {
150 var exc = new AtomicReference<Throwable>();
151 try {
152 Thread.ofVirtual().start(() -> {
153 try {
154 task.run();
155 fail();
156 } catch (Throwable e) {
157 exc.set(e);
158 }
159 }).join();
160 } catch (InterruptedException e) {
161 fail();
162 }
163 assertTrue(exc.get() instanceof WrongThreadException);
164 };
165 var scheduler = adapt(executor);
166 VThreadScheduler.virtualThreadBuilder(scheduler).start(LockSupport::park);
167 }
168
169 /**
170 * Test parking with the virtual thread interrupt set, should not leak to the
171 * carrier thread when the task completes.
172 */
173 @Test
174 void testParkWithInterruptSet() {
175 Thread carrier = Thread.currentThread();
176 assumeFalse(carrier.isVirtual(), "Main thread is a virtual thread");
177 try {
178 var scheduler = adapt(Runnable::run);
179 Thread vthread = VThreadScheduler.virtualThreadBuilder(scheduler).start(() -> {
180 Thread.currentThread().interrupt();
181 Thread.yield();
182 });
183 assertTrue(vthread.isInterrupted());
184 assertFalse(carrier.isInterrupted());
185 } finally {
186 Thread.interrupted();
187 }
188 }
189
190 /**
191 * Test terminating with the virtual thread interrupt set, should not leak to
192 * the carrier thread when the task completes.
193 */
194 @Test
195 void testTerminateWithInterruptSet() {
196 Thread carrier = Thread.currentThread();
197 assumeFalse(carrier.isVirtual(), "Main thread is a virtual thread");
198 try {
199 var scheduler = adapt(Runnable::run);
200 Thread vthread = VThreadScheduler.virtualThreadBuilder(scheduler).start(() -> {
201 Thread.currentThread().interrupt();
202 });
203 assertTrue(vthread.isInterrupted());
204 assertFalse(carrier.isInterrupted());
205 } finally {
206 Thread.interrupted();
207 }
208 }
209
210 /**
211 * Test running task with the carrier's interrupted status set.
212 */
213 @Test
214 void testRunWithInterruptSet() throws Exception {
215 assumeFalse(Thread.currentThread().isVirtual(), "Main thread is a virtual thread");
216 var scheduler = adapt(task -> {
217 Thread.currentThread().interrupt();
218 task.run();
219 });
220 try {
221 AtomicBoolean interrupted = new AtomicBoolean();
222 Thread vthread = VThreadScheduler.virtualThreadBuilder(scheduler).start(() -> {
223 interrupted.set(Thread.currentThread().isInterrupted());
224 });
225 assertFalse(vthread.isInterrupted());
226 } finally {
227 Thread.interrupted();
228 }
229 }
230
231 /**
232 * Test custom scheduler throwing OOME when starting a thread.
233 */
234 @Test
235 void testThreadStartOOME() throws Exception {
236 var scheduler = adapt(task -> {
237 System.err.println("OutOfMemoryError");
238 throw new OutOfMemoryError();
239 });
240 Thread thread = VThreadScheduler.virtualThreadBuilder(scheduler).unstarted(() -> { });
241 assertThrows(OutOfMemoryError.class, thread::start);
242 }
243
244 /**
245 * Test custom scheduler throwing OOME when unparking a thread.
246 */
247 @Test
248 void testThreadUnparkOOME() throws Exception {
249 try (ExecutorService executor = Executors.newFixedThreadPool(1)) {
250 AtomicInteger counter = new AtomicInteger();
251 var scheduler = adapt(task -> {
252 switch (counter.getAndIncrement()) {
253 case 0 -> executor.execute(task); // Thread.start
254 case 1, 2 -> { // unpark attempt 1+2
255 System.err.println("OutOfMemoryError");
256 throw new OutOfMemoryError();
257 }
258 default -> executor.execute(task);
259 }
260 executor.execute(task);
261 });
262
263 // start thread and wait for it to park
264 var thread = VThreadScheduler.virtualThreadBuilder(scheduler).start(LockSupport::park);
265 await(thread, Thread.State.WAITING);
266
267 // unpark thread, this should retry until OOME is not thrown
268 LockSupport.unpark(thread);
269 thread.join();
270 }
271 }
272
273 /**
274 * Returns the scheduler for the current virtual thread.
275 */
276 private static Thread.VirtualThreadScheduler currentScheduler() {
277 return VThreadScheduler.scheduler(Thread.currentThread());
278 }
279
280 private static Thread.VirtualThreadScheduler adapt(Executor executor) {
281 return new Thread.VirtualThreadScheduler() {
282 @Override
283 public void onStart(Thread.VirtualThreadTask task) {
284 executor.execute(task);
285 }
286 @Override
287 public void onContinue(Thread.VirtualThreadTask task) {
288 executor.execute(task);
289 }
290 };
291 }
292
293 /**
294 * Test virtual thread mounted on carrier executing class initializer
295 */
296 @Test
297 void testMountedAtKlassInit() throws Exception {
298 assumeFalse(Thread.currentThread().isVirtual(), "Main thread is a virtual thread");
299 class CarrierHelper {
300 static boolean passed;
301 static {
302 try {
303 Carrier.testEmpty();
304 Carrier.testParking();
305 Carrier.testMonitorEnter();
306 passed = true;
307 } catch (Exception e) {}
308 }
309 }
310 assertTrue(CarrierHelper.passed);
311 }
312
313 class Carrier {
314 static ThreadFactory factory = VThreadScheduler.virtualThreadBuilder(r -> r.run()).factory();
315
316 static void testEmpty() throws Exception {
317 Thread vthread = factory.newThread(() -> {});
318 vthread.start();
319 vthread.join();
320 }
321
322 static void testParking() throws Exception {
323 Thread vthread = factory.newThread(LockSupport::park);
324 vthread.start();
325 assertTrue(vthread.getState() == Thread.State.WAITING);
326 LockSupport.unpark(vthread);
327 vthread.join();
328 }
329
330 static void testMonitorEnter() throws Exception {
331 Object lock = new Object();
332 AtomicBoolean done = new AtomicBoolean();
333 AtomicInteger counter = new AtomicInteger();
334 CountDownLatch started = new CountDownLatch(1);
335
336 Thread contender = Thread.ofPlatform().start(() -> {
337 synchronized (lock) {
338 started.countDown();
339 while (!done.get()) {}
340 }
341 });
342 started.await();
343 Thread vthread = factory.newThread(() -> {
344 synchronized (lock) {
345 counter.getAndIncrement();
346 }
347 });
348 vthread.start();
349 assertTrue(vthread.getState() == Thread.State.BLOCKED);
350
351 done.set(true);
352 contender.join();
353 vthread.join();
354 }
355 }
356
357 /**
358 * Waits for the given thread to reach a given state.
359 */
360 private static void await(Thread thread, Thread.State expectedState) throws InterruptedException {
361 Thread.State state = thread.getState();
362 while (state != expectedState) {
363 assertTrue(state != Thread.State.TERMINATED, "Thread has terminated");
364 Thread.sleep(10);
365 state = thread.getState();
366 }
367 }
368 }