diff --git a/src/DistributedLock.Core/Internal/Data/ConnectionMonitor.cs b/src/DistributedLock.Core/Internal/Data/ConnectionMonitor.cs index 41d796f..97fe43c 100644 --- a/src/DistributedLock.Core/Internal/Data/ConnectionMonitor.cs +++ b/src/DistributedLock.Core/Internal/Data/ConnectionMonitor.cs @@ -277,11 +277,7 @@ private void CloseOrCancelMonitoringHandleRegistrationsNoLock(bool isCancel) if (isCancel) { // cancel in a background thread in case we have hangs or errors - Task.Run(() => - { - try { cancellationTokenSource.Cancel(); } - finally { cancellationTokenSource.Dispose(); } - }); + _ = CancelAndDisposeAsync(cancellationTokenSource); } else { @@ -327,11 +323,21 @@ private void FireStateChangedNoLock() // it is still safer and easier to reason about not to have that happen. This also ensures // that FireStateChangedNoLock() always returns quickly, even if the monitoring loop // were to do some synchronous work on the continuation thread. - Task.Run(() => + _ = CancelAndDisposeAsync(monitorStateChangedTokenSource); + } + + private static async Task CancelAndDisposeAsync( + CancellationTokenSource cancellationTokenSource) + { + try { - try { monitorStateChangedTokenSource.Cancel(); } - finally { monitorStateChangedTokenSource.Dispose(); } - }); + await cancellationTokenSource.CancelAsync() + .ConfigureAwait(false); + } + finally + { + cancellationTokenSource.Dispose(); + } } private async Task MonitorWorkerLoop() diff --git a/src/DistributedLock.Core/Internal/Helpers.cs b/src/DistributedLock.Core/Internal/Helpers.cs index 84a2916..2644005 100644 --- a/src/DistributedLock.Core/Internal/Helpers.cs +++ b/src/DistributedLock.Core/Internal/Helpers.cs @@ -106,6 +106,14 @@ public static bool TryGetValue(this T? nullable, out T value) value = nullable.GetValueOrDefault(); return nullable.HasValue; } + +#if !NET8_0_OR_GREATER + public static Task CancelAsync( + this CancellationTokenSource cancellationTokenSource) + { + return Task.Run(cancellationTokenSource.Cancel); + } +#endif } /// diff --git a/src/DistributedLock.Core/Internal/LeaseMonitor.cs b/src/DistributedLock.Core/Internal/LeaseMonitor.cs index 7154912..6b5107b 100644 --- a/src/DistributedLock.Core/Internal/LeaseMonitor.cs +++ b/src/DistributedLock.Core/Internal/LeaseMonitor.cs @@ -107,7 +107,7 @@ private static async Task RunMonitoringLoopIterationAsync(WeakReference monitor._cancellationTask = Task.Run(() => monitor._handleLostSource.Cancel()); + void OnHandleLost() => monitor._cancellationTask = monitor._handleLostSource.CancelAsync(); } private async Task CheckLeaseAsync() diff --git a/src/DistributedLock.Tests/Tests/Core/HelpersTest.cs b/src/DistributedLock.Tests/Tests/Core/HelpersTest.cs index abd7961..3ef7a7e 100644 --- a/src/DistributedLock.Tests/Tests/Core/HelpersTest.cs +++ b/src/DistributedLock.Tests/Tests/Core/HelpersTest.cs @@ -77,4 +77,33 @@ async Task TryAwaitFailedTask() await task.TryAwait(); } } + + [Test] + public async Task TestCancelAsyncDoesNotBlockOnCallbacks() + { + using var cts = new CancellationTokenSource(); + + var callbackStarted = new TaskCompletionSource( + TaskCreationOptions.RunContinuationsAsynchronously); + + var releaseCallback = new TaskCompletionSource( + TaskCreationOptions.RunContinuationsAsynchronously); + + using var registration = cts.Token.Register(() => + { + callbackStarted.SetResult(true); + + releaseCallback.Task.GetAwaiter().GetResult(); + }); + + var cancellationTask = cts.CancelAsync(); + + await callbackStarted.Task; + + Assert.That(cancellationTask.IsCompleted, Is.False); + + releaseCallback.SetResult(true); + + await cancellationTask; + } }