HHH-18326 Use instance identity to track immutable enhanced

entities

Author: mbellade

hibernate-core/src/main/java/org/hibernate/bytecode/enhance/internal/bytebuddy/EnhancerImpl.java

@@ -225,6 +225,14 @@ private DynamicType.Builder<?> doEnhance(Supplier<DynamicType.Builder<?>> builde
 					EnhancerConstants.USE_TRACKER_SETTER_NAME
 			);
 
+			builder = addFieldWithGetterAndSetter(
+					builder,
+					constants.TypeIntegerPrimitive,
+					EnhancerConstants.INSTANCE_ID_FIELD_NAME,
+					EnhancerConstants.INSTANCE_ID_GETTER_NAME,
+					EnhancerConstants.INSTANCE_ID_SETTER_NAME
+			);
+
 			builder = addInterceptorHandling( builder, managedCtClass );
 
 			if ( enhancementContext.doDirtyCheckingInline( managedCtClass ) ) {

hibernate-core/src/main/java/org/hibernate/bytecode/enhance/internal/bytebuddy/EnhancerImplConstants.java

@@ -74,6 +74,7 @@ public Class<? extends Annotation> annotationType() {
 	//Frequently used Types for method signatures:
 	final TypeDefinition TypeVoid = TypeDescription.ForLoadedType.of( void.class );
 	final TypeDefinition TypeBooleanPrimitive = TypeDescription.ForLoadedType.of( boolean.class );
+	final TypeDefinition TypeIntegerPrimitive = TypeDescription.ForLoadedType.of( int.class );
 	final TypeDefinition TypeManagedEntity = TypeDescription.ForLoadedType.of( ManagedEntity.class );
 	final TypeDefinition TypeEntityEntry = TypeDescription.ForLoadedType.of( EntityEntry.class );
 	final TypeDefinition TypePersistentAttributeInterceptor = TypeDescription.ForLoadedType.of( PersistentAttributeInterceptor.class );

hibernate-core/src/main/java/org/hibernate/bytecode/enhance/spi/EnhancerConstants.java

@@ -184,6 +184,10 @@ public final class EnhancerConstants {
 	public static final String USE_TRACKER_GETTER_NAME = "$$_hibernate_useTracker";
 	public static final String USE_TRACKER_SETTER_NAME = "$$_hibernate_setUseTracker";
 
+	public static final String INSTANCE_ID_FIELD_NAME = "$$_hibernate_instanceId";
+	public static final String INSTANCE_ID_GETTER_NAME = "$$_hibernate_getInstanceId";
+	public static final String INSTANCE_ID_SETTER_NAME = "$$_hibernate_setInstanceId";
+
 
 	private EnhancerConstants() {
 	}

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

@@ -14,6 +14,8 @@
 import org.hibernate.engine.spi.PersistentAttributeInterceptor;
 import org.hibernate.internal.CoreLogging;
 import org.hibernate.internal.CoreMessageLogger;
+import org.hibernate.internal.util.collections.InstanceIdentityMap;
+import org.hibernate.internal.util.collections.StandardStack;
 
 import java.io.IOException;
 import java.io.ObjectInputStream;
@@ -49,7 +51,9 @@ public class EntityEntryContext {
 
 	private final transient PersistenceContext persistenceContext;
 
-	private transient IdentityHashMap<ManagedEntity,ImmutableManagedEntityHolder> immutableManagedEntityXref;
+	private transient InstanceIdentityMap<ManagedEntity, ImmutableManagedEntityHolder> immutableManagedEntityXref;
+	private transient StandardStack<Integer> reusableInstanceIds;
+	private transient int currentInstanceId;
 
 	private transient ManagedEntity head;
 	private transient ManagedEntity tail;
@@ -107,13 +111,8 @@ public void addEntityEntry(Object entity, EntityEntry entityEntry) {
 				else {
 					// Create a holder for PersistenceContext-related data.
 					managedEntity = new ImmutableManagedEntityHolder( managed );
-					if ( immutableManagedEntityXref == null ) {
-						immutableManagedEntityXref = new IdentityHashMap<>();
-					}
-					immutableManagedEntityXref.put(
-							managed,
-							(ImmutableManagedEntityHolder) managedEntity
-					);
+					managed.$$_hibernate_setInstanceId( nextManagedEntityInstanceId() );
+					putImmutableManagedEntity( managed, (ImmutableManagedEntityHolder) managedEntity );
 				}
 			}
 			else {
@@ -171,9 +170,8 @@ private ManagedEntity getAssociatedManagedEntity(Object entity) {
 						: null;
 			}
 			else {
-				// if managedEntity is associated with this EntityEntryContext, then
-				// it will have an entry in immutableManagedEntityXref and its
-				// holder will be returned.
+				// if managedEntity is associated with this EntityEntryContext, it may have
+				// an entry in immutableManagedEntityXref and its holder will be returned.
 				return immutableManagedEntityXref != null
 						? immutableManagedEntityXref.get( managedEntity )
 						: null;
@@ -186,6 +184,19 @@ private ManagedEntity getAssociatedManagedEntity(Object entity) {
 		}
 	}
 
+	private int nextManagedEntityInstanceId() {
+		return reusableInstanceIds != null && !reusableInstanceIds.isEmpty() ?
+				reusableInstanceIds.pop() :
+				currentInstanceId++;
+	}
+
+	private void putImmutableManagedEntity(ManagedEntity managed, ImmutableManagedEntityHolder holder) {
+		if ( immutableManagedEntityXref == null ) {
+			immutableManagedEntityXref = new InstanceIdentityMap<>();
+		}
+		immutableManagedEntityXref.put( managed, holder );
+	}
+
 	private void checkNotAssociatedWithOtherPersistenceContextIfMutable(ManagedEntity managedEntity) {
 		// we only have to check mutable managedEntity
 		final AbstractEntityEntry entityEntry = (AbstractEntityEntry) managedEntity.$$_hibernate_getEntityEntry();
@@ -253,12 +264,15 @@ public EntityEntry removeEntityEntry(Object entity) {
 
 		dirty = true;
 
-		if (managedEntity instanceof ImmutableManagedEntityHolder) {
-			assert entity == ( (ImmutableManagedEntityHolder) managedEntity ).managedEntity;
+		if ( managedEntity instanceof ImmutableManagedEntityHolder ) {
+			assert entity == ((ImmutableManagedEntityHolder) managedEntity).managedEntity;
 			immutableManagedEntityXref.remove( entity );
-
+			if ( reusableInstanceIds == null ) {
+				reusableInstanceIds = new StandardStack<>();
+			}
+			reusableInstanceIds.push( managedEntity.$$_hibernate_getInstanceId() );
 		}
-		else if ( ! ( isManagedEntity( entity ) ) ) {
+		else if ( !isManagedEntity( entity ) ) {
 			nonEnhancedEntityXref.remove( entity );
 		}
 
@@ -387,6 +401,8 @@ public void clear() {
 		count = 0;
 
 		reentrantSafeEntries = null;
+		currentInstanceId = 0;
+		reusableInstanceIds = null;
 	}
 
 	private static void clearManagedEntity(final ManagedEntity node) {
@@ -483,16 +499,9 @@ public static EntityEntryContext deserialize(ObjectInputStream ois, StatefulPers
 				}
 				else {
 					final ManagedEntity castedEntity = asManagedEntity( entity );
+					castedEntity.$$_hibernate_setInstanceId( context.nextManagedEntityInstanceId() );
 					managedEntity = new ImmutableManagedEntityHolder( castedEntity );
-					if ( context.immutableManagedEntityXref == null ) {
-						context.immutableManagedEntityXref =
-								new IdentityHashMap<>();
-					}
-					context.immutableManagedEntityXref.put(
-							castedEntity,
-							(ImmutableManagedEntityHolder) managedEntity
-
-					);
+					context.putImmutableManagedEntity( castedEntity, (ImmutableManagedEntityHolder) managedEntity );
 				}
 			}
 			else {
@@ -605,6 +614,15 @@ public ManagedEntityImpl(Object entityInstance) {
 		public void $$_hibernate_setPreviousManagedEntity(ManagedEntity previous) {
 			this.previous = previous;
 		}
+
+		@Override
+		public int $$_hibernate_getInstanceId() {
+			return -1;
+		}
+
+		@Override
+		public void $$_hibernate_setInstanceId(int id) {
+		}
 	}
 
 	private static class ImmutableManagedEntityHolder implements ManagedEntity {
@@ -693,6 +711,16 @@ private boolean canClearEntityEntryReference() {
 			return !(entityEntry instanceof ImmutableEntityEntry)
 				|| !entityEntry.getPersister().canUseReferenceCacheEntries();
 		}
+
+		@Override
+		public int $$_hibernate_getInstanceId() {
+			return managedEntity.$$_hibernate_getInstanceId();
+		}
+
+		@Override
+		public void $$_hibernate_setInstanceId(int id) {
+			managedEntity.$$_hibernate_setInstanceId( id );
+		}
 	}
 
 	/**

hibernate-core/src/main/java/org/hibernate/engine/spi/ManagedEntity.java

@@ -21,7 +21,7 @@
  *
  * @author Steve Ebersole
  */
-public interface ManagedEntity extends Managed {
+public interface ManagedEntity extends Managed, InstanceIdentity {
 	/**
 	 * Obtain a reference to the entity instance.
 	 *

hibernate-core/src/test/java/org/hibernate/orm/test/bytecode/enhance/version/SimpleEntity.java

@@ -67,4 +67,14 @@ public class SimpleEntity implements ManagedEntity {
 	public boolean $$_hibernate_useTracker() {
 		return false;
 	}
+
+	@Override
+	public int $$_hibernate_getInstanceId() {
+		return 0;
+	}
+
+	@Override
+	public void $$_hibernate_setInstanceId(int id) {
+
+	}
 }

hibernate-core/src/test/java/org/hibernate/orm/test/bytecode/enhancement/basic/BasicSessionTest.java

@@ -81,6 +81,8 @@ static class MyEntity implements ManagedEntity {
 		private transient ManagedEntity previous;
 		@Transient
 		private transient ManagedEntity next;
+		@Transient
+		private transient int instanceId;
 
 		MyEntity() {
 		}
@@ -133,5 +135,15 @@ static class MyEntity implements ManagedEntity {
 		public boolean $$_hibernate_useTracker() {
 			return false;
 		}
+
+		@Override
+		public int $$_hibernate_getInstanceId() {
+			return instanceId;
+		}
+
+		@Override
+		public void $$_hibernate_setInstanceId(int id) {
+			this.instanceId = id;
+		}
 	}
 }

hibernate-core/src/test/java/org/hibernate/orm/test/cache/ByteCodeEnhancedImmutableReferenceCacheTest.java

@@ -238,6 +238,8 @@ public static class MyEnhancedReferenceData implements ManagedEntity {
 		private transient ManagedEntity previous;
 		@Transient
 		private transient ManagedEntity next;
+		@Transient
+		private transient int instanceId;
 
 		public MyEnhancedReferenceData(Integer id, String name, String theValue) {
 			this.id = id;
@@ -316,5 +318,15 @@ public void setTheValue(String theValue) {
 		public boolean $$_hibernate_useTracker() {
 			return false;
 		}
+
+		@Override
+		public int $$_hibernate_getInstanceId() {
+			return instanceId;
+		}
+
+		@Override
+		public void $$_hibernate_setInstanceId(int id) {
+			this.instanceId = id;
+		}
 	}
 }