Fast-path ExecutionContext for ThreadPool items (dotnet#20308)

Fast-path ExecutionContext for ThreadPool items

Maintain the ThreadPool threads on the Default contexts between work items.

Always restore the Default context on the ThreadPool Dispatch loop after a workitem has run (to clean up any ExecutionContext leakage from changes on flow suppressed workitems, or AsyncLocal change eventhandlers; as well as firing their notifications if they have them)

Store the `CurrentThread` as part of the thread static `ThreadPoolWorkQueueThreadLocals` which are already looked up at the start of the Dispatch loop to avoid an additional lookup via `Thread.CurrentThread`.

As workitems are started on the Default context and are returned to it `QueueUserWorkItemCallbackDefaultContext` items can just be run their callbacks directly rather than via `ExecutionContext.Run` (as the EC don't need to move to Default and is always moved back to Default).

As `QueueUserWorkItemCallbackDefaultContext` now runs items directly; flow suppressed callbacks can use the smaller `QueueUserWorkItemCallbackDefaultContext` rather than `QueueUserWorkItemCallback` with a null context; and handling for flow suppression can be removed from `QueueUserWorkItemCallback`.

As `AwaitTaskContinuation`'s `IThreadPoolWorkItem.Execute` doesn't preform additional work after it completes, it can run `m_action` directly for Default context in addition to the flow suppressed context, rather than going via `ExecutionContext.Run`.

Given that the items on the ThreadPool are always started on the threadpool and restored to it; we can introduce some faster paths than `ExecutionContext:RunInternal` (328 bytes asm).

Introduce `ExecutionContext:RunForThreadPoolUnsafe` (71 bytes asm), for `IThreadPoolWorkItem`s where they need to run on a provided context, but do not need to execute anything after they complete so can rely on the Dispatch loop restore. This includes `QueueUserWorkItemCallback`, `QueueUserWorkItemCallback<TState>` and `AwaitTaskContinuation`.

Introduce `ExecutionContext:RunFromThreadPoolDispatchLoop` (225 bytes asm), for items run from the ThreadPool, so don't need to capture current context (Default) to restore later, however need to do need to restore back to Default after execution as they then perform additional work. This includes 
`Task`/`AsyncStateMachineBox`/`AwaitTaskContinuation`/`Timer`. 

Change `Task.ExecuteFromThreadPool()` to take the thread `Task.ExecuteFromThreadPool(Thread threadPoolThread)` from the ThreadPool Dispatch loop so it can pass it into the `ExecutionContext:RunFromThreadPoolDispatchLoop` overload and avoid the `Thread.CurrentThread` lookup.

Perf test: dotnet#20308 (comment)

Resolves: dotnet/corefx#32695

Author: benaadams

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

@@ -12,6 +12,7 @@
 ===========================================================*/
 
 using System.Diagnostics;
+using System.Runtime.CompilerServices;
 using System.Runtime.ExceptionServices;
 using System.Runtime.Serialization;
 
@@ -136,6 +137,11 @@ internal static void RunInternal(ExecutionContext executionContext, ContextCallb
             Thread currentThread0 = Thread.CurrentThread;
             Thread currentThread = currentThread0;
             ExecutionContext previousExecutionCtx0 = currentThread0.ExecutionContext;
+            if (previousExecutionCtx0 != null && previousExecutionCtx0.m_isDefault)
+            {
+                // Default is a null ExecutionContext internally
+                previousExecutionCtx0 = null;
+            }
 
             // Store current ExecutionContext and SynchronizationContext as "previousXxx".
             // This allows us to restore them and undo any Context changes made in callback.Invoke
@@ -215,6 +221,11 @@ internal static void RunInternal<TState>(ExecutionContext executionContext, Cont
             Thread currentThread0 = Thread.CurrentThread;
             Thread currentThread = currentThread0;
             ExecutionContext previousExecutionCtx0 = currentThread0.ExecutionContext;
+            if (previousExecutionCtx0 != null && previousExecutionCtx0.m_isDefault)
+            {
+                // Default is a null ExecutionContext internally
+                previousExecutionCtx0 = null;
+            }
 
             // Store current ExecutionContext and SynchronizationContext as "previousXxx".
             // This allows us to restore them and undo any Context changes made in callback.Invoke
@@ -282,6 +293,115 @@ internal static void RunInternal<TState>(ExecutionContext executionContext, Cont
             edi?.Throw();
         }
 
+        internal static void RunFromThreadPoolDispatchLoop(Thread threadPoolThread, ExecutionContext executionContext, ContextCallback callback, object state)
+        {
+            Debug.Assert(threadPoolThread == Thread.CurrentThread);
+            CheckThreadPoolAndContextsAreDefault();
+            // ThreadPool starts on Default Context so we don't need to save the "previous" state as we know it is Default (null)
+
+            if (executionContext != null && executionContext.m_isDefault)
+            {
+                // Default is a null ExecutionContext internally
+                executionContext = null;
+            }
+            else if (executionContext != null)
+            {
+                // Non-Default context to restore
+                threadPoolThread.ExecutionContext = executionContext;
+                if (executionContext.HasChangeNotifications)
+                {
+                    // There are change notifications; trigger any affected
+                    OnValuesChanged(previousExecutionCtx: null, executionContext);
+                }
+            }
+
+            ExceptionDispatchInfo edi = null;
+            try
+            {
+                callback.Invoke(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);
+            }
+
+            // Enregister threadPoolThread as it crossed EH, and use enregistered variable
+            Thread currentThread = threadPoolThread;
+
+            ExecutionContext currentExecutionCtx = currentThread.ExecutionContext;
+
+            // Restore changed SynchronizationContext back to Default
+            currentThread.SynchronizationContext = null;
+            if (currentExecutionCtx != null)
+            {
+                // The EC always needs to be reset for this overload, as it will flow back to the caller if it performs 
+                // extra work prior to returning to the Dispatch loop. For example for Task-likes it will flow out of await points
+
+                // Restore to Default before Notifications, as the change can be observed in the handler.
+                currentThread.ExecutionContext = null;
+                if (currentExecutionCtx.HasChangeNotifications)
+                {
+                    // There are change notifications; trigger any affected
+                    OnValuesChanged(currentExecutionCtx, nextExecutionCtx: null);
+                }
+            }
+
+            // If exception was thrown by callback, rethrow it now original contexts are restored
+            edi?.Throw();
+        }
+
+        internal static void RunForThreadPoolUnsafe<TState>(ExecutionContext executionContext, Action<TState> callback, in TState state)
+        {
+            // We aren't running in try/catch as if an exception is directly thrown on the ThreadPool either process 
+            // will crash or its a ThreadAbortException. 
+
+            CheckThreadPoolAndContextsAreDefault();
+            Debug.Assert(executionContext != null && !executionContext.m_isDefault, "ExecutionContext argument is Default.");
+
+            Thread currentThread = Thread.CurrentThread;
+            // Restore Non-Default context
+            currentThread.ExecutionContext = executionContext;
+            if (executionContext.HasChangeNotifications)
+            {
+                OnValuesChanged(previousExecutionCtx: null, executionContext);
+            }
+
+            callback.Invoke(state);
+
+            // ThreadPoolWorkQueue.Dispatch will handle notifications and reset EC and SyncCtx back to default
+        }
+
+        // Inline as only called in one place and always called
+        [MethodImpl(MethodImplOptions.AggressiveInlining)]
+        internal static void ResetThreadPoolThread(Thread currentThread)
+        {
+            ExecutionContext currentExecutionCtx = currentThread.ExecutionContext;
+
+            // Reset to defaults
+            currentThread.SynchronizationContext = null;
+            currentThread.ExecutionContext = null;
+
+            if (currentExecutionCtx != null && currentExecutionCtx.HasChangeNotifications)
+            {
+                OnValuesChanged(currentExecutionCtx, nextExecutionCtx: null);
+
+                // Reset to defaults again without change notifications in case the Change handler changed the contexts
+                currentThread.SynchronizationContext = null;
+                currentThread.ExecutionContext = null;
+            }
+        }
+
+        [System.Diagnostics.Conditional("DEBUG")]
+        internal static void CheckThreadPoolAndContextsAreDefault()
+        {
+            Debug.Assert(Thread.CurrentThread.IsThreadPoolThread);
+            Debug.Assert(Thread.CurrentThread.ExecutionContext == null, "ThreadPool thread not on Default ExecutionContext.");
+            Debug.Assert(Thread.CurrentThread.SynchronizationContext == null, "ThreadPool thread not on Default SynchronizationContext.");
+        }
+
         internal static void OnValuesChanged(ExecutionContext previousExecutionCtx, ExecutionContext nextExecutionCtx)
         {
             Debug.Assert(previousExecutionCtx != nextExecutionCtx);

src/System.Private.CoreLib/shared/System/Threading/SemaphoreSlim.cs

@@ -80,7 +80,7 @@ private sealed class TaskNode : Task<bool>
             internal TaskNode Prev, Next;
             internal TaskNode() : base() { }
 
-            internal override void ExecuteFromThreadPool()
+            internal override void ExecuteFromThreadPool(Thread threadPoolThread)
             {
                 bool setSuccessfully = TrySetResult(true);
                 Debug.Assert(setSuccessfully, "Should have been able to complete task");

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

@@ -540,10 +540,12 @@ private class AsyncStateMachineBox<TStateMachine> : // SOS DumpAsync command dep
             /// <summary>A delegate to the <see cref="MoveNext"/> method.</summary>
             public Action MoveNextAction => _moveNextAction ?? (_moveNextAction = new Action(MoveNext));
 
-            internal sealed override void ExecuteFromThreadPool() => MoveNext();
+            internal sealed override void ExecuteFromThreadPool(Thread threadPoolThread) => MoveNext(threadPoolThread);
 
             /// <summary>Calls MoveNext on <see cref="StateMachine"/></summary>
-            public void MoveNext()
+            public void MoveNext() => MoveNext(threadPoolThread: null);
+
+            private void MoveNext(Thread threadPoolThread)
             {
                 Debug.Assert(!IsCompleted);
 
@@ -560,7 +562,14 @@ public void MoveNext()
                 }
                 else
                 {
-                    ExecutionContext.RunInternal(context, s_callback, this);
+                    if (threadPoolThread is null)
+                    {
+                        ExecutionContext.RunInternal(context, s_callback, this);
+                    }
+                    else
+                    {
+                        ExecutionContext.RunFromThreadPoolDispatchLoop(threadPoolThread, context, s_callback, this);
+                    }
                 }
 
                 if (IsCompleted)

src/System.Private.CoreLib/src/System/Threading/Tasks/Task.cs

@@ -2368,16 +2368,16 @@ internal bool ExecuteEntry()
         /// can override to customize their behavior, which is usually done by promises
         /// that want to reuse the same object as a queued work item.
         /// </summary>
-        internal virtual void ExecuteFromThreadPool() => ExecuteEntryUnsafe();
+        internal virtual void ExecuteFromThreadPool(Thread threadPoolThread) => ExecuteEntryUnsafe(threadPoolThread);
 
-        internal void ExecuteEntryUnsafe() // used instead of ExecuteEntry() when we don't have to worry about double-execution prevent
+        internal void ExecuteEntryUnsafe(Thread threadPoolThread) // used instead of ExecuteEntry() when we don't have to worry about double-execution prevent
         {
             // Remember that we started running the task delegate.
             m_stateFlags |= TASK_STATE_DELEGATE_INVOKED;
 
             if (!IsCancellationRequested & !IsCanceled)
             {
-                ExecuteWithThreadLocal(ref t_currentTask);
+                ExecuteWithThreadLocal(ref t_currentTask, threadPoolThread);
             }
             else
             {
@@ -2398,7 +2398,7 @@ internal void ExecuteEntryCancellationRequestedOrCanceled()
         }
 
         // A trick so we can refer to the TLS slot with a byref.
-        private void ExecuteWithThreadLocal(ref Task currentTaskSlot)
+        private void ExecuteWithThreadLocal(ref Task currentTaskSlot, Thread threadPoolThread = null)
         {
             // Remember the current task so we can restore it after running, and then
             Task previousTask = currentTaskSlot;
@@ -2439,7 +2439,14 @@ private void ExecuteWithThreadLocal(ref Task currentTaskSlot)
                     else
                     {
                         // Invoke it under the captured ExecutionContext
-                        ExecutionContext.RunInternal(ec, s_ecCallback, this);
+                        if (threadPoolThread is null)
+                        {
+                            ExecutionContext.RunInternal(ec, s_ecCallback, this);
+                        }
+                        else
+                        {
+                            ExecutionContext.RunFromThreadPoolDispatchLoop(threadPoolThread, ec, s_ecCallback, this);
+                        }
                     }
                 }
                 catch (Exception exn)

src/System.Private.CoreLib/src/System/Threading/Tasks/TaskContinuation.cs

@@ -646,16 +646,20 @@ void IThreadPoolWorkItem.Execute()
                 // We're not inside of a task, so t_currentTask doesn't need to be specially maintained.
                 // We're on a thread pool thread with no higher-level callers, so exceptions can just propagate.
 
-                // If there's no execution context, just invoke the delegate.
-                if (context == null)
+                ExecutionContext.CheckThreadPoolAndContextsAreDefault();
+                // If there's no execution context or Default, just invoke the delegate as ThreadPool is on Default context.
+                // We don't have to use ExecutionContext.Run for the Default context here as there is no extra processing after the delegate
+                if (context == null || context.IsDefault)
                 {
                     m_action();
                 }
                 // If there is an execution context, get the cached delegate and run the action under the context.
                 else
                 {
-                    ExecutionContext.RunInternal(context, GetInvokeActionCallback(), m_action);
+                    ExecutionContext.RunForThreadPoolUnsafe(context, s_invokeAction, m_action);
                 }
+
+                // ThreadPoolWorkQueue.Dispatch handles notifications and reset context back to default
             }
             finally
             {
@@ -667,19 +671,11 @@ void IThreadPoolWorkItem.Execute()
         }
 
         /// <summary>Cached delegate that invokes an Action passed as an object parameter.</summary>
-        private static ContextCallback s_invokeActionCallback;
-
-        /// <summary>Runs an action provided as an object parameter.</summary>
-        /// <param name="state">The Action to invoke.</param>
-        private static void InvokeAction(object state) { ((Action)state)(); }
+        private readonly static ContextCallback s_invokeContextCallback = (state) => ((Action)state)();
+        private readonly static Action<Action> s_invokeAction = (action) => action();
 
         [MethodImpl(MethodImplOptions.AggressiveInlining)]
-        protected static ContextCallback GetInvokeActionCallback()
-        {
-            ContextCallback callback = s_invokeActionCallback;
-            if (callback == null) { s_invokeActionCallback = callback = InvokeAction; } // lazily initialize SecurityCritical delegate
-            return callback;
-        }
+        protected static ContextCallback GetInvokeActionCallback() => s_invokeContextCallback;
 
         /// <summary>Runs the callback synchronously with the provided state.</summary>
         /// <param name="callback">The callback to run.</param>

src/System.Private.CoreLib/src/System/Threading/Tasks/ThreadPoolTaskScheduler.cs

@@ -33,7 +33,7 @@ internal ThreadPoolTaskScheduler()
         }
 
         // static delegate for threads allocated to handle LongRunning tasks.
-        private static readonly ParameterizedThreadStart s_longRunningThreadWork = s => ((Task)s).ExecuteEntryUnsafe();
+        private static readonly ParameterizedThreadStart s_longRunningThreadWork = s => ((Task)s).ExecuteEntryUnsafe(threadPoolThread: null);
 
         /// <summary>
         /// Schedules a task to the ThreadPool.
@@ -73,7 +73,7 @@ protected override bool TryExecuteTaskInline(Task task, bool taskWasPreviouslyQu
 
             try
             {
-                task.ExecuteEntryUnsafe(); // handles switching Task.Current etc.
+                task.ExecuteEntryUnsafe(threadPoolThread: null); // handles switching Task.Current etc.
             }
             finally
             {