Changes for HR

Author: DavideD

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

@@ -330,10 +330,7 @@ public Object[] getDatabaseSnapshot(Object id, EntityPersister persister) throws
 		}
 		else {
 			final Object[] snapshot = persister.getDatabaseSnapshot( id, session );
-			if ( entitySnapshotsByKey == null ) {
-				entitySnapshotsByKey = CollectionHelper.mapOfSize( INIT_COLL_SIZE );
-			}
-			entitySnapshotsByKey.put( key, snapshot == null ? NO_ROW : snapshot );
+			getOrInitializeEntitySnapshotsByKey().put( key, snapshot == null ? NO_ROW : snapshot );
 			return snapshot;
 		}
 	}
@@ -1187,7 +1184,8 @@ public void initializeNonLazyCollections() throws HibernateException {
 		initializeNonLazyCollections( PersistentCollection::forceInitialization );
 	}
 
-	protected void initializeNonLazyCollections(Consumer<PersistentCollection<?>> initializeAction ) {
+	@Override
+	public void initializeNonLazyCollections(Consumer<PersistentCollection<?>> initializeAction ) {
 		if ( loadCounter == 0 ) {
 			LOG.trace( "Initializing non-lazy collections" );
 
@@ -1314,6 +1312,19 @@ public Map<EntityKey,Object> getEntitiesByKey() {
 		return result;
 	}
 
+	@Override
+	public Map<EntityKey, Object> getEntitySnapshotsByKey() {
+		return entitySnapshotsByKey;
+	}
+
+	@Override
+	public Map<EntityKey, Object> getOrInitializeEntitySnapshotsByKey() {
+		if ( entitySnapshotsByKey == null ) {
+			entitySnapshotsByKey = CollectionHelper.mapOfSize( INIT_COLL_SIZE );
+		}
+		return entitySnapshotsByKey;
+	}
+
 	@Override
 	public Map<EntityKey, EntityHolder> getEntityHoldersByKey() {
 		//noinspection unchecked,rawtypes

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

@@ -444,6 +444,18 @@ CollectionEntry addInitializedCollection(
 	 */
 	void initializeNonLazyCollections() throws HibernateException;
 
+	/**
+	 * Force initialization of all non-lazy collections encountered during
+	 * the current two-phase load (actually, this is a no-op, unless this
+	 * is the "outermost" load) allowing to customize how the initialization
+	 *  should occur
+	 *
+	 * @see #initializeNonLazyCollections()
+	 * @param initializeAction the function that initialize the collection
+	 */
+	// Used by Hibernate Reactive
+	void initializeNonLazyCollections(Consumer<PersistentCollection<?>> initializeAction);
+
 	/**
 	 * Get the {@code PersistentCollection} object for an array
 	 */
@@ -533,6 +545,12 @@ EntityHolder claimEntityHolderIfPossible(
 	@Internal
 	Map<EntityKey,Object> getEntitiesByKey();
 
+	@Internal
+	Map<EntityKey,Object> getEntitySnapshotsByKey();
+
+	@Internal
+	Map<EntityKey,Object> getOrInitializeEntitySnapshotsByKey();
+
 	/**
 	 * Doubly internal
 	 */