Queue ValueTaskAwaiter IAsyncStateMachineBox directly to ThreadPool (#21159)

* Queue ValueTaskAwaiter IAsyncStateMachineBox directly to ThreadPool

* Invert the dependency

* Move to UnsafeQueueUserWorkItem

* MRVTSC queue null or Deafult EC to UnsafeQUWI

* Revert MRVTSC change

* Add comment and validation

* Use s_invokeAsyncStateMachineBox for AsTask

* nits

* nits 2

* Rever ValueTask

* nits

diff --git a/src/System.Private.CoreLib/shared/System/Runtime/CompilerServices/ConfiguredValueTaskAwaitable.cs b/src/System.Private.CoreLib/shared/System/Runtime/CompilerServices/ConfiguredValueTaskAwaitable.cs
index 8f7b0c8..b874af2 100644 (file)
--- a/src/System.Private.CoreLib/shared/System/Runtime/CompilerServices/ConfiguredValueTaskAwaitable.cs
+++ b/src/System.Private.CoreLib/shared/System/Runtime/CompilerServices/ConfiguredValueTaskAwaitable.cs
@@ -4,6 +4,7 @@
 
 using System.Diagnostics;
 using System.Runtime.InteropServices;
+using System.Threading;
 using System.Threading.Tasks;
 using System.Threading.Tasks.Sources;
 
@@ -111,7 +112,7 @@ namespace System.Runtime.CompilerServices
                 }
                 else if (obj != null)
                 {
-                    Unsafe.As<IValueTaskSource>(obj).OnCompleted(ValueTaskAwaiter.s_invokeAsyncStateMachineBox, box, _value._token,
+                    Unsafe.As<IValueTaskSource>(obj).OnCompleted(ThreadPoolGlobals.s_invokeAsyncStateMachineBox, box, _value._token,
                         _value._continueOnCapturedContext ? ValueTaskSourceOnCompletedFlags.UseSchedulingContext : ValueTaskSourceOnCompletedFlags.None);
                 }
                 else
@@ -222,7 +223,7 @@ namespace System.Runtime.CompilerServices
                 }
                 else if (obj != null)
                 {
-                    Unsafe.As<IValueTaskSource<TResult>>(obj).OnCompleted(ValueTaskAwaiter.s_invokeAsyncStateMachineBox, box, _value._token,
+                    Unsafe.As<IValueTaskSource<TResult>>(obj).OnCompleted(ThreadPoolGlobals.s_invokeAsyncStateMachineBox, box, _value._token,
                         _value._continueOnCapturedContext ? ValueTaskSourceOnCompletedFlags.UseSchedulingContext : ValueTaskSourceOnCompletedFlags.None);
                 }
                 else

diff --git a/src/System.Private.CoreLib/shared/System/Runtime/CompilerServices/ValueTaskAwaiter.cs b/src/System.Private.CoreLib/shared/System/Runtime/CompilerServices/ValueTaskAwaiter.cs
index 02b5910..db14806 100644 (file)
--- a/src/System.Private.CoreLib/shared/System/Runtime/CompilerServices/ValueTaskAwaiter.cs
+++ b/src/System.Private.CoreLib/shared/System/Runtime/CompilerServices/ValueTaskAwaiter.cs
@@ -3,6 +3,7 @@
 // See the LICENSE file in the project root for more information.
 
 using System.Diagnostics;
+using System.Threading;
 using System.Threading.Tasks;
 using System.Threading.Tasks.Sources;
 
@@ -101,25 +102,13 @@ namespace System.Runtime.CompilerServices
             }
             else if (obj != null)
             {
-                Unsafe.As<IValueTaskSource>(obj).OnCompleted(s_invokeAsyncStateMachineBox, box, _value._token, ValueTaskSourceOnCompletedFlags.UseSchedulingContext);
+                Unsafe.As<IValueTaskSource>(obj).OnCompleted(ThreadPoolGlobals.s_invokeAsyncStateMachineBox, box, _value._token, ValueTaskSourceOnCompletedFlags.UseSchedulingContext);
             }
             else
             {
                 TaskAwaiter.UnsafeOnCompletedInternal(Task.CompletedTask, box, continueOnCapturedContext: true);
             }
         }
-
-        /// <summary>Shim used to invoke <see cref="ITaskCompletionAction.Invoke"/> of the supplied <see cref="IAsyncStateMachineBox"/>.</summary>
-        internal static readonly Action<object> s_invokeAsyncStateMachineBox = state =>
-        {
-            if (!(state is IAsyncStateMachineBox box))
-            {
-                ThrowHelper.ThrowArgumentOutOfRangeException(ExceptionArgument.state);
-                return;
-            }
-
-            box.MoveNext();
-        };
 #endif
     }
 
@@ -201,7 +190,7 @@ namespace System.Runtime.CompilerServices
             }
             else if (obj != null)
             {
-                Unsafe.As<IValueTaskSource<TResult>>(obj).OnCompleted(ValueTaskAwaiter.s_invokeAsyncStateMachineBox, box, _value._token, ValueTaskSourceOnCompletedFlags.UseSchedulingContext);
+                Unsafe.As<IValueTaskSource<TResult>>(obj).OnCompleted(ThreadPoolGlobals.s_invokeAsyncStateMachineBox, box, _value._token, ValueTaskSourceOnCompletedFlags.UseSchedulingContext);
             }
             else
             {

diff --git a/src/System.Private.CoreLib/src/System/Threading/ThreadPool.cs b/src/System.Private.CoreLib/src/System/Threading/ThreadPool.cs
index 5026003..cf44d85 100644 (file)
--- a/src/System.Private.CoreLib/src/System/Threading/ThreadPool.cs
+++ b/src/System.Private.CoreLib/src/System/Threading/ThreadPool.cs
@@ -37,6 +37,18 @@ namespace System.Threading
         public static bool enableWorkerTracking;
 
         public static readonly ThreadPoolWorkQueue workQueue = new ThreadPoolWorkQueue();
+
+        /// <summary>Shim used to invoke <see cref="IAsyncStateMachineBox.MoveNext"/> of the supplied <see cref="IAsyncStateMachineBox"/>.</summary>
+        internal static readonly Action<object> s_invokeAsyncStateMachineBox = state =>
+        {
+            if (!(state is IAsyncStateMachineBox box))
+            {
+                ThrowHelper.ThrowArgumentOutOfRangeException(ExceptionArgument.state);
+                return;
+            }
+
+            box.MoveNext();
+        };
     }
 
     [StructLayout(LayoutKind.Sequential)] // enforce layout so that padding reduces false sharing
@@ -1333,6 +1345,24 @@ namespace System.Threading
                 ThrowHelper.ThrowArgumentNullException(ExceptionArgument.callBack);
             }
 
+            // If the callback is the runtime-provided invocation of an IAsyncStateMachineBox,
+            // then we can queue the Task state directly to the ThreadPool instead of 
+            // wrapping it in a QueueUserWorkItemCallback.
+            //
+            // This occurs when user code queues its provided continuation to the ThreadPool;
+            // internally we call UnsafeQueueUserWorkItemInternal directly for Tasks.
+            if (ReferenceEquals(callBack, ThreadPoolGlobals.s_invokeAsyncStateMachineBox))
+            {
+                if (!(state is IAsyncStateMachineBox))
+                {
+                    // The provided state must be the internal IAsyncStateMachineBox (Task) type
+                    ThrowHelper.ThrowArgumentOutOfRangeException(ExceptionArgument.state);
+                }
+
+                UnsafeQueueUserWorkItemInternal((object)state, preferLocal);
+                return true;
+            }
+
             EnsureVMInitialized();
 
             ThreadPoolGlobals.workQueue.Enqueue(