The entity enhancement process made defensive against concurrent enha…

…ncement of the same resources

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

@@ -0,0 +1,70 @@
+/*
+ * 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 java.util.concurrent.ConcurrentMap;
+
+import net.bytebuddy.description.type.TypeDescription;
+import net.bytebuddy.pool.TypePool;
+import net.bytebuddy.utility.nullability.MaybeNull;
+
+/**
+ * A CacheProvider for ByteBuddy which is specifically designed for
+ * our CoreTypePool: it ensures the cache content doesn't get tainted
+ * by model types or anything else which isn't responsibility of this
+ * particular type pool.
+ * @implNote This cache instance shares the same @{@link CorePrefixFilter}
+ * instance of the @{@link CoreTypePool} which is using it, and uses it
+ * to guard writes into its internal caches.
+ */
+class CoreCacheProvider implements TypePool.CacheProvider {
+
+	private final ConcurrentMap<String, TypePool.Resolution> storage = new ConcurrentHashMap<>();
+	private final CorePrefixFilter acceptedPrefixes;
+
+	CoreCacheProvider(final CorePrefixFilter acceptedPrefixes) {
+		this.acceptedPrefixes = Objects.requireNonNull( acceptedPrefixes );
+		register(
+				Object.class.getName(),
+				new TypePool.Resolution.Simple( TypeDescription.ForLoadedType.of( Object.class ) )
+		);
+	}
+
+	/**
+	 * {@inheritDoc}
+	 */
+	@MaybeNull
+	public TypePool.Resolution find(final String name) {
+		return storage.get( name );
+	}
+
+	/**
+	 * {@inheritDoc}
+	 */
+	public TypePool.Resolution register(String name, TypePool.Resolution resolution) {
+		//Ensure we DO NOT cache anything from a non-core namespace, to not leak application specific code:
+		if ( acceptedPrefixes.isCoreClassName( name ) ) {
+			TypePool.Resolution cached = storage.putIfAbsent( name, resolution );
+			return cached == null
+					? resolution
+					: cached;
+		}
+		else {
+			return resolution;
+		}
+	}
+
+	/**
+	 * {@inheritDoc}
+	 */
+	public void clear() {
+		storage.clear();
+	}
+
+}

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

@@ -0,0 +1,53 @@
+/*
+ * 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;
+
+/**
+ * We differentiate between core classes and application classes during symbol
+ * resolution for the purposes of entity enhancement.
+ * The discriminator is the prefix of the fully qualified classname, for
+ * example it could be package names.
+ * The "core classes" don't have to be comprehensively defined: we want a small
+ * set of prefixes for which we know with certainty that a)They won't be used
+ * in application code (assuming people honour reasonable package name rules)
+ * or any code that needs being enhanced. and b) frequently need to be looked up
+ * during the enhancement process.
+ * A great example is the "jakarta.persistence.Entity" annotation: we'll most likely
+ * need to load it when doing any form of introspection on user's code, but we expect
+ * the bytecode which represents the annotation to not be enhanced.
+ * We then benefit from caching such representations of object types which are frequently
+ * loaded; since caching end user code would lead to enhancement problems, it's best
+ * to keep the list conservative when in doubt.
+ * For example, don't consider all of {@code "org.hibernate."} prefixes as safe, as
+ * that would include entities used during our own testsuite and entities defined by Envers.
+ *
+ */
+public final class CorePrefixFilter {
+
+	private final String[] acceptedPrefixes;
+
+	public CorePrefixFilter() {
+		//By default optimise for jakarta annotations, java util collections, and Hibernate marker interfaces
+		this("jakarta.", "java.", "org.hibernate.annotations.", "org.hibernate.bytecode.enhance.spi.", "org.hibernate.engine.spi.");
+	}
+
+	public CorePrefixFilter(final String... acceptedPrefixes) {
+		this.acceptedPrefixes = Objects.requireNonNull( acceptedPrefixes );
+	}
+
+	boolean isCoreClassName(final String name) {
+		for ( String acceptedPrefix : this.acceptedPrefixes ) {
+			if ( name.startsWith( acceptedPrefix ) ) {
+				return true;
+			}
+		}
+		return false;
+	}
+
+}

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

@@ -6,6 +6,7 @@
  */
 package org.hibernate.bytecode.enhance.internal.bytebuddy;
 
+import java.util.Objects;
 import java.util.concurrent.ConcurrentHashMap;
 
 import net.bytebuddy.description.type.TypeDescription;
@@ -27,42 +28,39 @@ 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;
+	private final CorePrefixFilter acceptedPrefixes;
 
 	/**
-	 * Construct a new {@link CoreTypePool} with its default configuration:
-	 * to only load classes whose package names start with either "jakarta."
-	 * or "java."
+	 * Construct a new {@link CoreTypePool} with its default configuration.
+	 *
+	 * @see CorePrefixFilter
 	 */
 	public CoreTypePool() {
-		//By default optimise for jakarta annotations, java util collections, and Hibernate marker interfaces
-		this("jakarta.", "java.", "org.hibernate.annotations.", "org.hibernate.bytecode.enhance.spi.", "org.hibernate.engine.spi.");
+		this( new CorePrefixFilter() );
 	}
 
 	/**
 	 * Construct a new {@link CoreTypePool} with a choice of which prefixes
 	 * for fully qualified classnames will be loaded by this {@link TypePool}.
+	 *
+	 * @deprecated used by Quarkus
 	 */
+	@Deprecated
 	public CoreTypePool(final String... acceptedPrefixes) {
+		this( new CorePrefixFilter( acceptedPrefixes ) );
+	}
+
+	public CoreTypePool(CorePrefixFilter 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;
+		super( new CoreCacheProvider( acceptedPrefixes ) );
+		this.acceptedPrefixes = Objects.requireNonNull( acceptedPrefixes );
 	}
 
 	@Override
 	protected Resolution doDescribe(final String name) {
-		if ( isCoreClassName( name ) ) {
+		if ( acceptedPrefixes.isCoreClassName( name ) ) {
 			final Resolution resolution = resolutions.get( name );
 			if ( resolution != null ) {
 				return resolution;

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

@@ -165,6 +165,15 @@ public void discoverTypes(String className, byte[] originalBytes) {
 	}
 
 	private DynamicType.Builder<?> doEnhance(Supplier<DynamicType.Builder<?>> builderSupplier, TypeDescription managedCtClass) {
+		// skip if the class was already enhanced. This is very common in WildFly as classloading is highly concurrent.
+		// We need to ensure that no class is instrumented multiple times as that might result in incorrect bytecode.
+		// N.B. there is a second check below using a different approach: checking for the marker interfaces; I'm currently
+		// inclined to keep both checks for better backwards compatibility.
+		if ( managedCtClass.getDeclaredAnnotations().isAnnotationPresent( EnhancementInfo.class ) ) {
+			verifyVersions( managedCtClass, enhancementContext );
+			log.debugf( "Skipping enhancement of [%s]: it's already annotated with @EnhancementInfo", managedCtClass.getName() );
+			return null;
+		}
 		// can't effectively enhance interfaces
 		if ( managedCtClass.isInterface() ) {
 			log.debugf( "Skipping enhancement of [%s]: it's an interface", managedCtClass.getName() );

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

@@ -20,9 +20,11 @@ public class ModelTypePool extends TypePool.Default implements EnhancerClassLoca
 
 	private final ConcurrentHashMap<String, Resolution> resolutions = new ConcurrentHashMap<>();
 	private final OverridingClassFileLocator locator;
+	private final SafeCacheProvider poolCache;
 
-	private ModelTypePool(CacheProvider cacheProvider, OverridingClassFileLocator classFileLocator, CoreTypePool parent) {
+	private ModelTypePool(SafeCacheProvider cacheProvider, OverridingClassFileLocator classFileLocator, CoreTypePool parent) {
 		super( cacheProvider, classFileLocator, ReaderMode.FAST, parent );
+		this.poolCache = cacheProvider;
 		this.locator = classFileLocator;
 	}
 
@@ -62,7 +64,7 @@ public static EnhancerClassLocator buildModelTypePool(ClassFileLocator classFile
 	 * @return
 	 */
 	public static EnhancerClassLocator buildModelTypePool(ClassFileLocator classFileLocator, CoreTypePool coreTypePool) {
-		return buildModelTypePool( classFileLocator, coreTypePool, new TypePool.CacheProvider.Simple() );
+		return buildModelTypePool( classFileLocator, coreTypePool, new SafeCacheProvider() );
 	}
 
 	/**
@@ -72,7 +74,7 @@ public static EnhancerClassLocator buildModelTypePool(ClassFileLocator classFile
 	 * @param cacheProvider
 	 * @return
 	 */
-	public static EnhancerClassLocator buildModelTypePool(ClassFileLocator classFileLocator, CoreTypePool coreTypePool, CacheProvider cacheProvider) {
+	public static EnhancerClassLocator buildModelTypePool(ClassFileLocator classFileLocator, CoreTypePool coreTypePool, SafeCacheProvider cacheProvider) {
 		Objects.requireNonNull( classFileLocator );
 		Objects.requireNonNull( coreTypePool );
 		Objects.requireNonNull( cacheProvider );
@@ -92,6 +94,13 @@ protected Resolution doDescribe(final String name) {
 
 	@Override
 	public void registerClassNameAndBytes(final String className, final byte[] bytes) {
+		//Very important: ensure the registered override is actually effective in case this class
+		//was already resolved in the recent past; this could have happened for example as a side effect
+		//of symbol resolution during enhancement of a different class, or very simply when attempting
+		//to re-enhanced the same class - which happens frequently in WildFly because of the class transformers
+		//being triggered concurrently by multiple parallel deployments.
+		resolutions.remove( className );
+		poolCache.remove( className );
 		locator.put( className, new ClassFileLocator.Resolution.Explicit( Objects.requireNonNull( bytes ) ) );
 	}
 

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

@@ -15,19 +15,34 @@
 /**
  * Allows wrapping another ClassFileLocator to add the ability to define
  * resolution overrides for specific resources.
+ * This is useful when enhancing entities and we need to process an
+ * input byte array representing the bytecode of the entity, and some
+ * external party is providing the byte array explicitly, avoiding for
+ * us having to load it from the classloader.
+ * We'll still need to load several other symbols from a parent @{@link ClassFileLocator}
+ * (typically based on the classloader), but need to return the provided
+ * byte array for some selected resources.
+ * Any override is scoped to the current thread; this is to avoid
+ * interference among threads in systems which perform parallel
+ * class enhancement, for example containers requiring "on the fly"
+ * transformation of classes as they are being loaded will often
+ * perform such operations concurrently.
  */
-public final class OverridingClassFileLocator implements ClassFileLocator {
+final class OverridingClassFileLocator implements ClassFileLocator {
 
-	private final ConcurrentHashMap<String, Resolution> registeredResolutions = new ConcurrentHashMap<>();
+	private final ThreadLocal<ConcurrentHashMap<String, Resolution>> registeredResolutions = ThreadLocal.withInitial( () -> new ConcurrentHashMap<>() );
 	private final ClassFileLocator parent;
 
-	public OverridingClassFileLocator(final ClassFileLocator parent) {
+	/**
+	 * @param parent the @{@link ClassFileLocator} which will be used to load any resource which wasn't explicitly registered as an override.
+	 */
+	OverridingClassFileLocator(final ClassFileLocator parent) {
 		this.parent = Objects.requireNonNull( parent );
 	}
 
 	@Override
 	public Resolution locate(final String name) throws IOException {
-		final Resolution resolution = registeredResolutions.get( name );
+		final Resolution resolution = getLocalMap().get( name );
 		if ( resolution != null ) {
 			return resolution;
 		}
@@ -36,17 +51,30 @@ public Resolution locate(final String name) throws IOException {
 		}
 	}
 
+	private ConcurrentHashMap<String, Resolution> getLocalMap() {
+		return registeredResolutions.get();
+	}
+
 	@Override
-	public void close() throws IOException {
-		//Nothing to do: we're not responsible for parent
+	public void close() {
+		registeredResolutions.remove();
 	}
 
+	/**
+	 * Registers an explicit resolution override
+	 * @param className
+	 * @param explicit
+	 */
 	void put(String className, Resolution.Explicit explicit) {
-		registeredResolutions.put( className, explicit );
+		getLocalMap().put( className, explicit );
 	}
 
+	/**
+	 * Removes an explicit resolution override
+	 * @param className
+	 */
 	void remove(String className) {
-		registeredResolutions.remove( className );
+		getLocalMap().remove( className );
 	}
 
 }

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

@@ -0,0 +1,49 @@
+/*
+ * 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 java.util.concurrent.ConcurrentMap;
+
+import net.bytebuddy.pool.TypePool;
+
+/**
+ * An implementation of @{@link net.bytebuddy.pool.TypePool.CacheProvider} which scopes
+ * all state to the current thread, and allows to remove specific registrations.
+ * The threadscoping is necessary to resolve a race condition happening during concurrent entity enhancement:
+ * while one thread is resolving metadata about the entity which needs to be enhanced, other threads
+ * might be working on the same operation (or a different entity which needs this one to be resolved)
+ * and the resolution output - potentially shared across them - could be tainted as they do need
+ * to occasionally work on different input because of the specific overrides managed via @{@link OverridingClassFileLocator}.
+ */
+final class SafeCacheProvider implements TypePool.CacheProvider {
+
+	private final ThreadLocal<ConcurrentMap<String, TypePool.Resolution>> delegate = ThreadLocal.withInitial( () -> new ConcurrentHashMap<>() );
+
+	@Override
+	public TypePool.Resolution find(final String name) {
+		return delegate.get().get( name );
+	}
+
+	@Override
+	public TypePool.Resolution register(final String name, final TypePool.Resolution resolution) {
+		final TypePool.Resolution cached = delegate.get().putIfAbsent( name, resolution );
+		return cached == null
+				? resolution
+				: cached;
+	}
+
+	@Override
+	public void clear() {
+		delegate.get().clear();
+		delegate.remove();
+	}
+
+	public TypePool.Resolution remove(final String name) {
+		return delegate.get().remove( name );
+	}
+}

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

@@ -1317,7 +1317,7 @@ public String[] call() {
 	/**
 	 * 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.
+	 * Used in Quarkus.
 	 * @param enhancementContext
 	 * @param classLocator
 	 * @return