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
///
///
- public Task WaitAsync(TimeSpan? timeout = null, CancellationToken ct = default)
+ public async Task WaitAsync(TimeSpan? timeout = null, CancellationToken ct = default)
{
- lock (_waits)
+ CancellationTokenRegistration registration = default;
+ try
{
- if (_signaled)
+ Task waiter = _completed;
+ lock (_waits)
{
- if(_reset)
- _signaled = false;
- return _completed;
- }
- else
- {
- if (ct.IsCancellationRequested)
- return _completed;
-
- var tcs = new TaskCompletionSource(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(TaskCreationOptions.RunContinuationsAsynchronously);
+ if (timeout.HasValue)
{
- tcs.TrySetResult(false);
-
- // Not the cleanest but it works
- _waits = new Queue>(_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>(_waits.Where(i => i != tcs));
+ }
+ }, useSynchronizationContext: false);
+
+
+ _waits.Enqueue(tcs);
+ waiter = tcs.Task;
+ }
}
+
+ return await waiter.ConfigureAwait(false);
+ }
+ finally
+ {
+ registration.Dispose();
}
}