Fix RequestCollapserEngine caching lazy with exceptions as value indefinitely (#6)

* Lazy objects caches exception that are created by the factory. This will result in the RequestCollapser's Engine to cache the exception indefinitely.

* Address PR

Author: Allen-Xiao

src/Polly.Contrib.DuplicateRequestCollapser.Specs/RequestCollapserAsyncTResultSpecs.cs

@@ -1,4 +1,5 @@
-using System.Threading.Tasks;
+using System.Collections.Concurrent;
+using System.Threading.Tasks;
 using Xunit.Abstractions;
 
 namespace Polly.Contrib.DuplicateRequestCollapser.Specs
@@ -7,11 +8,16 @@ public class RequestCollapserAsyncTResultSpecs : RequestCollapserTResultSpecsBas
     {
         public RequestCollapserAsyncTResultSpecs(ITestOutputHelper testOutputHelper) : base(testOutputHelper) { }
 
+        private ConcurrentDictionary<(bool, IKeyStrategy, ISyncLockProvider), IsPolicy> PolicyCache = new ConcurrentDictionary<(bool, IKeyStrategy, ISyncLockProvider), IsPolicy>();
+
         protected override IsPolicy GetPolicy(bool useCollapser, IKeyStrategy overrideKeyStrategy = null, ISyncLockProvider lockProvider = null)
         {
-            return useCollapser ?
+            return PolicyCache.GetOrAdd((useCollapser, overrideKeyStrategy, lockProvider), _ =>
+            {
+                return useCollapser ?
                 AsyncRequestCollapserPolicy<ResultClass>.Create(overrideKeyStrategy ?? RequestCollapserPolicy.DefaultKeyStrategy, new AsyncWrapperLockProvider(lockProvider ?? RequestCollapserPolicy.GetDefaultLockProvider()))
                 : (IAsyncPolicy<ResultClass>)Policy.NoOpAsync<ResultClass>();
+            });
         }
 
         protected override Task ExecuteThroughPolicy(IsPolicy policy, Context context, int j, bool gated)

src/Polly.Contrib.DuplicateRequestCollapser.Specs/RequestCollapserSpecsBase.TestOrchestration.cs

@@ -50,7 +50,13 @@ protected void UnderlyingExpensiveWork(int i)
             ReleaseHoldingGate();
 
             // Wait for task completion.
-            Task.WaitAll(ConcurrentTasks);
+            try
+            {
+                Task.WaitAll(ConcurrentTasks);
+            }
+            catch
+            {
+            }
             testOutputHelper.WriteLine("All tasks completed.");
 
             // Return results to caller; the caller is responsible for asserting.

src/Polly.Contrib.DuplicateRequestCollapser.Specs/RequestCollapserTResultSpecs.cs

@@ -1,4 +1,5 @@
-using System.Threading.Tasks;
+using System.Collections.Concurrent;
+using System.Threading.Tasks;
 using Xunit.Abstractions;
 
 namespace Polly.Contrib.DuplicateRequestCollapser.Specs
@@ -7,11 +8,16 @@ public class RequestCollapserTResultSpecs : RequestCollapserTResultSpecsBase
     {
         public RequestCollapserTResultSpecs(ITestOutputHelper testOutputHelper) : base(testOutputHelper) { }
 
+        private ConcurrentDictionary<(bool, IKeyStrategy, ISyncLockProvider), IsPolicy> PolicyCache = new ConcurrentDictionary<(bool, IKeyStrategy, ISyncLockProvider), IsPolicy>();
+
         protected override IsPolicy GetPolicy(bool useCollapser, IKeyStrategy overrideKeyStrategy = null, ISyncLockProvider lockProvider = null)
         {
-            return useCollapser ?
-                RequestCollapserPolicy<ResultClass>.Create(overrideKeyStrategy ?? RequestCollapserPolicy.DefaultKeyStrategy, lockProvider ?? RequestCollapserPolicy.GetDefaultLockProvider())
-                : (ISyncPolicy<ResultClass>)Policy.NoOp<ResultClass>();
+            return PolicyCache.GetOrAdd((useCollapser, overrideKeyStrategy, lockProvider), _ =>
+            {
+                return useCollapser ?
+                    RequestCollapserPolicy<ResultClass>.Create(overrideKeyStrategy ?? RequestCollapserPolicy.DefaultKeyStrategy, lockProvider ?? RequestCollapserPolicy.GetDefaultLockProvider())
+                    : (ISyncPolicy<ResultClass>)Policy.NoOp<ResultClass>();
+            });
         }
 
         protected override Task ExecuteThroughPolicy(IsPolicy policy, Context context, int j, bool gated)

src/Polly.Contrib.DuplicateRequestCollapser.Specs/RequestCollapserTResultSpecsBase.cs

@@ -9,22 +9,24 @@ namespace Polly.Contrib.DuplicateRequestCollapser.Specs
 {
     public abstract class RequestCollapserTResultSpecsBase : RequestCollapserSpecsBase
     {
+        private Random _rng = new Random();
+
         protected RequestCollapserTResultSpecsBase(ITestOutputHelper testOutputHelper) : base(testOutputHelper) { }
 
-        private protected Func<ResultClass> ResultFactory => () => new ResultClass(ResultPrimitive.Good);
+        private protected Func<ResultClass> ResultFactory = () => new ResultClass(ResultPrimitive.Good);
 
         [Theory]
         [ClassData(typeof(RequestCollapserTestParallelisms))]
         public void Executing_concurrent_duplicate_task_through_CollapserPolicy_should_execute_only_once_and_return_same_single_result_instance(int parallelism)
         {
             (int actualInvocations, Task[] tasks) = Execute_parallel_delegates_through_policy_with_key_strategy(parallelism, useCollapser: true, sameKey: true);
-            
+
             actualInvocations.Should().Be(1);
 
             // All executions should have been handled the same single result instance.
-            ResultClass first = ((Task<ResultClass>) tasks[0]).Result;
+            ResultClass first = ((Task<ResultClass>)tasks[0]).Result;
 
-            tasks.All(t => Object.ReferenceEquals(((Task<ResultClass>) t).Result, first)).Should().BeTrue();
+            tasks.All(t => Object.ReferenceEquals(((Task<ResultClass>)t).Result, first)).Should().BeTrue();
         }
 
         [Theory]
@@ -39,8 +41,30 @@ public void Executing_concurrent_duplicate_task_not_through_CollapserPolicy_shou
             actualInvocations.Should().Be(parallelism);
 
             // All executions should have been handled the same single result instance.
-            tasks.Select(t => ((Task<ResultClass>) t).Result).Distinct().Count().Should().Be(parallelism);
+            tasks.Select(t => ((Task<ResultClass>)t).Result).Distinct().Count().Should().Be(parallelism);
         }
 
+        [Theory]
+        [ClassData(typeof(RequestCollapserTestParallelisms))]
+        public void Executing_concurrent_duplicate_faulted_task_through_CollapserPolicy_should_execute_only_once_and_return_same_single_result_instance(int parallelism)
+        {
+            ResultFactory = () => throw new Exception(_rng.Next().ToString());
+            (int actualInvocations, Task[] tasks) = Execute_parallel_delegates_through_policy_with_key_strategy(parallelism, useCollapser: true, sameKey: true);
+            actualInvocations.Should().Be(1);
+            tasks.First().IsFaulted.Should().BeTrue();
+
+            Exception first = tasks.First().Exception.InnerException;
+            // All executions should have been handed the same single result instance.
+            tasks.Select(x => x.Exception.InnerException).ShouldAllBeEquivalentTo(first);
+
+            (actualInvocations, tasks) = Execute_parallel_delegates_through_policy_with_key_strategy(parallelism, useCollapser: true, sameKey: true);
+            actualInvocations.Should().Be(2);
+            tasks.First().IsFaulted.Should().BeTrue();
+
+            Exception second = tasks.First().Exception.InnerException;
+            // The result of the second batch should not be the same as the first batch.
+            second.Message.Should().NotBe(first.Message);
+            tasks.Select(x => x.Exception.InnerException).ShouldAllBeEquivalentTo(second);
+        }
     }
 }

src/Polly.Contrib.DuplicateRequestCollapser/RequestCollapserEngine.cs

@@ -31,30 +31,33 @@ internal static TResult Implementation<TResult>(
                 lazy = collapser.GetOrAdd(key, new Lazy<object>(() => action(context, cancellationToken), LazyThreadSafetyMode.ExecutionAndPublication)); // Note: per documentation, LazyThreadSafetyMode.ExecutionAndPublication guarantees single execution, but means the executed code must not lock, as this risks deadlocks.  We should document.
             }
 
-            TResult result = (TResult)lazy.Value;
-
-            // As soon as the lazy has returned a result to one thread, the concurrent request set is over, so we evict the lazy from the ConcurrentDictionary.
-            // We need to evict within a lock, to be sure we are not, due to potential race with new threads populating, evicting a different lazy created by a different thread.
-            // To reduce lock contention, first check outside the lock whether we still need to remove it (we will double-check inside the lock).
-            if (collapser.TryGetValue(key, out Lazy<object> currentValue))
+            try
+            {
+                return (TResult)lazy.Value;
+            }
+            finally
             {
-                if (currentValue == lazy)
+                // As soon as the lazy has returned a result to one thread, the concurrent request set is over, so we evict the lazy from the ConcurrentDictionary.
+                // We need to evict within a lock, to be sure we are not, due to potential race with new threads populating, evicting a different lazy created by a different thread.
+                // To reduce lock contention, first check outside the lock whether we still need to remove it (we will double-check inside the lock).
+                if (collapser.TryGetValue(key, out Lazy<object> currentValue))
                 {
-                    using (lockProvider.AcquireLock(key, context, cancellationToken))
+                    if (currentValue == lazy)
                     {
-                        // Double-check that there has not been a race which updated the dictionary with a new value.
-                        if (collapser.TryGetValue(key, out Lazy<object> valueWithinLock))
+                        using (lockProvider.AcquireLock(key, context, cancellationToken))
                         {
-                            if (valueWithinLock == lazy)
+                            // Double-check that there has not been a race which updated the dictionary with a new value.
+                            if (collapser.TryGetValue(key, out Lazy<object> valueWithinLock))
                             {
-                                collapser.TryRemove(key, out _);
+                                if (valueWithinLock == lazy)
+                                {
+                                    collapser.TryRemove(key, out _);
+                                }
                             }
                         }
                     }
                 }
             }
-
-            return result;
         }
     }
 }

src/Polly.Contrib.DuplicateRequestCollapser/RequestCollapserEngineAsync.cs

@@ -33,31 +33,34 @@ internal static async Task<TResult> ImplementationAsync<TResult>(
                 lazy = collapser.GetOrAdd(key, new Lazy<Task<object>>(async () => await action(context, cancellationToken).ConfigureAwait(continueOnCapturedContext), LazyThreadSafetyMode.ExecutionAndPublication)); // Note: per documentation, LazyThreadSafetyMode.ExecutionAndPublication guarantees single execution, but means the executed code must not lock, as this risks deadlocks.  We should document.
             }
 
-            TResult result = (TResult)await lazy.Value.ConfigureAwait(continueOnCapturedContext);
-
-            // As soon as the lazy has returned a result to one thread, the concurrent request set is over, so we evict the lazy from the ConcurrentDictionary.
-            // We need to evict within a lock, to be sure we are not, due to potential race with new threads populating, evicting a different lazy created by a different thread.
-            // To reduce lock contention, first check outside the lock whether we still need to remove it (we will double-check inside the lock).
-            if (collapser.TryGetValue(key, out Lazy<Task<object>> currentValue))
+            try
+            {
+                return (TResult)await lazy.Value.ConfigureAwait(continueOnCapturedContext);
+            }
+            finally
             {
-                if (currentValue == lazy)
+                // As soon as the lazy has returned a result to one thread, the concurrent request set is over, so we evict the lazy from the ConcurrentDictionary.
+                // We need to evict within a lock, to be sure we are not, due to potential race with new threads populating, evicting a different lazy created by a different thread.
+                // To reduce lock contention, first check outside the lock whether we still need to remove it (we will double-check inside the lock).
+                if (collapser.TryGetValue(key, out Lazy<Task<object>> currentValue))
                 {
-                    await using (lockProvider.AcquireLockAsync(key, context, cancellationToken, continueOnCapturedContext)
-                        .ConfigureAwait(continueOnCapturedContext))
+                    if (currentValue == lazy)
                     {
-                        // Double-check that there has not been a race which updated the dictionary with a new value.
-                        if (collapser.TryGetValue(key, out Lazy<Task<object>> valueWithinLock))
+                        await using (lockProvider.AcquireLockAsync(key, context, cancellationToken, continueOnCapturedContext)
+                            .ConfigureAwait(continueOnCapturedContext))
                         {
-                            if (valueWithinLock == lazy)
+                            // Double-check that there has not been a race which updated the dictionary with a new value.
+                            if (collapser.TryGetValue(key, out Lazy<Task<object>> valueWithinLock))
                             {
-                                collapser.TryRemove(key, out _);
+                                if (valueWithinLock == lazy)
+                                {
+                                    collapser.TryRemove(key, out _);
+                                }
                             }
                         }
                     }
                 }
             }
-
-            return result;
         }
     }
 }