Speedup Hibernate ORM's enhancement of large models

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

@@ -2,18 +2,22 @@
 
 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;
+import net.bytebuddy.dynamic.ClassFileLocator;
 
 /**
  * Used to transform bytecode by registering to
@@ -29,44 +33,37 @@
  */
 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.persistence.", "java.",
+            "org.hibernate.annotations.",
+            "io.quarkus.hibernate.reactive.panache.", "io.quarkus.hibernate.orm.panache.",
+            "org.hibernate.search.mapper.pojo.mapping.definition.annotation.",
+            "jakarta.validation.constraints.");
+
+    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 +80,65 @@ 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);
+    }
+
+    private static class EnhancerHolder {
+
+        private volatile Enhancer actualEnhancer;
+
+        public Enhancer getEnhancer() {
+            //Lazily initialized as it's expensive and might not be necessary: these transformations are cacheable.
+            if (actualEnhancer == null) {
+                synchronized (this) {
+                    if (actualEnhancer == null) {
+                        actualEnhancer = PROVIDER.getEnhancer(QuarkusEnhancementContext.INSTANCE, new ThreadsafeLocator());
+                    }
+                }
             }
+            return actualEnhancer;
+        }
+    }
+
+    private static final class ThreadsafeLocator implements EnhancerClassLocator {
+
+        final ThreadLocal<EnhancerClassLocator> localLocator = ThreadLocal
+                .withInitial(() -> ModelTypePool.buildModelTypePool(QuarkusClassFileLocator.INSTANCE,
+                        CORE_POOL));
 
-        };
-        Enhancer enhancer = PROVIDER.getEnhancer(enhancementContext);
-        return enhancer.enhance(className, bytes);
+        @Override
+        public void registerClassNameAndBytes(String s, byte[] bytes) {
+            localLocator.get().registerClassNameAndBytes(s, bytes);
+        }
+
+        @Override
+        public void deregisterClassNameAndBytes(String s) {
+            localLocator.get().deregisterClassNameAndBytes(s);
+        }
+
+        @Override
+        public ClassFileLocator asClassFileLocator() {
+            return localLocator.get().asClassFileLocator();
+        }
+
+        @Override
+        public Resolution describe(String s) {
+            return localLocator.get().describe(s);
+        }
+
+        @Override
+        public void clear() {
+            //not essential as it gets discarded, but could help:
+            localLocator.get().clear();
+            localLocator.remove();
+        }
     }
+
 }

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

@@ -0,0 +1,43 @@
+package io.quarkus.hibernate.orm.deployment.integration;
+
+import java.io.IOException;
+
+import io.quarkus.deployment.util.IoUtil;
+import net.bytebuddy.dynamic.ClassFileLocator;
+
+/**
+ * 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();
+        final byte[] bytes = IoUtil.readClassAsBytes(classLoader, name);
+        if (bytes != null) {
+            return new Resolution.Explicit(bytes);
+        } 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");
+    }
+
+}