1 /*
  2  * Copyright (c) 2022, 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 package jdk.test.lib.thread;
 25 
 26 import java.lang.reflect.Field;
 27 import java.time.Duration;
 28 import java.util.concurrent.ForkJoinPool;
 29 import java.util.concurrent.atomic.AtomicReference;
 30 
 31 /**
 32  * Helper class to support tests running tasks a in virtual thread.
 33  */
 34 public class VThreadRunner {
 35     private VThreadRunner() { }
 36 
 37     /**
 38      * Characteristic value signifying that initial values for inheritable
 39      * thread locals are not inherited from the constructing thread.
 40      */
 41     public static final int NO_INHERIT_THREAD_LOCALS = 1 << 2;
 42 
 43     /**
 44      * Represents a task that does not return a result but may throw
 45      * an exception.
 46      */
 47     @FunctionalInterface
 48     public interface ThrowingRunnable {
 49         /**
 50          * Runs this operation.
 51          */
 52         void run() throws Exception;
 53     }
 54 
 55     /**
 56      * Run a task in a virtual thread and wait for it to terminate.
 57      * If the task completes with an exception then it is thrown by this method.
 58      * If the task throws an Error then it is wrapped in an RuntimeException.
 59      *
 60      * @param name thread name, can be null
 61      * @param characteristics thread characteristics
 62      * @param task the task to run
 63      * @throws Exception the exception thrown by the task
 64      */
 65     public static void run(String name,
 66                            int characteristics,
 67                            ThrowingRunnable task) throws Exception {
 68         AtomicReference<Exception> exc = new AtomicReference<>();
 69         Runnable target =  () -> {
 70             try {
 71                 task.run();
 72             } catch (Error e) {
 73                 exc.set(new RuntimeException(e));
 74             } catch (Exception e) {
 75                 exc.set(e);
 76             }
 77         };
 78 
 79         Thread.Builder builder = Thread.ofVirtual();
 80         if (name != null)
 81             builder.name(name);
 82         if ((characteristics & NO_INHERIT_THREAD_LOCALS) != 0)
 83             builder.inheritInheritableThreadLocals(false);
 84         Thread thread = builder.start(target);
 85 
 86         // wait for thread to terminate
 87         while (thread.join(Duration.ofSeconds(10)) == false) {
 88             System.out.println("-- " + thread + " --");
 89             for (StackTraceElement e : thread.getStackTrace()) {
 90                 System.out.println("  " + e);
 91             }
 92         }
 93 
 94         Exception e = exc.get();
 95         if (e != null) {
 96             throw e;
 97         }
 98     }
 99 
100     /**
101      * Run a task in a virtual thread and wait for it to terminate.
102      * If the task completes with an exception then it is thrown by this method.
103      * If the task throws an Error then it is wrapped in an RuntimeException.
104      *
105      * @param name thread name, can be null
106      * @param task the task to run
107      * @throws Exception the exception thrown by the task
108      */
109     public static void run(String name, ThrowingRunnable task) throws Exception {
110         run(name, 0, task);
111     }
112 
113     /**
114      * Run a task in a virtual thread and wait for it to terminate.
115      * If the task completes with an exception then it is thrown by this method.
116      * If the task throws an Error then it is wrapped in an RuntimeException.
117      *
118      * @param characteristics thread characteristics
119      * @param task the task to run
120      * @throws Exception the exception thrown by the task
121      */
122     public static void run(int characteristics, ThrowingRunnable task) throws Exception {
123         run(null, characteristics, task);
124     }
125 
126     /**
127      * Run a task in a virtual thread and wait for it to terminate.
128      * If the task completes with an exception then it is thrown by this method.
129      * If the task throws an Error then it is wrapped in an RuntimeException.
130      *
131      * @param task the task to run
132      * @throws Exception the exception thrown by the task
133      */
134     public static void run(ThrowingRunnable task) throws Exception {
135         run(null, 0, task);
136     }
137 
138     /**
139      * Returns the virtual thread scheduler.
140      */
141     private static ForkJoinPool defaultScheduler() {
142         try {
143             var clazz = Class.forName("java.lang.VirtualThread");
144             var field = clazz.getDeclaredField("DEFAULT_SCHEDULER");
145             field.setAccessible(true);
146             return (ForkJoinPool) field.get(null);
147         } catch (Exception e) {
148             throw new RuntimeException(e);
149         }
150     }
151 
152     /**
153      * Sets the virtual thread scheduler's target parallelism.
154      * @return the previous parallelism level
155      */
156     public static int setParallelism(int size) {
157         return defaultScheduler().setParallelism(size);
158     }
159 
160     /**
161      * Ensures that the virtual thread scheduler's target parallelism is at least
162      * the given size. If the target parallelism is less than the given size then
163      * it is changed to the given size.
164      * @return the previous parallelism level
165      */
166     public static int ensureParallelism(int size) {
167         ForkJoinPool pool = defaultScheduler();
168         int parallelism = pool.getParallelism();
169         if (size > parallelism) {
170             pool.setParallelism(size);
171         }
172         return parallelism;
173     }
174 }