[RuntimeAsync] Skip over nonuser continuations when searching for awaiter configuration (#121976)

Fixes: #121971

Sometimes we have transparent/infrastructure code between user code and
a ValueTask-returning method that wraps IValueTaskSource. An example is
an unboxing stub.

This is the scenario where we need to figure the configuration of the
await operation by looking at continuation that corresponds to the user
code (only those are configured).
In a case if we see transparent continuations, we should just continue
walking continuation chain.

Author: VSadov

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

@@ -576,6 +576,7 @@ public static void HandleSuspended<T, TOps>(T task) where T : Task, ITaskComplet
                     ContinuationFlags.ContinueOnCapturedSynchronizationContext |
                     ContinuationFlags.ContinueOnThreadPool |
                     ContinuationFlags.ContinueOnCapturedTaskScheduler;
+
                 Debug.Assert((headContinuation.Flags & continueFlags) == 0);
 
                 TOps.SetContinuationState(task, headContinuation);
@@ -598,7 +599,7 @@ public static void HandleSuspended<T, TOps>(T task) where T : Task, ITaskComplet
                     }
                     else if (vtsNotifier != null)
                     {
-                        // The awaiter must inform the ValueTaskSource source on whether the continuation
+                        // The awaiter must inform the ValueTaskSource 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)
@@ -607,8 +608,18 @@ public static void HandleSuspended<T, TOps>(T task) where T : Task, ITaskComplet
                         // 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;
 
+                        // Skip to a nontransparent/user continuation. Such continuaton must exist.
+                        // Since we see a VTS notifier, something was directly or indirectly
+                        // awaiting an async thunk for a ValueTask-returning method.
+                        // That can only happen in nontransparent/user code.
+                        Continuation nextUserContinuation = headContinuation.Next!;
+                        while ((nextUserContinuation.Flags & continueFlags) == 0 && nextUserContinuation.Next != null)
+                        {
+                            nextUserContinuation = nextUserContinuation.Next;
+                        }
+
+                        ContinuationFlags continuationFlags = nextUserContinuation.Flags;
                         const ContinuationFlags continueOnContextFlags =
                             ContinuationFlags.ContinueOnCapturedSynchronizationContext |
                             ContinuationFlags.ContinueOnCapturedTaskScheduler;
@@ -620,7 +631,7 @@ public static void HandleSuspended<T, TOps>(T task) where T : Task, ITaskComplet
                         }
 
                         // Clear continuation flags, so that continuation runs transparently
-                        headContinuation.Next!.Flags &= ~continueFlags;
+                        nextUserContinuation.Flags &= ~continueFlags;
                         TOps.ValueTaskSourceOnCompleted(task, vtsNotifier, configFlags);
                     }
                     else

src/tests/async/valuetask-source/valuetask-source-stub.cs

@@ -0,0 +1,67 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+
+using System;
+using System.Threading;
+using System.Threading.Tasks;
+using System.Threading.Tasks.Sources;
+using Xunit;
+
+public class ValueTaskSourceAndStubs
+{
+    [Fact]
+    public static void EntryPoint()
+    {
+        SynchronizationContext? original = SynchronizationContext.Current;
+        SynchronizationContext.SetSynchronizationContext(new MySyncContext());
+
+        try
+        {
+            new ValueTaskSourceAndStubs().TestAsync(new C()).GetAwaiter().GetResult();
+        }
+        finally
+        {
+            SynchronizationContext.SetSynchronizationContext(original);
+        }
+    }
+
+    private async Task TestAsync(IFace i)
+    {
+        await i.Foo<string>(0, 1, 2, 3, 4, 5, 6, 7, "value");
+    }
+
+    private struct C : IFace
+    {
+        public ValueTask Foo<T>(int a0, int a1, int a2, int a3, int a4, int a5, int a6, int a7, T value)
+        {
+            return new ValueTask(new Source(), 0);
+        }
+
+        private class Source : IValueTaskSource
+        {
+            public void GetResult(short token)
+            {
+            }
+
+            public ValueTaskSourceStatus GetStatus(short token)
+            {
+                return ValueTaskSourceStatus.Pending;
+            }
+
+            public void OnCompleted(Action<object> continuation, object state, short token, ValueTaskSourceOnCompletedFlags flags)
+            {
+                Assert.Equal(ValueTaskSourceOnCompletedFlags.UseSchedulingContext, flags);
+                ThreadPool.UnsafeQueueUserWorkItem(continuation, state, preferLocal: true);
+            }
+        }
+    }
+
+    private interface IFace
+    {
+        ValueTask Foo<T>(int a0, int a1, int a2, int a3, int a4, int a5, int a6, int a7, T value);
+    }
+
+    private class MySyncContext : SynchronizationContext
+    {
+    }
+}

src/tests/async/valuetask-source/valuetask-source-stub.csproj

@@ -0,0 +1,8 @@
+<Project Sdk="Microsoft.NET.Sdk.IL">
+  <PropertyGroup>
+    <Optimize>True</Optimize>
+  </PropertyGroup>
+  <ItemGroup>
+    <Compile Include="$(MSBuildProjectName).cs" />
+  </ItemGroup>
+</Project>