Use cached continuation in Await as well when used with ValueTaskSource

[VSadov]

Same idea as in #127973, but applied when the actual Await is used (not a thunk).

With the original repro that I was looking at, merging of #127973 helped to reduce Gen0 GCs from ~ 15/sec to ~ 12/sec.
Alocations went down but not completely, because of code like the following piece in the DoReceive() loop and possibly in other places. This pattern still churns continuations:

https://github.com/dotnet/aspnetcore/blob/c0c2230799a3f876aa673a66bc008bf9b803acac/src/Servers/Kestrel/Transport.Sockets/src/Internal/SocketConnection.cs#L173-L182

                var flushTask = Input.FlushAsync();

                var paused = !flushTask.IsCompleted;

                if (paused)
                {
                    SocketsLog.ConnectionPause(_logger, this);
                }

                var result = await flushTask;

Because flushTask is not a method, the await does not go through a thunk which could reuse the continuation. We call the actual Await in this case and it suspends/allocates.

The change in this PR uses the same approach as in the thunk, but applied to the Await, when awaiting a ValueTaskSource.

With this change we get to ~ 2 Gen0/sec in the repro - on par with net10.

[dotnet-policy-service]

Tagging subscribers to this area: @agocke
See info in area-owners.md if you want to be subscribed.

[VSadov]

CC: @jakobbotsch I think this could be generalized further, but handling just VTS is likely the most interesting/effective case.

[Copilot]

Pull request overview

This PR updates the runtime-async infrastructure in System.Private.CoreLib to reduce allocations when AsyncHelpers.Await(ValueTask/ValueTask<T>) suspends on IValueTaskSource-backed ValueTasks, by reusing a cached ValueTaskContinuation and capturing continuation-context info for the “actual Await” case (not just thunk-generated awaits).

Changes:

• Update AsyncHelpers.Await(ValueTask) / AsyncHelpers.Await<T>(ValueTask<T>) to special-case Task vs IValueTaskSource and route the latter through new ValueTaskSource-specific await helpers.
• Extend ValueTaskContinuation to carry a continuation-context slot and reset continuation-context-related flags/state when caching the continuation for reuse.
• Add AwaitValueTaskSource / AwaitValueTaskSourceOfT and adjust suspension handling assertions/flag propagation to accommodate captured continuation-context flags for IValueTaskSource awaits.

Reviewed changes

Copilot reviewed 3 out of 3 changed files in this pull request and generated 3 comments.

File	Description
src/libraries/System.Private.CoreLib/src/System/Runtime/CompilerServices/AsyncHelpers.cs	Routes ValueTask awaits through Task vs ValueTaskSource paths to enable cached continuation reuse.
src/coreclr/System.Private.CoreLib/src/System/Runtime/CompilerServices/AsyncHelpers.ValueTaskContinuation.cs	Adds continuation-context storage to the cached continuation and clears related state/flags on reuse.
src/coreclr/System.Private.CoreLib/src/System/Runtime/CompilerServices/AsyncHelpers.CoreCLR.cs	Introduces ValueTaskSource await helpers and updates suspension handling to work with captured continuation-context flags.

[Copilot]

Pull request overview

Copilot reviewed 3 out of 3 changed files in this pull request and generated 2 comments.

[Copilot]

Pull request overview

Copilot reviewed 3 out of 3 changed files in this pull request and generated 3 comments.

[Copilot]

Pull request overview

Copilot reviewed 3 out of 3 changed files in this pull request and generated no new comments.

Comments suppressed due to low confidence (1)

src/libraries/System.Private.CoreLib/src/System/Runtime/CompilerServices/AsyncHelpers.cs:128

• The XML doc comment for Await(ValueTask task) still refers to Task ("Awaits the specified Task...") even though the parameter and method are ValueTask. This is confusing for consumers (even if EditorBrowsable(Never)) and should be updated to reference ValueTask / "value task" consistently.

        /// <summary>
        /// Awaits the specified <see cref="Task"/> and throws any exception produced by the task.
        /// </summary>
        /// <param name="task">The task to await.</param>
        [Intrinsic]

[jakobbotsch]

LGTM!

[VSadov]

Thanks!