HHH-12425 Move afterInitialize() phase after the collection initialization

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

@@ -284,8 +284,6 @@ public final class TwoPhaseLoad {
 			persistenceContext.setEntryStatus( entityEntry, Status.MANAGED );
 		}
 
-		persister.afterInitialize( entity, session );
-
 		if ( debugEnabled ) {
 			LOG.debugf(
 					"Done materializing entity %s",
@@ -298,6 +296,22 @@ public final class TwoPhaseLoad {
 		}
 	}
 
+	/**
+	 * Perform the afterInitialize() step. This needs to be done after the collections have been properly initialized
+	 * thus a separate step.
+	 *
+	 * @param entity The entity being loaded
+	 * @param session The Session
+	 */
+	public static void afterInitialize(
+			final Object entity,
+			final SharedSessionContractImplementor session) {
+		final PersistenceContext persistenceContext = session.getPersistenceContext();
+		final EntityEntry entityEntry = persistenceContext.getEntry( entity );
+
+		entityEntry.getPersister().afterInitialize( entity, session );
+	}
+
 	/**
 	 * Check if eager of the association is overriden by anything.
 	 *

hibernate-core/src/main/java/org/hibernate/loader/Loader.java

@@ -1164,6 +1164,12 @@ public abstract class Loader {
 			}
 		}
 
+		if ( hydratedObjects != null ) {
+			for ( Object hydratedObject : hydratedObjects ) {
+				TwoPhaseLoad.afterInitialize( hydratedObject, session );
+			}
+		}
+
 		// Until this entire method is refactored w/ polymorphism, postLoad was
 		// split off from initializeEntity.  It *must* occur after
 		// endCollectionLoad to ensure the collection is in the

hibernate-core/src/main/java/org/hibernate/loader/plan/exec/process/internal/AbstractRowReader.java

@@ -211,6 +211,9 @@ public abstract class AbstractRowReader implements RowReader {
 		// now we can finalize loading collections
 		finishLoadingCollections( context );
 
+		// and trigger the afterInitialize() hooks
+		afterInitialize( context, hydratedEntityRegistrations );
+
 		// finally, perform post-load operations
 		postLoad( postLoadEvent, context, hydratedEntityRegistrations, afterLoadActionList );
 	}
@@ -250,6 +253,17 @@ public abstract class AbstractRowReader implements RowReader {
 		}
 	}
 
+	private void afterInitialize(ResultSetProcessingContextImpl context,
+			List<HydratedEntityRegistration> hydratedEntityRegistrations) {
+		if ( hydratedEntityRegistrations == null ) {
+			return;
+		}
+
+		for ( HydratedEntityRegistration registration : hydratedEntityRegistrations ) {
+			TwoPhaseLoad.afterInitialize( registration.getInstance(), context.getSession() );
+		}
+	}
+
 	private void postLoad(
 			PostLoadEvent postLoadEvent,
 			ResultSetProcessingContextImpl context,

hibernate-core/src/test/java/org/hibernate/test/bytecode/enhancement/lazy/LazyCollectionHandlingTest.java

@@ -0,0 +1,102 @@
+/*
+ * 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.test.bytecode.enhancement.lazy;
+
+import static org.hibernate.testing.transaction.TransactionUtil.doInHibernate;
+
+import java.util.LinkedHashSet;
+import java.util.Set;
+
+import javax.persistence.Entity;
+import javax.persistence.FetchType;
+import javax.persistence.GeneratedValue;
+import javax.persistence.Id;
+import javax.persistence.ManyToMany;
+import javax.persistence.MappedSuperclass;
+
+import org.hibernate.testing.TestForIssue;
+import org.hibernate.testing.bytecode.enhancement.BytecodeEnhancerRunner;
+import org.hibernate.testing.junit4.BaseCoreFunctionalTestCase;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+
+@TestForIssue(jiraKey = "")
+@RunWith(BytecodeEnhancerRunner.class)
+public class LazyCollectionHandlingTest extends BaseCoreFunctionalTestCase {
+
+	private Integer id;
+
+	@Override
+	protected Class<?>[] getAnnotatedClasses() {
+		return new Class[]{
+				JafSid.class, UserGroup.class
+		};
+	}
+
+	@Test
+	public void test() {
+		doInHibernate( this::sessionFactory, s -> {
+			JafSid sid = new JafSid();
+			s.save( sid );
+
+			s.flush();
+			s.clear();
+
+			this.id = sid.getId();
+		});
+
+		doInHibernate( this::sessionFactory, s -> {
+			s.get( JafSid.class, this.id );
+		} );
+	}
+
+	@MappedSuperclass
+	public abstract static class DatabaseEntity {
+		private int id;
+
+		@Id
+		@GeneratedValue
+		public int getId() {
+			return id;
+		}
+
+		public void setId(int id) {
+			this.id = id;
+		}
+
+	}
+
+	@Entity(name = "JafSid")
+	public static class JafSid extends DatabaseEntity {
+
+		private Set<UserGroup> groups = new LinkedHashSet<>();
+
+		@ManyToMany(mappedBy = "members", fetch = FetchType.EAGER)
+		public Set<UserGroup> getGroups() {
+			return groups;
+		}
+
+		public void setGroups(Set<UserGroup> groups) {
+			this.groups = groups;
+		}
+	}
+
+	@Entity(name = "UserGroup")
+	public static class UserGroup extends DatabaseEntity {
+
+		private Set<JafSid> members = new LinkedHashSet<>();
+
+		@ManyToMany
+		public Set<JafSid> getMembers() {
+			return members;
+		}
+
+		public void setMembers(Set<JafSid> members) {
+			this.members = members;
+		}
+	}
+}

hibernate-core/src/test/java/org/hibernate/test/legacy/CustomPersister.java

@@ -356,6 +356,7 @@ public class CustomPersister implements EntityPersister {
 					session,
 					new PreLoadEvent( (EventSource) session )
 			);
+			TwoPhaseLoad.afterInitialize( clone, session );
 			TwoPhaseLoad.postLoad( clone, session, new PostLoadEvent( (EventSource) session ) );
 		}
 		return clone;