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