1 /*
  2  * Copyright (c) 2022, 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 package jdk.test.lib.thread;
 25 
 26 import java.lang.management.ManagementFactory;
 27 import java.time.Duration;
 28 import java.util.concurrent.atomic.AtomicReference;
 29 import jdk.management.VirtualThreadSchedulerMXBean;
 30 
 31 /**
 32  * Helper class to support tests running tasks in a 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 an exception.
 45      */
 46     @FunctionalInterface
 47     public interface ThrowingRunnable<X extends Throwable> {
 48         void run() throws X;
 49     }
 50 
 51     /**
 52      * Run a task in a virtual thread and wait for it to terminate.
 53      * If the task completes with an exception then it is thrown by this method.
 54      *
 55      * @param name thread name, can be null
 56      * @param characteristics thread characteristics
 57      * @param task the task to run
 58      * @throws X the exception thrown by the task
 59      */
 60     public static <X extends Throwable> void run(String name,
 61                                                  int characteristics,
 62                                                  ThrowingRunnable<X> task) throws X {
 63         Thread.Builder.OfVirtual builder = Thread.ofVirtual();
 64         if (name != null)
 65             builder.name(name);
 66         if ((characteristics & NO_INHERIT_THREAD_LOCALS) != 0)
 67             builder.inheritInheritableThreadLocals(false);
 68         run(builder, task);
 69     }
 70 
 71     /**
 72      * Run a task in a virtual thread and wait for it to terminate.
 73      * If the task completes with an exception then it is thrown by this method.
 74      *
 75      * @param builder the builder to create the thread
 76      * @param task the task to run
 77      * @throws X the exception thrown by the task
 78      */
 79     public static <X extends Throwable> void run(Thread.Builder.OfVirtual builder,
 80                                                  ThrowingRunnable<X> task) throws X {
 81         var throwableRef = new AtomicReference<Throwable>();
 82         Runnable target = () -> {
 83             try {
 84                 task.run();
 85             } catch (Throwable ex) {
 86                 throwableRef.set(ex);
 87             }
 88         };
 89         Thread thread = builder.start(target);
 90 
 91         // wait for thread to terminate
 92         try {
 93             while (thread.join(Duration.ofSeconds(10)) == false) {
 94                 System.out.println("-- " + thread + " --");
 95                 for (StackTraceElement e : thread.getStackTrace()) {
 96                     System.out.println("  " + e);
 97                 }
 98             }
 99         } catch (InterruptedException e) {
100             throw new RuntimeException(e);
101         }
102 
103         Throwable ex = throwableRef.get();
104         if (ex != null) {
105             if (ex instanceof RuntimeException e)
106                 throw e;
107             if (ex instanceof Error e)
108                 throw e;
109             @SuppressWarnings("unchecked")
110             var x = (X) ex;
111             throw x;
112         }
113     }
114 
115     /**
116      * Run a task in a virtual thread and wait for it to terminate.
117      * If the task completes with an exception then it is thrown by this method.
118      *
119      * @param name thread name, can be null
120      * @param task the task to run
121      * @throws X the exception thrown by the task
122      */
123     public static <X extends Throwable> void run(String name, ThrowingRunnable<X> task) throws X {
124         run(name, 0, task);
125     }
126 
127     /**
128      * Run a task in a virtual thread and wait for it to terminate.
129      * If the task completes with an exception then it is thrown by this method.
130      *
131      * @param characteristics thread characteristics
132      * @param task the task to run
133      * @throws X the exception thrown by the task
134      */
135     public static <X extends Throwable> void run(int characteristics, ThrowingRunnable<X> task) throws X {
136         run(null, characteristics, task);
137     }
138 
139     /**
140      * Run a task in a virtual thread and wait for it to terminate.
141      * If the task completes with an exception then it is thrown by this method.
142      *
143      * @param task the task to run
144      * @throws X the exception thrown by the task
145      */
146     public static <X extends Throwable> void run(ThrowingRunnable<X> task) throws X {
147         run(null, 0, task);
148     }
149 
150     /**
151      * Sets the virtual thread scheduler's target parallelism.
152      *
153      * <p> Tests using this method should use "{@code @modules jdk.management}" to help
154      * test selection.
155      *
156      * @return the previous parallelism level
157      */
158     public static int setParallelism(int size) {
159         var bean = ManagementFactory.getPlatformMXBean(VirtualThreadSchedulerMXBean.class);
160         int parallelism = bean.getParallelism();
161         bean.setParallelism(size);
162         return parallelism;
163     }
164 
165     /**
166      * Ensures that the virtual thread scheduler's target parallelism is at least
167      * the given size. If the target parallelism is less than the given size then
168      * it is changed to the given size.
169      *
170      * <p> Tests using this method should use "{@code @modules jdk.management}" to help
171      * test selection.
172      *
173      * @return the previous parallelism level
174      */
175     public static int ensureParallelism(int size) {
176         var bean = ManagementFactory.getPlatformMXBean(VirtualThreadSchedulerMXBean.class);
177         int parallelism = bean.getParallelism();
178         if (size > parallelism) {
179             bean.setParallelism(size);
180         }
181         return parallelism;
182     }
183 }