Fix TaskExtensions.Unwrap stack dive

In strange but possible task configurations, it's possible to set up a series of Unwrap calls such that the continuations involved will form a very long chain of synchronously-executing calls.  In such cases, without extra measure, the implementation can overflow the stack.  The mscorlib Unwrap implementation takes precautions to avoid this, and ContinueWith (which was used in the previous, slower, standalone Unwrap implementation) has checks to avoid this as well, but currently await does not.  Since our new Unwrap implementation uses the await infrastructure, it's susceptible again.

This commit adds a check to the new Unwrap implementation to avoid the stack dive, as well as porting a test that triggers the stack dive prior to the fix.  Some local benchmarks show little-to-no impact from the additional checks, and it's very rare to hit the fallback path that triggers the async queueing: while the actual numbers will depend on the actual code being run, the size of stack frames, etc., in the test we only hit the fallback case ~8 times out of 12000 iterations, and that's in code trying very hard to force this condition.

Author: stephentoub

src/System.Threading.Tasks/src/System/Threading/Tasks/TaskExtensions.CoreCLR.cs

@@ -2,6 +2,7 @@
 // Licensed under the MIT license. See LICENSE file in the project root for full license information.
 
 using System.Diagnostics;
+using System.Runtime.CompilerServices;
 
 namespace System.Threading.Tasks
 {
@@ -29,7 +30,7 @@ public static Task Unwrap(this Task<Task> task)
         {
             if (task == null)
                 throw new ArgumentNullException("task");
-            
+
             // Fast path for an already successfully completed outer task: just return the inner one.
             // As in the subsequent slower path, a null inner task is special-cased to mean cancellation.
             if (task.Status == TaskStatus.RanToCompletion && (task.CreationOptions & TaskCreationOptions.AttachedToParent) == 0)
@@ -160,6 +161,29 @@ private static bool TrySetFromTask<TResult>(this TaskCompletionSource<TResult> c
         {
             Debug.Assert(task.IsCompleted);
 
+            // Before transferring the results, check to make sure we're not too deep on the stack.  Calling TrySet*
+            // will cause any synchronous continuations to be invoked, which is fine unless we're so deep that doing
+            // so overflows.  ContinueWith has built-in support for avoiding such stack dives, but that support is not
+            // (yet) part of await's infrastructure, so until it is we mimic it manually.  This matches the behavior
+            // employed by the Unwrap implementation in mscorlib. (If TryEnsureSufficientExecutionStack is made public 
+            // in the future, switch to using it to avoid the try/catch block and exception.)
+            try
+            {
+                RuntimeHelpers.EnsureSufficientExecutionStack();
+            }
+            catch (InsufficientExecutionStackException)
+            {
+                // This is very rare.  We're too deep to safely invoke
+                // TrySet* synchronously, so do so asynchronously instead.
+                Task.Factory.StartNew(s =>
+                {
+                    var t = (Tuple<TaskCompletionSource<TResult>, Task>)s;
+                    TrySetFromTask(t.Item1, t.Item2);
+                }, Tuple.Create(completionSource, task), CancellationToken.None, TaskCreationOptions.None, TaskScheduler.Default);
+                return true;
+            }
+
+            // Transfer the results from the supplied Task to the TaskCompletionSource.
             bool result = false;
             switch(task.Status)
             {

src/System.Threading.Tasks/tests/UnwrapTests.cs

@@ -444,6 +444,24 @@ public void Generic_DefaultSchedulerUsed()
             }, CancellationToken.None, TaskCreationOptions.None, scheduler).GetAwaiter().GetResult();
         }
 
+        /// <summary>
+        /// Test that a long chain of Unwraps can execute without overflowing the stack.
+        /// </summary>
+        [Fact]
+        public void RunStackGuardTests()
+        {
+            const int DiveDepth = 12000;
+
+            Func<int, Task<int>> func = null;
+            func = count =>
+                ++count < DiveDepth ?
+                    Task.Factory.StartNew(() => func(count), CancellationToken.None, TaskCreationOptions.None, TaskScheduler.Default).Unwrap() :
+                    Task.FromResult(count);
+
+            // This test will overflow if it fails.
+            Assert.Equal(DiveDepth, func(0).Result);
+        }
+
         /// <summary>Gets an enumerable of already completed non-generic tasks.</summary>
         public static IEnumerable<object[]> CompletedNonGenericTasks
         {