From c96c998b77fc4fbe5dd24e67d1d4c00feb2116b7 Mon Sep 17 00:00:00 2001
From: Davide D'Alto <davide@hibernate.org>
Date: Wed, 15 Apr 2020 14:50:44 +0100
Subject: [PATCH] HHH-13997 : Add methods to SessionImpl and
 StatefulPersistenceContext; make Loader method protected

* Change visibility for Loader#registerNonExists
* Methods for creating PersistenceContext and ActionQueue
* Add StatefulPersistenceContext#initializeNonLazyCollections(Consumer<PersistentCollection> initializeAction)
---
 .gradletasknamecache                                 |  0
 .../engine/internal/StatefulPersistenceContext.java  |  7 ++++++-
 .../java/org/hibernate/internal/SessionImpl.java     | 12 ++++++++++--
 .../src/main/java/org/hibernate/loader/Loader.java   |  2 +-
 4 files changed, 17 insertions(+), 4 deletions(-)
 create mode 100644 .gradletasknamecache

diff --git a/.gradletasknamecache b/.gradletasknamecache
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/hibernate-core/src/main/java/org/hibernate/engine/internal/StatefulPersistenceContext.java b/hibernate-core/src/main/java/org/hibernate/engine/internal/StatefulPersistenceContext.java
index 20e07a9300..df740d3f3a 100644
--- a/hibernate-core/src/main/java/org/hibernate/engine/internal/StatefulPersistenceContext.java
+++ b/hibernate-core/src/main/java/org/hibernate/engine/internal/StatefulPersistenceContext.java
@@ -23,6 +23,7 @@ import java.util.Map;
 import java.util.Map.Entry;
 import java.util.concurrent.ConcurrentMap;
 import java.util.function.BiConsumer;
+import java.util.function.Consumer;
 import java.util.function.Supplier;
 
 import org.hibernate.AssertionFailure;
@@ -984,6 +985,10 @@ public class StatefulPersistenceContext implements PersistenceContext {
 
 	@Override
 	public void initializeNonLazyCollections() throws HibernateException {
+		initializeNonLazyCollections( PersistentCollection::forceInitialization );
+	}
+
+	protected void initializeNonLazyCollections(Consumer<PersistentCollection> initializeAction ) {
 		if ( loadCounter == 0 ) {
 			LOG.trace( "Initializing non-lazy collections" );
 
@@ -994,7 +999,7 @@ public class StatefulPersistenceContext implements PersistenceContext {
 				int size;
 				while ( nonlazyCollections != null && ( size = nonlazyCollections.size() ) > 0 ) {
 					//note that each iteration of the loop may add new elements
-					nonlazyCollections.remove( size - 1 ).forceInitialization();
+					initializeAction.accept( nonlazyCollections.remove( size - 1 ) );
 				}
 			}
 			finally {
diff --git a/hibernate-core/src/main/java/org/hibernate/internal/SessionImpl.java b/hibernate-core/src/main/java/org/hibernate/internal/SessionImpl.java
index 423cba01e8..f33ad21018 100644
--- a/hibernate-core/src/main/java/org/hibernate/internal/SessionImpl.java
+++ b/hibernate-core/src/main/java/org/hibernate/internal/SessionImpl.java
@@ -221,8 +221,8 @@ public class SessionImpl
 	public SessionImpl(SessionFactoryImpl factory, SessionCreationOptions options) {
 		super( factory, options );
 
-		this.actionQueue = new ActionQueue( this );
-		this.persistenceContext = new StatefulPersistenceContext( this );
+		this.persistenceContext = createPersistenceContext();
+		this.actionQueue = createActionQueue();
 
 		this.autoClear = options.shouldAutoClear();
 		this.autoClose = options.shouldAutoClose();
@@ -269,6 +269,14 @@ public class SessionImpl
 		}
 	}
 
+	protected StatefulPersistenceContext createPersistenceContext() {
+		return new StatefulPersistenceContext( this );
+	}
+
+	protected ActionQueue createActionQueue() {
+		return new ActionQueue( this );
+	}
+
 	private LockOptions getLockOptionsForRead() {
 		return this.lockOptions == null ? fastSessionServices.defaultLockOptions : this.lockOptions;
 	}
diff --git a/hibernate-core/src/main/java/org/hibernate/loader/Loader.java b/hibernate-core/src/main/java/org/hibernate/loader/Loader.java
index 6b32fd46c2..2bf01103dd 100644
--- a/hibernate-core/src/main/java/org/hibernate/loader/Loader.java
+++ b/hibernate-core/src/main/java/org/hibernate/loader/Loader.java
@@ -1311,7 +1311,7 @@ public abstract class Loader {
 	 * result set, register the fact that the the object is missing with the
 	 * session.
 	 */
-	private void registerNonExists(
+	protected void registerNonExists(
 			final EntityKey[] keys,
 			final Loadable[] persisters,
 			final SharedSessionContractImplementor session) {