Support for ValueTaskSource

Author: VSadov

src/coreclr/System.Private.CoreLib/src/System/Runtime/CompilerServices/AsyncHelpers.CoreCLR.cs

@@ -11,6 +11,7 @@
 using System.Runtime.Versioning;
 using System.Threading;
 using System.Threading.Tasks;
+using System.Threading.Tasks.Sources;
 
 namespace System.Runtime.CompilerServices
 {
@@ -134,9 +135,14 @@ public static partial class AsyncHelpers
         private struct RuntimeAsyncAwaitState
         {
             public Continuation? SentinelContinuation;
+
+            // The following are the possible introducers of asynchrony into a chain of awaits.
+            // In other words - when we build a chain of continuations it would be logicaly attached
+            // to one of these notifiers.
             public ICriticalNotifyCompletion? CriticalNotifier;
             public INotifyCompletion? Notifier;
-            public Task? CalledTask;
+            public IValueTaskSourceNotifier? ValueTaskSourceNotifier;
+            public Task? TaskNotifier;
         }
 
         [ThreadStatic]
@@ -171,17 +177,35 @@ private static unsafe Continuation AllocContinuationClass(Continuation prevConti
             return newContinuation;
         }
 
+        /// <summary>
+        /// Used by internal thunks that implement awaiting on Task or a ValueTask.
+        /// A ValueTask may wrap:
+        /// - Completed result   (we never await this)
+        /// - Task
+        /// - ValueTaskSource
+        /// Therefore, when we are awaiting a ValueTask completion we are really
+        /// awaiting a completion of an underlying Task or ValueTaskSource.
+        /// </summary>
+        /// <param name="o"> Task or a ValueTaskNotifier whose completion we are awaiting.</param>
         [BypassReadyToRun]
         [MethodImpl(MethodImplOptions.NoInlining | MethodImplOptions.Async)]
         [RequiresPreviewFeatures]
-        private static void TransparentAwaitTask(Task t)
+        private static void TransparentAwait(object o)
         {
             ref RuntimeAsyncAwaitState state = ref t_runtimeAsyncAwaitState;
             Continuation? sentinelContinuation = state.SentinelContinuation;
             if (sentinelContinuation == null)
                 state.SentinelContinuation = sentinelContinuation = new Continuation();
 
-            state.CalledTask = t;
+            if (o is Task t)
+            {
+                state.TaskNotifier = t;
+            }
+            else
+            {
+                state.ValueTaskSourceNotifier = (IValueTaskSourceNotifier)o;
+            }
+
             AsyncSuspend(sentinelContinuation);
         }
 
@@ -419,13 +443,16 @@ public static unsafe void DispatchContinuations<T, TOps>(T task) where T : Task,
             public static void HandleSuspended<T, TOps>(T task) where T : Task, ITaskCompletionAction where TOps : IRuntimeAsyncTaskOps<T>
             {
                 ref RuntimeAsyncAwaitState state = ref t_runtimeAsyncAwaitState;
+
                 ICriticalNotifyCompletion? critNotifier = state.CriticalNotifier;
                 INotifyCompletion? notifier = state.Notifier;
-                Task? calledTask = state.CalledTask;
+                IValueTaskSourceNotifier? vtsNotifier = state.ValueTaskSourceNotifier;
+                Task? taskNotifier = state.TaskNotifier;
 
                 state.CriticalNotifier = null;
                 state.Notifier = null;
-                state.CalledTask = null;
+                state.ValueTaskSourceNotifier = null;
+                state.TaskNotifier = null;
 
                 Continuation sentinelContinuation = state.SentinelContinuation!;
                 Continuation headContinuation = sentinelContinuation.Next!;
@@ -447,16 +474,44 @@ public static void HandleSuspended<T, TOps>(T task) where T : Task, ITaskComplet
                     {
                         critNotifier.UnsafeOnCompleted(TOps.GetContinuationAction(task));
                     }
-                    else if (calledTask != null)
+                    else if (taskNotifier != null)
                     {
                         // Runtime async callable wrapper for task returning
                         // method. This implements the context transparent
                         // forwarding and makes these wrappers minimal cost.
-                        if (!calledTask.TryAddCompletionAction(task))
+                        if (!taskNotifier.TryAddCompletionAction(task))
                         {
                             ThreadPool.UnsafeQueueUserWorkItemInternal(task, preferLocal: true);
                         }
                     }
+                    else if (vtsNotifier != null)
+                    {
+                        // The awaiter must inform the ValueTaskSource source on whether the continuation
+                        // wants to run on a context, although the source may decide to ignore the suggestion.
+                        // Since the behavior of the source takes precedence, we clear the context flags of
+                        // the awaiting continuation (so it will run transparently on what the source decides)
+                        // and then tell the source if the awaiting frame prefers to continue on a context.
+                        // The reason why we do it here and not when the notifier is created is because
+                        // the continuation chain builds from the innermost frame out and at the time when the
+                        // notifier is created we do not know yet if the caller wants to continue on a context.
+                        ValueTaskSourceOnCompletedFlags configFlags = ValueTaskSourceOnCompletedFlags.None;
+                        ContinuationFlags continuationFlags = headContinuation.Next!.Flags;
+
+                        const ContinuationFlags continueOnContextFlags =
+                            ContinuationFlags.ContinueOnCapturedSynchronizationContext |
+                            ContinuationFlags.ContinueOnCapturedTaskScheduler;
+
+                        if ((continuationFlags & continueOnContextFlags) != 0)
+                        {
+                            // if await has captured some context, inform the source
+                            configFlags |= ValueTaskSourceOnCompletedFlags.UseSchedulingContext;
+                        }
+
+                        // Clear continuation flags, so that continuation runs transparently
+                        headContinuation.Next!.Flags &= ~continueFlags;
+
+                        vtsNotifier.OnCompleted(o => TOps.GetContinuationAction((T)o!).Invoke(), task, configFlags);
+                    }
                     else
                     {
                         Debug.Assert(notifier != null);

src/coreclr/jit/importer.cpp

@@ -4432,6 +4432,14 @@ bool Compiler::impIsImplicitTailCallCandidate(
         return false;
     }
 
+    // We keep finding cases where tailcalling awaits is problematic.
+    // Perhaps at the end there will be a set of conditions when tailcalling is ok.
+    // For now we will just forbid.
+    if (prefixFlags & PREFIX_IS_TASK_AWAIT)
+    {
+        return false;
+    }
+
 #if !FEATURE_TAILCALL_OPT_SHARED_RETURN
     // the block containing call is marked as BBJ_RETURN
     // We allow shared ret tail call optimization on recursive calls even under

src/coreclr/vm/asyncthunks.cpp

@@ -475,27 +475,30 @@ void MethodDesc::EmitAsyncMethodThunk(MethodDesc* pAsyncOtherVariant, MetaSig& m
     _ASSERTE(!pAsyncOtherVariant->IsVoid());
 
     // Implement IL that is effectively the following:
-    /*
-    {
-        Task task = other(arg);
-        if (!task.IsCompleted)
-        {
-            // Magic function which will suspend the current run of async methods
-            AsyncHelpers.TransparentAwaitTask(task);
-        }
-        return AsyncHelpers.CompletedTaskResult(task);
-    }
+    // {
+    //    Task task = other(arg);
+    //    if (!task.IsCompleted)
+    //    {
+    //        // Magic function which will suspend the current run of async methods
+    //        AsyncHelpers.TransparentAwait(task);
+    //    }
+    //    return AsyncHelpers.CompletedTaskResult(task);
+    // }
 
-    For ValueTask:
-    {
-        ValueTask vt = other(arg);
-        if (vt.IsCompleted)
-            return vt.Result/vt.ThrowIfCompletedUnsuccessfully();
+    // For ValueTask:
 
-        Task task = vt.AsTask();
-        <same code as above>
-    }
-    */
+    // {
+    //    ValueTask vt = other(arg);
+    //    if (!vt.IsCompleted)
+    //    {
+    //        taskOrNotifier = vt.AsTaskOrNotifier()
+
+    //        // Magic function which will suspend the current run of async methods
+    //        AsyncHelpers.TransparentAwait(taskOrNotifier);
+    //    }
+
+    //    return vt.Result/vt.ThrowIfCompletedUnsuccessfully();
+    // }
     ILCodeStream* pCode = pSL->NewCodeStream(ILStubLinker::kDispatch);
 
     int userFuncToken;
@@ -563,32 +566,27 @@ void MethodDesc::EmitAsyncMethodThunk(MethodDesc* pAsyncOtherVariant, MetaSig& m
         pCode->EmitLDARG(localArg++);
     }
 
+    // other(arg)
     pCode->EmitCALL(userFuncToken, localArg, 1);
 
     TypeHandle thLogicalRetType = msig.GetRetTypeHandleThrowing();
     if (IsValueTaskAsyncThunk())
     {
-        // Emit
-        // if (vtask.IsCompleted)
-        //   return vtask.Result/vtask.ThrowIfCompletedUnsuccessfully()
-        // task = vtask.AsTask()
-        //
-
         MethodTable* pMTValueTask;
         int isCompletedToken;
         int completionResultToken;
-        int asTaskToken;
+        int asTaskOrNotifierToken;
         if (msig.IsReturnTypeVoid())
         {
             pMTValueTask = CoreLibBinder::GetClass(CLASS__VALUETASK);
 
             MethodDesc* pMDValueTaskIsCompleted = CoreLibBinder::GetMethod(METHOD__VALUETASK__GET_ISCOMPLETED);
             MethodDesc* pMDCompletionResult = CoreLibBinder::GetMethod(METHOD__VALUETASK__THROW_IF_COMPLETED_UNSUCCESSFULLY);
-            MethodDesc* pMDAsTask = CoreLibBinder::GetMethod(METHOD__VALUETASK__AS_TASK);
+            MethodDesc* pMDAsTaskOrNotifier = CoreLibBinder::GetMethod(METHOD__VALUETASK__AS_TASK_OR_NOTIFIER);
 
             isCompletedToken = pCode->GetToken(pMDValueTaskIsCompleted);
             completionResultToken = pCode->GetToken(pMDCompletionResult);
-            asTaskToken = pCode->GetToken(pMDAsTask);
+            asTaskOrNotifierToken = pCode->GetToken(pMDAsTaskOrNotifier);
         }
         else
         {
@@ -597,73 +595,75 @@ void MethodDesc::EmitAsyncMethodThunk(MethodDesc* pAsyncOtherVariant, MetaSig& m
 
             MethodDesc* pMDValueTaskIsCompleted = CoreLibBinder::GetMethod(METHOD__VALUETASK_1__GET_ISCOMPLETED);
             MethodDesc* pMDCompletionResult = CoreLibBinder::GetMethod(METHOD__VALUETASK_1__GET_RESULT);
-            MethodDesc* pMDAsTask = CoreLibBinder::GetMethod(METHOD__VALUETASK_1__AS_TASK);
+            MethodDesc* pMDAsTaskOrNotifier = CoreLibBinder::GetMethod(METHOD__VALUETASK_1__AS_TASK_OR_NOTIFIER);
 
             pMDValueTaskIsCompleted = FindOrCreateAssociatedMethodDesc(pMDValueTaskIsCompleted, pMTValueTask, FALSE, Instantiation(), FALSE);
             pMDCompletionResult = FindOrCreateAssociatedMethodDesc(pMDCompletionResult, pMTValueTask, FALSE, Instantiation(), FALSE);
-            pMDAsTask = FindOrCreateAssociatedMethodDesc(pMDAsTask, pMTValueTask, FALSE, Instantiation(), FALSE);
+            pMDAsTaskOrNotifier = FindOrCreateAssociatedMethodDesc(pMDAsTaskOrNotifier, pMTValueTask, FALSE, Instantiation(), FALSE);
 
             isCompletedToken = GetTokenForGenericTypeMethodCallWithAsyncReturnType(pCode, pMDValueTaskIsCompleted);
             completionResultToken = GetTokenForGenericTypeMethodCallWithAsyncReturnType(pCode, pMDCompletionResult);
-            asTaskToken = GetTokenForGenericTypeMethodCallWithAsyncReturnType(pCode, pMDAsTask);
+            asTaskOrNotifierToken = GetTokenForGenericTypeMethodCallWithAsyncReturnType(pCode, pMDAsTaskOrNotifier);
         }
 
         LocalDesc valueTaskLocalDesc(pMTValueTask);
         DWORD valueTaskLocal = pCode->NewLocal(valueTaskLocalDesc);
-        ILCodeLabel* valueTaskNotCompletedLabel = pCode->NewCodeLabel();
+        ILCodeLabel* valueTaskCompletedLabel = pCode->NewCodeLabel();
 
         // Store value task returned by call to actual user func
         pCode->EmitSTLOC(valueTaskLocal);
-
         pCode->EmitLDLOCA(valueTaskLocal);
         pCode->EmitCALL(isCompletedToken, 1, 1);
-        pCode->EmitBRFALSE(valueTaskNotCompletedLabel);
+        pCode->EmitBRTRUE(valueTaskCompletedLabel);
 
         pCode->EmitLDLOCA(valueTaskLocal);
-        pCode->EmitCALL(completionResultToken, 1, msig.IsReturnTypeVoid() ? 0 : 1);
-        pCode->EmitRET();
+        pCode->EmitCALL(asTaskOrNotifierToken, 1, 1);
+        pCode->EmitCALL(METHOD__ASYNC_HELPERS__TRANSPARENT_AWAIT, 1, 0);
 
-        pCode->EmitLabel(valueTaskNotCompletedLabel);
+        pCode->EmitLabel(valueTaskCompletedLabel);
         pCode->EmitLDLOCA(valueTaskLocal);
-        pCode->EmitCALL(asTaskToken, 1, 1);
-    }
-
-    MethodTable* pMTTask;
-
-    int completedTaskResultToken;
-    if (msig.IsReturnTypeVoid())
-    {
-        pMTTask = CoreLibBinder::GetClass(CLASS__TASK);
-
-        MethodDesc* pMDCompletedTask = CoreLibBinder::GetMethod(METHOD__ASYNC_HELPERS__COMPLETED_TASK);
-        completedTaskResultToken = pCode->GetToken(pMDCompletedTask);
+        pCode->EmitCALL(completionResultToken, 1, msig.IsReturnTypeVoid() ? 0 : 1);
+        pCode->EmitRET();
     }
     else
     {
-        MethodTable* pMTTaskOpen = CoreLibBinder::GetClass(CLASS__TASK_1);
-        pMTTask = ClassLoader::LoadGenericInstantiationThrowing(pMTTaskOpen->GetModule(), pMTTaskOpen->GetCl(), Instantiation(&thLogicalRetType, 1)).GetMethodTable();
+        MethodTable* pMTTask;
 
-        MethodDesc* pMDCompletedTaskResult = CoreLibBinder::GetMethod(METHOD__ASYNC_HELPERS__COMPLETED_TASK_RESULT);
-        pMDCompletedTaskResult = FindOrCreateAssociatedMethodDesc(pMDCompletedTaskResult, pMDCompletedTaskResult->GetMethodTable(), FALSE, Instantiation(&thLogicalRetType, 1), FALSE);
-        completedTaskResultToken = GetTokenForGenericMethodCallWithAsyncReturnType(pCode, pMDCompletedTaskResult);
-    }
+        int completedTaskResultToken;
+        if (msig.IsReturnTypeVoid())
+        {
+            pMTTask = CoreLibBinder::GetClass(CLASS__TASK);
 
-    LocalDesc taskLocalDesc(pMTTask);
-    DWORD taskLocal = pCode->NewLocal(taskLocalDesc);
-    ILCodeLabel* pGetResultLabel = pCode->NewCodeLabel();
+            MethodDesc* pMDCompletedTask = CoreLibBinder::GetMethod(METHOD__ASYNC_HELPERS__COMPLETED_TASK);
+            completedTaskResultToken = pCode->GetToken(pMDCompletedTask);
+        }
+        else
+        {
+            MethodTable* pMTTaskOpen = CoreLibBinder::GetClass(CLASS__TASK_1);
+            pMTTask = ClassLoader::LoadGenericInstantiationThrowing(pMTTaskOpen->GetModule(), pMTTaskOpen->GetCl(), Instantiation(&thLogicalRetType, 1)).GetMethodTable();
 
-    // Store task returned by actual user func or by ValueTask.AsTask
-    pCode->EmitSTLOC(taskLocal);
+            MethodDesc* pMDCompletedTaskResult = CoreLibBinder::GetMethod(METHOD__ASYNC_HELPERS__COMPLETED_TASK_RESULT);
+            pMDCompletedTaskResult = FindOrCreateAssociatedMethodDesc(pMDCompletedTaskResult, pMDCompletedTaskResult->GetMethodTable(), FALSE, Instantiation(&thLogicalRetType, 1), FALSE);
+            completedTaskResultToken = GetTokenForGenericMethodCallWithAsyncReturnType(pCode, pMDCompletedTaskResult);
+        }
 
-    pCode->EmitLDLOC(taskLocal);
-    pCode->EmitCALL(METHOD__TASK__GET_ISCOMPLETED, 1, 1);
-    pCode->EmitBRTRUE(pGetResultLabel);
+        LocalDesc taskLocalDesc(pMTTask);
+        DWORD taskLocal = pCode->NewLocal(taskLocalDesc);
+        ILCodeLabel* pGetResultLabel = pCode->NewCodeLabel();
 
-    pCode->EmitLDLOC(taskLocal);
-    pCode->EmitCALL(METHOD__ASYNC_HELPERS__TRANSPARENT_AWAIT_TASK, 1, 0);
+        // Store task returned by actual user func or by ValueTask.AsTask
+        pCode->EmitSTLOC(taskLocal);
 
-    pCode->EmitLabel(pGetResultLabel);
-    pCode->EmitLDLOC(taskLocal);
-    pCode->EmitCALL(completedTaskResultToken, 1, msig.IsReturnTypeVoid() ? 0 : 1);
-    pCode->EmitRET();
+        pCode->EmitLDLOC(taskLocal);
+        pCode->EmitCALL(METHOD__TASK__GET_ISCOMPLETED, 1, 1);
+        pCode->EmitBRTRUE(pGetResultLabel);
+
+        pCode->EmitLDLOC(taskLocal);
+        pCode->EmitCALL(METHOD__ASYNC_HELPERS__TRANSPARENT_AWAIT, 1, 0);
+
+        pCode->EmitLabel(pGetResultLabel);
+        pCode->EmitLDLOC(taskLocal);
+        pCode->EmitCALL(completedTaskResultToken, 1, msig.IsReturnTypeVoid() ? 0 : 1);
+        pCode->EmitRET();
+    }
 }

src/coreclr/vm/corelib.h

@@ -346,7 +346,7 @@ DEFINE_METHOD(THREAD_START_EXCEPTION,EX_CTOR,               .ctor,
 DEFINE_CLASS(VALUETASK_1, Tasks, ValueTask`1)
 DEFINE_METHOD(VALUETASK_1, GET_ISCOMPLETED, get_IsCompleted, NoSig)
 DEFINE_METHOD(VALUETASK_1, GET_RESULT, get_Result, NoSig)
-DEFINE_METHOD(VALUETASK_1, AS_TASK, AsTask, IM_RetTaskOfT)
+DEFINE_METHOD(VALUETASK_1, AS_TASK_OR_NOTIFIER, AsTaskOrNotifier, IM_RetObj)
 
 DEFINE_CLASS(VALUETASK, Tasks, ValueTask)
 DEFINE_METHOD(VALUETASK, FROM_EXCEPTION, FromException, SM_Exception_RetValueTask)
@@ -355,7 +355,7 @@ DEFINE_METHOD(VALUETASK, FROM_RESULT_T, FromResult, GM_T_RetValueTaskOfT)
 DEFINE_METHOD(VALUETASK, GET_COMPLETED_TASK, get_CompletedTask, SM_RetValueTask)
 DEFINE_METHOD(VALUETASK, GET_ISCOMPLETED, get_IsCompleted, NoSig)
 DEFINE_METHOD(VALUETASK, THROW_IF_COMPLETED_UNSUCCESSFULLY, ThrowIfCompletedUnsuccessfully, NoSig)
-DEFINE_METHOD(VALUETASK, AS_TASK, AsTask, IM_RetTask)
+DEFINE_METHOD(VALUETASK, AS_TASK_OR_NOTIFIER, AsTaskOrNotifier, IM_RetObj)
 
 DEFINE_CLASS(TASK_1, Tasks, Task`1)
 DEFINE_METHOD(TASK_1, GET_RESULTONSUCCESS, get_ResultOnSuccess, NoSig)
@@ -741,9 +741,9 @@ DEFINE_METHOD(ASYNC_HELPERS,      VALUETASK_FROM_EXCEPTION, ValueTaskFromExcepti
 DEFINE_METHOD(ASYNC_HELPERS,      VALUETASK_FROM_EXCEPTION_1, ValueTaskFromException, GM_Exception_RetValueTaskOfT)
 
 DEFINE_METHOD(ASYNC_HELPERS,      UNSAFE_AWAIT_AWAITER_1,    UnsafeAwaitAwaiter, GM_T_RetVoid)
-DEFINE_METHOD(ASYNC_HELPERS,      TRANSPARENT_AWAIT_TASK,    TransparentAwaitTask, NoSig)
-DEFINE_METHOD(ASYNC_HELPERS,      COMPLETED_TASK_RESULT,    CompletedTaskResult, NoSig)
-DEFINE_METHOD(ASYNC_HELPERS,      COMPLETED_TASK,           CompletedTask, NoSig)
+DEFINE_METHOD(ASYNC_HELPERS,      TRANSPARENT_AWAIT,         TransparentAwait, NoSig)
+DEFINE_METHOD(ASYNC_HELPERS,      COMPLETED_TASK_RESULT,     CompletedTaskResult, NoSig)
+DEFINE_METHOD(ASYNC_HELPERS,      COMPLETED_TASK,            CompletedTask, NoSig)
 DEFINE_METHOD(ASYNC_HELPERS,      CAPTURE_EXECUTION_CONTEXT, CaptureExecutionContext, NoSig)
 DEFINE_METHOD(ASYNC_HELPERS,      RESTORE_EXECUTION_CONTEXT, RestoreExecutionContext, NoSig)
 DEFINE_METHOD(ASYNC_HELPERS,      CAPTURE_CONTINUATION_CONTEXT, CaptureContinuationContext, NoSig)