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> ExecuteCoreAsync ((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> ExecuteCoreAsync 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;