Fix CancellationTokenRegistration memory leak

Havret · Havret · commit c50db43ee90f · 2021-08-24T19:50:10.000+02:00
diff --git a/src/ActiveMQ.Artemis.Client/Builders/AnonymousProducerBuilder.cs b/src/ActiveMQ.Artemis.Client/Builders/AnonymousProducerBuilder.cs
@@ -29,7 +29,7 @@ public async Task<IAnonymousProducer> CreateAsync(AnonymousProducerConfiguration
             if (configuration == null) throw new ArgumentNullException(nameof(configuration));
             
             cancellationToken.ThrowIfCancellationRequested();
-            cancellationToken.Register(() => _tcs.TrySetCanceled());
+            using var _ = cancellationToken.Register(() => _tcs.TrySetCanceled());
 
             var target = new Target
             {
diff --git a/src/ActiveMQ.Artemis.Client/Builders/ConnectionBuilder.cs b/src/ActiveMQ.Artemis.Client/Builders/ConnectionBuilder.cs
@@ -25,7 +25,7 @@ public ConnectionBuilder(ILoggerFactory loggerFactory, Func<IMessageIdPolicy> me
         public async Task<IConnection> CreateAsync(Endpoint endpoint, CancellationToken cancellationToken)
         {
             cancellationToken.ThrowIfCancellationRequested();
-            cancellationToken.Register(() => _tcs.TrySetCanceled());
+            using var _ = cancellationToken.Register(() => _tcs.TrySetCanceled());
 
             var connectionFactory = new Amqp.ConnectionFactory();
             try
diff --git a/src/ActiveMQ.Artemis.Client/Builders/ConsumerBuilder.cs b/src/ActiveMQ.Artemis.Client/Builders/ConsumerBuilder.cs
@@ -30,7 +30,7 @@ public async Task<IConsumer> CreateAsync(ConsumerConfiguration configuration, Ca
             CheckConfiguration(configuration);
 
             cancellationToken.ThrowIfCancellationRequested();
-            cancellationToken.Register(() => _tcs.TrySetCanceled());
+            using var _ = cancellationToken.Register(() => _tcs.TrySetCanceled());
 
             var source = new Source
             {
diff --git a/src/ActiveMQ.Artemis.Client/Builders/ProducerBuilder.cs b/src/ActiveMQ.Artemis.Client/Builders/ProducerBuilder.cs
@@ -34,7 +34,7 @@ public async Task<IProducer> CreateAsync(ProducerConfiguration configuration, Ca
             if (string.IsNullOrWhiteSpace(configuration.Address)) throw new ArgumentNullException(nameof(configuration.Address), "The address cannot be empty.");
 
             cancellationToken.ThrowIfCancellationRequested();
-            cancellationToken.Register(() => _tcs.TrySetCanceled());
+            using var _ = cancellationToken.Register(() => _tcs.TrySetCanceled());
 
             var target = new Target
             {
diff --git a/src/ActiveMQ.Artemis.Client/Builders/RpcClientBuilder.cs b/src/ActiveMQ.Artemis.Client/Builders/RpcClientBuilder.cs
@@ -26,7 +26,8 @@ public async Task<RpcClient> CreateAsync(string address, CancellationToken cance
 
         private async Task<SenderLink> CreateSenderLink(string address, CancellationToken cancellationToken)
         {
-            var tcs = TaskUtil.CreateTaskCompletionSource<bool>(cancellationToken);
+            var (tcs, ctr) = TaskUtil.CreateTaskCompletionSource<bool>(ref cancellationToken);
+            using var _ = ctr;
             var senderLink = new SenderLink(_session, Guid.NewGuid().ToString(), new Target
             {
                 Address = address
@@ -36,7 +37,6 @@ private async Task<SenderLink> CreateSenderLink(string address, CancellationToke
             senderLink.Closed -= OnClosed;
             return senderLink;
 
-
             void OnAttached(ILink link, Attach attach)
             {
                 if (attach != null)
@@ -56,7 +56,8 @@ void OnClosed(IAmqpObject sender, Error error)
 
         private async Task<(ReceiverLink receiverLink, string address)> CreateReceiverLink(CancellationToken cancellationToken)
         {
-            var tcs = TaskUtil.CreateTaskCompletionSource<string>(cancellationToken);
+            var (tcs, ctr) = TaskUtil.CreateTaskCompletionSource<string>(ref cancellationToken);
+            using var _ = ctr;
             var receiverLink = new ReceiverLink(_session, Guid.NewGuid().ToString(), new Source
             {
                 Dynamic = true
@@ -68,7 +69,7 @@ void OnClosed(IAmqpObject sender, Error error)
 
             void OnAttached(ILink link, Attach attach)
             {
-                if (attach != null && attach.Source is Source source)
+                if (attach is { Source: Source source })
                 {
                     tcs.TrySetResult(source.Address);
                 }
diff --git a/src/ActiveMQ.Artemis.Client/Builders/SessionBuilder.cs b/src/ActiveMQ.Artemis.Client/Builders/SessionBuilder.cs
@@ -23,7 +23,7 @@ public SessionBuilder(Amqp.Connection connection)
         public async Task<Session> CreateAsync(CancellationToken cancellationToken)
         {
             cancellationToken.ThrowIfCancellationRequested();
-            cancellationToken.Register(() => _tcs.TrySetCanceled());
+            using var _ = cancellationToken.Register(() => _tcs.TrySetCanceled());
 
             var begin = new Begin
             {
diff --git a/src/ActiveMQ.Artemis.Client/Builders/TransactionCoordinatorBuilder.cs b/src/ActiveMQ.Artemis.Client/Builders/TransactionCoordinatorBuilder.cs
@@ -23,7 +23,7 @@ public TransactionCoordinatorBuilder(Session session)
         public async Task<TransactionCoordinator> CreateAsync(CancellationToken cancellationToken)
         {
             cancellationToken.ThrowIfCancellationRequested();
-            cancellationToken.Register(() => _tcs.TrySetCanceled());
+            using var _ = cancellationToken.Register(() => _tcs.TrySetCanceled());
 
             var attach = new Attach
             {
diff --git a/src/ActiveMQ.Artemis.Client/InternalUtilities/TaskUtil.cs b/src/ActiveMQ.Artemis.Client/InternalUtilities/TaskUtil.cs
@@ -1,19 +1,27 @@
-﻿using System.Threading;
+﻿using System;
+using System.Threading;
 using System.Threading.Tasks;
 
 namespace ActiveMQ.Artemis.Client.InternalUtilities
 {
     internal static class TaskUtil
     {
-        public static TaskCompletionSource<T> CreateTaskCompletionSource<T>(CancellationToken cancellationToken)
+        public static (TaskCompletionSource<T> tcs, CancellationTokenRegistration ctr) CreateTaskCompletionSource<T>(ref CancellationToken cancellationToken)
+        {
+            return CreateTaskCompletionSource<T>(ref cancellationToken, null);
+        }
+        
+        public static (TaskCompletionSource<T> tcs, CancellationTokenRegistration ctr) CreateTaskCompletionSource<T>(ref CancellationToken cancellationToken, Action cleanup)
         {
             var tcs = new TaskCompletionSource<T>(TaskCreationOptions.RunContinuationsAsynchronously);
-            if (cancellationToken != default)
+            var ctr = cancellationToken != default ? cancellationToken.Register(() =>
             {
-                cancellationToken.Register(() => tcs.TrySetCanceled());
-            }
-
-            return tcs;
+                if (tcs.TrySetCanceled())
+                {
+                    cleanup?.Invoke();
+                }
+            }) : default;
+            return (tcs, ctr);
         }
     }
 }
diff --git a/src/ActiveMQ.Artemis.Client/Management/RpcClient.cs b/src/ActiveMQ.Artemis.Client/Management/RpcClient.cs
@@ -41,7 +41,9 @@ public async Task<Message> SendAsync(Message message, CancellationToken cancella
             message.SetCorrelationId(correlationId);
             message.Properties.ReplyTo = _replyToAddress;
 
-            var tcs = TaskUtil.CreateTaskCompletionSource<Message>(cancellationToken);
+            var (tcs, ctr) = TaskUtil.CreateTaskCompletionSource<Message>(ref cancellationToken);
+            using var cancellationTokenRegistration = ctr;
+            cancellationToken.Register(() => tcs.TrySetCanceled());
             try
             {
                 _pendingRequests.TryAdd(correlationId, tcs);
diff --git a/src/ActiveMQ.Artemis.Client/ProducerBase.cs b/src/ActiveMQ.Artemis.Client/ProducerBase.cs
@@ -39,14 +39,11 @@ protected async Task SendInternalAsync(string address, RoutingType? routingType,
 
             var txnId = await _transactionsManager.GetTxnIdAsync(transaction, cancellationToken).ConfigureAwait(false);
             var transactionalState = txnId != null ? new TransactionalState { TxnId = txnId } : null;
-            var tcs = TaskUtil.CreateTaskCompletionSource<bool>(cancellationToken);
-            cancellationToken.Register(() =>
+            var (tcs, ctr) = TaskUtil.CreateTaskCompletionSource<bool>(ref cancellationToken, () =>
             {
-                if (tcs.TrySetCanceled())
-                {
-                    _senderLink.Cancel(message.InnerMessage);    
-                }
+                _senderLink.Cancel(message.InnerMessage);
             });
+            using var _ = ctr;
             message.DurabilityMode ??= _configuration.MessageDurabilityMode ?? DurabilityMode.Durable;
             Send(address, routingType, message, transactionalState, _onOutcome, tcs);
             await tcs.Task.ConfigureAwait(false);
diff --git a/src/ActiveMQ.Artemis.Client/Transactions/TransactionCoordinator.cs b/src/ActiveMQ.Artemis.Client/Transactions/TransactionCoordinator.cs
@@ -20,20 +20,22 @@ public TransactionCoordinator(SenderLink senderLink)
             _senderLink = senderLink;
         }
         
-        public Task<byte[]> DeclareAsync(CancellationToken cancellationToken)
+        public async Task<byte[]> DeclareAsync(CancellationToken cancellationToken)
         {
             var message = new Amqp.Message(new Declare());
-            var tcs = TaskUtil.CreateTaskCompletionSource<byte[]>(cancellationToken);
+            var (tcs, ctr) = TaskUtil.CreateTaskCompletionSource<byte[]>(ref cancellationToken);
+            using var _ = ctr;
             _senderLink.Send(message, null, _onDeclareOutcome, tcs);
-            return tcs.Task;
+            return await tcs.Task;
         }
 
-        public Task DischargeAsync(byte[] txnId, bool fail, CancellationToken cancellationToken)
+        public async Task DischargeAsync(byte[] txnId, bool fail, CancellationToken cancellationToken)
         {
             var message = new Amqp.Message(new Discharge { TxnId = txnId, Fail = fail });
-            var tcs = TaskUtil.CreateTaskCompletionSource<bool>(cancellationToken);
+            var (tcs, ctr) = TaskUtil.CreateTaskCompletionSource<bool>(ref cancellationToken);
+            using var _ = ctr;
             _senderLink.Send(message, null, _onDischargeOutcome, tcs);
-            return tcs.Task;
+            await tcs.Task;
         }
         
         private static void OnDeclareOutcome(ILink link, Amqp.Message message, Outcome outcome, object state)
diff --git a/test/ActiveMQ.Artemis.Client.UnitTests/ActiveMQNetSpec.cs b/test/ActiveMQ.Artemis.Client.UnitTests/ActiveMQNetSpec.cs
@@ -86,7 +86,7 @@ protected static async Task DisposeHostAndWaitUntilConnectionNotified(TestContai
         {
             var tcs = new TaskCompletionSource<bool>();
             using var cts = new CancellationTokenSource(Timeout);
-            cts.Token.Register(() => tcs.TrySetCanceled());
+            await using var _ = cts.Token.Register(() => tcs.TrySetCanceled());
             connection.ConnectionClosed += (_, _) =>
             {
                 tcs.TrySetResult(true);
@@ -99,7 +99,7 @@ protected static async Task WaitUntilConnectionRecovered(IConnection connection)
         {
             var tcs = new TaskCompletionSource<bool>();
             using var cts = new CancellationTokenSource(Timeout);
-            cts.Token.Register(() => tcs.TrySetCanceled());
+            await using var _ = cts.Token.Register(() => tcs.TrySetCanceled());
             connection.ConnectionRecovered += (_, _) =>
             {
                 tcs.TrySetResult(true);

Original file line number	Diff line number	Diff line change
`@@ -29,7 +29,7 @@ public async Task<IAnonymousProducer> CreateAsync(AnonymousProducerConfiguration`
`29`	`29`	`if (configuration == null) throw new ArgumentNullException(nameof(configuration));`
`30`	`30`
`31`	`31`	`cancellationToken.ThrowIfCancellationRequested();`
`32`		`- cancellationToken.Register(() => _tcs.TrySetCanceled());`
	`32`	`+ using var _ = cancellationToken.Register(() => _tcs.TrySetCanceled());`
`33`	`33`
`34`	`34`	`var target = new Target`
`35`	`35`	`{`
Original file line number	Diff line number	Diff line change
`@@ -25,7 +25,7 @@ public ConnectionBuilder(ILoggerFactory loggerFactory, Func<IMessageIdPolicy> me`
`25`	`25`	`public async Task<IConnection> CreateAsync(Endpoint endpoint, CancellationToken cancellationToken)`
`26`	`26`	`{`
`27`	`27`	`cancellationToken.ThrowIfCancellationRequested();`
`28`		`- cancellationToken.Register(() => _tcs.TrySetCanceled());`
	`28`	`+ using var _ = cancellationToken.Register(() => _tcs.TrySetCanceled());`
`29`	`29`
`30`	`30`	`var connectionFactory = new Amqp.ConnectionFactory();`
`31`	`31`	`try`
Original file line number	Diff line number	Diff line change
`@@ -30,7 +30,7 @@ public async Task<IConsumer> CreateAsync(ConsumerConfiguration configuration, Ca`
`30`	`30`	`CheckConfiguration(configuration);`
`31`	`31`
`32`	`32`	`cancellationToken.ThrowIfCancellationRequested();`
`33`		`- cancellationToken.Register(() => _tcs.TrySetCanceled());`
	`33`	`+ using var _ = cancellationToken.Register(() => _tcs.TrySetCanceled());`
`34`	`34`
`35`	`35`	`var source = new Source`
`36`	`36`	`{`
Original file line number	Diff line number	Diff line change
`@@ -34,7 +34,7 @@ public async Task<IProducer> CreateAsync(ProducerConfiguration configuration, Ca`
`34`	`34`	`if (string.IsNullOrWhiteSpace(configuration.Address)) throw new ArgumentNullException(nameof(configuration.Address), "The address cannot be empty.");`
`35`	`35`
`36`	`36`	`cancellationToken.ThrowIfCancellationRequested();`
`37`		`- cancellationToken.Register(() => _tcs.TrySetCanceled());`
	`37`	`+ using var _ = cancellationToken.Register(() => _tcs.TrySetCanceled());`
`38`	`38`
`39`	`39`	`var target = new Target`
`40`	`40`	`{`
Original file line number	Diff line number	Diff line change
`@@ -26,7 +26,8 @@ public async Task<RpcClient> CreateAsync(string address, CancellationToken cance`
`26`	`26`
`27`	`27`	`private async Task<SenderLink> CreateSenderLink(string address, CancellationToken cancellationToken)`
`28`	`28`	`{`
`29`		`- var tcs = TaskUtil.CreateTaskCompletionSource<bool>(cancellationToken);`
	`29`	`+ var (tcs, ctr) = TaskUtil.CreateTaskCompletionSource<bool>(ref cancellationToken);`
	`30`	`+ using var _ = ctr;`
`30`	`31`	`var senderLink = new SenderLink(_session, Guid.NewGuid().ToString(), new Target`
`31`	`32`	`{`
`32`	`33`	`Address = address`
`@@ -36,7 +37,6 @@ private async Task<SenderLink> CreateSenderLink(string address, CancellationToke`
`36`	`37`	`senderLink.Closed -= OnClosed;`
`37`	`38`	`return senderLink;`
`38`	`39`
`39`		`-`
`40`	`40`	`void OnAttached(ILink link, Attach attach)`
`41`	`41`	`{`
`42`	`42`	`if (attach != null)`
`@@ -56,7 +56,8 @@ void OnClosed(IAmqpObject sender, Error error)`
`56`	`56`
`57`	`57`	`private async Task<(ReceiverLink receiverLink, string address)> CreateReceiverLink(CancellationToken cancellationToken)`
`58`	`58`	`{`
`59`		`- var tcs = TaskUtil.CreateTaskCompletionSource<string>(cancellationToken);`
	`59`	`+ var (tcs, ctr) = TaskUtil.CreateTaskCompletionSource<string>(ref cancellationToken);`
	`60`	`+ using var _ = ctr;`
`60`	`61`	`var receiverLink = new ReceiverLink(_session, Guid.NewGuid().ToString(), new Source`
`61`	`62`	`{`
`62`	`63`	`Dynamic = true`
`@@ -68,7 +69,7 @@ void OnClosed(IAmqpObject sender, Error error)`
`68`	`69`
`69`	`70`	`void OnAttached(ILink link, Attach attach)`
`70`	`71`	`{`
`71`		`- if (attach != null && attach.Source is Source source)`
	`72`	`+ if (attach is { Source: Source source })`
`72`	`73`	`{`
`73`	`74`	`tcs.TrySetResult(source.Address);`
`74`	`75`	`}`
Original file line number	Diff line number	Diff line change
`@@ -23,7 +23,7 @@ public SessionBuilder(Amqp.Connection connection)`
`23`	`23`	`public async Task<Session> CreateAsync(CancellationToken cancellationToken)`
`24`	`24`	`{`
`25`	`25`	`cancellationToken.ThrowIfCancellationRequested();`
`26`		`- cancellationToken.Register(() => _tcs.TrySetCanceled());`
	`26`	`+ using var _ = cancellationToken.Register(() => _tcs.TrySetCanceled());`
`27`	`27`
`28`	`28`	`var begin = new Begin`
`29`	`29`	`{`
Original file line number	Diff line number	Diff line change
`@@ -23,7 +23,7 @@ public TransactionCoordinatorBuilder(Session session)`
`23`	`23`	`public async Task<TransactionCoordinator> CreateAsync(CancellationToken cancellationToken)`
`24`	`24`	`{`
`25`	`25`	`cancellationToken.ThrowIfCancellationRequested();`
`26`		`- cancellationToken.Register(() => _tcs.TrySetCanceled());`
	`26`	`+ using var _ = cancellationToken.Register(() => _tcs.TrySetCanceled());`
`27`	`27`
`28`	`28`	`var attach = new Attach`
`29`	`29`	`{`
Original file line number	Diff line number	Diff line change
`@@ -1,19 +1,27 @@`
`1`		`-using System.Threading;`
	`1`	`+using System;`
	`2`	`+using System.Threading;`
`2`	`3`	`using System.Threading.Tasks;`
`3`	`4`
`4`	`5`	`namespace ActiveMQ.Artemis.Client.InternalUtilities`
`5`	`6`	`{`
`6`	`7`	`internal static class TaskUtil`
`7`	`8`	`{`
`8`		`- public static TaskCompletionSource<T> CreateTaskCompletionSource<T>(CancellationToken cancellationToken)`
	`9`	`+ public static (TaskCompletionSource<T> tcs, CancellationTokenRegistration ctr) CreateTaskCompletionSource<T>(ref CancellationToken cancellationToken)`
	`10`	`+ {`
	`11`	`+ return CreateTaskCompletionSource<T>(ref cancellationToken, null);`
	`12`	`+ }`
	`13`	`+`
	`14`	`+ public static (TaskCompletionSource<T> tcs, CancellationTokenRegistration ctr) CreateTaskCompletionSource<T>(ref CancellationToken cancellationToken, Action cleanup)`
`9`	`15`	`{`
`10`	`16`	`var tcs = new TaskCompletionSource<T>(TaskCreationOptions.RunContinuationsAsynchronously);`
`11`		`- if (cancellationToken != default)`
	`17`	`+ var ctr = cancellationToken != default ? cancellationToken.Register(() =>`
`12`	`18`	`{`
`13`		`- cancellationToken.Register(() => tcs.TrySetCanceled());`
`14`		`- }`
`15`		`-`
`16`		`- return tcs;`
	`19`	`+ if (tcs.TrySetCanceled())`
	`20`	`+ {`
	`21`	`+ cleanup?.Invoke();`
	`22`	`+ }`
	`23`	`+ }) : default;`
	`24`	`+ return (tcs, ctr);`
`17`	`25`	`}`
`18`	`26`	`}`
`19`	`27`	`}`