Merge pull request #1228 from limebell/chore/dispose-register

Dispose `CancellationTokenRegistration`s
planetarium · Apr 5, 2021 · 05b537f · 05b537f
2 parents 8515e82 + 7868c4c
commit 05b537f
Show file tree

Hide file tree

Showing 14 changed files with 66 additions and 33 deletions.
diff --git a/CHANGES.md b/CHANGES.md
@@ -8,18 +8,28 @@ To be released.
 
 ### Backward-incompatible API changes
 
+ -  Removed `StunMessage.Parse(Stream)` method.  [[#1228]]
+
 ### Backward-incompatible network protocol changes
 
 ### Backward-incompatible storage format changes
 
 ### Added APIs
 
+ -  Added `StunMessage.ParseAsync(Stream, CancellationToken)` method.
+    [[#1228]]
+
 ### Behavioral changes
 
 ### Bug fixes
 
+ -  Fixed memory leak due to undisposed `CancellationTokenRegistration`s.
+    [[#1228]]
+
 ### CLI tools
 
+[#1228]: https://github.com/planetarium/libplanet/pull/1218
+
 
 Version 0.11.0
 --------------
@@ -319,7 +329,6 @@ Released on March 30, 2021.
 [#1218]: https://github.com/planetarium/libplanet/pull/1218
 
 
-
 Version 0.10.3
 --------------
 

diff --git a/Libplanet.Stun.Tests/Stun/Messages/AllocateErrorResponseTest.cs b/Libplanet.Stun.Tests/Stun/Messages/AllocateErrorResponseTest.cs
@@ -27,7 +27,7 @@ public async Task ParseBytes()
             using (var stream = new MemoryStream(bytes))
             {
                 var response =
-                    (AllocateErrorResponse)await StunMessage.Parse(stream);
+                    (AllocateErrorResponse)await StunMessage.ParseAsync(stream);
                 Assert.Equal(
                     new byte[]
                     {

diff --git a/Libplanet.Stun.Tests/Stun/Messages/AllocateSuccessResponseTest.cs b/Libplanet.Stun.Tests/Stun/Messages/AllocateSuccessResponseTest.cs
@@ -28,7 +28,7 @@ public async Task ParseBytes()
             using (var stream = new MemoryStream(bytes))
             {
                 var response =
-                    (AllocateSuccessResponse)await StunMessage.Parse(stream);
+                    (AllocateSuccessResponse)await StunMessage.ParseAsync(stream);
                 Assert.Equal(
                     new byte[]
                     {

diff --git a/Libplanet.Stun.Tests/Stun/Messages/BindingRequestSucessTest.cs b/Libplanet.Stun.Tests/Stun/Messages/BindingRequestSucessTest.cs
@@ -22,7 +22,7 @@ public async Task ParseBytes()
             using (var stream = new MemoryStream(bytes))
             {
                 var response =
-                    (BindingSuccessResponse)await StunMessage.Parse(stream);
+                    (BindingSuccessResponse)await StunMessage.ParseAsync(stream);
                 Assert.Equal(
                     new IPEndPoint(IPAddress.Parse("211.176.129.90"), 54141),
                     response.MappedAddress);

diff --git a/Libplanet.Stun.Tests/Stun/Messages/ConnectionAttemptTest.cs b/Libplanet.Stun.Tests/Stun/Messages/ConnectionAttemptTest.cs
@@ -23,7 +23,7 @@ public async Task ParseBytes()
             using (var stream = new MemoryStream(bytes))
             {
                 var response =
-                    (ConnectionAttempt)await StunMessage.Parse(stream);
+                    (ConnectionAttempt)await StunMessage.ParseAsync(stream);
                 Assert.Equal(
                     new byte[] { 0x21, 0x5a, 0x86, 0x01 },
                     response.ConnectionId);

diff --git a/Libplanet.Stun.Tests/Stun/Messages/ConnectionBindSuccessResponseTest.cs b/Libplanet.Stun.Tests/Stun/Messages/ConnectionBindSuccessResponseTest.cs
@@ -22,7 +22,7 @@ public async Task ParseBytes()
 
             using (var stream = new MemoryStream(bytes))
             {
-                var response = await StunMessage.Parse(stream);
+                var response = await StunMessage.ParseAsync(stream);
                 Assert.Equal(
                     new byte[]
                     {

diff --git a/Libplanet.Stun.Tests/Stun/Messages/CreatePermissionSuccessResponseTest.cs b/Libplanet.Stun.Tests/Stun/Messages/CreatePermissionSuccessResponseTest.cs
@@ -22,7 +22,7 @@ public async Task ParseBytes()
 
             using (var stream = new MemoryStream(bytes))
             {
-                var response = await StunMessage.Parse(stream);
+                var response = await StunMessage.ParseAsync(stream);
                 Assert.Equal(
                     new byte[]
                     {

diff --git a/Libplanet.Stun.Tests/Stun/Messages/RefreshErrorResponseTest.cs b/Libplanet.Stun.Tests/Stun/Messages/RefreshErrorResponseTest.cs
@@ -23,7 +23,7 @@ public async Task ParseBytes()
 
             using (var stream = new MemoryStream(bytes))
             {
-                var response = (RefreshErrorResponse)await StunMessage.Parse(stream);
+                var response = (RefreshErrorResponse)await StunMessage.ParseAsync(stream);
                 Assert.Equal(
                     new byte[]
                     {

diff --git a/Libplanet.Stun.Tests/Stun/Messages/RefreshSuccessResponseTest.cs b/Libplanet.Stun.Tests/Stun/Messages/RefreshSuccessResponseTest.cs
@@ -24,7 +24,7 @@ public async Task ParseBytes()
             using (var stream = new MemoryStream(bytes))
             {
                 var response =
-                    (RefreshSuccessResponse)await StunMessage.Parse(stream);
+                    (RefreshSuccessResponse)await StunMessage.ParseAsync(stream);
                 Assert.Equal(
                     new byte[]
                     {

diff --git a/Libplanet.Stun/Stun/Messages/StunMessage.cs b/Libplanet.Stun/Stun/Messages/StunMessage.cs
@@ -2,6 +2,7 @@
 using System.IO;
 using System.Linq;
 using System.Security.Cryptography;
+using System.Threading;
 using System.Threading.Tasks;
 using Libplanet.Stun.Attributes;
 
@@ -80,13 +81,19 @@ public enum MessageMethod : ushort
         /// Parses <see cref="StunMessage"/> from <paramref name="stream"/>.
         /// </summary>
         /// <param name="stream">A view of a sequence of STUN packet's bytes.</param>
+        /// <param name="cancellationToken">
+        /// A cancellation token used to propagate notification that this
+        /// operation should be canceled.
+        /// </param>
         /// <returns>A <see cref="StunMessage"/> derived on
         /// bytes read from <paramref name="stream"/>.
         /// </returns>
-        public static async Task<StunMessage> Parse(Stream stream)
+        public static async Task<StunMessage> ParseAsync(
+            Stream stream,
+            CancellationToken cancellationToken = default(CancellationToken))
         {
             var header = new byte[20];
-            await stream.ReadAsync(header, 0, 20);
+            await stream.ReadAsync(header, 0, 20, cancellationToken);
 
             MessageMethod method = ParseMethod(header[0], header[1]);
             MessageClass @class = ParseClass(header[0], header[1]);
@@ -98,7 +105,7 @@ public static async Task<StunMessage> Parse(Stream stream)
             System.Array.Copy(header, 8, transactionId, 0, 12);
 
             var body = new byte[length.ToUShort()];
-            await stream.ReadAsync(body, 0, body.Length);
+            await stream.ReadAsync(body, 0, body.Length, cancellationToken);
             IEnumerable<Attribute> attributes = ParseAttributes(
                 body,
                 transactionId

diff --git a/Libplanet.Stun/Stun/TurnClient.cs b/Libplanet.Stun/Stun/TurnClient.cs
@@ -135,7 +135,7 @@ public async Task ReconnectTurn(int listenPort, CancellationToken cancellationTo
             {
                 var request = new AllocateRequest((int)lifetime.TotalSeconds);
                 await SendMessageAsync(stream, request, cancellationToken);
-                response = await ReceiveMessage(request.TransactionId);
+                response = await ReceiveMessageAsync(request.TransactionId, cancellationToken);
 
                 if (response is AllocateErrorResponse allocError)
                 {
@@ -164,7 +164,9 @@ public async Task ReconnectTurn(int listenPort, CancellationToken cancellationTo
             NetworkStream stream = _control.GetStream();
             var request = new CreatePermissionRequest(peerAddress);
             await SendMessageAsync(stream, request, cancellationToken);
-            StunMessage response = await ReceiveMessage(request.TransactionId);
+            StunMessage response = await ReceiveMessageAsync(
+                request.TransactionId,
+                cancellationToken);
 
             if (response is CreatePermissionErrorResponse)
             {
@@ -191,7 +193,8 @@ public async Task ReconnectTurn(int listenPort, CancellationToken cancellationTo
                 try
                 {
                     await SendMessageAsync(relayedStream, bindRequest, cancellationToken);
-                    StunMessage bindResponse = await StunMessage.Parse(relayedStream);
+                    StunMessage bindResponse =
+                        await StunMessage.ParseAsync(relayedStream, cancellationToken);
 
                     if (bindResponse is ConnectionBindSuccessResponse)
                     {
@@ -216,7 +219,9 @@ public async Task ReconnectTurn(int listenPort, CancellationToken cancellationTo
             NetworkStream stream = _control.GetStream();
             var request = new BindingRequest();
             await SendMessageAsync(stream, request, cancellationToken);
-            StunMessage response = await ReceiveMessage(request.TransactionId);
+            StunMessage response = await ReceiveMessageAsync(
+                request.TransactionId,
+                cancellationToken);
 
             if (response is BindingSuccessResponse success)
             {
@@ -237,7 +242,9 @@ public async Task ReconnectTurn(int listenPort, CancellationToken cancellationTo
             var request = new RefreshRequest((int)lifetime.TotalSeconds);
             await SendMessageAsync(stream, request, cancellationToken);
 
-            StunMessage response = await ReceiveMessage(request.TransactionId);
+            StunMessage response = await ReceiveMessageAsync(
+                request.TransactionId,
+                cancellationToken);
             if (response is RefreshSuccessResponse success)
             {
                 return TimeSpan.FromSeconds(success.Lifetime);
@@ -275,7 +282,7 @@ public async Task ReconnectTurn(int listenPort, CancellationToken cancellationTo
                 var request = new BindingRequest();
                 var asBytes = request.Encode(this);
                 await stream.WriteAsync(asBytes, 0, asBytes.Length, cancellationToken);
-                await StunMessage.Parse(stream);
+                await StunMessage.ParseAsync(stream, cancellationToken);
 
                 return true;
             }
@@ -366,7 +373,6 @@ private async Task RefreshAllocate(CancellationToken cancellationToken)
             CancellationToken cancellationToken)
         {
             var tcs = new TaskCompletionSource<StunMessage>();
-            cancellationToken.Register(() => tcs.TrySetCanceled());
             _responses[message.TransactionId] = tcs;
             var asBytes = message.Encode(this);
             await stream.WriteAsync(
@@ -387,18 +393,13 @@ private async Task ProcessMessage(CancellationToken cancellationToken)
                     StunMessage message;
                     try
                     {
-                        message = await StunMessage.Parse(stream);
+                        message = await StunMessage.ParseAsync(stream, cancellationToken);
                         _logger.Debug("Stun Message is: {message}", message);
                     }
                     catch (TurnClientException e)
                     {
                         _logger.Error(e, "Failed to parse StunMessage. {e}", e);
-                        foreach (TaskCompletionSource<StunMessage> tcs in _responses.Values)
-                        {
-                            tcs.TrySetCanceled();
-                        }
-
-                        _responses.Clear();
+                        ClearResponses();
                         break;
                     }
 
@@ -426,9 +427,14 @@ private async Task ProcessMessage(CancellationToken cancellationToken)
             _logger.Debug($"{nameof(ProcessMessage)} is ended. Connected: {_control.Connected}");
         }
 
-        private async Task<StunMessage> ReceiveMessage(byte[] transactionId)
+        private async Task<StunMessage> ReceiveMessageAsync(
+            byte[] transactionId,
+            CancellationToken cancellationToken)
         {
-            StunMessage response = await _responses[transactionId].Task;
+            TaskCompletionSource<StunMessage> tcs = _responses[transactionId];
+            using CancellationTokenRegistration ctr =
+                cancellationToken.Register(() => tcs.TrySetCanceled());
+            StunMessage response = await tcs.Task;
             _responses.Remove(transactionId);
 
             return response;
@@ -440,14 +446,24 @@ private void ClearSession()
             _turnTaskCts.Cancel();
             _turnTaskCts.Dispose();
             _turnTasks.Clear();
-            _responses.Clear();
+            ClearResponses();
 
             foreach (TcpClient relays in _relayedClients)
             {
                 relays.Dispose();
             }
         }
 
+        private void ClearResponses()
+        {
+            foreach (TaskCompletionSource<StunMessage> tcs in _responses.Values)
+            {
+                tcs.TrySetCanceled();
+            }
+
+            _responses.Clear();
+        }
+
         private class ByteArrayComparer : IEqualityComparer<byte[]>
         {
             public bool Equals(byte[] x, byte[] y)

diff --git a/Libplanet/Net/BlockCompletion.cs b/Libplanet/Net/BlockCompletion.cs
@@ -348,7 +348,7 @@ private int Demand(IEnumerable<HashDigest<SHA256>> blockHashes, bool retry)
                         timeoutToken,
                         ct
                     );
-                    timeoutToken.Register(() =>
+                    using CancellationTokenRegistration ctr = timeoutToken.Register(() =>
                         _logger.Debug("Timed out to wait a response from {Peer}.", peer)
                     );
                     CancellationToken linkedToken = linkedTokenSource.Token;
@@ -486,7 +486,7 @@ public async Task WaitAll(CancellationToken cancellationToken = default)
                     if (tasks.Any())
                     {
                         var tcs = new TaskCompletionSource<object>();
-                        cancellationToken.Register(
+                        using CancellationTokenRegistration ctr = cancellationToken.Register(
                             () => tcs.TrySetCanceled(),
                             useSynchronizationContext: false
                         );

diff --git a/Libplanet/Net/NetMQTransport.cs b/Libplanet/Net/NetMQTransport.cs
@@ -411,7 +411,8 @@ public void Dispose()
                 Interlocked.Increment(ref _requestCount);
 
                 // FIXME should we also cancel tcs sender side too?
-                cancellationToken.Register(() => tcs.TrySetCanceled());
+                using CancellationTokenRegistration ctr =
+                    cancellationToken.Register(() => tcs.TrySetCanceled());
                 await _requests.AddAsync(
                     new MessageRequest(reqId, message, peer, now, timeout, expectedResponses, tcs),
                     cancellationToken

diff --git a/Libplanet/Net/Swarm.cs b/Libplanet/Net/Swarm.cs
@@ -525,7 +525,7 @@ CancellationToken cancellationToken
             CancellationToken cancellationToken = default(CancellationToken)
         )
         {
-            cancellationToken.Register(() =>
+            using CancellationTokenRegistration ctr = cancellationToken.Register(() =>
                 _logger.Information("Preloading is requested to be cancelled.")
             );