Move InitializeEntitiesAndCollections to ProcessResults

Author: bahusoid

src/AsyncGenerator.yml

@@ -5,6 +5,10 @@
   applyChanges: true
   analyzation:
     methodConversion:
+#TODO 6.0: Remove ignore rule for IQueryBatchItem.ProcessResults  
+    - conversion: Ignore
+      name: ProcessResults
+      containingTypeName: IQueryBatchItem
     - conversion: Ignore
       name: PostProcessInsert
       containingTypeName: HqlSqlWalker

src/NHibernate/Async/Multi/QueryBatchItemBase.cs

@@ -108,8 +108,9 @@ public async Task<int> ProcessResultsSetAsync(DbDataReader reader, CancellationT
 					await (reader.NextResultAsync(cancellationToken)).ConfigureAwait(false);
 				}
 
-				await (InitializeEntitiesAndCollectionsAsync(reader, hydratedObjects, cancellationToken)).ConfigureAwait(false);
-
+				StopLoadingCollections(reader);
+				_reader = reader;
+				_hydratedObjects = hydratedObjects;
 				return rowCount;
 			}
 		}
@@ -123,19 +124,5 @@ public async Task ExecuteNonBatchedAsync(CancellationToken cancellationToken)
 		}
 
 		protected abstract Task<IList<TResult>> GetResultsNonBatchedAsync(CancellationToken cancellationToken);
-
-		private async Task InitializeEntitiesAndCollectionsAsync(DbDataReader reader, List<object>[] hydratedObjects, CancellationToken cancellationToken)
-		{
-			cancellationToken.ThrowIfCancellationRequested();
-			for (var i = 0; i < _queryInfos.Count; i++)
-			{
-				var queryInfo = _queryInfos[i];
-				if (queryInfo.IsResultFromCache)
-					continue;
-				await (queryInfo.Loader.InitializeEntitiesAndCollectionsAsync(
-					hydratedObjects[i], reader, Session, queryInfo.Parameters.IsReadOnly(Session),
-					queryInfo.CacheBatcher, cancellationToken)).ConfigureAwait(false);
-			}
-		}
 	}
 }

src/NHibernate/Engine/Loading/CollectionLoadContext.cs

@@ -124,6 +124,8 @@ public IPersistentCollection GetLoadingCollection(ICollectionPersister persister
 			}
 			else
 			{
+				if (loadingCollectionEntry.StopLoading)
+					return null;
 				if (loadingCollectionEntry.ResultSet == resultSet)
 				{
 					log.Debug("found loading collection bound to current result set processing; reading row");
@@ -398,5 +400,17 @@ public override string ToString()
 		{
 			return base.ToString() + "<rs=" + ResultSet + ">";
 		}
+
+		internal void StopLoadingCollections(ICollectionPersister[] collectionPersisters)
+		{
+			foreach (var collectionKey in localLoadingCollectionKeys)
+			{
+				var loadingCollectionEntry = LoadContext.LocateLoadingCollectionEntry(collectionKey);
+				if (loadingCollectionEntry != null && Array.IndexOf(collectionPersisters, loadingCollectionEntry.Persister) >= 0)
+				{
+					loadingCollectionEntry.StopLoading = true;
+				}
+			}
+		}
 	}
 }

src/NHibernate/Engine/Loading/LoadingCollectionEntry.cs

@@ -44,6 +44,8 @@ public IPersistentCollection Collection
 			get { return collection; }
 		}
 
+		public bool StopLoading { get; set; }
+
 		public override string ToString()
 		{
 			return GetType().FullName + "<rs=" + ResultSet + ", coll=" + MessageHelper.InfoString(Persister.Role, Key) + ">@" + Convert.ToString(GetHashCode(), 16);

src/NHibernate/Loader/Loader.cs

@@ -672,6 +672,15 @@ internal void InitializeEntitiesAndCollections(
 			}
 		}
 
+		internal void StopLoadingCollections(ISessionImplementor session, DbDataReader reader)
+		{
+			var collectionPersisters = CollectionPersisters;
+			if (collectionPersisters == null || collectionPersisters.Length == 0)
+				return;
+
+			session.PersistenceContext.LoadContexts.GetCollectionLoadContext(reader).StopLoadingCollections(collectionPersisters);
+		}
+
 		private void EndCollectionLoad(DbDataReader reader, ISessionImplementor session, ICollectionPersister collectionPersister)
 		{
 			//this is a query and we are loading multiple instances of the same collection role

src/NHibernate/Multi/QueryBatchItemBase.cs

@@ -21,6 +21,8 @@ public abstract partial class QueryBatchItemBase<TResult> : IQueryBatchItem<TRes
 		private List<QueryInfo> _queryInfos;
 		private CacheMode? _cacheMode;
 		private IList<TResult> _finalResults;
+		private DbDataReader _reader;
+		private List<object>[] _hydratedObjects;
 
 		protected class QueryInfo : ICachingInformation
 		{
@@ -247,8 +249,9 @@ public int ProcessResultsSet(DbDataReader reader)
 					reader.NextResult();
 				}
 
-				InitializeEntitiesAndCollections(reader, hydratedObjects);
-
+				StopLoadingCollections(reader);
+				_reader = reader;
+				_hydratedObjects = hydratedObjects;
 				return rowCount;
 			}
 		}
@@ -258,6 +261,7 @@ public void ProcessResults()
 		{
 			ThrowIfNotInitialized();
 
+			InitializeEntitiesAndCollections(_reader, _hydratedObjects);
 			for (var i = 0; i < _queryInfos.Count; i++)
 			{
 				var queryInfo = _queryInfos[i];
@@ -329,6 +333,17 @@ private void InitializeEntitiesAndCollections(DbDataReader reader, List<object>[
 			}
 		}
 
+		private void StopLoadingCollections(DbDataReader reader)
+		{
+			for (var i = 0; i < _queryInfos.Count; i++)
+			{
+				var queryInfo = _queryInfos[i];
+				if (queryInfo.IsResultFromCache)
+					continue;
+				queryInfo.Loader.StopLoadingCollections(Session, reader);
+			}
+		}
+
 		private void ThrowIfNotInitialized()
 		{
 			if (_queryInfos == null)