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