Special case runtime async in ThreadPool.UnsafeQueueUserWorkItem<TState>

src/coreclr/System.Private.CoreLib/src/System/Runtime/CompilerServices/AsyncHelpers.CoreCLR.cs

@@ -374,7 +374,7 @@ public RuntimeAsyncTask()
                 m_stateFlags |= (int)InternalTaskOptions.HiddenState;
             }
 
-            internal override void ExecuteFromThreadPool(Thread threadPoolThread)
+            internal override void ExecuteDirectly(Thread? threadPoolThread)
             {
                 DispatchContinuations();
             }
@@ -491,7 +491,7 @@ internal unsafe bool HandleSuspended(ref RuntimeAsyncAwaitState state)
 
                         // Clear continuation flags, so that continuation runs transparently
                         nextUserContinuation.Flags &= ~continueFlags;
-                        valueTaskSourceNotifier.OnCompleted(s_runContinuationAction, this, configFlags);
+                        valueTaskSourceNotifier.OnCompleted(ThreadPool.s_dispatchRuntimeAsyncContinuationsCallback, this, configFlags);
                     }
                     else
                     {
@@ -851,12 +851,6 @@ private bool QueueContinuationFollowUpActionIfNecessary(Continuation continuatio
                 Debug.Assert(state is RuntimeAsyncTask<T>);
                 ((RuntimeAsyncTask<T>)state).DispatchContinuations();
             };
-
-            private static readonly Action<object?> s_runContinuationAction = static state =>
-            {
-                Debug.Assert(state is RuntimeAsyncTask<T>);
-                ((RuntimeAsyncTask<T>)state).DispatchContinuations();
-            };
         }
 
         private static void InstrumentedFinalizeRuntimeAsyncTask<T>(RuntimeAsyncTask<T> task, ref RuntimeAsyncAwaitState state, AsyncInstrumentation.Flags flags)

src/libraries/System.Private.CoreLib/src/System/Runtime/CompilerServices/AsyncTaskMethodBuilderT.cs

@@ -344,7 +344,7 @@ public ref ExecutionContext? Context
                 }
             }
 
-            internal sealed override void ExecuteFromThreadPool(Thread threadPoolThread) => MoveNext(threadPoolThread);
+            internal sealed override void ExecuteDirectly(Thread? threadPoolThread) => MoveNext(threadPoolThread);
 
             /// <summary>Calls MoveNext on <see cref="StateMachine"/></summary>
             public void MoveNext() => MoveNext(threadPoolThread: null);

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

@@ -2423,12 +2423,13 @@ internal bool ExecuteEntry()
         }
 
         /// <summary>
-        /// ThreadPool's entry point into the Task.  The base behavior is simply to
-        /// use the entry point that's not protected from double-invoke; derived internal tasks
-        /// can override to customize their behavior, which is usually done by promises
-        /// that want to reuse the same object as a queued work item.
+        /// This is used internally to execute the Task directly. ThreadPool uses this,
+        /// and it is also used to invoke runtime async tasks directly.
+        /// The base behavior is simply to use the entry point that's not protected from
+        /// double-invoke; derived internal tasks 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(Thread threadPoolThread) => ExecuteEntryUnsafe(threadPoolThread);
+        internal virtual void ExecuteDirectly(Thread? threadPoolThread) => ExecuteEntryUnsafe(threadPoolThread);
 
         internal void ExecuteEntryUnsafe(Thread? threadPoolThread) // used instead of ExecuteEntry() when we don't have to worry about double-execution prevent
         {

src/libraries/System.Private.CoreLib/src/System/Threading/ThreadPoolWorkQueue.cs

@@ -1047,7 +1047,7 @@ private static void DispatchWorkItem(object workItem, Thread currentThread)
             {
                 // Task workitems catch their exceptions for later observation
                 // We do not need to pass unhandled ones to ExceptionHandling.s_handler
-                task.ExecuteFromThreadPool(currentThread);
+                task.ExecuteDirectly(currentThread);
             }
             else
             {
@@ -1428,6 +1428,19 @@ public static partial class ThreadPool
             }
         };
 
+        /// <summary>Shim used to invoke <see cref="Task.ExecuteDirectly"/> of a supplied <see cref="Task"/>.</summary>
+        internal static readonly Action<object?> s_dispatchRuntimeAsyncContinuationsCallback = static state =>
+        {
+            if (state is Task t)
+            {
+                t.ExecuteDirectly(null);
+            }
+            else
+            {
+                ThrowHelper.ThrowUnexpectedStateForKnownCallback(state);
+            }
+        };
+
         internal static bool EnableWorkerTracking => IsWorkerTrackingEnabledInConfig && EventSource.IsSupported;
 
 #if !FEATURE_WASM_MANAGED_THREADS
@@ -1647,6 +1660,20 @@ public static bool UnsafeQueueUserWorkItem<TState>(Action<TState> callBack, TSta
                 return true;
             }
 
+            // Similarly, for runtime async, user code may call with the
+            // runtime async callback directly.
+            if (ReferenceEquals(callBack, s_dispatchRuntimeAsyncContinuationsCallback))
+            {
+                if (state is not Task)
+                {
+                    // The provided state must be the internal RuntimeAsyncTask (Task)
+                    ThrowHelper.ThrowUnexpectedStateForKnownCallback(state);
+                }
+
+                UnsafeQueueUserWorkItemInternal((object)state, preferLocal);
+                return true;
+            }
+
             s_workQueue.Enqueue(
                 new QueueUserWorkItemCallbackDefaultContext<TState>(callBack, state), forceGlobal: !preferLocal);