Resolve collection owner correctly in case of property-ref

Root cause of NH-3480

Author: fredericDelaporte

src/NHibernate.Test/Async/NHSpecificTest/NH3480/Fixture.cs

@@ -26,7 +26,15 @@ protected override void OnSetUp()
 			{
 				using (ITransaction transaction = session.BeginTransaction())
 				{
-					var parent1 = new Entity { Id = new Entity.Key() { Id = Guid.NewGuid() }, Name = "Bob", OtherId = 20 };
+					var parent1 = new Entity
+					{
+						Id = new Entity.Key { Id = Guid.NewGuid() },
+						Name = "Bob",
+						OtherId = 20,
+						YetAnotherOtherId = 21
+					};
+					parent1.Elements.Add(1);
+					parent1.Elements.Add(2);
 					session.Save(parent1);
 
 					var child1 = new Child { Name = "Bob1", Parent = parent1 };
@@ -72,5 +80,21 @@ public async Task TestAsync()
 				}
 			}
 		}
+
+		[Test]
+		public async Task TestOwnerAsync()
+		{
+			using (var session = OpenSession())
+			using (var t = session.BeginTransaction())
+			{
+				var entity = await (session.Query<Entity>().SingleAsync(e => e.Name == "Bob"));
+
+				// The Elements collection is mapped with a custom type which assert the owner
+				// is not null.
+				await (NHibernateUtil.InitializeAsync(entity.Elements));
+				Assert.That(entity.Elements, Has.Count.GreaterThan(0));
+				await (t.CommitAsync());
+			}
+		}
 	}
 }

src/NHibernate.Test/NHSpecificTest/NH3480/Entity.cs

@@ -1,5 +1,10 @@
 using System;
 using System.Collections.Generic;
+using System.Data.Common;
+using NHibernate.Engine;
+using NHibernate.SqlTypes;
+using NHibernate.UserTypes;
+using NUnit.Framework;
 
 namespace NHibernate.Test.NHSpecificTest.NH3480
 {
@@ -13,7 +18,9 @@ public Entity()
 		public virtual Key Id { get; set; }
 		public virtual string Name { get; set; }
 		public virtual int OtherId { get; set; }
+		public virtual int YetAnotherOtherId { get; set; }
 		public virtual ISet<Child> Children { get; set; }
+		public virtual ISet<int> Elements { get; set; } = new HashSet<int>();
 
 		public override bool Equals(object obj)
 		{
@@ -58,4 +65,77 @@ class Child
 		public virtual string Name { get; set; }
 		public virtual Entity Parent { get; set; }
 	}
-}
+
+	// This user type is the Elements collection element type
+	public class SimpleCustomType : IUserType
+	{
+		private static readonly SqlType[] ReturnSqlTypes = { SqlTypeFactory.Int32 };
+
+		#region IUserType Members
+
+		public new bool Equals(object x, object y)
+		{
+			if (ReferenceEquals(x, y))
+			{
+				return true;
+			}
+			if (x == null || y == null)
+			{
+				return false;
+			}
+
+			return x.Equals(y);
+		}
+
+		public int GetHashCode(object x)
+		{
+			return x?.GetHashCode() ?? 0;
+		}
+
+		public SqlType[] SqlTypes => ReturnSqlTypes;
+
+		public object DeepCopy(object value)
+		{
+			return value;
+		}
+
+		public void NullSafeSet(DbCommand cmd, object value, int index, ISessionImplementor session)
+		{
+			cmd.Parameters[index].Value = value ?? DBNull.Value;
+		}
+
+		public System.Type ReturnedType => typeof(int);
+
+		public object NullSafeGet(DbDataReader rs, string[] names, ISessionImplementor session, object owner)
+		{
+			// Detect if the root cause of the bug is still there: if yes, the owner will be null.
+			Assert.That(owner, Is.Not.Null);
+
+			var index = rs.GetOrdinal(names[0]);
+			if (rs.IsDBNull(index))
+			{
+				return null;
+			}
+			return rs.GetInt32(index);
+		}
+
+		public bool IsMutable => false;
+
+		public object Replace(object original, object target, object owner)
+		{
+			return original;
+		}
+
+		public object Assemble(object cached, object owner)
+		{
+			return cached;
+		}
+
+		public object Disassemble(object value)
+		{
+			return value;
+		}
+
+		#endregion
+	}
+}

src/NHibernate.Test/NHSpecificTest/NH3480/Fixture.cs

@@ -15,7 +15,15 @@ protected override void OnSetUp()
 			{
 				using (ITransaction transaction = session.BeginTransaction())
 				{
-					var parent1 = new Entity { Id = new Entity.Key() { Id = Guid.NewGuid() }, Name = "Bob", OtherId = 20 };
+					var parent1 = new Entity
+					{
+						Id = new Entity.Key { Id = Guid.NewGuid() },
+						Name = "Bob",
+						OtherId = 20,
+						YetAnotherOtherId = 21
+					};
+					parent1.Elements.Add(1);
+					parent1.Elements.Add(2);
 					session.Save(parent1);
 
 					var child1 = new Child { Name = "Bob1", Parent = parent1 };
@@ -61,5 +69,21 @@ public void Test()
 				}
 			}
 		}
+
+		[Test]
+		public void TestOwner()
+		{
+			using (var session = OpenSession())
+			using (var t = session.BeginTransaction())
+			{
+				var entity = session.Query<Entity>().Single(e => e.Name == "Bob");
+
+				// The Elements collection is mapped with a custom type which assert the owner
+				// is not null.
+				NHibernateUtil.Initialize(entity.Elements);
+				Assert.That(entity.Elements, Has.Count.GreaterThan(0));
+				t.Commit();
+			}
+		}
 	}
 }

src/NHibernate.Test/NHSpecificTest/NH3480/Mappings.hbm.xml

@@ -1,16 +1,21 @@
 <?xml version="1.0" encoding="utf-8" ?>
 <hibernate-mapping xmlns="urn:nhibernate-mapping-2.2" assembly="NHibernate.Test" namespace="NHibernate.Test.NHSpecificTest.NH3480">
 
-	<class name="Entity">
+    <class name="Entity">
         <composite-id name="Id">
             <key-property name="Id" />
         </composite-id>
-		<property name="Name" not-null="true" />
+        <property name="Name" not-null="true" />
         <property name="OtherId" unique="true" not-null="true" />
+        <property name="YetAnotherOtherId" unique="true" not-null="true" />
         <set name="Children" cascade="all-delete-orphan" inverse="true">
             <key property-ref="OtherId" column="Parent" />
             <one-to-many class="Child" />
         </set>
+        <set name="Elements">
+            <key property-ref="YetAnotherOtherId" column="Parent" />
+            <element type="NHibernate.Test.NHSpecificTest.NH3480.SimpleCustomType, NHibernate.Test"/>
+        </set>
     </class>
 
     <class name="Child">
@@ -19,4 +24,4 @@
         <many-to-one name="Parent" property-ref="OtherId" not-null="true" />
     </class>
 
-</hibernate-mapping>
+</hibernate-mapping>

src/NHibernate/Async/Engine/BatchFetchQueue.cs

@@ -47,18 +47,18 @@ public Task<object[]> GetCollectionBatchAsync(ICollectionPersister collectionPer
 		/// Get a batch of uninitialized collection keys for a given role
 		/// </summary>
 		/// <param name="collectionPersister">The persister for the collection role.</param>
-		/// <param name="id">A key that must be included in the batch fetch</param>
+		/// <param name="key">A key that must be included in the batch fetch</param>
 		/// <param name="batchSize">the maximum number of keys to return</param>
 		/// <param name="checkCache">Whether to check the cache for uninitialized collection keys.</param>
 		/// <param name="collectionEntries">An array that will be filled with collection entries if set.</param>
 		/// <param name="cancellationToken">A cancellation token that can be used to cancel the work</param>
 		/// <returns>An array of collection keys, of length <paramref name="batchSize"/> (padded with nulls)</returns>
-		internal async Task<object[]> GetCollectionBatchAsync(ICollectionPersister collectionPersister, object id, int batchSize, bool checkCache,
+		internal async Task<object[]> GetCollectionBatchAsync(ICollectionPersister collectionPersister, object key, int batchSize, bool checkCache,
 		                                     CollectionEntry[] collectionEntries, CancellationToken cancellationToken)
 		{
 			cancellationToken.ThrowIfCancellationRequested();
 			var keys = new object[batchSize];
-			keys[0] = id; // The first element of array is reserved for the actual instance we are loading
+			keys[0] = key; // The first element of array is reserved for the actual instance we are loading
 			var i = 1; // The current index of keys array
 			int? keyIndex = null; // The index of the demanding key in the linked hash set
 			var checkForEnd = false; // Stores whether we found the demanded collection and reached the batchSize
@@ -159,7 +159,7 @@ async Task<bool> ProcessKeyAsync(KeyValuePair<CollectionEntry, IPersistentCollec
 				{
 					return true;
 				}
-				if (collectionPersister.KeyType.IsEqual(id, ce.LoadedKey, collectionPersister.Factory))
+				if (collectionPersister.KeyType.IsEqual(key, ce.LoadedKey, collectionPersister.Factory))
 				{
 					if (collectionEntries != null)
 					{

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

@@ -73,10 +73,11 @@ public virtual async Task OnInitializeCollectionAsync(InitializeCollectionEvent
 		}
 
 		/// <summary> Try to initialize a collection from the cache</summary>
-		private async Task<bool> InitializeCollectionFromCacheAsync(object id, ICollectionPersister persister, IPersistentCollection collection, ISessionImplementor source, CancellationToken cancellationToken)
+		private async Task<bool> InitializeCollectionFromCacheAsync(
+			object collectionKey, ICollectionPersister persister, IPersistentCollection collection,
+			ISessionImplementor source, CancellationToken cancellationToken)
 		{
 			cancellationToken.ThrowIfCancellationRequested();
-
 			if (!(source.EnabledFilters.Count == 0) && persister.IsAffectedByEnabledFilters(source))
 			{
 				log.Debug("disregarding cached version (if any) of collection due to enabled filters ");
@@ -96,7 +97,7 @@ private async Task<bool> InitializeCollectionFromCacheAsync(object id, ICollecti
 				var collectionEntries = new CollectionEntry[batchSize];
 				// The first item in the array is the item that we want to load
 				var collectionBatch = await (source.PersistenceContext.BatchFetchQueue
-				                            .GetCollectionBatchAsync(persister, id, batchSize, false, collectionEntries, cancellationToken)).ConfigureAwait(false);
+				                            .GetCollectionBatchAsync(persister, collectionKey, batchSize, false, collectionEntries, cancellationToken)).ConfigureAwait(false);
 				// Ignore null values as the retrieved batch may contains them when there are not enough
 				// uninitialized collection in the queue
 				var keys = new List<CacheKey>(batchSize);
@@ -115,16 +116,17 @@ private async Task<bool> InitializeCollectionFromCacheAsync(object id, ICollecti
 					var coll = source.PersistenceContext.BatchFetchQueue.GetBatchLoadableCollection(persister, collectionEntries[i]);
 					await (AssembleAsync(keys[i], cachedObjects[i], persister, source, coll, collectionBatch[i], false, cancellationToken)).ConfigureAwait(false);
 				}
-				return await (AssembleAsync(keys[0], cachedObjects[0], persister, source, collection, id, true, cancellationToken)).ConfigureAwait(false);
+				return await (AssembleAsync(keys[0], cachedObjects[0], persister, source, collection, collectionKey, true, cancellationToken)).ConfigureAwait(false);
 			}
 
-			var cacheKey = source.GenerateCacheKey(id, persister.KeyType, persister.Role);
+			var cacheKey = source.GenerateCacheKey(collectionKey, persister.KeyType, persister.Role);
 			var cachedObject = await (persister.Cache.GetAsync(cacheKey, source.Timestamp, cancellationToken)).ConfigureAwait(false);
-			return await (AssembleAsync(cacheKey, cachedObject, persister, source, collection, id, true, cancellationToken)).ConfigureAwait(false);
+			return await (AssembleAsync(cacheKey, cachedObject, persister, source, collection, collectionKey, true, cancellationToken)).ConfigureAwait(false);
 		}
 
-		private async Task<bool> AssembleAsync(CacheKey ck, object ce, ICollectionPersister persister,  ISessionImplementor source,
-							  IPersistentCollection collection,  object id, bool alterStatistics, CancellationToken cancellationToken)
+		private async Task<bool> AssembleAsync(
+			CacheKey ck, object ce, ICollectionPersister persister, ISessionImplementor source,
+			IPersistentCollection collection, object collectionKey, bool alterStatistics, CancellationToken cancellationToken)
 		{
 			cancellationToken.ThrowIfCancellationRequested();
 			ISessionFactoryImplementor factory = source.Factory;
@@ -158,7 +160,7 @@ private async Task<bool> AssembleAsync(CacheKey ck, object ce, ICollectionPersis
 				IPersistenceContext persistenceContext = source.PersistenceContext;
 
 				CollectionCacheEntry cacheEntry = (CollectionCacheEntry) persister.CacheEntryStructure.Destructure(ce, factory);
-				await (cacheEntry.AssembleAsync(collection, persister, persistenceContext.GetCollectionOwner(id, persister), cancellationToken)).ConfigureAwait(false);
+				await (cacheEntry.AssembleAsync(collection, persister, persistenceContext.GetCollectionOwner(collectionKey, persister), cancellationToken)).ConfigureAwait(false);
 
 				persistenceContext.GetCollectionEntry(collection).PostInitialize(collection, persistenceContext);
 

src/NHibernate/Async/Loader/Loader.cs

@@ -646,9 +646,46 @@ private async Task<object[]> GetRowAsync(DbDataReader rs, ILoadable[] persisters
 						obj =
 							await (InstanceNotYetLoadedAsync(rs, i, persister, key, lockModes[i], optionalObjectKey,
 												 optionalObject, hydratedObjects, session, cancellationToken)).ConfigureAwait(false);
+
+						// IUniqueKeyLoadable.CacheByUniqueKeys caches all unique keys of the entity, regardless of
+						// associations loaded by the query. So if the entity is already loaded, it has forcibly already
+						// been cached too for all its unique keys, provided its persister implement it. With this new
+						// way of caching unique keys, it is no more needed to handle caching for alreadyLoaded path
+						// too.
+						var cacheByUniqueKeysTask = (persister as IUniqueKeyLoadable)?.CacheByUniqueKeysAsync(obj, session, cancellationToken);
+
+						// IUniqueKeyLoadable.CacheByUniqueKeys caches all unique keys of the entity, regardless of
+						// associations loaded by the query. So if the entity is already loaded, it has forcibly already
+						// been cached too for all its unique keys, provided its persister implement it. With this new
+						// way of caching unique keys, it is no more needed to handle caching for alreadyLoaded path
+						// too.
+						if (cacheByUniqueKeysTask != null)
+
+						// IUniqueKeyLoadable.CacheByUniqueKeys caches all unique keys of the entity, regardless of
+						// associations loaded by the query. So if the entity is already loaded, it has forcibly already
+						// been cached too for all its unique keys, provided its persister implement it. With this new
+						// way of caching unique keys, it is no more needed to handle caching for alreadyLoaded path
+						// too.
+						{
+
+							// IUniqueKeyLoadable.CacheByUniqueKeys caches all unique keys of the entity, regardless of
+							// associations loaded by the query. So if the entity is already loaded, it has forcibly already
+							// been cached too for all its unique keys, provided its persister implement it. With this new
+							// way of caching unique keys, it is no more needed to handle caching for alreadyLoaded path
+							// too.
+							await (cacheByUniqueKeysTask).ConfigureAwait(false);
+
+						// IUniqueKeyLoadable.CacheByUniqueKeys caches all unique keys of the entity, regardless of
+						// associations loaded by the query. So if the entity is already loaded, it has forcibly already
+						// been cached too for all its unique keys, provided its persister implement it. With this new
+						// way of caching unique keys, it is no more needed to handle caching for alreadyLoaded path
+						// too.
+						}
 					}
-					// #1226: Even if it is already loaded, if it can be loaded from an association with a property ref, make
-					// sure it is also cached by its unique key.
+					// 6.0 TODO: this call is nor more needed for up-to-date persisters, remove once CacheByUniqueKeys
+					// is merged in IUniqueKeyLoadable interface instead of being an extension method
+					// #1226 old fix: Even if it is already loaded, if it can be loaded from an association with a property ref,
+					// make sure it is also cached by its unique key.
 					CacheByUniqueKey(i, persister, obj, session, alreadyLoaded);
 				}
 

src/NHibernate/Async/Persister/Entity/AbstractEntityPersister.cs

@@ -259,6 +259,25 @@ public Task<object> LoadByUniqueKeyAsync(string propertyName, object uniqueKey,
 			}
 		}
 
+		public async Task CacheByUniqueKeysAsync(object entity, ISessionImplementor session, CancellationToken cancellationToken)
+		{
+			cancellationToken.ThrowIfCancellationRequested();
+			for (var i = 0; i < PropertySpan; i++)
+			{
+				if (!propertyUniqueness[i])
+					continue;
+
+				// The caching is done by semi-resolved values.
+				var propertyValue = session.PersistenceContext.GetEntry(entity).LoadedState[i];
+				if (propertyValue == null)
+					continue;
+				var type = PropertyTypes[i].GetSemiResolvedType(session.Factory);
+				propertyValue = await (type.SemiResolveAsync(propertyValue, session, entity, cancellationToken)).ConfigureAwait(false);
+				var euk = new EntityUniqueKey(EntityName, PropertyNames[i], propertyValue, type, session.Factory);
+				session.PersistenceContext.AddEntity(euk, entity);
+			}
+		}
+
 		protected Task<int> DehydrateAsync(object id, object[] fields, bool[] includeProperty, bool[][] includeColumns, int j, DbCommand st, ISessionImplementor session, CancellationToken cancellationToken)
 		{
 			if (cancellationToken.IsCancellationRequested)