Add YieldAwaiter support to the async method builder delegate optimization (#17441)

We added an optimization to async methods that lets the builder avoid allocating the Action delegate when it recognizes the awaiter being used.  Previously this was enabled for all of the publicly exposed awaiters in corelib, with the exception of YieldAwaiter (what's used with Task.YIeld).  It was enabled by having the awaiter implement an internal interface.  This commit just generalizes that interface name and then implements it on YIeldAwaiter to complete the picture.

Signed-off-by: dotnet-bot-corefx-mirror <dotnet-bot@microsoft.com>

Author: stephentoub

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

@@ -33,7 +33,7 @@ public readonly struct ConfiguredValueTaskAwaitable
         [StructLayout(LayoutKind.Auto)]
         public readonly struct ConfiguredValueTaskAwaiter : ICriticalNotifyCompletion
 #if CORECLR
-            , IValueTaskAwaiter
+            , IStateMachineBoxAwareAwaiter
 #endif
         {
             /// <summary>The value being awaited.</summary>
@@ -94,7 +94,7 @@ public void UnsafeOnCompleted(Action continuation)
             }
 
 #if CORECLR
-            void IValueTaskAwaiter.AwaitUnsafeOnCompleted(IAsyncStateMachineBox box)
+            void IStateMachineBoxAwareAwaiter.AwaitUnsafeOnCompleted(IAsyncStateMachineBox box)
             {
                 if (_value.ObjectIsTask)
                 {
@@ -135,7 +135,7 @@ public readonly struct ConfiguredValueTaskAwaitable<TResult>
         [StructLayout(LayoutKind.Auto)]
         public readonly struct ConfiguredValueTaskAwaiter : ICriticalNotifyCompletion
 #if CORECLR
-            , IValueTaskAwaiter
+            , IStateMachineBoxAwareAwaiter
 #endif
         {
             /// <summary>The value being awaited.</summary>
@@ -196,7 +196,7 @@ public void UnsafeOnCompleted(Action continuation)
             }
 
 #if CORECLR
-            void IValueTaskAwaiter.AwaitUnsafeOnCompleted(IAsyncStateMachineBox box)
+            void IStateMachineBoxAwareAwaiter.AwaitUnsafeOnCompleted(IAsyncStateMachineBox box)
             {
                 if (_value.ObjectIsTask)
                 {

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

@@ -11,7 +11,7 @@ namespace System.Runtime.CompilerServices
     /// <summary>Provides an awaiter for a <see cref="ValueTask"/>.</summary>
     public readonly struct ValueTaskAwaiter : ICriticalNotifyCompletion
 #if CORECLR
-            , IValueTaskAwaiter
+            , IStateMachineBoxAwareAwaiter
 #endif
     {
         /// <summary>Shim used to invoke an <see cref="Action"/> passed as the state argument to a <see cref="Action{Object}"/>.</summary>
@@ -80,7 +80,7 @@ public void UnsafeOnCompleted(Action continuation)
         }
 
 #if CORECLR
-        void IValueTaskAwaiter.AwaitUnsafeOnCompleted(IAsyncStateMachineBox box)
+        void IStateMachineBoxAwareAwaiter.AwaitUnsafeOnCompleted(IAsyncStateMachineBox box)
         {
             if (_value.ObjectIsTask)
             {
@@ -113,7 +113,7 @@ void IValueTaskAwaiter.AwaitUnsafeOnCompleted(IAsyncStateMachineBox box)
     /// <summary>Provides an awaiter for a <see cref="ValueTask{TResult}"/>.</summary>
     public readonly struct ValueTaskAwaiter<TResult> : ICriticalNotifyCompletion
 #if CORECLR
-            , IValueTaskAwaiter
+            , IStateMachineBoxAwareAwaiter
 #endif
     {
         /// <summary>The value being awaited.</summary>
@@ -171,7 +171,7 @@ public void UnsafeOnCompleted(Action continuation)
         }
 
 #if CORECLR
-        void IValueTaskAwaiter.AwaitUnsafeOnCompleted(IAsyncStateMachineBox box)
+        void IStateMachineBoxAwareAwaiter.AwaitUnsafeOnCompleted(IAsyncStateMachineBox box)
         {
             if (_value.ObjectIsTask)
             {
@@ -190,8 +190,8 @@ void IValueTaskAwaiter.AwaitUnsafeOnCompleted(IAsyncStateMachineBox box)
     }
 
 #if CORECLR
-    /// <summary>Internal interface used to enable optimizations from <see cref="AsyncTaskMethodBuilder"/> on <see cref="ValueTask"/>.</summary>>
-    internal interface IValueTaskAwaiter
+    /// <summary>Internal interface used to enable optimizations from <see cref="AsyncTaskMethodBuilder"/>.</summary>>
+    internal interface IStateMachineBoxAwareAwaiter
     {
         /// <summary>Invoked to set <see cref="ITaskCompletionAction.Invoke"/> of the <paramref name="box"/> as the awaiter's continuation.</summary>
         /// <param name="box">The box object.</param>

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

@@ -47,6 +47,9 @@ public readonly struct YieldAwaitable
         /// <summary>Provides an awaiter that switches into a target environment.</summary>
         /// <remarks>This type is intended for compiler use only.</remarks>
         public readonly struct YieldAwaiter : ICriticalNotifyCompletion
+#if CORECLR
+            , IStateMachineBoxAwareAwaiter
+#endif
         {
             /// <summary>Gets whether a yield is not required.</summary>
             /// <remarks>This property is intended for compiler user rather than use directly in code.</remarks>
@@ -115,6 +118,41 @@ private static void QueueContinuation(Action continuation, bool flowContext)
                 }
             }
 
+#if CORECLR
+            void IStateMachineBoxAwareAwaiter.AwaitUnsafeOnCompleted(IAsyncStateMachineBox box)
+            {
+                Debug.Assert(box != null);
+
+                // If tracing is enabled, delegate the Action-based implementation.
+                if (TplEtwProvider.Log.IsEnabled())
+                {
+                    QueueContinuation(box.MoveNextAction, flowContext: false);
+                    return;
+                }
+
+                // Otherwise, this is the same logic as in QueueContinuation, except using
+                // an IAsyncStateMachineBox instead of an Action, and only for flowContext:false.
+
+                SynchronizationContext syncCtx = SynchronizationContext.Current;
+                if (syncCtx != null && syncCtx.GetType() != typeof(SynchronizationContext))
+                {
+                    syncCtx.Post(s => ((IAsyncStateMachineBox)s).MoveNext(), box);
+                }
+                else
+                {
+                    TaskScheduler scheduler = TaskScheduler.Current;
+                    if (scheduler == TaskScheduler.Default)
+                    {
+                        ThreadPool.UnsafeQueueUserWorkItem(s => ((IAsyncStateMachineBox)s).MoveNext(), box);
+                    }
+                    else
+                    {
+                        Task.Factory.StartNew(s => ((IAsyncStateMachineBox)s).MoveNext(), box, default, TaskCreationOptions.PreferFairness, scheduler);
+                    }
+                }
+            }
+#endif
+
             private static Action OutputCorrelationEtwEvent(Action continuation)
             {
 #if CORERT
@@ -153,7 +191,6 @@ private static Action OutputCorrelationEtwEvent(Action continuation)
             private static readonly WaitCallback s_waitCallbackRunAction = RunAction;
             /// <summary>SendOrPostCallback that invokes the Action supplied as object state.</summary>
             private static readonly SendOrPostCallback s_sendOrPostCallbackRunAction = RunAction;
-
             /// <summary>Runs an Action delegate provided as state.</summary>
             /// <param name="state">The Action delegate to invoke.</param>
             private static void RunAction(object state) { ((Action)state)(); }