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");
+ }
+
+}