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