Merge pull request #10201 from stephentoub/valuetask_builder

Add AsyncValueTaskMethodBuilder in support of C# feature

src/System.Threading.Tasks.Extensions/src/System.Threading.Tasks.Extensions.csproj

@@ -4,7 +4,7 @@
   <PropertyGroup>
     <ProjectGuid>{F24D3391-2928-4E83-AADE-B34423498750}</ProjectGuid>
     <AssemblyName>System.Threading.Tasks.Extensions</AssemblyName>
-    <AssemblyVersion>4.0.1.0</AssemblyVersion>
+    <AssemblyVersion>4.1.0.0</AssemblyVersion>
     <DocumentationFile>$(OutputPath)$(AssemblyName).xml</DocumentationFile>
     <UseOpenKey Condition="'$(UseOpenKey)'==''">true</UseOpenKey>
   </PropertyGroup>
@@ -22,6 +22,7 @@
   <PropertyGroup Condition="'$(Configuration)|$(Platform)' == 'Release|AnyCPU'" />
   <ItemGroup>
     <Compile Include="System\Runtime\CompilerServices\ConfiguredValueTaskAwaitable.cs" />
+    <Compile Include="System\Runtime\CompilerServices\AsyncValueTaskMethodBuilder.cs" />
     <Compile Include="System\Runtime\CompilerServices\ValueTaskAwaiter.cs" />
     <Compile Include="System\Threading\Tasks\ValueTask.cs" />
   </ItemGroup>

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

@@ -0,0 +1,103 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+// See the LICENSE file in the project root for more information.
+
+using System.Runtime.InteropServices;
+using System.Security;
+using System.Threading.Tasks;
+
+namespace System.Runtime.CompilerServices
+{
+    /// <summary>Represents a builder for asynchronous methods that returns a <see cref="ValueTask{TResult}"/>.</summary>
+    /// <typeparam name="TResult">The type of the result.</typeparam>
+    [StructLayout(LayoutKind.Auto)]
+    public struct AsyncValueTaskMethodBuilder<TResult>
+    {
+        /// <summary>The <see cref="AsyncTaskMethodBuilder{TResult}"/> to which most operations are delegated.</summary>
+        private AsyncTaskMethodBuilder<TResult> _methodBuilder;
+        /// <summary>The result for this builder, if it's completed before any awaits occur.</summary>
+        private TResult _result;
+        /// <summary>true if <see cref="_result"/> contains the synchronous result for the async method; otherwise, false.</summary>
+        private bool _haveResult;
+        /// <summary>true if the builder should be used for setting/getting the result; otherwise, false.</summary>
+        private bool _useBuilder;
+
+        /// <summary>Creates an instance of the <see cref="AsyncValueTaskMethodBuilder{TResult}"/> struct.</summary>
+        /// <returns>The initialized instance.</returns>
+        public static AsyncValueTaskMethodBuilder<TResult> Create() =>
+            new AsyncValueTaskMethodBuilder<TResult>() { _methodBuilder = AsyncTaskMethodBuilder<TResult>.Create() };
+
+        /// <summary>Begins running the builder with the associated state machine.</summary>
+        /// <typeparam name="TStateMachine">The type of the state machine.</typeparam>
+        /// <param name="stateMachine">The state machine instance, passed by reference.</param>
+        public void Start<TStateMachine>(ref TStateMachine stateMachine) where TStateMachine : IAsyncStateMachine
+        {
+            _methodBuilder.Start(ref stateMachine); // will provide the right ExecutionContext semantics
+        }
+
+        /// <summary>Associates the builder with the specified state machine.</summary>
+        /// <param name="stateMachine">The state machine instance to associate with the builder.</param>
+        public void SetStateMachine(IAsyncStateMachine stateMachine) => _methodBuilder.SetStateMachine(stateMachine);
+
+        /// <summary>Marks the task as successfully completed.</summary>
+        /// <param name="result">The result to use to complete the task.</param>
+        public void SetResult(TResult result)
+        {
+            if (_useBuilder)
+            {
+                _methodBuilder.SetResult(result);
+            }
+            else
+            {
+                _result = result;
+                _haveResult = true;
+            }
+        }
+
+        /// <summary>Marks the task as failed and binds the specified exception to the task.</summary>
+        /// <param name="exception">The exception to bind to the task.</param>
+        public void SetException(Exception exception) => _methodBuilder.SetException(exception);
+
+        /// <summary>Gets the task for this builder.</summary>
+        public ValueTask<TResult> Task
+        {
+            get
+            {
+                if (_haveResult)
+                {
+                    return new ValueTask<TResult>(_result);
+                }
+                else
+                {
+                    _useBuilder = true;
+                    return new ValueTask<TResult>(_methodBuilder.Task);
+                }
+            }
+        }
+
+        /// <summary>Schedules the state machine to proceed to the next action when the specified awaiter completes.</summary>
+        /// <typeparam name="TAwaiter">The type of the awaiter.</typeparam>
+        /// <typeparam name="TStateMachine">The type of the state machine.</typeparam>
+        /// <param name="awaiter">the awaiter</param>
+        /// <param name="stateMachine">The state machine.</param>
+        public void AwaitOnCompleted<TAwaiter, TStateMachine>(ref TAwaiter awaiter, ref TStateMachine stateMachine)
+            where TAwaiter : INotifyCompletion
+            where TStateMachine : IAsyncStateMachine
+        {
+            _methodBuilder.AwaitOnCompleted(ref awaiter, ref stateMachine);
+        }
+
+        /// <summary>Schedules the state machine to proceed to the next action when the specified awaiter completes.</summary>
+        /// <typeparam name="TAwaiter">The type of the awaiter.</typeparam>
+        /// <typeparam name="TStateMachine">The type of the state machine.</typeparam>
+        /// <param name="awaiter">the awaiter</param>
+        /// <param name="stateMachine">The state machine.</param>
+        [SecuritySafeCritical]
+        public void AwaitUnsafeOnCompleted<TAwaiter, TStateMachine>(ref TAwaiter awaiter, ref TStateMachine stateMachine)
+            where TAwaiter : ICriticalNotifyCompletion 
+            where TStateMachine : IAsyncStateMachine
+        {
+            _methodBuilder.AwaitUnsafeOnCompleted(ref awaiter, ref stateMachine);
+        }
+    }
+}

src/System.Threading.Tasks.Extensions/src/System/Threading/Tasks/ValueTask.cs

@@ -3,6 +3,7 @@
 // See the LICENSE file in the project root for more information.
 
 using System.Collections.Generic;
+using System.ComponentModel;
 using System.Runtime.CompilerServices;
 using System.Runtime.InteropServices;
 
@@ -174,5 +175,10 @@ public override string ToString()
                     string.Empty;
             }
         }
+
+        /// <summary>Creates a method builder for use with an async method.</summary>
+        /// <returns>The created builder.</returns>
+        [EditorBrowsable(EditorBrowsableState.Never)] // intended only for compiler consumption
+        public static AsyncValueTaskMethodBuilder<TResult> CreateAsyncMethodBuilder() => AsyncValueTaskMethodBuilder<TResult>.Create();
     }
 }

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

@@ -0,0 +1,165 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+// See the LICENSE file in the project root for more information.
+
+using System.Runtime.CompilerServices;
+using Xunit;
+
+namespace System.Threading.Tasks.Tests
+{
+    public class AsyncValueTaskMethodBuilderTests
+    {
+        [Fact]
+        public void Create_ReturnsDefaultInstance()
+        {
+            AsyncValueTaskMethodBuilder<int> b = ValueTask<int>.CreateAsyncMethodBuilder();
+            Assert.Equal(default(AsyncValueTaskMethodBuilder<int>), b); // implementation detail being verified
+        }
+
+        [Fact]
+        public void SetResult_BeforeAccessTask_ValueTaskContainsValue()
+        {
+            AsyncValueTaskMethodBuilder<int> b = ValueTask<int>.CreateAsyncMethodBuilder();
+            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.Equal(42, vt.Result);
+        }
+
+        [Fact]
+        public void SetResult_AfterAccessTask_ValueTaskContainsValue()
+        {
+            AsyncValueTaskMethodBuilder<int> b = ValueTask<int>.CreateAsyncMethodBuilder();
+            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.Equal(42, vt.Result);
+        }
+
+        [Fact]
+        public void SetException_BeforeAccessTask_FaultsTask()
+        {
+            AsyncValueTaskMethodBuilder<int> b = ValueTask<int>.CreateAsyncMethodBuilder();
+            var e = new FormatException();
+            b.SetException(e);
+            ValueTask<int> vt = b.Task;
+            Assert.True(vt.IsFaulted);
+            Assert.Same(e, Assert.Throws<FormatException>(() => vt.GetAwaiter().GetResult()));
+        }
+
+        [Fact]
+        public void SetException_AfterAccessTask_FaultsTask()
+        {
+            AsyncValueTaskMethodBuilder<int> b = ValueTask<int>.CreateAsyncMethodBuilder();
+            var e = new FormatException();
+            ValueTask<int> vt = b.Task;
+            b.SetException(e);
+            Assert.True(vt.IsFaulted);
+            Assert.Same(e, Assert.Throws<FormatException>(() => vt.GetAwaiter().GetResult()));
+        }
+
+        [Fact]
+        public void SetException_OperationCanceledException_CancelsTask()
+        {
+            AsyncValueTaskMethodBuilder<int> b = ValueTask<int>.CreateAsyncMethodBuilder();
+            var e = new OperationCanceledException();
+            ValueTask<int> vt = b.Task;
+            b.SetException(e);
+            Assert.True(vt.IsCanceled);
+            Assert.Same(e, Assert.Throws<OperationCanceledException>(() => vt.GetAwaiter().GetResult()));
+        }
+
+        [Fact]
+        public void Start_InvokesMoveNext()
+        {
+            AsyncValueTaskMethodBuilder<int> b = ValueTask<int>.CreateAsyncMethodBuilder();
+            int invokes = 0;
+            var dsm = new DelegateStateMachine { MoveNextDelegate = () => invokes++ };
+            b.Start(ref dsm);
+            Assert.Equal(1, invokes);
+        }
+
+        [Theory]
+        [InlineData(false)]
+        [InlineData(true)]
+        public async Task AwaitOnCompleted_InvokesStateMachineMethods(bool awaitUnsafe)
+        {
+            AsyncValueTaskMethodBuilder<int> b = ValueTask<int>.CreateAsyncMethodBuilder();
+            var ignored = b.Task;
+
+            var callbackCompleted = new TaskCompletionSource<bool>();
+            IAsyncStateMachine foundSm = null;
+            var dsm = new DelegateStateMachine
+            {
+                MoveNextDelegate = () => callbackCompleted.SetResult(true),
+                SetStateMachineDelegate = sm => foundSm = sm
+            };
+
+            TaskAwaiter t = Task.CompletedTask.GetAwaiter();
+            if (awaitUnsafe)
+            {
+                b.AwaitUnsafeOnCompleted(ref t, ref dsm);
+            }
+            else
+            {
+                b.AwaitOnCompleted(ref t, ref dsm);
+            }
+
+            await callbackCompleted.Task;
+            Assert.Equal(dsm, foundSm);
+        }
+
+        [Fact]
+        public void SetStateMachine_InvalidArgument_ThrowsException()
+        {
+            AsyncValueTaskMethodBuilder<int> b = ValueTask<int>.CreateAsyncMethodBuilder();
+            Assert.Throws<ArgumentNullException>("stateMachine", () => b.SetStateMachine(null));
+            b.SetStateMachine(new DelegateStateMachine());
+        }
+
+        [Fact]
+        public void Start_ExecutionContextChangesInMoveNextDontFlowOut()
+        {
+            var al = new AsyncLocal<int> { Value = 0 };
+            int calls = 0;
+
+            var dsm = new DelegateStateMachine
+            {
+                MoveNextDelegate = () =>
+                {
+                    al.Value++;
+                    calls++;
+                }
+            };
+
+            dsm.MoveNext();
+            Assert.Equal(1, al.Value);
+            Assert.Equal(1, calls);
+
+            dsm.MoveNext();
+            Assert.Equal(2, al.Value);
+            Assert.Equal(2, calls);
+
+            AsyncValueTaskMethodBuilder<int> b = ValueTask<int>.CreateAsyncMethodBuilder();
+            b.Start(ref dsm);
+            Assert.Equal(2, al.Value); // change should not be visible
+            Assert.Equal(3, calls);
+
+            // 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());
+        }
+
+        private struct DelegateStateMachine : IAsyncStateMachine
+        {
+            internal Action MoveNextDelegate;
+            public void MoveNext() => MoveNextDelegate?.Invoke();
+
+            internal Action<IAsyncStateMachine> SetStateMachineDelegate;
+            public void SetStateMachine(IAsyncStateMachine stateMachine) => SetStateMachineDelegate?.Invoke(stateMachine);
+        }
+    }
+}

src/System.Threading.Tasks.Extensions/tests/System.Threading.Tasks.Extensions.Tests.csproj

@@ -1,4 +1,4 @@
-<?xml version="1.0" encoding="utf-8"?>
+<?xml version="1.0" encoding="utf-8"?>
 <Project ToolsVersion="12.0" DefaultTargets="Build" xmlns="http://schemas.microsoft.com/developer/msbuild/2003">
   <Import Project="$([MSBuild]::GetDirectoryNameOfFileAbove($(MSBuildThisFileDirectory), dir.props))\dir.props" />
   <PropertyGroup>
@@ -13,6 +13,7 @@
   <PropertyGroup Condition="'$(Configuration)|$(Platform)' == 'Debug|AnyCPU'" />
   <PropertyGroup Condition="'$(Configuration)|$(Platform)' == 'Release|AnyCPU'" />
   <ItemGroup>
+    <Compile Include="AsyncValueTaskMethodBuilderTests.cs" />
     <Compile Include="ValueTaskTests.cs" />
   </ItemGroup>
   <ItemGroup>
@@ -25,4 +26,4 @@
     <None Include="project.json" />
   </ItemGroup>
   <Import Project="$([MSBuild]::GetDirectoryNameOfFileAbove($(MSBuildThisFileDirectory), dir.targets))\dir.targets" />
-</Project>
+</Project>