1414import org .hibernate .engine .spi .PersistentAttributeInterceptor ;
1515import org .hibernate .internal .CoreLogging ;
1616import org .hibernate .internal .CoreMessageLogger ;
17+ import org .hibernate .internal .util .collections .InstanceIdentityStore ;
18+ import org .hibernate .persister .entity .EntityPersister ;
1719
1820import java .io .IOException ;
1921import java .io .ObjectInputStream ;
@@ -49,7 +51,8 @@ public class EntityEntryContext {
4951
5052 private final transient PersistenceContext persistenceContext ;
5153
52- private transient IdentityHashMap <ManagedEntity ,ImmutableManagedEntityHolder > immutableManagedEntityXref ;
54+ private transient InstanceIdentityStore <ImmutableManagedEntityHolder > immutableManagedEntityXref ;
55+ private transient int currentInstanceId ;
5356
5457 private transient ManagedEntity head ;
5558 private transient ManagedEntity tail ;
@@ -107,21 +110,19 @@ public void addEntityEntry(Object entity, EntityEntry entityEntry) {
107110 else {
108111 // Create a holder for PersistenceContext-related data.
109112 managedEntity = new ImmutableManagedEntityHolder ( managed );
110- if ( immutableManagedEntityXref == null ) {
111- immutableManagedEntityXref = new IdentityHashMap <>();
113+ if ( !isReferenceCachingEnabled ( entityEntry .getPersister () ) ) {
114+ managed .$$_hibernate_setInstanceId ( nextManagedEntityInstanceId () );
115+ putImmutableManagedEntity ( managed , (ImmutableManagedEntityHolder ) managedEntity );
116+ }
117+ else {
118+ // When reference caching is enabled we cannot set the instance-id on the entity instance
119+ putManagedEntity ( entity , managedEntity );
112120 }
113- immutableManagedEntityXref .put (
114- managed ,
115- (ImmutableManagedEntityHolder ) managedEntity
116- );
117121 }
118122 }
119123 else {
120- if ( nonEnhancedEntityXref == null ) {
121- nonEnhancedEntityXref = new IdentityHashMap <>();
122- }
123124 managedEntity = new ManagedEntityImpl ( entity );
124- nonEnhancedEntityXref . put ( entity , managedEntity );
125+ putManagedEntity ( entity , managedEntity );
125126 }
126127 }
127128
@@ -155,6 +156,12 @@ public void addEntityEntry(Object entity, EntityEntry entityEntry) {
155156 }
156157 }
157158
159+ private static boolean isReferenceCachingEnabled (EntityPersister persister ) {
160+ // Immutable entities which can use reference caching are treated as non-enhanced entities, as setting
161+ // the instance-id on them would be problematic in different sessions
162+ return persister .canUseReferenceCacheEntries () && persister .canReadFromCache ();
163+ }
164+
158165 private ManagedEntity getAssociatedManagedEntity (Object entity ) {
159166 if ( isManagedEntity ( entity ) ) {
160167 final ManagedEntity managedEntity = asManagedEntity ( entity );
@@ -170,20 +177,35 @@ private ManagedEntity getAssociatedManagedEntity(Object entity) {
170177 ? managedEntity // it is associated
171178 : null ;
172179 }
173- else {
174- // if managedEntity is associated with this EntityEntryContext, then
175- // it will have an entry in immutableManagedEntityXref and its
176- // holder will be returned.
180+ else if ( !isReferenceCachingEnabled ( entityEntry .getPersister () ) ) {
181+ // if managedEntity is associated with this EntityEntryContext, it may have
182+ // an entry in immutableManagedEntityXref and its holder will be returned.
177183 return immutableManagedEntityXref != null
178- ? immutableManagedEntityXref .get ( managedEntity )
184+ ? immutableManagedEntityXref .get ( managedEntity . $$_hibernate_getInstanceId (), managedEntity )
179185 : null ;
180186 }
181187 }
182- else {
183- return nonEnhancedEntityXref != null
184- ? nonEnhancedEntityXref .get ( entity )
185- : null ;
188+ return nonEnhancedEntityXref != null
189+ ? nonEnhancedEntityXref .get ( entity )
190+ : null ;
191+ }
192+
193+ private void putManagedEntity (Object entity , ManagedEntity managedEntity ) {
194+ if ( nonEnhancedEntityXref == null ) {
195+ nonEnhancedEntityXref = new IdentityHashMap <>();
186196 }
197+ nonEnhancedEntityXref .put ( entity , managedEntity );
198+ }
199+
200+ private int nextManagedEntityInstanceId () {
201+ return currentInstanceId ++;
202+ }
203+
204+ private void putImmutableManagedEntity (ManagedEntity managed , ImmutableManagedEntityHolder holder ) {
205+ if ( immutableManagedEntityXref == null ) {
206+ immutableManagedEntityXref = new InstanceIdentityStore <>();
207+ }
208+ immutableManagedEntityXref .put ( managed , holder );
187209 }
188210
189211 private void checkNotAssociatedWithOtherPersistenceContextIfMutable (ManagedEntity managedEntity ) {
@@ -253,12 +275,16 @@ public EntityEntry removeEntityEntry(Object entity) {
253275
254276 dirty = true ;
255277
256- if (managedEntity instanceof ImmutableManagedEntityHolder ) {
257- assert entity == ( (ImmutableManagedEntityHolder ) managedEntity ).managedEntity ;
258- immutableManagedEntityXref .remove ( entity );
259-
278+ if ( managedEntity instanceof ImmutableManagedEntityHolder holder ) {
279+ assert entity == holder .managedEntity ;
280+ if ( !isReferenceCachingEnabled ( holder .$$_hibernate_getEntityEntry ().getPersister () ) ) {
281+ immutableManagedEntityXref .remove ( managedEntity .$$_hibernate_getInstanceId (), entity );
282+ }
283+ else {
284+ nonEnhancedEntityXref .remove ( entity );
285+ }
260286 }
261- else if ( ! ( isManagedEntity ( entity ) ) ) {
287+ else if ( !isManagedEntity ( entity ) ) {
262288 nonEnhancedEntityXref .remove ( entity );
263289 }
264290
@@ -387,6 +413,7 @@ public void clear() {
387413 count = 0 ;
388414
389415 reentrantSafeEntries = null ;
416+ currentInstanceId = 0 ;
390417 }
391418
392419 private static void clearManagedEntity (final ManagedEntity node ) {
@@ -478,30 +505,26 @@ public static EntityEntryContext deserialize(ObjectInputStream ois, StatefulPers
478505
479506 final ManagedEntity managedEntity ;
480507 if ( isEnhanced ) {
508+ final ManagedEntity castedEntity = asManagedEntity ( entity );
481509 if ( entry .getPersister ().isMutable () ) {
482- managedEntity = ManagedTypeHelper . asManagedEntity ( entity ) ;
510+ managedEntity = castedEntity ;
483511 }
484512 else {
485- final ManagedEntity castedEntity = asManagedEntity ( entity );
486513 managedEntity = new ImmutableManagedEntityHolder ( castedEntity );
487- if ( context .immutableManagedEntityXref == null ) {
488- context .immutableManagedEntityXref =
489- new IdentityHashMap <>();
514+ if ( !isReferenceCachingEnabled ( entry .getPersister () ) ) {
515+ castedEntity .$$_hibernate_setInstanceId ( context .nextManagedEntityInstanceId () );
516+ context .putImmutableManagedEntity ( castedEntity , (ImmutableManagedEntityHolder ) managedEntity );
517+ }
518+ else {
519+ context .putManagedEntity ( entity , castedEntity );
490520 }
491- context .immutableManagedEntityXref .put (
492- castedEntity ,
493- (ImmutableManagedEntityHolder ) managedEntity
494-
495- );
496521 }
497522 }
498523 else {
499524 managedEntity = new ManagedEntityImpl ( entity );
500- if ( context .nonEnhancedEntityXref == null ) {
501- context .nonEnhancedEntityXref = new IdentityHashMap <>();
502- }
503- context .nonEnhancedEntityXref .put ( entity , managedEntity );
525+ context .putManagedEntity ( entity , managedEntity );
504526 }
527+
505528 managedEntity .$$_hibernate_setEntityEntry ( entry );
506529
507530 if ( previous == null ) {
@@ -605,6 +628,15 @@ public ManagedEntityImpl(Object entityInstance) {
605628 public void $$_hibernate_setPreviousManagedEntity (ManagedEntity previous ) {
606629 this .previous = previous ;
607630 }
631+
632+ @ Override
633+ public int $$_hibernate_getInstanceId () {
634+ return -1 ;
635+ }
636+
637+ @ Override
638+ public void $$_hibernate_setInstanceId (int id ) {
639+ }
608640 }
609641
610642 private static class ImmutableManagedEntityHolder implements ManagedEntity {
@@ -693,6 +725,16 @@ private boolean canClearEntityEntryReference() {
693725 return !(entityEntry instanceof ImmutableEntityEntry )
694726 || !entityEntry .getPersister ().canUseReferenceCacheEntries ();
695727 }
728+
729+ @ Override
730+ public int $$_hibernate_getInstanceId () {
731+ return managedEntity .$$_hibernate_getInstanceId ();
732+ }
733+
734+ @ Override
735+ public void $$_hibernate_setInstanceId (int id ) {
736+ managedEntity .$$_hibernate_setInstanceId ( id );
737+ }
696738 }
697739
698740 /**
0 commit comments