From da8a36d6071e13f0980ad3e6090d9d8199d36604 Mon Sep 17 00:00:00 2001 From: Lifeng Lu Date: Fri, 29 Mar 2024 11:59:58 -0700 Subject: [PATCH] Prevent creating two task waiting chains per AsyncLazy.GetValueAsync(CancellationToken) call. Those appeared a lot during the solution open time, because some global AsyncLazy doesn't provide value until the blocking loading time ends. And cancellation token is often used to handle the case loading is aborted. The reason why we didn't eliminate the duplication earlier is the JTF.JoinAsync and task.WithCancellation handles in-middle task continuation differently. The JTF one generally will capture the SynchorizationContext, while the task one would not. The exact performance charcter of the two choices is different. The JTF one prevent additional dependency to the thread pool in some cases, but may add unnecessary dependency to the UI thread in other cases, while the other one is on the reverse side. So eliminating the two chains could impact perf, because the accidental UI thread dependency is more likely to cause performance problems than the other one. It is impossible to decide which behavior is better, because the best choice should be aligned to the ConfiguredAwaiter option to wait the task later. But provide this extra option would make API more complex. In this case, we try to make the behavior matching the exisiting GetValueAsync API. Capturing the context makes little sense for the task which is forgotten in the original code path. --- .../AsyncLazy`1.cs | 7 +- .../JoinableTask`1.cs | 76 +++++++++++-------- 2 files changed, 45 insertions(+), 38 deletions(-) diff --git a/src/Microsoft.VisualStudio.Threading/AsyncLazy`1.cs b/src/Microsoft.VisualStudio.Threading/AsyncLazy`1.cs index 0b98fc7ad..24cafee01 100644 --- a/src/Microsoft.VisualStudio.Threading/AsyncLazy`1.cs +++ b/src/Microsoft.VisualStudio.Threading/AsyncLazy`1.cs @@ -245,12 +245,7 @@ public Task 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); } /// diff --git a/src/Microsoft.VisualStudio.Threading/JoinableTask`1.cs b/src/Microsoft.VisualStudio.Threading/JoinableTask`1.cs index 1ebead4ee..b7748174d 100644 --- a/src/Microsoft.VisualStudio.Threading/JoinableTask`1.cs +++ b/src/Microsoft.VisualStudio.Threading/JoinableTask`1.cs @@ -49,6 +49,46 @@ internal JoinableTask(JoinableTaskFactory owner, bool synchronouslyBlocking, str /// A cancellation token that will exit this method before the task is completed. /// A task that completes after the asynchronous operation completes and the join is reverted, with the result of the operation. public new Task JoinAsync(CancellationToken cancellationToken = default(CancellationToken)) + { + return this.JoinAsync(continueOnCapturedContext: AwaitShouldCaptureSyncContext, cancellationToken); + } + + /// + /// Synchronously blocks the calling thread until the operation has completed. + /// If the calling thread is the Main thread, deadlocks are mitigated. + /// + /// A cancellation token that will exit this method before the task is completed. + /// The result of the asynchronous operation. + public new T Join(CancellationToken cancellationToken = default(CancellationToken)) + { + base.Join(cancellationToken); + Assumes.True(this.Task.IsCompleted); + return this.Task.Result; + } + + /// + /// Gets an awaiter that is equivalent to calling . + /// + /// A task whose result is the result of the asynchronous operation. + public new TaskAwaiter GetAwaiter() + { + return this.JoinAsync().GetAwaiter(); + } + + internal new T CompleteOnCurrentThread() + { + base.CompleteOnCurrentThread(); + return this.Task.GetAwaiter().GetResult(); + } + + /// + /// 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. + /// + /// A value indicating whether *internal* continuations required to respond to cancellation should run on the current . + /// A cancellation token that will exit this method before the task is completed. + /// A task that completes after the asynchronous operation completes and the join is reverted, with the result of the operation. + internal Task 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 JoinSlowAsync(JoinableTask me, CancellationToken cancellationToken) + static async Task JoinSlowAsync(JoinableTask 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 JoinSlowAsync(JoinableTask me, CancellationToken cancell } } - /// - /// Synchronously blocks the calling thread until the operation has completed. - /// If the calling thread is the Main thread, deadlocks are mitigated. - /// - /// A cancellation token that will exit this method before the task is completed. - /// The result of the asynchronous operation. - public new T Join(CancellationToken cancellationToken = default(CancellationToken)) - { - base.Join(cancellationToken); - Assumes.True(this.Task.IsCompleted); - return this.Task.Result; - } - - /// - /// Gets an awaiter that is equivalent to calling . - /// - /// A task whose result is the result of the asynchronous operation. - public new TaskAwaiter GetAwaiter() - { - return this.JoinAsync().GetAwaiter(); - } - - internal new T CompleteOnCurrentThread() - { - base.CompleteOnCurrentThread(); - return this.Task.GetAwaiter().GetResult(); - } - /// internal override object CreateTaskCompletionSource() => new TaskCompletionSourceWithoutInlining(allowInliningContinuations: false);