finish big refactor of AbstractMultiIdEntityLoader and children

Signed-off-by: Gavin King <[email protected]>

Author: gavinking

hibernate-core/src/main/java/org/hibernate/internal/util/collections/CollectionHelper.java

@@ -262,7 +262,7 @@ public static <T> Set<T> makeCopy(Set<T> source) {
 		return copy;
 	}
 
-	public static boolean isEmpty(Collection collection) {
+	public static boolean isEmpty(Collection<?> collection) {
 		return collection == null || collection.isEmpty();
 	}
 

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

@@ -4,6 +4,7 @@
  */
 package org.hibernate.loader.ast.internal;
 
+import org.checkerframework.checker.nullness.qual.NonNull;
 import org.hibernate.LockMode;
 import org.hibernate.LockOptions;
 import org.hibernate.engine.jdbc.spi.JdbcServices;
@@ -20,10 +21,12 @@
 import org.hibernate.sql.exec.spi.JdbcSelectExecutor;
 import org.hibernate.type.descriptor.java.JavaType;
 
+import java.lang.reflect.Array;
 import java.util.ArrayList;
 import java.util.List;
 
 import static org.hibernate.internal.util.collections.CollectionHelper.arrayList;
+import static org.hibernate.internal.util.collections.CollectionHelper.isEmpty;
 import static org.hibernate.loader.ast.internal.CacheEntityLoaderHelper.loadFromSessionCacheStatic;
 import static org.hibernate.loader.ast.internal.LoaderHelper.getReadOnlyFromLoadQueryInfluencers;
 import static org.hibernate.loader.ast.internal.MultiKeyLoadLogging.MULTI_KEY_LOAD_LOGGER;
@@ -37,11 +40,13 @@ public abstract class AbstractMultiIdEntityLoader<T> implements MultiIdEntityLoa
 	private final EntityMappingType entityDescriptor;
 	private final SessionFactoryImplementor sessionFactory;
 	private final EntityIdentifierMapping identifierMapping;
+	protected final Object[] idArray;
 
 	public AbstractMultiIdEntityLoader(EntityMappingType entityDescriptor, SessionFactoryImplementor sessionFactory) {
 		this.entityDescriptor = entityDescriptor;
 		this.sessionFactory = sessionFactory;
 		identifierMapping = getLoadable().getIdentifierMapping();
+		idArray = (Object[]) Array.newInstance( identifierMapping.getJavaType().getJavaTypeClass(), 0 );
 	}
 
 	protected EntityMappingType getEntityDescriptor() {
@@ -76,31 +81,43 @@ public EntityMappingType getLoadable() {
 	@Override
 	public final <K> List<T> load(K[] ids, MultiIdLoadOptions loadOptions, EventSource session) {
 		assert ids != null;
-		if ( loadOptions.isOrderReturnEnabled() ) {
-			return performOrderedMultiLoad( ids, loadOptions, session );
-		}
-		else {
-			return performUnorderedMultiLoad( ids, loadOptions, session );
+		return loadOptions.isOrderReturnEnabled()
+				? performOrderedMultiLoad( ids, loadOptions, session )
+				: performUnorderedMultiLoad( ids, loadOptions, session );
+	}
+
+	private List<T> performUnorderedMultiLoad(
+			Object[] ids,
+			MultiIdLoadOptions loadOptions,
+			EventSource session) {
+		assert !loadOptions.isOrderReturnEnabled();
+		assert ids != null;
+		if ( MULTI_KEY_LOAD_LOGGER.isTraceEnabled() ) {
+			MULTI_KEY_LOAD_LOGGER.tracef( "#performUnorderedMultiLoad(`%s`, ..)", getLoadable().getEntityName() );
 		}
+		return unorderedMultiLoad( ids, loadOptions, lockOptions( loadOptions ), session );
 	}
 
 	protected List<T> performOrderedMultiLoad(
 			Object[] ids,
 			MultiIdLoadOptions loadOptions,
 			EventSource session) {
+		assert loadOptions.isOrderReturnEnabled();
+		assert ids != null;
 		if ( MULTI_KEY_LOAD_LOGGER.isTraceEnabled() ) {
 			MULTI_KEY_LOAD_LOGGER.tracef( "#performOrderedMultiLoad(`%s`, ..)", getLoadable().getEntityName() );
 		}
+		return orderedMultiLoad( ids, loadOptions, lockOptions( loadOptions ), session );
+	}
 
-		assert loadOptions.isOrderReturnEnabled();
-
-		final boolean coerce = !getSessionFactory().getJpaMetamodel().getJpaCompliance().isLoadByIdComplianceEnabled();
+	private List<T> orderedMultiLoad(
+			Object[] ids,
+			MultiIdLoadOptions loadOptions,
+			LockOptions lockOptions,
+			EventSource session) {
+		final boolean idCoercionEnabled = isIdCoercionEnabled();
 		final JavaType<?> idType = getLoadable().getIdentifierMapping().getJavaType();
 
-		final LockOptions lockOptions = loadOptions.getLockOptions() == null
-				? new LockOptions( LockMode.NONE )
-				: loadOptions.getLockOptions();
-
 		final int maxBatchSize = maxBatchSize( ids, loadOptions );
 
 		final List<Object> result = arrayList( ids.length );
@@ -109,10 +126,10 @@ protected List<T> performOrderedMultiLoad(
 		final List<Integer> elementPositionsLoadedByBatch = new ArrayList<>();
 
 		for ( int i = 0; i < ids.length; i++ ) {
-			final Object id = coerce ? idType.coerce( ids[i], session ) : ids[i];
+			final Object id = idCoercionEnabled ? idType.coerce( ids[i], session ) : ids[i];
 			final EntityKey entityKey = new EntityKey( id, getLoadable().getEntityPersister() );
 
-			if ( !loadFromCaches( loadOptions, session, id, lockOptions, entityKey, result, i ) ) {
+			if ( !loadFromEnabledCaches( loadOptions, session, id, lockOptions, entityKey, result, i ) ) {
 				// if we did not hit any of the continues above,
 				// then we need to batch load the entity state.
 				idsInBatch.add( id );
@@ -142,13 +159,27 @@ protected List<T> performOrderedMultiLoad(
 		return (List<T>) result;
 	}
 
+	protected static LockOptions lockOptions(MultiIdLoadOptions loadOptions) {
+		return loadOptions.getLockOptions() == null
+				? new LockOptions( LockMode.NONE )
+				: loadOptions.getLockOptions();
+	}
+
 	protected abstract int maxBatchSize(Object[] ids, MultiIdLoadOptions loadOptions);
 
-	protected abstract void handleResults(MultiIdLoadOptions loadOptions, EventSource session, List<Integer> elementPositionsLoadedByBatch, List<Object> result);
+	protected abstract void handleResults(
+			MultiIdLoadOptions loadOptions,
+			EventSource session,
+			List<Integer> elementPositionsLoadedByBatch,
+			List<Object> result);
 
-	protected abstract void loadEntitiesById(List<Object> idsInBatch, LockOptions lockOptions, MultiIdLoadOptions loadOptions, EventSource session);
+	protected abstract void loadEntitiesById(
+			List<Object> idsInBatch,
+			LockOptions lockOptions,
+			MultiIdLoadOptions loadOptions,
+			EventSource session);
 
-	protected boolean loadFromCaches(
+	protected boolean loadFromEnabledCaches(
 			MultiIdLoadOptions loadOptions,
 			EventSource session,
 			Object id,
@@ -157,48 +188,189 @@ protected boolean loadFromCaches(
 			List<Object> result,
 			int i) {
 		if ( loadOptions.isSessionCheckingEnabled() || loadOptions.isSecondLevelCacheCheckingEnabled() ) {
-			final LoadEvent loadEvent = new LoadEvent(
-					id,
-					getLoadable().getJavaType().getJavaTypeClass().getName(),
-					lockOptions,
-					session,
-					getReadOnlyFromLoadQueryInfluencers( session )
+			return loadFromCaches( loadOptions, entityKey, result, i,
+					new LoadEvent(
+							id,
+							getLoadable().getJavaType().getJavaTypeClass().getName(),
+							lockOptions,
+							session,
+							getReadOnlyFromLoadQueryInfluencers( session )
+					)
 			);
+		}
+		else {
+			return false;
+		}
+	}
 
-			Object managedEntity = null;
+	private boolean loadFromCaches(
+			MultiIdLoadOptions loadOptions,
+			EntityKey entityKey,
+			List<Object> result,
+			int i,
+			LoadEvent loadEvent) {
+		Object managedEntity = null;
+
+		if ( loadOptions.isSessionCheckingEnabled() ) {
+			// look for it in the Session first
+			final CacheEntityLoaderHelper.PersistenceContextEntry persistenceContextEntry =
+					loadFromSessionCacheStatic( loadEvent, entityKey, LoadEventListener.GET );
+			managedEntity = persistenceContextEntry.getEntity();
+
+			if ( managedEntity != null
+					&& !loadOptions.isReturnOfDeletedEntitiesEnabled()
+					&& !persistenceContextEntry.isManaged() ) {
+				// put a null in the result
+				result.add( i, null );
+				return true;
+			}
+		}
 
-			if ( loadOptions.isSessionCheckingEnabled() ) {
-				// look for it in the Session first
-				final CacheEntityLoaderHelper.PersistenceContextEntry persistenceContextEntry =
-						loadFromSessionCacheStatic( loadEvent, entityKey, LoadEventListener.GET );
-				managedEntity = persistenceContextEntry.getEntity();
+		if ( managedEntity == null
+				&& loadOptions.isSecondLevelCacheCheckingEnabled() ) {
+			// look for it in the SessionFactory
+			managedEntity = CacheEntityLoaderHelper.INSTANCE.loadFromSecondLevelCache(
+					loadEvent,
+					getLoadable().getEntityPersister(),
+					entityKey
+			);
+		}
 
-				if ( managedEntity != null
-						&& !loadOptions.isReturnOfDeletedEntitiesEnabled()
-						&& !persistenceContextEntry.isManaged() ) {
-					// put a null in the result
-					result.add( i, null );
-					return true;
-				}
-			}
+		if ( managedEntity != null ) {
+			result.add( i, managedEntity );
+			return true;
+		}
+		else {
+			return false;
+		}
+	}
+
+	protected List<T> unorderedMultiLoad(
+			Object[] ids,
+			MultiIdLoadOptions loadOptions,
+			LockOptions lockOptions,
+			EventSource session) {
+		final List<T> result = arrayList( ids.length );
+		final Object[] unresolvableIds =
+				resolveInCachesIfEnabled( ids, loadOptions, lockOptions, session,
+						(position, entityKey, resolvedRef) -> result.add( (T) resolvedRef ) );
+		if ( !isEmpty( unresolvableIds ) ) {
+			loadEntitiesWithUnresolvedIds( loadOptions, lockOptions, session, unresolvableIds, result );
+		}
+		return result;
+	}
+
+	protected abstract void loadEntitiesWithUnresolvedIds(
+			MultiIdLoadOptions loadOptions,
+			LockOptions lockOptions,
+			EventSource session,
+			Object[] unresolvableIds,
+			List<T> result);
+
+	protected final <R> Object[] resolveInCachesIfEnabled(
+			Object[] ids,
+			@NonNull MultiIdLoadOptions loadOptions,
+			@NonNull LockOptions lockOptions,
+			EventSource session,
+			ResolutionConsumer<R> resolutionConsumer) {
+		return loadOptions.isSessionCheckingEnabled() || loadOptions.isSecondLevelCacheCheckingEnabled()
+				// 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
+				// 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 )
+				: ids;
+	}
+
+	protected final <R> Object[] resolveInCaches(
+			Object[] ids,
+			MultiIdLoadOptions loadOptions,
+			LockOptions lockOptions,
+			EventSource session,
+			ResolutionConsumer<R> resolutionConsumer) {
+
+		final boolean idCoercionEnabled = isIdCoercionEnabled();
+		final JavaType<?> idType = getLoadable().getIdentifierMapping().getJavaType();
 
-			if ( managedEntity == null && loadOptions.isSecondLevelCacheCheckingEnabled() ) {
-				// look for it in the SessionFactory
-				managedEntity = CacheEntityLoaderHelper.INSTANCE.loadFromSecondLevelCache(
-						loadEvent,
-						getLoadable().getEntityPersister(),
-						entityKey
-				);
+		List<Object> unresolvedIds = null;
+		for ( int i = 0; i < ids.length; i++ ) {
+			final Object id = idCoercionEnabled ? idType.coerce( ids[i], session ) : ids[i];
+			final EntityKey entityKey = new EntityKey( id, getLoadable().getEntityPersister() );
+			unresolvedIds = loadFromCaches( id, entityKey, i, unresolvedIds, loadOptions, resolutionConsumer,
+					new LoadEvent(
+							id,
+							getLoadable().getJavaType().getJavaTypeClass().getName(),
+							lockOptions,
+							session,
+							getReadOnlyFromLoadQueryInfluencers( session )
+					)
+			);
+		}
+
+		if ( isEmpty( unresolvedIds ) ) {
+			// all the given ids were already associated with the Session
+			return null;
+		}
+		else if ( unresolvedIds.size() == ids.length ) {
+			// we need to load all the ids
+			return ids;
+		}
+		else {
+			// we need to load only some the ids
+			return unresolvedIds.toArray( idArray );
+		}
+	}
+
+	private boolean isIdCoercionEnabled() {
+		return !getSessionFactory().getJpaMetamodel().getJpaCompliance().isLoadByIdComplianceEnabled();
+	}
+
+	public interface ResolutionConsumer<T> {
+		void consume(int position, EntityKey entityKey, T resolvedRef);
+	}
+	private <R, K> List<K> loadFromCaches(
+			K id, EntityKey entityKey, int i,
+			List<K> unresolvedIds,
+			MultiIdLoadOptions loadOptions,
+			ResolutionConsumer<R> resolutionConsumer,
+			LoadEvent loadEvent) {
+		Object cachedEntity = null;
+
+		// look for it in the Session first
+		final CacheEntityLoaderHelper.PersistenceContextEntry persistenceContextEntry =
+				loadFromSessionCacheStatic( loadEvent, entityKey, LoadEventListener.GET );
+		if ( loadOptions.isSessionCheckingEnabled() ) {
+			cachedEntity = persistenceContextEntry.getEntity();
+
+			if ( cachedEntity != null
+					&& !loadOptions.isReturnOfDeletedEntitiesEnabled()
+					&& !persistenceContextEntry.isManaged() ) {
+				resolutionConsumer.consume( i, entityKey, null );
+				return unresolvedIds;
 			}
+		}
 
-			if ( managedEntity != null ) {
-				result.add( i, managedEntity );
-				return true;
+		if ( cachedEntity == null && loadOptions.isSecondLevelCacheCheckingEnabled() ) {
+			cachedEntity = CacheEntityLoaderHelper.INSTANCE.loadFromSecondLevelCache(
+					loadEvent,
+					getLoadable().getEntityPersister(),
+					entityKey
+			);
+		}
+
+		if ( cachedEntity != null ) {
+			//noinspection unchecked
+			resolutionConsumer.consume( i, entityKey, (R) cachedEntity);
+		}
+		else {
+			if ( unresolvedIds == null ) {
+				unresolvedIds = new ArrayList<>();
 			}
+			unresolvedIds.add( id );
 		}
-		return false;
+		return unresolvedIds;
 	}
 
-	protected abstract <K> List<T> performUnorderedMultiLoad(K[] ids, MultiIdLoadOptions loadOptions, EventSource session);
-
 }