HHH-13654 Make AbstractFlushingEventListener#entitiesByKey also lazily initialized

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

@@ -95,7 +95,7 @@ public class StatefulPersistenceContext implements PersistenceContext {
 	private SharedSessionContractImplementor session;
 
 	// Loaded entity instances, by EntityKey
-	private Map<EntityKey, Object> entitiesByKey;
+	private HashMap<EntityKey, Object> entitiesByKey;
 
 	// Loaded entity instances, by EntityUniqueKey
 	private Map<EntityUniqueKey, Object> entitiesByUniqueKey;
@@ -155,8 +155,6 @@ public class StatefulPersistenceContext implements PersistenceContext {
 	 */
 	public StatefulPersistenceContext(SharedSessionContractImplementor session) {
 		this.session = session;
-
-		entitiesByKey = new HashMap<>( INIT_COLL_SIZE );
 		entityEntryContext = new EntityEntryContext( this );
 	}
 
@@ -239,7 +237,7 @@ public class StatefulPersistenceContext implements PersistenceContext {
 		}
 
 		arrayHolders = null;
-		entitiesByKey.clear();
+		entitiesByKey = null;
 		entitiesByUniqueKey = null;
 		entityEntryContext.clear();
 		parentsByChild = null;
@@ -383,26 +381,31 @@ public class StatefulPersistenceContext implements PersistenceContext {
 
 	@Override
 	public void addEntity(EntityKey key, Object entity) {
+		if ( entitiesByKey == null ) {
+			entitiesByKey = new HashMap<>( INIT_COLL_SIZE );
+		}
 		entitiesByKey.put( key, entity );
 		final BatchFetchQueue fetchQueue = this.batchFetchQueue;
 		if ( fetchQueue != null ) {
-			fetchQueue.removeBatchLoadableEntityKey(key);
+			fetchQueue.removeBatchLoadableEntityKey( key );
 		}
 	}
 
 	@Override
 	public Object getEntity(EntityKey key) {
-		return entitiesByKey.get( key );
+		return entitiesByKey == null ? null : entitiesByKey.get( key );
 	}
 
 	@Override
 	public boolean containsEntity(EntityKey key) {
-		return entitiesByKey.containsKey( key );
+		return entitiesByKey == null ? false : entitiesByKey.containsKey( key );
 	}
 
 	@Override
 	public Object removeEntity(EntityKey key) {
-		final Object entity = entitiesByKey.remove( key );
+		final Object entity;
+		if ( entitiesByKey != null ) {
+			entity = entitiesByKey.remove( key );
 			if ( entitiesByUniqueKey != null ) {
 				final Iterator itr = entitiesByUniqueKey.values().iterator();
 				while ( itr.hasNext() ) {
@@ -411,6 +414,10 @@ public class StatefulPersistenceContext implements PersistenceContext {
 					}
 				}
 			}
+		}
+		else {
+			entity = null;
+		}
 
 		// Clear all parent cache
 		parentsByChild = null;
@@ -757,6 +764,9 @@ public class StatefulPersistenceContext implements PersistenceContext {
 
 	@Override
 	public void addEnhancedProxy(EntityKey key, PersistentAttributeInterceptable entity) {
+		if ( entitiesByKey == null ) {
+			entitiesByKey = new HashMap<>( INIT_COLL_SIZE );
+		}
 		entitiesByKey.put( key, entity );
 	}
 
@@ -1062,9 +1072,25 @@ public class StatefulPersistenceContext implements PersistenceContext {
 		return nullifiableEntityKeys;
 	}
 
+	/**
+	 * @deprecated this will be removed: it provides too wide access, making it hard to optimise the internals
+	 * for specific access needs. Consider using #iterateEntities instead.
+	 * @return
+	 */
+	@Deprecated
 	@Override
 	public Map getEntitiesByKey() {
-		return entitiesByKey;
+		return entitiesByKey == null ? Collections.emptyMap() : entitiesByKey;
+	}
+
+	@Override
+	public Iterator managedEntitiesIterator() {
+		if ( entitiesByKey == null ) {
+			return Collections.emptyIterator();
+		}
+		else {
+			return entitiesByKey.values().iterator();
+		}
 	}
 
 	@Override
@@ -1200,13 +1226,9 @@ public class StatefulPersistenceContext implements PersistenceContext {
 
 	@Override
 	public String toString() {
-		if ( collectionsByKey == null ) {
+		final String entityKeySet = entitiesByKey == null ? "[]" :  entitiesByKey.keySet().toString();
-			return "PersistenceContext[entityKeys=" + entitiesByKey.keySet() + ",collectionKeys=[]]";
+		final String collectionsKeySet = collectionsByKey == null ? "[]" : collectionsByKey.keySet().toString();
-		}
+		return "PersistenceContext[entityKeys=" + entityKeySet + ", collectionKeys=" + collectionsKeySet + "]";
-		else {
-			return "PersistenceContext[entityKeys=" + entitiesByKey.keySet()
-				+ ",collectionKeys=" + collectionsByKey.keySet() + "]";
-		}
 	}
 
 	@Override
@@ -1503,7 +1525,7 @@ public class StatefulPersistenceContext implements PersistenceContext {
 
 	@Override
 	public void replaceDelayedEntityIdentityInsertKeys(EntityKey oldKey, Serializable generatedId) {
-		final Object entity = entitiesByKey.remove( oldKey );
+		final Object entity = entitiesByKey == null ? null : entitiesByKey.remove( oldKey );
 		final EntityEntry oldEntry = entityEntryContext.removeEntityEntry( entity );
 		this.parentsByChild = null;
 
@@ -1536,6 +1558,10 @@ public class StatefulPersistenceContext implements PersistenceContext {
 		oos.writeBoolean( defaultReadOnly );
 		oos.writeBoolean( hasNonReadOnlyEntities );
 
+		if ( entitiesByKey == null ) {
+			oos.writeInt( 0 );
+		}
+		else {
 			oos.writeInt( entitiesByKey.size() );
 			if ( LOG.isTraceEnabled() ) {
 				LOG.trace( "Starting serialization of [" + entitiesByKey.size() + "] entitiesByKey entries" );
@@ -1544,6 +1570,7 @@ public class StatefulPersistenceContext implements PersistenceContext {
 				entry.getKey().serialize( oos );
 				oos.writeObject( entry.getValue() );
 			}
+		}
 
 		if ( entitiesByUniqueKey == null ) {
 			oos.writeInt( 0 );

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

@@ -9,6 +9,7 @@ package org.hibernate.engine.spi;
 import java.io.Serializable;
 import java.util.Collection;
 import java.util.HashSet;
+import java.util.Iterator;
 import java.util.Map;
 import java.util.function.BiConsumer;
 import java.util.function.Supplier;
@@ -486,7 +487,10 @@ public interface PersistenceContext {
 
 	/**
 	 * Get the mapping from key value to entity instance
+	 * @deprecated this will be removed: it provides too wide access, making it hard to optimise the internals
+	 * for specific access needs. Consider using #iterateEntities instead.
 	 */
+	@Deprecated
 	Map getEntitiesByKey();
 
 	/**
@@ -793,6 +797,12 @@ public interface PersistenceContext {
 	 */
 	void removeCollectionByKey(CollectionKey collectionKey);
 
+	/**
+	 * A read-only iterator on all entities managed by this persistence context
+	 * @return
+	 */
+	Iterator managedEntitiesIterator();
+
 	/**
 	 * Provides centralized access to natural-id-related functionality.
 	 */

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

@@ -37,7 +37,6 @@ import org.hibernate.event.spi.FlushEntityEventListener;
 import org.hibernate.event.spi.FlushEvent;
 import org.hibernate.internal.CoreMessageLogger;
 import org.hibernate.internal.util.EntityPrinter;
-import org.hibernate.internal.util.collections.LazyIterator;
 import org.hibernate.persister.entity.EntityPersister;
 
 import org.jboss.logging.Logger;
@@ -77,7 +76,7 @@ public abstract class AbstractFlushingEventListener implements JpaBootstrapSensi
 		EventSource session = event.getSession();
 
 		final PersistenceContext persistenceContext = session.getPersistenceContextInternal();
-		session.getInterceptor().preFlush( new LazyIterator( persistenceContext.getEntitiesByKey() ) );
+		session.getInterceptor().preFlush( persistenceContext.managedEntitiesIterator() );
 
 		prepareEntityFlushes( session, persistenceContext );
 		// we could move this inside if we wanted to
@@ -397,6 +396,7 @@ public abstract class AbstractFlushingEventListener implements JpaBootstrapSensi
 	}
 
 	protected void postPostFlush(SessionImplementor session) {
-		session.getInterceptor().postFlush( new LazyIterator( session.getPersistenceContextInternal().getEntitiesByKey() ) );
+		session.getInterceptor().postFlush( session.getPersistenceContextInternal().managedEntitiesIterator() );
 	}
+
 }

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

@@ -1,40 +0,0 @@
-/*
- * Hibernate, Relational Persistence for Idiomatic Java
- *
- * License: GNU Lesser General Public License (LGPL), version 2.1 or later.
- * See the lgpl.txt file in the root directory or <http://www.gnu.org/licenses/lgpl-2.1.html>.
- */
-package org.hibernate.internal.util.collections;
-
-import java.util.Iterator;
-import java.util.Map;
-
-public final class LazyIterator implements Iterator {
-	
-	private final Map map;
-	private Iterator iterator;
-	
-	private Iterator getIterator() {
-		if (iterator==null) {
-			iterator = map.values().iterator();
-		}
-		return iterator;
-	}
-
-	public LazyIterator(Map map) {
-		this.map = map;
-	}
-	
-	public boolean hasNext() {
-		return getIterator().hasNext();
-	}
-
-	public Object next() {
-		return getIterator().next();
-	}
-
-	public void remove() {
-		throw new UnsupportedOperationException();
-	}
-
-}

hibernate-core/src/test/java/org/hibernate/test/stateless/StatelessSessionPersistentContextTest.java

@@ -6,6 +6,7 @@
  */
 package org.hibernate.test.stateless;
 
+import java.util.Collections;
 import java.util.function.Consumer;
 import javax.persistence.Entity;
 import javax.persistence.GeneratedValue;
@@ -115,10 +116,18 @@ public class StatelessSessionPersistentContextTest extends BaseCoreFunctionalTes
 				"StatelessSession: PersistenceContext has not been cleared",
 				persistenceContextInternal.getEntitiesByKey().isEmpty()
 		);
+		assertTrue(
+				"StatelessSession: PersistenceContext has not been cleared",
+				persistenceContextInternal.managedEntitiesIterator() == Collections.emptyIterator()
+		);
 		assertTrue(
 				"StatelessSession: PersistenceContext has not been cleared",
 				persistenceContextInternal.getCollectionsByKey().isEmpty()
 		);
+		assertTrue(
+				"StatelessSession: PersistenceContext has not been cleared",
+				persistenceContextInternal.getCollectionsByKey() == Collections.emptyMap()
+		);
 	}
 
 	@Entity(name = "TestEntity")