CancellationToken.Combine: disposing the returned source does not unregister it from the source tokens (memory leak) #302

Summary

The static CancellationToken.Combine(...) (added in 25c891a; reproduced on DotNext.Threading 6.8.0, the current release) returns an IMultiplexedCancellationTokenSource. Disposing it does not remove the registrations it made on the source tokens.

If one of the tokens is long-lived, every combined source ever created stays reachable from it, so memory grows with each call and is never released.

CancellationTokenMultiplexer.Combine(...) on a multiplexer instance and the BCL's CancellationTokenSource.CreateLinkedTokenSource(...) do not have this problem: in the loop below both retain nothing.

Reproduction

DotNext.Threading 6.8.0, .NET SDK 10.0.400, as a file-based app (dotnet run -c Release repro.cs):

#:package DotNext.Threading@6.8.0
using System.Globalization;
using System.Runtime.CompilerServices;
using DotNext.Threading;

CultureInfo.CurrentCulture = CultureInfo.InvariantCulture;

// A token that outlives many short operations, e.g. IHostApplicationLifetime.ApplicationStopping.
using var longLived = new CancellationTokenSource();

// 1. A single combine + dispose: the short-lived source stays reachable from the long-lived one.
var shortLived = CombineAndDispose(longLived.Token);
GC.Collect();
GC.WaitForPendingFinalizers();
GC.Collect();
Console.WriteLine($"short-lived source still reachable after Dispose: {shortLived.TryGetTarget(out _)}");

// 2. The same in a loop: memory grows with every iteration.
const int iterations = 1_000_000;
var before = GC.GetTotalMemory(forceFullCollection: true);
for (var i = 0; i < iterations; i++) {
	using var perOperation = new CancellationTokenSource();
	using var linked = CancellationToken.Combine(perOperation.Token, longLived.Token);
}
var after = GC.GetTotalMemory(forceFullCollection: true);
Console.WriteLine($"retained after {iterations:N0} iterations: {(after - before) / 1048576.0:F1} MiB ({(after - before) / (double)iterations:F0} bytes per iteration)");

// 3. Every leaked registration is still invoked when the long-lived token is finally canceled.
var disposedExceptions = 0;
AppDomain.CurrentDomain.FirstChanceException += (_, e) => {
	if (e.Exception is ObjectDisposedException)
		disposedExceptions++;
};
longLived.Cancel();
Console.WriteLine($"ObjectDisposedExceptions thrown and swallowed on Cancel: {disposedExceptions:N0}");

[MethodImpl(MethodImplOptions.NoInlining)]
static WeakReference<CancellationTokenSource> CombineAndDispose(CancellationToken longLived) {
	var shortLived = new CancellationTokenSource();
	CancellationToken.Combine(shortLived.Token, longLived).Dispose();
	return new(shortLived);
}

Output:

short-lived source still reachable after Dispose: True
retained after 1,000,000 iterations: 569.1 MiB (597 bytes per iteration)
ObjectDisposedExceptions thrown and swallowed on Cancel: 1,000,001

Cause

CancellationTokenSourceFactory.Create in CancellationTokenMultiplexer.Factory.cs creates a PooledCancellationTokenSource and calls AddRange(tokens), which registers a callback on each token. PooledCancellationTokenSource.Dispose(bool) in CancellationTokenMultiplexer.CTS.cs then only clears extraTokens:

protected override void Dispose(bool disposing)
{
    if (disposing)
    {
        extraTokens = null; // help GC
    }

    base.Dispose(disposing);
}

The registrations are removed only by DetachLinkedTokens(), which is called from Scope.Dispose() / Scope.DisposeAsync() and from IValueTaskSource<CancellationToken>.GetResult. The static Combine path goes through neither. Each registration therefore stays in the source token's callback list and keeps the PooledCancellationTokenSource alive, and through the registrations it stores, the CancellationTokenSource objects behind the other tokens too.

For comparison, Linked2CancellationTokenSource.Dispose in the same library does unregister both of its registrations.

Still present on master at 34ebb0c.

Impact

• Unbounded growth when a short-lived token is combined with a long-lived one. I hit this in ASP.NET Core middleware that links HttpContext.RequestAborted to IHostApplicationLifetime.ApplicationStopping for each request: a heap dump of the server held one PooledCancellationTokenSource and one CancellationTokenSource+CallbackNode per request served.
• A burst of exceptions at cancellation. When the long-lived token is eventually canceled, every leaked registration fires and calls Cancel on an already disposed source. Each call throws an ObjectDisposedException that OnCanceled swallows (1,000,001 of them in the repro above).

Proposed fix

Detach in Dispose:

 protected override void Dispose(bool disposing)
 {
     if (disposing)
     {
+        DetachLinkedTokens();
+        count = 0;
         extraTokens = null; // help GC
     }

     base.Dispose(disposing);
 }

Notes on the change:

• count = 0 is needed for a second Dispose(). Double dispose is fine today. With only the DetachLinkedTokens() line added, a second Dispose() on a source that has more than three tokens indexes into the already cleared extraTokens and throws IndexOutOfRangeException.
• The pooled paths should be unaffected. Scope.Return calls Reset(), which zeroes count, before it can reach Dispose(), and GetResult detaches and resets without disposing. The added call is a no-op in both.
• Dispose() versus Unregister(). DetachLinkedTokens() uses CancellationTokenRegistration.Dispose(), which waits for a callback that is running on another thread. Linked2CancellationTokenSource uses Unregister(), which does not wait. Either removes the leak; which is right here is your call.

What I verified and what I did not: against the published 6.8.0 package, invoking DetachLinkedTokens() through reflection before Dispose() makes the short-lived source collectable and drops the retained memory in the loop to zero, and a second detach after Dispose() with four tokens throws IndexOutOfRangeException. I have not built the library with the patch or run its test suite.

Suggested regression test

In the style of MultiplexedCancellationTokenSourceTests. It fails on 6.8.0 (this is step 1 of the repro) and should pass with the change above:

[Fact]
public static void DisposeDetachesSourceTokens()
{
    using var longLived = new CancellationTokenSource();
    var shortLived = CombineAndDispose(longLived.Token);

    GC.Collect();
    GC.WaitForPendingFinalizers();
    GC.Collect();

    False(shortLived.TryGetTarget(out _));

    [MethodImpl(MethodImplOptions.NoInlining)]
    static WeakReference<CancellationTokenSource> CombineAndDispose(CancellationToken longLived)
    {
        var shortLived = new CancellationTokenSource();
        CancellationToken.Combine(shortLived.Token, longLived).Dispose();
        return new(shortLived);
    }
}