Fix hang in AsyncLazy<T>.DisposeValue() when T : IAsyncDisposable

src/Microsoft.VisualStudio.Threading/AsyncLazy`1.cs

@@ -41,14 +41,14 @@ public class AsyncLazy<T>
     private AsyncLocal<object>? recursiveFactoryCheck;
 
     /// <summary>
-    /// The function to invoke to produce the task.
+    /// An optional means to avoid deadlocks when synchronous APIs are called that must invoke async methods in user code.
     /// </summary>
-    private Func<Task<T>>? valueFactory;
+    private readonly JoinableTaskFactory? jobFactory;
 
     /// <summary>
-    /// The async pump to Join on calls to <see cref="GetValueAsync(CancellationToken)"/>.
+    /// The function to invoke to produce the task.
     /// </summary>
-    private JoinableTaskFactory? jobFactory;
+    private Func<Task<T>>? valueFactory;
 
     /// <summary>
     /// The result of the value factory.
@@ -64,7 +64,10 @@ public class AsyncLazy<T>
     /// Initializes a new instance of the <see cref="AsyncLazy{T}"/> class.
     /// </summary>
     /// <param name="valueFactory">The async function that produces the value.  To be invoked at most once.</param>
-    /// <param name="joinableTaskFactory">The factory to use when invoking the value factory in <see cref="GetValueAsync(CancellationToken)"/> to avoid deadlocks when the main thread is required by the value factory.</param>
+    /// <param name="joinableTaskFactory">
+    /// The <see cref="JoinableTaskFactory" /> to use for avoiding deadlocks when the <paramref name="valueFactory"/>
+    /// or the constructed value's <see cref="System.IAsyncDisposable.DisposeAsync"/> method may require the main thread in the process.
+    /// </param>
     public AsyncLazy(Func<Task<T>> valueFactory, JoinableTaskFactory? joinableTaskFactory = null)
     {
         Requires.NotNull(valueFactory, nameof(valueFactory));
@@ -203,7 +206,6 @@ public Task<T> GetValueAsync(CancellationToken cancellationToken)
                         }
                         finally
                         {
-                            this.jobFactory = null;
                             this.joinableTask = null;
                         }
                     };
@@ -285,11 +287,8 @@ public T GetValue(CancellationToken cancellationToken)
         }
         else
         {
-            // Capture the factory as a local before comparing and dereferencing it since
-            // the field can transition to null and we want to gracefully handle that race condition.
-            JoinableTaskFactory? factory = this.jobFactory;
-            return factory is object
-                ? factory.Run(() => this.GetValueAsync(cancellationToken))
+            return this.jobFactory is JoinableTaskFactory jtf
+                ? jtf.Run(() => this.GetValueAsync(cancellationToken))
                 : this.GetValueAsync(cancellationToken).GetAwaiter().GetResult();
         }
     }

test/Microsoft.VisualStudio.Threading.Tests/AsyncLazyTests.cs

@@ -897,6 +897,44 @@ public async Task Dispose_Disposable_Incomplete(DisposeStyle variety)
         await value.Disposed.WithCancellation(this.TimeoutToken);
     }
 
+    [Fact]
+    public void DisposeValue_AsyncDisposableValueRequiresMainThread()
+    {
+        JoinableTaskContext context = this.InitializeJTCAndSC();
+        SingleThreadedTestSynchronizationContext.IFrame frame = SingleThreadedTestSynchronizationContext.NewFrame();
+
+        AsyncLazy<SystemAsyncDisposable> lazy = new(
+            delegate
+            {
+                return Task.FromResult(new SystemAsyncDisposable { YieldDuringDispose = true });
+            },
+            context.Factory);
+        lazy.GetValue();
+
+        TaskCompletionSource<bool> delegateResult = new();
+        SynchronizationContext.Current!.Post(
+            delegate
+            {
+                try
+                {
+                    lazy.DisposeValue();
+
+                    delegateResult.SetResult(true);
+                }
+                catch (Exception ex)
+                {
+                    delegateResult.SetException(ex);
+                }
+                finally
+                {
+                    frame.Continue = false;
+                }
+            },
+            null);
+        SingleThreadedTestSynchronizationContext.PushFrame(SynchronizationContext.Current!, frame);
+        delegateResult.Task.GetAwaiter().GetResult(); // rethrow any exceptions
+    }
+
     [Fact]
     public async Task Dispose_NonDisposable_Incomplete()
     {
@@ -1066,10 +1104,15 @@ private class Disposable : DisposableBase, IDisposableObservable
 
     private class SystemAsyncDisposable : DisposableBase, System.IAsyncDisposable
     {
-        public ValueTask DisposeAsync()
+        internal bool YieldDuringDispose { get; set; }
+
+        public async ValueTask DisposeAsync()
         {
             this.disposalEvent.Set();
-            return default;
+            if (this.YieldDuringDispose)
+            {
+                await Task.Yield();
+            }
         }
     }