Merge pull request #23390 from stephentoub/port23152

[release/2.2] Add stack depth check to all Task continuations

Author: wtgodbe

src/mscorlib/src/System/Threading/Tasks/Task.cs

@@ -134,8 +134,6 @@ public class Task : IThreadPoolWorkItem, IAsyncResult, IDisposable
     {
         [ThreadStatic]
         internal static Task t_currentTask;  // The currently executing task.
-        [ThreadStatic]
-        private static StackGuard t_stackGuard;  // The stack guard object for this thread
 
         internal static int s_taskIdCounter; //static counter used to generate unique task IDs
 
@@ -1276,23 +1274,6 @@ internal static Task InternalCurrentIfAttached(TaskCreationOptions creationOptio
             return (creationOptions & TaskCreationOptions.AttachedToParent) != 0 ? InternalCurrent : null;
         }
 
-        /// <summary>
-        /// Gets the StackGuard object assigned to the current thread.
-        /// </summary>
-        internal static StackGuard CurrentStackGuard
-        {
-            get
-            {
-                StackGuard sg = t_stackGuard;
-                if (sg == null)
-                {
-                    t_stackGuard = sg = new StackGuard();
-                }
-                return sg;
-            }
-        }
-
-
         /// <summary>
         /// Gets the <see cref="T:System.AggregateException">Exception</see> that caused the <see
         /// cref="Task">Task</see> to end prematurely. If the <see
@@ -3263,7 +3244,8 @@ private void RunContinuations(object continuationObject) // separated out of Fin
             // Skip synchronous execution of continuations if this task's thread was aborted
             bool canInlineContinuations = !(((m_stateFlags & TASK_STATE_THREAD_WAS_ABORTED) != 0) ||
                                              (RuntimeThread.CurrentThread.ThreadState == ThreadState.AbortRequested) ||
-                                             ((m_stateFlags & (int)TaskCreationOptions.RunContinuationsAsynchronously) != 0));
+                                             ((m_stateFlags & (int)TaskCreationOptions.RunContinuationsAsynchronously) != 0))
+                                          && RuntimeHelpers.TryEnsureSufficientExecutionStack();
 
             switch (continuationObject)
             {
@@ -6415,56 +6397,6 @@ public enum TaskContinuationOptions
         ExecuteSynchronously = 0x80000
     }
 
-    /// <summary>
-    /// Internal helper class to keep track of stack depth and decide whether we should inline or not.
-    /// </summary>
-    internal class StackGuard
-    {
-        // current thread's depth of nested inline task executions
-        private int m_inliningDepth = 0;
-
-        // For relatively small inlining depths we don't want to get into the business of stack probing etc.
-        // This clearly leaves a window of opportunity for the user code to SO. However a piece of code
-        // that can SO in 20 inlines on a typical 1MB stack size probably needs to be revisited anyway.
-        private const int MAX_UNCHECKED_INLINING_DEPTH = 20;
-
-        /// <summary>
-        /// This method needs to be called before attempting inline execution on the current thread. 
-        /// If false is returned, it means we are too close to the end of the stack and should give up inlining.
-        /// Each call to TryBeginInliningScope() that returns true must be matched with a 
-        /// call to EndInliningScope() regardless of whether inlining actually took place.
-        /// </summary>
-        internal bool TryBeginInliningScope()
-        {
-            // If we're still under the 'safe' limit we'll just skip the stack probe to save p/invoke calls
-            if (m_inliningDepth < MAX_UNCHECKED_INLINING_DEPTH || CheckForSufficientStack())
-            {
-                m_inliningDepth++;
-                return true;
-            }
-            else
-                return false;
-        }
-
-        /// <summary>
-        /// This needs to be called once for each previous successful TryBeginInliningScope() call after
-        /// inlining related logic runs.
-        /// </summary>
-        internal void EndInliningScope()
-        {
-            m_inliningDepth--;
-            Debug.Assert(m_inliningDepth >= 0, "Inlining depth count should never go negative.");
-
-            // do the right thing just in case...
-            if (m_inliningDepth < 0) m_inliningDepth = 0;
-        }
-
-        private unsafe bool CheckForSufficientStack()
-        {
-            return RuntimeHelpers.TryEnsureSufficientExecutionStack();
-        }
-    }
-
     // Special internal struct that we use to signify that we are not interested in
     // a Task<VoidTaskResult>'s result.
     internal struct VoidTaskResult { }
@@ -6546,18 +6478,17 @@ public UnwrapPromise(Task outerTask, bool lookForOce)
         // For ITaskCompletionAction 
         public void Invoke(Task completingTask)
         {
-            // Check the current stack guard.  If we're ok to inline,
-            // process the task, and reset the guard when we're done.
-            var sg = Task.CurrentStackGuard;
-            if (sg.TryBeginInliningScope())
+            // If we're ok to inline, process the task. Otherwise, we're too deep on the stack, and
+            // we shouldn't run the continuation chain here, so queue a work item to call back here
+            // to Invoke asynchronously.
+            if (RuntimeHelpers.TryEnsureSufficientExecutionStack())
+            {
+                InvokeCore(completingTask);
+            }
+            else
             {
-                try { InvokeCore(completingTask); }
-                finally { sg.EndInliningScope(); }
+                InvokeCoreAsync(completingTask);
             }
-            // Otherwise, we're too deep on the stack, and
-            // we shouldn't run the continuation chain here, so queue a work
-            // item to call back here to Invoke asynchronously.
-            else InvokeCoreAsync(completingTask);
         }
 
         /// <summary>

src/mscorlib/src/System/Threading/Tasks/TaskScheduler.cs

@@ -186,41 +186,33 @@ internal bool TryRunInline(Task task, bool taskWasPreviouslyQueued)
             // Delegate cross-scheduler inlining requests to target scheduler
             if (ets != this && ets != null) return ets.TryRunInline(task, taskWasPreviouslyQueued);
 
-            StackGuard currentStackGuard;
             if ((ets == null) ||
                 (task.m_action == null) ||
                 task.IsDelegateInvoked ||
                 task.IsCanceled ||
-                (currentStackGuard = Task.CurrentStackGuard).TryBeginInliningScope() == false)
+                !RuntimeHelpers.TryEnsureSufficientExecutionStack())
             {
                 return false;
             }
 
             // Task class will still call into TaskScheduler.TryRunInline rather than TryExecuteTaskInline() so that 
             // 1) we can adjust the return code from TryExecuteTaskInline in case a buggy custom scheduler lies to us
             // 2) we maintain a mechanism for the TLS lookup optimization that we used to have for the ConcRT scheduler (will potentially introduce the same for TP)
-            bool bInlined = false;
-            try
-            {
-                if (TplEtwProvider.Log.IsEnabled())
-                {
-                    task.FireTaskScheduledIfNeeded(this);
-                }
-                bInlined = TryExecuteTaskInline(task, taskWasPreviouslyQueued);
-            }
-            finally
+            if (TplEtwProvider.Log.IsEnabled())
             {
-                currentStackGuard.EndInliningScope();
+                task.FireTaskScheduledIfNeeded(this);
             }
 
+            bool inlined = TryExecuteTaskInline(task, taskWasPreviouslyQueued);
+
             // If the custom scheduler returned true, we should either have the TASK_STATE_DELEGATE_INVOKED or TASK_STATE_CANCELED bit set
             // Otherwise the scheduler is buggy
-            if (bInlined && !(task.IsDelegateInvoked || task.IsCanceled))
+            if (inlined && !(task.IsDelegateInvoked || task.IsCanceled))
             {
                 throw new InvalidOperationException(SR.TaskScheduler_InconsistentStateAfterTryExecuteTaskInline);
             }
 
-            return bInlined;
+            return inlined;
         }
 
         /// <summary>