HHH-19278 fix some logic in MultiIdEntityLoaders that just looks wrong

I'm not going to make a habit of "fixing" things I don't have a test for, I promise

also various cleanups to MultiIdEntityLoaders and EntityBatchLoaders

Author: gavinking

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

@@ -931,8 +931,8 @@ public Object proxyFor(EntityPersister persister, EntityKey key, Object impl) th
 
 	@Override
 	public Object proxyFor(Object impl) throws HibernateException {
-		final EntityEntry e = getEntry( impl );
-		return e == null ? impl : proxyFor( e.getPersister(), e.getEntityKey(), impl );
+		final EntityEntry entry = getEntry( impl );
+		return entry == null ? impl : proxyFor( entry.getPersister(), entry.getEntityKey(), impl );
 	}
 
 	@Override
@@ -1026,15 +1026,12 @@ public Object getLoadedCollectionOwnerOrNull(PersistentCollection<?> collection)
 		if ( ce == null || ce.getLoadedPersister() == null ) {
 			return null;
 		}
-
-		Object loadedOwner = null;
-		// TODO: an alternative is to check if the owner has changed; if it hasn't then
-		// return collection.getOwner()
-		final Object entityId = getLoadedCollectionOwnerIdOrNull( ce );
-		if ( entityId != null ) {
-			loadedOwner = getCollectionOwner( entityId, ce.getLoadedPersister() );
+		else {
+			// TODO: an alternative is to check if the owner has changed; if it hasn't then
+			// return collection.getOwner()
+			final Object entityId = getLoadedCollectionOwnerIdOrNull( ce );
+			return entityId != null ? getCollectionOwner( entityId, ce.getLoadedPersister() ) : null;
 		}
-		return loadedOwner;
 	}
 
 	@Override
@@ -1052,9 +1049,12 @@ private Object getLoadedCollectionOwnerIdOrNull(CollectionEntry ce) {
 		if ( ce == null || ce.getLoadedKey() == null || ce.getLoadedPersister() == null ) {
 			return null;
 		}
-		// TODO: an alternative is to check if the owner has changed; if it hasn't then
-		// get the ID from collection.getOwner()
-		return ce.getLoadedPersister().getCollectionType().getIdOfOwnerOrNull( ce.getLoadedKey(), session );
+		else {
+			// TODO: an alternative is to check if the owner has changed; if it hasn't then
+			// get the ID from collection.getOwner()
+			return ce.getLoadedPersister().getCollectionType()
+					.getIdOfOwnerOrNull( ce.getLoadedKey(), session );
+		}
 	}
 
 	@Override

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

@@ -8,7 +8,9 @@
 import org.hibernate.LockMode;
 import org.hibernate.LockOptions;
 import org.hibernate.engine.jdbc.spi.JdbcServices;
+import org.hibernate.engine.spi.BatchFetchQueue;
 import org.hibernate.engine.spi.EntityKey;
+import org.hibernate.engine.spi.PersistenceContext;
 import org.hibernate.engine.spi.SessionFactoryImplementor;
 import org.hibernate.event.spi.EventSource;
 import org.hibernate.loader.ast.spi.MultiIdEntityLoader;
@@ -164,11 +166,32 @@ protected static LockOptions lockOptions(MultiIdLoadOptions loadOptions) {
 
 	protected abstract int maxBatchSize(Object[] ids, MultiIdLoadOptions loadOptions);
 
-	protected abstract void handleResults(
+	protected void handleResults(
 			MultiIdLoadOptions loadOptions,
 			EventSource session,
 			List<Integer> elementPositionsLoadedByBatch,
-			List<Object> result);
+			List<Object> results) {
+		final PersistenceContext persistenceContext = session.getPersistenceContext();
+		for ( Integer position : elementPositionsLoadedByBatch ) {
+			// the element value at this position in the results List should be
+			// the EntityKey for that entity - reuse it
+			final EntityKey entityKey = (EntityKey) results.get( position );
+			session.getPersistenceContextInternal().getBatchFetchQueue().removeBatchLoadableEntityKey( entityKey );
+			final Object entity = persistenceContext.getEntity( entityKey );
+			final Object result;
+			if ( entity == null
+				// the entity is locally deleted, and the options ask that we not return such entities
+				|| !loadOptions.isReturnOfDeletedEntitiesEnabled()
+					&& persistenceContext.getEntry( entity ).getStatus().isDeletedOrGone() ) {
+				result = null;
+			}
+			else {
+				result = persistenceContext.proxyFor( entity );
+			}
+			results.set( position, result );
+		}
+	}
+
 
 	protected abstract void loadEntitiesById(
 			List<Object> idsInBatch,
@@ -192,39 +215,36 @@ private boolean isLoadFromCaches(
 			MultiIdLoadOptions loadOptions,
 			EntityKey entityKey,
 			LockOptions lockOptions,
-			List<Object> result, int i,
+			List<Object> results, int i,
 			EventSource session) {
-		Object managedEntity = null;
 
 		if ( loadOptions.isSessionCheckingEnabled() ) {
 			// look for it in the Session first
-			final PersistenceContextEntry persistenceContextEntry =
+			final PersistenceContextEntry entry =
 					loadFromSessionCache( entityKey, lockOptions, GET, session );
-			managedEntity = persistenceContextEntry.entity();
-
-			if ( managedEntity != null
-					&& !loadOptions.isReturnOfDeletedEntitiesEnabled()
-					&& !persistenceContextEntry.isManaged() ) {
-				// put a null in the result
-				result.add( i, null );
+			final Object entity = entry.entity();
+			if ( entity != null ) {
+				// put a null in the results
+				final Object result =
+						loadOptions.isReturnOfDeletedEntitiesEnabled()
+							|| entry.isManaged()
+								? entity : null;
+				results.add( i, result );
 				return true;
 			}
 		}
 
-		if ( managedEntity == null
-				&& loadOptions.isSecondLevelCacheCheckingEnabled() ) {
-			// look for it in the SessionFactory
-			final EntityPersister persister = getLoadable().getEntityPersister();
-			managedEntity = session.loadFromSecondLevelCache( persister, entityKey, null, lockOptions.getLockMode() );
+		if ( loadOptions.isSecondLevelCacheCheckingEnabled() ) {
+			// look for it in the second-level cache
+			final Object entity =
+					loadFromSecondLevelCache( entityKey, lockOptions, session );
+			if ( entity != null ) {
+				results.add( i, entity );
+				return true;
+			}
 		}
 
-		if ( managedEntity != null ) {
-			result.add( i, managedEntity );
-			return true;
-		}
-		else {
-			return false;
-		}
+		return false;
 	}
 
 	protected List<T> unorderedMultiLoad(
@@ -238,6 +258,16 @@ protected List<T> unorderedMultiLoad(
 						(position, entityKey, resolvedRef) -> result.add( (T) resolvedRef ) );
 		if ( !isEmpty( unresolvableIds ) ) {
 			loadEntitiesWithUnresolvedIds( loadOptions, lockOptions, session, unresolvableIds, result );
+			final BatchFetchQueue batchFetchQueue = session.getPersistenceContextInternal().getBatchFetchQueue();
+			final EntityPersister persister = getLoadable().getEntityPersister();
+			for ( Object id : unresolvableIds ) {
+				// skip any of the null padded ids
+				// (actually we could probably even break on the first null)
+				if ( id != null ) {
+					// found or not, remove the key from the batch-fetch queue
+					batchFetchQueue.removeBatchLoadableEntityKey( session.generateEntityKey( id, persister ) );
+				}
+			}
 		}
 		return result;
 	}
@@ -247,7 +277,7 @@ protected abstract void loadEntitiesWithUnresolvedIds(
 			LockOptions lockOptions,
 			EventSource session,
 			Object[] unresolvableIds,
-			List<T> result);
+			List<T> results);
 
 	protected final <R> Object[] resolveInCachesIfEnabled(
 			Object[] ids,
@@ -259,7 +289,7 @@ protected final <R> Object[] resolveInCachesIfEnabled(
 				// the user requested that we exclude ids corresponding to already managed
 				// entities from the generated load SQL. So here we will iterate all
 				// incoming id values and see whether it corresponds to an existing
-				// entity associated with the PC. If it does, we add it to the result
+				// entity associated with the PC. If it does, we add it to the results
 				// list immediately and remove its id from the group of ids to load.
 				// we'll load all of them from the database
 				? resolveInCaches( ids, loadOptions, lockOptions, session, resolutionConsumer )
@@ -319,14 +349,14 @@ private <R> List<Object> loadFromCaches(
 			EventSource session) {
 
 		// look for it in the Session first
-		final PersistenceContextEntry persistenceContextEntry =
+		final PersistenceContextEntry entry =
 				loadFromSessionCache( entityKey, lockOptions, GET, session );
 		final Object sessionEntity;
 		if ( loadOptions.isSessionCheckingEnabled() ) {
-			sessionEntity = persistenceContextEntry.entity();
+			sessionEntity = entry.entity();
 			if ( sessionEntity != null
 					&& !loadOptions.isReturnOfDeletedEntitiesEnabled()
-					&& !persistenceContextEntry.isManaged() ) {
+					&& !entry.isManaged() ) {
 				resolutionConsumer.consume( i, entityKey, null );
 				return unresolvedIds;
 			}
@@ -335,13 +365,10 @@ private <R> List<Object> loadFromCaches(
 			sessionEntity = null;
 		}
 
-		final Object cachedEntity;
-		if ( sessionEntity == null && loadOptions.isSecondLevelCacheCheckingEnabled() ) {
-			cachedEntity = session.loadFromSecondLevelCache( getLoadable().getEntityPersister(), entityKey, null, lockOptions.getLockMode() );
-		}
-		else {
-			cachedEntity = sessionEntity;
-		}
+		final Object cachedEntity =
+				sessionEntity == null && loadOptions.isSecondLevelCacheCheckingEnabled()
+						? loadFromSecondLevelCache( entityKey, lockOptions, session )
+						: sessionEntity;
 
 		if ( cachedEntity != null ) {
 			//noinspection unchecked
@@ -356,4 +383,8 @@ private <R> List<Object> loadFromCaches(
 		return unresolvedIds;
 	}
 
+	private Object loadFromSecondLevelCache(EntityKey entityKey, LockOptions lockOptions, EventSource session) {
+		final EntityPersister persister = getLoadable().getEntityPersister();
+		return session.loadFromSecondLevelCache( persister, entityKey, null, lockOptions.getLockMode() );
+	}
 }

hibernate-core/src/main/java/org/hibernate/loader/ast/internal/AbstractMultiIdEntityLoader.java

@@ -9,6 +9,7 @@
 import java.util.Locale;
 
 import org.hibernate.LockOptions;
+import org.hibernate.engine.spi.BatchFetchQueue;
 import org.hibernate.engine.spi.LoadQueryInfluencers;
 import org.hibernate.engine.spi.SharedSessionContractImplementor;
 import org.hibernate.internal.build.AllowReflection;
@@ -17,14 +18,14 @@
 import org.hibernate.metamodel.mapping.EntityIdentifierMapping;
 import org.hibernate.metamodel.mapping.EntityMappingType;
 import org.hibernate.metamodel.mapping.JdbcMapping;
+import org.hibernate.persister.entity.EntityPersister;
 import org.hibernate.query.spi.QueryOptions;
 import org.hibernate.sql.ast.tree.expression.JdbcParameter;
 import org.hibernate.sql.ast.tree.select.SelectStatement;
 import org.hibernate.sql.exec.internal.JdbcParameterImpl;
 import org.hibernate.sql.exec.spi.JdbcOperationQuerySelect;
 import org.hibernate.sql.exec.spi.JdbcParameterBindings;
 
-import static org.hibernate.engine.internal.BatchFetchQueueHelper.removeBatchLoadableEntityKey;
 import static org.hibernate.loader.ast.internal.MultiKeyLoadHelper.trimIdBatch;
 import static org.hibernate.loader.ast.internal.MultiKeyLoadLogging.MULTI_KEY_LOAD_LOGGER;
 
@@ -40,7 +41,6 @@ public class EntityBatchLoaderArrayParam<T>
 		implements SqlArrayMultiKeyLoader {
 	private final int domainBatchSize;
 
-	private final LoadQueryInfluencers loadQueryInfluencers;
 	private final BasicEntityIdentifierMapping identifierMapping;
 	private final JdbcMapping arrayJdbcMapping;
 	private final JdbcParameter jdbcParameter;
@@ -64,7 +64,6 @@ public EntityBatchLoaderArrayParam(
 			EntityMappingType entityDescriptor,
 			LoadQueryInfluencers loadQueryInfluencers) {
 		super( entityDescriptor, loadQueryInfluencers );
-		this.loadQueryInfluencers = loadQueryInfluencers;
 		this.domainBatchSize = domainBatchSize;
 
 		if ( MULTI_KEY_LOAD_LOGGER.isDebugEnabled() ) {
@@ -133,10 +132,12 @@ protected void initializeEntities(
 					getLoadable().getEntityName(), id, Arrays.toString(idsToInitialize) );
 		}
 
+		final BatchFetchQueue batchFetchQueue = session.getPersistenceContextInternal().getBatchFetchQueue();
+		final EntityPersister persister = getLoadable().getEntityPersister();
 		for ( Object initializedId : idsToInitialize ) {
 			if ( initializedId != null ) {
 				// found or not, remove the key from the batch-fetch queue
-				removeBatchLoadableEntityKey( initializedId, getLoadable(), session );
+				batchFetchQueue.removeBatchLoadableEntityKey( session.generateEntityKey( initializedId, persister ) );
 			}
 		}