Suppress ExecutionContext flow in more awaiters

microsoft · Sep 30, 2020 · 44d9db7 · 44d9db7
1 parent babb79a
commit 44d9db7
Show file tree

Hide file tree

Showing 5 changed files with 65 additions and 5 deletions.
diff --git a/src/Microsoft.VisualStudio.Threading/AsyncReaderWriterLock.cs b/src/Microsoft.VisualStudio.Threading/AsyncReaderWriterLock.cs
@@ -2525,9 +2525,10 @@ private void OnCompleted(Action continuation, bool flowExecutionContext)
                 }
 
                 bool restoreFlow = !flowExecutionContext && !ExecutionContext.IsFlowSuppressed();
+                AsyncFlowControl flowControl = default;
                 if (restoreFlow)
                 {
-                    ExecutionContext.SuppressFlow();
+                    flowControl = ExecutionContext.SuppressFlow();
                 }
 
                 try
@@ -2549,7 +2550,7 @@ private void OnCompleted(Action continuation, bool flowExecutionContext)
                 {
                     if (restoreFlow)
                     {
-                        ExecutionContext.RestoreFlow();
+                        flowControl.Dispose();
                     }
                 }
             }

diff --git a/src/Microsoft.VisualStudio.Threading/AwaitExtensions.cs b/src/Microsoft.VisualStudio.Threading/AwaitExtensions.cs
@@ -423,8 +423,21 @@ public void UnsafeOnCompleted(Action continuation)
                 }
                 else
                 {
-                    // There is no API for scheduling a Task without capturing the ExecutionContext.
+#if NETFRAMEWORK // Only bother suppressing flow on .NET Framework where the perf would improve from doing so.
+                    if (ExecutionContext.IsFlowSuppressed())
+                    {
+                        Task.Factory.StartNew(continuation, CancellationToken.None, TaskCreationOptions.None, this.scheduler);
+                    }
+                    else
+                    {
+                        using (ExecutionContext.SuppressFlow())
+                        {
+                            Task.Factory.StartNew(continuation, CancellationToken.None, TaskCreationOptions.None, this.scheduler);
+                        }
+                    }
+#else
                     Task.Factory.StartNew(continuation, CancellationToken.None, TaskCreationOptions.None, this.scheduler);
+#endif
                 }
             }
 
@@ -567,7 +580,7 @@ public ExecuteContinuationSynchronouslyAwaitable(Task antecedent)
         /// A Task awaiter that has affinity to executing callbacks synchronously on the completing callstack.
         /// </summary>
         [SuppressMessage("Microsoft.Performance", "CA1815:OverrideEqualsAndOperatorEqualsOnValueTypes", Justification = "Awaitables are not compared.")]
-        public readonly struct ExecuteContinuationSynchronouslyAwaiter : INotifyCompletion
+        public readonly struct ExecuteContinuationSynchronouslyAwaiter : ICriticalNotifyCompletion
         {
             /// <summary>
             /// The task whose completion will execute the continuation.
@@ -609,6 +622,26 @@ public void OnCompleted(Action continuation)
                     TaskContinuationOptions.ExecuteSynchronously,
                     TaskScheduler.Default);
             }
+
+            /// <inheritdoc cref="OnCompleted(Action)"/>
+            public void UnsafeOnCompleted(Action continuation)
+            {
+#if NETFRAMEWORK // Only bother suppressing flow on .NET Framework where the perf would improve from doing so.
+                if (ExecutionContext.IsFlowSuppressed())
+                {
+                    this.OnCompleted(continuation);
+                }
+                else
+                {
+                    using (ExecutionContext.SuppressFlow())
+                    {
+                        this.OnCompleted(continuation);
+                    }
+                }
+#else
+                this.OnCompleted(continuation);
+#endif
+            }
         }
 
         /// <summary>
@@ -645,7 +678,7 @@ public ExecuteContinuationSynchronouslyAwaitable(Task<T> antecedent)
         /// </summary>
         /// <typeparam name="T">The type of value returned by the awaited <see cref="Task"/>.</typeparam>
         [SuppressMessage("Microsoft.Performance", "CA1815:OverrideEqualsAndOperatorEqualsOnValueTypes", Justification = "Awaitables are not compared.")]
-        public readonly struct ExecuteContinuationSynchronouslyAwaiter<T> : INotifyCompletion
+        public readonly struct ExecuteContinuationSynchronouslyAwaiter<T> : ICriticalNotifyCompletion
         {
             /// <summary>
             /// The task whose completion will execute the continuation.
@@ -687,6 +720,26 @@ public void OnCompleted(Action continuation)
                     TaskContinuationOptions.ExecuteSynchronously,
                     TaskScheduler.Default);
             }
+
+            /// <inheritdoc cref="OnCompleted(Action)"/>
+            public void UnsafeOnCompleted(Action continuation)
+            {
+#if NETFRAMEWORK // Only bother suppressing flow on .NET Framework where the perf would improve from doing so.
+                if (ExecutionContext.IsFlowSuppressed())
+                {
+                    this.OnCompleted(continuation);
+                }
+                else
+                {
+                    using (ExecutionContext.SuppressFlow())
+                    {
+                        this.OnCompleted(continuation);
+                    }
+                }
+#else
+                this.OnCompleted(continuation);
+#endif
+            }
         }
 
         /// <summary>

diff --git a/src/Microsoft.VisualStudio.Threading/net472/PublicAPI.Unshipped.txt b/src/Microsoft.VisualStudio.Threading/net472/PublicAPI.Unshipped.txt
@@ -8,4 +8,6 @@ Microsoft.VisualStudio.Threading.AwaitExtensions.AggregateExceptionAwaiter.GetRe
 Microsoft.VisualStudio.Threading.AwaitExtensions.AggregateExceptionAwaiter.IsCompleted.get -> bool
 Microsoft.VisualStudio.Threading.AwaitExtensions.AggregateExceptionAwaiter.OnCompleted(System.Action! continuation) -> void
 Microsoft.VisualStudio.Threading.AwaitExtensions.AggregateExceptionAwaiter.UnsafeOnCompleted(System.Action! continuation) -> void
+Microsoft.VisualStudio.Threading.AwaitExtensions.ExecuteContinuationSynchronouslyAwaiter.UnsafeOnCompleted(System.Action! continuation) -> void
+Microsoft.VisualStudio.Threading.AwaitExtensions.ExecuteContinuationSynchronouslyAwaiter<T>.UnsafeOnCompleted(System.Action! continuation) -> void
 static Microsoft.VisualStudio.Threading.AwaitExtensions.ConfigureAwaitForAggregateException(this System.Threading.Tasks.Task! task, bool continueOnCapturedContext = true) -> Microsoft.VisualStudio.Threading.AwaitExtensions.AggregateExceptionAwaitable
diff --git a/src/Microsoft.VisualStudio.Threading/netcoreapp3.0/PublicAPI.Unshipped.txt b/src/Microsoft.VisualStudio.Threading/netcoreapp3.0/PublicAPI.Unshipped.txt
@@ -8,4 +8,6 @@ Microsoft.VisualStudio.Threading.AwaitExtensions.AggregateExceptionAwaiter.GetRe
 Microsoft.VisualStudio.Threading.AwaitExtensions.AggregateExceptionAwaiter.IsCompleted.get -> bool
 Microsoft.VisualStudio.Threading.AwaitExtensions.AggregateExceptionAwaiter.OnCompleted(System.Action! continuation) -> void
 Microsoft.VisualStudio.Threading.AwaitExtensions.AggregateExceptionAwaiter.UnsafeOnCompleted(System.Action! continuation) -> void
+Microsoft.VisualStudio.Threading.AwaitExtensions.ExecuteContinuationSynchronouslyAwaiter.UnsafeOnCompleted(System.Action! continuation) -> void
+Microsoft.VisualStudio.Threading.AwaitExtensions.ExecuteContinuationSynchronouslyAwaiter<T>.UnsafeOnCompleted(System.Action! continuation) -> void
 static Microsoft.VisualStudio.Threading.AwaitExtensions.ConfigureAwaitForAggregateException(this System.Threading.Tasks.Task! task, bool continueOnCapturedContext = true) -> Microsoft.VisualStudio.Threading.AwaitExtensions.AggregateExceptionAwaitable
diff --git a/src/Microsoft.VisualStudio.Threading/netstandard2.0/PublicAPI.Unshipped.txt b/src/Microsoft.VisualStudio.Threading/netstandard2.0/PublicAPI.Unshipped.txt
@@ -8,4 +8,6 @@ Microsoft.VisualStudio.Threading.AwaitExtensions.AggregateExceptionAwaiter.GetRe
 Microsoft.VisualStudio.Threading.AwaitExtensions.AggregateExceptionAwaiter.IsCompleted.get -> bool
 Microsoft.VisualStudio.Threading.AwaitExtensions.AggregateExceptionAwaiter.OnCompleted(System.Action! continuation) -> void
 Microsoft.VisualStudio.Threading.AwaitExtensions.AggregateExceptionAwaiter.UnsafeOnCompleted(System.Action! continuation) -> void
+Microsoft.VisualStudio.Threading.AwaitExtensions.ExecuteContinuationSynchronouslyAwaiter.UnsafeOnCompleted(System.Action! continuation) -> void
+Microsoft.VisualStudio.Threading.AwaitExtensions.ExecuteContinuationSynchronouslyAwaiter<T>.UnsafeOnCompleted(System.Action! continuation) -> void
 static Microsoft.VisualStudio.Threading.AwaitExtensions.ConfigureAwaitForAggregateException(this System.Threading.Tasks.Task! task, bool continueOnCapturedContext = true) -> Microsoft.VisualStudio.Threading.AwaitExtensions.AggregateExceptionAwaitable