HHH-11970 (5.1)

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

@@ -0,0 +1,86 @@
+/*
+ * Hibernate, Relational Persistence for Idiomatic Java
+ *
+ * License: GNU Lesser General Public License (LGPL), version 2.1 or later.
+ * See the lgpl.txt file in the root directory or <http://www.gnu.org/licenses/lgpl-2.1.html>.
+ */
+package org.hibernate.engine.internal;
+
+import java.io.Serializable;
+import java.util.Arrays;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Set;
+
+import org.hibernate.engine.spi.BatchFetchQueue;
+import org.hibernate.engine.spi.EntityKey;
+import org.hibernate.engine.spi.SessionImplementor;
+import org.hibernate.internal.CoreMessageLogger;
+import org.hibernate.persister.entity.EntityPersister;
+
+import org.jboss.logging.Logger;
+
+/**
+ * @author Gail Badner
+ */
+public class BatchFetchQueueHelper {
+	private static final CoreMessageLogger LOG = Logger.getMessageLogger(
+			CoreMessageLogger.class,
+			BatchFetchQueueHelper.class.getName()
+	);
+
+	private BatchFetchQueueHelper(){
+	}
+
+	/**
+	 * Finds the IDs for entities that were not found when the batch was loaded, and removes
+	 * the corresponding entity keys from the {@link BatchFetchQueue}.
+	 *
+	 * @param ids - the IDs for the entities that were batch loaded
+	 * @param results - the results from loading the batch
+	 * @param persister - the entity persister for the entities in batch
+	 * @param session - the session
+	 */
+	public static void removeNotFoundBatchLoadableEntityKeys(
+			Serializable[] ids,
+			List<?> results,
+			EntityPersister persister,
+			SessionImplementor session) {
+		if ( !persister.isBatchLoadable() ) {
+			return;
+		}
+		if ( ids.length == results.size() ) {
+			return;
+		}
+		LOG.debug( "Not all entities were loaded." );
+		Set<Serializable> idSet = new HashSet<Serializable>( Arrays.asList( ids ) );
+		for ( Object result : results ) {
+			// All results should be in the PersistenceContext
+			idSet.remove( session.getContextEntityIdentifier( result ) );
+		}
+		assert idSet.size() == ids.length - results.size();
+		if ( LOG.isDebugEnabled() ) {
+			LOG.debug( "Entities of type [" + persister.getEntityName() + "] not found; IDs: " + idSet );
+		}
+		for ( Serializable id : idSet ) {
+			removeBatchLoadableEntityKey( id, persister, session );
+		}
+	}
+
+	/**
+	 * Remove the entity key with the specified {@code id} and {@code persister} from
+	 * the batch loadable entities {@link BatchFetchQueue}.
+	 *
+	 * @param id - the ID for the entity to be removed
+	 * @param persister - the entity persister
+	 * @param session - the session
+	 */
+	public static void removeBatchLoadableEntityKey(
+			Serializable id,
+			EntityPersister persister,
+			SessionImplementor session) {
+		final EntityKey entityKey = session.generateEntityKey( id, persister );
+		final BatchFetchQueue batchFetchQueue = session.getPersistenceContext().getBatchFetchQueue();
+		batchFetchQueue.removeBatchLoadableEntityKey( entityKey );
+	}
+}

hibernate-core/src/main/java/org/hibernate/loader/entity/BatchingEntityLoader.java

@@ -12,6 +12,7 @@
 import java.util.List;
 
 import org.hibernate.LockOptions;
+import org.hibernate.engine.internal.BatchFetchQueueHelper;
 import org.hibernate.engine.spi.QueryParameters;
 import org.hibernate.engine.spi.SessionImplementor;
 import org.hibernate.loader.Loader;
@@ -97,6 +98,15 @@ protected Object doBatchLoad(
 		try {
 			final List results = loaderToUse.doQueryAndInitializeNonLazyCollections( session, qp, false );
 			log.debug( "Done entity batch load" );
+			// The EntityKey for any entity that is not found will remain in the batch.
+			// Explicitly remove the EntityKeys for entities that were not found to
+			// avoid including them in future batches that get executed.
+			BatchFetchQueueHelper.removeNotFoundBatchLoadableEntityKeys(
+					ids,
+					results,
+					persister(),
+					session
+			);
 			return getObjectFromList(results, id, session);
 		}
 		catch ( SQLException sqle ) {

hibernate-core/src/main/java/org/hibernate/loader/entity/DynamicBatchingEntityLoaderBuilder.java

@@ -18,6 +18,7 @@
 import org.hibernate.LockMode;
 import org.hibernate.LockOptions;
 import org.hibernate.dialect.pagination.LimitHelper;
+import org.hibernate.engine.internal.BatchFetchQueueHelper;
 import org.hibernate.engine.spi.EntityEntry;
 import org.hibernate.engine.spi.EntityKey;
 import org.hibernate.engine.spi.LoadQueryInfluencers;
@@ -343,7 +344,13 @@ public Object load(
 
 			final int numberOfIds = ArrayHelper.countNonNull( batch );
 			if ( numberOfIds <= 1 ) {
-				return singleKeyLoader.load( id, optionalObject, session );
+				final Object result =  singleKeyLoader.load( id, optionalObject, session );
+				if ( result == null ) {
+					// There was no entity with the specified ID. Make sure the EntityKey does not remain
+					// in the batch to avoid including it in future batches that get executed.
+					BatchFetchQueueHelper.removeBatchLoadableEntityKey( id, persister(), session );
+				}
+				return result;
 			}
 
 			final Serializable[] idsToLoad = new Serializable[numberOfIds];
@@ -355,6 +362,12 @@ public Object load(
 
 			QueryParameters qp = buildQueryParameters( id, idsToLoad, optionalObject, lockOptions );
 			List results = dynamicLoader.doEntityBatchFetch( session, qp, idsToLoad );
+
+			// The EntityKey for any entity that is not found will remain in the batch.
+			// Explicitly remove the EntityKeys for entities that were not found to
+			// avoid including them in future batches that get executed.
+			BatchFetchQueueHelper.removeNotFoundBatchLoadableEntityKeys( idsToLoad, results, persister(), session );
+
 			return getObjectFromList( results, id, session );
 		}
 	}

hibernate-core/src/main/java/org/hibernate/loader/entity/LegacyBatchingEntityLoaderBuilder.java

@@ -11,6 +11,7 @@
 
 import org.hibernate.LockMode;
 import org.hibernate.LockOptions;
+import org.hibernate.engine.internal.BatchFetchQueueHelper;
 import org.hibernate.engine.spi.LoadQueryInfluencers;
 import org.hibernate.engine.spi.SessionFactoryImplementor;
 import org.hibernate.engine.spi.SessionImplementor;
@@ -100,10 +101,25 @@ public Object load(Serializable id, Object optionalObject, SessionImplementor se
 							persister(),
 							lockOptions
 					);
+					// The EntityKey for any entity that is not found will remain in the batch.
+					// Explicitly remove the EntityKeys for entities that were not found to
+					// avoid including them in future batches that get executed.
+					BatchFetchQueueHelper.removeNotFoundBatchLoadableEntityKeys(
+							smallBatch,
+							results,
+							persister(),
+							session
+					);
 					return getObjectFromList(results, id, session); //EARLY EXIT
 				}
 			}
-			return ( (UniqueEntityLoader) loaders[batchSizes.length-1] ).load(id, optionalObject, session);
+			final Object result = ( (UniqueEntityLoader) loaders[batchSizes.length-1] ).load(id, optionalObject, session);
+			if ( result == null ) {
+				// There was no entity with the specified ID. Make sure the EntityKey does not remain
+				// in the batch to avoid including it in future batches that get executed.
+				BatchFetchQueueHelper.removeBatchLoadableEntityKey( id, persister(), session );
+			}
+			return result;
 		}
 	}
 

hibernate-core/src/main/java/org/hibernate/loader/entity/PaddedBatchingEntityLoaderBuilder.java

@@ -11,6 +11,7 @@
 import org.hibernate.HibernateException;
 import org.hibernate.LockMode;
 import org.hibernate.LockOptions;
+import org.hibernate.engine.internal.BatchFetchQueueHelper;
 import org.hibernate.engine.spi.LoadQueryInfluencers;
 import org.hibernate.engine.spi.SessionFactoryImplementor;
 import org.hibernate.engine.spi.SessionImplementor;
@@ -96,7 +97,13 @@ public Object load(Serializable id, Object optionalObject, SessionImplementor se
 
 			final int numberOfIds = ArrayHelper.countNonNull( batch );
 			if ( numberOfIds <= 1 ) {
-				return ( (UniqueEntityLoader) loaders[batchSizes.length-1] ).load( id, optionalObject, session );
+				final Object result =  ( (UniqueEntityLoader) loaders[batchSizes.length-1] ).load( id, optionalObject, session );
+				if ( result == null ) {
+					// There was no entity with the specified ID. Make sure the EntityKey does not remain
+					// in the batch to avoid including it in future batches that get executed.
+					BatchFetchQueueHelper.removeBatchLoadableEntityKey( id, persister(), session );
+				}
+				return result;
 			}
 
 			// Uses the first batch-size bigger than the number of actual ids in the batch

hibernate-core/src/main/java/org/hibernate/loader/entity/plan/LegacyBatchingEntityLoaderBuilder.java

@@ -11,6 +11,7 @@
 
 import org.hibernate.LockMode;
 import org.hibernate.LockOptions;
+import org.hibernate.engine.internal.BatchFetchQueueHelper;
 import org.hibernate.engine.spi.LoadQueryInfluencers;
 import org.hibernate.engine.spi.SessionFactoryImplementor;
 import org.hibernate.engine.spi.SessionImplementor;
@@ -106,12 +107,27 @@ public Object load(Serializable id, Object optionalObject, SessionImplementor se
 							persister(),
 							lockOptions
 					);
+					// The EntityKey for any entity that is not found will remain in the batch.
+					// Explicitly remove the EntityKeys for entities that were not found to
+					// avoid including them in future batches that get executed.
+					BatchFetchQueueHelper.removeNotFoundBatchLoadableEntityKeys(
+							smallBatch,
+							results,
+							persister(),
+							session
+					);
+
 					//EARLY EXIT
 					return getObjectFromList( results, id, session );
 				}
 			}
-			return ( loaders[batchSizes.length-1] ).load( id, optionalObject, session, lockOptions );
+			final Object result = ( loaders[batchSizes.length-1] ).load( id, optionalObject, session, lockOptions );
+			if ( result == null ) {
+				// There was no entity with the specified ID. Make sure the EntityKey does not remain
+				// in the batch to avoid including it in future batches that get executed.
+				BatchFetchQueueHelper.removeBatchLoadableEntityKey( id, persister(), session );
+			}
+			return result;
 		}
 	}
-
 }