update async code

AndyAyersMS · Oct 24, 2017 · 435eb72 · 435eb72
1 parent 40f02f7
commit 435eb72
Show file tree

Hide file tree

Showing 4 changed files with 35 additions and 55 deletions.
diff --git a/src/mscorlib/shared/System/Runtime/CompilerServices/ConfiguredValueTaskAwaitable.cs b/src/mscorlib/shared/System/Runtime/CompilerServices/ConfiguredValueTaskAwaitable.cs
@@ -34,7 +34,7 @@ public ConfiguredValueTaskAwaiter GetAwaiter() =>
 
         /// <summary>Provides an awaiter for a <see cref="ConfiguredValueTaskAwaitable{TResult}"/>.</summary>
         [StructLayout(LayoutKind.Auto)]
-        public struct ConfiguredValueTaskAwaiter : ICriticalNotifyCompletion
+        public struct ConfiguredValueTaskAwaiter : ICriticalNotifyCompletion, IValueTaskAwaiter
         {
             /// <summary>The value being awaited.</summary>
             private ValueTask<TResult> _value; // Methods are called on this; avoid making it readonly so as to avoid unnecessary copies
@@ -69,6 +69,12 @@ public void UnsafeOnCompleted(Action continuation) =>
 
             /// <summary>Gets the task underlying <see cref="_value"/>.</summary>
             internal Task<TResult> AsTask() => _value.AsTask();
+
+            public Task GetTask(ref bool continueOnCapturedContent)
+            {
+                continueOnCapturedContent = _continueOnCapturedContext;
+                return AsTask();
+            }
         }
     }
 }
diff --git a/src/mscorlib/shared/System/Runtime/CompilerServices/ValueTaskAwaiter.cs b/src/mscorlib/shared/System/Runtime/CompilerServices/ValueTaskAwaiter.cs
@@ -8,7 +8,7 @@
 namespace System.Runtime.CompilerServices
 {
     /// <summary>Provides an awaiter for a <see cref="ValueTask{TResult}"/>.</summary>
-    public struct ValueTaskAwaiter<TResult> : ICriticalNotifyCompletion
+    public struct ValueTaskAwaiter<TResult> : ICriticalNotifyCompletion, IValueTaskAwaiter
     {
         /// <summary>The value being awaited.</summary>
         private ValueTask<TResult> _value; // Methods are called on this; avoid making it readonly so as to avoid unnecessary copies
@@ -20,7 +20,7 @@ public struct ValueTaskAwaiter<TResult> : ICriticalNotifyCompletion
         /// <summary>Gets whether the <see cref="ValueTask{TResult}"/> has completed.</summary>
         public bool IsCompleted => _value.IsCompleted;
 
-        /// <summary>Gets the result of the ValueTask.</summary>
+        /// <summary>Gets the result of the ValueTask.t</summary>
         public TResult GetResult() =>
             _value._task == null ? 
                 _value._result : 
@@ -36,5 +36,11 @@ public void UnsafeOnCompleted(Action continuation) =>
 
         /// <summary>Gets the task underlying <see cref="_value"/>.</summary>
         internal Task<TResult> AsTask() => _value.AsTask();
+
+        public Task GetTask(ref bool continueOnCapturedContent)
+        {
+            continueOnCapturedContent = true;
+            return AsTask();
+        }
     }
 }
diff --git a/src/mscorlib/src/System/Runtime/CompilerServices/AsyncMethodBuilder.cs b/src/mscorlib/src/System/Runtime/CompilerServices/AsyncMethodBuilder.cs
@@ -393,7 +393,7 @@ public void AwaitUnsafeOnCompleted<TAwaiter, TStateMachine>(
         {
             IAsyncStateMachineBox box = GetStateMachineBox(ref stateMachine);
 
-            // TThe null tests here ensure that the jit can optimize away the interface
+            // The null tests here ensure that the jit can optimize away the interface
             // tests when TAwaiter is is a ref type.
             if ((null != (object)default(TAwaiter)) && (awaiter is ITaskAwaiter))
             {
@@ -405,62 +405,23 @@ public void AwaitUnsafeOnCompleted<TAwaiter, TStateMachine>(
                 ref ConfiguredTaskAwaitable.ConfiguredTaskAwaiter ta = ref Unsafe.As<TAwaiter, ConfiguredTaskAwaitable.ConfiguredTaskAwaiter>(ref awaiter);
                 TaskAwaiter.UnsafeOnCompletedInternal(ta.m_task, box, ta.m_continueOnCapturedContext);
             }
-
-            // Handle common {Configured}ValueTaskAwaiter<T> types.  Unfortunately these need to be special-cased
-            // individually, as we don't have good way to extract the task from a ValueTaskAwaiter<T> when we don't
-            // know what the T is; we could make ValueTaskAwaiter<T> implement an IValueTaskAwaiter interface, but
-            // calling a GetTask method on that would end up boxing the awaiter.  This hard-coded list here is
-            // somewhat arbitrary and is based on types currently in use with ValueTask<T> in coreclr/corefx.
-            else if (typeof(TAwaiter) == typeof(ValueTaskAwaiter<int>))
-            {
-                var vta = (ValueTaskAwaiter<int>)(object)awaiter;
-                TaskAwaiter.UnsafeOnCompletedInternal(vta.AsTask(), box, continueOnCapturedContext: true);
-            }
-            else if (typeof(TAwaiter) == typeof(ConfiguredValueTaskAwaitable<int>.ConfiguredValueTaskAwaiter))
-            {
-                var vta = (ConfiguredValueTaskAwaitable<int>.ConfiguredValueTaskAwaiter)(object)awaiter;
-                TaskAwaiter.UnsafeOnCompletedInternal(vta.AsTask(), box, vta._continueOnCapturedContext);
-            }
-            else if (typeof(TAwaiter) == typeof(ConfiguredValueTaskAwaitable<System.IO.Stream>.ConfiguredValueTaskAwaiter))
-            {
-                var vta = (ConfiguredValueTaskAwaitable<System.IO.Stream>.ConfiguredValueTaskAwaiter)(object)awaiter;
-                TaskAwaiter.UnsafeOnCompletedInternal(vta.AsTask(), box, vta._continueOnCapturedContext);
-            }
-            else if (typeof(TAwaiter) == typeof(ConfiguredValueTaskAwaitable<ArraySegment<byte>>.ConfiguredValueTaskAwaiter))
+            else if ((null != (object)default(TAwaiter)) && (awaiter is IValueTaskAwaiter))
             {
-                var vta = (ConfiguredValueTaskAwaitable<ArraySegment<byte>>.ConfiguredValueTaskAwaiter)(object)awaiter;
-                TaskAwaiter.UnsafeOnCompletedInternal(vta.AsTask(), box, vta._continueOnCapturedContext);
+                bool continueOnCapturedContent = true;
+                Task task = ((IValueTaskAwaiter)awaiter).GetTask(ref continueOnCapturedContent);
+                TaskAwaiter.UnsafeOnCompletedInternal(task, box, continueOnCapturedContent);
             }
-            else if (typeof(TAwaiter) == typeof(ConfiguredValueTaskAwaitable<object>.ConfiguredValueTaskAwaiter))
-            {
-                var vta = (ConfiguredValueTaskAwaitable<object>.ConfiguredValueTaskAwaiter)(object)awaiter;
-                TaskAwaiter.UnsafeOnCompletedInternal(vta.AsTask(), box, vta._continueOnCapturedContext);
-            }
-
             // The awaiter isn't specially known. Fall back to doing a normal await.
             else
             {
-                // TODO https://github.com/dotnet/coreclr/issues/14177:
-                // Move the code back into this method once the JIT is able to
-                // elide it successfully when one of the previous branches is hit.
-                AwaitArbitraryAwaiterUnsafeOnCompleted(ref awaiter, box);
-            }
-        }
-
-        /// <summary>Schedules the specified state machine to be pushed forward when the specified awaiter completes.</summary>
-        /// <typeparam name="TAwaiter">Specifies the type of the awaiter.</typeparam>
-        /// <param name="awaiter">The awaiter.</param>
-        /// <param name="box">The state machine box.</param>
-        private static void AwaitArbitraryAwaiterUnsafeOnCompleted<TAwaiter>(ref TAwaiter awaiter, IAsyncStateMachineBox box)
-            where TAwaiter : ICriticalNotifyCompletion
-        {
-            try
-            {
-                awaiter.UnsafeOnCompleted(box.MoveNextAction);
-            }
-            catch (Exception e)
-            {
-                AsyncMethodBuilderCore.ThrowAsync(e, targetContext: null);
+                try
+                {
+                    awaiter.UnsafeOnCompleted(box.MoveNextAction);
+                }
+                catch (Exception e)
+                {
+                    AsyncMethodBuilderCore.ThrowAsync(e, targetContext: null);
+                }
             }
         }
 

diff --git a/src/mscorlib/src/System/Runtime/CompilerServices/TaskAwaiter.cs b/src/mscorlib/src/System/Runtime/CompilerServices/TaskAwaiter.cs
@@ -386,6 +386,13 @@ internal interface ITaskAwaiter { }
     /// </summary>
     internal interface IConfiguredTaskAwaiter { }
 
+    /// <summary>
+    /// Interface used by various ValueTaskAwaiters. 
+    /// </summary>
+    internal interface IValueTaskAwaiter {
+        Task GetTask(ref bool continueOnCapturedContent);
+    }
+
     /// <summary>Provides an awaitable object that allows for configured awaits on <see cref="System.Threading.Tasks.Task"/>.</summary>
     /// <remarks>This type is intended for compiler use only.</remarks>
     public struct ConfiguredTaskAwaitable