Implicit Async Context ("AsyncLocal")

Occasionally, someone will ask about support for some kind of implicit “context” that will flow with async code. The most common use case is for logging, where you can maintain a stack of logical operations so every time you log you capture (and log) the state of that stack (e.g., CorrelationManager for TraceSource, or the Nested Diagnostic Context for log4net). Jon Skeet has a great blog entry on this kind of implicit “context” with several possible uses.

Logical CallContext

There is a solution for this problem: the “logical call context”, which you can access by CallContext.LogicalGetData and CallContext.LogicalSetData. The regular call context (CallContext.GetData and CallContext.SetData) acts just like thread-local storage, which of course doesn’t work for async methods.

Here’s how logical call context works with asynchronous code.

Logical call context data flows with ExecutionContext. This means that it’s not affected by ConfigureAwait(continueOnCapturedContext: false); you can’t “opt-out” of the logical call context. So the logical call context at the beginning of an async method will always flow through to its continuations.

When an async method starts, it notifies its logical call context to activate copy-on-write behavior. This means the current logical call context is not actually changed, but it is marked so that if your code does call CallContext.LogicalSetData, the logical call context data is copied into a new current logical call context before it is changed. Note: the copy-on-write behavior of logical call contexts is only available on .NET 4.5.

This “copying” of the logical call context data is a shallow copy. You can think of the logical call context data as an IDictionary<string, object> of name/value pairs. When it’s copied, it creates a new dictionary and copies all the name/value pairs into the new dictionary. Both dictionaries then refer to all the same actual object instances; there’s no “deep cloning” of any of your data being done.

Because the references are shared, it’s important not to mutate any values retrieved from the logical call context. If you need to change a logical call context value, update the actual value using CallContext.LogicalSetData. You should only use immutable types as logical call context data values.

Also note that the design is heavily optimized for the common case: when there is no logical call context data at all. When you start adding “implicit context”, you’re going to start adding overhead. Probably not too much, though, since everything is shallow-copied at worst.

As a final note, remember that you can end up sharing data two different ways: .NET before 4.5 did not have the copy-on-write behavior, and the copies are shallow. For async code, any sharing of data like this will get you in trouble as soon as you start doing fork/join (e.g., Task.WhenAll). So follow the two rules: only .NET 4.5 and immutable data!

An Example

Let’s take a simple example. We want to keep an implicit stack of logical operations, and when we call our “log” method we want to output the stack as part of that log.

First, let’s define a strongly-typed accessor for the logical call context data we’ll be using. Remember, we’re only storing immutable data, so we’ll use the new immutable collections:

public static partial class MyStack
{
    private static readonly string name = Guid.NewGuid().ToString("N");

    private sealed class Wrapper : MarshalByRefObject
    {
        public ImmutableStack<string> Value { get; set; }
    }

    private static ImmutableStack<string> CurrentContext
    {
        get
        {
            var ret = CallContext.LogicalGetData(name) as Wrapper;
            return ret == null ? ImmutableStack.Create<string>() : ret.Value;
        }

        set
        {
            CallContext.LogicalSetData(name, new Wrapper { Value = value });
        }
    }
}

So far, so good. Now we have a strongly-typed property we can use to get (or update) the current stack. Next we’ll start defining our public API. We want the ability to “push” a string onto the stack, and get back a disposable that will pop that string back off the stack when disposed. Simple enough:

public static partial class MyStack
{
    public static IDisposable Push([CallerMemberName] string context = "")
    {
        CurrentContext = CurrentContext.Push(context);
        return new PopWhenDisposed();
    }

    private static void Pop()
    {
        CurrentContext = CurrentContext.Pop();
    }

    private sealed class PopWhenDisposed : IDisposable
    {
        private bool disposed;

        public void Dispose()
        {
            if (disposed)
                return;
            Pop();
            disposed = true;
        }
    }
}

The final part of our public API is a method that returns the current stack. I’ll just return it as a string:

public static partial class MyStack
{
    public static string CurrentStack
    {
        get
        {
            return string.Join(" ", CurrentContext.Reverse());
        }
    }
}

Now let’s turn our attention to the code that will be using MyStack. First, our “log” method:

partial class Program
{
    static void Log(string message)
    {
        Console.WriteLine(MyStack.CurrentStack + ": " + message);
    }
}

Yeah, that was pretty easy.

Our test code is going to be a bit more complex. First, I’m going to push a “Main” value onto the stack that will last for the entire program. Then, I’ll start two separate (concurrent) pieces of work called “1” and “2”. Each of those are going to log when they start and finish, and they’ll each do some more work called “A” and “B” (sequentially). So we should end up with some interleaving of this output from “1”:

Main 1: <SomeWork>
Main 1 A: <MoreWork>
Main 1 A: </MoreWork>
Main 1 B: <MoreWork>
Main 1 B: </MoreWork>
Main 1: </SomeWork>

with this output from “2”:

Main 2: <SomeWork>
Main 2 A: <MoreWork>
Main 2 A: </MoreWork>
Main 2 B: <MoreWork>
Main 2 B: </MoreWork>
Main 2: </SomeWork>

Remember, “1” and “2” are concurrent, so there’s no one right answer for the output. As long as all the messages above are present and in the correct (relative) order, it’s acceptable.

The code, without further ado:

partial class Program
{
    static void Main(string[] args)
    {
        using (MyStack.Push("Main"))
        {
            Task.WhenAll(SomeWork("1"), SomeWork("2")).Wait();
        }

        Console.ReadKey();
    }

    static async Task SomeWork(string stackName)
    {
        using (MyStack.Push(stackName))
        {
            Log("<SomeWork>");
            await MoreWork("A");
            await MoreWork("B");
            Log("</SomeWork>");
        }
    }

    static async Task MoreWork(string stackName)
    {
        using (MyStack.Push(stackName))
        {
            Log("<MoreWork>");
            await Task.Delay(10);
            Log("</MoreWork>");
        }
    }
}

One sample run from my machine is:

Main 1: <SomeWork>
Main 1 A: <MoreWork>
Main 2: <SomeWork>
Main 2 A: <MoreWork>
Main 2 A: </MoreWork>
Main 1 A: </MoreWork>
Main 1 B: <MoreWork>
Main 2 B: <MoreWork>
Main 2 B: </MoreWork>
Main 2: </SomeWork>
Main 1 B: </MoreWork>
Main 1: </SomeWork>

If you sort out the “1” and the “2” messages, you’ll see that each set is in the correct order and that the stacks are nicely laid out as expected.

Similar code will compile targeting .NET 4.0 (with Microsoft.Bcl.Async); however, it will not work correctly (unless it happens to run on .NET 4.5) because the logical call context is not copied at the right time. In that situation, different parts of different async methods will end up sharing the same stack (and overwriting each other’s stack).

As a final reminder: this will work only on .NET 4.5 and it only works because we stored immutable data.