Search code examples
c#asynchronousasync-awaitiasyncenumerableasynclocal

Why is AsyncLocal value lost between IAsyncEnumerable yield points


Consider the following example, and in particular the "Wrap function":

await foreach (var item in Wrap(Sequence()))
    continue;


static IAsyncEnumerable<int> Sequence()
{
    return Core();

    static async IAsyncEnumerable<int> Core([EnumeratorCancellation] CancellationToken ct = default)
    {
        yield return 123;
        await Task.Delay(TimeSpan.FromSeconds(5), ct);
        yield return 456;
    }
}

static IAsyncEnumerable<int> Wrap(IAsyncEnumerable<int> source)
{
    return Core(source);
    
    static async IAsyncEnumerable<int> Core(IAsyncEnumerable<int> source, [EnumeratorCancellation] CancellationToken ct = default)
    {
        Test.Context.Value = "Hello world";
        await foreach (var value in source.WithCancellation(ct))
        {
            var before = Test.Context.Value; // Surely this should always be "Hello world", right?
            yield return value;
            var after = Test.Context.Value; // .. and surely this too?
            Console.WriteLine($"Before={before}, after={after}");
        }
    }
}
    
static class Test
{
    public static readonly AsyncLocal<string> Context = new();
}

What I expected was that both the captued before and after values would contain the previously set "Hello world".

What happes, however, is that for the first iteration, before captures "Hello world", but after is then null. For the second iteration of this loop, both before and after are null.

  1. Why is this the case? It's not quite to me what or why the value of the AsyncLocal is cleared out after that first yield point
  2. What would be the best way to re-formule the Wrap function so that the AsyncLocal value is preserved for the entirety of the enumeration?

Solution

  • It's because on the second iteration, or let's say "dive" into the state machines we don't execute Test.Context.Value = "Hello world"; so it actually is the value from the calling function.

    Short repro:

    async Task Main() {
        Test.Context.Value = "Initial";
        await foreach (var hello in Wrap(Sequence())) {
            continue;
        }
    }
    
    static async IAsyncEnumerable<int> Wrap(IAsyncEnumerable<int> wrapped) {
        Test.Context.Value = "Hello world";
        await foreach (var value in wrapped) {
            var before = Test.Context.Value; // Surely this should always be "Hello world", right?
            yield return value;
            var after = Test.Context.Value; // .. and surely this too?
            Console.WriteLine($"Before={before}, after={after}");
        }
    }
    
    static async IAsyncEnumerable<int> Sequence() {
        yield return 3;
        await Task.Delay(1000);
        yield return 4;
    }
    
    
    static class Test {
        public static AsyncLocal<string> Context = new();
    }
    

    This writes:

    Before=Hello world, after=Initial
    Before=Initial, after=Initial
    

    For it to work, you need to just set the value in the beginning of Wrap before you return the Core static local:

    static IAsyncEnumerable<int> Wrap(IAsyncEnumerable<int> wrapped) {
    
        Test.Context.Value = "Hello world";
        return Core(wrapped);
    
        static async IAsyncEnumerable<int> Core(IAsyncEnumerable<int> source, [EnumeratorCancellation] CancellationToken ct = default) {
            await foreach (var value in source.WithCancellation(ct)) {
                var before = Test.Context.Value; // Surely this should always be "Hello world", right?
                yield return value;
                var after = Test.Context.Value; // .. and surely this too?
                Console.WriteLine($"Before={before}, after={after}");
            }
        }
    }