Pool CancellationTokenSources

App-vNext · Apr 12, 2023 · 22a7ea7 · 22a7ea7
1 parent 97c2307
commit 22a7ea7
Show file tree

Hide file tree

Showing 4 changed files with 75 additions and 8 deletions.
diff --git a/src/Polly.Core.Benchmarks/README.md b/src/Polly.Core.Benchmarks/README.md
@@ -28,10 +28,10 @@ LaunchCount=2  WarmupCount=10
 
 ## TIMEOUT
 
-|            Method |     Mean |   Error |  StdDev | Ratio | RatioSD |   Gen0 | Allocated | Alloc Ratio |
-|------------------ |---------:|--------:|--------:|------:|--------:|-------:|----------:|------------:|
-| ExecuteTimeout_V7 | 281.5 ns | 5.76 ns | 8.08 ns |  1.00 |    0.00 | 0.0868 |     728 B |        1.00 |
-| ExecuteTimeout_V8 | 268.9 ns | 3.86 ns | 5.53 ns |  0.96 |    0.04 | 0.0257 |     216 B |        0.30 |
+|            Method |     Mean |   Error |   StdDev | Ratio | RatioSD |   Gen0 | Allocated | Alloc Ratio |
+|------------------ |---------:|--------:|---------:|------:|--------:|-------:|----------:|------------:|
+| ExecuteTimeout_V7 | 304.9 ns | 7.53 ns | 11.27 ns |  1.00 |    0.00 | 0.0868 |     728 B |        1.00 |
+| ExecuteTimeout_V8 | 266.5 ns | 5.95 ns |  8.72 ns |  0.88 |    0.04 |      - |         - |        0.00 |
 
 ## RETRY
 
@@ -51,5 +51,5 @@ LaunchCount=2  WarmupCount=10
 
 |                     Method |     Mean |     Error |    StdDev | Ratio | RatioSD |   Gen0 | Allocated | Alloc Ratio |
 |--------------------------- |---------:|----------:|----------:|------:|--------:|-------:|----------:|------------:|
-| ExecuteStrategyPipeline_V7 | 1.321 us | 0.0355 us | 0.0520 us |  1.00 |    0.00 | 0.2861 |    2400 B |        1.00 |
-| ExecuteStrategyPipeline_V8 | 1.126 us | 0.0193 us | 0.0283 us |  0.85 |    0.03 | 0.0763 |     640 B |        0.27 |
+| ExecuteStrategyPipeline_V7 | 1.265 us | 0.0372 us | 0.0558 us |  1.00 |    0.00 | 0.2861 |    2400 B |        1.00 |
+| ExecuteStrategyPipeline_V8 | 1.032 us | 0.0165 us | 0.0236 us |  0.82 |    0.04 | 0.0076 |      64 B |        0.03 |
diff --git a/src/Polly.Core.Tests/Utils/CancellationTokenSourcePoolTests.cs b/src/Polly.Core.Tests/Utils/CancellationTokenSourcePoolTests.cs
@@ -0,0 +1,33 @@
+using Polly.Utils;
+
+namespace Polly.Core.Tests.Utils;
+
+public class CancellationTokenSourcePoolTests
+{
+    [Fact]
+    public void RentReturn_Reusable_EnsureProperBehavior()
+    {
+        var cts = CancellationTokenSourcePool.Get();
+        CancellationTokenSourcePool.Return(cts);
+
+        var cts2 = CancellationTokenSourcePool.Get();
+#if NET6_0_OR_GREATER
+        cts2.Should().BeSameAs(cts);
+#else
+        cts2.Should().NotBeSameAs(cts);
+#endif
+    }
+
+    [Fact]
+    public void RentReturn_NotReusable_EnsureProperBehavior()
+    {
+        var cts = CancellationTokenSourcePool.Get();
+        cts.Cancel();
+        CancellationTokenSourcePool.Return(cts);
+
+        cts.Invoking(c => c.Token).Should().Throw<ObjectDisposedException>();
+
+        var cts2 = CancellationTokenSourcePool.Get();
+        cts2.Token.Should().NotBeNull();
+    }
+}
diff --git a/src/Polly.Core/Timeout/TimeoutResilienceStrategy.cs b/src/Polly.Core/Timeout/TimeoutResilienceStrategy.cs
@@ -35,7 +35,7 @@ protected internal override async ValueTask<TResult> ExecuteCoreAsync<TResult, T
         }
 
         var previousToken = context.CancellationToken;
-        using var cancellationSource = new CancellationTokenSource();
+        var cancellationSource = CancellationTokenSourcePool.Get();
         _timeProvider.CancelAfter(cancellationSource, timeout);
         context.CancellationToken = cancellationSource.Token;
 
@@ -56,6 +56,8 @@ protected internal override async ValueTask<TResult> ExecuteCoreAsync<TResult, T
         }
         catch (OperationCanceledException e) when (cancellationSource.IsCancellationRequested && !previousToken.IsCancellationRequested)
         {
+            context.CancellationToken = previousToken;
+
             _telemetry.Report(TimeoutConstants.OnTimeoutEvent, context);
 
             if (OnTimeout != null)
@@ -73,10 +75,11 @@ protected internal override async ValueTask<TResult> ExecuteCoreAsync<TResult, T
         finally
         {
             context.CancellationToken = previousToken;
+            CancellationTokenSourcePool.Return(cancellationSource);
         }
     }
 
-    internal async Task<TimeSpan> GetTimeoutAsync(ResilienceContext context)
+    internal async ValueTask<TimeSpan> GetTimeoutAsync(ResilienceContext context)
     {
         if (TimeoutGenerator == null)
         {

diff --git a/src/Polly.Core/Utils/CancellationTokenSourcePool.cs b/src/Polly.Core/Utils/CancellationTokenSourcePool.cs
@@ -0,0 +1,31 @@
+namespace Polly.Utils
+{
+    internal static class CancellationTokenSourcePool
+    {
+#if NET6_0_OR_GREATER
+        private static readonly ObjectPool<CancellationTokenSource> Pool = new(
+            static () => new CancellationTokenSource(),
+            static cts => true);
+#endif
+        public static CancellationTokenSource Get()
+        {
+#if NET6_0_OR_GREATER
+            return Pool.Get();
+#else
+            return new CancellationTokenSource();
+#endif
+        }
+
+        public static void Return(CancellationTokenSource source)
+        {
+#if NET6_0_OR_GREATER
+            if (source.TryReset())
+            {
+                Pool.Return(source);
+                return;
+            }
+#endif
+            source.Dispose();
+        }
+    }
+}