From 4562b5b650437967925738110c22fa6170394016 Mon Sep 17 00:00:00 2001
From: Solution <egorshugaev6@gmail.com>
Date: Wed, 12 Aug 2026 04:51:01 +0300
Subject: [PATCH] Add a CancellationTokenSource.CancelAsync shim

---
 .../Internal/Data/ConnectionMonitor.cs        | 24 +++++++++------
 src/DistributedLock.Core/Internal/Helpers.cs  |  8 +++++
 .../Internal/LeaseMonitor.cs                  |  2 +-
 .../Tests/Core/HelpersTest.cs                 | 29 +++++++++++++++++++
 4 files changed, 53 insertions(+), 10 deletions(-)

diff --git a/src/DistributedLock.Core/Internal/Data/ConnectionMonitor.cs b/src/DistributedLock.Core/Internal/Data/ConnectionMonitor.cs
index 41d796f5..97fe43cf 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 84a2916b..26440053 100644
--- a/src/DistributedLock.Core/Internal/Helpers.cs
+++ b/src/DistributedLock.Core/Internal/Helpers.cs
@@ -106,6 +106,14 @@ public static bool TryGetValue<T>(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
 }
 
 /// <summary>
diff --git a/src/DistributedLock.Core/Internal/LeaseMonitor.cs b/src/DistributedLock.Core/Internal/LeaseMonitor.cs
index 71549128..6b5107bf 100644
--- a/src/DistributedLock.Core/Internal/LeaseMonitor.cs
+++ b/src/DistributedLock.Core/Internal/LeaseMonitor.cs
@@ -107,7 +107,7 @@ private static async Task<bool> RunMonitoringLoopIterationAsync(WeakReference<Le
         }
 
         // offload cancel to a background thread to avoid hangs or errors
-        void OnHandleLost() => monitor._cancellationTask = Task.Run(() => monitor._handleLostSource.Cancel());
+        void OnHandleLost() => monitor._cancellationTask = monitor._handleLostSource.CancelAsync();
     }
 
     private async Task<LeaseState> CheckLeaseAsync()
diff --git a/src/DistributedLock.Tests/Tests/Core/HelpersTest.cs b/src/DistributedLock.Tests/Tests/Core/HelpersTest.cs
index abd79610..3ef7a7ec 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<bool>(
+            TaskCreationOptions.RunContinuationsAsynchronously);
+
+        var releaseCallback = new TaskCompletionSource<bool>(
+            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;
+    }
 }