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 }
--- EOF ---