HHH-18011 Parallel entity enhancement causes a classloader cache stampede [6.5]

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