From dbb22b83e7350ca3e4acca7e44c394f54e3aa584 Mon Sep 17 00:00:00 2001 From: Tim Lovell-Smith Date: Wed, 6 Sep 2023 09:15:13 -0700 Subject: [PATCH] Make LinkedCancellationTokenSource register itself with a weak reference, so that it doesn't memory leak by default when you forget to Dispose it. --- .../src/System/Threading/CancellationTokenSource.cs | 13 ++++++++----- 1 file changed, 8 insertions(+), 5 deletions(-) diff --git a/src/libraries/System.Private.CoreLib/src/System/Threading/CancellationTokenSource.cs b/src/libraries/System.Private.CoreLib/src/System/Threading/CancellationTokenSource.cs index 49214ef9b4d1a4..408f3d629663d5 100644 --- a/src/libraries/System.Private.CoreLib/src/System/Threading/CancellationTokenSource.cs +++ b/src/libraries/System.Private.CoreLib/src/System/Threading/CancellationTokenSource.cs @@ -883,7 +883,7 @@ private sealed class Linked1CancellationTokenSource : CancellationTokenSource internal Linked1CancellationTokenSource(CancellationToken token1) { - _reg1 = token1.UnsafeRegister(LinkedNCancellationTokenSource.s_linkedTokenCancelDelegate, this); + _reg1 = token1.UnsafeRegister(LinkedNCancellationTokenSource.s_linkedTokenCancelDelegate, new WeakReference(this)); } protected override void Dispose(bool disposing) @@ -905,8 +905,8 @@ private sealed class Linked2CancellationTokenSource : CancellationTokenSource internal Linked2CancellationTokenSource(CancellationToken token1, CancellationToken token2) { - _reg1 = token1.UnsafeRegister(LinkedNCancellationTokenSource.s_linkedTokenCancelDelegate, this); - _reg2 = token2.UnsafeRegister(LinkedNCancellationTokenSource.s_linkedTokenCancelDelegate, this); + _reg1 = token1.UnsafeRegister(LinkedNCancellationTokenSource.s_linkedTokenCancelDelegate, new WeakReference(this)); + _reg2 = token2.UnsafeRegister(LinkedNCancellationTokenSource.s_linkedTokenCancelDelegate, new WeakReference(this)); } protected override void Dispose(bool disposing) @@ -924,8 +924,11 @@ protected override void Dispose(bool disposing) private sealed class LinkedNCancellationTokenSource : CancellationTokenSource { - internal static readonly Action s_linkedTokenCancelDelegate = static s => + internal static readonly Action s_linkedTokenCancelDelegate = static wr => { + Debug.Assert(wr is WeakReference, $"Expected {typeof(WeakReference)}, got {wr}"); + object? s = ((WeakReference)wr).Target; + if (s is null) return; Debug.Assert(s is CancellationTokenSource, $"Expected {typeof(CancellationTokenSource)}, got {s}"); ((CancellationTokenSource)s).NotifyCancellation(throwOnFirstException: false); // skip ThrowIfDisposed() check in Cancel() }; @@ -939,7 +942,7 @@ internal LinkedNCancellationTokenSource(CancellationToken[] tokens) { if (tokens[i].CanBeCanceled) { - _linkingRegistrations[i] = tokens[i].UnsafeRegister(s_linkedTokenCancelDelegate, this); + _linkingRegistrations[i] = tokens[i].UnsafeRegister(s_linkedTokenCancelDelegate, new WeakReference(this)); } // Empty slots in the array will be default(CancellationTokenRegistration), which are nops to Dispose. // Based on usage patterns, such occurrences should also be rare, such that it's not worth resizing