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         var throwableRef = new AtomicReference<Throwable>();
 64         Runnable target = () -> {
 65             try {
 66                 task.run();
 67             } catch (Throwable ex) {
 68                 throwableRef.set(ex);
 69             }
 70         };
 71 
 72         Thread.Builder builder = Thread.ofVirtual();
 73         if (name != null)
 74             builder.name(name);
 75         if ((characteristics & NO_INHERIT_THREAD_LOCALS) != 0)
 76             builder.inheritInheritableThreadLocals(false);
 77         Thread thread = builder.start(target);
 78 
 79         // wait for thread to terminate
 80         try {
 81             while (thread.join(Duration.ofSeconds(10)) == false) {
 82                 System.out.println("-- " + thread + " --");
 83                 for (StackTraceElement e : thread.getStackTrace()) {
 84                     System.out.println("  " + e);
 85                 }
 86             }
 87         } catch (InterruptedException e) {
 88             throw new RuntimeException(e);
 89         }
 90 
 91         Throwable ex = throwableRef.get();
 92         if (ex != null) {
 93             if (ex instanceof RuntimeException e)
 94                 throw e;
 95             if (ex instanceof Error e)
 96                 throw e;
 97             @SuppressWarnings("unchecked")
 98             var x = (X) ex;
 99             throw x;
100         }
101     }
102 
103     /**
104      * Run a task in a virtual thread and wait for it to terminate.
105      * If the task completes with an exception then it is thrown by this method.
106      *
107      * @param name thread name, can be null
108      * @param task the task to run
109      * @throws X the exception thrown by the task
110      */
111     public static <X extends Throwable> void run(String name, ThrowingRunnable<X> task) throws X {
112         run(name, 0, task);
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 characteristics thread characteristics
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(int characteristics, ThrowingRunnable<X> task) throws X {
124         run(null, characteristics, 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 task the task to run
132      * @throws X the exception thrown by the task
133      */
134     public static <X extends Throwable> void run(ThrowingRunnable<X> task) throws X {
135         run(null, 0, task);
136     }
137 
138     /**
139      * Sets the virtual thread scheduler's target parallelism.
140      *
141      * <p> Tests using this method should use "{@code @modules jdk.management}" to help
142      * test selection.
143      *
144      * @return the previous parallelism level
145      */
146     public static int setParallelism(int size) {
147         var bean = ManagementFactory.getPlatformMXBean(VirtualThreadSchedulerMXBean.class);
148         int parallelism = bean.getParallelism();
149         bean.setParallelism(size);
150         return parallelism;
151     }
152 
153     /**
154      * Ensures that the virtual thread scheduler's target parallelism is at least
155      * the given size. If the target parallelism is less than the given size then
156      * it is changed to the given size.
157      *
158      * <p> Tests using this method should use "{@code @modules jdk.management}" to help
159      * test selection.
160      *
161      * @return the previous parallelism level
162      */
163     public static int ensureParallelism(int size) {
164         var bean = ManagementFactory.getPlatformMXBean(VirtualThreadSchedulerMXBean.class);
165         int parallelism = bean.getParallelism();
166         if (size > parallelism) {
167             bean.setParallelism(size);
168         }
169         return parallelism;
170     }
171 }