HHH-17489 Fix concurrent modifications during post load processing

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

@@ -407,7 +407,7 @@ public class StatefulPersistenceContext implements PersistenceContext {
 		}
 		assert holder.entityInitializer == null || holder.entityInitializer == initializer;
 		holder.entityInitializer = initializer;
-		holder.processingState = processingState;
+		processingState.registerLoadingEntityHolder( holder );
 		return holder;
 	}
 
@@ -423,57 +423,59 @@ public class StatefulPersistenceContext implements PersistenceContext {
 
 	@Override
 	public void postLoad(JdbcValuesSourceProcessingState processingState, Consumer<EntityHolder> holderConsumer) {
-		if ( entitiesByKey == null ) {
+		final Callback callback = processingState.getExecutionContext().getCallback();
+		if ( processingState.getLoadingEntityHolders() != null ) {
+			final EventListenerGroup<PostLoadEventListener> listenerGroup = getSession().getFactory()
+					.getFastSessionServices()
+					.eventListenerGroup_POST_LOAD;
+			final PostLoadEvent postLoadEvent = processingState.getPostLoadEvent();
+			for ( final EntityHolder holder : processingState.getLoadingEntityHolders() ) {
+				processLoadedEntityHolder( holder, listenerGroup, postLoadEvent, callback, holderConsumer );
+			}
+			processingState.getLoadingEntityHolders().clear();
+		}
+		if ( processingState.getReloadedEntityHolders() != null ) {
+			for ( final EntityHolder holder : processingState.getReloadedEntityHolders() ) {
+				processLoadedEntityHolder( holder, null, null, callback, holderConsumer );
+			}
+			processingState.getReloadedEntityHolders().clear();
+		}
+	}
+
+	private void processLoadedEntityHolder(
+			EntityHolder holder,
+			EventListenerGroup<PostLoadEventListener> listenerGroup,
+			PostLoadEvent postLoadEvent,
+			Callback callback,
+			Consumer<EntityHolder> holderConsumer) {
+		if ( holderConsumer != null ) {
+			holderConsumer.accept( holder );
+		}
+		if ( holder.getEntity() == null ) {
+			// It's possible that we tried to load an entity and found out it doesn't exist,
+			// in which case we added an entry with a null proxy and entity.
+			// Remove that empty entry on post load to avoid unwanted side effects
+			entitiesByKey.remove( holder.getEntityKey() );
 			return;
 		}
-		final Callback callback = processingState.getExecutionContext().getCallback();
+		if ( postLoadEvent != null ) {
-		final EventListenerGroup<PostLoadEventListener> listenerGroup = getSession().getFactory()
+			postLoadEvent.reset();
-				.getFastSessionServices()
+			postLoadEvent.setEntity( holder.getEntity() )
-				.eventListenerGroup_POST_LOAD;
+					.setId( holder.getEntityKey().getIdentifier() )
-		final PostLoadEvent postLoadEvent = getSession().isEventSource() ? new PostLoadEvent( getSession().asEventSource() ) : null;
+					.setPersister( holder.getDescriptor() );
-		for ( Iterator<EntityHolderImpl> iterator = entitiesByKey.values().iterator(); iterator.hasNext(); ) {
+			listenerGroup.fireEventOnEachListener(
-			final EntityHolderImpl holder = iterator.next();
+					postLoadEvent,
-			if ( holder.processingState == processingState ) {
+					PostLoadEventListener::onPostLoad
-				if ( holderConsumer != null ) {
+			);
-					holderConsumer.accept( holder );
-				}
-				if ( holder.entity == null ) {
-					// It's possible that we tried to load an entity and found out it doesn't exist,
-					// in which case we added an entry with a null proxy and entity.
-					// Remove that empty entry on post load to avoid unwanted side effects
-					iterator.remove();
-					continue;
-				}
-				if ( postLoadEvent != null ) {
-					postLoadEvent.reset();
-					postLoadEvent.setEntity( holder.entity )
-							.setId( holder.entityKey.getIdentifier() )
-							.setPersister( holder.getDescriptor() );
-					listenerGroup.fireEventOnEachListener(
-							postLoadEvent,
-							PostLoadEventListener::onPostLoad
-					);
-				}
-
-				if ( callback != null ) {
-					callback.invokeAfterLoadActions(
-							holder.getEntity(),
-							holder.getDescriptor(),
-							getSession()
-					);
-					if ( holder.reloaded ) {
-						callback.invokeAfterLoadActions(
-								holder.getEntity(),
-								holder.getDescriptor(),
-								getSession()
-						);
-					}
-				}
-				holder.processingState = null;
-				holder.entityInitializer = null;
-				holder.reloaded = false;
-			}
 		}
+		if ( callback != null ) {
+			callback.invokeAfterLoadActions(
+					holder.getEntity(),
+					holder.getDescriptor(),
+					getSession()
+			);
+		}
+		( (EntityHolderImpl) holder ).entityInitializer = null;
 	}
 
 	@Override
@@ -2133,8 +2135,6 @@ public class StatefulPersistenceContext implements PersistenceContext {
 		Object entity;
 		Object proxy;
 		EntityInitializer entityInitializer;
-		JdbcValuesSourceProcessingState processingState;
-		boolean reloaded;
 		EntityHolderState state;
 
 		private EntityHolderImpl(EntityKey entityKey, EntityPersister descriptor, Object entity, Object proxy) {
@@ -2173,9 +2173,7 @@ public class StatefulPersistenceContext implements PersistenceContext {
 
 		@Override
 		public void markAsReloaded(JdbcValuesSourceProcessingState processingState) {
-			assert this.processingState == null;
+			processingState.registerReloadedEntityHolder( this );
-			this.reloaded = true;
-			this.processingState = processingState;
 		}
 
 		@Override

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

@@ -40,6 +40,8 @@ public class JdbcValuesSourceProcessingStateStandardImpl implements JdbcValuesSo
 	private final ExecutionContext executionContext;
 	private final JdbcValuesSourceProcessingOptions processingOptions;
 
+	private List<EntityHolder> loadingEntityHolders;
+	private List<EntityHolder> reloadedEntityHolders;
 	private Map<EntityUniqueKey, Initializer> initializerByUniquKeyMap;
 	private Map<CollectionKey, LoadingCollectionEntry> loadingCollectionMap;
 	private List<CollectionInitializer> arrayInitializers;
@@ -89,6 +91,32 @@ public class JdbcValuesSourceProcessingStateStandardImpl implements JdbcValuesSo
 		return postLoadEvent;
 	}
 
+	@Override
+	public void registerLoadingEntityHolder(EntityHolder holder) {
+		if ( loadingEntityHolders == null ) {
+			loadingEntityHolders = new ArrayList<>();
+		}
+		loadingEntityHolders.add( holder );
+	}
+
+	@Override
+	public List<EntityHolder> getLoadingEntityHolders() {
+		return loadingEntityHolders;
+	}
+
+	@Override
+	public void registerReloadedEntityHolder(EntityHolder holder) {
+		if ( reloadedEntityHolders == null ) {
+			reloadedEntityHolders = new ArrayList<>();
+		}
+		reloadedEntityHolders.add( holder );
+	}
+
+	@Override
+	public List<EntityHolder> getReloadedEntityHolders() {
+		return reloadedEntityHolders;
+	}
+
 	@Override
 	public void registerInitializer(EntityUniqueKey entityKey, Initializer initializer) {
 		if ( initializerByUniquKeyMap == null ) {

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

@@ -6,7 +6,10 @@
  */
 package org.hibernate.sql.results.jdbc.spi;
 
+import java.util.List;
+
 import org.hibernate.engine.spi.CollectionKey;
+import org.hibernate.engine.spi.EntityHolder;
 import org.hibernate.engine.spi.EntityUniqueKey;
 import org.hibernate.engine.spi.SharedSessionContractImplementor;
 import org.hibernate.event.spi.PostLoadEvent;
@@ -39,13 +42,20 @@ public interface JdbcValuesSourceProcessingState {
 	PreLoadEvent getPreLoadEvent();
 	PostLoadEvent getPostLoadEvent();
 
+	void registerLoadingEntityHolder(EntityHolder holder);
+
+	List<EntityHolder> getLoadingEntityHolders();
+
+	void registerReloadedEntityHolder(EntityHolder holder);
+
+	List<EntityHolder> getReloadedEntityHolders();
+
 	void registerInitializer(
 			EntityUniqueKey entityKey,
 			Initializer initializer);
 
 	Initializer findInitializer(EntityUniqueKey entityKey);
 
-
 	/**
 	 * Find a LoadingCollectionEntry locally to this context.
 	 *

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

@@ -6,44 +6,57 @@
  */
 package org.hibernate.orm.test.jpa.collection;
 
-import java.util.HashSet;
+import java.util.List;
-import java.util.Set;
 
-import org.hibernate.testing.TestForIssue;
 import org.hibernate.testing.orm.junit.EntityManagerFactoryScope;
+import org.hibernate.testing.orm.junit.Jira;
 import org.hibernate.testing.orm.junit.Jpa;
+import org.junit.jupiter.api.AfterAll;
+import org.junit.jupiter.api.BeforeAll;
 import org.junit.jupiter.api.Test;
 
-import static org.junit.Assert.assertEquals;
+import jakarta.persistence.criteria.CriteriaQuery;
+import jakarta.persistence.criteria.Root;
 
-@TestForIssue(jiraKey = "HHH-6043")
+import static org.junit.jupiter.api.Assertions.assertEquals;
-@Jpa(
+
-		annotatedClasses = { Child.class, Parent.class }
+@Jpa( annotatedClasses = { Child.class, Parent.class } )
-)
 public class PostLoadTest {
+	@BeforeAll
+	public void setUp(EntityManagerFactoryScope scope) {
+		scope.inTransaction( entityManager -> {
+			final Parent parent1 = new Parent();
+			parent1.setId( 1 );
+			final Child child1 = new Child();
+			child1.setId( 1 );
+			child1.setDaddy( parent1 );
+			entityManager.persist( parent1 );
+			entityManager.persist( child1 );
+			final Parent parent2 = new Parent();
+			parent2.setId( 2 );
+			final Child child2 = new Child();
+			child2.setId( 2 );
+			child2.setDaddy( parent2 );
+			entityManager.persist( parent2 );
+			entityManager.persist( child2 );
+		} );
+	}
+
+	@AfterAll
+	public void tearDown(EntityManagerFactoryScope scope) {
+		scope.inTransaction( entityManager -> {
+			entityManager.createQuery( "delete from Child" ).executeUpdate();
+			entityManager.createQuery( "delete from Parent" ).executeUpdate();
+		} );
+	}
 
 	/**
 	 * Load an entity with a collection of associated entities, that uses a @PostLoad method to
 	 * access the association.
 	 */
 	@Test
+	@Jira( "https://hibernate.atlassian.net/browse/HHH-6043" )
 	public void testAccessAssociatedSetInPostLoad(EntityManagerFactoryScope scope) {
-
-		scope.inTransaction(
-				entityManager -> {
-					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.persist( daddy );
-				}
-		);
-
 		scope.inEntityManager(
 				entityManager -> {
 					Parent daddy = entityManager.find( Parent.class, 1 );
@@ -52,4 +65,16 @@ public class PostLoadTest {
 		);
 	}
 
+	@Test
+	@Jira( "https://hibernate.atlassian.net/browse/HHH-17489" )
+	public void testAccessAssociatedSetInPostLoadQuery(EntityManagerFactoryScope scope) {
+		scope.inEntityManager( entityManager -> {
+			final CriteriaQuery<Parent> cq = entityManager.getCriteriaBuilder().createQuery( Parent.class );
+			final Root<Parent> from = cq.from( Parent.class );
+			final List<Parent> parents = entityManager.createQuery( cq.select( from ) ).getResultList();
+			assertEquals( 2, parents.size() );
+			assertEquals( 1, parents.get( 0 ).getNrOfChildren() );
+			assertEquals( 1, parents.get( 1 ).getNrOfChildren() );
+		} );
+	}
 }