Reduce Execution Context Save+Restore (dotnet/coreclr#15629)

* Reduce Save+Restore for ExecutionContext

* Use flag rather than comparison to static

* Skip null check for pre-checked EC.Run

* Feedback

* Add static helper lookup for default context for TP

* Add note for enregistering

* Return to Default context when no values

* Default + FlowSuppressed Context

* Move AsyncMethodBuilder.Start to static non-generic

* Feedback

Signed-off-by: dotnet-bot-corefx-mirror <[email protected]>

src/Common/src/CoreLib/System.Private.CoreLib.Shared.projitems

@@ -363,6 +363,7 @@
     <Compile Include="$(MSBuildThisFileDirectory)System\Resources\SatelliteContractVersionAttribute.cs" />
     <Compile Include="$(MSBuildThisFileDirectory)System\Resources\UltimateResourceFallbackLocation.cs" />
     <Compile Include="$(MSBuildThisFileDirectory)System\Runtime\CompilerServices\AccessedThroughPropertyAttribute.cs" />
+    <Compile Include="$(MSBuildThisFileDirectory)System\Runtime\CompilerServices\AsyncMethodBuilder.cs" />
     <Compile Include="$(MSBuildThisFileDirectory)System\Runtime\CompilerServices\AsyncMethodBuilderAttribute.cs" />
     <Compile Include="$(MSBuildThisFileDirectory)System\Runtime\CompilerServices\AsyncStateMachineAttribute.cs" />
     <Compile Include="$(MSBuildThisFileDirectory)System\Runtime\CompilerServices\AsyncValueTaskMethodBuilder.cs" />

src/Common/src/CoreLib/System/Runtime/CompilerServices/AsyncMethodBuilder.cs

@@ -0,0 +1,68 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+// See the LICENSE file in the project root for more information.
+
+using System;
+using System.Diagnostics;
+using System.Threading;
+
+namespace System.Runtime.CompilerServices
+{
+    internal static partial class AsyncMethodBuilder
+    {
+        /// <summary>Initiates the builder's execution with the associated state machine.</summary>
+        /// <typeparam name="TStateMachine">Specifies the type of the state machine.</typeparam>
+        /// <param name="stateMachine">The state machine instance, passed by reference.</param>
+        [DebuggerStepThrough]
+        public static void Start<TStateMachine>(ref TStateMachine stateMachine) where TStateMachine : IAsyncStateMachine
+        {
+            if (stateMachine == null) // TStateMachines are generally non-nullable value types, so this check will be elided
+            {
+                ThrowHelper.ThrowArgumentNullException(ExceptionArgument.stateMachine);
+            }
+
+            // enregistrer variables with 0 post-fix so they can be used in registers without EH forcing them to stack
+            // Capture references to Thread Contexts
+            Thread currentThread0 = Thread.CurrentThread;
+            Thread currentThread = currentThread0;
+            ExecutionContext previousExecutionCtx0 = currentThread0.ExecutionContext;
+
+            // Store current ExecutionContext and SynchronizationContext as "previousXxx".
+            // This allows us to restore them and undo any Context changes made in stateMachine.MoveNext
+            // so that they won't "leak" out of the first await.
+            ExecutionContext previousExecutionCtx = previousExecutionCtx0;
+            SynchronizationContext previousSyncCtx = currentThread0.SynchronizationContext;
+
+            try
+            {
+                stateMachine.MoveNext();
+            }
+            finally
+            {
+                // Re-enregistrer variables post EH with 1 post-fix so they can be used in registers rather than from stack
+                SynchronizationContext previousSyncCtx1 = previousSyncCtx;
+                Thread currentThread1 = currentThread;
+                // The common case is that these have not changed, so avoid the cost of a write barrier if not needed.
+                if (previousSyncCtx1 != currentThread1.SynchronizationContext)
+                {
+                    // Restore changed SynchronizationContext back to previous
+                    currentThread1.SynchronizationContext = previousSyncCtx1;
+                }
+
+                ExecutionContext previousExecutionCtx1 = previousExecutionCtx;
+                ExecutionContext currentExecutionCtx1 = currentThread1.ExecutionContext;
+                if (previousExecutionCtx1 != currentExecutionCtx1)
+                {
+                    // Restore changed ExecutionContext back to previous
+                    currentThread1.ExecutionContext = previousExecutionCtx1;
+                    if ((currentExecutionCtx1 != null && currentExecutionCtx1.HasChangeNotifications) ||
+                        (previousExecutionCtx1 != null && previousExecutionCtx1.HasChangeNotifications))
+                    {
+                        // There are change notifications; trigger any affected
+                        ExecutionContext.OnValuesChanged(currentExecutionCtx1, previousExecutionCtx1);
+                    }
+                }
+            }
+        }
+    }
+}

src/Common/src/CoreLib/System/Runtime/CompilerServices/AsyncValueTaskMethodBuilder.cs

@@ -38,8 +38,9 @@ public static AsyncValueTaskMethodBuilder<TResult> Create() =>
         /// <summary>Begins running the builder with the associated state machine.</summary>
         /// <typeparam name="TStateMachine">The type of the state machine.</typeparam>
         /// <param name="stateMachine">The state machine instance, passed by reference.</param>
+        [MethodImpl(MethodImplOptions.AggressiveInlining)]
         public void Start<TStateMachine>(ref TStateMachine stateMachine) where TStateMachine : IAsyncStateMachine =>
-            _methodBuilder.Start(ref stateMachine); // will provide the right ExecutionContext semantics
+            AsyncMethodBuilder.Start(ref stateMachine); // will provide the right ExecutionContext semantics
 
         /// <summary>Associates the builder with the specified state machine.</summary>
         /// <param name="stateMachine">The state machine instance to associate with the builder.</param>

src/Common/src/CoreLib/System/Threading/AsyncLocal.cs

@@ -125,7 +125,7 @@ internal static class AsyncLocalValueMap
         public static IAsyncLocalValueMap Empty { get; } = new EmptyAsyncLocalValueMap();
 
         // Instance without any key/value pairs.  Used as a singleton/
-        private sealed class EmptyAsyncLocalValueMap : IAsyncLocalValueMap
+        internal sealed class EmptyAsyncLocalValueMap : IAsyncLocalValueMap
         {
             public IAsyncLocalValueMap Set(IAsyncLocal key, object value)
             {
@@ -144,7 +144,7 @@ public bool TryGetValue(IAsyncLocal key, out object value)
         }
 
         // Instance with one key/value pair.
-        private sealed class OneElementAsyncLocalValueMap : IAsyncLocalValueMap
+        internal sealed class OneElementAsyncLocalValueMap : IAsyncLocalValueMap
         {
             private readonly IAsyncLocal _key1;
             private readonly object _value1;