From 9622f1965b6c7f69394032deba66636778167a90 Mon Sep 17 00:00:00 2001
From: Christian Beikov <christian.beikov@gmail.com>
Date: Mon, 7 Oct 2024 17:06:05 +0200
Subject: [PATCH] HHH-18689 Maintain proxy targets when converting cache
 entries

---
 .../engine/internal/TwoPhaseLoad.java         | 21 ++++-
 .../ast/internal/CacheEntityLoaderHelper.java | 81 +++++++++++--------
 2 files changed, 65 insertions(+), 37 deletions(-)

diff --git a/hibernate-core/src/main/java/org/hibernate/engine/internal/TwoPhaseLoad.java b/hibernate-core/src/main/java/org/hibernate/engine/internal/TwoPhaseLoad.java
index 61427bc973..5094e4bca6 100644
--- a/hibernate-core/src/main/java/org/hibernate/engine/internal/TwoPhaseLoad.java
+++ b/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 );
+		}
 	}
 }
diff --git a/hibernate-core/src/main/java/org/hibernate/loader/ast/internal/CacheEntityLoaderHelper.java b/hibernate-core/src/main/java/org/hibernate/loader/ast/internal/CacheEntityLoaderHelper.java
index 59af5d9303..27c1d73706 100644
--- a/hibernate-core/src/main/java/org/hibernate/loader/ast/internal/CacheEntityLoaderHelper.java
+++ b/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 ( holder != null && holder.getEntity() != null ) {
+			if ( oldHolder != null && oldHolder.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(
-				entityKey,
-				entity,
-				subclassPersister,
-				LockMode.NONE,
-				entry.getVersion(),
-				source
+		final EntityHolder holder = persistenceContext.addEntityHolder( entityKey, entity );
+		final Object proxy = holder.getProxy();
+		final boolean isReadOnly;
+		if ( proxy != null ) {
+			// there is already a proxy for this impl
+			// only set the status to read-only if the proxy is read-only
+			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();
-
-		final Object proxy = persistenceContext.getProxy( entityKey );
-		final boolean isReadOnly;
-		if ( proxy != null ) {
-			// there is already a proxy for this impl
-			// only set the status to read-only if the proxy is read-only
-			isReadOnly = extractLazyInitializer( proxy ).isReadOnly();
-		}
-		else {
-			isReadOnly = source.isDefaultReadOnly();
-		}
-
-		final EntityEntry entityEntry = persistenceContext.addEntry(
-				entity,
-				isReadOnly ? Status.READ_ONLY : Status.MANAGED,
-				values,
-				null,
-				entityId,
-				version,
-				LockMode.NONE,
-				true,
-				subclassPersister,
-				false
+		holder.setEntityEntry(
+				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();