Allow bytecode transformers to run in parallel when using shamrock:run

This drops startup time from 17s to 12s in the 1k entity test on my laptop

core/deployment/src/main/java/org/jboss/shamrock/deployment/BuildTimeGenerator.java

@@ -34,6 +34,7 @@
 import java.util.ServiceLoader;
 import java.util.Set;
 import java.util.concurrent.atomic.AtomicInteger;
+import java.util.function.BiFunction;
 import java.util.function.Function;
 import java.util.logging.Level;
 import java.util.logging.Logger;
@@ -86,7 +87,7 @@ public class BuildTimeGenerator {
     private final DeploymentProcessorInjection injection;
     private final ClassLoader classLoader;
     private final boolean useStaticInit;
-    private final List<Function<String, Function<ClassVisitor, ClassVisitor>>> bytecodeTransformers = new ArrayList<>();
+    private final Map<String, List<BiFunction<String, ClassVisitor, ClassVisitor>>> byteCodeTransformers = new HashMap<>();
     private final Set<String> applicationArchiveMarkers;
     private final ArchiveContextBuilder archiveContextBuilder;
     private final Set<String> capabilities;
@@ -110,8 +111,8 @@ public BuildTimeGenerator(ClassOutput classOutput, ClassLoader cl, boolean useSt
         this.capabilities = new HashSet<>(setupContext.capabilities);
     }
 
-    public List<Function<String, Function<ClassVisitor, ClassVisitor>>> getBytecodeTransformers() {
-        return bytecodeTransformers;
+    public Map<String, List<BiFunction<String, ClassVisitor, ClassVisitor>>> getByteCodeTransformers() {
+        return byteCodeTransformers;
     }
 
     public void run(Path root) throws IOException {
@@ -342,8 +343,8 @@ public void createResource(String name, byte[] data) throws IOException {
         }
 
         @Override
-        public void addByteCodeTransformer(Function<String, Function<ClassVisitor, ClassVisitor>> visitorFunction) {
-            bytecodeTransformers.add(visitorFunction);
+        public void addByteCodeTransformer(String className, BiFunction<String, ClassVisitor, ClassVisitor> visitorFunction) {
+            byteCodeTransformers.computeIfAbsent(className, (e) -> new ArrayList<>()).add(visitorFunction);
         }
 
         @Override
@@ -438,7 +439,6 @@ void writeMainClass() throws IOException {
                 ResultHandle dup = mv.newInstance(ofConstructor(holder.className));
                 mv.invokeInterfaceMethod(ofMethod(StartupTask.class, "deploy", void.class, StartupContext.class), dup, startupContext);
             }
-
             mv.invokeStaticMethod(ofMethod(Timing.class, "printStartupTime", void.class));
             mv.returnValue(null);
 

core/deployment/src/main/java/org/jboss/shamrock/deployment/ProcessorContext.java

@@ -3,6 +3,7 @@
 import java.io.IOException;
 import java.lang.reflect.Field;
 import java.lang.reflect.Method;
+import java.util.function.BiFunction;
 import java.util.function.Function;
 
 import org.jboss.jandex.FieldInfo;
@@ -41,9 +42,9 @@ public interface ProcessorContext {
      * This method is used to indicate that a given class requires reflection.
      * <p>
      * It is used in the graal output to allow the class to be reflected when running on substrate VM
-     *
+     * <p>
      * This will add all constructors, as well as all methods and fields if the appropriate fields are set.
-     *
+     * <p>
      * Where possible consider using the more fine grained addReflective* variants
      *
      * @param className The class name
@@ -60,56 +61,56 @@ public interface ProcessorContext {
 
     /**
      * Attempts to register a complete type hierarchy for reflection.
-     *
+     * <p>
      * This is intended to be used to register types that are going to be serialized,
      * e.g. by Jackson or some other JSON mapper.
-     *
+     * <p>
      * This will do 'smart discovery' and in addition to registering the type itself it will also attempt to
      * register the following:
-     *
+     * <p>
      * - Superclasses
      * - Component types of collections
      * - Types used in bean properties if (if method reflection is enabled)
      * - Field types (if field reflection is enabled)
-     *
+     * <p>
      * This discovery is applied recursively, so any additional types that are registered will also have their dependencies
      * discovered
-     *
      */
     void addReflectiveHierarchy(Type type);
 
 
     /**
-     *
      * @param applicationClass If this class should be loaded by the application class loader when in runtime mode
-     * @param name The class name
-     * @param classData The class bytes
+     * @param name             The class name
+     * @param classData        The class bytes
      * @throws IOException
      */
     void addGeneratedClass(boolean applicationClass, String name, byte[] classData) throws IOException;
 
     /**
      * Creates a resources with the provided contents
+     *
      * @param name
      * @param data
      * @throws IOException
      */
     void createResource(String name, byte[] data) throws IOException;
 
     /**
-     * Adds a bytecode transformer that can transform application classes.
+     * Adds a bytecode transformer that can transform application classes
+     * <p>
+     * This is added on a per-class basis, by specifying the class name. The transformer is a function that
+     * can be used to wrap an ASM {@link ClassVisitor}.
      * <p>
-     * This takes the form of a function that takes the class name as a String, and returns a Function that wraps an
-     * ASM {@link ClassVisitor}.
-     *
-     * If this function returns null then no transform is applied. If it returns a function then it will be transformed.
-     *
      * The transformation is applied by calling each function that has been registered it turn to create a chain
      * of visitors. These visitors are then applied and the result is saved to the output.
-     *
+     * <p>
      * At present these transformations are only applied to application classes, not classes provided by dependencies
+     * <p>
+     * These transformations may be run concurrently in multiple threads, so if a function is registered for multiple
+     * classes it must be thread safe
      */
-    void addByteCodeTransformer(Function<String, Function<ClassVisitor, ClassVisitor>> visitorFunction);
+    void addByteCodeTransformer(String classToTransform, BiFunction<String, ClassVisitor, ClassVisitor> visitorFunction);
 
     /**
      * Adds a resource to the image that will be accessible when running under substrate.
@@ -126,25 +127,24 @@ public interface ProcessorContext {
      *
      * @param classes The classes to lazily init
      */
-    void addRuntimeInitializedClasses(String ... classes);
+    void addRuntimeInitializedClasses(String... classes);
 
     /**
      * Adds a proxy definition to allow proxies to be created using {@link java.lang.reflect.Proxy}
      *
      * @param proxyClasses The interface names that this proxy will implement
      */
-    void addProxyDefinition(String ... proxyClasses);
+    void addProxyDefinition(String... proxyClasses);
 
     /**
      * Set a system property to be passed in to the native image tool.
      *
-     * @param name the property name (must not be {@code null})
+     * @param name  the property name (must not be {@code null})
      * @param value the property value
      */
     void addNativeImageSystemProperty(String name, String value);
 
     /**
-     *
      * @param capability
      * @return if the given capability is present
      * @see Capabilities

core/deployment/src/main/java/org/jboss/shamrock/runner/RuntimeClassLoader.java

@@ -11,33 +11,33 @@
 import java.net.URLStreamHandler;
 import java.nio.file.Files;
 import java.nio.file.Path;
-import java.util.ArrayList;
 import java.util.Collections;
 import java.util.Enumeration;
 import java.util.HashMap;
 import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
+import java.util.function.BiFunction;
 import java.util.function.Consumer;
-import java.util.function.Function;
 
 import org.jboss.shamrock.deployment.ClassOutput;
 import org.objectweb.asm.ClassReader;
 import org.objectweb.asm.ClassVisitor;
 import org.objectweb.asm.ClassWriter;
 
-public class RuntimeClassLoader extends ClassLoader implements ClassOutput, Consumer<List<Function<String, Function<ClassVisitor, ClassVisitor>>>> {
+public class RuntimeClassLoader extends ClassLoader implements ClassOutput, Consumer<Map<String, List<BiFunction<String, ClassVisitor, ClassVisitor>>>> {
 
     private final Map<String, byte[]> appClasses = new HashMap<>();
     private final Set<String> frameworkClasses = new HashSet<>();
 
     private final Map<String, byte[]> resources = new HashMap<>();
 
-    private volatile List<Function<String, Function<ClassVisitor, ClassVisitor>>> functions = null;
+    private volatile Map<String, List<BiFunction<String, ClassVisitor, ClassVisitor>>> bytecodeTransformers = null;
 
     private final Path applicationClasses;
     private final Path frameworkClassesPath;
+
     static {
         registerAsParallelCapable();
     }
@@ -51,7 +51,7 @@ public RuntimeClassLoader(ClassLoader parent, Path applicationClasses, Path fram
     @Override
     public Enumeration<URL> getResources(String nm) throws IOException {
         String name;
-        if(nm.startsWith("/")) {
+        if (nm.startsWith("/")) {
             name = nm.substring(1);
         } else {
             name = nm;
@@ -83,7 +83,7 @@ public InputStream getInputStream() throws IOException {
     @Override
     public URL getResource(String nm) {
         String name;
-        if(nm.startsWith("/")) {
+        if (nm.startsWith("/")) {
             name = nm.substring(1);
         } else {
             name = nm;
@@ -94,7 +94,7 @@ public URL getResource(String nm) {
     @Override
     public InputStream getResourceAsStream(String nm) {
         String name;
-        if(nm.startsWith("/")) {
+        if (nm.startsWith("/")) {
             name = nm.substring(1);
         } else {
             name = nm;
@@ -111,7 +111,7 @@ protected Class<?> loadClass(String name, boolean resolve) throws ClassNotFoundE
         if (appClasses.containsKey(name)) {
             return findClass(name);
         }
-        if(frameworkClasses.contains(name)) {
+        if (frameworkClasses.contains(name)) {
             return super.loadClass(name, resolve);
         }
 
@@ -130,31 +130,33 @@ protected Class<?> loadClass(String name, boolean resolve) throws ClassNotFoundE
             }
             byte[] bytes = out.toByteArray();
             bytes = handleTransform(name, bytes);
-            return defineClass(name, bytes, 0, bytes.length);
+            try {
+                return defineClass(name, bytes, 0, bytes.length);
+            } catch (Error e) {
+                //potential race conditions if another thread is loading the same class
+                ex = findLoadedClass(name);
+                if(ex != null) {
+                    return ex;
+                }
+                throw e;
+            }
         }
         return super.loadClass(name, resolve);
     }
 
     private byte[] handleTransform(String name, byte[] bytes) {
-        if (functions == null || functions.isEmpty()) {
+        if (bytecodeTransformers == null || bytecodeTransformers.isEmpty()) {
             return bytes;
         }
-        List<Function<ClassVisitor, ClassVisitor>> transformers = new ArrayList<>();
-        for (Function<String, Function<ClassVisitor, ClassVisitor>> function : this.functions) {
-            Function<ClassVisitor, ClassVisitor> res = function.apply(name);
-            if (res != null) {
-                transformers.add(res);
-            }
-        }
-        if (transformers.isEmpty()) {
+        List<BiFunction<String, ClassVisitor, ClassVisitor>> transformers = bytecodeTransformers.get(name);
+        if (transformers == null) {
             return bytes;
         }
-
         ClassReader cr = new ClassReader(bytes);
         ClassWriter writer = new ClassWriter(cr, ClassWriter.COMPUTE_FRAMES | ClassWriter.COMPUTE_MAXS);
         ClassVisitor visitor = writer;
-        for (Function<ClassVisitor, ClassVisitor> i : transformers) {
-            visitor = i.apply(visitor);
+        for (BiFunction<String, ClassVisitor, ClassVisitor> i : transformers) {
+            visitor = i.apply(name, visitor);
         }
         cr.accept(visitor, 0);
         byte[] data = writer.toByteArray();
@@ -163,11 +165,24 @@ private byte[] handleTransform(String name, byte[] bytes) {
 
     @Override
     protected Class<?> findClass(String name) throws ClassNotFoundException {
+        Class<?> existing = findLoadedClass(name);
+        if(existing != null) {
+            return existing;
+        }
         byte[] bytes = appClasses.get(name);
         if (bytes == null) {
             throw new ClassNotFoundException(name);
         }
-        return defineClass(name, bytes, 0, bytes.length);
+        try {
+            return defineClass(name, bytes, 0, bytes.length);
+        } catch (Error e) {
+            //potential race conditions if another thread is loading the same class
+            existing = findLoadedClass(name);
+            if(existing != null) {
+                return existing;
+            }
+            throw e;
+        }
     }
 
     @Override
@@ -184,7 +199,7 @@ public void writeClass(boolean applicationClass, String className, byte[] data)
             Path fileName = frameworkClassesPath.resolve(className.replace(".", "/") + ".class");
             try {
                 Files.createDirectories(fileName.getParent());
-                try(FileOutputStream out = new FileOutputStream(fileName.toFile())) {
+                try (FileOutputStream out = new FileOutputStream(fileName.toFile())) {
                     out.write(data);
                 }
             } catch (IOException e) {
@@ -193,8 +208,8 @@ public void writeClass(boolean applicationClass, String className, byte[] data)
         }
     }
 
-    public void accept(List<Function<String, Function<ClassVisitor, ClassVisitor>>> functions) {
-        this.functions = functions;
+    public void accept(Map<String, List<BiFunction<String, ClassVisitor, ClassVisitor>>> functions) {
+        this.bytecodeTransformers = functions;
     }
 
     public void writeResource(String name, byte[] data) throws IOException {

core/deployment/src/main/java/org/jboss/shamrock/runner/RuntimeRunner.java

@@ -5,9 +5,17 @@
 import java.lang.reflect.InvocationTargetException;
 import java.lang.reflect.Method;
 import java.nio.file.Path;
+import java.util.List;
+import java.util.Map;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.function.BiFunction;
+import java.util.function.Function;
 
 import org.jboss.shamrock.deployment.ArchiveContextBuilder;
 import org.jboss.shamrock.deployment.BuildTimeGenerator;
+import org.jboss.shamrock.runtime.Timing;
+import org.objectweb.asm.ClassVisitor;
 
 /**
  * Class that can be used to run shamrock directly, ececuting the build and runtime
@@ -38,7 +46,29 @@ public void run() {
         try {
             BuildTimeGenerator buildTimeGenerator = new BuildTimeGenerator(loader, loader, false, archiveContextBuilder);
             buildTimeGenerator.run(target);
-            loader.accept(buildTimeGenerator.getBytecodeTransformers());
+            loader.accept(buildTimeGenerator.getByteCodeTransformers());
+            if(!buildTimeGenerator.getByteCodeTransformers().isEmpty()) {
+                //transformation can be slow, and classes that are transformed are generally always loaded on startup
+                //to speed this along we eagerly load the classes in parallel
+
+                ExecutorService executorService = Executors.newFixedThreadPool(Runtime.getRuntime().availableProcessors());
+                for(Map.Entry<String, List<BiFunction<String, ClassVisitor, ClassVisitor>>> entry : buildTimeGenerator.getByteCodeTransformers().entrySet()) {
+                    executorService.submit(new Runnable() {
+                        @Override
+                        public void run() {
+                            try {
+                                loader.loadClass(entry.getKey(), true);
+                            } catch (ClassNotFoundException e) {
+                                //ignore
+                                //this will show up at runtime anyway
+                            }
+                        }
+                    });
+                }
+                executorService.shutdown();
+
+            }
+
             Class<?> mainClass = loader.findClass(BuildTimeGenerator.MAIN_CLASS);
             ClassLoader old = Thread.currentThread().getContextClassLoader();
             try {