Merge pull request #1340 from AArnott/libtemplateUpdate

Fix hang in `AsyncLazy<T>.DisposeValueAsync`
microsoft · Jul 19, 2024 · 98acf4e · 98acf4e
2 parents 678cc14 + 0cc0cea
commit 98acf4e
Show file tree

Hide file tree

Showing 2 changed files with 33 additions and 2 deletions.
diff --git a/src/Microsoft.VisualStudio.Threading/AsyncLazy`1.cs b/src/Microsoft.VisualStudio.Threading/AsyncLazy`1.cs
@@ -384,7 +384,7 @@ public void DisposeValue()
     /// </returns>
     /// <remarks>
     /// <para>Calling this method will put this object into a disposed state where future calls to obtain the value will throw <see cref="ObjectDisposedException"/>.</para>
-    /// <para>If the value has already been produced and implements <see cref="IDisposable"/>, <see cref="IAsyncDisposable"/>,  or <see cref="System.IAsyncDisposable"/> it will be disposed of.
+    /// <para>If the value has already been produced and implements <see cref="IDisposable"/>, <see cref="IAsyncDisposable"/>, or <see cref="System.IAsyncDisposable"/> it will be disposed of.
     /// If the value factory has already started but has not yet completed, its value will be disposed of when the value factory completes.</para>
     /// <para>If prior calls to obtain the value are in flight when this method is called, those calls <em>may</em> complete and their callers may obtain the value, although <see cref="IDisposable.Dispose"/>
     /// may have been or will soon be called on the value, leading those users to experience a <see cref="ObjectDisposedException"/>.</para>
@@ -394,6 +394,7 @@ public void DisposeValue()
     /// </remarks>
     public async Task DisposeValueAsync()
     {
+        JoinableTask<T>? localJoinableTask = null;
         Task<T>? localValueTask = null;
         object? localValue = default;
         lock (this.syncObject)
@@ -417,6 +418,7 @@ public async Task DisposeValueAsync()
                     // We'll schedule the value for disposal outside the lock so it can be synchronous with the value factory,
                     // but will not execute within our lock.
                     localValueTask = this.value;
+                    localJoinableTask = this.joinableTask;
                     break;
             }
 
@@ -431,7 +433,11 @@ public async Task DisposeValueAsync()
             this.valueFactory = null;
         }
 
-        if (localValueTask is not null)
+        if (localJoinableTask is not null)
+        {
+            localValue = await localJoinableTask;
+        }
+        else if (localValueTask is not null)
         {
             localValue = await localValueTask.ConfigureAwait(false);
         }

diff --git a/test/Microsoft.VisualStudio.Threading.Tests/AsyncLazyTests.cs b/test/Microsoft.VisualStudio.Threading.Tests/AsyncLazyTests.cs
@@ -988,6 +988,31 @@ public async Task Dispose_CalledTwice_Disposable_Completed()
         await this.AssertDisposedLazyAsync(lazy);
     }
 
+    [Fact]
+    public void DisposeValue_MidFactoryThatContestsForMainThread()
+    {
+        JoinableTaskContext context = this.InitializeJTCAndSC();
+
+        AsyncLazy<object> lazy = new(
+            async delegate
+            {
+                // Ensure the caller keeps control of the UI thread,
+                // so that the request for the main thread comes in when it's controlled by others.
+                await Task.Yield();
+                await context.Factory.SwitchToMainThreadAsync(this.TimeoutToken);
+                return new();
+            },
+            context.Factory);
+
+        Task<object> lazyFactory = lazy.GetValueAsync(this.TimeoutToken);
+
+        // Make a JTF blocking call on the main thread that won't return until the factory completes.
+        context.Factory.Run(async delegate
+        {
+            await lazy.DisposeValueAsync().WithCancellation(this.TimeoutToken);
+        });
+    }
+
     [Fact(Skip = "Hangs. This test documents a deadlock scenario that is not fixed (by design, IIRC).")]
     public async Task ValueFactoryRequiresReadLockHeldByOther()
     {