Hidden catch when using linked CancellationTokenSource


Today’s short post was inspired by a recent memory leak in Nancy. I thought it’s worth to describe it in detail as the reason why the memory was leaking was not so obvious and many of us could commit the same mistake as Nancy authors. The leak was present in the NancyEngine class, which is the central point in the Nancy request handling process. In other words: each request served by the Nancy application must pass through this class instance. NancyEngine processes the requests asynchronously and accepts as a parameter a CancellationToken instance – thus making it possible to cancel the request from the “outside”. At the same time it uses an internal CancellationTokenSource instance which cancels the current requests when the engine is getting disposed. As you see there are two cancellation tokens involved and the HandleRequest method needs to respect their statuses. For such an occasion there is a method in the CancellationTokenSource class in .NET Framework which creates for you a special “linked” CancellationTokenSource that will depend on values in the related tokens. From now on you don’t need to worry about the other tokens as whenever they get cancelled your linked token will become cancelled too. With this introduction the prologue of the HandleRequest becomes clear:

public Task<NancyContext> HandleRequest(Request request, Func<NancyContext, NancyContext> preRequest,
        CancellationToken cancellationToken)
{
    var cts = CancellationTokenSource.CreateLinkedTokenSource(
            this.engineDisposedCts.Token, cancellationToken);
    ...
    // cts.Token used
}

Let’s leave Nancy for a second and focus on the cancellation tokens. Have a look at the following XUnit test and guess whether it will pass or fail:

[Fact]
public static void RunCancellableTask()
{
    WeakReference wref = null;
    M1(ref wref);
    GC.Collect();

    Assert.False(wref.IsAlive);

    var cts1 = new CancellationTokenSource();
    var cts2 = new CancellationTokenSource();
    M2(ref wref, cts1.Token, cts2.Token);
    GC.Collect();

    Assert.False(wref.IsAlive);
}

private static void M1(ref WeakReference wref)
{
    var cts = new CancellationTokenSource();
    wref = new WeakReference(cts);
    Task.Delay(1000).Wait(cts.Token);
}

private static void M2(ref WeakReference wref, CancellationToken token1, CancellationToken token2)
{
    var cts = CancellationTokenSource.CreateLinkedTokenSource(token1, token2);
    wref = new WeakReference(cts);
    Task.Delay(1000).Wait(cts.Token);
}

Well, it will fail on the second assert. Now the question is why the first cancellation token (from the M1 method) got disposed and the second one (from the M2 method) not. Let’s stop on the second assert and check who keeps references to the wref.Target:

cancel-tokens

As you can see the linked CancellationTokenSource hasn’t been reclaimed as there are references to it from the callbacks arrays in the linked CancellationTokens. This reveals the mechanism upon which linked tokens are built. When you create a linked CancellationTokenSource it initiates internally a new CancellationToken and registers a callback method for each of the linked tokens. In this callback method it cancels its internal token whenever any of the other tokens gets cancelled. As we now know what is the leak source let’s see how it can be fixed. Nancy authors already did that in the 1.4.2 version of the Nancy framework:

public Task<NancyContext> HandleRequest(Request request, Func<NancyContext, NancyContext> preRequest,
        CancellationToken cancellationToken)
{
    using (var cts = CancellationTokenSource.CreateLinkedTokenSource(
            this.engineDisposedCts.Token, cancellationToken)) {
        ...
        // cts.Token used
    }
}

That’s it. When linked CancellationTokenSource is getting disposed it unregisters itself from the linked tokens callbacks tables and no leak is present. To conclude: always dispose CancellationTokenSource instances when you are done using them.

Leave a Reply

Fill in your details below or click an icon to log in:

WordPress.com Logo

You are commenting using your WordPress.com account. Log Out / Change )

Twitter picture

You are commenting using your Twitter account. Log Out / Change )

Facebook photo

You are commenting using your Facebook account. Log Out / Change )

Google+ photo

You are commenting using your Google+ account. Log Out / Change )

Connecting to %s