HHH-18689 Maintain proxy targets when converting cache entries

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

@@ -5,10 +5,16 @@
 package org.hibernate.engine.internal;
 
 import org.hibernate.LockMode;
+import org.hibernate.engine.spi.EntityEntry;
+import org.hibernate.engine.spi.EntityHolder;
 import org.hibernate.engine.spi.EntityKey;
+import org.hibernate.engine.spi.PersistenceContext;
 import org.hibernate.engine.spi.SharedSessionContractImplementor;
 import org.hibernate.engine.spi.Status;
 import org.hibernate.persister.entity.EntityPersister;
+import org.hibernate.proxy.LazyInitializer;
+
+import static org.hibernate.proxy.HibernateProxy.extractLazyInitializer;
 
 /**
  * Functionality relating to the Hibernate two-phase loading process, that may be reused by persisters
@@ -37,16 +43,27 @@ public final class TwoPhaseLoad {
 			final LockMode lockMode,
 			final Object version,
 			final SharedSessionContractImplementor session) {
-		session.getPersistenceContextInternal().addEntity(
+		final PersistenceContext persistenceContext = session.getPersistenceContextInternal();
+		final EntityHolder entityHolder = persistenceContext.addEntityHolder( key, object );
+		final EntityEntry entityEntry = persistenceContext.addEntry(
 				object,
 				Status.LOADING,
 				null,
-				key,
+				null,
+				key.getIdentifier(),
 				version,
 				lockMode,
 				true,
 				persister,
 				false
 		);
+		entityHolder.setEntityEntry( entityEntry );
+		final Object proxy = entityHolder.getProxy();
+		if ( proxy != null ) {
+			// there is already a proxy for this impl
+			final LazyInitializer lazyInitializer = extractLazyInitializer( proxy );
+			assert lazyInitializer != null;
+			lazyInitializer.setImplementation( object );
+		}
 	}
 }

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

@@ -31,6 +31,7 @@ import org.hibernate.internal.CoreLogging;
 import org.hibernate.internal.CoreMessageLogger;
 import org.hibernate.persister.entity.EntityPersister;
 import org.hibernate.pretty.MessageHelper;
+import org.hibernate.proxy.LazyInitializer;
 import org.hibernate.sql.results.LoadingLogger;
 import org.hibernate.stat.internal.StatsHelper;
 import org.hibernate.stat.spi.StatisticsImplementor;
@@ -377,16 +378,17 @@ public class CacheEntityLoaderHelper {
 		final EntityPersister subclassPersister =
 				factory.getRuntimeMetamodels().getMappingMetamodel()
 						.getEntityDescriptor( entry.getSubclass() );
+		final PersistenceContext persistenceContext = source.getPersistenceContextInternal();
+		final EntityHolder oldHolder = persistenceContext.getEntityHolder( entityKey );
 
 		final Object entity;
 		if ( instanceToLoad != null ) {
 			entity = instanceToLoad;
 		}
 		else {
-			final EntityHolder holder = source.getPersistenceContextInternal().getEntityHolder( entityKey );
+			if ( oldHolder != null && oldHolder.getEntity() != null ) {
-			if ( holder != null && holder.getEntity() != null ) {
 				// Use the entity which might already be
-				entity = holder.getEntity();
+				entity = oldHolder.getEntity();
 			}
 			else {
 				entity = source.instantiate( subclassPersister, entityId );
@@ -407,13 +409,34 @@ public class CacheEntityLoaderHelper {
 		}
 
 		// make it circular-reference safe
-		TwoPhaseLoad.addUninitializedCachedEntity(
+		final EntityHolder holder = persistenceContext.addEntityHolder( entityKey, entity );
-				entityKey,
+		final Object proxy = holder.getProxy();
-				entity,
+		final boolean isReadOnly;
-				subclassPersister,
+		if ( proxy != null ) {
-				LockMode.NONE,
+			// there is already a proxy for this impl
-				entry.getVersion(),
+			// only set the status to read-only if the proxy is read-only
-				source
+			final LazyInitializer lazyInitializer = extractLazyInitializer( proxy );
+			assert lazyInitializer != null;
+			lazyInitializer.setImplementation( entity );
+
+			isReadOnly = lazyInitializer.isReadOnly();
+		}
+		else {
+			isReadOnly = source.isDefaultReadOnly();
+		}
+		holder.setEntityEntry(
+				persistenceContext.addEntry(
+						entity,
+						Status.LOADING,
+						null,
+						null,
+						entityKey.getIdentifier(),
+						entry.getVersion(),
+						LockMode.NONE,
+						true,
+						persister,
+						false
+				)
 		);
 
 		final Type[] types = subclassPersister.getPropertyTypes();
@@ -438,32 +461,20 @@ public class CacheEntityLoaderHelper {
 		final Object version = getVersion( values, subclassPersister );
 		LOG.tracef( "Cached Version : %s", version );
 
-		final PersistenceContext persistenceContext = source.getPersistenceContext();
+		holder.setEntityEntry(
-
+				persistenceContext.addEntry(
-		final Object proxy = persistenceContext.getProxy( entityKey );
+						entity,
-		final boolean isReadOnly;
+						isReadOnly ? Status.READ_ONLY : Status.MANAGED,
-		if ( proxy != null ) {
+						values,
-			// there is already a proxy for this impl
+						null,
-			// only set the status to read-only if the proxy is read-only
+						entityId,
-			isReadOnly = extractLazyInitializer( proxy ).isReadOnly();
+						version,
-		}
+						LockMode.NONE,
-		else {
+						true,
-			isReadOnly = source.isDefaultReadOnly();
+						subclassPersister,
-		}
+						false
-
+				)
-		final EntityEntry entityEntry = persistenceContext.addEntry(
-				entity,
-				isReadOnly ? Status.READ_ONLY : Status.MANAGED,
-				values,
-				null,
-				entityId,
-				version,
-				LockMode.NONE,
-				true,
-				subclassPersister,
-				false
 		);
-		persistenceContext.getEntityHolder( entityKey ).setEntityEntry( entityEntry );
 		subclassPersister.afterInitialize( entity, source );
 		persistenceContext.initializeNonLazyCollections();