fixes

App-vNext · May 23, 2023 · 2aa54de · 2aa54de
1 parent 86b70af
commit 2aa54de
Showing 1 changed file with 8 additions and 15 deletions.
diff --git 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;