Avoid delegate/work item allocations when setting async continuation (#22373)

When awaiting a task, there's a race between seeing whether the task has completed (in which case we just continue running synchronously), finding the task hasn't completed (in which case we hook up a continuation), and then by the time we try to hook up the continuation finding the task has already completed.  In that final case, we don't want to just execute the callback synchronously, as we risk a stack dive, so we queue it.  That queueing currently entails two allocations in the common case: one for the work item object, and one for the Action delegate we force into existence for the state machine box's MoveNext method (in the common case it's now never allocated because you only await Tasks and ValueTasks known to the runtime, which bypasses its creation).  We can instead just queue the box itself, and avoid both allocations.

1 files changed, 11 insertions, 3 deletions

diff --git a/src/System.Private.CoreLib/shared/System/Threading/Tasks/Task.cs b/src/System.Private.CoreLib/shared/System/Threading/Tasks/Task.cs
index 61b83086a1..d6b5da2b6e 100644
--- a/src/System.Private.CoreLib/shared/System/Threading/Tasks/Task.cs
+++ b/src/System.Private.CoreLib/shared/System/Threading/Tasks/Task.cs
@@ -2618,6 +2618,13 @@ namespace System.Threading.Tasks
         {
             Debug.Assert(stateMachineBox != null);
 
+            // This code path doesn't emit all expected TPL-related events, such as for continuations.
+            // It's expected that all callers check whether events are enabled before calling this function,
+            // and only call it if they're not, so we assert. However, as events can be dynamically turned
+            // on and off, it's possible this assert could fire even when used correctly.  If it becomes
+            // noisy, it can be deleted.
+            Debug.Assert(!TplEventSource.Log.IsEnabled());
+
             // If the caller wants to continue on the current context/scheduler and there is one,
             // fall back to using the state machine's delegate.
             if (continueOnCapturedContext)
@@ -2647,11 +2654,12 @@ namespace System.Threading.Tasks
                 }
             }
 
-            // Otherwise, add the state machine box directly as the ITaskCompletionAction continuation.
-            // If we're unable to because the task has already completed, queue the delegate.
+            // Otherwise, add the state machine box directly as the continuation.
+            // If we're unable to because the task has already completed, queue it.
             if (!AddTaskContinuation(stateMachineBox, addBeforeOthers: false))
             {
-                AwaitTaskContinuation.UnsafeScheduleAction(stateMachineBox.MoveNextAction, this);
+                Debug.Assert(stateMachineBox is Task, "Every state machine box should derive from Task");
+                ThreadPool.UnsafeQueueUserWorkItemInternal(stateMachineBox, preferLocal: true);
             }
         }