HHH-18326 Also reset collection instance ids

Author: mbellade

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

@@ -134,7 +134,7 @@ the following fields are used in all circumstances, and are not worth (or not su
 	private InstanceIdentityMap<PersistentCollection<?>, CollectionEntry> collectionEntries;
 
 	// Current collection instance id
-	private transient int currentCollectionInstanceId;
+	private transient int currentCollectionInstanceId = 1;
 
 	// Collection wrappers, by the CollectionKey
 	private HashMap<CollectionKey, PersistentCollection<?>> collectionsByKey;
@@ -258,7 +258,10 @@ public void clear() {
 
 		final SharedSessionContractImplementor session = getSession();
 		if ( collectionEntries != null ) {
-			collectionEntries.forEach( (k, v) -> k.unsetSession( session ) );
+			collectionEntries.forEach( (k, v) -> {
+				k.$$_hibernate_setInstanceId( 0 );
+				k.unsetSession( session );
+			} );
 		}
 
 		arrayHolders = null;
@@ -270,7 +273,7 @@ public void clear() {
 		collectionsByKey = null;
 		nonlazyCollections = null;
 		collectionEntries = null;
-		currentCollectionInstanceId = 0;
+		currentCollectionInstanceId = 1;
 		unownedCollections = null;
 		nullifiableEntityKeys = null;
 		deletedUnloadedEntityKeys = null;
@@ -1133,6 +1136,7 @@ private void putCollectionEntry(PersistentCollection<?> collection, CollectionEn
 		if ( this.collectionEntries == null ) {
 			this.collectionEntries = new InstanceIdentityMap<>();
 		}
+		assert collection.$$_hibernate_getInstanceId() == 0;
 		collection.$$_hibernate_setInstanceId( nextCollectionInstanceId() );
 		this.collectionEntries.put( collection, entry );
 	}
@@ -2203,7 +2207,9 @@ public int getCollectionEntriesSize() {
 	@Override
 	public CollectionEntry removeCollectionEntry(PersistentCollection<?> collection) {
 		if ( collectionEntries != null ) {
-			return collectionEntries.remove( collection.$$_hibernate_getInstanceId(), collection );
+			final int instanceId = collection.$$_hibernate_getInstanceId();
+			collection.$$_hibernate_setInstanceId( 0 );
+			return collectionEntries.remove( instanceId, collection );
 		}
 		else {
 			return null;

hibernate-core/src/main/java/org/hibernate/internal/StatelessSessionImpl.java

@@ -1033,6 +1033,7 @@ else if ( association instanceof PersistentCollection<?> collection ) {
 					}
 				}
 				finally {
+					collection.$$_hibernate_setInstanceId( 0 );
 					collection.unsetSession( this );
 					if ( persistenceContext.isLoadFinished() ) {
 						persistenceContext.clear();

hibernate-core/src/main/java/org/hibernate/internal/util/collections/InstanceIdentityMap.java

@@ -27,6 +27,9 @@
  * but, contrary to the store, it initializes {@link Map.Entry}s eagerly to optimize iteration
  * performance and avoid type-pollution issues when checking the type of contained objects.
  * <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>
  * Methods accessing / modifying the map with {@link Object} typed parameters will need
  * to type check against the instance identity interface which might be inefficient,
  * so it's recommended to use the position (int) based variant of those methods.
@@ -107,9 +110,21 @@ private boolean containsMapping(Object key, Object value) {
 	 * equality check ({@code ==}) with the provided key to ensure it corresponds to the mapped one
 	 */
 	public @Nullable V get(int instanceId, Object key) {
-		final Entry<K, V> entry = get( instanceId );
-		if ( entry != null && entry.getKey() == key ) {
-			return entry.getValue();
+		if ( instanceId <= 0 ) {
+			return null;
+		}
+
+		final Entry<K, V> entry = get( instanceId - 1 );
+		if ( entry != null ) {
+			if ( entry.getKey() == key ) {
+				return entry.getValue();
+			}
+			else {
+				throw new ConcurrentModificationException(
+						"Found a different instance corresponding to instanceId [" + instanceId +
+						"], this might indicate a concurrent access to this persistence context."
+				);
+			}
 		}
 		return null;
 	}
@@ -133,8 +148,12 @@ private boolean containsMapping(Object key, Object value) {
 			throw new NullPointerException( "This map does not support null keys" );
 		}
 
-		final int instanceId = key.$$_hibernate_getInstanceId();
-		final Map.Entry<K, V> old = set( instanceId, new AbstractMap.SimpleImmutableEntry<>( key, value ) );
+		final int index = key.$$_hibernate_getInstanceId() - 1;
+		if ( index < 0 ) {
+			throw new IllegalArgumentException( "Instance ID must be a positive value" );
+		}
+
+		final Map.Entry<K, V> old = set( index, new AbstractMap.SimpleImmutableEntry<>( key, value ) );
 		if ( old == null ) {
 			size++;
 			return null;
@@ -154,9 +173,14 @@ private boolean containsMapping(Object key, Object value) {
 	 * equality check ({@code ==}) with the provided key to ensure it corresponds to the mapped one
 	 */
 	public @Nullable V remove(int instanceId, Object key) {
-		final Page<Map.Entry<K, V>> page = getPage( instanceId );
+		if ( instanceId <= 0 ) {
+			return null;
+		}
+
+		final int index = instanceId - 1;
+		final Page<Map.Entry<K, V>> page = getPage( index );
 		if ( page != null ) {
-			final int pageOffset = toPageOffset( instanceId );
+			final int pageOffset = toPageOffset( index );
 			final Map.Entry<K, V> entry = page.set( pageOffset, null );
 			// Check that the provided instance really matches with the key contained in the map
 			if ( entry != null ) {
@@ -165,8 +189,10 @@ private boolean containsMapping(Object key, Object value) {
 					return entry.getValue();
 				}
 				else {
-					// If it doesn't, reset the array value to the old key
-					page.set( pageOffset, entry );
+					throw new ConcurrentModificationException(
+							"Found a different instance corresponding to instanceId [" + instanceId +
+							"], this might indicate a concurrent access to this persistence context."
+					);
 				}
 			}
 		}

hibernate-core/src/test/java/org/hibernate/orm/test/util/InstanceIdentityMapTest.java

@@ -70,15 +70,15 @@ public void testSimpleMapOperations() {
 
 	@Test
 	public void testMapIteration() {
-		for ( int i = 0; i < 100; i++ ) {
+		for ( int i = 1; i <= 100; i++ ) {
 			testMap.put( new TestInstance( i ), "instance_" + i );
 		}
 
 		assertThat( testMap ).hasSize( 100 );
 
 		final AtomicInteger count = new AtomicInteger();
 		testMap.forEach( (k, v) -> {
-			assertThat( k.$$_hibernate_getInstanceId() ).isBetween( 0, 99 );
+			assertThat( k.$$_hibernate_getInstanceId() ).isBetween( 1, 100 );
 			assertThat( v ).isEqualTo( "instance_" + k.$$_hibernate_getInstanceId() );
 			count.getAndIncrement();
 		} );
@@ -88,7 +88,7 @@ public void testMapIteration() {
 	@Test
 	public void testSets() {
 		final Map<TestInstance, String> map = new HashMap<TestInstance, String>();
-		for ( int i = 0; i < 100; i++ ) {
+		for ( int i = 1; i <= 100; i++ ) {
 			map.put( new TestInstance( i ), "instance_" + i );
 		}