HHH-19080 Allow Hibernate Reactive to initialize the non lazy collections via PersistenceContext

Hibernate Reactive needs a reactive version of the PersistenceContext to
load non-lazy collections. But, now that `SatefulPersistenceContext` is
package-private, Hibernate Reactive cannot extend the class anymore.

These change make it possible to access the internal snapshots, so that the
reactive implementation can implement the interface and doesn't need to
extends `SatefulPersistenceContext`.

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,21 @@ public Map<EntityKey,Object> getEntitiesByKey() {
 		return result;
 	}
 
+	// Used by Hibernate Reactive
+	@Override
+	public Map<EntityKey, Object> getEntitySnapshotsByKey() {
+		return entitySnapshotsByKey;
+	}
+
+	// Used by Hibernate Reactive
+	@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,14 @@ EntityHolder claimEntityHolderIfPossible(
 	@Internal
 	Map<EntityKey,Object> getEntitiesByKey();
 
+	// Used by Hibernate Reactive
+	@Internal
+	Map<EntityKey,Object> getEntitySnapshotsByKey();
+
+	// Used by Hibernate Reactive
+	@Internal
+	Map<EntityKey,Object> getOrInitializeEntitySnapshotsByKey();
+
 	/**
 	 * Doubly internal
 	 */