HHH-18011 Extract DefaultEnhancerClassFileLocator and allow using a different implementation

hibernate-core/src/main/java/org/hibernate/bytecode/enhance/internal/bytebuddy/ByteBuddyEnhancementContext.java

@@ -7,6 +7,7 @@
 package org.hibernate.bytecode.enhance.internal.bytebuddy;
 
 import java.util.Map;
+import java.util.Objects;
 import java.util.Optional;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.function.Function;
@@ -35,10 +36,16 @@ class ByteBuddyEnhancementContext {
 	private final ConcurrentHashMap<TypeDescription, Map<String, MethodDescription>> getterByTypeMap = new ConcurrentHashMap<>();
 	private final ConcurrentHashMap<String, Object> locksMap = new ConcurrentHashMap<>();
 
-	ByteBuddyEnhancementContext(EnhancementContext enhancementContext) {
+	ByteBuddyEnhancementContext(final EnhancementContext enhancementContext) {
-		this.enhancementContext = enhancementContext;
+		this.enhancementContext = Objects.requireNonNull( enhancementContext );
 	}
 
+	/**
+	 * @deprecated as it's currently unused and we're not always actually sourcing the classes to be transformed
+	 * from a classloader, so this getter can't always be honoured correctly.
+	 * @return the ClassLoader provided by the underlying EnhancementContext. Might be otherwise ignored.
+	 */
+	@Deprecated(forRemoval = true)
 	public ClassLoader getLoadingClassLoader() {
 		return enhancementContext.getLoadingClassLoader();
 	}

hibernate-core/src/main/java/org/hibernate/bytecode/enhance/internal/bytebuddy/CoreTypePool.java

@@ -0,0 +1,98 @@
+/*
+ * Hibernate, Relational Persistence for Idiomatic Java
+ *
+ * License: GNU Lesser General Public License (LGPL), version 2.1 or later.
+ * See the lgpl.txt file in the root directory or <http://www.gnu.org/licenses/lgpl-2.1.html>.
+ */
+package org.hibernate.bytecode.enhance.internal.bytebuddy;
+
+import java.util.concurrent.ConcurrentHashMap;
+
+import net.bytebuddy.description.type.TypeDescription;
+import net.bytebuddy.pool.TypePool;
+
+/**
+ * A TypePool which only loads, and caches, types whose package
+ * name starts with certain chosen prefixes.
+ * The default is to only load classes whose package names start with
+ * either "jakarta." or "java.".
+ * This allows to reuse these caches independently from application
+ * code and classloader changes, as during enhancement we frequently
+ * encounter such symbols as well, for example triggered by JPA annotations
+ * or properties mapped via standard java types and collections.
+ * Symbols resolved by this pool are backed by loaded classes from
+ * ORM's classloader.
+ */
+public class CoreTypePool extends TypePool.AbstractBase implements TypePool {
+
+	private final ClassLoader hibernateClassLoader = CoreTypePool.class.getClassLoader();
+	private final ConcurrentHashMap<String, Resolution> resolutions = new ConcurrentHashMap<>();
+	private final String[] acceptedPrefixes;
+
+	/**
+	 * Construct a new {@link CoreTypePool} with its default configuration:
+	 * to only load classes whose package names start with either "jakarta."
+	 * or "java."
+	 */
+	public CoreTypePool() {
+		//By default optimise for jakarta annotations, and java util collections
+		this("jakarta.", "java.", "org.hibernate.annotations.");
+	}
+
+	/**
+	 * Construct a new {@link CoreTypePool} with a choice of which prefixes
+	 * for fully qualified classnames will be loaded by this {@link TypePool}.
+	 */
+	public CoreTypePool(final String... acceptedPrefixes) {
+		//While we implement a cache in this class we also want to enable
+		//ByteBuddy's default caching mechanism as it will cache the more
+		//useful output of the parsing and introspection of such types.
+		super( new TypePool.CacheProvider.Simple() );
+		this.acceptedPrefixes = acceptedPrefixes;
+	}
+
+	private boolean isCoreClassName(final String name) {
+		for ( String acceptedPrefix : this.acceptedPrefixes ) {
+			if ( name.startsWith( acceptedPrefix ) ) {
+				return true;
+			}
+		}
+		return false;
+	}
+
+	@Override
+	protected Resolution doDescribe(final String name) {
+		if ( isCoreClassName( name ) ) {
+			final Resolution resolution = resolutions.get( name );
+			if ( resolution != null ) {
+				return resolution;
+			}
+			else {
+				//We implement this additional layer of caching, which is on top of
+				//ByteBuddy's default caching, so as to prevent resolving the same
+				//types concurrently from the classloader.
+				//This is merely an efficiency improvement and will NOT provide a
+				//strict guarantee of symbols being resolved exactly once as there
+				//is no SPI within ByteBuddy which would allow this: the point is to
+				//make it exceptionally infrequent, which greatly helps with
+				//processing of large models.
+				return resolutions.computeIfAbsent( name, this::actualResolve );
+			}
+		}
+		else {
+			//These are not cached to not leak references to application code names
+			return new Resolution.Illegal( name );
+		}
+	}
+
+	private Resolution actualResolve(final String name) {
+		try {
+			final Class<?> aClass = Class.forName( name, false, hibernateClassLoader );
+			return new TypePool.Resolution.Simple( TypeDescription.ForLoadedType.of( aClass ) );
+		}
+		catch ( ClassNotFoundException e ) {
+			return new Resolution.Illegal( name );
+		}
+	}
+
+}

hibernate-core/src/main/java/org/hibernate/bytecode/enhance/internal/bytebuddy/EnhancerClassLocator.java

@@ -0,0 +1,36 @@
+/*
+ * Hibernate, Relational Persistence for Idiomatic Java
+ *
+ * License: GNU Lesser General Public License (LGPL), version 2.1 or later.
+ * See the lgpl.txt file in the root directory or <http://www.gnu.org/licenses/lgpl-2.1.html>.
+ */
+package org.hibernate.bytecode.enhance.internal.bytebuddy;
+
+import net.bytebuddy.dynamic.ClassFileLocator;
+import net.bytebuddy.pool.TypePool;
+
+/**
+ * Extends the TypePool contract of ByteBuddy with our additional needs.
+ */
+public interface EnhancerClassLocator extends TypePool {
+
+	/**
+	 * Register a new class to the locator explicitly.
+	 * @param className
+	 * @param originalBytes
+	 */
+	void registerClassNameAndBytes(String className, byte[] originalBytes);
+
+	/**
+	 * This can optionally be used to remove an explicit mapping when it's no longer
+	 * essential to retain it.
+	 * The underlying implementation might ignore the operation.
+	 * @param className
+	 */
+	void deregisterClassNameAndBytes(String className);
+
+	/**
+	 * @return the underlying {@link ClassFileLocator}
+	 */
+	ClassFileLocator asClassFileLocator();
+}

hibernate-core/src/main/java/org/hibernate/bytecode/enhance/internal/bytebuddy/EnhancerImpl.java

@@ -6,7 +6,6 @@
  */
 package org.hibernate.bytecode.enhance.internal.bytebuddy;
 
-import java.io.IOException;
 import java.lang.annotation.Annotation;
 import java.lang.reflect.Modifier;
 import java.util.ArrayList;
@@ -14,8 +13,8 @@ import java.util.Collection;
 import java.util.Collections;
 import java.util.List;
 import java.util.Map;
+import java.util.Objects;
 import java.util.Optional;
-import java.util.concurrent.ConcurrentHashMap;
 import java.util.function.Supplier;
 
 import org.hibernate.Version;
@@ -66,7 +65,6 @@ import net.bytebuddy.implementation.FieldAccessor;
 import net.bytebuddy.implementation.FixedValue;
 import net.bytebuddy.implementation.Implementation;
 import net.bytebuddy.implementation.StubMethod;
-import net.bytebuddy.pool.TypePool;
 
 import static net.bytebuddy.matcher.ElementMatchers.isDefaultFinalizer;
 
@@ -93,9 +91,7 @@ public class EnhancerImpl implements Enhancer {
 
 	protected final ByteBuddyEnhancementContext enhancementContext;
 	private final ByteBuddyState byteBuddyState;
-
+	private final EnhancerClassLocator typePool;
-	private final EnhancerClassFileLocator classFileLocator;
-	private final TypePool typePool;
 
 	/**
 	 * Extract the following constants so that enhancement on large projects
@@ -126,10 +122,20 @@ public class EnhancerImpl implements Enhancer {
 	 * @param byteBuddyState refers to the ByteBuddy instance to use
 	 */
 	public EnhancerImpl(final EnhancementContext enhancementContext, final ByteBuddyState byteBuddyState) {
+		this( enhancementContext, byteBuddyState, ModelTypePool.buildModelTypePool( enhancementContext.getLoadingClassLoader() ) );
+	}
+
+	/**
+	 * Expert level constructor, this allows for more control of state and bytecode loading,
+	 * which allows integrators to optimise for particular contexts of use.
+	 * @param enhancementContext
+	 * @param byteBuddyState
+	 * @param classLocator
+	 */
+	public EnhancerImpl(final EnhancementContext enhancementContext, final ByteBuddyState byteBuddyState, final EnhancerClassLocator classLocator) {
 		this.enhancementContext = new ByteBuddyEnhancementContext( enhancementContext );
-		this.byteBuddyState = byteBuddyState;
+		this.byteBuddyState = Objects.requireNonNull( byteBuddyState );
-		this.classFileLocator = new EnhancerClassFileLocator( enhancementContext.getLoadingClassLoader() );
+		this.typePool = Objects.requireNonNull( classLocator );
-		this.typePool = buildTypePool( classFileLocator );
 	}
 
 	/**
@@ -147,13 +153,13 @@ public class EnhancerImpl implements Enhancer {
 	public byte[] enhance(String className, byte[] originalBytes) throws EnhancementException {
 		//Classpool#describe does not accept '/' in the description name as it expects a class name. See HHH-12545
 		final String safeClassName = className.replace( '/', '.' );
-		classFileLocator.registerClassNameAndBytes( safeClassName, originalBytes );
+		typePool.registerClassNameAndBytes( safeClassName, originalBytes );
 		try {
 			final TypeDescription typeDescription = typePool.describe( safeClassName ).resolve();
 
 			return byteBuddyState.rewrite( typePool, safeClassName, byteBuddy -> doEnhance(
 					() -> byteBuddy.ignore( isDefaultFinalizer() )
-							.redefine( typeDescription, ClassFileLocator.Simple.of( safeClassName, originalBytes ) )
+							.redefine( typeDescription, typePool.asClassFileLocator() )
 							.annotateType( HIBERNATE_VERSION_ANNOTATION ),
 					typeDescription
 			) );
@@ -165,14 +171,14 @@ public class EnhancerImpl implements Enhancer {
 			throw new EnhancementException( "Failed to enhance class " + className, e );
 		}
 		finally {
-			classFileLocator.deregisterClassNameAndBytes( safeClassName );
+			typePool.deregisterClassNameAndBytes( safeClassName );
 		}
 	}
 
 	@Override
 	public void discoverTypes(String className, byte[] originalBytes) {
 		if ( originalBytes != null ) {
-			classFileLocator.registerClassNameAndBytes( className, originalBytes );
+			typePool.registerClassNameAndBytes( className, originalBytes );
 		}
 		try {
 			final TypeDescription typeDescription = typePool.describe( className ).resolve();
@@ -183,14 +189,10 @@ public class EnhancerImpl implements Enhancer {
 			throw new EnhancementException( "Failed to discover types for class " + className, e );
 		}
 		finally {
-			classFileLocator.deregisterClassNameAndBytes( className );
+			typePool.deregisterClassNameAndBytes( className );
 		}
 	}
 
-	private TypePool buildTypePool(final ClassFileLocator classFileLocator) {
-		return TypePool.Default.WithLazyResolution.of( classFileLocator );
-	}
-
 	private DynamicType.Builder<?> doEnhance(Supplier<DynamicType.Builder<?>> builderSupplier, TypeDescription managedCtClass) {
 		// can't effectively enhance interfaces
 		if ( managedCtClass.isInterface() ) {
@@ -652,39 +654,4 @@ public class EnhancerImpl implements Enhancer {
 		}
 	}
 
-	private static class EnhancerClassFileLocator extends ClassFileLocator.ForClassLoader {
-		private final ConcurrentHashMap<String, Resolution> resolutions = new ConcurrentHashMap<>();
-
-		/**
-		 * Creates a new class file locator for the given class loader.
-		 *
-		 * @param classLoader The class loader to query which must not be the bootstrap class loader, i.e. {@code null}.
-		 */
-		protected EnhancerClassFileLocator(ClassLoader classLoader) {
-			super( classLoader );
-		}
-
-		@Override
-		public Resolution locate(String className) throws IOException {
-			assert className != null;
-			final Resolution resolution = resolutions.get( className );
-			if ( resolution != null ) {
-				return resolution;
-			}
-			else {
-				return super.locate( className );
-			}
-		}
-
-		void registerClassNameAndBytes(String className, byte[] bytes) {
-			assert className != null;
-			assert bytes != null;
-			resolutions.put( className, new Resolution.Explicit( bytes ) );
-		}
-
-		void deregisterClassNameAndBytes(String className) {
-			resolutions.remove( className );
-		}
-	}
-
 }

hibernate-core/src/main/java/org/hibernate/bytecode/enhance/internal/bytebuddy/ModelTypePool.java

@@ -0,0 +1,108 @@
+/*
+ * Hibernate, Relational Persistence for Idiomatic Java
+ *
+ * License: GNU Lesser General Public License (LGPL), version 2.1 or later.
+ * See the lgpl.txt file in the root directory or <http://www.gnu.org/licenses/lgpl-2.1.html>.
+ */
+package org.hibernate.bytecode.enhance.internal.bytebuddy;
+
+import java.util.Objects;
+import java.util.concurrent.ConcurrentHashMap;
+
+import net.bytebuddy.dynamic.ClassFileLocator;
+import net.bytebuddy.pool.TypePool;
+
+/**
+ * A TypePool suitable for loading user's classes,
+ * potentially in parallel operations.
+ */
+public class ModelTypePool extends TypePool.Default implements EnhancerClassLocator {
+
+	private final ConcurrentHashMap<String, Resolution> resolutions = new ConcurrentHashMap<>();
+	private final OverridingClassFileLocator locator;
+
+	private ModelTypePool(CacheProvider cacheProvider, OverridingClassFileLocator classFileLocator, CoreTypePool parent) {
+		super( cacheProvider, classFileLocator, ReaderMode.FAST, parent );
+		this.locator = classFileLocator;
+	}
+
+	/**
+	 * Creates a new empty EnhancerClassLocator instance which will load any application
+	 * classes that need being reflected on from the ClassLoader passed as parameter.
+	 * This TypePool will delegate, parent first, to a newly constructed empty instance
+	 * of CoreTypePool; this parent pool will be used to load non-application types from
+	 * the Hibernate classloader instead, not the one specified as argument.
+	 * @see CoreTypePool
+	 * @param classLoader
+	 * @return the newly created EnhancerClassLocator
+	 */
+	public static EnhancerClassLocator buildModelTypePool(ClassLoader classLoader) {
+		return buildModelTypePool( ClassFileLocator.ForClassLoader.of( classLoader ) );
+	}
+
+	/**
+	 * Similar to {@link #buildModelTypePool(ClassLoader)} except the application classes
+	 * are not necessarily sourced from a standard classloader: it accepts a {@link ClassFileLocator},
+	 * which offers some more flexibility.
+	 * @param classFileLocator
+	 * @return the newly created EnhancerClassLocator
+	 */
+	public static EnhancerClassLocator buildModelTypePool(ClassFileLocator classFileLocator) {
+		return buildModelTypePool( classFileLocator, new CoreTypePool() );
+	}
+
+	/**
+	 * Similar to {@link #buildModelTypePool(ClassFileLocator)} but allows specifying an existing
+	 * {@link CoreTypePool} to be used as parent pool.
+	 * This forms allows constructing a custom CoreTypePool and also separated the cache of the parent pool,
+	 * which might be useful to reuse for multiple enhancement processes while desiring a clean new
+	 * state for the {@link ModelTypePool}.
+	 * @param classFileLocator
+	 * @param coreTypePool
+	 * @return
+	 */
+	public static EnhancerClassLocator buildModelTypePool(ClassFileLocator classFileLocator, CoreTypePool coreTypePool) {
+		return buildModelTypePool( classFileLocator, coreTypePool, new TypePool.CacheProvider.Simple() );
+	}
+
+	/**
+	 * The more advanced strategy to construct a new ModelTypePool, allowing customization of all its aspects.
+	 * @param classFileLocator
+	 * @param coreTypePool
+	 * @param cacheProvider
+	 * @return
+	 */
+	public static EnhancerClassLocator buildModelTypePool(ClassFileLocator classFileLocator, CoreTypePool coreTypePool, CacheProvider cacheProvider) {
+		Objects.requireNonNull( classFileLocator );
+		Objects.requireNonNull( coreTypePool );
+		Objects.requireNonNull( cacheProvider );
+		return new ModelTypePool( cacheProvider, new OverridingClassFileLocator( classFileLocator ), coreTypePool );
+	}
+
+	@Override
+	protected Resolution doDescribe(final String name) {
+		final Resolution resolution = resolutions.get( name );
+		if ( resolution != null ) {
+			return resolution;
+		}
+		else {
+			return resolutions.computeIfAbsent( name, super::doDescribe );
+		}
+	}
+
+	@Override
+	public void registerClassNameAndBytes(final String className, final byte[] bytes) {
+		locator.put( className, new ClassFileLocator.Resolution.Explicit( Objects.requireNonNull( bytes ) ) );
+	}
+
+	@Override
+	public void deregisterClassNameAndBytes(final String className) {
+		locator.remove( className );
+	}
+
+	@Override
+	public ClassFileLocator asClassFileLocator() {
+		return locator;
+	}
+
+}

hibernate-core/src/main/java/org/hibernate/bytecode/enhance/internal/bytebuddy/OverridingClassFileLocator.java

@@ -0,0 +1,52 @@
+/*
+ * Hibernate, Relational Persistence for Idiomatic Java
+ *
+ * License: GNU Lesser General Public License (LGPL), version 2.1 or later.
+ * See the lgpl.txt file in the root directory or <http://www.gnu.org/licenses/lgpl-2.1.html>.
+ */
+package org.hibernate.bytecode.enhance.internal.bytebuddy;
+
+import java.io.IOException;
+import java.util.Objects;
+import java.util.concurrent.ConcurrentHashMap;
+
+import net.bytebuddy.dynamic.ClassFileLocator;
+
+/**
+ * Allows wrapping another ClassFileLocator to add the ability to define
+ * resolution overrides for specific resources.
+ */
+public final class OverridingClassFileLocator implements ClassFileLocator {
+
+	private final ConcurrentHashMap<String, Resolution> registeredResolutions = new ConcurrentHashMap<>();
+	private final ClassFileLocator parent;
+
+	public OverridingClassFileLocator(final ClassFileLocator parent) {
+		this.parent = Objects.requireNonNull( parent );
+	}
+
+	@Override
+	public Resolution locate(final String name) throws IOException {
+		final Resolution resolution = registeredResolutions.get( name );
+		if ( resolution != null ) {
+			return resolution;
+		}
+		else {
+			return parent.locate( name );
+		}
+	}
+
+	@Override
+	public void close() throws IOException {
+		//Nothing to do: we're not responsible for parent
+	}
+
+	void put(String className, Resolution.Explicit explicit) {
+		registeredResolutions.put( className, explicit );
+	}
+
+	void remove(String className) {
+		registeredResolutions.remove( className );
+	}
+
+}

hibernate-core/src/main/java/org/hibernate/bytecode/internal/bytebuddy/BytecodeProviderImpl.java

@@ -21,6 +21,7 @@ import java.util.Map;
 import java.util.concurrent.Callable;
 
 import org.hibernate.HibernateException;
+import org.hibernate.bytecode.enhance.internal.bytebuddy.EnhancerClassLocator;
 import org.hibernate.bytecode.enhance.internal.bytebuddy.EnhancerImpl;
 import org.hibernate.bytecode.enhance.spi.EnhancementContext;
 import org.hibernate.bytecode.enhance.spi.Enhancer;
@@ -67,6 +68,7 @@ import net.bytebuddy.jar.asm.Opcodes;
 import net.bytebuddy.jar.asm.Type;
 import net.bytebuddy.matcher.ElementMatcher;
 import net.bytebuddy.matcher.ElementMatchers;
+import net.bytebuddy.pool.TypePool;
 import org.checkerframework.checker.nullness.qual.Nullable;
 
 public class BytecodeProviderImpl implements BytecodeProvider {
@@ -1312,6 +1314,18 @@ public class BytecodeProviderImpl implements BytecodeProvider {
 		return new EnhancerImpl( enhancementContext, byteBuddyState );
 	}
 
+	/**
+	 * Similar to {@link #getEnhancer(EnhancementContext)} but intended for advanced users who wish
+	 * to customize how ByteBuddy is locating the class files and caching the types.
+	 * Possibly used in Quarkus in a future version.
+	 * @param enhancementContext
+	 * @param classLocator
+	 * @return
+	 */
+	public @Nullable Enhancer getEnhancer(EnhancementContext enhancementContext, EnhancerClassLocator classLocator) {
+		return new EnhancerImpl( enhancementContext, byteBuddyState, classLocator );
+	}
+
 	@Override
 	public void resetCaches() {
 		byteBuddyState.clearState();