Improved collection batch fetching (HHH-1775)

Author: gliljas

src/NHibernate.Test/NHSpecificTest/SetFixture.cs

@@ -348,6 +348,8 @@ public object NotFoundObject
 			get { throw new NotImplementedException(); }
 		}
 
+		public int BatchSize { get; } = 0;
+
 		public ISessionFactoryImplementor Factory
 		{
 			get { return null; }

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,23 +41,42 @@ 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)
+					{
+						// should never happen
+						if (log.IsWarnEnabled())
+						{
+							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);
+					var isEqual = collectionPersister.KeyType.IsEqual(
+						id,
+						ce.LoadedKey,
+						collectionPersister.Factory
+					);
 
 					if (isEqual)
 					{
@@ -79,6 +99,7 @@ public async Task<object[]> GetCollectionBatchAsync(ICollectionPersister collect
 					}
 				}
 			}
+			
 			return keys; //we ran out of keys to try
 		}
 
@@ -101,9 +122,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))//TODO: this needn't exclude subclasses...
 			{
-				if (key.EntityName.Equals(persister.EntityName))
+				foreach (var key in set)
 				{
 					//TODO: this needn't exclude subclasses...
 					if (checkForEnd && i == end)

src/NHibernate/Collection/AbstractPersistentCollection.cs

@@ -233,7 +233,7 @@ public virtual bool RowUpdatePossible
 		}
 
 		/// <summary></summary>
-		protected virtual ISessionImplementor Session
+		public virtual ISessionImplementor Session
 		{
 			get { return session; }
 		}

src/NHibernate/Engine/BatchFetchQueue.cs

@@ -5,31 +5,32 @@
 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>
 		/// keyed by the <see cref="EntityKey" /> against which the descriptor is
 		/// registered.
 		/// </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,23 +187,42 @@ 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)
+					{
+						// should never happen
+						if (log.IsWarnEnabled())
+						{
+							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);
+					var isEqual = collectionPersister.KeyType.IsEqual(
+						id,
+						ce.LoadedKey,
+						collectionPersister.Factory
+					);
 
 					if (isEqual)
 					{
@@ -182,6 +245,7 @@ public object[] GetCollectionBatch(ICollectionPersister collectionPersister, obj
 					}
 				}
 			}
+			
 			return keys; //we ran out of keys to try
 		}
 
@@ -202,9 +266,9 @@ public object[] GetEntityBatch(IEntityPersister persister,object id,int batchSiz
 			int end = -1;
 			bool checkForEnd = false;
 
-			foreach (EntityKey key in batchLoadableEntityKeys.Keys)
+			if (batchLoadableEntityKeys.TryGetValue(persister.EntityName,out var set))//TODO: this needn't exclude subclasses...
 			{
-				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

@@ -302,6 +302,10 @@ public void PostInitialize(IPersistentCollection collection)
 		{
 			snapshot = LoadedPersister.IsMutable ? collection.GetSnapshot(LoadedPersister) : null;
 			collection.SetSnapshot(loadedKey, role, snapshot);
+			if (LoadedPersister.BatchSize > 1)
+			{
+				((AbstractPersistentCollection) collection).Session.PersistenceContext.BatchFetchQueue.RemoveBatchLoadableCollection(this);
+			}
 		}
 
 		/// <summary>

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.BatchSize > 0)
+			{
+				batchFetchQueue.AddBatchLoadableCollection(collection, ce);
+			}
 		}
 
 		/// <summary> add a detached uninitialized collection</summary>

src/NHibernate/Event/Default/EvictVisitor.cs

@@ -53,6 +53,10 @@ private void EvictCollection(IPersistentCollection collection)
 			Session.PersistenceContext.CollectionEntries.Remove(collection);
 			if (log.IsDebugEnabled())
 				log.Debug("evicting collection: {0}", MessageHelper.CollectionInfoString(ce.LoadedPersister, collection, ce.LoadedKey, Session));
+			if (ce.LoadedPersister?.BatchSize > 1)
+			{
+				Session.PersistenceContext.BatchFetchQueue.RemoveBatchLoadableCollection(ce);
+			}
 			if (ce.LoadedPersister != null && ce.LoadedKey != null)
 			{
 				//TODO: is this 100% correct?

src/NHibernate/Persister/Collection/AbstractCollectionPersister.cs

@@ -2043,6 +2043,11 @@ public string[] RootTableKeyColumnNames
 			get { return new string[] {IdentifierColumnName}; }
 		}
 
+		public int BatchSize
+		{
+			get { return batchSize; }
+		}
+
 		public SqlString GetSelectByUniqueKeyString(string propertyName)
 		{
 			return

src/NHibernate/Persister/Collection/BasicCollectionPersister.cs

@@ -281,7 +281,7 @@ private string ManyToManySelectFragment(IJoinable rhs, string rhsAlias, string l
 		/// </summary>
 		protected override ICollectionInitializer CreateCollectionInitializer(IDictionary<string, IFilter> enabledFilters)
 		{
-			return BatchingCollectionInitializer.CreateBatchingCollectionInitializer(this, batchSize, Factory, enabledFilters);
+			return BatchingCollectionInitializer.CreateBatchingCollectionInitializer(this, BatchSize, Factory, enabledFilters);
 		}
 
 		public override SqlString FromJoinFragment(string alias, bool innerJoin, bool includeSubclasses)

src/NHibernate/Persister/Collection/ICollectionPersister.cs

@@ -282,5 +282,7 @@ public partial interface ICollectionPersister
 		/// A place-holder to inform that the data-reader was empty.
 		/// </summary>
 		object NotFoundObject { get; }
+
+		int BatchSize { get; }
 	}
 }

src/NHibernate/Persister/Collection/OneToManyPersister.cs

@@ -343,7 +343,7 @@ protected override SelectFragment GenerateSelectFragment(string alias, string co
 		/// </summary>
 		protected override ICollectionInitializer CreateCollectionInitializer(IDictionary<string, IFilter> enabledFilters)
 		{
-			return BatchingCollectionInitializer.CreateBatchingOneToManyInitializer(this, batchSize, Factory, enabledFilters);
+			return BatchingCollectionInitializer.CreateBatchingOneToManyInitializer(this, BatchSize, Factory, enabledFilters);
 		}
 
 		public override SqlString FromJoinFragment(string alias, bool innerJoin, bool includeSubclasses)

src/NHibernate/Type/ManyToOneType.cs

@@ -101,7 +101,7 @@ private void ScheduleBatchLoadIfNeeded(object id, ISessionImplementor session)
 			{
 				IEntityPersister persister = session.Factory.GetEntityPersister(GetAssociatedEntityName());
 				EntityKey entityKey = session.GenerateEntityKey(id, persister);
-				if (!session.PersistenceContext.ContainsEntity(entityKey))
+				if (entityKey.IsBatchLoadable && !session.PersistenceContext.ContainsEntity(entityKey))
 				{
 					session.PersistenceContext.BatchFetchQueue.AddBatchLoadableEntityKey(entityKey);
 				}