Prevent creating two task waiting chains per AsyncLazy.GetValueAsync(CancellationToken)

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

@@ -245,12 +245,7 @@ public Task<T> GetValueAsync(CancellationToken cancellationToken)
             resumableAwaiter?.Resume();
         }
 
-        if (!this.value.IsCompleted)
-        {
-            this.joinableTask?.JoinAsync(cancellationToken).Forget();
-        }
-
-        return this.value.WithCancellation(cancellationToken);
+        return this.joinableTask?.JoinAsync(continueOnCapturedContext: false, cancellationToken) ?? this.value.WithCancellation(cancellationToken);
     }
 
     /// <summary>

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

@@ -49,6 +49,46 @@ internal JoinableTask(JoinableTaskFactory owner, bool synchronouslyBlocking, str
         /// <param name="cancellationToken">A cancellation token that will exit this method before the task is completed.</param>
         /// <returns>A task that completes after the asynchronous operation completes and the join is reverted, with the result of the operation.</returns>
         public new Task<T> JoinAsync(CancellationToken cancellationToken = default(CancellationToken))
+        {
+            return this.JoinAsync(continueOnCapturedContext: AwaitShouldCaptureSyncContext, cancellationToken);
+        }
+
+        /// <summary>
+        /// Synchronously blocks the calling thread until the operation has completed.
+        /// If the calling thread is the Main thread, deadlocks are mitigated.
+        /// </summary>
+        /// <param name="cancellationToken">A cancellation token that will exit this method before the task is completed.</param>
+        /// <returns>The result of the asynchronous operation.</returns>
+        public new T Join(CancellationToken cancellationToken = default(CancellationToken))
+        {
+            base.Join(cancellationToken);
+            Assumes.True(this.Task.IsCompleted);
+            return this.Task.Result;
+        }
+
+        /// <summary>
+        /// Gets an awaiter that is equivalent to calling <see cref="JoinAsync(CancellationToken)"/>.
+        /// </summary>
+        /// <returns>A task whose result is the result of the asynchronous operation.</returns>
+        public new TaskAwaiter<T> GetAwaiter()
+        {
+            return this.JoinAsync().GetAwaiter();
+        }
+
+        internal new T CompleteOnCurrentThread()
+        {
+            base.CompleteOnCurrentThread();
+            return this.Task.GetAwaiter().GetResult();
+        }
+
+        /// <summary>
+        /// Joins any main thread affinity of the caller with the asynchronous operation to avoid deadlocks
+        /// in the event that the main thread ultimately synchronously blocks waiting for the operation to complete.
+        /// </summary>
+        /// <param name="continueOnCapturedContext">A value indicating whether *internal* continuations required to respond to cancellation should run on the current <see cref="SynchronizationContext"/>.</param>
+        /// <param name="cancellationToken">A cancellation token that will exit this method before the task is completed.</param>
+        /// <returns>A task that completes after the asynchronous operation completes and the join is reverted, with the result of the operation.</returns>
+        internal Task<T> JoinAsync(bool continueOnCapturedContext, CancellationToken cancellationToken = default(CancellationToken))
         {
             cancellationToken.ThrowIfCancellationRequested();
             if (this.IsCompleted)
@@ -65,18 +105,18 @@ internal JoinableTask(JoinableTaskFactory owner, bool synchronouslyBlocking, str
             }
             else
             {
-                return JoinSlowAsync(this, cancellationToken);
+                return JoinSlowAsync(this, continueOnCapturedContext, cancellationToken);
             }
 
-            static async Task<T> JoinSlowAsync(JoinableTask<T> me, CancellationToken cancellationToken)
+            static async Task<T> JoinSlowAsync(JoinableTask<T> me, bool continueOnCapturedContext, CancellationToken cancellationToken)
             {
                 // No need to dispose of this except in cancellation case.
                 JoinableTaskCollection.JoinRelease dependency = me.AmbientJobJoinsThis();
 
                 try
                 {
-                    await me.Task.WithCancellation(continueOnCapturedContext: AwaitShouldCaptureSyncContext, cancellationToken).ConfigureAwait(AwaitShouldCaptureSyncContext);
-                    return await me.Task.ConfigureAwait(AwaitShouldCaptureSyncContext);
+                    await me.Task.WithCancellation(continueOnCapturedContext, cancellationToken).ConfigureAwait(continueOnCapturedContext);
+                    return await me.Task.ConfigureAwait(continueOnCapturedContext);
                 }
                 catch (OperationCanceledException)
                 {
@@ -86,34 +126,6 @@ static async Task<T> JoinSlowAsync(JoinableTask<T> me, CancellationToken cancell
             }
         }
 
-        /// <summary>
-        /// Synchronously blocks the calling thread until the operation has completed.
-        /// If the calling thread is the Main thread, deadlocks are mitigated.
-        /// </summary>
-        /// <param name="cancellationToken">A cancellation token that will exit this method before the task is completed.</param>
-        /// <returns>The result of the asynchronous operation.</returns>
-        public new T Join(CancellationToken cancellationToken = default(CancellationToken))
-        {
-            base.Join(cancellationToken);
-            Assumes.True(this.Task.IsCompleted);
-            return this.Task.Result;
-        }
-
-        /// <summary>
-        /// Gets an awaiter that is equivalent to calling <see cref="JoinAsync"/>.
-        /// </summary>
-        /// <returns>A task whose result is the result of the asynchronous operation.</returns>
-        public new TaskAwaiter<T> GetAwaiter()
-        {
-            return this.JoinAsync().GetAwaiter();
-        }
-
-        internal new T CompleteOnCurrentThread()
-        {
-            base.CompleteOnCurrentThread();
-            return this.Task.GetAwaiter().GetResult();
-        }
-
         /// <inheritdoc/>
         internal override object CreateTaskCompletionSource() => new TaskCompletionSourceWithoutInlining<T>(allowInliningContinuations: false);