Fix initializer post load call timing

hibernate-core/src/main/java/org/hibernate/sql/results/graph/entity/AbstractEntityInitializer.java

@@ -28,8 +28,6 @@ import org.hibernate.event.service.spi.EventListenerGroup;
 import org.hibernate.event.service.spi.EventListenerRegistry;
 import org.hibernate.event.spi.EventSource;
 import org.hibernate.event.spi.EventType;
-import org.hibernate.event.spi.PostLoadEvent;
-import org.hibernate.event.spi.PostLoadEventListener;
 import org.hibernate.event.spi.PreLoadEvent;
 import org.hibernate.event.spi.PreLoadEventListener;
 import org.hibernate.internal.util.StringHelper;
@@ -540,19 +538,18 @@ public abstract class AbstractEntityInitializer extends AbstractFetchParentAcces
 						instance
 				);
 			}
-
-			final LoadingEntityEntry loadingEntry = new LoadingEntityEntry(
-					this,
-					entityKey,
-					concreteDescriptor,
-					instance
-			);
-
-			rowProcessingState.getJdbcValuesSourceProcessingState().registerLoadingEntity(
-					entityKey,
-					loadingEntry
-			);
 		}
+
+		final LoadingEntityEntry loadingEntry = new LoadingEntityEntry(
+				this,
+				entityKey,
+				concreteDescriptor,
+				instance
+		);
+		rowProcessingState.getJdbcValuesSourceProcessingState().registerLoadingEntity(
+				entityKey,
+				loadingEntry
+		);
 		return instance;
 	}
 
@@ -578,11 +575,9 @@ public abstract class AbstractEntityInitializer extends AbstractFetchParentAcces
 			);
 			intializeEntity( instance, rowProcessingState, session, persistenceContext );
 			hibernateLazyInitializer.setImplementation( instance );
-			postLoad( instance, rowProcessingState );
 		}
 		else {
 			intializeEntity( entityInstance, rowProcessingState, session, persistenceContext );
-			postLoad( entityInstance, rowProcessingState );
 		}
 	}
 
@@ -808,29 +803,6 @@ public abstract class AbstractEntityInitializer extends AbstractFetchParentAcces
 		}
 	}
 
-	private void postLoad(Object instance,RowProcessingState rowProcessingState) {
-		final SharedSessionContractImplementor session = rowProcessingState.getJdbcValuesSourceProcessingState().getSession();
-
-		if ( session instanceof EventSource ) {
-			final PostLoadEvent postLoadEvent = rowProcessingState.getJdbcValuesSourceProcessingState().getPostLoadEvent();
-			assert postLoadEvent != null;
-
-			postLoadEvent.reset();
-
-			postLoadEvent.setEntity( instance )
-					.setId( entityKey.getIdentifier() )
-					.setPersister( concreteDescriptor );
-
-			final EventListenerGroup<PostLoadEventListener> listenerGroup = entityDescriptor.getFactory()
-					.getServiceRegistry()
-					.getService( EventListenerRegistry.class )
-					.getEventListenerGroup( EventType.POST_LOAD );
-			for ( PostLoadEventListener listener : listenerGroup.listeners() ) {
-				listener.onPostLoad( postLoadEvent );
-			}
-		}
-	}
-
 	@Override
 	public EntityPersister getConcreteDescriptor() {
 		return concreteDescriptor;

hibernate-core/src/main/java/org/hibernate/sql/results/internal/StandardRowReader.java

@@ -181,9 +181,6 @@ public class StandardRowReader<T> implements RowReader<T> {
 		for ( int i = 0; i < initializers.size(); i++ ) {
 			initializers.get( i ).endLoading( processingState.getExecutionContext() );
 		}
-
-		// todo : use Callback to execute AfterLoadActions
-		// todo : another option is to use Callback to execute the AfterLoadActions after each row
 	}
 
 	@Override

hibernate-core/src/main/java/org/hibernate/sql/results/jdbc/internal/JdbcValuesSourceProcessingStateStandardImpl.java

@@ -15,8 +15,12 @@ import java.util.function.BiConsumer;
 import org.hibernate.engine.spi.CollectionKey;
 import org.hibernate.engine.spi.EntityKey;
 import org.hibernate.engine.spi.SharedSessionContractImplementor;
+import org.hibernate.event.service.spi.EventListenerGroup;
+import org.hibernate.event.service.spi.EventListenerRegistry;
 import org.hibernate.event.spi.EventSource;
+import org.hibernate.event.spi.EventType;
 import org.hibernate.event.spi.PostLoadEvent;
+import org.hibernate.event.spi.PostLoadEventListener;
 import org.hibernate.event.spi.PreLoadEvent;
 import org.hibernate.sql.results.graph.Initializer;
 import org.hibernate.sql.results.graph.collection.internal.ArrayInitializer;
@@ -169,12 +173,37 @@ public class JdbcValuesSourceProcessingStateStandardImpl implements JdbcValuesSo
 
 			// now we can finalize loading collections
 			finishLoadingCollections();
+
+			postLoad();
 		}
 		finally {
 			executionContext.getSession().getPersistenceContext().getLoadContexts().deregister( this );
 		}
 	}
 
+	private void postLoad() {
+		if ( loadingEntityMap == null ) {
+			return;
+		}
+		final EventListenerGroup<PostLoadEventListener> listenerGroup = executionContext.getSession().getFactory()
+				.getServiceRegistry()
+				.getService( EventListenerRegistry.class )
+				.getEventListenerGroup( EventType.POST_LOAD );
+
+		loadingEntityMap.forEach(
+				(entityKey, loadingEntityEntry) -> {
+					postLoadEvent.reset();
+					postLoadEvent.setEntity( loadingEntityEntry.getEntityInstance() )
+							.setId( entityKey.getIdentifier() )
+							.setPersister( loadingEntityEntry.getDescriptor() );
+
+					for ( PostLoadEventListener listener : listenerGroup.listeners() ) {
+						listener.onPostLoad( postLoadEvent );
+					}
+				}
+		);
+	}
+
 	private void finishLoadingArrays() {
 		if ( arrayInitializers != null ) {
 			for ( CollectionInitializer collectionInitializer : arrayInitializers ) {
@@ -188,7 +217,6 @@ public class JdbcValuesSourceProcessingStateStandardImpl implements JdbcValuesSo
 		if ( loadingEntityMap == null ) {
 			return;
 		}
-
 		log.tracev( "Total objects hydrated: {0}", loadingEntityMap.size() );
 	}
 

hibernate-core/src/main/java/org/hibernate/sql/results/spi/ListResultsConsumer.java

@@ -79,11 +79,11 @@ public class ListResultsConsumer<R> implements ResultsConsumer<List<R>, R> {
 				rowProcessingState.finishRowProcessing();
 			}
 			persistenceContext.initializeNonLazyCollections();
+			jdbcValuesSourceProcessingState.finishUp( );
 			return results;
 		}
 		finally {
 			rowReader.finishUp( jdbcValuesSourceProcessingState );
-			jdbcValuesSourceProcessingState.finishUp();
 			jdbcValues.finishUp( session );
 		}
 	}

hibernate-core/src/test/java/org/hibernate/jpa/test/collection/PostLoadTest.java

@@ -1,55 +0,0 @@
-/*
- * 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.jpa.test.collection;
-
-import static org.junit.Assert.assertEquals;
-
-import java.util.HashSet;
-import java.util.Set;
-
-import javax.persistence.EntityManager;
-
-import org.hibernate.jpa.test.BaseEntityManagerFunctionalTestCase;
-import org.hibernate.testing.TestForIssue;
-import org.junit.Test;
-
-@TestForIssue( jiraKey="HHH-6043" )
-public class PostLoadTest extends BaseEntityManagerFunctionalTestCase {
-
-    /**
-     * Load an entity with a collection of associated entities, that uses a @PostLoad method to
-     * access the association.
-     */
-	@Test
-    public void testAccessAssociatedSetInPostLoad() {
-    	Child child = new Child();
-        child.setId(1);
-        Parent daddy = new Parent();
-        daddy.setId(1);
-        child.setDaddy(daddy);
-        Set<Child> children = new HashSet<Child>();
-        children.add(child);
-        daddy.setChildren(children);
-
-        EntityManager em = getOrCreateEntityManager();
-
-        em.getTransaction().begin();
-        em.persist(daddy);
-        em.getTransaction().commit();
-        em.clear();
-
-        daddy = em.find(Parent.class, 1);
-        assertEquals(1, daddy.getNrOfChildren());
-    }
-	
-	@Override
-	protected Class<?>[] getAnnotatedClasses() {
-		return new Class[] { Child.class, Parent.class };
-	}
-}
-
-

hibernate-core/src/test/java/org/hibernate/orm/test/jpa/collection/Child.java

@@ -4,7 +4,7 @@
  * 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.jpa.test.collection;
+package org.hibernate.orm.test.jpa.collection;
 
 import javax.persistence.Entity;
 import javax.persistence.Id;
@@ -27,6 +27,7 @@ public class Child {
     public void setId(Integer id) {
         this.id = id;
     }
+
     @ManyToOne
     public Parent getDaddy() {
         return daddy;

hibernate-core/src/test/java/org/hibernate/orm/test/jpa/collection/Parent.java

@@ -4,7 +4,7 @@
  * 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.jpa.test.collection;
+package org.hibernate.orm.test.jpa.collection;
 
 import java.util.HashSet;
 import java.util.Set;
@@ -35,6 +35,7 @@ public class Parent {
     public void setId(Integer id) {
         this.id = id;
     }
+
     @OneToMany(mappedBy="daddy", fetch=FetchType.EAGER, cascade=CascadeType.ALL)
     public Set<Child> getChildren() {
         return children;

hibernate-core/src/test/java/org/hibernate/orm/test/jpa/collection/PostLoadTest.java

@@ -0,0 +1,61 @@
+/*
+ * 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.orm.test.jpa.collection;
+
+import java.util.HashSet;
+import java.util.Set;
+
+import org.hibernate.testing.TestForIssue;
+import org.hibernate.testing.orm.junit.EntityManagerFactoryScope;
+import org.hibernate.testing.orm.junit.Jpa;
+import org.junit.jupiter.api.Test;
+
+import static org.junit.Assert.assertEquals;
+
+@TestForIssue(jiraKey = "HHH-6043")
+@Jpa(
+		annotatedClasses = { Child.class, Parent.class }
+)
+public class PostLoadTest {
+
+	/**
+	 * Load an entity with a collection of associated entities, that uses a @PostLoad method to
+	 * access the association.
+	 */
+	@Test
+	public void testAccessAssociatedSetInPostLoad(EntityManagerFactoryScope scope) {
+		scope.inEntityManager(
+				entityManager -> {
+					try {
+						Child child = new Child();
+						child.setId( 1 );
+						Parent daddy = new Parent();
+						daddy.setId( 1 );
+						child.setDaddy( daddy );
+						Set<Child> children = new HashSet<>();
+						children.add( child );
+						daddy.setChildren( children );
+
+						entityManager.getTransaction().begin();
+						entityManager.persist( daddy );
+						entityManager.getTransaction().commit();
+						entityManager.clear();
+
+						daddy = entityManager.find( Parent.class, 1 );
+						assertEquals( 1, daddy.getNrOfChildren() );
+					}
+					catch (Exception e) {
+						if ( entityManager.getTransaction().isActive() ) {
+							entityManager.getTransaction().rollback();
+						}
+						throw e;
+					}
+				}
+		);
+	}
+
+}