Optimize batchable cache calls for cached queries (nhibernate#1955)

Author: maca88

src/NHibernate.Test/Async/CacheTest/BatchableCacheFixture.cs

@@ -13,6 +13,7 @@
 using System.Collections.Generic;
 using System.Linq;
 using NHibernate.Cfg;
+using NHibernate.Collection;
 using NHibernate.Engine;
 using NHibernate.Persister.Collection;
 using NHibernate.Type;
@@ -208,44 +209,86 @@ public async Task<IList[]> GetManyAsync(
 			var persistenceContext = session.PersistenceContext;
 			var defaultReadOnlyOrig = persistenceContext.DefaultReadOnly;
 			var results = new IList[keys.Length];
-			for (var i = 0; i < keys.Length; i++)
+			var finalReturnTypes = new ICacheAssembler[keys.Length][];
+			try
 			{
-				var cacheable = cacheables[i];
-				if (cacheable == null)
-					continue;
+				session.PersistenceContext.BatchFetchQueue.InitializeQueryCacheQueue();
 
-				var key = keys[i];
-				if (checkedSpacesIndexes.Contains(i) && !upToDates[upToDatesIndex++])
+				for (var i = 0; i < keys.Length; i++)
 				{
-					Log.Debug("cached query results were not up to date for: {0}", key);
-					continue;
+					var cacheable = cacheables[i];
+					if (cacheable == null)
+						continue;
+
+					var key = keys[i];
+					if (checkedSpacesIndexes.Contains(i) && !upToDates[upToDatesIndex++])
+					{
+						Log.Debug("cached query results were not up to date for: {0}", key);
+						continue;
+					}
+
+					var queryParams = queryParameters[i];
+					if (queryParams.IsReadOnlyInitialized)
+						persistenceContext.DefaultReadOnly = queryParams.ReadOnly;
+					else
+						queryParams.ReadOnly = persistenceContext.DefaultReadOnly;
+
+					Log.Debug("returning cached query results for: {0}", key);
+
+					finalReturnTypes[i] = GetReturnTypes(key, returnTypes[i], cacheable);
+					await (PerformBeforeAssembleAsync(finalReturnTypes[i], session, cacheable, cancellationToken)).ConfigureAwait(false);
 				}
 
-				var queryParams = queryParameters[i];
-				if (queryParams.IsReadOnlyInitialized)
-					persistenceContext.DefaultReadOnly = queryParams.ReadOnly;
-				else
-					queryParams.ReadOnly = persistenceContext.DefaultReadOnly;
+				for (var i = 0; i < keys.Length; i++)
+				{
+					if (finalReturnTypes[i] == null)
+					{
+						continue;
+					}
 
-				// Adjust the session cache mode, as GetResultFromCacheable assemble types which may cause
-				// entity loads, which may interact with the cache.
-				using (session.SwitchCacheMode(queryParams.CacheMode))
+					var queryParams = queryParameters[i];
+					// Adjust the session cache mode, as PerformAssemble assemble types which may cause
+					// entity loads, which may interact with the cache.
+					using (session.SwitchCacheMode(queryParams.CacheMode))
+					{
+						try
+						{
+							results[i] = await (PerformAssembleAsync(keys[i], finalReturnTypes[i], queryParams.NaturalKeyLookup, session, cacheables[i], cancellationToken)).ConfigureAwait(false);
+						}
+						finally
+						{
+							persistenceContext.DefaultReadOnly = defaultReadOnlyOrig;
+						}
+					}
+				}
+
+				for (var i = 0; i < keys.Length; i++)
 				{
-					try
+					if (finalReturnTypes[i] == null)
 					{
-						results[i] = await (GetResultFromCacheableAsync(
-							key,
-							returnTypes[i],
-							queryParams.NaturalKeyLookup,
-							session,
-							cacheable, cancellationToken)).ConfigureAwait(false);
+						continue;
 					}
-					finally
+
+					var queryParams = queryParameters[i];
+					// Adjust the session cache mode, as InitializeCollections will initialize collections,
+					// which may interact with the cache.
+					using (session.SwitchCacheMode(queryParams.CacheMode))
 					{
-						persistenceContext.DefaultReadOnly = defaultReadOnlyOrig;
+						try
+						{
+							await (InitializeCollectionsAsync(finalReturnTypes[i], session, results[i], cacheables[i], cancellationToken)).ConfigureAwait(false);
+						}
+						finally
+						{
+							persistenceContext.DefaultReadOnly = defaultReadOnlyOrig;
+						}
 					}
 				}
 			}
+			finally
+			{
+				session.PersistenceContext.BatchFetchQueue.TerminateQueryCacheQueue();
+			}
 
 			return results;
 		}
@@ -275,20 +318,40 @@ private static async Task<List<object>> GetCacheableResultAsync(
 			return cacheable;
 		}
 
-		private async Task<IList> GetResultFromCacheableAsync(
-			QueryKey key,
+		private static async Task PerformBeforeAssembleAsync(
 			ICacheAssembler[] returnTypes,
-			bool isNaturalKeyLookup,
 			ISessionImplementor session,
 			IList cacheable, CancellationToken cancellationToken)
 		{
 			cancellationToken.ThrowIfCancellationRequested();
-			Log.Debug("returning cached query results for: {0}", key);
-			if (key.ResultTransformer?.AutoDiscoverTypes == true && cacheable.Count > 0)
+			if (returnTypes.Length == 1)
 			{
-				returnTypes = GuessTypes(cacheable);
+				var returnType = returnTypes[0];
+
+				// Skip first element, it is the timestamp
+				for (var i = 1; i < cacheable.Count; i++)
+				{
+					await (returnType.BeforeAssembleAsync(cacheable[i], session, cancellationToken)).ConfigureAwait(false);
+				}
 			}
+			else
+			{
+				// Skip first element, it is the timestamp
+				for (var i = 1; i < cacheable.Count; i++)
+				{
+					await (TypeHelper.BeforeAssembleAsync((object[]) cacheable[i], returnTypes, session, cancellationToken)).ConfigureAwait(false);
+				}
+			}
+		}
 
+		private async Task<IList> PerformAssembleAsync(
+			QueryKey key,
+			ICacheAssembler[] returnTypes,
+			bool isNaturalKeyLookup,
+			ISessionImplementor session,
+			IList cacheable, CancellationToken cancellationToken)
+		{
+			cancellationToken.ThrowIfCancellationRequested();
 			try
 			{
 				var result = new List<object>(cacheable.Count - 1);
@@ -297,53 +360,27 @@ private async Task<IList> GetResultFromCacheableAsync(
 					var returnType = returnTypes[0];
 
 					// Skip first element, it is the timestamp
-					for (var i = 1; i < cacheable.Count; i++)
-					{
-						await (returnType.BeforeAssembleAsync(cacheable[i], session, cancellationToken)).ConfigureAwait(false);
-					}
-
 					for (var i = 1; i < cacheable.Count; i++)
 					{
 						result.Add(await (returnType.AssembleAsync(cacheable[i], session, null, cancellationToken)).ConfigureAwait(false));
 					}
 				}
 				else
 				{
-					var collectionIndexes = new Dictionary<int, ICollectionPersister>();
 					var nonCollectionTypeIndexes = new List<int>();
 					for (var i = 0; i < returnTypes.Length; i++)
 					{
-						if (returnTypes[i] is CollectionType collectionType)
-						{
-							collectionIndexes.Add(i, session.Factory.GetCollectionPersister(collectionType.Role));
-						}
-						else
+						if (!(returnTypes[i] is CollectionType))
 						{
 							nonCollectionTypeIndexes.Add(i);
 						}
 					}
 
 					// Skip first element, it is the timestamp
-					for (var i = 1; i < cacheable.Count; i++)
-					{
-						await (TypeHelper.BeforeAssembleAsync((object[]) cacheable[i], returnTypes, session, cancellationToken)).ConfigureAwait(false);
-					}
-
 					for (var i = 1; i < cacheable.Count; i++)
 					{
 						result.Add(await (TypeHelper.AssembleAsync((object[]) cacheable[i], returnTypes, nonCollectionTypeIndexes, session, cancellationToken)).ConfigureAwait(false));
 					}
-
-					// Initialization of the fetched collection must be done at the end in order to be able to batch fetch them
-					// from the cache or database. The collections were already created in the previous for statement so we only
-					// have to initialize them.
-					if (collectionIndexes.Count > 0)
-					{
-						for (var i = 1; i < cacheable.Count; i++)
-						{
-							await (TypeHelper.InitializeCollectionsAsync((object[]) cacheable[i], (object[]) result[i - 1], collectionIndexes, session, cancellationToken)).ConfigureAwait(false);
-						}
-					}
 				}
 
 				return result;
@@ -367,6 +404,66 @@ private async Task<IList> GetResultFromCacheableAsync(
 			}
 		}
 
+		private static async Task InitializeCollectionsAsync(
+			ICacheAssembler[] returnTypes,
+			ISessionImplementor session,
+			IList assembleResult,
+			IList cacheResult, CancellationToken cancellationToken)
+		{
+			cancellationToken.ThrowIfCancellationRequested();
+			var collectionIndexes = new Dictionary<int, ICollectionPersister>();
+			for (var i = 0; i < returnTypes.Length; i++)
+			{
+				if (returnTypes[i] is CollectionType collectionType)
+				{
+					collectionIndexes.Add(i, session.Factory.GetCollectionPersister(collectionType.Role));
+				}
+			}
+
+			if (collectionIndexes.Count == 0)
+			{
+				return;
+			}
+
+			// Skip first element, it is the timestamp
+			for (var i = 1; i < cacheResult.Count; i++)
+			{
+				// Initialization of the fetched collection must be done at the end in order to be able to batch fetch them
+				// from the cache or database. The collections were already created when their owners were assembled so we only
+				// have to initialize them.
+				await (TypeHelper.InitializeCollectionsAsync(
+					(object[]) cacheResult[i],
+					(object[]) assembleResult[i - 1],
+					collectionIndexes,
+					session, cancellationToken)).ConfigureAwait(false);
+			}
+		}
+
+		private async Task<IList> GetResultFromCacheableAsync(
+			QueryKey key,
+			ICacheAssembler[] returnTypes,
+			bool isNaturalKeyLookup,
+			ISessionImplementor session,
+			IList cacheable, CancellationToken cancellationToken)
+		{
+			cancellationToken.ThrowIfCancellationRequested();
+			Log.Debug("returning cached query results for: {0}", key);
+			returnTypes = GetReturnTypes(key, returnTypes, cacheable);
+			try
+			{
+				session.PersistenceContext.BatchFetchQueue.InitializeQueryCacheQueue();
+
+				await (PerformBeforeAssembleAsync(returnTypes, session, cacheable, cancellationToken)).ConfigureAwait(false);
+				var result = await (PerformAssembleAsync(key, returnTypes, isNaturalKeyLookup, session, cacheable, cancellationToken)).ConfigureAwait(false);
+				await (InitializeCollectionsAsync(returnTypes, session, result, cacheable, cancellationToken)).ConfigureAwait(false);
+				return result;
+			}
+			finally
+			{
+				session.PersistenceContext.BatchFetchQueue.TerminateQueryCacheQueue();
+			}
+		}
+
 		protected virtual Task<bool> IsUpToDateAsync(ISet<string> spaces, long timestamp, CancellationToken cancellationToken)
 		{
 			if (cancellationToken.IsCancellationRequested)

src/NHibernate.Test/CacheTest/BatchableCacheFixture.cs

@@ -15,6 +15,7 @@
 using NHibernate.Persister.Entity;
 using NHibernate.Util;
 using System.Collections.Generic;
+using System.Linq;
 using Iesi.Collections.Generic;
 
 namespace NHibernate.Engine
@@ -395,6 +396,13 @@ private async Task<bool[]> AreCachedAsync(List<KeyValuePair<EntityKey, int>> ent
 			{
 				return result;
 			}
+
+			// Do not check the cache when disassembling entities from the cached query that were already checked
+			if (QueryCacheQueue != null && entityKeys.All(o => QueryCacheQueue.WasEntityKeyChecked(persister, o.Key)))
+			{
+				return result;
+			}
+
 			var cacheKeys = new object[keyIndexes.Length];
 			var i = 0;
 			foreach (var index in keyIndexes)
@@ -434,6 +442,13 @@ private async Task<bool[]> AreCachedAsync(List<KeyValuePair<KeyValuePair<Collect
 			{
 				return result;
 			}
+
+			// Do not check the cache when disassembling collections from the cached query that were already checked
+			if (QueryCacheQueue != null && collectionKeys.All(o => QueryCacheQueue.WasCollectionEntryChecked(persister, o.Key.Key)))
+			{
+				return result;
+			}
+
 			var cacheKeys = new object[keyIndexes.Length];
 			var i = 0;
 			foreach (var index in keyIndexes)