From 1c86d94a15f34c2e0bebf9612f23f5142248c971 Mon Sep 17 00:00:00 2001
From: Christian Beikov <christian.beikov@gmail.com>
Date: Tue, 15 Aug 2023 15:34:42 +0200
Subject: [PATCH] HHH-17078 Register/Deregister load context on every scroll
 operation

---
 .../FetchingScrollableResultsImpl.java        | 54 ++++++++++++-------
 .../internal/ScrollableResultsImpl.java       | 10 ++--
 .../spi/ScrollableResultsConsumer.java        |  1 -
 .../orm/test/iterate/ScrollTest.java          |  2 +
 4 files changed, 41 insertions(+), 26 deletions(-)

diff --git a/hibernate-core/src/main/java/org/hibernate/internal/FetchingScrollableResultsImpl.java b/hibernate-core/src/main/java/org/hibernate/internal/FetchingScrollableResultsImpl.java
index 7462c297c5..26d51b7d6c 100644
--- a/hibernate-core/src/main/java/org/hibernate/internal/FetchingScrollableResultsImpl.java
+++ b/hibernate-core/src/main/java/org/hibernate/internal/FetchingScrollableResultsImpl.java
@@ -8,6 +8,7 @@ package org.hibernate.internal;
 
 import org.hibernate.HibernateException;
 import org.hibernate.engine.spi.EntityKey;
+import org.hibernate.engine.spi.PersistenceContext;
 import org.hibernate.engine.spi.SharedSessionContractImplementor;
 import org.hibernate.sql.results.graph.Initializer;
 import org.hibernate.sql.results.graph.entity.EntityInitializer;
@@ -15,6 +16,7 @@ import org.hibernate.sql.results.internal.RowProcessingStateStandardImpl;
 import org.hibernate.sql.results.jdbc.internal.JdbcValuesSourceProcessingStateStandardImpl;
 import org.hibernate.sql.results.jdbc.spi.JdbcValues;
 import org.hibernate.sql.results.jdbc.spi.JdbcValuesSourceProcessingOptions;
+import org.hibernate.sql.results.spi.LoadContexts;
 import org.hibernate.sql.results.spi.RowReader;
 
 /**
@@ -299,8 +301,9 @@ public class FetchingScrollableResultsImpl<R> extends AbstractScrollableResults<
 	}
 
 	private boolean prepareCurrentRow() {
-		if ( getRowProcessingState().isBeforeFirst() ) {
-			getRowProcessingState().next();
+		final RowProcessingStateStandardImpl rowProcessingState = getRowProcessingState();
+		if ( rowProcessingState.isBeforeFirst() ) {
+			rowProcessingState.next();
 		}
 
 		final RowReader<R> rowReader = getRowReader();
@@ -309,32 +312,43 @@ public class FetchingScrollableResultsImpl<R> extends AbstractScrollableResults<
 		boolean resultProcessed = false;
 
 		final EntityKey entityKey = getEntityKey();
+		final PersistenceContext persistenceContext = rowProcessingState.getSession().getPersistenceContext();
+		final LoadContexts loadContexts = persistenceContext.getLoadContexts();
 
-		currentRow = rowReader.readRow( getRowProcessingState(), getProcessingOptions() );
+		loadContexts.register( getJdbcValuesSourceProcessingState() );
+		persistenceContext.beforeLoad();
+		try {
+			currentRow = rowReader.readRow( rowProcessingState, getProcessingOptions() );
 
-		getRowProcessingState().finishRowProcessing();
+			rowProcessingState.finishRowProcessing();
 
-		while ( !resultProcessed ) {
-			if ( getRowProcessingState().next() ) {
-				final EntityKey entityKey2 = getEntityKey();
-				if ( !entityKey.equals( entityKey2 ) ) {
-					resultInitializer.finishUpRow( getRowProcessingState() );
-					resultProcessed = true;
-					afterLast = false;
+			while ( !resultProcessed ) {
+				if ( rowProcessingState.next() ) {
+					final EntityKey entityKey2 = getEntityKey();
+					if ( !entityKey.equals( entityKey2 ) ) {
+						resultInitializer.finishUpRow( rowProcessingState );
+						resultProcessed = true;
+						afterLast = false;
+					}
+					else {
+						rowReader.readRow( rowProcessingState, getProcessingOptions() );
+						rowProcessingState.finishRowProcessing();
+					}
 				}
 				else {
-					rowReader.readRow( getRowProcessingState(), getProcessingOptions() );
-					getRowProcessingState().finishRowProcessing();
+					afterLast = true;
+					resultProcessed = true;
 				}
-			}
-			else {
-				afterLast = true;
-				resultProcessed = true;
-			}
 
+			}
+			getJdbcValuesSourceProcessingState().finishUp();
 		}
-		getJdbcValuesSourceProcessingState().finishUp();
-		getRowProcessingState().getSession().getPersistenceContext().initializeNonLazyCollections();
+		finally {
+			persistenceContext.afterLoad();
+			loadContexts.deregister( getJdbcValuesSourceProcessingState() );
+		}
+		persistenceContext.initializeNonLazyCollections();
+		afterScrollOperation();
 		return afterLast;
 	}
 
diff --git a/hibernate-core/src/main/java/org/hibernate/internal/ScrollableResultsImpl.java b/hibernate-core/src/main/java/org/hibernate/internal/ScrollableResultsImpl.java
index 0c7304e6a0..c23b2e3fd4 100644
--- a/hibernate-core/src/main/java/org/hibernate/internal/ScrollableResultsImpl.java
+++ b/hibernate-core/src/main/java/org/hibernate/internal/ScrollableResultsImpl.java
@@ -13,6 +13,7 @@ import org.hibernate.sql.results.internal.RowProcessingStateStandardImpl;
 import org.hibernate.sql.results.jdbc.internal.JdbcValuesSourceProcessingStateStandardImpl;
 import org.hibernate.sql.results.jdbc.spi.JdbcValues;
 import org.hibernate.sql.results.jdbc.spi.JdbcValuesSourceProcessingOptions;
+import org.hibernate.sql.results.spi.LoadContexts;
 import org.hibernate.sql.results.spi.RowReader;
 
 /**
@@ -124,19 +125,18 @@ public class ScrollableResultsImpl<R> extends AbstractScrollableResults<R> {
 		}
 
 		final PersistenceContext persistenceContext = getPersistenceContext().getPersistenceContext();
-
+		final LoadContexts loadContexts = persistenceContext.getLoadContexts();
+		loadContexts.register( getJdbcValuesSourceProcessingState() );
 		persistenceContext.beforeLoad();
 		try {
-			currentRow = getRowReader().readRow(
-					getRowProcessingState(),
-					getProcessingOptions()
-			);
+			currentRow = getRowReader().readRow( getRowProcessingState(), getProcessingOptions() );
 
 			getRowProcessingState().finishRowProcessing();
 			getJdbcValuesSourceProcessingState().finishUp();
 		}
 		finally {
 			persistenceContext.afterLoad();
+			loadContexts.deregister( getJdbcValuesSourceProcessingState() );
 		}
 		persistenceContext.initializeNonLazyCollections();
 
diff --git a/hibernate-core/src/main/java/org/hibernate/sql/results/spi/ScrollableResultsConsumer.java b/hibernate-core/src/main/java/org/hibernate/sql/results/spi/ScrollableResultsConsumer.java
index 8c2bd15888..382b033268 100644
--- a/hibernate-core/src/main/java/org/hibernate/sql/results/spi/ScrollableResultsConsumer.java
+++ b/hibernate-core/src/main/java/org/hibernate/sql/results/spi/ScrollableResultsConsumer.java
@@ -42,7 +42,6 @@ public class ScrollableResultsConsumer<R> implements ResultsConsumer<ScrollableR
 			JdbcValuesSourceProcessingStateStandardImpl jdbcValuesSourceProcessingState,
 			RowProcessingStateStandardImpl rowProcessingState,
 			RowReader<R> rowReader) {
-		session.getPersistenceContext().getLoadContexts().register( jdbcValuesSourceProcessingState );
 		if ( containsCollectionFetches( jdbcValues.getValuesMapping() ) ) {
 			return new FetchingScrollableResultsImpl<>(
 					jdbcValues,
diff --git a/hibernate-core/src/test/java/org/hibernate/orm/test/iterate/ScrollTest.java b/hibernate-core/src/test/java/org/hibernate/orm/test/iterate/ScrollTest.java
index a6a844d1ad..79172f3a5a 100644
--- a/hibernate-core/src/test/java/org/hibernate/orm/test/iterate/ScrollTest.java
+++ b/hibernate-core/src/test/java/org/hibernate/orm/test/iterate/ScrollTest.java
@@ -64,6 +64,8 @@ public class ScrollTest {
 						s.delete( i1 );
 						s.delete( i2 );
 					}
+
+					assertTrue( s.getPersistenceContext().getLoadContexts().isLoadingFinished() );
 				}
 		);