Fix hang in AsyncLazy<T> when value factory synchronously blocks for UI thread

src/Microsoft.VisualStudio.Threading.Shared/AsyncLazy.cs

@@ -133,6 +133,7 @@ public Task<T> GetValueAsync(CancellationToken cancellationToken)
                     Verify.FailOperation(Strings.ValueFactoryReentrancy);
                 }
 
+                InlineResumable resumableAwaiter = null;
                 lock (this.syncObject)
                 {
                     // Note that if multiple threads hit GetValueAsync() before
@@ -143,12 +144,18 @@ public Task<T> GetValueAsync(CancellationToken cancellationToken)
                     if (this.value == null)
                     {
                         cancellationToken.ThrowIfCancellationRequested();
+                        resumableAwaiter = new InlineResumable();
+                        Func<Task<T>> originalValueFactory = this.valueFactory;
+                        this.valueFactory = null;
+                        Func<Task<T>> valueFactory = async delegate
+                        {
+                            await resumableAwaiter;
+                            return await originalValueFactory().ConfigureAwait(continueOnCapturedContext: false);
+                        };
+
                         this.recursiveFactoryCheck.Value = RecursiveCheckSentinel;
                         try
                         {
-                            var valueFactory = this.valueFactory;
-                            this.valueFactory = null;
-
                             if (this.jobFactory != null)
                             {
                                 // Wrapping with RunAsync allows a future caller
@@ -171,18 +178,15 @@ public Task<T> GetValueAsync(CancellationToken cancellationToken)
                                 this.value = valueFactory();
                             }
                         }
-                        catch (Exception ex)
-                        {
-                            var tcs = new TaskCompletionSource<T>();
-                            tcs.SetException(ex);
-                            this.value = tcs.Task;
-                        }
                         finally
                         {
                             this.recursiveFactoryCheck.Value = null;
                         }
                     }
                 }
+
+                // Allow the original value factory to actually run.
+                resumableAwaiter?.Resume();
             }
 
             if (!this.value.IsCompleted)

src/Microsoft.VisualStudio.Threading.Shared/InlineResumable.cs

@@ -0,0 +1,80 @@
+/********************************************************
+*                                                        *
+*   © Copyright (C) Microsoft. All rights reserved.      *
+*                                                        *
+*********************************************************/
+
+namespace Microsoft.VisualStudio.Threading
+{
+    using System;
+    using System.Runtime.CompilerServices;
+    using System.Threading;
+
+    /// <summary>
+    /// An awaiter that can be pre-created, and later immediately execute its one scheduled continuation.
+    /// </summary>
+    internal class InlineResumable : INotifyCompletion
+    {
+        /// <summary>
+        /// The continuation that has been scheduled.
+        /// </summary>
+        private Action continuation;
+
+        /// <summary>
+        /// The current <see cref="SynchronizationContext"/> as of when the continuation was scheduled.
+        /// </summary>
+        private SynchronizationContext capturedSynchronizationContext;
+
+        /// <summary>
+        /// Whether <see cref="Resume"/> has been called already.
+        /// </summary>
+        private bool resumed;
+
+        /// <summary>
+        /// Gets a value indicating whether an awaiting expression should yield.
+        /// </summary>
+        /// <value>Always <c>false</c></value>
+        public bool IsCompleted => this.resumed;
+
+        /// <summary>
+        /// Does and returns nothing.
+        /// </summary>
+        [System.Diagnostics.CodeAnalysis.SuppressMessage("Microsoft.Performance", "CA1822:MarkMembersAsStatic", Justification = "Awaiter pattern requires this to be an instance method.")]
+        public void GetResult()
+        {
+        }
+
+        /// <summary>
+        /// Stores the continuation for later execution when <see cref="Resume"/> is invoked.
+        /// </summary>
+        /// <param name="continuation">The delegate to execute later.</param>
+        public void OnCompleted(Action continuation)
+        {
+            Requires.NotNull(continuation, nameof(continuation));
+            Assumes.Null(this.continuation); // Only one continuation is supported.
+
+            this.capturedSynchronizationContext = SynchronizationContext.Current;
+            this.continuation = continuation;
+        }
+
+        /// <summary>
+        /// Gets this instance. This method makes this awaiter double as its own awaitable.
+        /// </summary>
+        /// <returns>This instance.</returns>
+        public InlineResumable GetAwaiter() => this;
+
+        /// <summary>
+        /// Executes the continuation immediately, on the caller's thread.
+        /// </summary>
+        public void Resume()
+        {
+            this.resumed = true;
+            var continuation = this.continuation;
+            this.continuation = null;
+            using (this.capturedSynchronizationContext.Apply())
+            {
+                continuation?.Invoke();
+            }
+        }
+    }
+}

src/Microsoft.VisualStudio.Threading.Shared/Microsoft.VisualStudio.Threading.Shared.projitems

@@ -31,6 +31,7 @@
     <Compile Include="$(MSBuildThisFileDirectory)HangReportContribution.cs" />
     <Compile Include="$(MSBuildThisFileDirectory)IAsyncDisposable.cs" />
     <Compile Include="$(MSBuildThisFileDirectory)IHangReportContributor.cs" />
+    <Compile Include="$(MSBuildThisFileDirectory)InlineResumable.cs" />
     <Compile Include="$(MSBuildThisFileDirectory)InternalUtilities.cs" />
     <Compile Include="$(MSBuildThisFileDirectory)JoinableTask+DependentSynchronousTask.cs" />
     <Compile Include="$(MSBuildThisFileDirectory)JoinableTask+ExecutionQueue.cs" />

src/Microsoft.VisualStudio.Threading.Tests.Shared/AsyncLazyTests.cs

@@ -412,6 +412,97 @@ async delegate
             Assert.True(resultTask.Wait(AsyncDelay));
         }
 
+        /// <summary>
+        /// Verifies that no deadlock occurs if the value factory synchronously blocks while switching to the UI thread.
+        /// </summary>
+        [Theory]
+        [CombinatorialData]
+        public void ValueFactoryRequiresMainThreadHeldByOtherSync(bool passJtfToLazyCtor)
+        {
+            var ctxt = SingleThreadedSynchronizationContext.New();
+            SynchronizationContext.SetSynchronizationContext(ctxt);
+            var context = new JoinableTaskContext();
+            var asyncPump = context.Factory;
+            var originalThread = Thread.CurrentThread;
+
+            var evt = new AsyncManualResetEvent();
+            var lazy = new AsyncLazy<object>(
+                async delegate
+                {
+                    // It is important that no await appear before this JTF.Run call, since
+                    // we're testing that the value factory is not invoked while the AsyncLazy
+                    // holds a private lock that would deadlock when called from another thread.
+                    asyncPump.Run(async delegate
+                    {
+                        await asyncPump.SwitchToMainThreadAsync(this.TimeoutToken);
+                    });
+                    await Task.Yield();
+                    return new object();
+                },
+                passJtfToLazyCtor ? asyncPump : null); // mix it up to exercise all the code paths in the ctor.
+
+            var backgroundRequest = Task.Run(async delegate
+            {
+                return await lazy.GetValueAsync();
+            });
+
+            Thread.Sleep(AsyncDelay); // Give the background thread time to call GetValueAsync(), but it doesn't yield (when the test was written).
+            var foregroundRequest = lazy.GetValueAsync();
+
+            var frame = SingleThreadedSynchronizationContext.NewFrame();
+            var combinedTask = Task.WhenAll(foregroundRequest, backgroundRequest);
+            combinedTask.WithTimeout(UnexpectedTimeout).ContinueWith(_ => frame.Continue = false, TaskScheduler.Default);
+            SingleThreadedSynchronizationContext.PushFrame(ctxt, frame);
+
+            // Ensure that the test didn't simply timeout, and that the individual tasks did not throw.
+            Assert.True(foregroundRequest.IsCompleted);
+            Assert.True(backgroundRequest.IsCompleted);
+            Assert.Same(foregroundRequest.GetAwaiter().GetResult(), backgroundRequest.GetAwaiter().GetResult());
+        }
+
+        /// <summary>
+        /// Verifies that no deadlock occurs if the value factory synchronously blocks while switching to the UI thread
+        /// and the UI thread then starts a JTF.Run that wants it too.
+        /// </summary>
+        [Fact]
+        public void ValueFactoryRequiresMainThreadHeldByOtherInJTFRun()
+        {
+            var ctxt = SingleThreadedSynchronizationContext.New();
+            SynchronizationContext.SetSynchronizationContext(ctxt);
+            var context = new JoinableTaskContext();
+            var asyncPump = context.Factory;
+            var originalThread = Thread.CurrentThread;
+
+            var evt = new AsyncManualResetEvent();
+            var lazy = new AsyncLazy<object>(
+                async delegate
+                {
+                    // It is important that no await appear before this JTF.Run call, since
+                    // we're testing that the value factory is not invoked while the AsyncLazy
+                    // holds a private lock that would deadlock when called from another thread.
+                    asyncPump.Run(async delegate
+                    {
+                        await asyncPump.SwitchToMainThreadAsync(this.TimeoutToken);
+                    });
+                    await Task.Yield();
+                    return new object();
+                },
+                asyncPump);
+
+            var backgroundRequest = Task.Run(async delegate
+            {
+                return await lazy.GetValueAsync();
+            });
+
+            Thread.Sleep(AsyncDelay); // Give the background thread time to call GetValueAsync(), but it doesn't yield (when the test was written).
+            asyncPump.Run(async delegate
+            {
+                var foregroundValue = await lazy.GetValueAsync(this.TimeoutToken);
+                var backgroundValue = await backgroundRequest;
+                Assert.Same(foregroundValue, backgroundValue);
+            });
+        }
+
         [Fact(Skip = "Hangs. This test documents a deadlock scenario that is not fixed (by design, IIRC).")]
         public async Task ValueFactoryRequiresReadLockHeldByOther()
         {