HHH-18901 The entity enhancement process made defensive against concurrent enhancement

In particular, made it defensive against concurrent enhancement of the same resource,
or of resources that might trigger loading symbols which another thread is re-processing.

Author: Sanne

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

@@ -0,0 +1,69 @@
+/*
+ * SPDX-License-Identifier: LGPL-2.1-or-later
+ * Copyright Red Hat Inc. and Hibernate Authors
+ */
+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;
+
+/**
+ * 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}
+	 */
+	@Override
+	public TypePool.Resolution find(final String name) {
+		return storage.get( name );
+	}
+
+	/**
+	 * {@inheritDoc}
+	 */
+	@Override
+	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}
+	 */
+	@Override
+	public void clear() {
+		storage.clear();
+	}
+
+}

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

@@ -0,0 +1,55 @@
+/*
+ * SPDX-License-Identifier: LGPL-2.1-or-later
+ * Copyright Red Hat Inc. and Hibernate Authors
+ */
+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 static final CorePrefixFilter DEFAULT_INSTANCE = new CorePrefixFilter();
+
+	/**
+	 * Do not invoke: use DEFAULT_INSTANCE
+	 */
+	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 );
+	}
+
+	public 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

@@ -4,6 +4,7 @@
  */
 package org.hibernate.bytecode.enhance.internal.bytebuddy;
 
+import java.util.Objects;
 import java.util.concurrent.ConcurrentHashMap;
 
 import net.bytebuddy.description.type.TypeDescription;
@@ -25,42 +26,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, and java util collections
-		this("jakarta.", "java.", "org.hibernate.annotations.");
+		this( CorePrefixFilter.DEFAULT_INSTANCE );
 	}
 
 	/**
 	 * 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

@@ -163,6 +163,18 @@ 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,
+		// which does not address the case of extended bytecode enhancement
+		// (because it enhances classes that do not end up with these marker interfaces).
+		// I'm currently inclined to keep both checks, as one is safer and the other has 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() );
@@ -179,7 +191,7 @@ private DynamicType.Builder<?> doEnhance(Supplier<DynamicType.Builder<?>> builde
 		if ( alreadyEnhanced( managedCtClass ) ) {
 			verifyVersions( managedCtClass, enhancementContext );
 
-			log.debugf( "Skipping enhancement of [%s]: already enhanced", managedCtClass.getName() );
+			log.debugf( "Skipping enhancement of [%s]: it's already implementing 'Managed'", managedCtClass.getName() );
 			return null;
 		}
 

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

@@ -18,9 +18,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;
 	}
 
@@ -60,7 +62,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() );
 	}
 
 	/**
@@ -70,7 +72,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 );
@@ -90,6 +92,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

@@ -5,27 +5,42 @@
 package org.hibernate.bytecode.enhance.internal.bytebuddy;
 
 import java.io.IOException;
+import java.util.HashMap;
 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.
+ * 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<HashMap<String, Resolution>> registeredResolutions = ThreadLocal.withInitial( () -> new HashMap<>() );
 	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;
 		}
@@ -34,17 +49,32 @@ public Resolution locate(final String name) throws IOException {
 		}
 	}
 
+	private HashMap<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,47 @@
+/*
+ * SPDX-License-Identifier: LGPL-2.1-or-later
+ * Copyright Red Hat Inc. and Hibernate Authors
+ */
+package org.hibernate.bytecode.enhance.internal.bytebuddy;
+
+import java.util.HashMap;
+import java.util.Map;
+
+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<Map<String, TypePool.Resolution>> delegate = ThreadLocal.withInitial( () -> new HashMap<>() );
+
+	@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 );
+	}
+}