Async-streams: Consider optimizing handling of yield directly following yield

[jcouv]

Reported by @stephentoub:

Right now, every time we yield an item, we incur an interlocked operation as part of SetResult on the MRVTSC. That’s because, from the MRVTSC’s perspective, it needs to be able to handle the case where someone might be concurrently calling OnCompleted, and so it can’t avoid the interlocked if it sees that there’s currently no delegate hooked up (in the corelib implementation, it will avoid the interlocked in that other case).

However, the state machine could be taught to understand whether we’ve yet await’d something as part of this MoveNextAsync call: if we haven’t, then we’re still in the synchronous call to MoveNextAsync, we haven’t handed out the awaiter yet, and so there’s no way someone could concurrently be calling OnCompleted.

In such a case, we could avoid the interlocked. It might be worth thinking through what it would look like to make that happen, and whether we’d need to add anything to the MRVTSC to enable it (e.g. a DangerousSetResult that would trust the caller to say no interlocked is needed) or whether the state machine could bypass the MRVTSC completely in that case.

If it were tracking that information, it could also then potentially be used to avoid doing a Reset on the MRVTSC until the first time we actually await something such that we’re going to yield, saving a few more cycles.

Update (1/9): From LDM discussion, we're supportive of this optimization. Adds a bool and some maintenance for that bool. No additional API needed.

[jcouv]

There may be a way to store this information (whether the last suspension was a yield return) in the state. Currently, resuming will set the state to -1 for both yield return and await, but we could use different values.
That said, I feel I don't fully understand the problem. Let's chat some time next week.

[stephentoub]

That said, I feel I don't fully understand the problem. Let's chat some time next week.

Sounds good. In advance of chatting, just to try to explain it differently...

Right now, the compiler generates code that always invokes ManualResetValueTaskSourceCore.SetResult when yielding, regardless of whether it's already awaited as part of the current MoveNextAsync call. That's functionally fine, and results in simpler compiler-generated code, which is of course good. The issue, though, is that SetResult is implemented to work regardless of whether you've awaited as part of that MoveNextAsync. If that MoveNextAsync has already awaited, then there's a race condition between the awaiter hooking up an OnCompleted callback and the awaited operation completing and causing the call to SetResult. That means SetResult needs to coordinate with OnCompleted, which means it needs to synchronize, and to do that it uses Interlocked.CompareExchange:
https://github.com/dotnet/coreclr/blob/a81680cc473a6620625815eb643c2f2fea771f2e/src/System.Private.CoreLib/shared/System/Threading/Tasks/Sources/ManualResetValueTaskSourceCore.cs#L209
CompareExchange is not particularly cheap, even when not contended, and so if we could avoid doing that CompareExchange when it's not necessary, that would help with throughput. In theory we can avoid doing it in the synchronously completing case, because we know that aforementioned race condition won't exist.

[jcouv]

I was having trouble with this bit: "... whether it's already awaited as part of the current MoveNextAsync call."
Then I realized you're probably talking about a MoveNextAsync with the other optimization you suggested (issue #31246). With that MoveNextAsync optimization in place, we know that yield return causes no await.

What if there was an additional call I could do in MoveNextAsync (after calling GetStatus and GetResult) on the MRVTSC to indicate that I'm done with it and don't care about awaiting?

Something like:

...
f (<>v__promiseOfValueOrEnd.GetStatus(version) == ValueTaskSourceStatus.Succeeded)
{
    bool result = <>v__promiseOfValueOrEnd.GetResult(version);
    <>v__promiseOfValueOrEnd.DiscardOnCompleted(); // tell the MRVTSC that it won't be awaited
    return new ValueTask<bool>(result);
}
return new ValueTask<bool>(this, version);

[stephentoub]

I was having trouble with this bit

What I meant was:

1. Code calls MoveNextAsync.
2. Inside of the resulting MoveNext, we execute code associated with the state machine.
3. As part of that execution, we hit the code associated with a yield return.

We didn't encounter an await during that MoveNextAsync call. Therefore, there's no race condition because it's not possible the calling code could have hooked up an OnCompleted callback, because we never returned to the calling code prior to the yield return. So the yield return doesn't need to synchronize.

What if there was an additional call I could do in MoveNextAsync

That wouldn't help; it's too late at that point to avoid the costs I'm talking about, as they were already incurred. The state machine code has already called v__promiseOfValueOrEnd.SetResult(true), which is what incurs the Interlocked.CompareExchange.

[stephentoub]

FYI, I included this in the manual implementation for System.Threading.Channels here:
dotnet/corefx#34756

Specifically, you can see that there's a field that stores whether we're using the ManualResetValueTaskSourceCore or not:
https://github.com/dotnet/corefx/pull/34756/files#diff-db44d7215ae1b1e97957843ee3840fcaR52
which starts off as false, and then is set to true when we end up awaiting as part of the MoveNextAsync:
https://github.com/dotnet/corefx/pull/34756/files#diff-db44d7215ae1b1e97957843ee3840fcaR127
Then when we have an item to yield, we only complete the MRVTSC if that field is true:
https://github.com/dotnet/corefx/pull/34756/files#diff-db44d7215ae1b1e97957843ee3840fcaR163
and in MoveNextAsync, we check that field to determine whether to use the MRVTSC or not:
https://github.com/dotnet/corefx/pull/34756/files#diff-db44d7215ae1b1e97957843ee3840fcaR102

In a benchmark of that ChannelReader.ReadAllAsync async enumerator, when I populate an unbounded channel and then use the iterator to read through all of those items, this optimization improves throughput by ~20%.