Fix inlining of IAsyncStateMachineBox (#16830)

Prior to .NET Core 2.1, if a task was awaited in a default context (or if ConfigureAwait(false) was used), and if the task was then completed on a non-default context (e.g. code running a non-default TaskScheduler calling SetResult on a TaskCompletionSource), its await continuations are not allowed to be inlined, due to concerns that we could be running an arbitrary amount of unknown code in a scheduling environment that's not amenable to it, like a UI thread. The optimizations added in 2.1 erroneously ended up bypassing that IsValidLocationForInlining check, leading to continuations getting inlined in places they weren't previously.

This commit fixes that.  Previously we'd made the IAsyncStateMachineBox interface inherit ITaskCompletionAction, with its Invoke interface method just delegating to MoveNext; then the box could be added to a task as a continuation.  But that ITaskCompletionAction logic isn't specific to async/await and doesn't invoke IsValidLocationForInlining.  We instead need to follow the Action logic that is meant to be used for async/await.  I've removed the ITaskCompletionAction from IAsyncStateMachineBox, replacing it just with a MoveNext interface method, and added a case for IAsyncStateMachineBox to Task's RunContinuations logic.  I then duplicated the AwaitTaskContinuation.RunOrScheduleAction logic that's there for Action and tweaked it for IAsyncStateMachineBox.  In the process I cleaned up a little code to use some newer C# features.

4 files changed, 112 insertions, 92 deletions

diff --git a/src/mscorlib/shared/System/Runtime/CompilerServices/ValueTaskAwaiter.cs b/src/mscorlib/shared/System/Runtime/CompilerServices/ValueTaskAwaiter.cs
index e019f66131..12bac0e828 100644
--- a/src/mscorlib/shared/System/Runtime/CompilerServices/ValueTaskAwaiter.cs
+++ b/src/mscorlib/shared/System/Runtime/CompilerServices/ValueTaskAwaiter.cs
@@ -105,7 +105,7 @@ namespace System.Runtime.CompilerServices
                 return;
             }
 
-            box.Invoke(null);
+            box.MoveNext();
         };
 #endif
     }
diff --git a/src/mscorlib/src/System/Runtime/CompilerServices/AsyncMethodBuilder.cs b/src/mscorlib/src/System/Runtime/CompilerServices/AsyncMethodBuilder.cs
index 4603ee75c1..f23a4aa337 100644
--- a/src/mscorlib/src/System/Runtime/CompilerServices/AsyncMethodBuilder.cs
+++ b/src/mscorlib/src/System/Runtime/CompilerServices/AsyncMethodBuilder.cs
@@ -567,17 +567,6 @@ namespace System.Runtime.CompilerServices
                 }
             }
 
-            /// <summary>
-            /// Calls MoveNext on <see cref="StateMachine"/>. Implements ITaskCompletionAction.Invoke so
-            /// that the state machine object may be queued directly as a continuation into a Task's
-            /// continuation slot/list.
-            /// </summary>
-            /// <param name="ignored">The completing task that caused this method to be invoked, if there was one.</param>
-            void ITaskCompletionAction.Invoke(Task ignored) => MoveNext();
-
-            /// <summary>Signals to Task's continuation logic that <see cref="Invoke"/> runs arbitrary user code via MoveNext.</summary>
-            bool ITaskCompletionAction.InvokeMayRunArbitraryCode => true;
-
             /// <summary>Gets the state machine as a boxed object.  This should only be used for debugging purposes.</summary>
             IAsyncStateMachine IAsyncStateMachineBox.GetStateMachineObject() => StateMachine; // likely boxes, only use for debugging
         }
@@ -895,8 +884,11 @@ namespace System.Runtime.CompilerServices
     /// <summary>
     /// An interface implemented by all <see cref="AsyncStateMachineBox{TStateMachine, TResult}"/> instances, regardless of generics.
     /// </summary>
-    internal interface IAsyncStateMachineBox : ITaskCompletionAction
+    internal interface IAsyncStateMachineBox
     {
+        /// <summary>Move the state machine forward.</summary>
+        void MoveNext();
+
         /// <summary>
         /// Gets an action for moving forward the contained state machine.
         /// This will lazily-allocate the delegate as needed.
diff --git a/src/mscorlib/src/System/Threading/Tasks/Task.cs b/src/mscorlib/src/System/Threading/Tasks/Task.cs
index 8e4115a5d0..e4ca132a63 100644
--- a/src/mscorlib/src/System/Threading/Tasks/Task.cs
+++ b/src/mscorlib/src/System/Threading/Tasks/Task.cs
@@ -3261,42 +3261,38 @@ namespace System.Threading.Tasks
                 AsyncCausalityTracer.TraceSynchronousWorkStart(CausalityTraceLevel.Required, this.Id, CausalitySynchronousWork.CompletionNotification);
 
             // Skip synchronous execution of continuations if this task's thread was aborted
-            bool bCanInlineContinuations = !(((m_stateFlags & TASK_STATE_THREAD_WAS_ABORTED) != 0) ||
-                                              (RuntimeThread.CurrentThread.ThreadState == ThreadState.AbortRequested) ||
-                                              ((m_stateFlags & (int)TaskCreationOptions.RunContinuationsAsynchronously) != 0));
+            bool canInlineContinuations = !(((m_stateFlags & TASK_STATE_THREAD_WAS_ABORTED) != 0) ||
+                                             (RuntimeThread.CurrentThread.ThreadState == ThreadState.AbortRequested) ||
+                                             ((m_stateFlags & (int)TaskCreationOptions.RunContinuationsAsynchronously) != 0));
+
+            switch (continuationObject)
+            {
+                // Handle the single IAsyncStateMachineBox case.  This could be handled as part of the ITaskCompletionAction
+                // but we want to ensure that inlining is properly handled in the face of schedulers, so its behavior
+                // needs to be customized ala raw Actions.  This is also the most important case, as it represents the
+                // most common form of continuation, so we check it first.
+                case IAsyncStateMachineBox stateMachineBox:
+                    AwaitTaskContinuation.RunOrScheduleAction(stateMachineBox, canInlineContinuations);
+                    LogFinishCompletionNotification();
+                    return;
 
-            // Handle the single-Action case
-            Action singleAction = continuationObject as Action;
-            if (singleAction != null)
-            {
-                AwaitTaskContinuation.RunOrScheduleAction(singleAction, bCanInlineContinuations, ref t_currentTask);
-                LogFinishCompletionNotification();
-                return;
-            }
+                // Handle the single Action case.
+                case Action action:
+                    AwaitTaskContinuation.RunOrScheduleAction(action, canInlineContinuations);
+                    LogFinishCompletionNotification();
+                    return;
 
-            // Handle the single-ITaskCompletionAction case
-            ITaskCompletionAction singleTaskCompletionAction = continuationObject as ITaskCompletionAction;
-            if (singleTaskCompletionAction != null)
-            {
-                if (bCanInlineContinuations || !singleTaskCompletionAction.InvokeMayRunArbitraryCode)
-                {
-                    singleTaskCompletionAction.Invoke(this);
-                }
-                else
-                {
-                    ThreadPool.UnsafeQueueCustomWorkItem(new CompletionActionInvoker(singleTaskCompletionAction, this), forceGlobal: false);
-                }
-                LogFinishCompletionNotification();
-                return;
-            }
+                // Handle the single TaskContinuation case.
+                case TaskContinuation tc:
+                    tc.Run(this, canInlineContinuations);
+                    LogFinishCompletionNotification();
+                    return;
 
-            // Handle the single-TaskContinuation case
-            TaskContinuation singleTaskContinuation = continuationObject as TaskContinuation;
-            if (singleTaskContinuation != null)
-            {
-                singleTaskContinuation.Run(this, bCanInlineContinuations);
-                LogFinishCompletionNotification();
-                return;
+                // Handle the single ITaskCompletionAction case.
+                case ITaskCompletionAction completionAction:
+                    RunOrQueueCompletionAction(completionAction, canInlineContinuations);
+                    LogFinishCompletionNotification();
+                    return;
             }
 
             // Not a single; it must be a list.
@@ -3315,69 +3311,70 @@ namespace System.Threading.Tasks
             {
                 // Synchronous continuation tasks will have the ExecuteSynchronously option,
                 // and we're looking for asynchronous tasks...
-                var tc = continuations[i] as StandardTaskContinuation;
-                if (tc != null && (tc.m_options & TaskContinuationOptions.ExecuteSynchronously) == 0)
+                if (continuations[i] is StandardTaskContinuation tc &&
+                    (tc.m_options & TaskContinuationOptions.ExecuteSynchronously) == 0)
                 {
                     if (tplEtwProviderLoggingEnabled)
                     {
                         etw.RunningContinuationList(Id, i, tc);
                     }
                     continuations[i] = null; // so that we can skip this later
-                    tc.Run(this, bCanInlineContinuations);
+                    tc.Run(this, canInlineContinuations);
                 }
             }
 
             // ... and then fire the synchronous continuations (if there are any).
-            // This includes ITaskCompletionAction, AwaitTaskContinuations, and
-            // Action delegates, which are all by default implicitly synchronous.
+            // This includes ITaskCompletionAction, AwaitTaskContinuations, IAsyncStateMachineBox,
+            // and Action delegates, which are all by default implicitly synchronous.
             for (int i = 0; i < continuationCount; i++)
             {
                 object currentContinuation = continuations[i];
-                if (currentContinuation == null) continue;
+                if (currentContinuation == null)
+                {
+                    continue;
+                }
                 continuations[i] = null; // to enable free'ing up memory earlier
                 if (tplEtwProviderLoggingEnabled)
                 {
                     etw.RunningContinuationList(Id, i, currentContinuation);
                 }
 
-                // If the continuation is an Action delegate, it came from an await continuation,
-                // and we should use AwaitTaskContinuation to run it.
-                Action ad = currentContinuation as Action;
-                if (ad != null)
+                switch (currentContinuation)
                 {
-                    AwaitTaskContinuation.RunOrScheduleAction(ad, bCanInlineContinuations, ref t_currentTask);
-                }
-                else
-                {
-                    // If it's a TaskContinuation object of some kind, invoke it.
-                    TaskContinuation tc = currentContinuation as TaskContinuation;
-                    if (tc != null)
-                    {
-                        // We know that this is a synchronous continuation because the
-                        // asynchronous ones have been weeded out
-                        tc.Run(this, bCanInlineContinuations);
-                    }
-                    // Otherwise, it must be an ITaskCompletionAction, so invoke it.
-                    else
-                    {
-                        Debug.Assert(currentContinuation is ITaskCompletionAction, "Expected continuation element to be Action, TaskContinuation, or ITaskContinuationAction");
-                        var action = (ITaskCompletionAction)currentContinuation;
-
-                        if (bCanInlineContinuations || !action.InvokeMayRunArbitraryCode)
-                        {
-                            action.Invoke(this);
-                        }
-                        else
-                        {
-                            ThreadPool.UnsafeQueueCustomWorkItem(new CompletionActionInvoker(action, this), forceGlobal: false);
-                        }
-                    }
+                    case IAsyncStateMachineBox stateMachineBox:
+                        AwaitTaskContinuation.RunOrScheduleAction(stateMachineBox, canInlineContinuations);
+                        break;
+
+                    case Action action:
+                        AwaitTaskContinuation.RunOrScheduleAction(action, canInlineContinuations);
+                        break;
+
+                    case TaskContinuation tc:
+                        tc.Run(this, canInlineContinuations);
+                        break;
+
+                    default:
+                        Debug.Assert(currentContinuation is ITaskCompletionAction);
+                        RunOrQueueCompletionAction((ITaskCompletionAction)currentContinuation, canInlineContinuations);
+                        break;
                 }
             }
 
             LogFinishCompletionNotification();
         }
 
+        private void RunOrQueueCompletionAction(ITaskCompletionAction completionAction, bool allowInlining)
+        {
+            if (allowInlining || !completionAction.InvokeMayRunArbitraryCode)
+            {
+                completionAction.Invoke(this);
+            }
+            else
+            {
+                ThreadPool.UnsafeQueueCustomWorkItem(new CompletionActionInvoker(completionAction, this), forceGlobal: false);
+            }
+        }
+
         private void LogFinishCompletionNotification()
         {
             if (AsyncCausalityTracer.LoggingOn)
diff --git a/src/mscorlib/src/System/Threading/Tasks/TaskContinuation.cs b/src/mscorlib/src/System/Threading/Tasks/TaskContinuation.cs
index cb18c66d1c..3f72500ce4 100644
--- a/src/mscorlib/src/System/Threading/Tasks/TaskContinuation.cs
+++ b/src/mscorlib/src/System/Threading/Tasks/TaskContinuation.cs
@@ -737,30 +737,25 @@ namespace System.Threading.Tasks
         /// <param name="allowInlining">
         /// true to allow inlining, or false to force the action to run asynchronously.
         /// </param>
-        /// <param name="currentTask">
-        /// A reference to the t_currentTask thread static value.
-        /// This is passed by-ref rather than accessed in the method in order to avoid
-        /// unnecessary thread-static writes.
-        /// </param>
         /// <remarks>
         /// No ExecutionContext work is performed used.  This method is only used in the
         /// case where a raw Action continuation delegate was stored into the Task, which
         /// only happens in Task.SetContinuationForAwait if execution context flow was disabled
         /// via using TaskAwaiter.UnsafeOnCompleted or a similar path.
         /// </remarks>
-        internal static void RunOrScheduleAction(Action action, bool allowInlining, ref Task currentTask)
+        internal static void RunOrScheduleAction(Action action, bool allowInlining)
         {
-            Debug.Assert(currentTask == Task.t_currentTask);
+            ref Task currentTask = ref Task.t_currentTask;
+            Task prevCurrentTask = currentTask;
 
             // If we're not allowed to run here, schedule the action
             if (!allowInlining || !IsValidLocationForInlining)
             {
-                UnsafeScheduleAction(action, currentTask);
+                UnsafeScheduleAction(action, prevCurrentTask);
                 return;
             }
 
             // Otherwise, run it, making sure that t_currentTask is null'd out appropriately during the execution
-            Task prevCurrentTask = currentTask;
             try
             {
                 if (prevCurrentTask != null) currentTask = null;
@@ -776,6 +771,42 @@ namespace System.Threading.Tasks
             }
         }
 
+        /// <summary>Invokes or schedules the action to be executed.</summary>
+        /// <param name="box">The <see cref="IAsyncStateMachineBox"/> that needs to be invoked or queued.</param>
+        /// <param name="allowInlining">
+        /// true to allow inlining, or false to force the box's action to run asynchronously.
+        /// </param>
+        internal static void RunOrScheduleAction(IAsyncStateMachineBox box, bool allowInlining)
+        {
+            // Same logic as in the RunOrScheduleAction(Action, ...) overload, except invoking
+            // box.Invoke instead of action().
+
+            ref Task currentTask = ref Task.t_currentTask;
+            Task prevCurrentTask = currentTask;
+
+            // If we're not allowed to run here, schedule the action
+            if (!allowInlining || !IsValidLocationForInlining)
+            {
+                UnsafeScheduleAction(box.MoveNextAction, prevCurrentTask);
+                return;
+            }
+
+            // Otherwise, run it, making sure that t_currentTask is null'd out appropriately during the execution
+            try
+            {
+                if (prevCurrentTask != null) currentTask = null;
+                box.MoveNext();
+            }
+            catch (Exception exception)
+            {
+                ThrowAsyncIfNecessary(exception);
+            }
+            finally
+            {
+                if (prevCurrentTask != null) currentTask = prevCurrentTask;
+            }
+        }
+
         /// <summary>Schedules the action to be executed.  No ExecutionContext work is performed used.</summary>
         /// <param name="action">The action to invoke or queue.</param>
         internal static void UnsafeScheduleAction(Action action, Task task)