Simplify CacheTagHelper.ProcessAsync (#6506)

* Simplify CacheTagHelper.ProcessAsync

Modify ProcessAsync to cause all tasks to fail when the ongoing GetChildContentAsync call that they're
awaiting on fails.

src/Microsoft.AspNetCore.Mvc.TagHelpers/Cache/CacheTagKey.cs

@@ -33,7 +33,6 @@ namespace Microsoft.AspNetCore.Mvc.TagHelpers.Cache
         private const string VaryByCookieName = "VaryByCookie";
         private const string VaryByUserName = "VaryByUser";
 
-        private readonly string _key;
         private readonly string _prefix;
         private readonly string _varyBy;
         private readonly DateTimeOffset? _expiresOn;
@@ -58,7 +57,7 @@ namespace Microsoft.AspNetCore.Mvc.TagHelpers.Cache
         public CacheTagKey(CacheTagHelper tagHelper, TagHelperContext context)
             : this(tagHelper)
         {
-            _key = context.UniqueId;
+            Key = context.UniqueId;
             _prefix = nameof(CacheTagHelper);
         }
 
@@ -70,7 +69,7 @@ namespace Microsoft.AspNetCore.Mvc.TagHelpers.Cache
         public CacheTagKey(DistributedCacheTagHelper tagHelper)
             : this((CacheTagHelperBase)tagHelper)
         {
-            _key = tagHelper.Name;
+            Key = tagHelper.Name;
             _prefix = nameof(DistributedCacheTagHelper);
         }
 
@@ -95,6 +94,9 @@ namespace Microsoft.AspNetCore.Mvc.TagHelpers.Cache
             }
         }
 
+        // Internal for unit testing.
+        internal string Key { get; }
+
         /// <summary>
         /// Creates a <see cref="string"/> representation of the key.
         /// </summary>
@@ -110,7 +112,7 @@ namespace Microsoft.AspNetCore.Mvc.TagHelpers.Cache
             var builder = new StringBuilder(_prefix);
             builder
                 .Append(CacheKeyTokenSeparator)
-                .Append(_key);
+                .Append(Key);
 
             if (!string.IsNullOrEmpty(_varyBy))
             {
@@ -174,7 +176,7 @@ namespace Microsoft.AspNetCore.Mvc.TagHelpers.Cache
         /// <inheritdoc />
         public bool Equals(CacheTagKey other)
         {
-            return string.Equals(other._key, _key, StringComparison.Ordinal) &&
+            return string.Equals(other.Key, Key, StringComparison.Ordinal) &&
                 other._expiresAfter == _expiresAfter &&
                 other._expiresOn == _expiresOn &&
                 other._expiresSliding == _expiresSliding &&
@@ -203,7 +205,7 @@ namespace Microsoft.AspNetCore.Mvc.TagHelpers.Cache
 
             var hashCodeCombiner = new HashCodeCombiner();
 
-            hashCodeCombiner.Add(_key, StringComparer.Ordinal);
+            hashCodeCombiner.Add(Key, StringComparer.Ordinal);
             hashCodeCombiner.Add(_expiresAfter);
             hashCodeCombiner.Add(_expiresOn);
             hashCodeCombiner.Add(_expiresSliding);

src/Microsoft.AspNetCore.Mvc.TagHelpers/CacheTagHelper.cs

@@ -63,80 +63,18 @@ namespace Microsoft.AspNetCore.Mvc.TagHelpers
                 throw new ArgumentNullException(nameof(output));
             }
 
-            IHtmlContent content = null;
+            IHtmlContent content;
-
             if (Enabled)
             {
                 var cacheKey = new CacheTagKey(this, context);
-
+                if (MemoryCache.TryGetValue(cacheKey, out Task<IHtmlContent> cachedResult))
-                while (content == null)
                 {
-                    Task<IHtmlContent> result;
+                    // There is either some value already cached (as a Task) or a worker processing the output.
-
+                    content = await cachedResult;
-                    if (!MemoryCache.TryGetValue(cacheKey, out result))
+                }
-                    {
+                else
-                        var tokenSource = new CancellationTokenSource();
+                {
-
+                    content = await CreateCacheEntry(cacheKey, output);
-                        // Create an entry link scope and flow it so that any tokens related to the cache entries
-                        // created within this scope get copied to this scope.
-
-                        var options = GetMemoryCacheEntryOptions();
-                        options.AddExpirationToken(new CancellationChangeToken(tokenSource.Token));
-
-                        var tcs = new TaskCompletionSource<IHtmlContent>();
-
-                        // The returned value is ignored, we only do this so that
-                        // the compiler doesn't complain about the returned task
-                        // not being awaited
-                        var localTcs = MemoryCache.Set(cacheKey, tcs.Task, options);
-
-                        try
-                        {
-                            // The entry is set instead of assigning a value to the
-                            // task so that the expiration options are not impacted
-                            // by the time it took to compute it.
-
-                            using (var entry = MemoryCache.CreateEntry(cacheKey))
-                            {
-                                // The result is processed inside an entry
-                                // such that the tokens are inherited.
-
-                                result = ProcessContentAsync(output);
-
-                                entry.SetOptions(options);
-                                entry.Value = result;
-
-                                content = await result;
-                            }
-                            
-                            tcs.SetResult(content);
-                        }
-                        catch
-                        {
-                            // Remove the worker task from the cache in case it can't complete.
-                            tokenSource.Cancel();
-                            
-                            // If an exception occurs, ensure the other awaiters
-                            // render the output by themselves.
-                            tcs.SetResult(null);
-                            throw;
-                        }
-                        finally
-                        {
-                            // The tokenSource needs to be disposed as the MemoryCache
-                            // will register a callback on the Token.
-                            tokenSource.Dispose();
-                        }
-                    }
-                    else
-                    {
-                        // There is either some value already cached (as a Task)
-                        // or a worker processing the output. In the case of a worker,
-                        // the result will be null, and the request will try to acquire
-                        // the result from memory another time.
-
-                        content = await result;
-                    }
                 }
             }
             else
@@ -146,10 +84,65 @@ namespace Microsoft.AspNetCore.Mvc.TagHelpers
 
             // Clear the contents of the "cache" element since we don't want to render it.
             output.SuppressOutput();
-
             output.Content.SetHtmlContent(content);
         }
 
+        private async Task<IHtmlContent> CreateCacheEntry(CacheTagKey cacheKey, TagHelperOutput output)
+        {
+            var tokenSource = new CancellationTokenSource();
+
+            var options = GetMemoryCacheEntryOptions();
+            options.AddExpirationToken(new CancellationChangeToken(tokenSource.Token));
+
+            var tcs = new TaskCompletionSource<IHtmlContent>();
+
+            // The returned value is ignored, we only do this so that
+            // the compiler doesn't complain about the returned task
+            // not being awaited
+            var localTcs = MemoryCache.Set(cacheKey, tcs.Task, options);
+
+            try
+            {
+                // The entry is set instead of assigning a value to the
+                // task so that the expiration options are not impacted
+                // by the time it took to compute it.
+
+                // Use the CreateEntry to ensure a cache scope is created that will copy expiration tokens from
+                // cache entries created from the GetChildContentAsync call to the current entry.
+                IHtmlContent content;
+                using (var entry = MemoryCache.CreateEntry(cacheKey))
+                {
+                    // The result is processed inside an entry
+                    // such that the tokens are inherited.
+
+                    var result = ProcessContentAsync(output);
+
+                    entry.SetOptions(options);
+                    entry.Value = result;
+
+                    content = await result;
+                }
+
+                tcs.SetResult(content);
+                return content;
+            }
+            catch (Exception ex)
+            {
+                // Remove the worker task from the cache in case it can't complete.
+                tokenSource.Cancel();
+
+                // Fail the TCS so other awaiters see the exception.
+                tcs.SetException(ex);
+                throw;
+            }
+            finally
+            {
+                // The tokenSource needs to be disposed as the MemoryCache
+                // will register a callback on the Token.
+                tokenSource.Dispose();
+            }
+        }
+
         // Internal for unit testing
         internal MemoryCacheEntryOptions GetMemoryCacheEntryOptions()
         {

test/Microsoft.AspNetCore.Mvc.TagHelpers.Test/CacheTagHelperTest.cs

@@ -459,8 +459,7 @@ namespace Microsoft.AspNetCore.Mvc.TagHelpers
                 new TagHelperAttributeList { { "attr", "value" } },
                 getChildContentAsync: (useCachedResult, encoder) =>
                 {
-                    TagHelperContent tagHelperContent;
+                    if (!cache.TryGetValue("key1", out TagHelperContent tagHelperContent))
-                    if (!cache.TryGetValue("key1", out tagHelperContent))
                     {
                         tagHelperContent = expectedContent;
                         cache.Set("key1", tagHelperContent, cacheEntryOptions);
@@ -480,8 +479,7 @@ namespace Microsoft.AspNetCore.Mvc.TagHelpers
 
             // Act - 1
             await cacheTagHelper.ProcessAsync(tagHelperContext, tagHelperOutput);
-            Task<IHtmlContent> cachedValue;
+            var result = cache.TryGetValue(cacheTagKey, out Task<IHtmlContent> cachedValue);
-            var result = cache.TryGetValue(cacheTagKey, out cachedValue);
 
             // Assert - 1
             Assert.Equal("HtmlEncode[[some-content]]", tagHelperOutput.Content.GetContent());
@@ -804,6 +802,74 @@ namespace Microsoft.AspNetCore.Mvc.TagHelpers
             Assert.Equal(encoder.Encode(expected), tagHelperOutput1.Content.GetContent());
         }
 
+        [Fact]
+        public async Task ProcessAsync_ThrowsExceptionForAwaiters_IfExecutorEncountersAnException()
+        {
+            // Arrange
+            var expected = new DivideByZeroException();
+            var cache = new TestMemoryCache();
+            var cacheTagHelper = new CacheTagHelper(cache, new HtmlTestEncoder())
+            {
+                ViewContext = GetViewContext(),
+                Enabled = true
+            };
+
+            var invokeCount = 0;
+            var tagHelperOutput = new TagHelperOutput(
+                "cache",
+                new TagHelperAttributeList(),
+                getChildContentAsync: (useCachedResult, _) =>
+                {
+                    invokeCount++;
+                    throw expected;
+                });
+
+            // Act
+            var task1 = Task.Run(() => cacheTagHelper.ProcessAsync(GetTagHelperContext(cache.Key1), tagHelperOutput));
+            var task2 = Task.Run(() => cacheTagHelper.ProcessAsync(GetTagHelperContext(cache.Key2), tagHelperOutput));
+
+            // Assert
+            await Assert.ThrowsAsync<DivideByZeroException>(() => task1);
+            await Assert.ThrowsAsync<DivideByZeroException>(() => task2);
+            Assert.Equal(1, invokeCount);
+        }
+
+        [Fact]
+        public async Task ProcessAsync_AwaitersUseTheResultOfExecutor()
+        {
+            // Arrange
+            var expected = "Hello world";
+            var cache = new TestMemoryCache();
+            var encoder = new HtmlTestEncoder();
+            var cacheTagHelper = new CacheTagHelper(cache, encoder)
+            {
+                ViewContext = GetViewContext(),
+                Enabled = true
+            };
+
+            var invokeCount = 0;
+            var tagHelperOutput = new TagHelperOutput(
+                "cache",
+                new TagHelperAttributeList(),
+                getChildContentAsync: (useCachedResult, _) =>
+                {
+                    invokeCount++;
+
+                    var content = new DefaultTagHelperContent();
+                    content.SetContent(expected);
+                    return Task.FromResult<TagHelperContent>(content);
+                });
+
+            // Act
+            var task1 = Task.Run(() => cacheTagHelper.ProcessAsync(GetTagHelperContext(cache.Key1), tagHelperOutput));
+            var task2 = Task.Run(() => cacheTagHelper.ProcessAsync(GetTagHelperContext(cache.Key2), tagHelperOutput));
+
+            // Assert
+            await Task.WhenAll(task1, task2);
+            Assert.Equal(encoder.Encode(expected), tagHelperOutput.Content.GetContent());
+            Assert.Equal(1, invokeCount);
+        }
+
         private static ViewContext GetViewContext()
         {
             var actionContext = new ActionContext(new DefaultHttpContext(), new RouteData(), new ActionDescriptor());
@@ -842,5 +908,90 @@ namespace Microsoft.AspNetCore.Mvc.TagHelpers
                     return Task.FromResult<TagHelperContent>(tagHelperContent);
                 });
         }
+
+        private class TestCacheEntry : ICacheEntry
+        {
+            public object Key { get; set; }
+            public object Value { get; set; }
+            public DateTimeOffset? AbsoluteExpiration { get; set; }
+            public TimeSpan? AbsoluteExpirationRelativeToNow { get; set; }
+            public TimeSpan? SlidingExpiration { get; set; }
+
+            public IList<IChangeToken> ExpirationTokens { get; } = new List<IChangeToken>();
+
+            public IList<PostEvictionCallbackRegistration> PostEvictionCallbacks { get; } =
+                new List<PostEvictionCallbackRegistration>();
+
+            public CacheItemPriority Priority { get; set; }
+
+            public Action DisposeCallback { get; set; }
+
+            public void Dispose() => DisposeCallback();
+        }
+
+        // Simulates the scenario where a call to CacheTagHelper.ProcessAsync appears immediately after a prior one has assigned
+        // a TaskCancellationSource as an entry. We want to ensure that both calls use the results of the TCS as their output.
+        private class TestMemoryCache : IMemoryCache
+        {
+            private const int WaitTimeout = 3000;
+            public readonly string Key1 = "Key1";
+            public readonly string Key2 = "Key2";
+            public readonly ManualResetEventSlim ManualResetEvent1 = new ManualResetEventSlim();
+            public readonly ManualResetEventSlim ManualResetEvent2 = new ManualResetEventSlim();
+            public TestCacheEntry Entry { get; private set; }
+
+            public ICacheEntry CreateEntry(object key)
+            {
+                if (Entry != null)
+                {
+                    // We're being invoked in the inner "CreateEntry" call where the TCS is replaced by the GetChildContentAsync
+                    // Task. Wait for the other concurrent Task to grab the TCS before we proceed.
+                    Assert.True(ManualResetEvent1.Wait(WaitTimeout));
+                }
+
+                var cacheKey = Assert.IsType<CacheTagKey>(key);
+                Assert.Equal(Key1, cacheKey.Key);
+
+                Entry = new TestCacheEntry
+                {
+                    Key = key,
+                    DisposeCallback = ManualResetEvent2.Set,
+                };
+
+                return Entry;
+            }
+
+            public void Dispose()
+            {
+            }
+
+            public void Remove(object key)
+            {
+            }
+
+            public bool TryGetValue(object key, out object value)
+            {
+                var cacheKey = Assert.IsType<CacheTagKey>(key);
+                if (cacheKey.Key == Key2)
+                {
+                    Assert.True(ManualResetEvent2.Wait(WaitTimeout));
+
+                    Assert.NotNull(Entry);
+                    value = Entry.Value;
+                    Assert.NotNull(value);
+
+                    ManualResetEvent1.Set();
+
+                    return true;
+                }
+                else if (cacheKey.Key == Key1)
+                {
+                    value = null;
+                    return false;
+                }
+
+                throw new Exception();
+            }
+        }
     }
 }