Fixed Lazy loader performance issue (#35835)

henriquewr · web-flow · commit 63e55f1c3f38 · 2025-03-25T14:40:11.000-07:00
Use static AsyncLocal
diff --git a/src/EFCore/Infrastructure/Internal/LazyLoader.cs b/src/EFCore/Infrastructure/Internal/LazyLoader.cs
@@ -23,7 +23,8 @@ public class LazyLoader : ILazyLoader, IInjectableService
     private bool _detached;
     private IDictionary<string, bool>? _loadedStates;
     private readonly Lock _isLoadingLock = new Lock();
-    private readonly Dictionary<(object Entity, string NavigationName), (TaskCompletionSource TaskCompletionSource, AsyncLocal<int> Depth)> _isLoading = new(NavEntryEqualityComparer.Instance);
+    private readonly Dictionary<(object Entity, string NavigationName), TaskCompletionSource> _isLoading = new(NavEntryEqualityComparer.Instance);
+    private static readonly AsyncLocal<int> _isLoadingCallDepth = new();
     private HashSet<string>? _nonLazyNavigations;
 
     /// <summary>
@@ -112,26 +113,28 @@ public virtual void Load(object entity, [CallerMemberName] string navigationName
         var navEntry = (entity, navigationName);
 
         bool exists;
-        (TaskCompletionSource TaskCompletionSource, AsyncLocal<int> Depth) isLoadingValue;
+        TaskCompletionSource isLoadingValue;
 
         lock (_isLoadingLock)
         {
             ref var refIsLoadingValue = ref CollectionsMarshal.GetValueRefOrAddDefault(_isLoading, navEntry, out exists);
             if (!exists)
             {
-                refIsLoadingValue = (new(), new());
+                refIsLoadingValue = new();
             }
+            _isLoadingCallDepth.Value++;
             isLoadingValue = refIsLoadingValue!;
-            isLoadingValue.Depth.Value++;
         }
 
         if (exists)
         {
-            // Only waits for the outermost call on the call stack. See  #35528.
-            if (isLoadingValue.Depth.Value == 1)
+            // Only waits for the outermost call on the call stack. See #35528.
+            // if _isLoadingCallDepth.Value > 1 the call is recursive, waiting probably generates a deadlock See #35832.
+            if (_isLoadingCallDepth.Value == 1)
             {
-                isLoadingValue.TaskCompletionSource.Task.Wait();
+                isLoadingValue.Task.Wait();
             }
+            _isLoadingCallDepth.Value--;
             return;
         }
 
@@ -156,7 +159,8 @@ public virtual void Load(object entity, [CallerMemberName] string navigationName
         }
         finally
         {
-            isLoadingValue.TaskCompletionSource.TrySetResult();
+            isLoadingValue.TrySetResult();
+            _isLoadingCallDepth.Value--;
             lock (_isLoadingLock)
             {
                 _isLoading.Remove(navEntry);
@@ -181,26 +185,28 @@ public virtual async Task LoadAsync(
         var navEntry = (entity, navigationName);
 
         bool exists;
-        (TaskCompletionSource TaskCompletionSource, AsyncLocal<int> Depth) isLoadingValue;
+        TaskCompletionSource isLoadingValue;
 
         lock (_isLoadingLock)
         {
             ref var refIsLoadingValue = ref CollectionsMarshal.GetValueRefOrAddDefault(_isLoading, navEntry, out exists);
             if (!exists)
             {
-                refIsLoadingValue = (new(), new());
+                refIsLoadingValue = new();
             }
+            _isLoadingCallDepth.Value++;
             isLoadingValue = refIsLoadingValue!;
-            isLoadingValue.Depth.Value++;
         }
 
         if (exists)
         {
-            // Only waits for the outermost call on the call stack. See  #35528.
-            if (isLoadingValue.Depth.Value == 1)
+            // Only waits for the outermost call on the call stack. See #35528.
+            // if _isLoadingCallDepth.Value > 1 the call is recursive, waiting probably generates a deadlock See #35832.
+            if (_isLoadingCallDepth.Value == 1)
             {
-                await isLoadingValue.TaskCompletionSource.Task.WaitAsync(cancellationToken).ConfigureAwait(false);
+                await isLoadingValue.Task.WaitAsync(cancellationToken).ConfigureAwait(false);
             }
+            _isLoadingCallDepth.Value--;
             return;
         }
 
@@ -226,7 +232,8 @@ await entry.LoadAsync(
         }
         finally
         {
-            isLoadingValue.TaskCompletionSource.TrySetResult();
+            isLoadingValue.TrySetResult();
+            _isLoadingCallDepth.Value--;
             lock (_isLoadingLock)
             {
                 _isLoading.Remove(navEntry);
diff --git a/test/EFCore.Specification.Tests/LoadTestBase.cs b/test/EFCore.Specification.Tests/LoadTestBase.cs
@@ -5059,14 +5059,31 @@ public virtual async Task Lazy_loading_is_thread_safe(bool noTracking, bool asyn
         var parent = query.Single();
 
         var children = (await parent.LazyLoadChildren(async))?.Select(x => x.Id).OrderBy(x => x).ToList();
+        var childrenInvert = (await parent.LazyLoadChildren(!async))?.Select(x => x.Id).OrderBy(x => x).ToList();
+
         var singlePkToPk = (await parent.LazyLoadSinglePkToPk(async))?.Id;
+        var singlePkToPkInvert = (await parent.LazyLoadSinglePkToPk(!async))?.Id;
+
         var single = (await parent.LazyLoadSingle(async))?.Id;
+        var singleInvert = (await parent.LazyLoadSingle(!async))?.Id;
+
         var childrenAk = (await parent.LazyLoadChildrenAk(async))?.Select(x => x.Id).OrderBy(x => x).ToList();
+        var childrenAkInvert = (await parent.LazyLoadChildrenAk(!async))?.Select(x => x.Id).OrderBy(x => x).ToList();
+
         var singleAk = (await parent.LazyLoadSingleAk(async))?.Id;
+        var singleAkInvert = (await parent.LazyLoadSingleAk(!async))?.Id;
+
         var childrenShadowFk = (await parent.LazyLoadChildrenShadowFk(async))?.Select(x => x.Id).OrderBy(x => x).ToList();
+        var childrenShadowFkInvert = (await parent.LazyLoadChildrenShadowFk(!async))?.Select(x => x.Id).OrderBy(x => x).ToList();
+
         var singleShadowFk = (await parent.LazyLoadSingleShadowFk(async))?.Id;
+        var singleShadowFkInvert = (await parent.LazyLoadSingleShadowFk(!async))?.Id;
+
         var childrenCompositeKey = (await parent.LazyLoadChildrenCompositeKey(async))?.Select(x => x.Id).OrderBy(x => x).ToList();
+        var childrenCompositeKeyInvert = (await parent.LazyLoadChildrenCompositeKey(!async))?.Select(x => x.Id).OrderBy(x => x).ToList();
+
         var singleCompositeKey = (await parent.LazyLoadSingleCompositeKey(async))?.Id;
+        var singleCompositeKeyInvert = (await parent.LazyLoadSingleCompositeKey(!async))?.Id;
 
         var parent2 = query2.Single();
 
@@ -5075,18 +5092,57 @@ public virtual async Task Lazy_loading_is_thread_safe(bool noTracking, bool asyn
             MaxDegreeOfParallelism = Environment.ProcessorCount * 500
         };
 
-        await Parallel.ForAsync(0, 50000, parallelOptions, async (i, ct) =>
-        {
-            Assert.Equal(children, (await parent2.LazyLoadChildren(async))?.Select(x => x.Id).OrderBy(x => x).ToList());
-            Assert.Equal(singlePkToPk, (await parent2.LazyLoadSinglePkToPk(async))?.Id);
-            Assert.Equal(single, (await parent2.LazyLoadSingle(async))?.Id);
-            Assert.Equal(childrenAk, (await parent2.LazyLoadChildrenAk(async))?.Select(x => x.Id).OrderBy(x => x).ToList());
-            Assert.Equal(singleAk, (await parent2.LazyLoadSingleAk(async))?.Id);
-            Assert.Equal(childrenShadowFk, (await parent2.LazyLoadChildrenShadowFk(async))?.Select(x => x.Id).OrderBy(x => x).ToList());
-            Assert.Equal(singleShadowFk, (await parent2.LazyLoadSingleShadowFk(async))?.Id);
-            Assert.Equal(childrenCompositeKey, (await parent2.LazyLoadChildrenCompositeKey(async))?.Select(x => x.Id).OrderBy(x => x).ToList());
-            Assert.Equal(singleCompositeKey, (await parent2.LazyLoadSingleCompositeKey(async))?.Id);
+        await Parallel.ForAsync(0, 10000, parallelOptions, async (i, ct) =>
+        {
+            await Task.WhenAll(
+                AssertEqual(
+                    (children, async () => (await parent2.LazyLoadChildren(async))?.Select(x => x.Id).OrderBy(x => x).ToList()),
+                    (childrenInvert, async () => (await parent2.LazyLoadChildren(!async))?.Select(x => x.Id).OrderBy(x => x).ToList())
+                ),
+                AssertEqual(
+                    (singlePkToPk, async () => (await parent2.LazyLoadSinglePkToPk(async))?.Id),
+                    (singlePkToPkInvert, async () => (await parent2.LazyLoadSinglePkToPk(!async))?.Id)
+                ),
+                AssertEqual(
+                    (single, async () => (await parent2.LazyLoadSingle(async))?.Id),
+                    (singleInvert, async () => (await parent2.LazyLoadSingle(!async))?.Id)
+                ),
+                AssertEqual(
+                    (childrenAk, async () => (await parent2.LazyLoadChildrenAk(async))?.Select(x => x.Id).OrderBy(x => x).ToList()),
+                    (childrenAkInvert, async () => (await parent2.LazyLoadChildrenAk(!async))?.Select(x => x.Id).OrderBy(x => x).ToList())
+                ),
+                AssertEqual(
+                    (singleAk, async () => (await parent2.LazyLoadSingleAk(async))?.Id),
+                    (singleAkInvert, async () => (await parent2.LazyLoadSingleAk(!async))?.Id)
+                ),
+                AssertEqual(
+                    (childrenShadowFk, async () => (await parent2.LazyLoadChildrenShadowFk(async))?.Select(x => x.Id).OrderBy(x => x).ToList()),
+                    (childrenShadowFkInvert, async () => (await parent2.LazyLoadChildrenShadowFk(!async))?.Select(x => x.Id).OrderBy(x => x).ToList())
+                ),
+                AssertEqual(
+                    (singleShadowFk, async () => (await parent2.LazyLoadSingleShadowFk(async))?.Id),
+                    (singleShadowFkInvert, async () => (await parent2.LazyLoadSingleShadowFk(!async))?.Id)
+                ),
+                AssertEqual(
+                    (childrenCompositeKey, async () => (await parent2.LazyLoadChildrenCompositeKey(async))?.Select(x => x.Id).OrderBy(x => x).ToList()),
+                    (childrenCompositeKeyInvert, async () => (await parent2.LazyLoadChildrenCompositeKey(!async))?.Select(x => x.Id).OrderBy(x => x).ToList())
+                ),
+                AssertEqual(
+                    (singleCompositeKey, async () => (await parent2.LazyLoadSingleCompositeKey(async))?.Id),
+                    (singleCompositeKeyInvert, async () => (await parent2.LazyLoadSingleCompositeKey(!async))?.Id)
+                )
+            );
         });
+
+        static async Task AssertEqual<T>((T Data, Func<Task<T>> Expected) data, (T Data, Func<Task<T>> Expected) dataInvert)
+        {
+            //Do the processing at the same time
+            var dataTask = data.Expected();
+            var dataInvertTask = dataInvert.Expected();
+
+            Assert.Equal(data.Data, await dataTask);
+            Assert.Equal(dataInvert.Data, await dataInvertTask);
+        }
     }
 
     private static void SetState(

Original file line number	Diff line number	Diff line change
`@@ -23,7 +23,8 @@ public class LazyLoader : ILazyLoader, IInjectableService`
`23`	`23`	`private bool _detached;`
`24`	`24`	`private IDictionary<string, bool>? _loadedStates;`
`25`	`25`	`private readonly Lock _isLoadingLock = new Lock();`
`26`		`- private readonly Dictionary<(object Entity, string NavigationName), (TaskCompletionSource TaskCompletionSource, AsyncLocal<int> Depth)> _isLoading = new(NavEntryEqualityComparer.Instance);`
	`26`	`+ private readonly Dictionary<(object Entity, string NavigationName), TaskCompletionSource> _isLoading = new(NavEntryEqualityComparer.Instance);`
	`27`	`+ private static readonly AsyncLocal<int> _isLoadingCallDepth = new();`
`27`	`28`	`private HashSet<string>? _nonLazyNavigations;`
`28`	`29`
`29`	`30`	`/// <summary>`
`@@ -112,26 +113,28 @@ public virtual void Load(object entity, [CallerMemberName] string navigationName`
`112`	`113`	`var navEntry = (entity, navigationName);`
`113`	`114`
`114`	`115`	`bool exists;`
`115`		`- (TaskCompletionSource TaskCompletionSource, AsyncLocal<int> Depth) isLoadingValue;`
	`116`	`+ TaskCompletionSource isLoadingValue;`
`116`	`117`
`117`	`118`	`lock (_isLoadingLock)`
`118`	`119`	`{`
`119`	`120`	`ref var refIsLoadingValue = ref CollectionsMarshal.GetValueRefOrAddDefault(_isLoading, navEntry, out exists);`
`120`	`121`	`if (!exists)`
`121`	`122`	`{`
`122`		`- refIsLoadingValue = (new(), new());`
	`123`	`+ refIsLoadingValue = new();`
`123`	`124`	`}`
	`125`	`+ _isLoadingCallDepth.Value++;`
`124`	`126`	`isLoadingValue = refIsLoadingValue!;`
`125`		`- isLoadingValue.Depth.Value++;`
`126`	`127`	`}`
`127`	`128`
`128`	`129`	`if (exists)`
`129`	`130`	`{`
`130`		`- // Only waits for the outermost call on the call stack. See #35528.`
`131`		`- if (isLoadingValue.Depth.Value == 1)`
	`131`	`+ // Only waits for the outermost call on the call stack. See #35528.`
	`132`	`+ // if _isLoadingCallDepth.Value > 1 the call is recursive, waiting probably generates a deadlock See #35832.`
	`133`	`+ if (_isLoadingCallDepth.Value == 1)`
`132`	`134`	`{`
`133`		`- isLoadingValue.TaskCompletionSource.Task.Wait();`
	`135`	`+ isLoadingValue.Task.Wait();`
`134`	`136`	`}`
	`137`	`+ _isLoadingCallDepth.Value--;`
`135`	`138`	`return;`
`136`	`139`	`}`
`137`	`140`
`@@ -156,7 +159,8 @@ public virtual void Load(object entity, [CallerMemberName] string navigationName`
`156`	`159`	`}`
`157`	`160`	`finally`
`158`	`161`	`{`
`159`		`- isLoadingValue.TaskCompletionSource.TrySetResult();`
	`162`	`+ isLoadingValue.TrySetResult();`
	`163`	`+ _isLoadingCallDepth.Value--;`
`160`	`164`	`lock (_isLoadingLock)`
`161`	`165`	`{`
`162`	`166`	`_isLoading.Remove(navEntry);`
`@@ -181,26 +185,28 @@ public virtual async Task LoadAsync(`
`181`	`185`	`var navEntry = (entity, navigationName);`
`182`	`186`
`183`	`187`	`bool exists;`
`184`		`- (TaskCompletionSource TaskCompletionSource, AsyncLocal<int> Depth) isLoadingValue;`
	`188`	`+ TaskCompletionSource isLoadingValue;`
`185`	`189`
`186`	`190`	`lock (_isLoadingLock)`
`187`	`191`	`{`
`188`	`192`	`ref var refIsLoadingValue = ref CollectionsMarshal.GetValueRefOrAddDefault(_isLoading, navEntry, out exists);`
`189`	`193`	`if (!exists)`
`190`	`194`	`{`
`191`		`- refIsLoadingValue = (new(), new());`
	`195`	`+ refIsLoadingValue = new();`
`192`	`196`	`}`
	`197`	`+ _isLoadingCallDepth.Value++;`
`193`	`198`	`isLoadingValue = refIsLoadingValue!;`
`194`		`- isLoadingValue.Depth.Value++;`
`195`	`199`	`}`
`196`	`200`
`197`	`201`	`if (exists)`
`198`	`202`	`{`
`199`		`- // Only waits for the outermost call on the call stack. See #35528.`
`200`		`- if (isLoadingValue.Depth.Value == 1)`
	`203`	`+ // Only waits for the outermost call on the call stack. See #35528.`
	`204`	`+ // if _isLoadingCallDepth.Value > 1 the call is recursive, waiting probably generates a deadlock See #35832.`
	`205`	`+ if (_isLoadingCallDepth.Value == 1)`
`201`	`206`	`{`
`202`		`- await isLoadingValue.TaskCompletionSource.Task.WaitAsync(cancellationToken).ConfigureAwait(false);`
	`207`	`+ await isLoadingValue.Task.WaitAsync(cancellationToken).ConfigureAwait(false);`
`203`	`208`	`}`
	`209`	`+ _isLoadingCallDepth.Value--;`
`204`	`210`	`return;`
`205`	`211`	`}`
`206`	`212`
`@@ -226,7 +232,8 @@ await entry.LoadAsync(`
`226`	`232`	`}`
`227`	`233`	`finally`
`228`	`234`	`{`
`229`		`- isLoadingValue.TaskCompletionSource.TrySetResult();`
	`235`	`+ isLoadingValue.TrySetResult();`
	`236`	`+ _isLoadingCallDepth.Value--;`
`230`	`237`	`lock (_isLoadingLock)`
`231`	`238`	`{`
`232`	`239`	`_isLoading.Remove(navEntry);`