[RuntimeAsync] Fix for configured ValueTask sources

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);
         }
 
@@ -192,6 +216,7 @@ private interface IRuntimeAsyncTaskOps<T>
             static abstract void SetContinuationState(T task, Continuation value);
             static abstract bool SetCompleted(T task);
             static abstract void PostToSyncContext(T task, SynchronizationContext syncCtx);
+            static abstract void ValueTaskSourceOnCompleted(T task, IValueTaskSourceNotifier vtsNotifier, ValueTaskSourceOnCompletedFlags configFlags);
             static abstract ref byte GetResultStorage(T task);
         }
 
@@ -238,6 +263,12 @@ void ITaskCompletionAction.Invoke(Task completingTask)
                 ((RuntimeAsyncTask<T>)state).MoveNext();
             };
 
+            public static readonly Action<object?> s_runContinuationAction = static state =>
+            {
+                Debug.Assert(state is RuntimeAsyncTask<T>);
+                ((RuntimeAsyncTask<T>)state).MoveNext();
+            };
+
             private struct Ops : IRuntimeAsyncTaskOps<RuntimeAsyncTask<T>>
             {
                 public static Action GetContinuationAction(RuntimeAsyncTask<T> task) => (Action)task.m_action!;
@@ -257,6 +288,11 @@ public static void PostToSyncContext(RuntimeAsyncTask<T> task, SynchronizationCo
                     syncContext.Post(s_postCallback, task);
                 }
 
+                public static void ValueTaskSourceOnCompleted(RuntimeAsyncTask<T> task, IValueTaskSourceNotifier vtsNotifier, ValueTaskSourceOnCompletedFlags configFlags)
+                {
+                    vtsNotifier.OnCompleted(s_runContinuationAction, task, configFlags);
+                }
+
                 public static ref byte GetResultStorage(RuntimeAsyncTask<T> task) => ref Unsafe.As<T?, byte>(ref task.m_result);
             }
         }
@@ -304,6 +340,12 @@ void ITaskCompletionAction.Invoke(Task completingTask)
                 ((RuntimeAsyncTask)state).MoveNext();
             };
 
+            public static readonly Action<object?> s_runContinuationAction = static state =>
+            {
+                Debug.Assert(state is RuntimeAsyncTask);
+                ((RuntimeAsyncTask)state).MoveNext();
+            };
+
             private struct Ops : IRuntimeAsyncTaskOps<RuntimeAsyncTask>
             {
                 public static Action GetContinuationAction(RuntimeAsyncTask task) => (Action)task.m_action!;
@@ -323,6 +365,11 @@ public static void PostToSyncContext(RuntimeAsyncTask task, SynchronizationConte
                     syncContext.Post(s_postCallback, task);
                 }
 
+                public static void ValueTaskSourceOnCompleted(RuntimeAsyncTask task, IValueTaskSourceNotifier vtsNotifier, ValueTaskSourceOnCompletedFlags configFlags)
+                {
+                    vtsNotifier.OnCompleted(s_runContinuationAction, task, configFlags);
+                }
+
                 public static ref byte GetResultStorage(RuntimeAsyncTask task) => ref Unsafe.NullRef<byte>();
             }
         }
@@ -419,13 +466,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 +497,43 @@ 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;
+                        TOps.ValueTaskSourceOnCompleted(task, vtsNotifier, configFlags);
+                    }
                     else
                     {
                         Debug.Assert(notifier != null);

src/coreclr/jit/importer.cpp

@@ -4432,6 +4432,15 @@ bool Compiler::impIsImplicitTailCallCandidate(
         return false;
     }
 
+    // We cannot tailcall ValueTask returning methods as we need to preserve
+    // the Continuation instance for ValueTaskSource handling (the BCL needs
+    // to look at continuation.Next). We cannot easily differentiate between
+    // ValueTask and Task here, so we just disable it more generally.
+    if ((prefixFlags & PREFIX_IS_TASK_AWAIT) != 0)
+    {
+        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