Fixed lazy loading thread safety

src/EFCore/Infrastructure/Internal/LazyLoader.cs

@@ -4,6 +4,8 @@
 using System.Collections.Concurrent;
 using System.Diagnostics.CodeAnalysis;
 using System.Runtime.CompilerServices;
+using System.Runtime.InteropServices;
+using Microsoft.EntityFrameworkCore.ChangeTracking;
 using Microsoft.EntityFrameworkCore.Internal;
 
 namespace Microsoft.EntityFrameworkCore.Infrastructure.Internal;
@@ -20,7 +22,8 @@ public class LazyLoader : ILazyLoader, IInjectableService
     private bool _disposed;
     private bool _detached;
     private IDictionary<string, bool>? _loadedStates;
-    private readonly ConcurrentDictionary<(object Entity, string NavigationName), bool> _isLoading = new(NavEntryEqualityComparer.Instance);
+    private readonly Lock _isLoadingLock = new Lock();
+    private readonly Dictionary<(object Entity, string NavigationName), (TaskCompletionSource TaskCompletionSource, int ThreadId)> _isLoading = new(NavEntryEqualityComparer.Instance);
     private HashSet<string>? _nonLazyNavigations;
 
     /// <summary>
@@ -107,30 +110,55 @@ public virtual void Load(object entity, [CallerMemberName] string navigationName
         Check.NotEmpty(navigationName, nameof(navigationName));
 
         var navEntry = (entity, navigationName);
-        if (_isLoading.TryAdd(navEntry, true))
+
+        bool exists;
+        (TaskCompletionSource TaskCompletionSource, int ThreadId) isLoadingValue;
+        var currentThreadId = Environment.CurrentManagedThreadId;
+
+        lock (_isLoadingLock)
+        {
+            ref var refIsLoadingValue = ref CollectionsMarshal.GetValueRefOrAddDefault(_isLoading, navEntry, out exists);
+            if (!exists)
+            {
+                refIsLoadingValue = (new TaskCompletionSource(), currentThreadId);
+            }
+            isLoadingValue = refIsLoadingValue!;
+        }
+
+        if (exists)
+        {
+            if (isLoadingValue.ThreadId != currentThreadId)
+            {
+                isLoadingValue.TaskCompletionSource.Task.Wait();
+            }
+            return;
+        }
+
+        try
         {
-            try
+            // ShouldLoad is called after _isLoading.Add because it could attempt to load the property. See #13138.
+            if (ShouldLoad(entity, navigationName, out var entry))
             {
-                // ShouldLoad is called after _isLoading.Add because it could attempt to load the property. See #13138.
-                if (ShouldLoad(entity, navigationName, out var entry))
+                try
                 {
-                    try
-                    {
-                        entry.Load(
-                            _queryTrackingBehavior == QueryTrackingBehavior.NoTrackingWithIdentityResolution
-                                ? LoadOptions.ForceIdentityResolution
-                                : LoadOptions.None);
-                    }
-                    catch
-                    {
-                        entry.IsLoaded = false;
-                        throw;
-                    }
+                    entry.Load(
+                        _queryTrackingBehavior == QueryTrackingBehavior.NoTrackingWithIdentityResolution
+                            ? LoadOptions.ForceIdentityResolution
+                            : LoadOptions.None);
+                }
+                catch
+                {
+                    entry.IsLoaded = false;
+                    throw;
                 }
             }
-            finally
+        }
+        finally
+        {
+            isLoadingValue.TaskCompletionSource.TrySetResult();
+            lock (_isLoadingLock)
             {
-                _isLoading.TryRemove(navEntry, out _);
+                _isLoading.Remove(navEntry);
             }
         }
     }
@@ -150,31 +178,56 @@ public virtual async Task LoadAsync(
         Check.NotEmpty(navigationName, nameof(navigationName));
 
         var navEntry = (entity, navigationName);
-        if (_isLoading.TryAdd(navEntry, true))
+
+        bool exists;
+        (TaskCompletionSource TaskCompletionSource, int ThreadId) isLoadingValue;
+        var currentThreadId = Environment.CurrentManagedThreadId;
+
+        lock (_isLoadingLock)
+        {
+            ref var refIsLoadingValue = ref CollectionsMarshal.GetValueRefOrAddDefault(_isLoading, navEntry, out exists);
+            if (!exists)
+            {
+                refIsLoadingValue = (new TaskCompletionSource(), currentThreadId);
+            }
+            isLoadingValue = refIsLoadingValue!;
+        }
+
+        if (exists)
+        {
+            if(isLoadingValue.ThreadId != currentThreadId)
+            {
+                await isLoadingValue.TaskCompletionSource.Task.WaitAsync(cancellationToken).ConfigureAwait(false);
+            }
+            return;
+        }
+
+        try
         {
-            try
+            // ShouldLoad is called after _isLoading.Add because it could attempt to load the property. See #13138.
+            if (ShouldLoad(entity, navigationName, out var entry))
             {
-                // ShouldLoad is called after _isLoading.Add because it could attempt to load the property. See #13138.
-                if (ShouldLoad(entity, navigationName, out var entry))
+                try
                 {
-                    try
-                    {
-                        await entry.LoadAsync(
-                            _queryTrackingBehavior == QueryTrackingBehavior.NoTrackingWithIdentityResolution
-                                ? LoadOptions.ForceIdentityResolution
-                                : LoadOptions.None,
-                            cancellationToken).ConfigureAwait(false);
-                    }
-                    catch
-                    {
-                        entry.IsLoaded = false;
-                        throw;
-                    }
+                    await entry.LoadAsync(
+                               _queryTrackingBehavior == QueryTrackingBehavior.NoTrackingWithIdentityResolution
+                                   ? LoadOptions.ForceIdentityResolution
+                                   : LoadOptions.None,
+                               cancellationToken).ConfigureAwait(false);
+                }
+                catch
+                {
+                    entry.IsLoaded = false;
+                    throw;
                 }
             }
-            finally
+        }
+        finally
+        {
+            isLoadingValue.TaskCompletionSource.TrySetResult();
+            lock (_isLoadingLock)
             {
-                _isLoading.TryRemove(navEntry, out _);
+                _isLoading.Remove(navEntry);
             }
         }
     }

test/EFCore.Specification.Tests/LazyLoadProxyTestBase.cs

@@ -1,6 +1,7 @@
 // Licensed to the .NET Foundation under one or more agreements.
 // The .NET Foundation licenses this file to you under the MIT license.
 
+using System.Collections.Concurrent;
 using System.ComponentModel.DataAnnotations.Schema;
 using System.Text.Json;
 using System.Text.Json.Serialization;
@@ -67,6 +68,70 @@ public virtual void Can_use_proxies_from_multiple_threads_when_navigations_alrea
         Task.WaitAll(tests.Select(Task.Run).ToArray());
     }
 
+    [ConditionalTheory] // Issue #35528
+    [InlineData(false)]
+    [InlineData(true)]
+    public virtual void Lazy_loading_is_thread_safe(bool noTracking)
+    {
+        using var context = CreateContext(lazyLoadingEnabled: true);
+
+        //Creating another context to avoid caches
+        using var context2 = CreateContext(lazyLoadingEnabled: true);
+
+        IQueryable<Parent> query = context.Set<Parent>();
+        IQueryable<Parent> query2 = context2.Set<Parent>();
+
+        if (noTracking)
+        {
+            query = query.AsNoTracking();
+            query2 = query2.AsNoTracking();
+        }
+
+        var parent = query.Single();
+
+        var children = parent.Children!.Select(x => x.Id).OrderBy(x => x).ToList();
+        var singlePkToPk = parent.SinglePkToPk!.Id;
+        var single = parent.Single!.Id;
+        var childrenAk = parent.ChildrenAk!.Select(x => x.Id).OrderBy(x => x).ToList();
+        var singleAk = parent.SingleAk!.Id;
+        var childrenShadowFk = parent.ChildrenShadowFk!.Select(x => x.Id).OrderBy(x => x).ToList();
+        var singleShadowFk = parent.SingleShadowFk!.Id;
+        var childrenCompositeKey = parent.ChildrenCompositeKey!.Select(x => x.Id).OrderBy(x => x).ToList();
+        var singleCompositeKey = parent.SingleCompositeKey!.Id;
+        var withRecursiveProperty = parent.WithRecursiveProperty!.Id;
+        var manyChildren = parent.ManyChildren!.Select(x => x.Id).OrderBy(x => x).ToList();
+
+        var parent2 = query2.Single();
+
+        var parallelOptions = new ParallelOptions
+        {
+            MaxDegreeOfParallelism = Environment.ProcessorCount * 500
+        };
+
+        try
+        {
+            Parallel.For(0, 50000, parallelOptions, i =>
+            {
+                Assert.Equal(children, parent2.Children!.Select(x => x.Id).OrderBy(x => x).ToList());
+                Assert.Equal(singlePkToPk, parent2.SinglePkToPk!.Id);
+                Assert.Equal(single, parent2.Single!.Id);
+                Assert.Equal(childrenAk, parent2.ChildrenAk!.Select(x => x.Id).OrderBy(x => x).ToList());
+                Assert.Equal(singleAk, parent2.SingleAk!.Id);
+                Assert.Equal(childrenShadowFk, parent2.ChildrenShadowFk!.Select(x => x.Id).OrderBy(x => x).ToList());
+                Assert.Equal(singleShadowFk, parent2.SingleShadowFk!.Id);
+                Assert.Equal(childrenCompositeKey, parent2.ChildrenCompositeKey!.Select(x => x.Id).OrderBy(x => x).ToList());
+                Assert.Equal(singleCompositeKey, parent2.SingleCompositeKey!.Id);
+                Assert.Equal(withRecursiveProperty, parent2.WithRecursiveProperty!.Id);
+                Assert.Equal(manyChildren, parent2.ManyChildren!.Select(x => x.Id).OrderBy(x => x).ToList());
+            });
+        }
+        catch (Exception ex)
+        {
+            Assert.Fail($"Lazy loading is not thread safe exception: {ex.Message}");
+            throw;
+        }
+    }
+
     [ConditionalFact]
     public virtual void Detected_principal_reference_navigation_changes_are_detected_and_marked_loaded()
     {