HHH-13654 Defer initialization of StatefulPersistenceContext#collectionsByKey

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

@@ -160,7 +160,6 @@ public class StatefulPersistenceContext implements PersistenceContext {
 		entitySnapshotsByKey = new HashMap<>( INIT_COLL_SIZE );
 
 		entityEntryContext = new EntityEntryContext( this );
-		collectionsByKey = new HashMap<>( INIT_COLL_SIZE );
 	}
 
 	private ConcurrentMap<EntityKey, Object> getOrInitializeProxiesByKey() {
@@ -247,7 +246,7 @@ public class StatefulPersistenceContext implements PersistenceContext {
 		entityEntryContext.clear();
 		parentsByChild = null;
 		entitySnapshotsByKey.clear();
-		collectionsByKey.clear();
+		collectionsByKey = null;
 		nonlazyCollections = null;
 		collectionEntries = null;
 		unownedCollections = null;
@@ -890,7 +889,7 @@ public class StatefulPersistenceContext implements PersistenceContext {
 	private void addCollection(PersistentCollection coll, CollectionEntry entry, Serializable key) {
 		getOrInitializeCollectionEntries().put( coll, entry );
 		final CollectionKey collectionKey = new CollectionKey( entry.getLoadedPersister(), key );
-		final PersistentCollection old = collectionsByKey.put( collectionKey, coll );
+		final PersistentCollection old = addCollectionByKey( collectionKey, coll );
 		if ( old != null ) {
 			if ( old == coll ) {
 				throw new AssertionFailure( "bug adding collection twice" );
@@ -947,7 +946,7 @@ public class StatefulPersistenceContext implements PersistenceContext {
 
 	@Override
 	public PersistentCollection getCollection(CollectionKey collectionKey) {
-		return collectionsByKey.get( collectionKey );
+		return collectionsByKey == null ? null : collectionsByKey.get( collectionKey );
 	}
 
 	@Override
@@ -1099,7 +1098,12 @@ public class StatefulPersistenceContext implements PersistenceContext {
 
 	@Override
 	public Map getCollectionsByKey() {
-		return collectionsByKey;
+		if ( collectionsByKey == null ) {
+			return Collections.emptyMap();
+		}
+		else {
+			return collectionsByKey;
+		}
 	}
 
 	@Override
@@ -1190,8 +1194,13 @@ public class StatefulPersistenceContext implements PersistenceContext {
 
 	@Override
 	public String toString() {
-		return "PersistenceContext[entityKeys=" + entitiesByKey.keySet()
+		if ( collectionsByKey == null ) {
+			return "PersistenceContext[entityKeys=" + entitiesByKey.keySet() + ",collectionKeys=[]]";
+		}
+		else {
+			return "PersistenceContext[entityKeys=" + entitiesByKey.keySet()
 				+ ",collectionKeys=" + collectionsByKey.keySet() + "]";
+		}
 	}
 
 	@Override
@@ -1569,13 +1578,18 @@ public class StatefulPersistenceContext implements PersistenceContext {
 
 		entityEntryContext.serialize( oos );
 
-		oos.writeInt( collectionsByKey.size() );
+		if ( collectionsByKey == null ) {
-		if ( LOG.isTraceEnabled() ) {
+			oos.writeInt( 0 );
-			LOG.trace( "Starting serialization of [" + collectionsByKey.size() + "] collectionsByKey entries" );
 		}
-		for ( Map.Entry<CollectionKey,PersistentCollection> entry : collectionsByKey.entrySet() ) {
+		else {
-			entry.getKey().serialize( oos );
+			oos.writeInt( collectionsByKey.size() );
-			oos.writeObject( entry.getValue() );
+			if ( LOG.isTraceEnabled() ) {
+				LOG.trace( "Starting serialization of [" + collectionsByKey.size() + "] collectionsByKey entries" );
+			}
+			for ( Map.Entry<CollectionKey, PersistentCollection> entry : collectionsByKey.entrySet() ) {
+				entry.getKey().serialize( oos );
+				oos.writeObject( entry.getValue() );
+			}
 		}
 
 		if ( collectionEntries == null ) {
@@ -1836,6 +1850,32 @@ public class StatefulPersistenceContext implements PersistenceContext {
 		}
 	}
 
+	@Override
+	public void clearCollectionsByKey() {
+		if ( collectionsByKey != null ) {
+			//A valid alternative would be to set this to null, like we do on close.
+			//The difference being that in this case we expect the collection will be used again, so we bet that clear()
+			//might allow us to skip having to re-allocate the collection.
+			collectionsByKey.clear();
+		}
+	}
+
+	@Override
+	public PersistentCollection addCollectionByKey(CollectionKey collectionKey, PersistentCollection persistentCollection) {
+		if ( collectionsByKey == null ) {
+			collectionsByKey = new HashMap<>( INIT_COLL_SIZE );
+		}
+		final PersistentCollection old = collectionsByKey.put( collectionKey, persistentCollection );
+		return old;
+	}
+
+	@Override
+	public void removeCollectionByKey(CollectionKey collectionKey) {
+		if ( collectionsByKey != null ) {
+			collectionsByKey.remove( collectionKey );
+		}
+	}
+
 	private void cleanUpInsertedKeysAfterTransaction() {
 		if ( insertedKeysMap != null ) {
 			insertedKeysMap.clear();

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

@@ -525,7 +525,12 @@ public interface PersistenceContext {
 
 	/**
 	 * Get the mapping from collection key to collection instance
+	 * @deprecated this method should be removed; alternative methods are available that better express the intent, allowing
+	 * for better optimisations. Not aggressively removing this as it's an SPI, but also useful for testing and other
+	 * contexts which are not performance sensitive.
+	 * N.B. This might return an immutable map: do not use for mutations!
 	 */
+	@Deprecated
 	Map getCollectionsByKey();
 
 	/**
@@ -769,6 +774,25 @@ public interface PersistenceContext {
 	 */
 	CollectionEntry removeCollectionEntry(PersistentCollection collection);
 
+	/**
+	 * Remove all state of the collections-by-key map.
+	 */
+	void clearCollectionsByKey();
+
+	/**
+	 * Adds a collection in the collections-by-key map.
+	 * @param collectionKey
+	 * @param persistentCollection
+	 * @return the previous collection, it the key was already mapped.
+	 */
+	PersistentCollection addCollectionByKey(CollectionKey collectionKey, PersistentCollection persistentCollection);
+
+	/**
+	 * Remove a collection-by-key mapping.
+	 * @param collectionKey the key to clear
+	 */
+	void removeCollectionByKey(CollectionKey collectionKey);
+
 	/**
 	 * Provides centralized access to natural-id-related functionality.
 	 */

hibernate-core/src/main/java/org/hibernate/event/internal/AbstractFlushingEventListener.java

@@ -369,7 +369,7 @@ public abstract class AbstractFlushingEventListener implements JpaBootstrapSensi
 		LOG.trace( "Post flush" );
 
 		final PersistenceContext persistenceContext = session.getPersistenceContextInternal();
-		persistenceContext.getCollectionsByKey().clear();
+		persistenceContext.clearCollectionsByKey();
 		
 		// the database has changed now, so the subselect results need to be invalidated
 		// the batch fetching queues should also be cleared - especially the collection batch fetching one
@@ -390,7 +390,7 @@ public abstract class AbstractFlushingEventListener implements JpaBootstrapSensi
 								collectionEntry.getLoadedPersister(),
 								collectionEntry.getLoadedKey()
 						);
-						persistenceContext.getCollectionsByKey().put( collectionKey, persistentCollection );
+						persistenceContext.addCollectionByKey( collectionKey, persistentCollection );
 					}
 				}, true
 		);

hibernate-core/src/main/java/org/hibernate/event/internal/EvictVisitor.java

@@ -81,9 +81,7 @@ public class EvictVisitor extends AbstractVisitor {
 		}
 		if ( ce.getLoadedPersister() != null && ce.getLoadedKey() != null ) {
 			//TODO: is this 100% correct?
-			persistenceContext.getCollectionsByKey().remove(
+			persistenceContext.removeCollectionByKey( new CollectionKey( ce.getLoadedPersister(), ce.getLoadedKey() ) );
-					new CollectionKey( ce.getLoadedPersister(), ce.getLoadedKey() )
-			);
 		}
 	}