HHH-19070 HHH-18326 Reset entity instance ids, use setPersistenceInfo

mbellade · mbellade · commit cc76f7595283 · 2025-02-10T17:12:57.000+01:00
diff --git a/hibernate-core/src/main/java/org/hibernate/engine/internal/EntityEntryContext.java b/hibernate-core/src/main/java/org/hibernate/engine/internal/EntityEntryContext.java
@@ -52,7 +52,7 @@ public class EntityEntryContext {
 	private final transient PersistenceContext persistenceContext;
 
 	private transient InstanceIdentityStore<ImmutableManagedEntityHolder> immutableManagedEntityXref;
-	private transient int currentInstanceId;
+	private transient int currentInstanceId = 1;
 
 	private transient ManagedEntity head;
 	private transient ManagedEntity tail;
@@ -97,10 +97,13 @@ public void addEntityEntry(Object entity, EntityEntry entityEntry) {
 		// Throw an exception if entity is a mutable ManagedEntity that is associated with a different
 		// PersistenceContext.
 		ManagedEntity managedEntity = getAssociatedManagedEntity( entity );
+
+		int instanceId = nextManagedEntityInstanceId();
 		final boolean alreadyAssociated = managedEntity != null;
 		if ( !alreadyAssociated ) {
 			if ( isManagedEntity( entity ) ) {
 				final ManagedEntity managed = asManagedEntity( entity );
+				assert managed.$$_hibernate_getInstanceId() == 0;
 				if ( entityEntry.getPersister().isMutable() ) {
 					managedEntity = managed;
 					// We know that managedEntity is not associated with the same PersistenceContext.
@@ -111,11 +114,11 @@ public void addEntityEntry(Object entity, EntityEntry entityEntry) {
 					// Create a holder for PersistenceContext-related data.
 					managedEntity = new ImmutableManagedEntityHolder( managed );
 					if ( !isReferenceCachingEnabled( entityEntry.getPersister() ) ) {
-						managed.$$_hibernate_setInstanceId( nextManagedEntityInstanceId() );
-						putImmutableManagedEntity( managed, (ImmutableManagedEntityHolder) managedEntity );
+						putImmutableManagedEntity( managed, instanceId, (ImmutableManagedEntityHolder) managedEntity );
 					}
 					else {
 						// When reference caching is enabled we cannot set the instance-id on the entity instance
+						instanceId = 0;
 						putManagedEntity( entity, managedEntity );
 					}
 				}
@@ -126,34 +129,33 @@ public void addEntityEntry(Object entity, EntityEntry entityEntry) {
 			}
 		}
 
-		// associate the EntityEntry with the entity
-		managedEntity.$$_hibernate_setEntityEntry( entityEntry );
-
 		if ( alreadyAssociated ) {
 			// if the entity was already associated with the context, skip the linking step.
+			managedEntity.$$_hibernate_setEntityEntry( entityEntry );
 			return;
 		}
 
 		// TODO: can dirty be set to true here?
 
 		// finally, set up linking and count
+		final ManagedEntity previous;
 		if ( tail == null ) {
 			assert head == null;
-			// Protect against stale data in the ManagedEntity and nullify previous/next references.
-			managedEntity.$$_hibernate_setPreviousManagedEntity( null );
-			managedEntity.$$_hibernate_setNextManagedEntity( null );
+			// Protect against stale data in the ManagedEntity and nullify previous reference.
+			previous = null;
 			head = managedEntity;
 			tail = head;
 			count = 1;
 		}
 		else {
 			tail.$$_hibernate_setNextManagedEntity( managedEntity );
-			managedEntity.$$_hibernate_setPreviousManagedEntity( tail );
-			// Protect against stale data left in the ManagedEntity nullify next reference.
-			managedEntity.$$_hibernate_setNextManagedEntity( null );
+			previous = tail;
 			tail = managedEntity;
 			count++;
 		}
+
+		// Protect against stale data left in the ManagedEntity nullify next reference.
+		managedEntity.$$_hibernate_setPersistenceInfo( entityEntry, previous, null, instanceId );
 	}
 
 	private static boolean isReferenceCachingEnabled(EntityPersister persister) {
@@ -201,11 +203,11 @@ private int nextManagedEntityInstanceId() {
 		return currentInstanceId++;
 	}
 
-	private void putImmutableManagedEntity(ManagedEntity managed, ImmutableManagedEntityHolder holder) {
+	private void putImmutableManagedEntity(ManagedEntity managed, int instanceId, ImmutableManagedEntityHolder holder) {
 		if ( immutableManagedEntityXref == null ) {
 			immutableManagedEntityXref = new InstanceIdentityStore<>();
 		}
-		immutableManagedEntityXref.put( managed, holder );
+		immutableManagedEntityXref.put( managed, instanceId, holder );
 	}
 
 	private void checkNotAssociatedWithOtherPersistenceContextIfMutable(ManagedEntity managedEntity) {
@@ -288,12 +290,6 @@ else if ( !isManagedEntity( entity ) ) {
 			nonEnhancedEntityXref.remove( entity );
 		}
 
-		// prepare for re-linking...
-		final ManagedEntity previous = managedEntity.$$_hibernate_getPreviousManagedEntity();
-		final ManagedEntity next = managedEntity.$$_hibernate_getNextManagedEntity();
-		managedEntity.$$_hibernate_setPreviousManagedEntity( null );
-		managedEntity.$$_hibernate_setNextManagedEntity( null );
-
 		// re-link
 		count--;
 
@@ -302,11 +298,13 @@ else if ( !isManagedEntity( entity ) ) {
 			head = null;
 			tail = null;
 
-			assert previous == null;
-			assert next == null;
+			assert managedEntity.$$_hibernate_getPreviousManagedEntity() == null;
+			assert managedEntity.$$_hibernate_getNextManagedEntity() == null;
 		}
 		else {
 			// otherwise, previous or next (or both) should be non-null
+			final ManagedEntity previous = managedEntity.$$_hibernate_getPreviousManagedEntity();
+			final ManagedEntity next = managedEntity.$$_hibernate_getNextManagedEntity();
 			if ( previous == null ) {
 				// we are removing head
 				assert managedEntity == head;
@@ -327,9 +325,7 @@ else if ( !isManagedEntity( entity ) ) {
 		}
 
 		// finally clean out the ManagedEntity and return the associated EntityEntry
-		final EntityEntry theEntityEntry = managedEntity.$$_hibernate_getEntityEntry();
-		managedEntity.$$_hibernate_setEntityEntry( null );
-		return theEntityEntry;
+		return clearManagedEntity( managedEntity );
 	}
 
 	/**
@@ -379,7 +375,7 @@ private void clearAllReferencesFromManagedEntities() {
 			nextManagedEntity = current.$$_hibernate_getNextManagedEntity();
 			Object toProcess = current.$$_hibernate_getEntityInstance();
 			unsetSession( asPersistentAttributeInterceptableOrNull( toProcess ) );
-			clearManagedEntity( current );//careful this also unlinks from the "next" entry in the list
+			clearManagedEntity( current ); //careful this also unlinks from the "next" entry in the list
 		}
 	}
 
@@ -413,13 +409,17 @@ public void clear() {
 		count = 0;
 
 		reentrantSafeEntries = null;
-		currentInstanceId = 0;
+		currentInstanceId = 1;
 	}
 
-	private static void clearManagedEntity(final ManagedEntity node) {
-		node.$$_hibernate_setEntityEntry( null );
-		node.$$_hibernate_setPreviousManagedEntity( null );
-		node.$$_hibernate_setNextManagedEntity( null );
+	/**
+	 * Resets the persistence information in a managed entity, and returns its previous {@link EntityEntry}
+	 *
+	 * @param node the managed entity to clear
+	 * @return the previous {@link EntityEntry} contained in the managed entity
+	 */
+	private static EntityEntry clearManagedEntity(final ManagedEntity node) {
+		return node.$$_hibernate_setPersistenceInfo( null, null, null, 0 );
 	}
 
 	/**
@@ -503,6 +503,7 @@ public static EntityEntryContext deserialize(ObjectInputStream ois, StatefulPers
 
 			final EntityEntry entry = deserializeEntityEntry( entityEntryClassNameArr, ois, rtn );
 
+			final int instanceId = context.nextManagedEntityInstanceId();
 			final ManagedEntity managedEntity;
 			if ( isEnhanced ) {
 				final ManagedEntity castedEntity = asManagedEntity( entity );
@@ -512,8 +513,7 @@ public static EntityEntryContext deserialize(ObjectInputStream ois, StatefulPers
 				else {
 					managedEntity = new ImmutableManagedEntityHolder( castedEntity );
 					if ( !isReferenceCachingEnabled( entry.getPersister() ) ) {
-						castedEntity.$$_hibernate_setInstanceId( context.nextManagedEntityInstanceId() );
-						context.putImmutableManagedEntity( castedEntity, (ImmutableManagedEntityHolder) managedEntity );
+						context.putImmutableManagedEntity( castedEntity, instanceId, (ImmutableManagedEntityHolder) managedEntity );
 					}
 					else {
 						context.putManagedEntity( entity, castedEntity );
@@ -525,16 +525,15 @@ public static EntityEntryContext deserialize(ObjectInputStream ois, StatefulPers
 				context.putManagedEntity( entity, managedEntity );
 			}
 
-			managedEntity.$$_hibernate_setEntityEntry( entry );
-
 			if ( previous == null ) {
 				context.head = managedEntity;
 			}
 			else {
 				previous.$$_hibernate_setNextManagedEntity( managedEntity );
-				managedEntity.$$_hibernate_setPreviousManagedEntity( previous );
 			}
 
+			managedEntity.$$_hibernate_setPersistenceInfo( entry, previous, null, instanceId );
+
 			previous = managedEntity;
 		}
 
@@ -631,12 +630,21 @@ public ManagedEntityImpl(Object entityInstance) {
 
 		@Override
 		public int $$_hibernate_getInstanceId() {
-			return -1;
+			return 0;
 		}
 
 		@Override
 		public void $$_hibernate_setInstanceId(int id) {
 		}
+
+		@Override
+		public EntityEntry $$_hibernate_setPersistenceInfo(EntityEntry entityEntry, ManagedEntity previous, ManagedEntity next, int instanceId) {
+			final EntityEntry oldEntry = this.entityEntry;
+			this.entityEntry = entityEntry;
+			this.previous = previous;
+			this.next = next;
+			return oldEntry;
+		}
 	}
 
 	private static class ImmutableManagedEntityHolder implements ManagedEntity {
@@ -662,7 +670,7 @@ public ImmutableManagedEntityHolder(ManagedEntity immutableManagedEntity) {
 		public void $$_hibernate_setEntityEntry(EntityEntry entityEntry) {
 			// need to think about implications for memory leaks here if we don't removed reference to EntityEntry
 			if ( entityEntry == null ) {
-				if ( canClearEntityEntryReference() ) {
+				if ( canClearEntityEntryReference( managedEntity.$$_hibernate_getEntityEntry() ) ) {
 					managedEntity.$$_hibernate_setEntityEntry( null );
 				}
 				// otherwise, do nothing.
@@ -720,10 +728,9 @@ public ImmutableManagedEntityHolder(ManagedEntity immutableManagedEntity) {
 
 		// Check instance type of EntityEntry and if type is ImmutableEntityEntry,
 		// check to see if entity is referenced cached in the second level cache
-		private boolean canClearEntityEntryReference() {
-			EntityEntry entityEntry = managedEntity.$$_hibernate_getEntityEntry();
+		private static boolean canClearEntityEntryReference(EntityEntry entityEntry) {
 			return !(entityEntry instanceof ImmutableEntityEntry)
-				|| !entityEntry.getPersister().canUseReferenceCacheEntries();
+				|| !isReferenceCachingEnabled( entityEntry.getPersister() );
 		}
 
 		@Override
@@ -735,6 +742,25 @@ private boolean canClearEntityEntryReference() {
 		public void $$_hibernate_setInstanceId(int id) {
 			managedEntity.$$_hibernate_setInstanceId( id );
 		}
+
+		@Override
+		public EntityEntry $$_hibernate_setPersistenceInfo(EntityEntry entityEntry, ManagedEntity previous, ManagedEntity next, int instanceId) {
+			final EntityEntry oldEntry;
+			if ( entityEntry == null ) {
+				oldEntry = managedEntity.$$_hibernate_getEntityEntry();
+				if ( canClearEntityEntryReference( oldEntry ) ) {
+					managedEntity.$$_hibernate_setEntityEntry( null );
+				}
+			}
+			else {
+				managedEntity.$$_hibernate_setEntityEntry( entityEntry );
+				oldEntry = null; // no need to retrieve the previous entity entry
+			}
+			this.previous = previous;
+			this.next = next;
+			managedEntity.$$_hibernate_setInstanceId( instanceId );
+			return oldEntry;
+		}
 	}
 
 	/**
diff --git a/hibernate-core/src/main/java/org/hibernate/engine/spi/ManagedEntity.java b/hibernate-core/src/main/java/org/hibernate/engine/spi/ManagedEntity.java
@@ -128,4 +128,26 @@ default ManagedEntity asManagedEntity() {
 		return this;
 	}
 
+	/**
+	 * Utility method that allows injecting all persistence-related information on the managed entity at once.
+	 *
+	 * @param entityEntry the {@link EntityEntry} associated with this entity instance
+	 * @param previous the previous entry
+	 * @param next the next entry
+	 * @param instanceId unique identifier for this instance
+	 * @return the previous {@link EntityEntry} contained in this managed entity, or {@code null}
+	 * @see #$$_hibernate_setEntityEntry(EntityEntry)
+	 * @see #$$_hibernate_setPreviousManagedEntity(ManagedEntity)
+	 * @see #$$_hibernate_setNextManagedEntity(ManagedEntity)
+	 * @see #$$_hibernate_setInstanceId(int)
+	 * @since 7.0
+	 */
+	default EntityEntry $$_hibernate_setPersistenceInfo(EntityEntry entityEntry, ManagedEntity previous, ManagedEntity next, int instanceId) {
+		final EntityEntry oldEntry = $$_hibernate_getEntityEntry();
+		$$_hibernate_setEntityEntry( entityEntry );
+		$$_hibernate_setPreviousManagedEntity( previous );
+		$$_hibernate_setNextManagedEntity( next );
+		$$_hibernate_setInstanceId( instanceId );
+		return oldEntry;
+	}
 }
diff --git a/hibernate-core/src/main/java/org/hibernate/internal/util/collections/InstanceIdentityStore.java b/hibernate-core/src/main/java/org/hibernate/internal/util/collections/InstanceIdentityStore.java
@@ -7,11 +7,16 @@
 import org.checkerframework.checker.nullness.qual.Nullable;
 import org.hibernate.engine.spi.InstanceIdentity;
 
+import java.util.ConcurrentModificationException;
+
 /**
  * Utility collection that takes advantage of {@link InstanceIdentity}'s identifier to store objects.
  * The store is based on {@link AbstractPagedArray} and it stores element using their instance-id
  * as index.
  * <p>
+ * Instance ids are considered to start from 1, and if two instances are found to have the
+ * same identifier a {@link java.util.ConcurrentModificationException will be thrown}.
+ * <p>
  * Both keys and values are stored in this array, requiring very few allocations to keep track of the pair.
  * The downside to this is we cannot easily access the key, especially if asking for a specific type, since
  * that would cause type-pollution issues at the call site that would degrade performance.
@@ -26,7 +31,7 @@ public class InstanceIdentityStore<V> extends AbstractPagedArray<Object> {
 	 * @return the index of the corresponding key instance in the array
 	 */
 	private static int toKeyIndex(int instanceId) {
-		return instanceId * 2;
+		return (instanceId - 1) * 2;
 	}
 
 	/**
@@ -41,6 +46,10 @@ private static int toKeyIndex(int instanceId) {
 	 * equality check ({@code ==}) with the provided key to ensure it corresponds to the mapped one
 	 */
 	public @Nullable V get(int instanceId, Object key) {
+		if ( instanceId <= 0 ) {
+			return null;
+		}
+
 		final int keyIndex = toKeyIndex( instanceId );
 		final Page<Object> page = getPage( keyIndex );
 		if ( page != null ) {
@@ -50,6 +59,12 @@ private static int toKeyIndex(int instanceId) {
 				//noinspection unchecked
 				return (V) page.get( offset + 1 );
 			}
+			else {
+				throw new ConcurrentModificationException(
+						"Found a different instance corresponding to instanceId [" + instanceId +
+						"], this might indicate a concurrent access to this persistence context."
+				);
+			}
 		}
 		return null;
 	}
@@ -61,12 +76,14 @@ private static int toKeyIndex(int instanceId) {
 	 * @param key key with which the specified value is to be associated
 	 * @param value value to be associated with the specified key
 	 */
-	public <K extends InstanceIdentity> void put(K key, V value) {
+	public void put(Object key, int instanceId, V value) {
 		if ( key == null ) {
 			throw new NullPointerException( "This store does not support null keys" );
 		}
+		else if ( instanceId <= 0 ) {
+			throw new IllegalArgumentException( "Instance ID must be a positive value" );
+		}
 
-		final int instanceId = key.$$_hibernate_getInstanceId();
 		final int keyIndex = toKeyIndex( instanceId );
 		final Page<Object> page = getOrCreateEntryPage( keyIndex );
 		final int pageOffset = toPageOffset( keyIndex );
@@ -83,6 +100,10 @@ public <K extends InstanceIdentity> void put(K key, V value) {
 	 * equality check ({@code ==}) with the provided key to ensure it corresponds to the mapped one
 	 */
 	public void remove(int instanceId, Object key) {
+		if ( instanceId <= 0 ) {
+			return;
+		}
+
 		final int keyIndex = toKeyIndex( instanceId );
 		final Page<Object> page = getPage( keyIndex );
 		if ( page != null ) {
@@ -93,8 +114,10 @@ public void remove(int instanceId, Object key) {
 				page.set( pageOffset + 1, null );
 			}
 			else {
-				// If it doesn't, reset the array value to the old key
-				page.set( pageOffset, k );
+				throw new ConcurrentModificationException(
+						"Found a different instance corresponding to instanceId [" + instanceId +
+						"], this might indicate a concurrent access to this persistence context."
+				);
 			}
 		}
 	}
