From efb8ef33e6d718c5cf8fb554d43d9639efeec2ac Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?S=C3=A9bastien=20Ros?= <sebastienros@gmail.com>
Date: Wed, 6 Apr 2016 09:42:24 -0700
Subject: [PATCH] Fixing race in distributed cache tag helper

Fixes #4407
---
 .../Cache/DistributedCacheTagHelperService.cs | 20 +++++++++----------
 .../DistributedCacheTagHelper.cs              |  4 ++--
 .../DistributedCacheTagHelperTest.cs          |  2 +-
 3 files changed, 13 insertions(+), 13 deletions(-)

diff --git a/src/Microsoft.AspNetCore.Mvc.TagHelpers/Cache/DistributedCacheTagHelperService.cs b/src/Microsoft.AspNetCore.Mvc.TagHelpers/Cache/DistributedCacheTagHelperService.cs
index 94653073e9..dd3bb4f57c 100644
--- a/src/Microsoft.AspNetCore.Mvc.TagHelpers/Cache/DistributedCacheTagHelperService.cs
+++ b/src/Microsoft.AspNetCore.Mvc.TagHelpers/Cache/DistributedCacheTagHelperService.cs
@@ -95,7 +95,7 @@ namespace Microsoft.AspNetCore.Mvc.TagHelpers.Cache
                         var serializedKey = Encoding.UTF8.GetBytes(key.GenerateKey());
                         var storageKey = key.GenerateHashedKey();
                         var value = await _storage.GetAsync(storageKey);
-                                                
+
                         if (value == null)
                         {
                             // The value is not cached, we need to render the tag helper output
@@ -126,7 +126,7 @@ namespace Microsoft.AspNetCore.Mvc.TagHelpers.Cache
                             // The value was found in the storage, decode and ensure
                             // there is no cache key hash collision
                             byte[] decodedValue = Decode(value, serializedKey);
-                            
+
                             try
                             {
                                 if (decodedValue != null)
@@ -147,21 +147,21 @@ namespace Microsoft.AspNetCore.Mvc.TagHelpers.Cache
                                 }
                             }
                         }
-
-                        // Notify all other awaiters of the final content
-                        tcs.TrySetResult(content);
                     }
                     catch
                     {
-                        // Notify all other awaiters to render the content
-                        tcs.TrySetResult(null);
+                        content = null;
                         throw;
                     }
                     finally
                     {
-                        // Remove the worker task from the in-memory cache
-                        Task<IHtmlContent> worker;
-                        _workers.TryRemove(key, out worker);
+                        // Remove the worker task before setting the result.
+                        // If the result is null, other threads would potentially
+                        // acquire it otherwise.
+                        _workers.TryRemove(key, out result);
+
+                        // Notify all other awaiters to render the content
+                        tcs.TrySetResult(content);
                     }
                 }
                 else
diff --git a/src/Microsoft.AspNetCore.Mvc.TagHelpers/DistributedCacheTagHelper.cs b/src/Microsoft.AspNetCore.Mvc.TagHelpers/DistributedCacheTagHelper.cs
index 04edc01b1e..55d399258c 100644
--- a/src/Microsoft.AspNetCore.Mvc.TagHelpers/DistributedCacheTagHelper.cs
+++ b/src/Microsoft.AspNetCore.Mvc.TagHelpers/DistributedCacheTagHelper.cs
@@ -67,7 +67,7 @@ namespace Microsoft.AspNetCore.Mvc.TagHelpers
 
             IHtmlContent content = null;
 
-            // Create a cancellation token that will be used 
+            // Create a cancellation token that will be used
             // to release the task from the memory cache.
             var tokenSource = new CancellationTokenSource();
 
@@ -109,6 +109,6 @@ namespace Microsoft.AspNetCore.Mvc.TagHelpers
 
             return options;
         }
-        
+
     }
 }
\ No newline at end of file
diff --git a/test/Microsoft.AspNetCore.Mvc.TagHelpers.Test/DistributedCacheTagHelperTest.cs b/test/Microsoft.AspNetCore.Mvc.TagHelpers.Test/DistributedCacheTagHelperTest.cs
index 6d05e1373e..ad1a176678 100644
--- a/test/Microsoft.AspNetCore.Mvc.TagHelpers.Test/DistributedCacheTagHelperTest.cs
+++ b/test/Microsoft.AspNetCore.Mvc.TagHelpers.Test/DistributedCacheTagHelperTest.cs
@@ -28,7 +28,7 @@ namespace Microsoft.AspNetCore.Mvc.TagHelpers
 {
     public class DistributedCacheTagHelperTest
     {
-        
+
         [Fact]
         public async Task ProcessAsync_DoesNotCache_IfDisabled()
         {