Add stack depth check to all Task continuations (#23152)

Currently Task has a stack depth check that avoids stack overflows on very deep stack continuation chains, but it only applies to Task.ContinueWith, not to other kinds of continuations.  This changes that to have it apply to all.

As part of this, this also deletes the current StackGuard type used to achieve the check.  The type was meant to avoid expensive calls to check where we are on the stack, but now that we're using TryEnsureSufficientExecutionStack, it's actually faster to just call that rather than access the current StackGuard from a ThreadLocal.  This then also cleans up the call sites nicely, as they no longer need finally blocks to undo the increment performed on the StackGuard.

2 files changed, 18 insertions, 90 deletions

diff --git a/src/System.Private.CoreLib/shared/System/Threading/Tasks/Task.cs b/src/System.Private.CoreLib/shared/System/Threading/Tasks/Task.cs
index 4d790f483d..5d893c5cee 100644
--- a/src/System.Private.CoreLib/shared/System/Threading/Tasks/Task.cs
+++ b/src/System.Private.CoreLib/shared/System/Threading/Tasks/Task.cs
@@ -131,8 +131,6 @@ namespace System.Threading.Tasks
     {
         [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
 
@@ -1258,23 +1256,6 @@ namespace System.Threading.Tasks
         }
 
         /// <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
         /// cref="Task">Task</see> completed successfully or has not yet thrown any
@@ -3336,7 +3317,9 @@ namespace System.Threading.Tasks
             if (AsyncCausalityTracer.LoggingOn)
                 AsyncCausalityTracer.TraceSynchronousWorkStart(this, CausalitySynchronousWork.CompletionNotification);
 
-            bool canInlineContinuations = (m_stateFlags & (int)TaskCreationOptions.RunContinuationsAsynchronously) == 0;
+            bool canInlineContinuations = 
+                (m_stateFlags & (int)TaskCreationOptions.RunContinuationsAsynchronously) == 0 &&
+                RuntimeHelpers.TryEnsureSufficientExecutionStack();
 
             switch (continuationObject)
             {
@@ -6485,51 +6468,6 @@ namespace System.Threading.Tasks
         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 || RuntimeHelpers.TryEnsureSufficientExecutionStack())
-            {
-                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;
-        }
-    }
-
     // Special internal struct that we use to signify that we are not interested in
     // a Task<VoidTaskResult>'s result.
     internal struct VoidTaskResult { }
@@ -6609,18 +6547,17 @@ namespace System.Threading.Tasks
         // 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>
diff --git a/src/System.Private.CoreLib/shared/System/Threading/Tasks/TaskScheduler.cs b/src/System.Private.CoreLib/shared/System/Threading/Tasks/TaskScheduler.cs
index f274f20558..1d1a581b0b 100644
--- a/src/System.Private.CoreLib/shared/System/Threading/Tasks/TaskScheduler.cs
+++ b/src/System.Private.CoreLib/shared/System/Threading/Tasks/TaskScheduler.cs
@@ -179,12 +179,11 @@ namespace System.Threading.Tasks
             // 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;
             }
@@ -192,27 +191,19 @@ namespace System.Threading.Tasks
             // 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 (TplEventSource.Log.IsEnabled())
-                    task.FireTaskScheduledIfNeeded(this);
+            if (TplEventSource.Log.IsEnabled())
+                task.FireTaskScheduledIfNeeded(this);
 
-                bInlined = TryExecuteTaskInline(task, taskWasPreviouslyQueued);
-            }
-            finally
-            {
-                currentStackGuard.EndInliningScope();
-            }
+            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>