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