Fix ValueTask behavior as async return type

There's a bug in ValueTask<T>'s builder based on a misunderstanding of what the codegen would be from the compiler.  I'd thought the compiler would codegen access to the builder's Task property, but that was based on a very old model; the model actualy in place relies on the Await*OnCompleted methods doing any initialization that's necessary prior to the await.  That's an important distinction, because the builder for ValueTask needs defer to the wrapped task builder's Task property once any await has been issued.  Which means that rather than just delegating to the wrapped builder's Await*OnCompleted methods, before doing so the ValueTask's builder must remember that these methods have been called, such that a subsequent access will defer to the wrapped Task, even if a result is set in the interim.

Author: stephentoub

src/System.Threading.Tasks.Extensions/src/System/Runtime/CompilerServices/AsyncValueTaskMethodBuilder.cs

@@ -84,6 +84,7 @@ public void AwaitOnCompleted<TAwaiter, TStateMachine>(ref TAwaiter awaiter, ref
             where TAwaiter : INotifyCompletion
             where TStateMachine : IAsyncStateMachine
         {
+            _useBuilder = true;
             _methodBuilder.AwaitOnCompleted(ref awaiter, ref stateMachine);
         }
 
@@ -97,6 +98,7 @@ public void AwaitUnsafeOnCompleted<TAwaiter, TStateMachine>(ref TAwaiter awaiter
             where TAwaiter : ICriticalNotifyCompletion 
             where TStateMachine : IAsyncStateMachine
         {
+            _useBuilder = true;
             _methodBuilder.AwaitUnsafeOnCompleted(ref awaiter, ref stateMachine);
         }
     }

src/System.Threading.Tasks.Extensions/tests/AsyncValueTaskMethodBuilderTests.cs

@@ -23,7 +23,7 @@ public void SetResult_BeforeAccessTask_ValueTaskContainsValue()
             b.SetResult(42);
             ValueTask<int> vt = b.Task;
             Assert.True(vt.IsCompletedSuccessfully);
-            Assert.NotSame(vt.AsTask(), vt.AsTask()); // will be different if completed synchronously
+            Assert.False(WrapsTask(vt));
             Assert.Equal(42, vt.Result);
         }
 
@@ -34,7 +34,7 @@ public void SetResult_AfterAccessTask_ValueTaskContainsValue()
             ValueTask<int> vt = b.Task;
             b.SetResult(42);
             Assert.True(vt.IsCompletedSuccessfully);
-            Assert.Same(vt.AsTask(), vt.AsTask()); // will be safe if completed asynchronously
+            Assert.True(WrapsTask(vt));
             Assert.Equal(42, vt.Result);
         }
 
@@ -111,6 +111,37 @@ public async Task AwaitOnCompleted_InvokesStateMachineMethods(bool awaitUnsafe)
             Assert.Equal(dsm, foundSm);
         }
 
+        [Theory]
+        [InlineData(1, false)]
+        [InlineData(2, false)]
+        [InlineData(1, true)]
+        [InlineData(2, true)]
+        public void AwaitOnCompleted_ForcesTaskCreation(int numAwaits, bool awaitUnsafe)
+        {
+            AsyncValueTaskMethodBuilder<int> b = ValueTask<int>.CreateAsyncMethodBuilder();
+
+            var dsm = new DelegateStateMachine();
+            TaskAwaiter<int> t = new TaskCompletionSource<int>().Task.GetAwaiter();
+
+            Assert.InRange(numAwaits, 1, int.MaxValue);
+            for (int i = 1; i <= numAwaits; i++)
+            {
+                if (awaitUnsafe)
+                {
+                    b.AwaitUnsafeOnCompleted(ref t, ref dsm);
+                }
+                else
+                {
+                    b.AwaitOnCompleted(ref t, ref dsm);
+                }
+            }
+
+            b.SetResult(42);
+
+            Assert.True(WrapsTask(b.Task));
+            Assert.Equal(42, b.Task.Result);
+        }
+
         [Fact]
         public void SetStateMachine_InvalidArgument_ThrowsException()
         {
@@ -150,9 +181,12 @@ public void Start_ExecutionContextChangesInMoveNextDontFlowOut()
             // Make sure we've not caused the Task to be allocated
             b.SetResult(42);
             ValueTask<int> vt = b.Task;
-            Assert.NotSame(vt.AsTask(), vt.AsTask());
+            Assert.False(WrapsTask(vt));
         }
 
+        /// <summary>Gets whether the ValueTask has a non-null Task.</summary>
+        private static bool WrapsTask<T>(ValueTask<T> vt) => ReferenceEquals(vt.AsTask(), vt.AsTask());
+
         private struct DelegateStateMachine : IAsyncStateMachine
         {
             internal Action MoveNextDelegate;