AsyncLocal<T> does not survive across `yield return` statement

[CXuesong]

Description

I'm not sure this is bug or by-design. Consider the following use case of AsyncLocal<T>. Should I expect stringLocal.Value to be "Test" or null?

var stringLocal = new AsyncLocal<string>();

async IAsyncEnumerable<int> IntSequence()
{
    stringLocal.Value = "Test";
    yield return 10;
    // Now what should happen if we do not explicitly capture the execution context???
    Console.WriteLine("stringLocal.Value: {0}", stringLocal.Value);
}

From what I can observe currently, stringLocal.Value after yield return is null.

Configuration

.NET 5 / .NET 6.0.100-preview.1.21103.9

Regression?

N / A

Other information

Here is an extensive test around the persistence of AsyncLocal<T>. [.NET Fiddle]

using System;
using System.Collections.Generic;
using System.Diagnostics;
using System.Threading;
using System.Threading.Tasks;

var stringLocal = new AsyncLocal<string>();

async IAsyncEnumerable<int> IntSequence(string localValue)
{
    stringLocal.Value = localValue;
    Console.WriteLine("Set up stringLocal: {0}.", localValue);
    await Task.Delay(100).ConfigureAwait(false);
    // await didn't affect AsyncLocal flow
    Console.WriteLine("stringLocal.Value: {0}", stringLocal.Value);

    using var execContext = ExecutionContext.Capture();
    Debug.Assert(execContext != null);      // i.e. ExecutionContext.IsFlowSuppressed() == false
    yield return 10;
    ExecutionContext.Restore(execContext);
    // This gives us the correct value, as we have restored execution context.
    Console.WriteLine("stringLocal.Value: {0}", stringLocal.Value);

    // Now what should happen if we do not explicitly capture the execution context???
    yield return 20;
    Console.WriteLine("stringLocal.Value: {0}", stringLocal.Value);
    yield return 30;
    Console.WriteLine("stringLocal.Value: {0}", stringLocal.Value);
}

async Task PrintSequence1()
{
    await foreach (var i in IntSequence("value 1"))
    {
        Console.WriteLine("Sequence: {0}", i);
        await Task.Delay(100);
    }
}

async Task PrintSequence2()
{
    await using var it = IntSequence("value 2").GetAsyncEnumerator();
    while (await it.MoveNextAsync())
    {
        Console.WriteLine("Sequence: {0}", it.Current);
    }
}

PrintSequence1().Wait();
Console.WriteLine();
PrintSequence2().Wait();

Here is the output

Set up stringLocal: value 1.
stringLocal.Value: value 1
Sequence: 10
stringLocal.Value: value 1
Sequence: 20
stringLocal.Value: 
Sequence: 30
stringLocal.Value: 

Set up stringLocal: value 2.
stringLocal.Value: value 2
Sequence: 10
stringLocal.Value: value 2
Sequence: 20
stringLocal.Value: 
Sequence: 30
stringLocal.Value: 

I couldn't figure out the best area label to add to this issue. If you have write-permissions please help me learn by adding exactly one area label.

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

Issue Details

---

Description

I'm not sure this is bug or by-design. Consider the following use case of AsyncLocal<T>. Should I expect stringLocal.Value to be "Test" or null?

var stringLocal = new AsyncLocal<string>();

async IAsyncEnumerable<int> IntSequence()
{
    stringLocal.Value = "Test";
    yield return 10;
    // Now what should happen if we do not explicitly capture the execution context???
    Console.WriteLine("stringLocal.Value: {0}", stringLocal.Value);
}

From what I can observe currently, stringLocal.Value after yield return is null.

Configuration

.NET 5 / .NET 6.0.100-preview.1.21103.9

Regression?

N / A

Other information

Here is an extensive test around the persistence of AsyncLocal<T>. [.NET Fiddle]

using System;
using System.Collections.Generic;
using System.Diagnostics;
using System.Threading;
using System.Threading.Tasks;

var stringLocal = new AsyncLocal<string>();

async IAsyncEnumerable<int> IntSequence(string localValue)
{
    stringLocal.Value = localValue;
    Console.WriteLine("Set up stringLocal: {0}.", localValue);
    await Task.Delay(100).ConfigureAwait(false);
    // await didn't affect AsyncLocal flow
    Console.WriteLine("stringLocal.Value: {0}", stringLocal.Value);

    using var execContext = ExecutionContext.Capture();
    Debug.Assert(execContext != null);      // i.e. ExecutionContext.IsFlowSuppressed() == false
    yield return 10;
    ExecutionContext.Restore(execContext);
    // This gives us the correct value, as we have restored execution context.
    Console.WriteLine("stringLocal.Value: {0}", stringLocal.Value);

    // Now what should happen if we do not explicitly capture the execution context???
    yield return 20;
    Console.WriteLine("stringLocal.Value: {0}", stringLocal.Value);
    yield return 30;
    Console.WriteLine("stringLocal.Value: {0}", stringLocal.Value);
}

async Task PrintSequence1()
{
    await foreach (var i in IntSequence("value 1"))
    {
        Console.WriteLine("Sequence: {0}", i);
        await Task.Delay(100);
    }
}

async Task PrintSequence2()
{
    await using var it = IntSequence("value 2").GetAsyncEnumerator();
    while (await it.MoveNextAsync())
    {
        Console.WriteLine("Sequence: {0}", it.Current);
    }
}

PrintSequence1().Wait();
Console.WriteLine();
PrintSequence2().Wait();

Here is the output

Set up stringLocal: value 1.
stringLocal.Value: value 1
Sequence: 10
stringLocal.Value: value 1
Sequence: 20
stringLocal.Value: 
Sequence: 30
stringLocal.Value: 

Set up stringLocal: value 2.
stringLocal.Value: value 2
Sequence: 10
stringLocal.Value: value 2
Sequence: 20
stringLocal.Value: 
Sequence: 30
stringLocal.Value: 

Author:	CXuesong
Assignees:	-
Labels:	area-System.Threading.Tasks, untriaged
Milestone:	-

[CXuesong]

Note that this means there would be problem if you are using MEL scoped logging in async stream function.

Consider the following usage scenario

using System;
using System.Collections.Generic;
using Microsoft.Extensions.Logging;

async IAsyncEnumerable<int> IntSequence(ILogger logger)
{
    using var scope = logger.BeginScope("My scope");
    logger.LogWarning("Log 1");
    yield return 10;
    logger.LogWarning("Log 2");
    yield return 20;
    logger.LogWarning("Log 3");
    yield return 30;
}

using var lf = LoggerFactory.Create(b => b.AddSimpleConsole(o => o.IncludeScopes = true));
await foreach (var i in IntSequence(lf.CreateLogger("Root")))
{
    Console.WriteLine(i);
}

The expected output is

10
20
30
warn: Root[0]
      => My scope
      Log 1
warn: Root[0]
      => My scope
      Log 2
warn: Root[0]
      => My scope
      Log 3

The actual output is

10
20
30
warn: Root[0]
      => My scope
      Log 1
warn: Root[0]
      Log 2
warn: Root[0]
      Log 3

Scopes are lost for Log 2 and Log 3.

[stephentoub]

It's "by design" in that it was done this way on purpose, with ExecutionContext flowing across await points just as in regular async methods and nothing special being done about yield returns, which logically exit and transfer control to something else (not exactly the same, but putting state into a thread local for example prior to a yield return in a normal sync iterator is in no way guaranteed to be there upon re-entering the iterator).

Changing that would theoretically be possible but also non-trivial. Either the Roslyn compiler-generated state machine would need to do the ExecutionContext capture and restore around yields, storing the ExecutionContext into the state machine, or AsyncIteratorMethodBuilder would somehow need to be informed whether the compiler-generated IAsyncStateMachine.MoveNext method was returning because of a yield or an await, which would require both compiler and runtime changes (and probably significant ones at that). It would also have a potential performance impact, adding an ExecutionContext.Capture() at every yield (right now we only pay that cost when awaiting a task that hasn't yet completed).
cc: @jcouv

[CXuesong]

Thanks for your explanation @stephentoub, and that makes some sense.

However, aside from the MEL logging scope problem that I have bumped into, I'd like to comment on the following statement

It would also have a potential performance impact, adding an ExecutionContext.Capture() at every yield

I've checked the source code of ExecutionContext.Capture (or the CaptureForRestore method lying beside), it looks trivial and even might get inlined during runtime. The only place that may introduce overhead is the access to CurrentThread.

I think the point is, we should have well-defined (and consistent, unsurprising, if possible) behavior on AsycLocal. For now I fail to see any caution of usage with async stream on docs, and from the fashion MEL loggers use (or continue to use) AsyncLocal, I guess they didn't see this coming, either.

Btw, as of now, what is the recommended pattern for me to temporarily mitigate the issue? I have to call ExecutionContext.Capture before yield return, but should I either use ExecutionContext.Restore or ExecutionContext.Run after reentry? When (or whether) should I dispose the captured ExecutionContext?

[stephentoub]

I've checked the source code of ExecutionContext.Capture (or the CaptureForRestore method lying beside), it looks trivial and even might get inlined during runtime. The only place that may introduce overhead is the access to CurrentThread.

TLS is not free. For example, here's a simple benchmark comparing it against a (well predicted) interface dispatch:

    [ThreadStatic]
    public static object t_obj;

    public IList<int> _list = new List<int>();

    [Benchmark] public object Get() => t_obj;
    [Benchmark] public int Count() => _list.Count;

which on my machine shows it to be twice as expensive:

Method	Mean	Error	StdDev	Code Size
Get	2.557 ns	0.0218 ns	0.0194 ns	32 B
Count	1.335 ns	0.0068 ns	0.0064 ns	30 B

[CXuesong]

TLS is not free.

I see.

Then perhaps the correct approach to address the ILogger scope loss problem is to expect Microsoft.Extensions.Logging implement LoggerFactoryScopeProvider without AsyncLocal… Do I need to open another issue for this extension library (in the same repo)?

runtime/src/libraries/Microsoft.Extensions.Logging/src/LoggerFactoryScopeProvider.cs

Lines 14 to 21 in f93fee4

	/// <summary>
	/// Default implementation of <see cref="IExternalScopeProvider"/>
	/// </summary>
	internal class LoggerFactoryScopeProvider : IExternalScopeProvider
	{
	private readonly AsyncLocal<Scope> _currentScope = new AsyncLocal<Scope>();
	private readonly ActivityTrackingOptions _activityTrackingOption;

[CXuesong]

Also, it seems that ExecutionContext.Restore is available only on .NET 5.0. I'll have to go with ExecutionContext.Run on .NET Core 3 / .NET Standard 2.1.

[stephentoub]

I do agree the behavior is unintuitive. I will look into what it would take to change it at the runtime level and what impact that would have.

[davidfowl]

@CXuesong fixing it in logging is the wrong layer. It should be done as part of the runtime. Which @stephentoub is going to look at.