From b13cff5a953124d61b719e31b8302d542d2a611c Mon Sep 17 00:00:00 2001
From: James Carter <james@loveelectronics.co.uk>
Date: Mon, 24 Feb 2025 07:33:26 +0000
Subject: [PATCH] Fix memory leak in AsyncAutoResetEvent (#229)

* Fix memory leak in AsyncAutoResetEvent

CancellationTokenRegistration MUST be disposed, as the CancellationToken passed is saved for the lifetime of the Client, and registrations build up forever.
---
 .../AsyncResetEventTests.cs                   |  1 +
 .../Objects/AsyncAutoResetEvent.cs            | 63 ++++++++++---------
 2 files changed, 36 insertions(+), 28 deletions(-)

diff --git a/CryptoExchange.Net.UnitTests/AsyncResetEventTests.cs b/CryptoExchange.Net.UnitTests/AsyncResetEventTests.cs
index db829e3..992d2df 100644
--- a/CryptoExchange.Net.UnitTests/AsyncResetEventTests.cs
+++ b/CryptoExchange.Net.UnitTests/AsyncResetEventTests.cs
@@ -106,6 +106,7 @@ namespace CryptoExchange.Net.UnitTests
             for(var i = 1; i <= 10; i++)
             {
                 evnt.Set();
+                await Task.Delay(1); // Wait for the continuation.
                 Assert.That(10 - i == waiters.Count(w => w.Status != TaskStatus.RanToCompletion));
             }
 
diff --git a/CryptoExchange.Net/Objects/AsyncAutoResetEvent.cs b/CryptoExchange.Net/Objects/AsyncAutoResetEvent.cs
index 9c6f77e..ff4895d 100644
--- a/CryptoExchange.Net/Objects/AsyncAutoResetEvent.cs
+++ b/CryptoExchange.Net/Objects/AsyncAutoResetEvent.cs
@@ -32,44 +32,51 @@ namespace CryptoExchange.Net.Objects
         /// Wait for the AutoResetEvent to be set
         /// </summary>
         /// <returns></returns>
-        public Task<bool> WaitAsync(TimeSpan? timeout = null, CancellationToken ct = default)
+        public async Task<bool> WaitAsync(TimeSpan? timeout = null, CancellationToken ct = default)
         {
-            lock (_waits)
+            CancellationTokenRegistration registration = default;
+            try
             {
-                if (_signaled)
+                Task<bool> waiter = _completed;
+                lock (_waits)
                 {
-                    if(_reset)
-                        _signaled = false;
-                    return _completed;
-                }
-                else
-                {
-                    if (ct.IsCancellationRequested)
-                        return _completed;
-
-                    var tcs = new TaskCompletionSource<bool>(TaskCreationOptions.RunContinuationsAsynchronously);
-                    if (timeout.HasValue)
+                    if (_signaled)
                     {
-                        var timeoutSource = new CancellationTokenSource(timeout.Value);
-                        var cancellationSource = CancellationTokenSource.CreateLinkedTokenSource(timeoutSource.Token, ct);
-                        ct = cancellationSource.Token;
+                        if (_reset)
+                            _signaled = false;
                     }
-
-                    var registration = ct.Register(() =>
+                    else if (!ct.IsCancellationRequested)
                     {
-                        lock (_waits)
+                        var tcs = new TaskCompletionSource<bool>(TaskCreationOptions.RunContinuationsAsynchronously);
+                        if (timeout.HasValue)
                         {
-                            tcs.TrySetResult(false);
-
-                            // Not the cleanest but it works
-                            _waits = new Queue<TaskCompletionSource<bool>>(_waits.Where(i => i != tcs));
+                            var timeoutSource = new CancellationTokenSource(timeout.Value);
+                            var cancellationSource = CancellationTokenSource.CreateLinkedTokenSource(timeoutSource.Token, ct);
+                            ct = cancellationSource.Token;
                         }
-                    }, useSynchronizationContext: false);
-                    
 
-                    _waits.Enqueue(tcs);
-                    return tcs.Task;
+                        registration = ct.Register(() =>
+                        {
+                            lock (_waits)
+                            {
+                                tcs.TrySetResult(false);
+
+                                // Not the cleanest but it works
+                                _waits = new Queue<TaskCompletionSource<bool>>(_waits.Where(i => i != tcs));
+                            }
+                        }, useSynchronizationContext: false);
+
+
+                        _waits.Enqueue(tcs);
+                        waiter = tcs.Task;
+                    }
                 }
+
+                return await waiter.ConfigureAwait(false);
+            }
+            finally
+            {
+                registration.Dispose();
             }
         }