Fix for configured ValueTask sources

VSadov · VSadov · commit 5ad464f207e7 · 2025-10-13T14:45:20.000-07:00
diff --git a/src/coreclr/System.Private.CoreLib/src/System/Runtime/CompilerServices/AsyncHelpers.CoreCLR.cs b/src/coreclr/System.Private.CoreLib/src/System/Runtime/CompilerServices/AsyncHelpers.CoreCLR.cs
@@ -10,6 +10,7 @@
 using System.Runtime.Versioning;
 using System.Threading;
 using System.Threading.Tasks;
+using System.Threading.Tasks.Sources;
 
 namespace System.Runtime.CompilerServices
 {
@@ -491,6 +492,49 @@ public static void HandleSuspended<T, TOps>(T task) where T : Task, ITaskComplet
                     }
                     else if (calledTask != null)
                     {
+                        if (calledTask is IValueTaskAsTask vtTask)
+                        {
+                            if (!vtTask.IsConfigured)
+                            {
+                                // ValueTaskSource can be configured to use scheduling context or to ignore it.
+                                // The awaiter must inform the source whether it wants to continue on a context,
+                                // but the source may decide to ignore the suggestion. Since the behavior of
+                                // the source takes precedence, we clear the context flags from the awaiting
+                                // continuation (so it will run transparently on what source decides) and tell
+                                // the source that awaiting frame prefers to continue on a context.
+                                // The reason why we do it here is because the continuation chain builds from the
+                                // innermost frame out and when the leaf thunk links the head continuation,
+                                // it does not know if the caller wants to continue in a context. Thus the thunk
+                                // creates an "unconfigured" IValueTaskAsTask and we configure it here.
+                                ValueTaskSourceOnCompletedFlags configFlags = ValueTaskSourceOnCompletedFlags.None;
+                                CorInfoContinuationFlags continuationFlags = headContinuation.Next!.Flags;
+
+                                const CorInfoContinuationFlags continueOnContextFlags =
+                                    CorInfoContinuationFlags.CORINFO_CONTINUATION_CONTINUE_ON_CAPTURED_SYNCHRONIZATION_CONTEXT |
+                                    CorInfoContinuationFlags.CORINFO_CONTINUATION_CONTINUE_ON_CAPTURED_TASK_SCHEDULER;
+
+                                if ((continuationFlags & continueOnContextFlags) != 0)
+                                {
+                                    // effectively move context flags from headContinuation to the source config.
+                                    configFlags |= ValueTaskSourceOnCompletedFlags.UseSchedulingContext;
+                                    headContinuation.Next!.Flags &= ~continueOnContextFlags;
+                                }
+
+                                vtTask.Configure(configFlags);
+
+                                if (!calledTask.TryAddCompletionAction(task))
+                                {
+                                    // calledTask has already completed and we need to schedule
+                                    // our code for execution ourselves.
+                                    // Restore the continuation flags before doing that.
+                                    headContinuation.Next!.Flags = continuationFlags;
+                                    ThreadPool.UnsafeQueueUserWorkItemInternal(task, preferLocal: true);
+                                }
+
+                                return;
+                            }
+                        }
+
                         // Runtime async callable wrapper for task returning
                         // method. This implements the context transparent
                         // forwarding and makes these wrappers minimal cost.
diff --git a/src/coreclr/vm/asyncthunks.cpp b/src/coreclr/vm/asyncthunks.cpp
@@ -591,7 +591,7 @@ void MethodDesc::EmitAsyncMethodThunk(MethodDesc* pAsyncOtherVariant, MetaSig& m
 
             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* pMDAsTask = CoreLibBinder::GetMethod(METHOD__VALUETASK__AS_UNCONFIGURED_TASK);
 
             isCompletedToken = pCode->GetToken(pMDValueTaskIsCompleted);
             completionResultToken = pCode->GetToken(pMDCompletionResult);
@@ -604,7 +604,7 @@ 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* pMDAsTask = CoreLibBinder::GetMethod(METHOD__VALUETASK_1__AS_UNCONFIGURED_TASK);
 
             pMDValueTaskIsCompleted = FindOrCreateAssociatedMethodDesc(pMDValueTaskIsCompleted, pMTValueTask, FALSE, Instantiation(), FALSE);
             pMDCompletionResult = FindOrCreateAssociatedMethodDesc(pMDCompletionResult, pMTValueTask, FALSE, Instantiation(), FALSE);
diff --git a/src/coreclr/vm/corelib.h b/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_UNCONFIGURED_TASK, AsUnconfiguredTask, IM_RetTaskOfT)
 
 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_UNCONFIGURED_TASK, AsUnconfiguredTask, IM_RetTask)
 
 DEFINE_CLASS(TASK_1, Tasks, Task`1)
 DEFINE_METHOD(TASK_1, GET_RESULTONSUCCESS, get_ResultOnSuccess, NoSig)
diff --git a/src/libraries/System.Private.CoreLib/src/System/Threading/Tasks/ValueTask.cs b/src/libraries/System.Private.CoreLib/src/System/Threading/Tasks/ValueTask.cs
@@ -176,7 +176,17 @@ public Task AsTask()
             return
                 obj == null ? Task.CompletedTask :
                 obj as Task ??
-                GetTaskForValueTaskSource(Unsafe.As<IValueTaskSource>(obj));
+                GetTaskForValueTaskSource(Unsafe.As<IValueTaskSource>(obj), ValueTaskSourceOnCompletedFlags.None);
+        }
+
+        internal Task AsUnconfiguredTask()
+        {
+            object? obj = _obj;
+            Debug.Assert(obj == null || obj is Task || obj is IValueTaskSource);
+            return
+                obj == null ? Task.CompletedTask :
+                obj as Task ??
+                GetTaskForValueTaskSource(Unsafe.As<IValueTaskSource>(obj), IValueTaskAsTask._unconfigured);
         }
 
         /// <summary>Gets a <see cref="ValueTask"/> that may be used at any point in the future.</summary>
@@ -187,7 +197,7 @@ obj as Task ??
         /// The <see cref="IValueTaskSource"/> is passed in rather than reading and casting <see cref="_obj"/>
         /// so that the caller can pass in an object it's already validated.
         /// </remarks>
-        private Task GetTaskForValueTaskSource(IValueTaskSource t)
+        private Task GetTaskForValueTaskSource(IValueTaskSource t, ValueTaskSourceOnCompletedFlags flags)
         {
             ValueTaskSourceStatus status = t.GetStatus(_token);
             if (status != ValueTaskSourceStatus.Pending)
@@ -225,11 +235,11 @@ private Task GetTaskForValueTaskSource(IValueTaskSource t)
                 }
             }
 
-            return new ValueTaskSourceAsTask(t, _token);
+            return new ValueTaskSourceAsTask(t, _token, flags);
         }
 
         /// <summary>Type used to create a <see cref="Task"/> to represent a <see cref="IValueTaskSource"/>.</summary>
-        private sealed class ValueTaskSourceAsTask : Task
+        private sealed class ValueTaskSourceAsTask : Task, IValueTaskAsTask
         {
             private static readonly Action<object?> s_completionAction = static state =>
             {
@@ -283,11 +293,28 @@ state is ValueTaskSourceAsTask vsts ?
             /// <summary>The token to pass through to operations on <see cref="_source"/></summary>
             private readonly short _token;
 
-            internal ValueTaskSourceAsTask(IValueTaskSource source, short token)
+            private bool _configured;
+
+            bool IValueTaskAsTask.IsConfigured => _configured;
+            void IValueTaskAsTask.Configure(ValueTaskSourceOnCompletedFlags flags) => Configure(flags);
+
+            private void Configure(ValueTaskSourceOnCompletedFlags flags)
+            {
+                Debug.Assert(!_configured);
+
+                _configured = true;
+                _source!.OnCompleted(s_completionAction, this, _token, flags);
+            }
+
+            internal ValueTaskSourceAsTask(IValueTaskSource source, short token, ValueTaskSourceOnCompletedFlags flags)
             {
                 _token = token;
                 _source = source;
-                source.OnCompleted(s_completionAction, this, token, ValueTaskSourceOnCompletedFlags.None);
+
+                if (flags != IValueTaskAsTask._unconfigured)
+                {
+                    Configure(flags);
+                }
             }
         }
 
@@ -585,7 +612,25 @@ public Task<TResult> AsTask()
                 return t;
             }
 
-            return GetTaskForValueTaskSource(Unsafe.As<IValueTaskSource<TResult>>(obj));
+            return GetTaskForValueTaskSource(Unsafe.As<IValueTaskSource<TResult>>(obj), ValueTaskSourceOnCompletedFlags.None);
+        }
+
+        internal Task<TResult> AsUnconfiguredTask()
+        {
+            object? obj = _obj;
+            Debug.Assert(obj == null || obj is Task<TResult> || obj is IValueTaskSource<TResult>);
+
+            if (obj == null)
+            {
+                return Task.FromResult(_result!);
+            }
+
+            if (obj is Task<TResult> t)
+            {
+                return t;
+            }
+
+            return GetTaskForValueTaskSource(Unsafe.As<IValueTaskSource<TResult>>(obj), IValueTaskAsTask._unconfigured);
         }
 
         /// <summary>Gets a <see cref="ValueTask{TResult}"/> that may be used at any point in the future.</summary>
@@ -596,7 +641,7 @@ public Task<TResult> AsTask()
         /// The <see cref="IValueTaskSource{TResult}"/> is passed in rather than reading and casting <see cref="_obj"/>
         /// so that the caller can pass in an object it's already validated.
         /// </remarks>
-        private Task<TResult> GetTaskForValueTaskSource(IValueTaskSource<TResult> t)
+        private Task<TResult> GetTaskForValueTaskSource(IValueTaskSource<TResult> t, ValueTaskSourceOnCompletedFlags flags)
         {
             ValueTaskSourceStatus status = t.GetStatus(_token);
             if (status != ValueTaskSourceStatus.Pending)
@@ -633,11 +678,11 @@ private Task<TResult> GetTaskForValueTaskSource(IValueTaskSource<TResult> t)
                 }
             }
 
-            return new ValueTaskSourceAsTask(t, _token);
+            return new ValueTaskSourceAsTask(t, _token, flags);
         }
 
         /// <summary>Type used to create a <see cref="Task{TResult}"/> to represent a <see cref="IValueTaskSource{TResult}"/>.</summary>
-        private sealed class ValueTaskSourceAsTask : Task<TResult>
+        private sealed class ValueTaskSourceAsTask : Task<TResult>, IValueTaskAsTask
         {
             private static readonly Action<object?> s_completionAction = static state =>
             {
@@ -690,11 +735,28 @@ state is ValueTaskSourceAsTask vsts ?
             /// <summary>The token to pass through to operations on <see cref="_source"/></summary>
             private readonly short _token;
 
-            public ValueTaskSourceAsTask(IValueTaskSource<TResult> source, short token)
+            private bool _configured;
+
+            bool IValueTaskAsTask.IsConfigured => _configured;
+            void IValueTaskAsTask.Configure(ValueTaskSourceOnCompletedFlags flags) => Configure(flags);
+
+            private void Configure(ValueTaskSourceOnCompletedFlags flags)
+            {
+                Debug.Assert(!_configured);
+
+                _configured = true;
+                _source!.OnCompleted(s_completionAction, this, _token, flags);
+            }
+
+            internal ValueTaskSourceAsTask(IValueTaskSource<TResult> source, short token, ValueTaskSourceOnCompletedFlags flags)
             {
-                _source = source;
                 _token = token;
-                source.OnCompleted(s_completionAction, this, token, ValueTaskSourceOnCompletedFlags.None);
+                _source = source;
+
+                if (flags != IValueTaskAsTask._unconfigured)
+                {
+                    Configure(flags);
+                }
             }
         }
 
@@ -848,4 +910,12 @@ public ConfiguredValueTaskAwaitable<TResult> ConfigureAwait(bool continueOnCaptu
             return string.Empty;
         }
     }
+
+    internal interface IValueTaskAsTask
+    {
+        internal const ValueTaskSourceOnCompletedFlags _unconfigured = (ValueTaskSourceOnCompletedFlags)(-1);
+
+        bool IsConfigured { get; }
+        void Configure(ValueTaskSourceOnCompletedFlags flags);
+    }
 }
