HHH-11147 - Integrate enhanced-proxy support with BatchFetchQueue

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

@@ -10,6 +10,7 @@ import org.hibernate.bytecode.enhance.spi.interceptor.BytecodeLazyAttributeInter
 import org.hibernate.bytecode.enhance.spi.interceptor.LazyAttributeLoadingInterceptor;
 import org.hibernate.bytecode.enhance.spi.interceptor.LazyAttributesMetadata;
 import org.hibernate.engine.spi.EntityKey;
+import org.hibernate.engine.spi.PersistentAttributeInterceptable;
 import org.hibernate.engine.spi.PersistentAttributeInterceptor;
 import org.hibernate.engine.spi.SharedSessionContractImplementor;
 
@@ -36,6 +37,14 @@ public interface BytecodeEnhancementMetadata {
 
 	LazyAttributesMetadata getLazyAttributesMetadata();
 
+	/**
+	 * Create an "enhancement as proxy" instance for the given entity
+	 *
+	 * @apiNote The `addEmptyEntry` parameter is used to avoid creation of `EntityEntry` instances when we
+	 * do not need them. - mainly from StatelessSession
+	 */
+	PersistentAttributeInterceptable createEnhancedProxy(EntityKey keyToLoad, boolean addEmptyEntry, SharedSessionContractImplementor session);
+
 	/**
 	 * Build and inject an interceptor instance into the enhanced entity.
 	 *

hibernate-core/src/main/java/org/hibernate/engine/internal/StatefulPersistenceContext.java

@@ -746,6 +746,11 @@ public class StatefulPersistenceContext implements PersistenceContext {
 		return proxyFor( e.getPersister(), e.getEntityKey(), impl );
 	}
 
+	@Override
+	public void addEnhancedProxy(EntityKey key, PersistentAttributeInterceptable entity) {
+		entitiesByKey.put( key, entity );
+	}
+
 	@Override
 	public Object getCollectionOwner(Serializable key, CollectionPersister collectionPersister) throws MappingException {
 		// todo : we really just need to add a split in the notions of:

hibernate-core/src/main/java/org/hibernate/engine/spi/BatchFetchQueue.java

@@ -127,12 +127,12 @@ public class BatchFetchQueue {
 	 */
 	public void addBatchLoadableEntityKey(EntityKey key) {
 		if ( key.isBatchLoadable() ) {
-			LinkedHashSet<EntityKey> set =  batchLoadableEntityKeys.get( key.getEntityName());
+			final LinkedHashSet<EntityKey> keysForEntity = batchLoadableEntityKeys.computeIfAbsent(
-			if (set == null) {
+					key.getEntityName(),
-				set = new LinkedHashSet<>( 8 );
+					k -> new LinkedHashSet<>( 8 )
-				batchLoadableEntityKeys.put( key.getEntityName(), set);
+			);
-			}
+
-			set.add(key);
+			keysForEntity.add( key );
 		}
 	}
 	

hibernate-core/src/main/java/org/hibernate/engine/spi/PersistenceContext.java

@@ -344,6 +344,12 @@ public interface PersistenceContext {
 	 */
 	Object proxyFor(Object impl) throws HibernateException;
 
+	/**
+	 * Cross between {@link #addEntity(EntityKey, Object)} and {@link #addProxy(EntityKey, Object)}
+	 * for use with enhancement-as-proxy
+	 */
+	void addEnhancedProxy(EntityKey key, PersistentAttributeInterceptable entity);
+
 	/**
 	 * Get the entity that owns this persistent collection
 	 */

hibernate-core/src/main/java/org/hibernate/event/internal/DefaultLoadEventListener.java

@@ -313,31 +313,12 @@ public class DefaultLoadEventListener extends AbstractLockUpgradeEventListener i
 				}
 			}
 
+			// Potentially add a batch-fetch entry into the queue for this entity
+			persistenceContext.getBatchFetchQueue().addBatchLoadableEntityKey( keyToLoad );
+
 			// This is the crux of HHH-11147
 			// create the (uninitialized) entity instance - has only id set
-			final Object entity = persister.getEntityTuplizer().instantiate(
+			return persister.getBytecodeEnhancementMetadata().createEnhancedProxy( keyToLoad, true, event.getSession() );
-					keyToLoad.getIdentifier(),
-					event.getSession()
-			);
-
-			// add the entity instance to the persistence context
-			persistenceContext.addEntity(
-					entity,
-					Status.MANAGED,
-					null,
-					keyToLoad,
-					null,
-					LockMode.NONE,
-					true,
-					persister,
-					true
-			);
-
-			persister.getEntityMetamodel()
-					.getBytecodeEnhancementMetadata()
-					.injectEnhancedEntityAsProxyInterceptor( entity, keyToLoad, event.getSession() );
-
-			return entity;
 		}
 		else {
 			if ( persister.hasProxy() ) {

hibernate-core/src/main/java/org/hibernate/internal/StatelessSessionImpl.java

@@ -303,14 +303,7 @@ public class StatelessSessionImpl extends AbstractSharedSessionContract implemen
 
 				// we cannot use bytecode proxy for entities with subclasses
 				if ( !persister.getEntityMetamodel().hasSubclasses() ) {
-					final Object entity = persister.getEntityTuplizer().instantiate( id, this );
+					return persister.getBytecodeEnhancementMetadata().createEnhancedProxy( entityKey, false, this );
-
-					persister.getEntityMetamodel()
-							.getBytecodeEnhancementMetadata()
-							.injectEnhancedEntityAsProxyInterceptor( entity, entityKey, this );
-
-					getPersistenceContext().addEntity( entityKey, entity );
-					return entity;
 				}
 			}
 

hibernate-core/src/main/java/org/hibernate/tuple/entity/BytecodeEnhancementMetadataNonPojoImpl.java

@@ -12,6 +12,7 @@ import org.hibernate.bytecode.enhance.spi.interceptor.LazyAttributesMetadata;
 import org.hibernate.bytecode.spi.BytecodeEnhancementMetadata;
 import org.hibernate.bytecode.spi.NotInstrumentedException;
 import org.hibernate.engine.spi.EntityKey;
+import org.hibernate.engine.spi.PersistentAttributeInterceptable;
 import org.hibernate.engine.spi.PersistentAttributeInterceptor;
 import org.hibernate.engine.spi.SharedSessionContractImplementor;
 
@@ -68,6 +69,11 @@ public class BytecodeEnhancementMetadataNonPojoImpl implements BytecodeEnhanceme
 		throw new NotInstrumentedException( errorMsg );
 	}
 
+	@Override
+	public PersistentAttributeInterceptable createEnhancedProxy(EntityKey keyToLoad, boolean addEmptyEntry, SharedSessionContractImplementor session) {
+		throw new NotInstrumentedException( errorMsg );
+	}
+
 	@Override
 	public LazyAttributeLoadingInterceptor extractInterceptor(Object entity) throws NotInstrumentedException {
 		throw new NotInstrumentedException( errorMsg );

hibernate-core/src/main/java/org/hibernate/tuple/entity/BytecodeEnhancementMetadataPojoImpl.java

@@ -6,9 +6,11 @@
  */
 package org.hibernate.tuple.entity;
 
+import java.io.Serializable;
 import java.util.Set;
 import java.util.function.Function;
 
+import org.hibernate.LockMode;
 import org.hibernate.bytecode.enhance.spi.interceptor.BytecodeLazyAttributeInterceptor;
 import org.hibernate.bytecode.enhance.spi.interceptor.EnhancementAsProxyLazinessInterceptor;
 import org.hibernate.bytecode.enhance.spi.interceptor.LazyAttributeLoadingInterceptor;
@@ -19,7 +21,9 @@ import org.hibernate.engine.spi.EntityKey;
 import org.hibernate.engine.spi.PersistentAttributeInterceptable;
 import org.hibernate.engine.spi.PersistentAttributeInterceptor;
 import org.hibernate.engine.spi.SharedSessionContractImplementor;
+import org.hibernate.engine.spi.Status;
 import org.hibernate.mapping.PersistentClass;
+import org.hibernate.persister.entity.EntityPersister;
 import org.hibernate.type.CompositeType;
 
 /**
@@ -130,6 +134,45 @@ public class BytecodeEnhancementMetadataPojoImpl implements BytecodeEnhancementM
 		return (LazyAttributeLoadingInterceptor) extractLazyInterceptor( entity );
 	}
 
+	@Override
+	public PersistentAttributeInterceptable createEnhancedProxy(EntityKey entityKey, boolean addEmptyEntry, SharedSessionContractImplementor session) {
+		final EntityPersister persister = entityKey.getPersister();
+		final Serializable identifier = entityKey.getIdentifier();
+
+		// first, instantiate the entity instance to use as the proxy
+		final PersistentAttributeInterceptable entity = (PersistentAttributeInterceptable) persister.getEntityTuplizer().instantiate( identifier, session );
+
+		// add the entity (proxy) instance to the PC
+		session.getPersistenceContext().addEnhancedProxy( entityKey, entity );
+
+		// if requested, add the "holder entry" to the PC
+		if ( addEmptyEntry ) {
+			session.getPersistenceContext().addEntry(
+					entity,
+					Status.MANAGED,
+					// loaded state
+					null,
+					// row-id
+					null,
+					identifier,
+					// version
+					null,
+					LockMode.NONE,
+					// we assume it exists in db
+					true,
+					persister,
+					true
+			);
+		}
+
+		// inject the interceptor
+		persister.getEntityMetamodel()
+				.getBytecodeEnhancementMetadata()
+				.injectEnhancedEntityAsProxyInterceptor( entity, entityKey, session );
+
+		return entity;
+	}
+
 	@Override
 	public LazyAttributeLoadingInterceptor injectInterceptor(
 			Object entity,