Use ExecutionContext.Restore rather than EC.Run callback

Author: benaadams

src/libraries/System.Private.CoreLib/src/System/Threading/ExecutionContext.cs

@@ -200,74 +200,27 @@ internal static void RunInternal(ExecutionContext? executionContext, ContextCall
             edi?.Throw();
         }
 
-        // Direct copy of the above RunInternal overload, except that it passes the state into the callback strongly-typed and by ref.
-        internal static void RunInternal<TState>(ExecutionContext? executionContext, ContextCallback<TState> callback, ref TState state)
+        internal static void Restore(ExecutionContext? executionContext)
         {
-            // Note: ExecutionContext.RunInternal is an extremely hot function and used by every await, ThreadPool execution, etc.
-            // Note: Manual enregistering may be addressed by "Exception Handling Write Through Optimization"
-            //       https://github.com/dotnet/runtime/blob/master/docs/design/features/eh-writethru.md
+            Thread currentThread = Thread.CurrentThread;
 
-            // Enregister 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;
-            if (previousExecutionCtx0 != null && previousExecutionCtx0.m_isDefault)
+            ExecutionContext? currentExecutionCtx = currentThread._executionContext;
+            if (currentExecutionCtx != null && currentExecutionCtx.m_isDefault)
             {
                 // Default is a null ExecutionContext internally
-                previousExecutionCtx0 = null;
+                currentExecutionCtx = null;
             }
 
-            // Store current ExecutionContext and SynchronizationContext as "previousXxx".
-            // This allows us to restore them and undo any Context changes made in callback.Invoke
-            // so that they won't "leak" back into caller.
-            // These variables will cross EH so be forced to stack
-            ExecutionContext? previousExecutionCtx = previousExecutionCtx0;
-            SynchronizationContext? previousSyncCtx = currentThread0._synchronizationContext;
-
             if (executionContext != null && executionContext.m_isDefault)
             {
                 // Default is a null ExecutionContext internally
                 executionContext = null;
             }
 
-            if (previousExecutionCtx0 != executionContext)
+            if (currentExecutionCtx != executionContext)
             {
-                RestoreChangedContextToThread(currentThread0, executionContext, previousExecutionCtx0);
+                RestoreChangedContextToThread(currentThread, executionContext, currentExecutionCtx);
             }
-
-            ExceptionDispatchInfo? edi = null;
-            try
-            {
-                callback.Invoke(ref state);
-            }
-            catch (Exception ex)
-            {
-                // Note: we have a "catch" rather than a "finally" because we want
-                // to stop the first pass of EH here.  That way we can restore the previous
-                // context before any of our callers' EH filters run.
-                edi = ExceptionDispatchInfo.Capture(ex);
-            }
-
-            // 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 (currentThread1._synchronizationContext != previousSyncCtx1)
-            {
-                // Restore changed SynchronizationContext back to previous
-                currentThread1._synchronizationContext = previousSyncCtx1;
-            }
-
-            ExecutionContext? previousExecutionCtx1 = previousExecutionCtx;
-            ExecutionContext? currentExecutionCtx1 = currentThread1._executionContext;
-            if (currentExecutionCtx1 != previousExecutionCtx1)
-            {
-                RestoreChangedContextToThread(currentThread1, previousExecutionCtx1, currentExecutionCtx1);
-            }
-
-            // If exception was thrown by callback, rethrow it now original contexts are restored
-            edi?.Throw();
         }
 
         internal static void RunFromThreadPoolDispatchLoop(Thread threadPoolThread, ExecutionContext executionContext, ContextCallback callback, object state)

src/libraries/System.Private.CoreLib/src/System/Threading/Tasks/Sources/ManualResetValueTaskSourceCore.cs

@@ -207,51 +207,104 @@ private void SignalCompletion()
             }
             _completed = true;
 
-            if (_continuation != null || Interlocked.CompareExchange(ref _continuation, ManualResetValueTaskSourceCoreShared.s_sentinel, null) != null)
+            if (_continuation is null && Interlocked.CompareExchange(ref _continuation, ManualResetValueTaskSourceCoreShared.s_sentinel, null) is null)
             {
-                if (_executionContext != null)
+                return;
+            }
+
+            if (_executionContext is null)
+            {
+                if (_capturedContext is null)
+                {
+                    if (RunContinuationsAsynchronously)
+                    {
+                        ThreadPool.UnsafeQueueUserWorkItem(_continuation, _continuationState, preferLocal: true);
+                    }
+                    else
+                    {
+                        _continuation(_continuationState);
+                    }
+
+                    return;
+                }
+
+                InvokeSchedulerContinuation();
+                return;
+            }
+
+            ExecutionContext? currentContext = ExecutionContext.Capture();
+            // Restore the captured ExecutionContext before executing anything.
+            ExecutionContext.Restore(_executionContext);
+
+            if (_capturedContext is null)
+            {
+                if (RunContinuationsAsynchronously)
                 {
-                    ExecutionContext.RunInternal(
-                        _executionContext,
-                        (ref ManualResetValueTaskSourceCore<TResult> s) => s.InvokeContinuation(),
-                        ref this);
+                    ThreadPool.QueueUserWorkItem(_continuation, _continuationState, preferLocal: true);
+                    // Restore the current ExecutionContext.
+                    ExecutionContext.Restore(currentContext);
                 }
                 else
                 {
-                    InvokeContinuation();
+                    // Running inline may throw; capture the edi if it does as we changed the ExecutionContext,
+                    // so need to restore it back before propagating the throw.
+                    ExceptionDispatchInfo? edi = InvokeInlineContinuation();
+                    // Restore the current ExecutionContext.
+                    ExecutionContext.Restore(currentContext);
+                    // Now rethrow the exception; if there is one.
+                    edi?.Throw();
                 }
+
+                return;
             }
+
+            InvokeSchedulerContinuation();
+            // Restore the current ExecutionContext.
+            ExecutionContext.Restore(currentContext);
         }
 
         /// <summary>
-        /// Invokes the continuation with the appropriate captured context / scheduler.
-        /// This assumes that if <see cref="_executionContext"/> is not null we're already
+        /// Invokes the continuation inline and captures any exception thrown.
+        /// This assumes that if <see cref="_continuation"/> is not null we're already
         /// running within that <see cref="ExecutionContext"/>.
         /// </summary>
-        private void InvokeContinuation()
+        private ExceptionDispatchInfo? InvokeInlineContinuation()
         {
             Debug.Assert(_continuation != null);
+            Debug.Assert(_capturedContext == null);
+            Debug.Assert(!RunContinuationsAsynchronously);
 
-            switch (_capturedContext)
+            ExceptionDispatchInfo? edi = null;
+            SynchronizationContext? syncContext = SynchronizationContext.Current;
+            try
             {
-                case null:
-                    if (RunContinuationsAsynchronously)
-                    {
-                        if (_executionContext != null)
-                        {
-                            ThreadPool.QueueUserWorkItem(_continuation, _continuationState, preferLocal: true);
-                        }
-                        else
-                        {
-                            ThreadPool.UnsafeQueueUserWorkItem(_continuation, _continuationState, preferLocal: true);
-                        }
-                    }
-                    else
-                    {
-                        _continuation(_continuationState);
-                    }
-                    break;
+                _continuation(_continuationState);
+            }
+            catch (Exception ex)
+            {
+                // Note: we have a "catch" rather than a "finally" because we want
+                // to stop the first pass of EH here.  That way we can restore the previous
+                // context before any of our callers' EH filters run.
+                edi = ExceptionDispatchInfo.Capture(ex);
+            }
+
+            // Set sync context back to what it was prior to coming in
+            SynchronizationContext.SetSynchronizationContext(syncContext);
+            return edi;
+        }
+
+        /// <summary>
+        /// Invokes the continuation with the appropriate scheduler.
+        /// This assumes that if <see cref="_continuation"/> is not null we're already
+        /// running within that <see cref="ExecutionContext"/>.
+        /// </summary>
+        private void InvokeSchedulerContinuation()
+        {
+            Debug.Assert(_capturedContext != null);
+            Debug.Assert(_continuation != null);
 
+            switch (_capturedContext)
+            {
                 case SynchronizationContext sc:
                     sc.Post(s =>
                     {