Improved collection batch fetching

src/NHibernate/Async/Engine/BatchFetchQueue.cs

@@ -15,6 +15,7 @@
 using NHibernate.Persister.Entity;
 using NHibernate.Util;
 using System.Collections.Generic;
+using Iesi.Collections.Generic;
 
 namespace NHibernate.Engine
 {
@@ -40,22 +41,33 @@ public async Task<object[]> GetCollectionBatchAsync(ICollectionPersister collect
 			int end = -1;
 			bool checkForEnd = false;
 
-			// this only works because collection entries are kept in a sequenced
-			// map by persistence context (maybe we should do like entities and
-			// keep a separate sequences set...)
-			foreach (DictionaryEntry me in context.CollectionEntries)
+			if (batchLoadableCollections.TryGetValue(collectionPersister.Role, out var map))
 			{
-				CollectionEntry ce = (CollectionEntry) me.Value;
-				IPersistentCollection collection = (IPersistentCollection) me.Key;
-				if (!collection.WasInitialized && ce.LoadedPersister == collectionPersister)
+				foreach (KeyValuePair<CollectionEntry, IPersistentCollection> me in map)
 				{
+					var ce = me.Key;
+					var collection = me.Value;
+					if (ce.LoadedKey == null)
+					{
+						// the LoadedKey of the CollectionEntry might be null as it might have been reset to null
+						// (see for example Collections.ProcessDereferencedCollection()
+						// and CollectionEntry.AfterAction())
+						// though we clear the queue on flush, it seems like a good idea to guard
+						// against potentially null LoadedKey:s
+						continue;
+					}
+
+					if (collection.WasInitialized)
+					{
+						log.Warn("Encountered initialized collection in BatchFetchQueue, this should not happen.");
+						continue;
+					}
+
 					if (checkForEnd && i == end)
 					{
 						return keys; //the first key found after the given key
 					}
 
-					//if ( end == -1 && count > batchSize*10 ) return keys; //try out ten batches, max
-
 					bool isEqual = collectionPersister.KeyType.IsEqual(id, ce.LoadedKey, collectionPersister.Factory);
 
 					if (isEqual)
@@ -79,6 +91,7 @@ public async Task<object[]> GetCollectionBatchAsync(ICollectionPersister collect
 					}
 				}
 			}
+
 			return keys; //we ran out of keys to try
 		}
 
@@ -92,7 +105,7 @@ public async Task<object[]> GetCollectionBatchAsync(ICollectionPersister collect
 		/// <param name="batchSize">The maximum number of keys to return</param>
 		/// <param name="cancellationToken">A cancellation token that can be used to cancel the work</param>
 		/// <returns>an array of identifiers, of length batchSize (possibly padded with nulls)</returns>
-		public async Task<object[]> GetEntityBatchAsync(IEntityPersister persister,object id,int batchSize, CancellationToken cancellationToken)
+		public async Task<object[]> GetEntityBatchAsync(IEntityPersister persister, object id, int batchSize, CancellationToken cancellationToken)
 		{
 			cancellationToken.ThrowIfCancellationRequested();
 			object[] ids = new object[batchSize];
@@ -101,9 +114,9 @@ public async Task<object[]> GetEntityBatchAsync(IEntityPersister persister,objec
 			int end = -1;
 			bool checkForEnd = false;
 
-			foreach (EntityKey key in batchLoadableEntityKeys.Keys)
+			if (batchLoadableEntityKeys.TryGetValue(persister.EntityName, out var set))
 			{
-				if (key.EntityName.Equals(persister.EntityName))
+				foreach (var key in set)
 				{
 					//TODO: this needn't exclude subclasses...
 					if (checkForEnd && i == end)

src/NHibernate/Async/Engine/Loading/CollectionLoadContext.cs

@@ -128,7 +128,9 @@ private async Task EndLoadingCollectionAsync(LoadingCollectionEntry lce, ICollec
 			{
 				log.Debug("ending loading collection [{0}]", lce);
 			}
-			ISessionImplementor session = LoadContext.PersistenceContext.Session;
+
+			var persistenceContext = LoadContext.PersistenceContext;
+			var session = persistenceContext.Session;
 
 			bool statsEnabled = session.Factory.Statistics.IsStatisticsEnabled;
 			var stopWath = new Stopwatch();
@@ -141,17 +143,17 @@ private async Task EndLoadingCollectionAsync(LoadingCollectionEntry lce, ICollec
 
 			if (persister.CollectionType.HasHolder())
 			{
-				LoadContext.PersistenceContext.AddCollectionHolder(lce.Collection);
+				persistenceContext.AddCollectionHolder(lce.Collection);
 			}
 
-			CollectionEntry ce = LoadContext.PersistenceContext.GetCollectionEntry(lce.Collection);
+			CollectionEntry ce = persistenceContext.GetCollectionEntry(lce.Collection);
 			if (ce == null)
 			{
-				ce = LoadContext.PersistenceContext.AddInitializedCollection(persister, lce.Collection, lce.Key);
+				ce = persistenceContext.AddInitializedCollection(persister, lce.Collection, lce.Key);
 			}
 			else
 			{
-				ce.PostInitialize(lce.Collection);
+				ce.PostInitialize(lce.Collection, persistenceContext);
 			}
 
 			bool addToCache = hasNoQueuedAdds && persister.HasCache && 

src/NHibernate/Async/Event/Default/DefaultInitializeCollectionEventListener.cs

@@ -127,7 +127,7 @@ private async Task<bool> InitializeCollectionFromCacheAsync(object id, ICollecti
 					CollectionCacheEntry cacheEntry = (CollectionCacheEntry)persister.CacheEntryStructure.Destructure(ce, factory);
 					await (cacheEntry.AssembleAsync(collection, persister, persistenceContext.GetCollectionOwner(id, persister), cancellationToken)).ConfigureAwait(false);
 
-					persistenceContext.GetCollectionEntry(collection).PostInitialize(collection);
+					persistenceContext.GetCollectionEntry(collection).PostInitialize(collection, persistenceContext);
 					return true;
 				}
 			}

src/NHibernate/Engine/BatchFetchQueue.cs

@@ -5,23 +5,23 @@
 using NHibernate.Persister.Entity;
 using NHibernate.Util;
 using System.Collections.Generic;
+using Iesi.Collections.Generic;
 
 namespace NHibernate.Engine
 {
 	public partial class BatchFetchQueue
 	{
-		private static readonly object Marker = new object();
+		private static readonly INHibernateLogger log = NHibernateLogger.For(typeof(BatchFetchQueue));
 
 		/// <summary>
-		/// Defines a sequence of <see cref="EntityKey" /> elements that are currently
-		/// eligible for batch fetching.
+		/// Used to hold information about the entities that are currently eligible for batch-fetching.  Ultimately
+		/// used by <see cref="GetEntityBatch" /> to build entity load batches.
 		/// </summary>
 		/// <remarks>
-		/// Even though this is a map, we only use the keys.  A map was chosen in
-		/// order to utilize a <see cref="LinkedHashMap{K, V}" /> to maintain sequencing
-		/// as well as uniqueness.
+		/// A Map structure is used to segment the keys by entity type since loading can only be done for a particular entity
+		/// type at a time.
 		/// </remarks>
-		private readonly IDictionary<EntityKey, object> batchLoadableEntityKeys = new LinkedHashMap<EntityKey, object>(8);
+		private readonly IDictionary<string, LinkedHashSet<EntityKey>> batchLoadableEntityKeys = new Dictionary<string, LinkedHashSet<EntityKey>>(8);
 
 		/// <summary>
 		/// A map of <see cref="SubselectFetch">subselect-fetch descriptors</see>
@@ -30,6 +30,7 @@ public partial class BatchFetchQueue
 		/// </summary>
 		private readonly IDictionary<EntityKey, SubselectFetch> subselectsByEntityKey = new Dictionary<EntityKey, SubselectFetch>(8);
 
+		private readonly IDictionary<string, LinkedHashMap<CollectionEntry, IPersistentCollection>> batchLoadableCollections = new Dictionary<string, LinkedHashMap<CollectionEntry, IPersistentCollection>>(8);
 		/// <summary>
 		/// The owning persistence context.
 		/// </summary>
@@ -50,6 +51,7 @@ public BatchFetchQueue(IPersistenceContext context)
 		public void Clear()
 		{
 			batchLoadableEntityKeys.Clear();
+			batchLoadableCollections.Clear();
 			subselectsByEntityKey.Clear();
 		}
 
@@ -113,7 +115,12 @@ public void AddBatchLoadableEntityKey(EntityKey key)
 		{
 			if (key.IsBatchLoadable)
 			{
-				batchLoadableEntityKeys[key] = Marker;
+				if (!batchLoadableEntityKeys.TryGetValue(key.EntityName, out var set))
+				{
+					set = new LinkedHashSet<EntityKey>();
+					batchLoadableEntityKeys.Add(key.EntityName, set);
+				}
+				set.Add(key);
 			}
 		}
 
@@ -125,7 +132,44 @@ public void AddBatchLoadableEntityKey(EntityKey key)
 		public void RemoveBatchLoadableEntityKey(EntityKey key)
 		{
 			if (key.IsBatchLoadable)
-				batchLoadableEntityKeys.Remove(key);
+			{
+				if (batchLoadableEntityKeys.TryGetValue(key.EntityName, out var set))
+				{
+					set.Remove(key);
+				}
+			}
+		}
+
+		/// <summary>
+		/// If a CollectionEntry represents a batch loadable collection, add
+		/// it to the queue.
+		/// </summary>
+		/// <param name="collection"></param>
+		/// <param name="ce"></param>
+		public void AddBatchLoadableCollection(IPersistentCollection collection, CollectionEntry ce)
+		{
+			var persister = ce.LoadedPersister;
+
+			if (!batchLoadableCollections.TryGetValue(persister.Role, out var map))
+			{
+				map = new LinkedHashMap<CollectionEntry, IPersistentCollection>();
+				batchLoadableCollections.Add(persister.Role, map);
+			}
+			map[ce] = collection;
+		}
+
+		/// <summary>
+		/// After a collection was initialized or evicted, we don't
+		/// need to batch fetch it anymore, remove it from the queue
+		/// if necessary
+		/// </summary>
+		/// <param name="ce"></param>
+		public void RemoveBatchLoadableCollection(CollectionEntry ce)
+		{
+			if (batchLoadableCollections.TryGetValue(ce.LoadedPersister.Role, out var map))
+			{
+				map.Remove(ce);
+			}
 		}
 
 		/// <summary>
@@ -143,22 +187,33 @@ public object[] GetCollectionBatch(ICollectionPersister collectionPersister, obj
 			int end = -1;
 			bool checkForEnd = false;
 
-			// this only works because collection entries are kept in a sequenced
-			// map by persistence context (maybe we should do like entities and
-			// keep a separate sequences set...)
-			foreach (DictionaryEntry me in context.CollectionEntries)
+			if (batchLoadableCollections.TryGetValue(collectionPersister.Role, out var map))
 			{
-				CollectionEntry ce = (CollectionEntry) me.Value;
-				IPersistentCollection collection = (IPersistentCollection) me.Key;
-				if (!collection.WasInitialized && ce.LoadedPersister == collectionPersister)
+				foreach (KeyValuePair<CollectionEntry, IPersistentCollection> me in map)
 				{
+					var ce = me.Key;
+					var collection = me.Value;
+					if (ce.LoadedKey == null)
+					{
+						// the LoadedKey of the CollectionEntry might be null as it might have been reset to null
+						// (see for example Collections.ProcessDereferencedCollection()
+						// and CollectionEntry.AfterAction())
+						// though we clear the queue on flush, it seems like a good idea to guard
+						// against potentially null LoadedKey:s
+						continue;
+					}
+
+					if (collection.WasInitialized)
+					{
+						log.Warn("Encountered initialized collection in BatchFetchQueue, this should not happen.");
+						continue;
+					}
+
 					if (checkForEnd && i == end)
 					{
 						return keys; //the first key found after the given key
 					}
 
-					//if ( end == -1 && count > batchSize*10 ) return keys; //try out ten batches, max
-
 					bool isEqual = collectionPersister.KeyType.IsEqual(id, ce.LoadedKey, collectionPersister.Factory);
 
 					if (isEqual)
@@ -182,6 +237,7 @@ public object[] GetCollectionBatch(ICollectionPersister collectionPersister, obj
 					}
 				}
 			}
+
 			return keys; //we ran out of keys to try
 		}
 
@@ -194,17 +250,17 @@ public object[] GetCollectionBatch(ICollectionPersister collectionPersister, obj
 		/// <param name="id">The identifier of the entity currently demanding load.</param>
 		/// <param name="batchSize">The maximum number of keys to return</param>
 		/// <returns>an array of identifiers, of length batchSize (possibly padded with nulls)</returns>
-		public object[] GetEntityBatch(IEntityPersister persister,object id,int batchSize)
+		public object[] GetEntityBatch(IEntityPersister persister, object id, int batchSize)
 		{
 			object[] ids = new object[batchSize];
 			ids[0] = id; //first element of array is reserved for the actual instance we are loading!
 			int i = 1;
 			int end = -1;
 			bool checkForEnd = false;
 
-			foreach (EntityKey key in batchLoadableEntityKeys.Keys)
+			if (batchLoadableEntityKeys.TryGetValue(persister.EntityName, out var set))
 			{
-				if (key.EntityName.Equals(persister.EntityName))
+				foreach (var key in set)
 				{
 					//TODO: this needn't exclude subclasses...
 					if (checkForEnd && i == end)

src/NHibernate/Engine/CollectionEntry.cs

@@ -298,12 +298,32 @@ public void PreFlush(IPersistentCollection collection)
 		/// has been initialized.
 		/// </summary>
 		/// <param name="collection">The initialized <see cref="AbstractPersistentCollection"/> that this Entry is for.</param>
+		//Since v5.1
+		[Obsolete("Please use PostInitialize(collection, persistenceContext) instead.")]
 		public void PostInitialize(IPersistentCollection collection)
 		{
 			snapshot = LoadedPersister.IsMutable ? collection.GetSnapshot(LoadedPersister) : null;
 			collection.SetSnapshot(loadedKey, role, snapshot);
 		}
 
+		/// <summary>
+		/// Updates the CollectionEntry to reflect that the <see cref="IPersistentCollection"/>
+		/// has been initialized.
+		/// </summary>
+		/// <param name="collection">The initialized <see cref="AbstractPersistentCollection"/> that this Entry is for.</param>
+		/// <param name="persistenceContext"></param>
+		public void PostInitialize(IPersistentCollection collection, IPersistenceContext persistenceContext)
+		{
+#pragma warning disable 618
+			//6.0 TODO: Inline PostInitialize here.
+			PostInitialize(collection);
+#pragma warning restore 618
+			if (LoadedPersister.GetBatchSize() > 1)
+			{
+				persistenceContext.BatchFetchQueue.RemoveBatchLoadableCollection(this);
+			}
+		}
+
 		/// <summary>
 		/// Updates the CollectionEntry to reflect that it is has been successfully flushed to the database.
 		/// </summary>

src/NHibernate/Engine/Loading/CollectionLoadContext.cs

@@ -234,7 +234,9 @@ private void EndLoadingCollection(LoadingCollectionEntry lce, ICollectionPersist
 			{
 				log.Debug("ending loading collection [{0}]", lce);
 			}
-			ISessionImplementor session = LoadContext.PersistenceContext.Session;
+
+			var persistenceContext = LoadContext.PersistenceContext;
+			var session = persistenceContext.Session;
 
 			bool statsEnabled = session.Factory.Statistics.IsStatisticsEnabled;
 			var stopWath = new Stopwatch();
@@ -247,17 +249,17 @@ private void EndLoadingCollection(LoadingCollectionEntry lce, ICollectionPersist
 
 			if (persister.CollectionType.HasHolder())
 			{
-				LoadContext.PersistenceContext.AddCollectionHolder(lce.Collection);
+				persistenceContext.AddCollectionHolder(lce.Collection);
 			}
 
-			CollectionEntry ce = LoadContext.PersistenceContext.GetCollectionEntry(lce.Collection);
+			CollectionEntry ce = persistenceContext.GetCollectionEntry(lce.Collection);
 			if (ce == null)
 			{
-				ce = LoadContext.PersistenceContext.AddInitializedCollection(persister, lce.Collection, lce.Key);
+				ce = persistenceContext.AddInitializedCollection(persister, lce.Collection, lce.Key);
 			}
 			else
 			{
-				ce.PostInitialize(lce.Collection);
+				ce.PostInitialize(lce.Collection, persistenceContext);
 			}
 
 			bool addToCache = hasNoQueuedAdds && persister.HasCache && 

src/NHibernate/Engine/StatefulPersistenceContext.cs

@@ -835,6 +835,10 @@ public void AddUninitializedCollection(ICollectionPersister persister, IPersiste
 		{
 			CollectionEntry ce = new CollectionEntry(collection, persister, id, flushing);
 			AddCollection(collection, ce, id);
+			if (persister.GetBatchSize() > 1)
+			{
+				batchFetchQueue.AddBatchLoadableCollection(collection, ce);
+			}
 		}
 
 		/// <summary> add a detached uninitialized collection</summary>
@@ -913,7 +917,7 @@ public CollectionEntry AddInitializedCollection(ICollectionPersister persister,
 																										object id)
 		{
 			CollectionEntry ce = new CollectionEntry(collection, persister, id, flushing);
-			ce.PostInitialize(collection);
+			ce.PostInitialize(collection, this);
 			AddCollection(collection, ce, id);
 			return ce;
 		}

src/NHibernate/Event/Default/DefaultInitializeCollectionEventListener.cs

@@ -115,7 +115,7 @@ private bool InitializeCollectionFromCache(object id, ICollectionPersister persi
 					CollectionCacheEntry cacheEntry = (CollectionCacheEntry)persister.CacheEntryStructure.Destructure(ce, factory);
 					cacheEntry.Assemble(collection, persister, persistenceContext.GetCollectionOwner(id, persister));
 
-					persistenceContext.GetCollectionEntry(collection).PostInitialize(collection);
+					persistenceContext.GetCollectionEntry(collection).PostInitialize(collection, persistenceContext);
 					return true;
 				}
 			}