Avoid AwaitTaskContinuation allocation in some awaits (dotnet/coreclr#20159)

In .NET Core 2.1, I added a bunch of optimizations to async methods that are based on reusing the async state machine object itself for other purposes in order to avoid related allocations.  One of those optimizations was using the boxed state machine itself as the continuation object that could be queued onto a Task, and in the common case where the continuation could be executed synchronously, there would then not be any further allocations.  However, if the continuation needed to be run asynchronously (e.g. because the Task required it via RunContinuationsAsynchronously), the code would allocate a new work item object and queue that to the thread pool to execute.  This then also forced the state machine object to lazily allocate the Action delegate for its MoveNext method. This PR extends the system slightly to also cover that asynchronous execution case, by making the state machine box itself being queueable to the thread pool.  In doing so, it avoids that AwaitTaskContinuation allocation and also avoids forcing the delegate into existence. (As is the case for other optimizations, this one is only employed when ETW logging isn't enabled; if it is enabled, we need to flow more information, and enabling that would penalize the non-logging case.)

Commit migrated from https://github.com/dotnet/coreclr/commit/03ead514a3aa414b26246a4160d0593811fd1044

diff --git a/src/coreclr/src/System.Private.CoreLib/src/System/Runtime/CompilerServices/AsyncMethodBuilder.cs b/src/coreclr/src/System.Private.CoreLib/src/System/Runtime/CompilerServices/AsyncMethodBuilder.cs
index 18888f9..2fb284a 100644 (file)
--- a/src/coreclr/src/System.Private.CoreLib/src/System/Runtime/CompilerServices/AsyncMethodBuilder.cs
+++ b/src/coreclr/src/System.Private.CoreLib/src/System/Runtime/CompilerServices/AsyncMethodBuilder.cs
@@ -534,6 +534,10 @@ namespace System.Runtime.CompilerServices
             /// <summary>A delegate to the <see cref="MoveNext"/> method.</summary>
             public Action MoveNextAction => _moveNextAction ?? (_moveNextAction = new Action(MoveNext));
 
+            /// <summary>Invokes <see cref="MoveNext"/> when the box is queued to the thread pool and executed as a work item.</summary>
+            void IThreadPoolWorkItem.ExecuteWorkItem() => MoveNext();
+            void IThreadPoolWorkItem.MarkAborted(ThreadAbortException tae) { /* nop */ }
+
             /// <summary>Calls MoveNext on <see cref="StateMachine"/></summary>
             public void MoveNext()
             {
@@ -884,7 +888,7 @@ namespace System.Runtime.CompilerServices
     /// <summary>
     /// An interface implemented by all <see cref="AsyncStateMachineBox{TStateMachine, TResult}"/> instances, regardless of generics.
     /// </summary>
-    internal interface IAsyncStateMachineBox
+    internal interface IAsyncStateMachineBox : IThreadPoolWorkItem
     {
         /// <summary>Move the state machine forward.</summary>
         void MoveNext();

diff --git a/src/coreclr/src/System.Private.CoreLib/src/System/Threading/Tasks/TaskContinuation.cs b/src/coreclr/src/System.Private.CoreLib/src/System/Threading/Tasks/TaskContinuation.cs
index dc63a67..7abe63c 100644 (file)
--- a/src/coreclr/src/System.Private.CoreLib/src/System/Threading/Tasks/TaskContinuation.cs
+++ b/src/coreclr/src/System.Private.CoreLib/src/System/Threading/Tasks/TaskContinuation.cs
@@ -787,7 +787,23 @@ namespace System.Threading.Tasks
             // If we're not allowed to run here, schedule the action
             if (!allowInlining || !IsValidLocationForInlining)
             {
-                UnsafeScheduleAction(box.MoveNextAction, prevCurrentTask);
+                // If logging is disabled, we can simply queue the box itself as a custom work
+                // item, and its work item execution will just invoke its MoveNext.  However, if
+                // logging is enabled, there is pre/post-work we need to do around logging to
+                // match what's done for other continuations, and that requires flowing additional
+                // information into the continuation, which we don't want to burden other cases of the
+                // box with... so, in that case we just delegate to the AwaitTaskContinuation-based
+                // path that already handles this, albeit at the expense of allocating the ATC
+                // object, and potentially forcing the box's delegate into existence, when logging
+                // is enabled.
+                if (TplEtwProvider.Log.IsEnabled())
+                {
+                    UnsafeScheduleAction(box.MoveNextAction, prevCurrentTask);
+                }
+                else
+                {
+                    ThreadPool.UnsafeQueueCustomWorkItem(box, forceGlobal: false);
+                }
                 return;
             }