diff --git a/src/coreclr/System.Private.CoreLib/src/System/Runtime/CompilerServices/AsyncHelpers.CoreCLR.cs b/src/coreclr/System.Private.CoreLib/src/System/Runtime/CompilerServices/AsyncHelpers.CoreCLR.cs index 64169a699065f4..efe5f8e9969326 100644 --- a/src/coreclr/System.Private.CoreLib/src/System/Runtime/CompilerServices/AsyncHelpers.CoreCLR.cs +++ b/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; } diff --git a/src/coreclr/System.Private.CoreLib/src/System/Runtime/CompilerServices/RuntimeAsyncTaskContinuation.cs b/src/coreclr/System.Private.CoreLib/src/System/Runtime/CompilerServices/RuntimeAsyncTaskContinuation.cs index faedf052025b72..a4928019e9a096 100644 --- a/src/coreclr/System.Private.CoreLib/src/System/Runtime/CompilerServices/RuntimeAsyncTaskContinuation.cs +++ b/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; } diff --git a/src/libraries/System.Private.CoreLib/src/System/Threading/Tasks/TaskContinuation.cs b/src/libraries/System.Private.CoreLib/src/System/Threading/Tasks/TaskContinuation.cs index a99b3425f9b4d9..ade4c8e5639730 100644 --- a/src/libraries/System.Private.CoreLib/src/System/Threading/Tasks/TaskContinuation.cs +++ b/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) + /// Inlines or schedules the action to run on the specified non-default scheduler. + /// The action to invoke. Must not be null. + /// The non-default scheduler with which to invoke the action. Must not be null and must not be the default scheduler. + /// The ExecutionContext with which to run the action, or null to not flow execution context. + /// true if inlining is permitted; otherwise, false. + 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 { /// The ExecutionContext with which to run the continuation. - private readonly ExecutionContext? m_capturedContext; + protected readonly ExecutionContext? m_capturedContext; /// The action to invoke. protected readonly Action m_action; @@ -543,8 +556,9 @@ internal AwaitTaskContinuation(Action action, bool flowExecutionContext) /// The action to run. Must not be null. /// The state to pass to the action. Must not be null. /// The scheduler to target. + /// The ExecutionContext with which to run the action, or null to not flow execution context. /// The created task. - protected Task CreateTask(Action action, object? state, TaskScheduler scheduler) + protected static Task CreateTask(Action action, object? state, TaskScheduler scheduler, ExecutionContext? capturedContext) { Debug.Assert(action != null); Debug.Assert(scheduler != null); @@ -553,7 +567,7 @@ protected Task CreateTask(Action action, object? state, TaskScheduler s action, state, null, default, TaskCreationOptions.None, InternalTaskOptions.QueuedByRuntime, scheduler) { - CapturedContext = m_capturedContext + CapturedContext = capturedContext }; }