Cleaning up CallContext in TPL

Question

Depending on whether I'm using async/await based code or TPL based code, I'm getting two different behaviors regarding the clean-up of logical CallContext.

I can set and clear logical CallContext exactly as I expect if I use the following async/await code:

class Program
{
    static async Task DoSomething()
    {
        CallContext.LogicalSetData("hello", "world");

        await Task.Run(() =>
            Debug.WriteLine(new
            {
                Place = "Task.Run",
                Id = Thread.CurrentThread.ManagedThreadId,
                Msg = CallContext.LogicalGetData("hello")
            }))
            .ContinueWith((t) =>
                CallContext.FreeNamedDataSlot("hello")
                );

        return;
    }

    static void Main(string[] args)
    {
        DoSomething().Wait();

        Debug.WriteLine(new
        {
            Place = "Main",
            Id = Thread.CurrentThread.ManagedThreadId,
            Msg = CallContext.LogicalGetData("hello")
        });

    }
}

The above outputs the following:

{ Place = Task.Run, Id = 9, Msg = world }
{ Place = Main, Id = 8, Msg = }

Notice the Msg = which indicates that CallContext on the main thread has been freed and is empty.

But when I switch to pure TPL / TAP code I can't achieve the same effect...

class Program
{
    static Task DoSomething()
    {
        CallContext.LogicalSetData("hello", "world");

        var result = Task.Run(() =>
            Debug.WriteLine(new
            {
                Place = "Task.Run",
                Id = Thread.CurrentThread.ManagedThreadId,
                Msg = CallContext.LogicalGetData("hello")
            }))
            .ContinueWith((t) =>
                CallContext.FreeNamedDataSlot("hello")
                );

        return result;
    }

    static void Main(string[] args)
    {
        DoSomething().Wait();

        Debug.WriteLine(new
        {
            Place = "Main",
            Id = Thread.CurrentThread.ManagedThreadId,
            Msg = CallContext.LogicalGetData("hello")
        });
    }
}

The above outputs the following:

{ Place = Task.Run, Id = 10, Msg = world }
{ Place = Main, Id = 9, Msg = world }

Is there anything I can do to coerce TPL to "free" the logical CallContext the same way as the async/await code does?

I am not interested in alternatives to CallContext.

I'm hoping to get the above TPL/TAP code fixed so that I can use it in projects targeting the .net 4.0 framework. If that is not possible in .net 4.0, I'm still curious if it can be done in .net 4.5.

Answer 1

In an async method the CallContext is copied on write:

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.

From Implicit Async Context ("AsyncLocal")

That means that in your async version the CallContext.FreeNamedDataSlot("hello") continuation is redundant as even without it:

static async Task DoSomething()
{
    CallContext.LogicalSetData("hello", "world");

    await Task.Run(() =>
        Console.WriteLine(new
        {
            Place = "Task.Run",
            Id = Thread.CurrentThread.ManagedThreadId,
            Msg = CallContext.LogicalGetData("hello")
        }));
}

The CallContext in Main wouldn't contain the "hello" slot:

{ Place = Task.Run, Id = 3, Msg = world }
{ Place = Main, Id = 1, Msg = }

In the TPL equivalent all code outside the Task.Run (which should be Task.Factory.StartNew as Task.Run was added in .Net 4.5) runs on the same thread with the same exact CallContext. If you want to clean it you need to do that on that context (and not in the continuation):

static Task DoSomething()
{
    CallContext.LogicalSetData("hello", "world");

    var result = Task.Factory.StartNew(() =>
        Debug.WriteLine(new
        {
            Place = "Task.Run",
            Id = Thread.CurrentThread.ManagedThreadId,
            Msg = CallContext.LogicalGetData("hello")
        }));

    CallContext.FreeNamedDataSlot("hello");
    return result;
}

You can even abstract a scope out of it to make sure you always clean up after yourself:

static Task DoSomething()
{
    using (CallContextScope.Start("hello", "world"))
    {
        return Task.Factory.StartNew(() =>
            Debug.WriteLine(new
            {
                Place = "Task.Run",
                Id = Thread.CurrentThread.ManagedThreadId,
                Msg = CallContext.LogicalGetData("hello")
            }));
    }
}

Using:

public static class CallContextScope
{
    public static IDisposable Start(string name, object data)
    {
        CallContext.LogicalSetData(name, data);
        return new Cleaner(name);
    }

    private class Cleaner : IDisposable
    {
        private readonly string _name;
        private bool _isDisposed;

        public Cleaner(string name)
        {
            _name = name;
        }

        public void Dispose()
        {
            if (_isDisposed)
            {
                return;
            }

            CallContext.FreeNamedDataSlot(_name);
            _isDisposed = true;
        }
    }
}

Answer 2

A good question. The await version may not work the way you may think it does here. Let's add another logging line inside DoSomething:

class Program
{
    static async Task DoSomething()
    {
        CallContext.LogicalSetData("hello", "world");

        await Task.Run(() =>
            Debug.WriteLine(new
            {
                Place = "Task.Run",
                Id = Thread.CurrentThread.ManagedThreadId,
                Msg = CallContext.LogicalGetData("hello")
            }))
            .ContinueWith((t) =>
                CallContext.FreeNamedDataSlot("hello")
                );

        Debug.WriteLine(new
        {
            Place = "after await",
            Id = Thread.CurrentThread.ManagedThreadId,
            Msg = CallContext.LogicalGetData("hello")
        });
    }

    static void Main(string[] args)
    {

        DoSomething().Wait();

        Debug.WriteLine(new
        {
            Place = "Main",
            Id = Thread.CurrentThread.ManagedThreadId,
            Msg = CallContext.LogicalGetData("hello")
        });

        Console.ReadLine();
    }
}

Output:

{ Place = Task.Run, Id = 10, Msg = world }
{ Place = after await, Id = 11, Msg = world }
{ Place = Main, Id = 9, Msg =  }

Note the "world" is still there after await, because it was there before await. And it is not there after DoSomething().Wait() because it wasn't there before it, in the first place.

Interestingly enough, the async version of DoSomething creates a copy-on-write clone of the LogicalCallContext for its scope, upon the first LogicalSetData. It does that even when there is no asynchrony inside it - try await Task.FromResult(0). I presume the whole ExecutionContext gets cloned for the scope of the async method, upon the 1st write operation.

OTOH, for the non-async version there is no "logical" scope and no outer ExecutionContext here, so the copy-on-write clone of ExecutionContext becomes current for the Main thread (but the continuations and the Task.Run lambdas still get their own clones). So, you'd either need to move CallContext.LogicalSetData("hello", "world") inside the Task.Run lambda, or clone the context manually:

static Task DoSomething()
{
    var ec = ExecutionContext.Capture();
    Task task = null;
    ExecutionContext.Run(ec, _ =>
    {
        CallContext.LogicalSetData("hello", "world");

        var result = Task.Run(() =>
            Debug.WriteLine(new
            {
                Place = "Task.Run",
                Id = Thread.CurrentThread.ManagedThreadId,
                Msg = CallContext.LogicalGetData("hello")
            }))
            .ContinueWith((t) =>
                CallContext.FreeNamedDataSlot("hello")
                );

        task = result;
    }, null);

    return task;
}