.NET TPL CancellationToken memory leak

Question

I've developed a library which implements a producer/consumer pattern for work items. Work is dequeued and a separate task with continuations for failure and success is spun up for each dequeued work item.

The task continuations re-queue the work item after it completed (or failed) its work.

The entire library shares one central CancellationTokenSource, which is triggered on application shutdown.

I now face a major memory leak. If the tasks are created with the cancellation token as a parameter, then the tasks seem to remain in memory until the cancellation source is triggered (and later disposed).

This can be reproduced in this sample code (VB.NET). The main task is the task which would wrap the work item and the continuation tasks would handle the rescheduling.

Dim oCancellationTokenSource As New CancellationTokenSource
Dim oToken As CancellationToken = oCancellationTokenSource.Token
Dim nActiveTasks As Integer = 0

Dim lBaseMemory As Long = GC.GetTotalMemory(True)

For iteration = 0 To 100 ' do this 101 times to see how much the memory increases

  Dim lMemory As Long = GC.GetTotalMemory(True)

  Console.WriteLine("Memory at iteration start: " & lMemory.ToString("N0"))
  Console.WriteLine("  to baseline: " & (lMemory - lBaseMemory).ToString("N0"))

  For i As Integer = 0 To 1000 ' 1001 iterations to get an immediate, measurable impact
    Interlocked.Increment(nActiveTasks)
    Dim outer As Integer = i
    Dim oMainTask As New Task(Sub()
                                ' perform some work
                                Interlocked.Decrement(nActiveTasks)
                              End Sub, oToken)
    Dim inner As Integer = 1
    Dim oFaulted As Task = oMainTask.ContinueWith(Sub()
                                                    Console.WriteLine("Failed " & outer & "." & inner)
                                                    ' if failed, do something with the work and re-queue it, if possible
                                                    ' (imagine code for re-queueing - essentially just a synchronized list.add)

                                                                                                            ' Does not help:
                                                    ' oMainTask.Dispose()
                                                  End Sub, oToken, TaskContinuationOptions.OnlyOnFaulted, TaskScheduler.Default)
    ' if not using token, does not cause increase in memory:
    'End Sub, TaskContinuationOptions.OnlyOnFaulted)

            ' Does not help:
    ' oFaulted.ContinueWith(Sub()
    '                         oFaulted.Dispose()
    '                       End Sub, TaskContinuationOptions.NotOnFaulted)


    Dim oSucceeded As Task = oMainTask.ContinueWith(Sub()
                                                      ' success
                                                      ' re-queue for next iteration
                                                      ' (imagine code for re-queueing - essentially just a synchronized list.add)

                                                                                                                ' Does not help:
                                                      ' oMainTask.Dispose()
                                                    End Sub, oToken, TaskContinuationOptions.OnlyOnRanToCompletion, TaskScheduler.Default)
    ' if not using token, does not cause increase in memory:
    'End Sub, TaskContinuationOptions.OnlyOnRanToCompletion)

            ' Does not help:
    ' oSucceeded.ContinueWith(Sub()
    '                           oSucceeded.Dispose()
    '                         End Sub, TaskContinuationOptions.NotOnFaulted)


    ' This does not help either and makes processing much slower due to the thrown exception (at least one of these tasks is cancelled)
    'Dim oDisposeTask As New Task(Sub()
    '                               Try
    '                                 Task.WaitAll({oMainTask, oFaulted, oSucceeded, oFaultedFaulted, oSuccededFaulted})
    '                               Catch ex As Exception

    '                               End Try
    '                               oMainTask.Dispose()
    '                               oFaulted.Dispose()
    '                               oSucceeded.Dispose()                                     
    '                             End Sub)

    oMainTask.Start()
    '  oDisposeTask.Start()
  Next

  Console.WriteLine("Memory after creating tasks: " & GC.GetTotalMemory(True).ToString("N0"))

  ' Wait until all main tasks are finished (may not mean that continuations finished)

  Dim previousActive As Integer = nActiveTasks
  While nActiveTasks > 0
    If previousActive <> nActiveTasks Then
      Console.WriteLine("Active: " & nActiveTasks)
      Thread.Sleep(500)
      previousActive = nActiveTasks
    End If

  End While

  Console.WriteLine("Memory after tasks finished: " & GC.GetTotalMemory(True).ToString("N0"))

Next

I measured the memory use with the ANTS Memory Profiler and saw a large increase in the System.Threading.ExecutionContext, which traces back to task continuations and CancellationCallbackInfo.

As you can see, I already tried to dispose the tasks which use the cancellation token, but this seems to have no effect.

Edit

I'm using .NET 4.0

Update

Even when just chaining the main task with a continuation on failure, the memory use continuously rises. The task continuation seems to prevent the de-registration from the cancellation token registration.

So if a task is chained with a continuation, which does not run (due to the TaskContinuationOptions), then there seems to be a memory leak. If there is only one continuation, which runs, then I did not observe a memory leak.

Workaround

As a workaround, I can do a single continuation without any TaskContinuationOptions and handle the state of the parent task there:

oMainTask.ContinueWith(Sub(t)
                     If t.IsCanceled Then
                       ' ignore
                     ElseIf t.IsCompleted Then
                       ' reschedule

                     ElseIf t.IsFaulted Then
                       ' error handling

                     End If
                   End Sub)

I'll have to check how this performs in case of a cancellation but this seems to do the trick. I almost suspect a bug in the .NET Framework. Task cancellations with mutual exclusive conditions aren't something which could be this rare.

Answer 1

Some observations

1. The potential leak only seems present in the case where there is a task "branch" that does not run. In your example, if you comment out the oFaulted task, the leak goes away for me. If you update your code to have the oMainTask fault, so that the oFaulted task runs and the oSucceeded task does not run, then commenting out oSucceeded prevents the leak.
2. Perhaps not helpful, but if you call oCancellationTokenSource.Cancel() after all the tasks have run, the memory frees. Dispose does not help, nor any combination of Disposing the cancellation source along with the tasks.
3. I took a look at http://referencesource.microsoft.com/ which is 4.5.2 (Is there a way to view early frameworks?) I know its not necessarily the same but its useful to know what types of things are going on. Basically when you pass a cancellation token to a task, the task registers itself with the cancellation token's cancellation source. So the cancellation source is holding references to all your tasks. I'm not clear yet on why your scenario appears to be leaking. I'll update after I've had a chance to look more in depth, if I find anything.

Workaround

Move your branching logic to a continuation that always runs.

Dim continuation As Task =
    oMainTask.ContinueWith(
        Sub(antecendent)
            If antecendent.Status = TaskStatus.Faulted Then
                'Handle errors
            ElseIf antecendent.Status = TaskStatus.RanToCompletion Then
                'Do something else
            End If
        End Sub,
        oToken,
        TaskContinuationOptions.None,
        TaskScheduler.Default)

There's a good chance this is lighter thant the other approach anyways. In both cases one continuation always runs, but with this code only 1 continuation task gets created instead of 2.