Speedup Hibernate ORM's enhancement of large models

extensions/hibernate-orm/deployment/src/main/java/io/quarkus/hibernate/orm/deployment/HibernateEntityEnhancer.java

@@ -2,17 +2,20 @@
 
 import java.util.function.BiFunction;
 
-import org.hibernate.bytecode.enhance.spi.DefaultEnhancementContext;
+import org.hibernate.bytecode.enhance.internal.bytebuddy.CoreTypePool;
+import org.hibernate.bytecode.enhance.internal.bytebuddy.EnhancerClassLocator;
+import org.hibernate.bytecode.enhance.internal.bytebuddy.ModelTypePool;
 import org.hibernate.bytecode.enhance.spi.Enhancer;
-import org.hibernate.bytecode.enhance.spi.UnloadedField;
-import org.hibernate.bytecode.spi.BytecodeProvider;
+import org.hibernate.bytecode.internal.bytebuddy.BytecodeProviderImpl;
 import org.objectweb.asm.ClassReader;
 import org.objectweb.asm.ClassVisitor;
 import org.objectweb.asm.ClassWriter;
 
 import io.quarkus.deployment.QuarkusClassVisitor;
 import io.quarkus.deployment.QuarkusClassWriter;
 import io.quarkus.gizmo.Gizmo;
+import io.quarkus.hibernate.orm.deployment.integration.QuarkusClassFileLocator;
+import io.quarkus.hibernate.orm.deployment.integration.QuarkusEnhancementContext;
 import net.bytebuddy.ClassFileVersion;
 
 /**
@@ -29,44 +32,34 @@
  */
 public final class HibernateEntityEnhancer implements BiFunction<String, ClassVisitor, ClassVisitor> {
 
-    private static final BytecodeProvider PROVIDER = new org.hibernate.bytecode.internal.bytebuddy.BytecodeProviderImpl(
+    private static final BytecodeProviderImpl PROVIDER = new org.hibernate.bytecode.internal.bytebuddy.BytecodeProviderImpl(
             ClassFileVersion.JAVA_V11);
 
+    //Choose this set to include Jakarta annotations, basic Java types such as String and Map, Hibernate annotations, and Panache supertypes:
+    private static final CoreTypePool CORE_POOL = new CoreTypePool("jakarta.", "java.", "org.hibernate.",
+            "io.quarkus.hibernate.");
+
+    private final EnhancerHolder enhancerHolder = new EnhancerHolder();
+
     @Override
     public ClassVisitor apply(String className, ClassVisitor outputClassVisitor) {
-        return new HibernateEnhancingClassVisitor(className, outputClassVisitor);
+        return new HibernateEnhancingClassVisitor(className, outputClassVisitor, enhancerHolder);
     }
 
     private static class HibernateEnhancingClassVisitor extends QuarkusClassVisitor {
 
         private final String className;
         private final ClassVisitor outputClassVisitor;
-        private final Enhancer enhancer;
+        private final EnhancerHolder enhancerHolder;
 
-        public HibernateEnhancingClassVisitor(String className, ClassVisitor outputClassVisitor) {
+        public HibernateEnhancingClassVisitor(String className, ClassVisitor outputClassVisitor,
+                EnhancerHolder enhancerHolder) {
             //Careful: the ASM API version needs to match the ASM version of Gizmo, not the one from Byte Buddy.
             //Most often these match - but occasionally they will diverge which is acceptable as Byte Buddy is shading ASM.
             super(Gizmo.ASM_API_VERSION, new QuarkusClassWriter(ClassWriter.COMPUTE_FRAMES | ClassWriter.COMPUTE_MAXS));
             this.className = className;
             this.outputClassVisitor = outputClassVisitor;
-            //note that as getLoadingClassLoader is resolved immediately this can't be created until transform time
-
-            DefaultEnhancementContext enhancementContext = new DefaultEnhancementContext() {
-
-                @Override
-                public boolean doBiDirectionalAssociationManagement(final UnloadedField field) {
-                    //Don't enable automatic association management as it's often too surprising.
-                    //Also, there's several cases in which its semantics are of unspecified,
-                    //such as what should happen when dealing with ordered collections.
-                    return false;
-                }
-
-                @Override
-                public ClassLoader getLoadingClassLoader() {
-                    return Thread.currentThread().getContextClassLoader();
-                }
-            };
-            this.enhancer = PROVIDER.getEnhancer(enhancementContext);
+            this.enhancerHolder = enhancerHolder;
         }
 
         @Override
@@ -83,21 +76,36 @@ public void visitEnd() {
         }
 
         private byte[] hibernateEnhancement(final String className, final byte[] originalBytes) {
-            final byte[] enhanced = enhancer.enhance(className, originalBytes);
+            final byte[] enhanced = enhancerHolder.getEnhancer().enhance(className, originalBytes);
             return enhanced == null ? originalBytes : enhanced;
         }
 
     }
 
     public byte[] enhance(String className, byte[] bytes) {
-        DefaultEnhancementContext enhancementContext = new DefaultEnhancementContext() {
-            @Override
-            public ClassLoader getLoadingClassLoader() {
-                return Thread.currentThread().getContextClassLoader();
-            }
+        return enhancerHolder.getEnhancer().enhance(className, bytes);
+    }
 
-        };
-        Enhancer enhancer = PROVIDER.getEnhancer(enhancementContext);
-        return enhancer.enhance(className, bytes);
+    private static class EnhancerHolder {
+
+        private volatile Enhancer actualEnhancer;
+
+        public Enhancer getEnhancer() {
+            //Lazily initialized for multiple reasons:
+            //1)it's expensive: try to skip it if we can; this is actually not unlikely to happen as the transformation is cacheable.
+            //2)We want the Advice loaders of the Hibernate ORM implementation to be initialized within the scope in which we
+            //have the transformation classloader installed in the thread's context.
+            if (actualEnhancer == null) {
+                synchronized (this) {
+                    if (actualEnhancer == null) {
+                        EnhancerClassLocator modelPool = ModelTypePool.buildModelTypePool(QuarkusClassFileLocator.INSTANCE,
+                                CORE_POOL);
+                        actualEnhancer = PROVIDER.getEnhancer(QuarkusEnhancementContext.INSTANCE, modelPool);
+                    }
+                }
+            }
+            return actualEnhancer;
+        }
     }
+
 }

extensions/hibernate-orm/deployment/src/main/java/io/quarkus/hibernate/orm/deployment/integration/QuarkusClassFileLocator.java

@@ -0,0 +1,48 @@
+package io.quarkus.hibernate.orm.deployment.integration;
+
+import java.io.IOException;
+import java.io.InputStream;
+
+import net.bytebuddy.dynamic.ClassFileLocator;
+import net.bytebuddy.utility.StreamDrainer;
+
+/**
+ * Custom implementation of a ClassFileLocator which will load resources
+ * from the context classloader which is set at the time of the locate()
+ * operation is being performed.
+ * Using a regular ForClassLoader implementation would capture the currently
+ * set ClassLoader and keep a reference to it, while we need it to look
+ * for a fresh copy during the enhancement.
+ * Additionally, we might be able to optimize how the resource is actually
+ * being loaded as we control the ClassLoader implementations
+ * (Such further optimisations are not implemented yet).
+ */
+public final class QuarkusClassFileLocator implements ClassFileLocator {
+
+    public static final QuarkusClassFileLocator INSTANCE = new QuarkusClassFileLocator();
+
+    private QuarkusClassFileLocator() {
+        //do not invoke, use INSTANCE
+    }
+
+    @Override
+    public Resolution locate(final String name) throws IOException {
+        final ClassLoader classLoader = Thread.currentThread().getContextClassLoader();
+        InputStream inputStream = classLoader.getResourceAsStream(name.replace('.', '/') + CLASS_FILE_EXTENSION);
+        if (inputStream != null) {
+            try {
+                return new Resolution.Explicit(StreamDrainer.DEFAULT.drain(inputStream));
+            } finally {
+                inputStream.close();
+            }
+        } else {
+            return new Resolution.Illegal(name);
+        }
+    }
+
+    @Override
+    public void close() {
+        //nothing to do
+    }
+
+}

extensions/hibernate-orm/deployment/src/main/java/io/quarkus/hibernate/orm/deployment/integration/QuarkusEnhancementContext.java

@@ -0,0 +1,29 @@
+package io.quarkus.hibernate.orm.deployment.integration;
+
+import org.hibernate.bytecode.enhance.spi.DefaultEnhancementContext;
+import org.hibernate.bytecode.enhance.spi.UnloadedField;
+
+public final class QuarkusEnhancementContext extends DefaultEnhancementContext {
+
+    public static final QuarkusEnhancementContext INSTANCE = new QuarkusEnhancementContext();
+
+    private QuarkusEnhancementContext() {
+        //do not invoke, use INSTANCE
+    }
+
+    @Override
+    public boolean doBiDirectionalAssociationManagement(final UnloadedField field) {
+        //Don't enable automatic association management as it's often too surprising.
+        //Also, there's several cases in which its semantics are of unspecified,
+        //such as what should happen when dealing with ordered collections.
+        return false;
+    }
+
+    @Override
+    public ClassLoader getLoadingClassLoader() {
+        //This shouldn't matter as we delegate resource location to QuarkusClassFileLocator;
+        //make sure of this:
+        throw new IllegalStateException("The Classloader of the EnhancementContext should not be used");
+    }
+
+}