HHH-16952 Discover embeddable types through @Embedded annotation for enhancement in a pre-discovery phase

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

@@ -16,11 +16,13 @@ import org.hibernate.bytecode.enhance.internal.bytebuddy.EnhancerImpl.AnnotatedF
 import org.hibernate.bytecode.enhance.spi.EnhancementContext;
 
 import jakarta.persistence.Embedded;
+import jakarta.persistence.metamodel.Type;
 import net.bytebuddy.description.field.FieldDescription;
 import net.bytebuddy.description.method.MethodDescription;
 import net.bytebuddy.description.type.TypeDescription;
 import net.bytebuddy.dynamic.scaffold.MethodGraph;
 import net.bytebuddy.matcher.ElementMatcher;
+import net.bytebuddy.pool.TypePool;
 
 import static net.bytebuddy.matcher.ElementMatchers.isGetter;
 
@@ -88,6 +90,50 @@ class ByteBuddyEnhancementContext {
 		return enhancementContext.doBiDirectionalAssociationManagement( field );
 	}
 
+	public boolean isDiscoveredType(TypeDescription typeDescription) {
+		return enhancementContext.isDiscoveredType( new UnloadedTypeDescription( typeDescription ) );
+	}
+
+	public void registerDiscoveredType(TypeDescription typeDescription, Type.PersistenceType type) {
+		enhancementContext.registerDiscoveredType( new UnloadedTypeDescription( typeDescription ), type );
+	}
+
+	public void discoverCompositeTypes(TypeDescription managedCtClass, TypePool typePool) {
+		if ( isDiscoveredType( managedCtClass ) ) {
+			return;
+		}
+		final Type.PersistenceType determinedPersistenceType;
+		if ( isEntityClass( managedCtClass ) ) {
+			determinedPersistenceType = Type.PersistenceType.ENTITY;
+		}
+		else if ( isCompositeClass( managedCtClass ) ) {
+			determinedPersistenceType = Type.PersistenceType.EMBEDDABLE;
+		}
+		else if ( isMappedSuperclassClass( managedCtClass ) ) {
+			determinedPersistenceType = Type.PersistenceType.MAPPED_SUPERCLASS;
+		}
+		else {
+			// Default to assuming a basic type if this is not a managed type
+			determinedPersistenceType = Type.PersistenceType.BASIC;
+		}
+		registerDiscoveredType( managedCtClass, determinedPersistenceType );
+		if ( determinedPersistenceType != Type.PersistenceType.BASIC ) {
+			final EnhancerImpl.AnnotatedFieldDescription[] enhancedFields = PersistentAttributeTransformer.collectPersistentFields(
+							managedCtClass,
+							this,
+							typePool
+					)
+					.getEnhancedFields();
+			for ( EnhancerImpl.AnnotatedFieldDescription enhancedField : enhancedFields ) {
+				final TypeDescription type = enhancedField.getType().asErasure();
+				if ( !type.isInterface() && enhancedField.hasAnnotation( Embedded.class ) ) {
+					registerDiscoveredType( type, Type.PersistenceType.EMBEDDABLE );
+				}
+				discoverCompositeTypes( type, typePool );
+			}
+		}
+	}
+
 	Optional<MethodDescription> resolveGetter(FieldDescription fieldDescription) {
 		Map<String, MethodDescription> getters = getterByTypeMap
 				.computeIfAbsent( fieldDescription.getDeclaringType().asErasure(), declaringType -> {

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

@@ -46,6 +46,7 @@ import org.hibernate.internal.CoreMessageLogger;
 import jakarta.persistence.Access;
 import jakarta.persistence.AccessType;
 import jakarta.persistence.Transient;
+import jakarta.persistence.metamodel.Type;
 import net.bytebuddy.asm.Advice;
 import net.bytebuddy.description.annotation.AnnotationDescription;
 import net.bytebuddy.description.annotation.AnnotationList;
@@ -161,6 +162,16 @@ public class EnhancerImpl implements Enhancer {
 		}
 	}
 
+	@Override
+	public void discoverTypes(String className, byte[] originalBytes) {
+		if ( originalBytes != null ) {
+			classFileLocator.setClassNameAndBytes( className, originalBytes );
+		}
+		final TypeDescription typeDescription = typePool.describe( className ).resolve();
+		enhancementContext.registerDiscoveredType( typeDescription, Type.PersistenceType.ENTITY );
+		enhancementContext.discoverCompositeTypes( typeDescription, typePool );
+	}
+
 	private TypePool buildTypePool(final ClassFileLocator classFileLocator) {
 		return TypePool.Default.WithLazyResolution.of( classFileLocator );
 	}

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

@@ -68,6 +68,8 @@ final class FieldAccessEnhancer implements AsmVisitorWrapper.ForDeclaredMethods.
 
 				TypeDescription declaredOwnerType = findDeclaredType( owner );
 				AnnotatedFieldDescription field = findField( declaredOwnerType, name, desc );
+				// try to discover composite types on the fly to support some testing scenarios
+				enhancementContext.discoverCompositeTypes( declaredOwnerType, typePool );
 
 				if ( ( enhancementContext.isEntityClass( declaredOwnerType.asErasure() )
 						|| enhancementContext.isCompositeClass( declaredOwnerType.asErasure() ) )

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

@@ -103,6 +103,10 @@ final class PersistentAttributeTransformer implements AsmVisitorWrapper.ForDecla
 		this.enhancedFields = enhancedFields;
 	}
 
+	public AnnotatedFieldDescription[] getEnhancedFields() {
+		return enhancedFields;
+	}
+
 	public static PersistentAttributeTransformer collectPersistentFields(
 			TypeDescription managedCtClass,
 			ByteBuddyEnhancementContext enhancementContext,

hibernate-core/src/main/java/org/hibernate/bytecode/enhance/spi/DefaultEnhancementContext.java

@@ -6,6 +6,8 @@
  */
 package org.hibernate.bytecode.enhance.spi;
 
+import java.util.concurrent.ConcurrentHashMap;
+
 import jakarta.persistence.Basic;
 import jakarta.persistence.Convert;
 import jakarta.persistence.ElementCollection;
@@ -15,6 +17,7 @@ import jakarta.persistence.ManyToMany;
 import jakarta.persistence.MappedSuperclass;
 import jakarta.persistence.OneToMany;
 import jakarta.persistence.Transient;
+import jakarta.persistence.metamodel.Type;
 
 /**
  * default implementation of EnhancementContext. May be sub-classed as needed.
@@ -23,6 +26,8 @@ import jakarta.persistence.Transient;
  */
 public class DefaultEnhancementContext implements EnhancementContext {
 
+	private final ConcurrentHashMap<String, Type.PersistenceType> discoveredTypes = new ConcurrentHashMap<>();
+
 	/**
 	 * @return the classloader for this class
 	 */
@@ -44,7 +49,8 @@ public class DefaultEnhancementContext implements EnhancementContext {
 	 */
 	@Override
 	public boolean isCompositeClass(UnloadedClass classDescriptor) {
-		return classDescriptor.hasAnnotation( Embeddable.class );
+		return classDescriptor.hasAnnotation( Embeddable.class )
+				|| discoveredTypes.get( classDescriptor.getName() ) == Type.PersistenceType.EMBEDDABLE;
 	}
 
 	/**
@@ -124,4 +130,14 @@ public class DefaultEnhancementContext implements EnhancementContext {
 	public UnloadedField[] order(UnloadedField[] persistentFields) {
 		return persistentFields;
 	}
+
+	@Override
+	public boolean isDiscoveredType(UnloadedClass classDescriptor) {
+		return discoveredTypes.containsKey( classDescriptor.getName() );
+	}
+
+	@Override
+	public void registerDiscoveredType(UnloadedClass classDescriptor, Type.PersistenceType type) {
+		discoveredTypes.put( classDescriptor.getName(), type );
+	}
 }

hibernate-core/src/main/java/org/hibernate/bytecode/enhance/spi/EnhancementContext.java

@@ -6,6 +6,8 @@
  */
 package org.hibernate.bytecode.enhance.spi;
 
+import jakarta.persistence.metamodel.Type;
+
 /**
  * The context for performing an enhancement.  Enhancement can happen in any number of ways:<ul>
  * <li>Build time, via Ant</li>
@@ -140,4 +142,8 @@ public interface EnhancementContext {
 	 * @return {@code true} if the field is mapped
 	 */
 	boolean isMappedCollection(UnloadedField field);
+
+	boolean isDiscoveredType(UnloadedClass classDescriptor);
+
+	void registerDiscoveredType(UnloadedClass classDescriptor, Type.PersistenceType type);
 }

hibernate-core/src/main/java/org/hibernate/bytecode/enhance/spi/EnhancementContextWrapper.java

@@ -6,6 +6,8 @@
  */
 package org.hibernate.bytecode.enhance.spi;
 
+import jakarta.persistence.metamodel.Type;
+
 public class EnhancementContextWrapper implements EnhancementContext {
 
 	private final ClassLoader loadingClassloader;
@@ -75,4 +77,14 @@ public class EnhancementContextWrapper implements EnhancementContext {
 	public boolean isMappedCollection(UnloadedField field) {
 		return wrappedContext.isMappedCollection( field );
 	}
+
+	@Override
+	public boolean isDiscoveredType(UnloadedClass classDescriptor) {
+		return wrappedContext.isDiscoveredType( classDescriptor );
+	}
+
+	@Override
+	public void registerDiscoveredType(UnloadedClass classDescriptor, Type.PersistenceType type) {
+		wrappedContext.registerDiscoveredType( classDescriptor, type );
+	}
 }

hibernate-core/src/main/java/org/hibernate/bytecode/enhance/spi/Enhancer.java

@@ -30,4 +30,6 @@ public interface Enhancer {
 	 * @throws EnhancementException Indicates a problem performing the enhancement
 	 */
 	byte[] enhance(String className, byte[] originalBytes) throws EnhancementException;
+
+	void discoverTypes(String className, byte[] originalBytes);
 }

hibernate-core/src/main/java/org/hibernate/bytecode/spi/ClassTransformer.java

@@ -42,4 +42,6 @@ public interface ClassTransformer extends jakarta.persistence.spi.ClassTransform
 			@Nullable Class<?> classBeingRedefined,
 			ProtectionDomain protectionDomain,
 			byte[] classfileBuffer)  throws TransformerException;
+
+	void discoverTypes(ClassLoader loader, String entityClassName);
 }

hibernate-core/src/main/java/org/hibernate/jpa/boot/internal/EntityManagerFactoryBuilderImpl.java

@@ -52,6 +52,7 @@ import org.hibernate.bytecode.enhance.spi.DefaultEnhancementContext;
 import org.hibernate.bytecode.enhance.spi.EnhancementContext;
 import org.hibernate.bytecode.enhance.spi.UnloadedClass;
 import org.hibernate.bytecode.enhance.spi.UnloadedField;
+import org.hibernate.bytecode.spi.ClassTransformer;
 import org.hibernate.cfg.AvailableSettings;
 import org.hibernate.cfg.Environment;
 import org.hibernate.engine.jdbc.connections.spi.ConnectionProvider;
@@ -70,6 +71,7 @@ import org.hibernate.jpa.boot.spi.StrategyRegistrationProviderList;
 import org.hibernate.jpa.boot.spi.TypeContributorList;
 import org.hibernate.jpa.internal.util.LogHelper;
 import org.hibernate.jpa.internal.util.PersistenceUnitTransactionTypeHelper;
+import org.hibernate.mapping.PersistentClass;
 import org.hibernate.proxy.EntityNotFoundDelegate;
 import org.hibernate.resource.transaction.backend.jdbc.internal.JdbcResourceLocalTransactionCoordinatorBuilderImpl;
 import org.hibernate.resource.transaction.backend.jta.internal.JtaTransactionCoordinatorBuilderImpl;
@@ -339,6 +341,20 @@ public class EntityManagerFactoryBuilderImpl implements EntityManagerFactoryBuil
 				);
 
 				persistenceUnit.pushClassTransformer( enhancementContext );
+
+				if ( !persistenceUnit.getClassTransformers().isEmpty() ) {
+					final ClassLoader classLoader = persistenceUnit.getTempClassLoader();
+					if ( classLoader == null ) {
+						throw persistenceException( "Enhancement requires a temp class loader, but none was given." );
+					}
+					for ( ClassTransformer classTransformer : persistenceUnit.getClassTransformers() ) {
+						for ( PersistentClass entityBinding : metadata.getEntityBindings() ) {
+							if ( entityBinding.getClassName() != null ) {
+								classTransformer.discoverTypes( classLoader, entityBinding.getClassName() );
+							}
+						}
+					}
+				}
 			}
 
 			// for the time being we want to revoke access to the temp ClassLoader if one was passed

hibernate-core/src/main/java/org/hibernate/jpa/boot/internal/ParsedPersistenceXmlDescriptor.java

@@ -9,6 +9,8 @@ package org.hibernate.jpa.boot.internal;
 import java.net.URL;
 import java.util.ArrayList;
 import java.util.Arrays;
+import java.util.Collection;
+import java.util.Collections;
 import java.util.List;
 import java.util.Properties;
 import jakarta.persistence.SharedCacheMode;
@@ -16,6 +18,7 @@ import jakarta.persistence.ValidationMode;
 import jakarta.persistence.spi.PersistenceUnitTransactionType;
 
 import org.hibernate.bytecode.enhance.spi.EnhancementContext;
+import org.hibernate.bytecode.spi.ClassTransformer;
 
 /**
  * Describes the information gleaned from a {@code <persistence-unit/>} element in a {@code persistence.xml} file
@@ -189,4 +192,9 @@ public class ParsedPersistenceXmlDescriptor implements org.hibernate.jpa.boot.sp
 	public void pushClassTransformer(EnhancementContext enhancementContext) {
 		// todo : log a message that this is currently not supported...
 	}
+
+	@Override
+	public Collection<ClassTransformer> getClassTransformers() {
+		return Collections.emptyList();
+	}
 }

hibernate-core/src/main/java/org/hibernate/jpa/boot/internal/PersistenceUnitInfoDescriptor.java

@@ -7,6 +7,8 @@
 package org.hibernate.jpa.boot.internal;
 
 import java.net.URL;
+import java.util.ArrayList;
+import java.util.Collection;
 import java.util.List;
 import java.util.Properties;
 import jakarta.persistence.SharedCacheMode;
@@ -15,6 +17,7 @@ import jakarta.persistence.spi.PersistenceUnitInfo;
 import jakarta.persistence.spi.PersistenceUnitTransactionType;
 
 import org.hibernate.bytecode.enhance.spi.EnhancementContext;
+import org.hibernate.bytecode.spi.ClassTransformer;
 import org.hibernate.jpa.boot.spi.PersistenceUnitDescriptor;
 import org.hibernate.jpa.internal.enhance.EnhancingClassTransformerImpl;
 
@@ -23,6 +26,7 @@ import org.hibernate.jpa.internal.enhance.EnhancingClassTransformerImpl;
  */
 public class PersistenceUnitInfoDescriptor implements PersistenceUnitDescriptor {
 	private final PersistenceUnitInfo persistenceUnitInfo;
+	private final ArrayList<ClassTransformer> classTransformers = new ArrayList<>();
 
 	public PersistenceUnitInfoDescriptor(PersistenceUnitInfo persistenceUnitInfo) {
 		this.persistenceUnitInfo = persistenceUnitInfo;
@@ -110,6 +114,16 @@ public class PersistenceUnitInfoDescriptor implements PersistenceUnitDescriptor
 
 	@Override
 	public void pushClassTransformer(EnhancementContext enhancementContext) {
-		persistenceUnitInfo.addTransformer( new EnhancingClassTransformerImpl( enhancementContext ) );
+		// During testing, we will return a null temp class loader in cases where we don't care about enhancement
+		if ( persistenceUnitInfo.getNewTempClassLoader() != null ) {
+			final EnhancingClassTransformerImpl classTransformer = new EnhancingClassTransformerImpl( enhancementContext );
+			classTransformers.add( classTransformer );
+			persistenceUnitInfo.addTransformer( classTransformer );
+		}
+	}
+
+	@Override
+	public Collection<ClassTransformer> getClassTransformers() {
+		return classTransformers;
 	}
 }

hibernate-core/src/main/java/org/hibernate/jpa/boot/spi/PersistenceUnitDescriptor.java

@@ -7,6 +7,7 @@
 package org.hibernate.jpa.boot.spi;
 
 import java.net.URL;
+import java.util.Collection;
 import java.util.List;
 import java.util.Properties;
 import jakarta.persistence.SharedCacheMode;
@@ -14,6 +15,7 @@ import jakarta.persistence.ValidationMode;
 import jakarta.persistence.spi.PersistenceUnitTransactionType;
 
 import org.hibernate.bytecode.enhance.spi.EnhancementContext;
+import org.hibernate.bytecode.spi.ClassTransformer;
 
 /**
  * Abstraction for dealing with either {@code <persistence-unit/>} information whether that comes from
@@ -85,4 +87,6 @@ public interface PersistenceUnitDescriptor {
 	ClassLoader getTempClassLoader();
 
 	void pushClassTransformer(EnhancementContext enhancementContext);
+
+	Collection<ClassTransformer> getClassTransformers();
 }

hibernate-core/src/main/java/org/hibernate/jpa/internal/enhance/EnhancingClassTransformerImpl.java

@@ -6,8 +6,10 @@
  */
 package org.hibernate.jpa.internal.enhance;
 
+import java.lang.ref.WeakReference;
 import java.security.ProtectionDomain;
 import java.util.Objects;
+import java.util.concurrent.locks.ReentrantLock;
 
 import org.hibernate.bytecode.enhance.spi.EnhancementContext;
 import org.hibernate.bytecode.enhance.spi.EnhancementContextWrapper;
@@ -15,7 +17,6 @@ import org.hibernate.bytecode.enhance.spi.Enhancer;
 import org.hibernate.bytecode.internal.BytecodeProviderInitiator;
 import org.hibernate.bytecode.spi.BytecodeProvider;
 import org.hibernate.bytecode.spi.ClassTransformer;
-import org.hibernate.cfg.Environment;
 
 import jakarta.persistence.spi.TransformerException;
 
@@ -27,6 +28,8 @@ public class EnhancingClassTransformerImpl implements ClassTransformer {
 
 	private final EnhancementContext enhancementContext;
 	private final BytecodeProvider bytecodeProvider;
+	private final ReentrantLock lock = new ReentrantLock();
+	private volatile WeakReference<Entry> entryReference;
 
 	public EnhancingClassTransformerImpl(EnhancementContext enhancementContext) {
 		Objects.requireNonNull( enhancementContext );
@@ -42,13 +45,8 @@ public class EnhancingClassTransformerImpl implements ClassTransformer {
 			ProtectionDomain protectionDomain,
 			byte[] classfileBuffer)  throws TransformerException {
 
-		// The first design had the enhancer as a class variable. That approach had some goods and bads.
-		// We don't have to create an enhancer for each class, but on the other end it would stay in memory forever.
-		// It also assumed that all calls come from the same class loader, which is fair, but this makes it more robust.
-
 		try {
-			Enhancer enhancer = bytecodeProvider.getEnhancer( new EnhancementContextWrapper( enhancementContext, loader ) );
+			return getEnhancer( loader ).enhance( className, classfileBuffer );
-			return enhancer.enhance( className, classfileBuffer );
 		}
 		catch (final Exception e) {
 			throw new TransformerException( "Error performing enhancement of " + className, e );
@@ -58,4 +56,51 @@ public class EnhancingClassTransformerImpl implements ClassTransformer {
 		}
 	}
 
+	@Override
+	public void discoverTypes(ClassLoader loader, String entityClassName) {
+		getEnhancer( loader ).discoverTypes( entityClassName, null );
+	}
+
+	private Enhancer getEnhancer(ClassLoader loader) {
+		Entry enhancerEntry = getEnhancerEntry( entryReference, loader );
+		if ( enhancerEntry == null ) {
+			lock.lock();
+			try {
+				enhancerEntry = getEnhancerEntry( entryReference, loader );
+				if ( enhancerEntry == null ) {
+					enhancerEntry = new Entry( loader, createEnhancer( loader ) );
+					entryReference = new WeakReference<>( enhancerEntry );
+				}
+			}
+			finally {
+				lock.unlock();
+			}
+		}
+		return enhancerEntry.enhancer;
+	}
+
+	private static Entry getEnhancerEntry(WeakReference<Entry> weakReference, ClassLoader loader) {
+		if ( weakReference == null ) {
+			return null;
+		}
+		final Entry entry = weakReference.get();
+		if ( entry == null || entry.classLoader != loader ) {
+			return null;
+		}
+		return entry;
+	}
+
+	private Enhancer createEnhancer(ClassLoader loader) {
+		return bytecodeProvider.getEnhancer( new EnhancementContextWrapper( enhancementContext, loader ) );
+	}
+
+	private static class Entry {
+		final ClassLoader classLoader;
+		final Enhancer enhancer;
+
+		public Entry(ClassLoader classLoader, Enhancer enhancer) {
+			this.classLoader = classLoader;
+			this.enhancer = enhancer;
+		}
+	}
 }

hibernate-core/src/test/java/org/hibernate/orm/test/bytecode/enhancement/dirty/DirtyTrackingEmbeddableTest.java

@@ -13,30 +13,44 @@ import org.junit.Test;
 import org.junit.runner.RunWith;
 
 import jakarta.persistence.Embeddable;
+import jakarta.persistence.Embedded;
 import jakarta.persistence.Entity;
 import jakarta.persistence.Id;
 
 @JiraKey( "HHH-16774" )
+@JiraKey( "HHH-16952" )
 @RunWith( BytecodeEnhancerRunner.class )
 public class DirtyTrackingEmbeddableTest {
 
     @Test
     public void test() {
         SimpleEntity entity = new SimpleEntity();
-        Address address = new Address();
+        Address1 address1 = new Address1();
-        entity.address = address;
+        entity.address1 = address1;
+        Address2 address2 = new Address2();
+        entity.address2 = address2;
         EnhancerTestUtils.clearDirtyTracking( entity );
 
         // testing composite object
-        address.city = "Arendal";
+        address1.city = "Arendal";
-        EnhancerTestUtils.checkDirtyTracking( entity, "address" );
+        address2.city = "Arendal";
+        EnhancerTestUtils.checkDirtyTracking( entity, "address1", "address2" );
         EnhancerTestUtils.clearDirtyTracking( entity );
     }
 
     // --- //
 
     @Embeddable
-    private static class Address {
+    private static class Address1 {
+        String street1;
+        String street2;
+        String city;
+        String state;
+        String zip;
+        String phone;
+    }
+
+    private static class Address2 {
         String street1;
         String street2;
         String city;
@@ -53,7 +67,9 @@ public class DirtyTrackingEmbeddableTest {
 
         String name;
 
-        Address address;
+        Address1 address1;
+        @Embedded
+        Address2 address2;
 
     }
 }

hibernate-core/src/test/java/org/hibernate/orm/test/exceptionhandling/BaseJpaOrNativeBootstrapFunctionalTestCase.java

@@ -9,6 +9,8 @@ package org.hibernate.orm.test.exceptionhandling;
 import java.net.URL;
 import java.util.ArrayList;
 import java.util.Arrays;
+import java.util.Collection;
+import java.util.Collections;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
@@ -22,6 +24,7 @@ import org.hibernate.boot.registry.BootstrapServiceRegistryBuilder;
 import org.hibernate.boot.registry.StandardServiceRegistryBuilder;
 import org.hibernate.boot.registry.internal.StandardServiceRegistryImpl;
 import org.hibernate.bytecode.enhance.spi.EnhancementContext;
+import org.hibernate.bytecode.spi.ClassTransformer;
 import org.hibernate.cfg.AvailableSettings;
 import org.hibernate.cfg.Configuration;
 import org.hibernate.cfg.Environment;
@@ -248,6 +251,11 @@ public abstract class BaseJpaOrNativeBootstrapFunctionalTestCase extends BaseUni
 		@Override
 		public void pushClassTransformer(EnhancementContext enhancementContext) {
 		}
+
+		@Override
+		public Collection<ClassTransformer> getClassTransformers() {
+			return Collections.emptyList();
+		}
 	}
 
 	private BootstrapServiceRegistry buildBootstrapServiceRegistry() {

hibernate-core/src/test/java/org/hibernate/orm/test/jpa/BaseEntityManagerFunctionalTestCase.java

@@ -9,6 +9,7 @@ package org.hibernate.orm.test.jpa;
 import java.net.URL;
 import java.util.ArrayList;
 import java.util.Arrays;
+import java.util.Collection;
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.List;
@@ -23,6 +24,7 @@ import jakarta.persistence.spi.PersistenceUnitTransactionType;
 import org.hibernate.boot.registry.internal.StandardServiceRegistryImpl;
 import org.hibernate.boot.spi.MetadataImplementor;
 import org.hibernate.bytecode.enhance.spi.EnhancementContext;
+import org.hibernate.bytecode.spi.ClassTransformer;
 import org.hibernate.cfg.AvailableSettings;
 import org.hibernate.cfg.Environment;
 import org.hibernate.dialect.Dialect;
@@ -187,6 +189,11 @@ public abstract class BaseEntityManagerFunctionalTestCase extends BaseUnitTestCa
 		@Override
 		public void pushClassTransformer(EnhancementContext enhancementContext) {
 		}
+
+		@Override
+		public Collection<ClassTransformer> getClassTransformers() {
+			return Collections.emptyList();
+		}
 	}
 
 	@SuppressWarnings("unchecked")

hibernate-core/src/test/java/org/hibernate/orm/test/jpa/xml/versions/JpaXsdVersionsTest.java

@@ -207,7 +207,7 @@ public class JpaXsdVersionsTest {
 		}
 
 		public ClassLoader getNewTempClassLoader() {
-			return getClassLoader();
+			return null;
 		}
 	}
 }

hibernate-testing/src/main/java/org/hibernate/testing/orm/jpa/PersistenceUnitDescriptorAdapter.java

@@ -7,6 +7,7 @@
 package org.hibernate.testing.orm.jpa;
 
 import java.net.URL;
+import java.util.Collection;
 import java.util.Collections;
 import java.util.List;
 import java.util.Properties;
@@ -16,6 +17,7 @@ import jakarta.persistence.spi.PersistenceUnitTransactionType;
 import javax.sql.DataSource;
 
 import org.hibernate.bytecode.enhance.spi.EnhancementContext;
+import org.hibernate.bytecode.spi.ClassTransformer;
 import org.hibernate.jpa.HibernatePersistenceProvider;
 import org.hibernate.jpa.boot.spi.PersistenceUnitDescriptor;
 
@@ -112,4 +114,9 @@ public class PersistenceUnitDescriptorAdapter implements PersistenceUnitDescript
 	@Override
 	public void pushClassTransformer(EnhancementContext enhancementContext) {
 	}
+
+	@Override
+	public Collection<ClassTransformer> getClassTransformers() {
+		return Collections.emptyList();
+	}
 }

hibernate-testing/src/main/java/org/hibernate/testing/orm/jpa/PersistenceUnitInfoAdapter.java

@@ -99,6 +99,6 @@ public class PersistenceUnitInfoAdapter implements PersistenceUnitInfo {
 	}
 
 	public ClassLoader getNewTempClassLoader() {
-		return Thread.currentThread().getContextClassLoader();
+		return null;
 	}
 }

hibernate-testing/src/main/java/org/hibernate/testing/orm/junit/EntityManagerFactoryBasedFunctionalTest.java

@@ -9,6 +9,8 @@ package org.hibernate.testing.orm.junit;
 import java.net.URL;
 import java.util.ArrayList;
 import java.util.Arrays;
+import java.util.Collection;
+import java.util.Collections;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
@@ -17,6 +19,7 @@ import java.util.function.Consumer;
 import java.util.function.Function;
 
 import org.hibernate.bytecode.enhance.spi.EnhancementContext;
+import org.hibernate.bytecode.spi.ClassTransformer;
 import org.hibernate.cfg.AvailableSettings;
 import org.hibernate.cfg.Environment;
 import org.hibernate.jpa.HibernatePersistenceProvider;
@@ -266,6 +269,11 @@ public class EntityManagerFactoryBasedFunctionalTest
 		@Override
 		public void pushClassTransformer(EnhancementContext enhancementContext) {
 		}
+
+		@Override
+		public Collection<ClassTransformer> getClassTransformers() {
+			return Collections.emptyList();
+		}
 	}
 
 	@AfterEach

hibernate-testing/src/main/java/org/hibernate/testing/util/jpa/DelegatingPersistenceUnitInfo.java

@@ -105,6 +105,6 @@ public class DelegatingPersistenceUnitInfo implements PersistenceUnitInfo {
 
 	@Override
 	public ClassLoader getNewTempClassLoader() {
-		return null;
+		return delegate.getNewTempClassLoader();
 	}
 }

hibernate-testing/src/main/java/org/hibernate/testing/util/jpa/PersistenceUnitInfoPropertiesWrapper.java

@@ -125,6 +125,6 @@ public class PersistenceUnitInfoPropertiesWrapper implements PersistenceUnitInfo
 	}
 
 	public ClassLoader getNewTempClassLoader() {
-		return Thread.currentThread().getContextClassLoader();
+		return null;
 	}
 }

tooling/hibernate-ant/src/main/java/org/hibernate/tool/enhance/EnhancementTask.java

@@ -169,6 +169,10 @@ public class EnhancementTask extends Task {
 		final BytecodeProvider bytecodeProvider = buildDefaultBytecodeProvider();
 		try {
 			Enhancer enhancer = bytecodeProvider.getEnhancer( enhancementContext );
+			for ( File file : sourceSet ) {
+				discoverTypes( file, enhancer );
+				log( "Successfully discovered types for class [" + file + "]", Project.MSG_INFO );
+			}
 			for ( File file : sourceSet ) {
 				byte[] enhancedBytecode = doEnhancement( file, enhancer );
 				if ( enhancedBytecode == null ) {
@@ -203,6 +207,35 @@ public class EnhancementTask extends Task {
 		return new URLClassLoader( urls.toArray( new URL[urls.size()] ), Enhancer.class.getClassLoader() );
 	}
 
+	private void discoverTypes(File javaClassFile, Enhancer enhancer) throws BuildException {
+		try {
+			String className = javaClassFile.getAbsolutePath().substring(
+					base.length() + 1,
+					javaClassFile.getAbsolutePath().length() - ".class".length()
+			).replace( File.separatorChar, '.' );
+			ByteArrayOutputStream originalBytes = new ByteArrayOutputStream();
+			FileInputStream fileInputStream = new FileInputStream( javaClassFile );
+			try {
+				byte[] buffer = new byte[1024];
+				int length;
+				while ( ( length = fileInputStream.read( buffer ) ) != -1 ) {
+					originalBytes.write( buffer, 0, length );
+				}
+			}
+			finally {
+				fileInputStream.close();
+			}
+			enhancer.discoverTypes( className, originalBytes.toByteArray() );
+		}
+		catch (Exception e) {
+			String msg = "Unable to discover types for class: " + javaClassFile.getName();
+			if ( failOnError ) {
+				throw new BuildException( msg, e );
+			}
+			log( msg, e, Project.MSG_WARN );
+		}
+	}
+
 	private byte[] doEnhancement(File javaClassFile, Enhancer enhancer) throws BuildException {
 		try {
 			String className = javaClassFile.getAbsolutePath().substring(

tooling/hibernate-enhance-maven-plugin/src/main/java/org/hibernate/orm/tooling/maven/MavenEnhancePlugin.java

@@ -160,6 +160,12 @@ public class MavenEnhancePlugin extends AbstractMojo {
 		try {
 			final Enhancer enhancer = bytecodeProvider.getEnhancer( enhancementContext );
 
+			for ( File file : sourceSet ) {
+				discoverTypes( file, enhancer );
+				if ( log.isDebugEnabled() ) {
+					log.debug( "Successfully discovered types for class [" + file + "]" );
+				}
+			}
 			for ( File file : sourceSet ) {
 
 				final byte[] enhancedBytecode = doEnhancement( file, enhancer );
@@ -252,6 +258,35 @@ public class MavenEnhancePlugin extends AbstractMojo {
 		}
 	}
 
+	private void discoverTypes(File javaClassFile, Enhancer enhancer) throws MojoExecutionException {
+		try {
+			String className = javaClassFile.getAbsolutePath().substring(
+					base.length() + 1,
+					javaClassFile.getAbsolutePath().length() - ".class".length()
+			).replace( File.separatorChar, '.' );
+			ByteArrayOutputStream originalBytes = new ByteArrayOutputStream();
+			FileInputStream fileInputStream = new FileInputStream( javaClassFile );
+			try {
+				byte[] buffer = new byte[1024];
+				int length;
+				while ( ( length = fileInputStream.read( buffer ) ) != -1 ) {
+					originalBytes.write( buffer, 0, length );
+				}
+			}
+			finally {
+				fileInputStream.close();
+			}
+			enhancer.discoverTypes( className, originalBytes.toByteArray() );
+		}
+		catch (Exception e) {
+			String msg = "Unable to discover types for class: " + javaClassFile.getName();
+			if ( failOnError ) {
+				throw new MojoExecutionException( msg, e );
+			}
+			buildContext.addMessage( javaClassFile, 0, 0, msg, BuildContext.SEVERITY_WARNING, e );
+		}
+	}
+
 	/**
 	 * Expects a directory.
 	 */

tooling/hibernate-gradle-plugin/src/main/java/org/hibernate/orm/tooling/gradle/enhance/EnhancementHelper.java

@@ -22,7 +22,6 @@ import org.hibernate.bytecode.enhance.spi.EnhancementContext;
 import org.hibernate.bytecode.enhance.spi.Enhancer;
 import org.hibernate.bytecode.enhance.spi.UnloadedClass;
 import org.hibernate.bytecode.enhance.spi.UnloadedField;
-import org.hibernate.cfg.Environment;
 import org.hibernate.orm.tooling.gradle.HibernateOrmSpec;
 
 import static org.hibernate.bytecode.internal.BytecodeProviderInitiator.buildDefaultBytecodeProvider;
@@ -49,13 +48,34 @@ public class EnhancementHelper {
 		}
 		final Enhancer enhancer = generateEnhancer( classLoader, ormDsl );
 
-		walk( classesDir, classesDir, enhancer, project );
+		discoverTypes( classesDir, classesDir, enhancer, project );
+		doEnhancement( classesDir, classesDir, enhancer, project );
 	}
 
-	private static void walk(File classesDir, File dir, Enhancer enhancer, Project project) {
+	private static void discoverTypes(File classesDir, File dir, Enhancer enhancer, Project project) {
 		for ( File subLocation : dir.listFiles() ) {
 			if ( subLocation.isDirectory() ) {
-				walk( classesDir, subLocation, enhancer, project );
+				discoverTypes( classesDir, subLocation, enhancer, project );
+			}
+			else if ( subLocation.isFile() && subLocation.getName().endsWith( ".class" ) ) {
+				final String className = determineClassName( classesDir, subLocation );
+				final long lastModified = subLocation.lastModified();
+
+				discoverTypes( subLocation, className, enhancer, project );
+
+				final boolean timestampReset = subLocation.setLastModified( lastModified );
+				if ( !timestampReset ) {
+					project.getLogger().debug( "`{}`.setLastModified failed", project.relativePath( subLocation ) );
+				}
+
+			}
+		}
+	}
+
+	private static void doEnhancement(File classesDir, File dir, Enhancer enhancer, Project project) {
+		for ( File subLocation : dir.listFiles() ) {
+			if ( subLocation.isDirectory() ) {
+				doEnhancement( classesDir, subLocation, enhancer, project );
 			}
 			else if ( subLocation.isFile() && subLocation.getName().endsWith( ".class" ) ) {
 				final String className = determineClassName( classesDir, subLocation );
@@ -72,6 +92,20 @@ public class EnhancementHelper {
 		}
 	}
 
+	private static void discoverTypes(
+			File javaClassFile,
+			String className,
+			Enhancer enhancer,
+			Project project) {
+		try {
+			enhancer.discoverTypes( className, Files.readAllBytes( javaClassFile.toPath() ) );
+			project.getLogger().info( "Successfully discovered types for class : " + className );
+		}
+		catch (Exception e) {
+			throw new GradleException( "Unable to discover types for class : " + className, e );
+		}
+	}
+
 	private static void enhance(
 			File javaClassFile,
 			String className,