From 2aa54def33e92b85f8668995dcee949713cfa152 Mon Sep 17 00:00:00 2001
From: Martin Tomka <martintomka@microsoft.com>
Date: Tue, 23 May 2023 08:26:38 +0200
Subject: [PATCH] fixes

---
 .../Timeout/TimeoutResilienceStrategy.cs      | 23 +++++++------------
 1 file changed, 8 insertions(+), 15 deletions(-)

diff --git a/src/Polly.Core/Timeout/TimeoutResilienceStrategy.cs b/src/Polly.Core/Timeout/TimeoutResilienceStrategy.cs
index 6f337b7d272..e1aa2f6e1c6 100644
--- a/src/Polly.Core/Timeout/TimeoutResilienceStrategy.cs
+++ b/src/Polly.Core/Timeout/TimeoutResilienceStrategy.cs
@@ -42,13 +42,11 @@ protected internal override async ValueTask<Outcome<TResult>> ExecuteCoreAsync<T
         var cancellationSource = _cancellationTokenSourcePool.Get(timeout);
         context.CancellationToken = cancellationSource.Token;
 
-        CancellationTokenRegistration? registration = null;
-
-        if (previousToken.CanBeCanceled)
-        {
-            registration = previousToken.Register(static state => ((CancellationTokenSource)state!).Cancel(), cancellationSource, useSynchronizationContext: false);
-        }
-
+#if NETCOREAPP
+        await using var registration = CreateRegistration(cancellationSource, previousToken).ConfigureAwait(context.ContinueOnCapturedContext);
+#else
+        using var registration = CreateRegistration(cancellationSource, previousToken);
+#endif
         try
         {
             var outcome = await callback(context, state).ConfigureAwait(context.ContinueOnCapturedContext);
@@ -77,7 +75,6 @@ protected internal override async ValueTask<Outcome<TResult>> ExecuteCoreAsync<T
         }
         finally
         {
-            await DisposeRegistration(registration).ConfigureAwait(context.ContinueOnCapturedContext);
             context.CancellationToken = previousToken;
             _cancellationTokenSourcePool.Return(cancellationSource);
         }
@@ -93,15 +90,11 @@ internal ValueTask<TimeSpan> GetTimeoutAsync(ResilienceContext context)
         return TimeoutGenerator(new TimeoutGeneratorArguments(context));
     }
 
-    private static ValueTask DisposeRegistration(CancellationTokenRegistration? registration)
+    private static CancellationTokenRegistration CreateRegistration(CancellationTokenSource cancellationSource, CancellationToken previousToken)
     {
-        if (registration.HasValue)
+        if (previousToken.CanBeCanceled)
         {
-#if NETCOREAPP
-            return registration.Value.DisposeAsync();
-#else
-            registration.Value.Dispose();
-#endif
+            return previousToken.Register(static state => ((CancellationTokenSource)state!).Cancel(), cancellationSource, useSynchronizationContext: false);
         }
 
         return default;