HHH-18011 Extract DefaultEnhancerClassFileLocator and allow using a different implementation

Author: Sanne

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) {
-		this.enhancementContext = enhancementContext;
+	ByteBuddyEnhancementContext(final 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,16 +6,15 @@
  */
 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;
 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.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<? extends Annotation> annotationType() {
 
 	protected final ByteBuddyEnhancementContext enhancementContext;
 	private final ByteBuddyState byteBuddyState;
-
-	private final EnhancerClassFileLocator classFileLocator;
-	private final TypePool typePool;
+	private final EnhancerClassLocator typePool;
 
 	/**
 	 * Extract the following constants so that enhancement on large projects
@@ -126,10 +122,20 @@ public Class<? extends Annotation> annotationType() {
 	 * @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.classFileLocator = new EnhancerClassFileLocator( enhancementContext.getLoadingClassLoader() );
-		this.typePool = buildTypePool( classFileLocator );
+		this.byteBuddyState = Objects.requireNonNull( byteBuddyState );
+		this.typePool = Objects.requireNonNull( classLocator );
 	}
 
 	/**
@@ -147,13 +153,13 @@ public EnhancerImpl(final EnhancementContext enhancementContext, final ByteBuddy
 	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 byte[] enhance(String className, byte[] originalBytes) throws Enhancement
 			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 void discoverTypes(String className, byte[] originalBytes) {
 			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 @@ else if ( access != null && access.load().value() == AccessType.FIELD ) {
 		}
 	}
 
-	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 );
-		}
-	}
-
 }