Address runtime async TaskScheduler TODO

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

@@ -1149,9 +1149,7 @@ private bool QueueContinuationFollowUpActionIfNecessary(Continuation continuatio
                     TaskScheduler sched = (TaskScheduler)continuationContext;
 
                     SetContinuationState(continuation);
-                    // TODO: We do not need TaskSchedulerAwaitTaskContinuation here, just need to refactor its Run method...
-                    var taskSchedCont = new TaskSchedulerAwaitTaskContinuation(sched, GetContinuationAction(), flowExecutionContext: false);
-                    taskSchedCont.Run(Task.CompletedTask, canInlineContinuationTask: true);
+                    TaskSchedulerAwaitTaskContinuation.RunOrScheduleAction(GetContinuationAction(), sched, capturedContext: null, allowInlining: true);
 
                     return true;
                 }

src/coreclr/System.Private.CoreLib/src/System/Runtime/CompilerServices/RuntimeAsyncTaskContinuation.cs

@@ -95,9 +95,7 @@ private bool QueueIfNecessary(bool canInline)
                 Debug.Assert(continuationContext is TaskScheduler { });
                 TaskScheduler sched = (TaskScheduler)continuationContext;
 
-                // TODO: We do not need TaskSchedulerAwaitTaskContinuation here, just need to refactor its Run method...
-                var taskSchedCont = new TaskSchedulerAwaitTaskContinuation(sched, (Action)RuntimeAsyncTask.m_action!, flowExecutionContext: false);
-                taskSchedCont.Run(Task.CompletedTask, canInlineContinuationTask: canInline);
+                TaskSchedulerAwaitTaskContinuation.RunOrScheduleAction((Action)RuntimeAsyncTask.m_action!, sched, capturedContext: null, allowInlining: canInline);
 
                 return true;
             }

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

@@ -479,39 +479,52 @@ internal sealed override void Run(Task ignored, bool canInlineContinuationTask)
             }
             else
             {
-                // We permit inlining if the caller allows us to, and
-                // either we're on a thread pool thread (in which case we're fine running arbitrary code)
-                // or we're already on the target scheduler (in which case we'll just ask the scheduler
-                // whether it's ok to run here).  We include the IsThreadPoolThread check here, whereas
-                // we don't in AwaitTaskContinuation.Run, since here it expands what's allowed as opposed
-                // to in AwaitTaskContinuation.Run where it restricts what's allowed.
-                bool inlineIfPossible = canInlineContinuationTask &&
-                    (TaskScheduler.InternalCurrent == m_scheduler || Thread.CurrentThread.IsThreadPoolThread);
-
-                // Create the continuation task task. If we're allowed to inline, try to do so.
-                // The target scheduler may still deny us from executing on this thread, in which case this'll be queued.
-                Task task = CreateTask(static state =>
-                {
-                    try
-                    {
-                        ((Action)state!)();
-                    }
-                    catch (Exception exception)
-                    {
-                        Task.ThrowAsync(exception, targetContext: null);
-                    }
-                }, m_action, m_scheduler);
+                RunOrScheduleAction(m_action, m_scheduler, m_capturedContext, canInlineContinuationTask);
+            }
+        }
 
-                if (inlineIfPossible)
+        /// <summary>Inlines or schedules the action to run on the specified non-default scheduler.</summary>
+        /// <param name="action">The action to invoke. Must not be null.</param>
+        /// <param name="scheduler">The non-default scheduler with which to invoke the action. Must not be null and must not be the default scheduler.</param>
+        /// <param name="capturedContext">The ExecutionContext with which to run the action, or null to not flow execution context.</param>
+        /// <param name="allowInlining">true if inlining is permitted; otherwise, false.</param>
+        internal static void RunOrScheduleAction(Action action, TaskScheduler scheduler, ExecutionContext? capturedContext, bool allowInlining)
+        {
+            Debug.Assert(action != null);
+            Debug.Assert(scheduler != null && scheduler != TaskScheduler.Default);
+
+            // We permit inlining if the caller allows us to, and
+            // either we're on a thread pool thread (in which case we're fine running arbitrary code)
+            // or we're already on the target scheduler (in which case we'll just ask the scheduler
+            // whether it's ok to run here).  We include the IsThreadPoolThread check here, whereas
+            // we don't in AwaitTaskContinuation.Run, since here it expands what's allowed as opposed
+            // to in AwaitTaskContinuation.Run where it restricts what's allowed.
+            bool inlineIfPossible = allowInlining &&
+                (TaskScheduler.InternalCurrent == scheduler || Thread.CurrentThread.IsThreadPoolThread);
+
+            // Create the continuation task task. If we're allowed to inline, try to do so.
+            // The target scheduler may still deny us from executing on this thread, in which case this'll be queued.
+            Task task = CreateTask(static state =>
+            {
+                try
                 {
-                    InlineIfPossibleOrElseQueue(task, needsProtection: false);
+                    ((Action)state!)();
                 }
-                else
+                catch (Exception exception)
                 {
-                    // We need to run asynchronously, so just schedule the task.
-                    try { task.ScheduleAndStart(needsProtection: false); }
-                    catch (TaskSchedulerException) { } // No further action is necessary, as ScheduleAndStart already transitioned task to faulted
+                    Task.ThrowAsync(exception, targetContext: null);
                 }
+            }, action, scheduler, capturedContext);
+
+            if (inlineIfPossible)
+            {
+                InlineIfPossibleOrElseQueue(task, needsProtection: false);
+            }
+            else
+            {
+                // We need to run asynchronously, so just schedule the task.
+                try { task.ScheduleAndStart(needsProtection: false); }
+                catch (TaskSchedulerException) { } // No further action is necessary, as ScheduleAndStart already transitioned task to faulted
             }
         }
     }
@@ -520,7 +533,7 @@ internal sealed override void Run(Task ignored, bool canInlineContinuationTask)
     internal class AwaitTaskContinuation : TaskContinuation, IThreadPoolWorkItem
     {
         /// <summary>The ExecutionContext with which to run the continuation.</summary>
-        private readonly ExecutionContext? m_capturedContext;
+        protected readonly ExecutionContext? m_capturedContext;
         /// <summary>The action to invoke.</summary>
         protected readonly Action m_action;
 
@@ -543,8 +556,9 @@ internal AwaitTaskContinuation(Action action, bool flowExecutionContext)
         /// <param name="action">The action to run. Must not be null.</param>
         /// <param name="state">The state to pass to the action. Must not be null.</param>
         /// <param name="scheduler">The scheduler to target.</param>
+        /// <param name="capturedContext">The ExecutionContext with which to run the action, or null to not flow execution context.</param>
         /// <returns>The created task.</returns>
-        protected Task CreateTask(Action<object?> action, object? state, TaskScheduler scheduler)
+        protected static Task CreateTask(Action<object?> action, object? state, TaskScheduler scheduler, ExecutionContext? capturedContext)
         {
             Debug.Assert(action != null);
             Debug.Assert(scheduler != null);
@@ -553,7 +567,7 @@ protected Task CreateTask(Action<object?> action, object? state, TaskScheduler s
                 action, state, null, default,
                 TaskCreationOptions.None, InternalTaskOptions.QueuedByRuntime, scheduler)
             {
-                CapturedContext = m_capturedContext
+                CapturedContext = capturedContext
             };
         }