1 /*
  2  * Copyright (c) 2020, 2023, 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  * @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.*;
 37 import java.util.concurrent.atomic.AtomicBoolean;
 38 import java.util.concurrent.atomic.AtomicReference;
 39 import java.util.concurrent.locks.LockSupport;
 40 
 41 import org.junit.jupiter.api.Test;
 42 import org.junit.jupiter.api.BeforeAll;
 43 import org.junit.jupiter.api.AfterAll;
 44 import static org.junit.jupiter.api.Assertions.*;
 45 import static org.junit.jupiter.api.Assumptions.*;
 46 
 47 class CustomScheduler {
 48     private static ExecutorService scheduler1;
 49     private static ExecutorService scheduler2;
 50 
 51     @BeforeAll
 52     static void setup() {
 53         scheduler1 = Executors.newFixedThreadPool(1);
 54         scheduler2 = Executors.newFixedThreadPool(1);
 55     }
 56 
 57     @AfterAll
 58     static void shutdown() {
 59         scheduler1.shutdown();
 60         scheduler2.shutdown();
 61     }
 62 
 63     /**
 64      * Test platform thread creating a virtual thread that uses a custom scheduler.
 65      */
 66     @Test
 67     void testCustomScheduler1() throws Exception {
 68         AtomicReference<Executor> ref = new AtomicReference<>();
 69         ThreadBuilders.virtualThreadBuilder(scheduler1).start(() -> {
 70             ref.set(scheduler(Thread.currentThread()));
 71         }).join();
 72         assertTrue(ref.get() == scheduler1);
 73     }
 74 
 75     /**
 76      * Test virtual thread creating a virtual thread that uses a custom scheduler.
 77      */
 78     @Test
 79     void testCustomScheduler2() throws Exception {
 80         AtomicReference<Executor> ref = new AtomicReference<>();
 81         Thread.ofVirtual().start(() -> {
 82             try {
 83                 ThreadBuilders.virtualThreadBuilder(scheduler1).start(() -> {
 84                     ref.set(scheduler(Thread.currentThread()));
 85                 }).join();
 86             } catch (Exception e) {
 87                 e.printStackTrace();
 88             }
 89         }).join();
 90         assertTrue(ref.get() == scheduler1);
 91     }
 92 
 93     /**
 94      * Test virtual thread using custom scheduler creating a virtual thread.
 95      * The scheduler should be inherited.
 96      */
 97     @Test
 98     void testCustomScheduler3() throws Exception {
 99         AtomicReference<Executor> ref = new AtomicReference<>();
100         ThreadBuilders.virtualThreadBuilder(scheduler1).start(() -> {
101             try {
102                 Thread.ofVirtual().start(() -> {
103                     ref.set(scheduler(Thread.currentThread()));
104                 }).join();
105             } catch (Exception e) {
106                 e.printStackTrace();
107             }
108         }).join();
109         assertTrue(ref.get() == scheduler1);
110     }
111 
112     /**
113      * Test virtual thread using custom scheduler creating a virtual thread
114      * that uses a different custom scheduler.
115      */
116     @Test
117     void testCustomScheduler4() throws Exception {
118         AtomicReference<Executor> ref = new AtomicReference<>();
119         ThreadBuilders.virtualThreadBuilder(scheduler1).start(() -> {
120             try {
121                 ThreadBuilders.virtualThreadBuilder(scheduler2).start(() -> {
122                     ref.set(scheduler(Thread.currentThread()));
123                 }).join();
124             } catch (Exception e) {
125                 e.printStackTrace();
126             }
127         }).join();
128         assertTrue(ref.get() == scheduler2);
129     }
130 
131     /**
132      * Test running task on a virtual thread, should thrown WrongThreadException.
133      */
134     @Test
135     void testBadCarrier() {
136         Executor scheduler = (task) -> {
137             var exc = new AtomicReference<Throwable>();
138             try {
139                 Thread.ofVirtual().start(() -> {
140                     try {
141                         task.run();
142                         fail();
143                     } catch (Throwable e) {
144                         exc.set(e);
145                     }
146                 }).join();
147             } catch (InterruptedException e) {
148                 fail();
149             }
150             assertTrue(exc.get() instanceof WrongThreadException);
151         };
152 
153         ThreadBuilders.virtualThreadBuilder(scheduler).start(LockSupport::park);
154     }
155 
156     /**
157      * Test parking with the virtual thread interrupt set, should not leak to the
158      * carrier thread when the task completes.
159      */
160     @Test
161     void testParkWithInterruptSet() {
162         Thread carrier = Thread.currentThread();
163         assumeFalse(carrier.isVirtual(), "Main thread is a virtual thread");
164         try {
165             var builder = ThreadBuilders.virtualThreadBuilder(Runnable::run);
166             Thread vthread = builder.start(() -> {
167                 Thread.currentThread().interrupt();
168                 Thread.yield();
169             });
170             assertTrue(vthread.isInterrupted());
171             assertFalse(carrier.isInterrupted());
172         } finally {
173             Thread.interrupted();
174         }
175     }
176 
177     /**
178      * Test terminating with the virtual thread interrupt set, should not leak to
179      * the carrier thread when the task completes.
180      */
181     @Test
182     void testTerminateWithInterruptSet() {
183         Thread carrier = Thread.currentThread();
184         assumeFalse(carrier.isVirtual(), "Main thread is a virtual thread");
185         try {
186             var builder = ThreadBuilders.virtualThreadBuilder(Runnable::run);
187             Thread vthread = builder.start(() -> {
188                 Thread.currentThread().interrupt();
189             });
190             assertTrue(vthread.isInterrupted());
191             assertFalse(carrier.isInterrupted());
192         } finally {
193             Thread.interrupted();
194         }
195     }
196 
197     /**
198      * Test running task with the carrier interrupt status set.
199      */
200     @Test
201     void testRunWithInterruptSet() throws Exception {
202         assumeFalse(Thread.currentThread().isVirtual(), "Main thread is a virtual thread");
203         Executor scheduler = (task) -> {
204             Thread.currentThread().interrupt();
205             task.run();
206         };
207         try {
208             AtomicBoolean interrupted = new AtomicBoolean();
209             Thread vthread = ThreadBuilders.virtualThreadBuilder(scheduler).start(() -> {
210                 interrupted.set(Thread.currentThread().isInterrupted());
211             });
212             assertFalse(vthread.isInterrupted());
213         } finally {
214             Thread.interrupted();
215         }
216     }
217 
218     /**
219      * Returns the scheduler for the given virtual thread.
220      */
221     private static Executor scheduler(Thread thread) {
222         if (!thread.isVirtual())
223             throw new IllegalArgumentException("Not a virtual thread");
224         try {
225             Field scheduler = Class.forName("java.lang.VirtualThread")
226                     .getDeclaredField("scheduler");
227             scheduler.setAccessible(true);
228             return (Executor) scheduler.get(thread);
229         } catch (Exception e) {
230             throw new RuntimeException(e);
231         }
232     }
233 }