HHH-17309 Introduce EntityHolder to unify proxy/entity management in single map

hibernate-core/src/main/java/org/hibernate/action/internal/EntityDeleteAction.java

@@ -172,8 +172,7 @@ public class EntityDeleteAction extends EntityAction {
 		}
 		entry.postDelete();
 		final EntityKey key = entry.getEntityKey();
-		persistenceContext.removeEntity( key );
+		persistenceContext.removeEntityHolder( key );
-		persistenceContext.removeProxy( key );
 		removeCacheItem( ck );
 		persistenceContext.getNaturalIdResolutions().removeSharedResolution( id, naturalIdValues, persister );
 		postDelete();

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

@@ -20,7 +20,7 @@ import java.util.IdentityHashMap;
 import java.util.Iterator;
 import java.util.Map;
 import java.util.Map.Entry;
-import java.util.concurrent.ConcurrentMap;
+import java.util.NoSuchElementException;
 import java.util.function.BiConsumer;
 import java.util.function.Consumer;
 import java.util.function.Supplier;
@@ -41,6 +41,7 @@ import org.hibernate.engine.spi.BatchFetchQueue;
 import org.hibernate.engine.spi.CollectionEntry;
 import org.hibernate.engine.spi.CollectionKey;
 import org.hibernate.engine.spi.EntityEntry;
+import org.hibernate.engine.spi.EntityHolder;
 import org.hibernate.engine.spi.EntityKey;
 import org.hibernate.engine.spi.EntityUniqueKey;
 import org.hibernate.engine.spi.NaturalIdResolutions;
@@ -53,7 +54,6 @@ import org.hibernate.engine.spi.SharedSessionContractImplementor;
 import org.hibernate.engine.spi.Status;
 import org.hibernate.internal.CoreMessageLogger;
 import org.hibernate.internal.util.collections.CollectionHelper;
-import org.hibernate.internal.util.collections.ConcurrentReferenceHashMap;
 import org.hibernate.internal.util.collections.IdentityMap;
 import org.hibernate.metamodel.spi.MappingMetamodelImplementor;
 import org.hibernate.persister.collection.CollectionPersister;
@@ -106,13 +106,11 @@ public class StatefulPersistenceContext implements PersistenceContext {
 	 */
 
 	// Loaded entity instances, by EntityKey
-	private HashMap<EntityKey, Object> entitiesByKey;
+	private HashMap<EntityKey, EntityHolderImpl> entitiesByKey;
 
 	// Loaded entity instances, by EntityUniqueKey
 	private HashMap<EntityUniqueKey, Object> entitiesByUniqueKey;
 
-	// Entity proxies, by EntityKey
-	private ConcurrentReferenceHashMap<EntityKey, Object> proxiesByKey;
 
 	// Snapshots of current database state for entities
 	// that have *not* been loaded
@@ -169,18 +167,11 @@ public class StatefulPersistenceContext implements PersistenceContext {
 		this.entityEntryContext = new EntityEntryContext( this );
 	}
 
-	private ConcurrentMap<EntityKey, Object> getOrInitializeProxiesByKey() {
+	private Map<EntityKey, EntityHolderImpl> getOrInitializeEntitiesByKey() {
-		if ( proxiesByKey == null ) {
+		if ( entitiesByKey == null ) {
-			proxiesByKey = new ConcurrentReferenceHashMap<>(
+			entitiesByKey = CollectionHelper.mapOfSize( INIT_COLL_SIZE );
-					INIT_COLL_SIZE,
-					.75f,
-					1,
-					ConcurrentReferenceHashMap.ReferenceType.STRONG,
-					ConcurrentReferenceHashMap.ReferenceType.WEAK,
-					null
-			);
 		}
-		return proxiesByKey;
+		return entitiesByKey;
 	}
 
 	@Override
@@ -229,11 +220,11 @@ public class StatefulPersistenceContext implements PersistenceContext {
 
 	@Override
 	public void clear() {
-		if ( proxiesByKey != null ) {
+		if ( entitiesByKey != null ) {
 			//Strictly avoid lambdas in this case
-			for ( Object value : proxiesByKey.values() ) {
+			for ( EntityHolderImpl value : entitiesByKey.values() ) {
-				if ( value != null) {
+				if ( value != null && value.proxy != null ) {
-					HibernateProxy.extractLazyInitializer( value ).unsetSession();
+					HibernateProxy.extractLazyInitializer( value.proxy ).unsetSession();
 				}
 			}
 		}
@@ -253,7 +244,6 @@ public class StatefulPersistenceContext implements PersistenceContext {
 		nonlazyCollections = null;
 		collectionEntries = null;
 		unownedCollections = null;
-		proxiesByKey = null;
 		nullifiableEntityKeys = null;
 		deletedUnloadedEntityKeys = null;
 		if ( batchFetchQueue != null ) {
@@ -380,11 +370,21 @@ public class StatefulPersistenceContext implements PersistenceContext {
 	}
 
 	@Override
-	public void addEntity(EntityKey key, Object entity) {
+	public EntityHolderImpl getEntityHolder(EntityKey key) {
-		if ( entitiesByKey == null ) {
+		return entitiesByKey == null ? null : entitiesByKey.get( key );
-			entitiesByKey = CollectionHelper.mapOfSize( INIT_COLL_SIZE );
+	}
+
+	@Override
+	public boolean containsEntityHolder(EntityKey key) {
+		return entitiesByKey != null && entitiesByKey.get( key ) != null;
+	}
+
+	@Override
+	public void addEntity(EntityKey key, Object entity) {
+		final EntityHolderImpl holder = getOrInitializeEntitiesByKey().putIfAbsent( key, EntityHolderImpl.forEntity( entity ) );
+		if ( holder != null ) {
+			holder.entity = entity;
 		}
-		entitiesByKey.put( key, entity );
 		final BatchFetchQueue fetchQueue = this.batchFetchQueue;
 		if ( fetchQueue != null ) {
 			fetchQueue.removeBatchLoadableEntityKey( key );
@@ -393,20 +393,36 @@ public class StatefulPersistenceContext implements PersistenceContext {
 
 	@Override
 	public Object getEntity(EntityKey key) {
-		return entitiesByKey == null ? null : entitiesByKey.get( key );
+		final EntityHolderImpl holder = entitiesByKey == null ? null : entitiesByKey.get( key );
+		return holder == null ? null : holder.entity;
 	}
 
 	@Override
 	public boolean containsEntity(EntityKey key) {
-		return entitiesByKey != null && entitiesByKey.containsKey( key );
+		final EntityHolderImpl holder = entitiesByKey == null ? null : entitiesByKey.get( key );
+		return holder != null && holder.entity != null;
 	}
 
 	@Override
 	public Object removeEntity(EntityKey key) {
-		final Object entity;
+		final EntityHolderImpl holder = removeEntityHolder( key );
+		if ( holder != null ) {
+			final Object entity = holder.entity;
+			if ( holder.proxy != null ) {
+				entitiesByKey.put( key, holder );
+			}
+			return entity;
+		}
+		return null;
+	}
+
+	@Override
+	public EntityHolderImpl removeEntityHolder(EntityKey key) {
+		final EntityHolderImpl holder;
 		if ( entitiesByKey != null ) {
-			entity = entitiesByKey.remove( key );
+			holder = entitiesByKey.remove( key );
 			if ( entitiesByUniqueKey != null ) {
+				final Object entity = holder == null ? null : holder.entity;
 				final Iterator<?> itr = entitiesByUniqueKey.values().iterator();
 				while ( itr.hasNext() ) {
 					if ( itr.next() == entity ) {
@@ -416,7 +432,7 @@ public class StatefulPersistenceContext implements PersistenceContext {
 			}
 		}
 		else {
-			entity = null;
+			holder = null;
 		}
 
 		// Clear all parent cache
@@ -432,7 +448,7 @@ public class StatefulPersistenceContext implements PersistenceContext {
 			fetchQueue.removeBatchLoadableEntityKey( key );
 			fetchQueue.removeSubselect( key );
 		}
-		return entity;
+		return holder;
 	}
 
 	@Override
@@ -578,7 +594,14 @@ public class StatefulPersistenceContext implements PersistenceContext {
 
 	@Override
 	public boolean containsProxy(Object entity) {
-		return proxiesByKey != null && proxiesByKey.containsValue( entity );
+		if ( entitiesByKey != null ) {
+			for ( EntityHolderImpl holder : entitiesByKey.values() ) {
+				if ( holder.proxy == entity ) {
+					return true;
+				}
+			}
+		}
+		return false;
 	}
 
 	@Override
@@ -631,7 +654,13 @@ public class StatefulPersistenceContext implements PersistenceContext {
 					.getEntityDescriptor( li.getEntityName() );
 			final EntityKey key = session.generateEntityKey( li.getInternalIdentifier(), persister );
 		  	// any earlier proxy takes precedence
-			getOrInitializeProxiesByKey().putIfAbsent( key, proxy );
+			final EntityHolderImpl holder = getOrInitializeEntitiesByKey().putIfAbsent(
+					key,
+					EntityHolderImpl.forProxy( proxy )
+			);
+			if ( holder != null && holder.proxy == null ) {
+				holder.proxy = proxy;
+			}
 			proxy.getHibernateLazyInitializer().setSession( session );
 		}
 	}
@@ -733,8 +762,11 @@ public class StatefulPersistenceContext implements PersistenceContext {
 	}
 
 	private Object removeProxyByKey(final EntityKey key) {
-		if ( proxiesByKey != null ) {
+		final EntityHolderImpl entityHolder;
-			return proxiesByKey.remove( key );
+		if ( entitiesByKey != null && ( entityHolder = entitiesByKey.get( key ) ) != null ) {
+			Object proxy = entityHolder.proxy;
+			entityHolder.proxy = null;
+			return proxy;
 		}
 		return null;
 	}
@@ -759,10 +791,10 @@ public class StatefulPersistenceContext implements PersistenceContext {
 
 	@Override
 	public void addEnhancedProxy(EntityKey key, PersistentAttributeInterceptable entity) {
-		if ( entitiesByKey == null ) {
+		final EntityHolderImpl holder = getOrInitializeEntitiesByKey().putIfAbsent( key, EntityHolderImpl.forEntity( entity ) );
-			entitiesByKey = CollectionHelper.mapOfSize( INIT_COLL_SIZE );
+		if ( holder != null ) {
+			holder.entity = entity;
 		}
-		entitiesByKey.put( key, entity );
 	}
 
 	@Override
@@ -1044,12 +1076,16 @@ public class StatefulPersistenceContext implements PersistenceContext {
 
 	@Override
 	public Object getProxy(EntityKey key) {
-		return proxiesByKey == null ? null : proxiesByKey.get( key );
+		final EntityHolderImpl holder = entitiesByKey == null ? null : entitiesByKey.get( key );
+		return holder == null ? null : holder.proxy;
 	}
 
 	@Override
 	public void addProxy(EntityKey key, Object proxy) {
-		getOrInitializeProxiesByKey().put( key, proxy );
+		final EntityHolderImpl holder = getOrInitializeEntitiesByKey().putIfAbsent( key, EntityHolderImpl.forProxy( proxy ) );
+		if ( holder != null ) {
+			holder.proxy = proxy;
+		}
 	}
 
 	@Override
@@ -1077,12 +1113,60 @@ public class StatefulPersistenceContext implements PersistenceContext {
 	@Deprecated
 	@Override
 	public Map<EntityKey,Object> getEntitiesByKey() {
-		return entitiesByKey == null ? Collections.emptyMap() : entitiesByKey;
+		if ( entitiesByKey == null ) {
+			return Collections.emptyMap();
+		}
+		final HashMap<EntityKey, Object> result = CollectionHelper.mapOfSize( entitiesByKey.size() );
+		for ( Entry<EntityKey, EntityHolderImpl> entry : entitiesByKey.entrySet() ) {
+			if ( entry.getValue().entity != null ) {
+				result.put( entry.getKey(), entry.getValue().entity );
+			}
+		}
+		return result;
+	}
+
+	@Override
+	public Map<EntityKey, EntityHolder> getEntityHoldersByKey() {
+		//noinspection unchecked,rawtypes
+		return (Map) entitiesByKey;
 	}
 
 	@Override
 	public Iterator<Object> managedEntitiesIterator() {
-		return entitiesByKey == null ? Collections.emptyIterator() : entitiesByKey.values().iterator();
+		if ( entitiesByKey == null ) {
+			return Collections.emptyIterator();
+		}
+		final Iterator<EntityHolderImpl> iterator = entitiesByKey.values().iterator();
+		final var iter = new Iterator<Object>() {
+			Object next;
+			void prepareNext() {
+				next = null;
+				while ( iterator.hasNext() ) {
+					final EntityHolderImpl next = iterator.next();
+					if ( next.entity != null ) {
+						this.next = next.entity;
+						break;
+					}
+				}
+			}
+
+			@Override
+			public boolean hasNext() {
+				return next != null;
+			}
+
+			@Override
+			public Object next() {
+				final Object next = this.next;
+				if ( next == null ) {
+					throw new NoSuchElementException();
+				}
+				prepareNext();
+				return next;
+			}
+		};
+		iter.prepareNext();
+		return iter;
 	}
 
 	@Override
@@ -1525,7 +1609,8 @@ public class StatefulPersistenceContext implements PersistenceContext {
 
 	@Override
 	public void replaceDelayedEntityIdentityInsertKeys(EntityKey oldKey, Object generatedId) {
-		final Object entity = entitiesByKey == null ? null : entitiesByKey.remove( oldKey );
+		final EntityHolderImpl holder = entitiesByKey == null ? null : entitiesByKey.remove( oldKey );
+		final Object entity = holder == null ? null : holder.entity;
 		final EntityEntry oldEntry = entityEntryContext.removeEntityEntry( entity );
 		this.parentsByChild = null;
 
@@ -1558,21 +1643,20 @@ public class StatefulPersistenceContext implements PersistenceContext {
 		oos.writeBoolean( defaultReadOnly );
 		oos.writeBoolean( hasNonReadOnlyEntities );
 
-		final Serializer<Entry<EntityKey, Object>> entityKeySerializer = (entry, stream) -> {
+		writeMapToStream( entitiesByKey, oos, "entitiesByKey", (entry, stream) -> {
+			entry.getKey().serialize( stream );
+			final EntityHolderImpl holder = entry.getValue();
+			stream.writeObject( holder.entity );
+			stream.writeObject( holder.proxy );
+		} );
+		writeMapToStream( entitiesByUniqueKey, oos, "entitiesByUniqueKey", (entry, stream) -> {
 			entry.getKey().serialize( stream );
 			stream.writeObject( entry.getValue() );
-		};
+		} );
-
+		writeMapToStream( entitySnapshotsByKey, oos, "entitySnapshotsByKey", (entry, stream) -> {
-		writeMapToStream( entitiesByKey, oos, "entitiesByKey", entityKeySerializer );
-		writeMapToStream(
-				entitiesByUniqueKey,
-				oos, "entitiesByUniqueKey", (entry, stream) -> {
 			entry.getKey().serialize( stream );
 			stream.writeObject( entry.getValue() );
-				}
+		} );
-		);
-		writeMapToStream( proxiesByKey, oos, "proxiesByKey", entityKeySerializer );
-		writeMapToStream( entitySnapshotsByKey, oos, "entitySnapshotsByKey", entityKeySerializer );
 
 		entityEntryContext.serialize( oos );
 		writeMapToStream(
@@ -1678,7 +1762,24 @@ public class StatefulPersistenceContext implements PersistenceContext {
 			}
 			rtn.entitiesByKey = CollectionHelper.mapOfSize(Math.max(count, INIT_COLL_SIZE));
 			for ( int i = 0; i < count; i++ ) {
-				rtn.entitiesByKey.put( EntityKey.deserialize( ois, sfi ), ois.readObject() );
+				final EntityKey ek = EntityKey.deserialize( ois, sfi );
+				final Object entity = ois.readObject();
+				final Object proxy = ois.readObject();
+				final EntityHolderImpl holder = EntityHolderImpl.forEntity( entity );
+				if ( proxy != null ) {
+					final LazyInitializer lazyInitializer = HibernateProxy.extractLazyInitializer( proxy );
+					if ( lazyInitializer != null ) {
+						lazyInitializer.setSession( session );
+						holder.proxy = proxy;
+					}
+					else {
+						// otherwise, the proxy was pruned during the serialization process
+						if ( traceEnabled ) {
+							LOG.trace( "Encountered pruned proxy" );
+						}
+					}
+				}
+				rtn.entitiesByKey.put( ek, holder );
 			}
 
 			count = ois.readInt();
@@ -1692,26 +1793,6 @@ public class StatefulPersistenceContext implements PersistenceContext {
 				}
 			}
 
-			count = ois.readInt();
-			if ( traceEnabled ) {
-				LOG.trace( "Starting deserialization of [" + count + "] proxiesByKey entries" );
-			}
-			for ( int i = 0; i < count; i++ ) {
-				final EntityKey ek = EntityKey.deserialize( ois, sfi );
-				final Object proxy = ois.readObject();
-				final LazyInitializer lazyInitializer = HibernateProxy.extractLazyInitializer( proxy );
-				if ( lazyInitializer != null ) {
-					lazyInitializer.setSession( session );
-					rtn.getOrInitializeProxiesByKey().put( ek, proxy );
-				}
-				else {
-					// otherwise, the proxy was pruned during the serialization process
-					if ( traceEnabled ) {
-						LOG.trace( "Encountered pruned proxy" );
-					}
-				}
-			}
-
 			count = ois.readInt();
 			if ( traceEnabled ) {
 				LOG.trace( "Starting deserialization of [" + count + "] entitySnapshotsByKey entries" );
@@ -1920,6 +2001,34 @@ public class StatefulPersistenceContext implements PersistenceContext {
 		}
 	}
 
+	private static class EntityHolderImpl implements EntityHolder, Serializable {
+		Object entity;
+		Object proxy;
+
+		private EntityHolderImpl(Object entity, Object proxy) {
+			this.entity = entity;
+			this.proxy = proxy;
+		}
+
+		@Override
+		public Object getEntity() {
+			return entity;
+		}
+
+		@Override
+		public Object getProxy() {
+			return proxy;
+		}
+
+		public static EntityHolderImpl forProxy(Object proxy) {
+			return new EntityHolderImpl( null, proxy );
+		}
+
+		public static EntityHolderImpl forEntity(Object entity) {
+			return new EntityHolderImpl( entity, null );
+		}
+	}
+
 	// NATURAL ID RESOLUTION HANDLING ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 
 	private NaturalIdResolutionsImpl naturalIdResolutions;

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

@@ -0,0 +1,12 @@
+/*
+ * 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.engine.spi;
+
+public interface EntityHolder {
+	Object getEntity();
+	Object getProxy();
+}

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

@@ -482,12 +482,24 @@ public interface PersistenceContext {
 //	@Deprecated
 //	HashSet getNullifiableEntityKeys();
 
+	EntityHolder getEntityHolder(EntityKey key);
+
+	boolean containsEntityHolder(EntityKey key);
+
+	EntityHolder removeEntityHolder(EntityKey key);
+
 	/**
 	 * Doubly internal
 	 */
 	@Internal
 	Map<EntityKey,Object> getEntitiesByKey();
 
+	/**
+	 * Doubly internal
+	 */
+	@Internal
+	Map<EntityKey,EntityHolder> getEntityHoldersByKey();
+
 	/**
 	 * Provides access to the entity/EntityEntry combos associated with the persistence context in a manner that
 	 * is safe from reentrant access.  Specifically, it is safe from additions/removals while iterating.

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

@@ -122,7 +122,7 @@ public abstract class AbstractFlushingEventListener implements JpaBootstrapSensi
 				persistenceContext.getCollectionEntriesSize()
 		);
 		new EntityPrinter( session.getFactory() ).toString(
-				persistenceContext.getEntitiesByKey().entrySet()
+				persistenceContext.getEntityHoldersByKey().entrySet()
 		);
 	}
 

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

@@ -11,6 +11,7 @@ import org.hibernate.engine.internal.Cascade;
 import org.hibernate.engine.internal.CascadePoint;
 import org.hibernate.engine.spi.CascadingActions;
 import org.hibernate.engine.spi.EntityEntry;
+import org.hibernate.engine.spi.EntityHolder;
 import org.hibernate.engine.spi.EntityKey;
 import org.hibernate.engine.spi.PersistenceContext;
 import org.hibernate.event.spi.EventSource;
@@ -57,9 +58,9 @@ public class DefaultEvictEventListener implements EvictEventListener {
 					.getMappingMetamodel()
 					.getEntityDescriptor( lazyInitializer.getEntityName() );
 			final EntityKey key = source.generateEntityKey( id, persister );
-			persistenceContext.removeProxy( key );
+			final EntityHolder holder = persistenceContext.removeEntityHolder( key );
 			if ( !lazyInitializer.isUninitialized() ) {
-				final Object entity = persistenceContext.removeEntity( key );
+				final Object entity = holder.getEntity();
 				if ( entity != null ) {
 					EntityEntry entry = persistenceContext.removeEntry( entity );
 					doEvict( entity, key, entry.getPersister(), event.getSession() );

hibernate-core/src/main/java/org/hibernate/internal/util/EntityPrinter.java

@@ -12,6 +12,7 @@ import java.util.Map;
 import org.hibernate.Hibernate;
 import org.hibernate.HibernateException;
 import org.hibernate.bytecode.enhance.spi.LazyPropertyInitializer;
+import org.hibernate.engine.spi.EntityHolder;
 import org.hibernate.engine.spi.EntityKey;
 import org.hibernate.engine.spi.SessionFactoryImplementor;
 import org.hibernate.engine.spi.TypedValue;
@@ -103,19 +104,23 @@ public final class EntityPrinter {
 	}
 
 	// Cannot use Map as an argument because it clashes with the previous method (due to type erasure)
-	public void toString(Iterable<Map.Entry<EntityKey, Object>> entitiesByEntityKey) throws HibernateException {
+	public void toString(Iterable<Map.Entry<EntityKey, EntityHolder>> entitiesByEntityKey) throws HibernateException {
 		if ( !LOG.isDebugEnabled() || !entitiesByEntityKey.iterator().hasNext() ) {
 			return;
 		}
 
 		LOG.debug( "Listing entities:" );
 		int i = 0;
-		for ( Map.Entry<EntityKey, Object> entityKeyAndEntity : entitiesByEntityKey ) {
+		for ( Map.Entry<EntityKey, EntityHolder> entityKeyAndEntity : entitiesByEntityKey ) {
+			final EntityHolder holder = entityKeyAndEntity.getValue();
+			if ( holder.getEntity() == null ) {
+				continue;
+			}
 			if ( i++ > 20 ) {
 				LOG.debug( "More......" );
 				break;
 			}
-			LOG.debug( toString( entityKeyAndEntity.getKey().getEntityName(), entityKeyAndEntity.getValue() ) );
+			LOG.debug( toString( entityKeyAndEntity.getKey().getEntityName(), holder.getEntity() ) );
 		}
 	}