From d9d7352098a9916483f53f1c408f562c0cb08df9 Mon Sep 17 00:00:00 2001
From: Sebastien Ros <sebastienros@gmail.com>
Date: Fri, 2 Sep 2022 18:09:36 -0700
Subject: [PATCH 01/11] Remove tagged keys as entries are evicted

---
 .../src/Memory/MemoryOutputCacheStore.cs      | 61 ++++++++++++++++---
 .../test/MemoryOutputCacheStoreTests.cs       | 36 +++++++++++
 .../test/OutputCacheMiddlewareTests.cs        |  6 +-
 3 files changed, 92 insertions(+), 11 deletions(-)

diff --git a/src/Middleware/OutputCaching/src/Memory/MemoryOutputCacheStore.cs b/src/Middleware/OutputCaching/src/Memory/MemoryOutputCacheStore.cs
index 3f1df73c7b80..238cb159bc74 100644
--- a/src/Middleware/OutputCaching/src/Memory/MemoryOutputCacheStore.cs
+++ b/src/Middleware/OutputCaching/src/Memory/MemoryOutputCacheStore.cs
@@ -1,6 +1,7 @@
 // Licensed to the .NET Foundation under one or more agreements.
 // The .NET Foundation licenses this file to you under the MIT license.
 
+using System.Linq;
 using Microsoft.Extensions.Caching.Memory;
 
 namespace Microsoft.AspNetCore.OutputCaching.Memory;
@@ -18,6 +19,9 @@ internal MemoryOutputCacheStore(IMemoryCache cache)
         _cache = cache;
     }
 
+    // For testing
+    internal Dictionary<string, HashSet<string>> TaggedEntries => _taggedEntries;
+
     public ValueTask EvictByTagAsync(string tag, CancellationToken cancellationToken)
     {
         ArgumentNullException.ThrowIfNull(tag);
@@ -26,12 +30,16 @@ public ValueTask EvictByTagAsync(string tag, CancellationToken cancellationToken
         {
             if (_taggedEntries.TryGetValue(tag, out var keys))
             {
-                foreach (var key in keys)
+                if (keys != null && keys.Count > 0)
                 {
-                    _cache.Remove(key);
-                }
+                    // Clone the tags as the collection might be updated by the post eviction callback
+                    var keysArray = keys.ToArray();
 
-                _taggedEntries.Remove(tag);
+                    foreach (var key in keysArray)
+                    {
+                        _cache.Remove(key);
+                    }
+                }
             }
         }
 
@@ -81,14 +89,49 @@ public ValueTask SetAsync(string key, byte[] value, string[]? tags, TimeSpan val
 
         void SetEntry()
         {
-            _cache.Set(
-            key,
-            value,
-            new MemoryCacheEntryOptions
+            var options = new MemoryCacheEntryOptions
             {
                 AbsoluteExpirationRelativeToNow = validFor,
                 Size = value.Length
-            });
+            };
+
+            if (tags != null && tags.Length > 0)
+            {
+                // Remove cache keys from tag lists when the entry is evicted
+                options.RegisterPostEvictionCallback(RemoveFromTags, tags);
+            }
+
+            _cache.Set(key, value, options);
+        }
+
+        void RemoveFromTags(object key, object? value, EvictionReason reason, object? state)
+        {
+            var tags = state as string[];
+
+            if (tags == null || tags.Length == 0)
+            {
+                return;
+            }
+
+            lock (_tagsLock)
+            {
+                foreach (var tag in tags)
+                {
+                    if (_taggedEntries.TryGetValue(tag, out var tagged) && tagged != null)
+                    {
+                        if (key is string s)
+                        {
+                            tagged.Remove(s);
+
+                            // Remove the collection if there is no more keys in it
+                            if (tagged.Count == 0)
+                            {
+                                _taggedEntries.Remove(tag);
+                            }
+                        }
+                    }
+                }
+            }
         }
 
         return ValueTask.CompletedTask;
diff --git a/src/Middleware/OutputCaching/test/MemoryOutputCacheStoreTests.cs b/src/Middleware/OutputCaching/test/MemoryOutputCacheStoreTests.cs
index e15f4515ca11..5c7a1952bd1d 100644
--- a/src/Middleware/OutputCaching/test/MemoryOutputCacheStoreTests.cs
+++ b/src/Middleware/OutputCaching/test/MemoryOutputCacheStoreTests.cs
@@ -70,8 +70,10 @@ public async Task EvictByTag_SingleTag_SingleEntry()
         await store.SetAsync(key, value, tags, TimeSpan.FromDays(1), default);
         await store.EvictByTagAsync("tag1", default);
         var result = await store.GetAsync(key, default);
+        var exists = store.TaggedEntries.TryGetValue("tag1", out _);
 
         Assert.Null(result);
+        Assert.False(exists);
     }
 
     [Fact]
@@ -140,6 +142,40 @@ public async Task EvictByTag_MultipleTags_MultipleEntries()
         Assert.Null(result2);
     }
 
+    [Fact]
+    public async Task ExpiredEntries_AreRemovedFromTags()
+    {
+        var testClock = new TestMemoryOptionsClock { UtcNow = DateTimeOffset.UtcNow };
+        var store = new MemoryOutputCacheStore(new MemoryCache(new MemoryCacheOptions { SizeLimit = 1000, Clock = testClock, ExpirationScanFrequency = TimeSpan.FromMilliseconds(1) }));
+        var value = "abc"u8.ToArray();
+
+        await store.SetAsync("a", value, new[] { "tag1" }, TimeSpan.FromMilliseconds(5), default);
+        await store.SetAsync("b", value, new[] { "tag1", "tag2" }, TimeSpan.FromMilliseconds(5), default);
+
+        testClock.Advance(TimeSpan.FromMilliseconds(10));
+
+        await store.SetAsync("c", value, new[] { "tag2" }, TimeSpan.FromMilliseconds(10), default);
+        await Task.Delay(10);
+
+        var resulta = await store.GetAsync("a", default);
+        var resultb = await store.GetAsync("b", default);
+        var resultc = await store.GetAsync("c", default);
+
+        var tag1s = store.TaggedEntries["tag1"];
+        var tag2s = store.TaggedEntries["tag2"];
+
+        // The hashset for tag1 should have been removed
+        var tag1Exists = store.TaggedEntries.TryGetValue("tag1", out _);
+
+        Assert.Null(resulta);
+        Assert.Null(resultb);
+        Assert.NotNull(resultc);
+
+        Assert.Empty(tag1s);
+        Assert.Single(tag2s);
+        Assert.False(tag1Exists);
+    }
+
     private class TestMemoryOptionsClock : Extensions.Internal.ISystemClock
     {
         public DateTimeOffset UtcNow { get; set; }
diff --git a/src/Middleware/OutputCaching/test/OutputCacheMiddlewareTests.cs b/src/Middleware/OutputCaching/test/OutputCacheMiddlewareTests.cs
index a4a030634caa..a228f003dc7b 100644
--- a/src/Middleware/OutputCaching/test/OutputCacheMiddlewareTests.cs
+++ b/src/Middleware/OutputCaching/test/OutputCacheMiddlewareTests.cs
@@ -1,6 +1,7 @@
 // Licensed to the .NET Foundation under one or more agreements.
 // The .NET Foundation licenses this file to you under the MIT license.
 
+using System.Globalization;
 using Microsoft.AspNetCore.Http;
 using Microsoft.AspNetCore.Http.Features;
 using Microsoft.AspNetCore.OutputCaching.Memory;
@@ -742,13 +743,14 @@ public async Task FinalizeCacheBody_DoNotCache_IfSizeTooBig()
             })));
         var context = TestUtils.CreateTestContext();
         middleware.TryGetRequestPolicies(context.HttpContext, out var policies);
-        middleware.ShimResponseStream(context);
 
         await context.HttpContext.Response.WriteAsync(new string('0', 101));
 
         context.CachedResponse = new OutputCacheEntry() { Headers = new HeaderDictionary() };
         context.CacheKey = "BaseKey";
-        context.CachedResponseValidFor = TimeSpan.FromSeconds(10);
+        context.CachedResponseValidFor = TimeSpan.FromSeconds(100);
+
+        middleware.ShimResponseStream(context);
 
         await middleware.FinalizeCacheBodyAsync(context);
 

From e4590208e01b0e01f99be8adf0c13273cebb62ec Mon Sep 17 00:00:00 2001
From: Sebastien Ros <sebastienros@gmail.com>
Date: Fri, 2 Sep 2022 18:17:19 -0700
Subject: [PATCH 02/11] Revert middleware tests changes

---
 .../OutputCaching/test/OutputCacheMiddlewareTests.cs        | 6 ++----
 1 file changed, 2 insertions(+), 4 deletions(-)

diff --git a/src/Middleware/OutputCaching/test/OutputCacheMiddlewareTests.cs b/src/Middleware/OutputCaching/test/OutputCacheMiddlewareTests.cs
index a228f003dc7b..a4a030634caa 100644
--- a/src/Middleware/OutputCaching/test/OutputCacheMiddlewareTests.cs
+++ b/src/Middleware/OutputCaching/test/OutputCacheMiddlewareTests.cs
@@ -1,7 +1,6 @@
 // Licensed to the .NET Foundation under one or more agreements.
 // The .NET Foundation licenses this file to you under the MIT license.
 
-using System.Globalization;
 using Microsoft.AspNetCore.Http;
 using Microsoft.AspNetCore.Http.Features;
 using Microsoft.AspNetCore.OutputCaching.Memory;
@@ -743,14 +742,13 @@ public async Task FinalizeCacheBody_DoNotCache_IfSizeTooBig()
             })));
         var context = TestUtils.CreateTestContext();
         middleware.TryGetRequestPolicies(context.HttpContext, out var policies);
+        middleware.ShimResponseStream(context);
 
         await context.HttpContext.Response.WriteAsync(new string('0', 101));
 
         context.CachedResponse = new OutputCacheEntry() { Headers = new HeaderDictionary() };
         context.CacheKey = "BaseKey";
-        context.CachedResponseValidFor = TimeSpan.FromSeconds(100);
-
-        middleware.ShimResponseStream(context);
+        context.CachedResponseValidFor = TimeSpan.FromSeconds(10);
 
         await middleware.FinalizeCacheBodyAsync(context);
 

From ae50e8bbabc130d8cca02eff2e321e9736177a90 Mon Sep 17 00:00:00 2001
From: Sebastien Ros <sebastienros@gmail.com>
Date: Tue, 6 Sep 2022 10:23:12 -0700
Subject: [PATCH 03/11] Fix unit test

---
 .../src/Memory/MemoryOutputCacheStore.cs      | 10 +++++-----
 .../test/MemoryOutputCacheStoreTests.cs       | 20 ++++++++++---------
 2 files changed, 16 insertions(+), 14 deletions(-)

diff --git a/src/Middleware/OutputCaching/src/Memory/MemoryOutputCacheStore.cs b/src/Middleware/OutputCaching/src/Memory/MemoryOutputCacheStore.cs
index 238cb159bc74..f45db8eb55ea 100644
--- a/src/Middleware/OutputCaching/src/Memory/MemoryOutputCacheStore.cs
+++ b/src/Middleware/OutputCaching/src/Memory/MemoryOutputCacheStore.cs
@@ -1,6 +1,7 @@
 // Licensed to the .NET Foundation under one or more agreements.
 // The .NET Foundation licenses this file to you under the MIT license.
 
+using System.Diagnostics;
 using System.Linq;
 using Microsoft.Extensions.Caching.Memory;
 
@@ -108,14 +109,12 @@ void RemoveFromTags(object key, object? value, EvictionReason reason, object? st
         {
             var tags = state as string[];
 
-            if (tags == null || tags.Length == 0)
-            {
-                return;
-            }
+            Debug.Assert(tags != null);
+            Debug.Assert(tags.Length > 0);
 
             lock (_tagsLock)
             {
-                foreach (var tag in tags)
+                foreach (var tag in tags!)
                 {
                     if (_taggedEntries.TryGetValue(tag, out var tagged) && tagged != null)
                     {
@@ -127,6 +126,7 @@ void RemoveFromTags(object key, object? value, EvictionReason reason, object? st
                             if (tagged.Count == 0)
                             {
                                 _taggedEntries.Remove(tag);
+                                break;
                             }
                         }
                     }
diff --git a/src/Middleware/OutputCaching/test/MemoryOutputCacheStoreTests.cs b/src/Middleware/OutputCaching/test/MemoryOutputCacheStoreTests.cs
index 5c7a1952bd1d..812ec5acb3ac 100644
--- a/src/Middleware/OutputCaching/test/MemoryOutputCacheStoreTests.cs
+++ b/src/Middleware/OutputCaching/test/MemoryOutputCacheStoreTests.cs
@@ -146,7 +146,8 @@ public async Task EvictByTag_MultipleTags_MultipleEntries()
     public async Task ExpiredEntries_AreRemovedFromTags()
     {
         var testClock = new TestMemoryOptionsClock { UtcNow = DateTimeOffset.UtcNow };
-        var store = new MemoryOutputCacheStore(new MemoryCache(new MemoryCacheOptions { SizeLimit = 1000, Clock = testClock, ExpirationScanFrequency = TimeSpan.FromMilliseconds(1) }));
+        var cache = new MemoryCache(new MemoryCacheOptions { SizeLimit = 1000, Clock = testClock, ExpirationScanFrequency = TimeSpan.FromMilliseconds(1) });
+        var store = new MemoryOutputCacheStore(cache);
         var value = "abc"u8.ToArray();
 
         await store.SetAsync("a", value, new[] { "tag1" }, TimeSpan.FromMilliseconds(5), default);
@@ -154,26 +155,27 @@ public async Task ExpiredEntries_AreRemovedFromTags()
 
         testClock.Advance(TimeSpan.FromMilliseconds(10));
 
+        // Force eviction of expired entries
+        cache.Compact(0);
+
         await store.SetAsync("c", value, new[] { "tag2" }, TimeSpan.FromMilliseconds(10), default);
-        await Task.Delay(10);
 
         var resulta = await store.GetAsync("a", default);
         var resultb = await store.GetAsync("b", default);
         var resultc = await store.GetAsync("c", default);
 
-        var tag1s = store.TaggedEntries["tag1"];
-        var tag2s = store.TaggedEntries["tag2"];
-
-        // The hashset for tag1 should have been removed
-        var tag1Exists = store.TaggedEntries.TryGetValue("tag1", out _);
+        store.TaggedEntries.TryGetValue("tag2", out var tag2s);
 
         Assert.Null(resulta);
         Assert.Null(resultb);
         Assert.NotNull(resultc);
 
-        Assert.Empty(tag1s);
+        // Wait for the hashset to be removed as it happens on a separate thread
+        var timeout = Task.Delay(1000);
+        while (store.TaggedEntries.TryGetValue("tag1", out _) && !timeout.IsCompleted) ;
+
+        Assert.False(timeout.IsCompleted);
         Assert.Single(tag2s);
-        Assert.False(tag1Exists);
     }
 
     private class TestMemoryOptionsClock : Extensions.Internal.ISystemClock

From 0fb6f644f399b3e58ef6259ebaf6572f11c5032c Mon Sep 17 00:00:00 2001
From: Sebastien Ros <sebastienros@gmail.com>
Date: Tue, 6 Sep 2022 11:17:20 -0700
Subject: [PATCH 04/11] Add braces

---
 .../OutputCaching/test/MemoryOutputCacheStoreTests.cs           | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/Middleware/OutputCaching/test/MemoryOutputCacheStoreTests.cs b/src/Middleware/OutputCaching/test/MemoryOutputCacheStoreTests.cs
index 812ec5acb3ac..dd1203e09fcf 100644
--- a/src/Middleware/OutputCaching/test/MemoryOutputCacheStoreTests.cs
+++ b/src/Middleware/OutputCaching/test/MemoryOutputCacheStoreTests.cs
@@ -172,7 +172,7 @@ public async Task ExpiredEntries_AreRemovedFromTags()
 
         // Wait for the hashset to be removed as it happens on a separate thread
         var timeout = Task.Delay(1000);
-        while (store.TaggedEntries.TryGetValue("tag1", out _) && !timeout.IsCompleted) ;
+        while (store.TaggedEntries.TryGetValue("tag1", out _) && !timeout.IsCompleted) { }
 
         Assert.False(timeout.IsCompleted);
         Assert.Single(tag2s);

From d9880c2dfc5ade8b1ddab734e9cb96ca301726d3 Mon Sep 17 00:00:00 2001
From: Sebastien Ros <sebastienros@gmail.com>
Date: Tue, 6 Sep 2022 15:47:00 -0700
Subject: [PATCH 05/11] Fix flaky tests

---
 .../test/MemoryOutputCacheStoreTests.cs       | 22 +++++++++++++------
 1 file changed, 15 insertions(+), 7 deletions(-)

diff --git a/src/Middleware/OutputCaching/test/MemoryOutputCacheStoreTests.cs b/src/Middleware/OutputCaching/test/MemoryOutputCacheStoreTests.cs
index dd1203e09fcf..c5ec300b6d09 100644
--- a/src/Middleware/OutputCaching/test/MemoryOutputCacheStoreTests.cs
+++ b/src/Middleware/OutputCaching/test/MemoryOutputCacheStoreTests.cs
@@ -70,10 +70,16 @@ public async Task EvictByTag_SingleTag_SingleEntry()
         await store.SetAsync(key, value, tags, TimeSpan.FromDays(1), default);
         await store.EvictByTagAsync("tag1", default);
         var result = await store.GetAsync(key, default);
-        var exists = store.TaggedEntries.TryGetValue("tag1", out _);
+
+        HashSet<string> tag1s;
+
+        // Wait for the hashset to be removed as it happens on a separate thread
+        var timeout = Task.Delay(1000);
+        while (store.TaggedEntries.TryGetValue("tag1", out tag1s) && !timeout.IsCompleted) { }
 
         Assert.Null(result);
-        Assert.False(exists);
+        Assert.Null(tag1s);
+        Assert.False(timeout.IsCompleted);
     }
 
     [Fact]
@@ -164,17 +170,19 @@ public async Task ExpiredEntries_AreRemovedFromTags()
         var resultb = await store.GetAsync("b", default);
         var resultc = await store.GetAsync("c", default);
 
-        store.TaggedEntries.TryGetValue("tag2", out var tag2s);
+        HashSet<string> tag1s, tag2s;
+
+        // Wait for the hashset to be removed as it happens on a separate thread
+        var timeout = Task.Delay(1000);
+        while (store.TaggedEntries.TryGetValue("tag1", out tag1s) && !timeout.IsCompleted) { }
+        while (store.TaggedEntries.TryGetValue("tag2", out tag2s) && tag2s.Count == 2 && !timeout.IsCompleted) { }
 
         Assert.Null(resulta);
         Assert.Null(resultb);
         Assert.NotNull(resultc);
 
-        // Wait for the hashset to be removed as it happens on a separate thread
-        var timeout = Task.Delay(1000);
-        while (store.TaggedEntries.TryGetValue("tag1", out _) && !timeout.IsCompleted) { }
-
         Assert.False(timeout.IsCompleted);
+        Assert.Null(tag1s);
         Assert.Single(tag2s);
     }
 

From 309a8de501c711fd1f9d3c9d00c6ccb54251b216 Mon Sep 17 00:00:00 2001
From: Sebastien Ros <sebastienros@gmail.com>
Date: Wed, 7 Sep 2022 10:56:50 -0700
Subject: [PATCH 06/11] Fix tests

---
 .../src/Memory/MemoryOutputCacheStore.cs      |  1 -
 .../test/MemoryOutputCacheStoreTests.cs       | 19 +++++++++----------
 2 files changed, 9 insertions(+), 11 deletions(-)

diff --git a/src/Middleware/OutputCaching/src/Memory/MemoryOutputCacheStore.cs b/src/Middleware/OutputCaching/src/Memory/MemoryOutputCacheStore.cs
index f45db8eb55ea..e2ab661323dc 100644
--- a/src/Middleware/OutputCaching/src/Memory/MemoryOutputCacheStore.cs
+++ b/src/Middleware/OutputCaching/src/Memory/MemoryOutputCacheStore.cs
@@ -126,7 +126,6 @@ void RemoveFromTags(object key, object? value, EvictionReason reason, object? st
                             if (tagged.Count == 0)
                             {
                                 _taggedEntries.Remove(tag);
-                                break;
                             }
                         }
                     }
diff --git a/src/Middleware/OutputCaching/test/MemoryOutputCacheStoreTests.cs b/src/Middleware/OutputCaching/test/MemoryOutputCacheStoreTests.cs
index c5ec300b6d09..5bee23df2cdd 100644
--- a/src/Middleware/OutputCaching/test/MemoryOutputCacheStoreTests.cs
+++ b/src/Middleware/OutputCaching/test/MemoryOutputCacheStoreTests.cs
@@ -158,28 +158,27 @@ public async Task ExpiredEntries_AreRemovedFromTags()
 
         await store.SetAsync("a", value, new[] { "tag1" }, TimeSpan.FromMilliseconds(5), default);
         await store.SetAsync("b", value, new[] { "tag1", "tag2" }, TimeSpan.FromMilliseconds(5), default);
+        await store.SetAsync("c", value, new[] { "tag2" }, TimeSpan.FromMilliseconds(20), default);
 
         testClock.Advance(TimeSpan.FromMilliseconds(10));
 
-        // Force eviction of expired entries
-        cache.Compact(0);
-
-        await store.SetAsync("c", value, new[] { "tag2" }, TimeSpan.FromMilliseconds(10), default);
+        // Background expiration checks are triggered by misc cache activity.
+        _ = cache.Get("a");
 
         var resulta = await store.GetAsync("a", default);
         var resultb = await store.GetAsync("b", default);
         var resultc = await store.GetAsync("c", default);
 
+        Assert.Null(resulta);
+        Assert.Null(resultb);
+        Assert.NotNull(resultc);
+
         HashSet<string> tag1s, tag2s;
 
         // Wait for the hashset to be removed as it happens on a separate thread
-        var timeout = Task.Delay(1000);
+        var timeout = Task.Delay(2000);
         while (store.TaggedEntries.TryGetValue("tag1", out tag1s) && !timeout.IsCompleted) { }
-        while (store.TaggedEntries.TryGetValue("tag2", out tag2s) && tag2s.Count == 2 && !timeout.IsCompleted) { }
-
-        Assert.Null(resulta);
-        Assert.Null(resultb);
-        Assert.NotNull(resultc);
+        while (store.TaggedEntries.TryGetValue("tag2", out tag2s) && tag2s.Count != 1 && !timeout.IsCompleted) { }
 
         Assert.False(timeout.IsCompleted);
         Assert.Null(tag1s);

From 1f019c1d9b64bc4a1c7a3643f300efbfc9d95943 Mon Sep 17 00:00:00 2001
From: Sebastien Ros <sebastienros@gmail.com>
Date: Thu, 8 Sep 2022 15:14:44 -0700
Subject: [PATCH 07/11] Feedback

---
 .../src/Memory/MemoryOutputCacheStore.cs      | 93 +++++++++++--------
 .../test/MemoryOutputCacheStoreTests.cs       | 26 ++++--
 2 files changed, 73 insertions(+), 46 deletions(-)

diff --git a/src/Middleware/OutputCaching/src/Memory/MemoryOutputCacheStore.cs b/src/Middleware/OutputCaching/src/Memory/MemoryOutputCacheStore.cs
index e2ab661323dc..116ed7e9e853 100644
--- a/src/Middleware/OutputCaching/src/Memory/MemoryOutputCacheStore.cs
+++ b/src/Middleware/OutputCaching/src/Memory/MemoryOutputCacheStore.cs
@@ -33,12 +33,25 @@ public ValueTask EvictByTagAsync(string tag, CancellationToken cancellationToken
             {
                 if (keys != null && keys.Count > 0)
                 {
-                    // Clone the tags as the collection might be updated by the post eviction callback
-                    var keysArray = keys.ToArray();
+                    // Since eviction callbacks run inline, iterating over keys could throw
+                    // To prevent allocating a copy of the keys we check if the eviction callback ran,
+                    // and if it did we restart the loop.
 
-                    foreach (var key in keysArray)
+                    var i = keys.Count;
+                    while (i > 0)
                     {
-                        _cache.Remove(key);
+                        var oldCount = keys.Count;
+                        foreach (var key in keys)
+                        {
+                            _cache.Remove(key);
+                            i--;
+                            if (oldCount != keys.Count)
+                            {
+                                // eviction callback ran inline, we need to restart the loop to avoid
+                                // "collection modified while iterating" errors
+                                break;
+                            }
+                        }
                     }
                 }
             }
@@ -80,59 +93,61 @@ public ValueTask SetAsync(string key, byte[] value, string[]? tags, TimeSpan val
                     keys.Add(key);
                 }
 
-                SetEntry();
+                SetEntry(key, value, tags, validFor);
             }
         }
         else
         {
-            SetEntry();
+            SetEntry(key, value, tags, validFor);
         }
 
-        void SetEntry()
-        {
-            var options = new MemoryCacheEntryOptions
-            {
-                AbsoluteExpirationRelativeToNow = validFor,
-                Size = value.Length
-            };
+        return ValueTask.CompletedTask;
+    }
 
-            if (tags != null && tags.Length > 0)
-            {
-                // Remove cache keys from tag lists when the entry is evicted
-                options.RegisterPostEvictionCallback(RemoveFromTags, tags);
-            }
+    void SetEntry(string key, byte[] value, string[]? tags, TimeSpan validFor)
+    {
+        var options = new MemoryCacheEntryOptions
+        {
+            AbsoluteExpirationRelativeToNow = validFor,
+            Size = value.Length
+        };
 
-            _cache.Set(key, value, options);
+        if (tags != null && tags.Length > 0)
+        {
+            // Remove cache keys from tag lists when the entry is evicted
+            options.RegisterPostEvictionCallback(RemoveFromTags, tags);
         }
 
-        void RemoveFromTags(object key, object? value, EvictionReason reason, object? state)
-        {
-            var tags = state as string[];
+        _cache.Set(key, value, options);
+    }
 
-            Debug.Assert(tags != null);
-            Debug.Assert(tags.Length > 0);
+    void RemoveFromTags(object key, object? value, EvictionReason reason, object? state)
+    {
+        var tags = state as string[];
 
-            lock (_tagsLock)
+        Debug.Assert(tags != null);
+        Debug.Assert(tags.Length > 0);
+
+        if (key is not string stringKey)
+        {
+            return;
+        }
+
+        lock (_tagsLock)
+        {
+            foreach (var tag in tags!)
             {
-                foreach (var tag in tags!)
+                if (_taggedEntries.TryGetValue(tag, out var tagged) && tagged != null)
                 {
-                    if (_taggedEntries.TryGetValue(tag, out var tagged) && tagged != null)
-                    {
-                        if (key is string s)
-                        {
-                            tagged.Remove(s);
+                    tagged.Remove(stringKey);
 
-                            // Remove the collection if there is no more keys in it
-                            if (tagged.Count == 0)
-                            {
-                                _taggedEntries.Remove(tag);
-                            }
-                        }
+                    // Remove the collection if there is no more keys in it
+                    if (tagged.Count == 0)
+                    {
+                        _taggedEntries.Remove(tag);
                     }
                 }
             }
         }
-
-        return ValueTask.CompletedTask;
     }
 }
diff --git a/src/Middleware/OutputCaching/test/MemoryOutputCacheStoreTests.cs b/src/Middleware/OutputCaching/test/MemoryOutputCacheStoreTests.cs
index 5bee23df2cdd..f3703fc0cd3f 100644
--- a/src/Middleware/OutputCaching/test/MemoryOutputCacheStoreTests.cs
+++ b/src/Middleware/OutputCaching/test/MemoryOutputCacheStoreTests.cs
@@ -74,12 +74,16 @@ public async Task EvictByTag_SingleTag_SingleEntry()
         HashSet<string> tag1s;
 
         // Wait for the hashset to be removed as it happens on a separate thread
-        var timeout = Task.Delay(1000);
-        while (store.TaggedEntries.TryGetValue("tag1", out tag1s) && !timeout.IsCompleted) { }
+
+        using var cts = new CancellationTokenSource(TimeSpan.FromSeconds(30));
+
+        while (store.TaggedEntries.TryGetValue("tag1", out tag1s) && !cts.IsCancellationRequested)
+        {
+            await Task.Yield();
+        }
 
         Assert.Null(result);
         Assert.Null(tag1s);
-        Assert.False(timeout.IsCompleted);
     }
 
     [Fact]
@@ -176,11 +180,19 @@ public async Task ExpiredEntries_AreRemovedFromTags()
         HashSet<string> tag1s, tag2s;
 
         // Wait for the hashset to be removed as it happens on a separate thread
-        var timeout = Task.Delay(2000);
-        while (store.TaggedEntries.TryGetValue("tag1", out tag1s) && !timeout.IsCompleted) { }
-        while (store.TaggedEntries.TryGetValue("tag2", out tag2s) && tag2s.Count != 1 && !timeout.IsCompleted) { }
 
-        Assert.False(timeout.IsCompleted);
+        using var cts = new CancellationTokenSource(TimeSpan.FromSeconds(30));
+
+        while (store.TaggedEntries.TryGetValue("tag1", out tag1s) && !cts.IsCancellationRequested)
+        {
+            await Task.Yield();
+        }
+
+        while (store.TaggedEntries.TryGetValue("tag2", out tag2s) && tag2s.Count != 1 && !cts.IsCancellationRequested)
+        {
+            await Task.Yield();
+        }
+
         Assert.Null(tag1s);
         Assert.Single(tag2s);
     }

From 7eb89619e064872b9cba12b12f9fc86eab72f19e Mon Sep 17 00:00:00 2001
From: Sebastien Ros <sebastienros@gmail.com>
Date: Thu, 8 Sep 2022 17:42:02 -0700
Subject: [PATCH 08/11] Feedback

---
 .../src/Memory/MemoryOutputCacheStore.cs      | 30 +++++++++++--------
 .../OutputCaching/src/Resources.resx          |  4 +--
 .../test/MemoryOutputCacheStoreTests.cs       | 12 ++++++++
 3 files changed, 32 insertions(+), 14 deletions(-)

diff --git a/src/Middleware/OutputCaching/src/Memory/MemoryOutputCacheStore.cs b/src/Middleware/OutputCaching/src/Memory/MemoryOutputCacheStore.cs
index 116ed7e9e853..a7c3b1f14136 100644
--- a/src/Middleware/OutputCaching/src/Memory/MemoryOutputCacheStore.cs
+++ b/src/Middleware/OutputCaching/src/Memory/MemoryOutputCacheStore.cs
@@ -2,18 +2,17 @@
 // The .NET Foundation licenses this file to you under the MIT license.
 
 using System.Diagnostics;
-using System.Linq;
 using Microsoft.Extensions.Caching.Memory;
 
 namespace Microsoft.AspNetCore.OutputCaching.Memory;
 
 internal sealed class MemoryOutputCacheStore : IOutputCacheStore
 {
-    private readonly IMemoryCache _cache;
+    private readonly MemoryCache _cache;
     private readonly Dictionary<string, HashSet<string>> _taggedEntries = new();
     private readonly object _tagsLock = new();
 
-    internal MemoryOutputCacheStore(IMemoryCache cache)
+    internal MemoryOutputCacheStore(MemoryCache cache)
     {
         ArgumentNullException.ThrowIfNull(cache);
 
@@ -33,7 +32,7 @@ public ValueTask EvictByTagAsync(string tag, CancellationToken cancellationToken
             {
                 if (keys != null && keys.Count > 0)
                 {
-                    // Since eviction callbacks run inline, iterating over keys could throw
+                    // If MemoryCache changed to run eviction callbacks run inline in Remove, iterating over keys could throw
                     // To prevent allocating a copy of the keys we check if the eviction callback ran,
                     // and if it did we restart the loop.
 
@@ -84,12 +83,21 @@ public ValueTask SetAsync(string key, byte[] value, string[]? tags, TimeSpan val
             {
                 foreach (var tag in tags)
                 {
+                    if (String.IsNullOrEmpty(tag))
+                    {
+                        throw new ArgumentException(Resources.TagCannotBeNullOrEmpty);
+                    }
+
+                    ArgumentNullException.ThrowIfNull(tag);
+
                     if (!_taggedEntries.TryGetValue(tag, out var keys))
                     {
                         keys = new HashSet<string>();
                         _taggedEntries[tag] = keys;
                     }
 
+                    Debug.Assert(keys != null);
+
                     keys.Add(key);
                 }
 
@@ -106,6 +114,8 @@ public ValueTask SetAsync(string key, byte[] value, string[]? tags, TimeSpan val
 
     void SetEntry(string key, byte[] value, string[]? tags, TimeSpan validFor)
     {
+        Debug.Assert(key != null);
+
         var options = new MemoryCacheEntryOptions
         {
             AbsoluteExpirationRelativeToNow = validFor,
@@ -127,19 +137,15 @@ void RemoveFromTags(object key, object? value, EvictionReason reason, object? st
 
         Debug.Assert(tags != null);
         Debug.Assert(tags.Length > 0);
-
-        if (key is not string stringKey)
-        {
-            return;
-        }
+        Debug.Assert(key is string);
 
         lock (_tagsLock)
         {
-            foreach (var tag in tags!)
+            foreach (var tag in tags)
             {
-                if (_taggedEntries.TryGetValue(tag, out var tagged) && tagged != null)
+                if (_taggedEntries.TryGetValue(tag, out var tagged))
                 {
-                    tagged.Remove(stringKey);
+                    tagged.Remove((string)key);
 
                     // Remove the collection if there is no more keys in it
                     if (tagged.Count == 0)
diff --git a/src/Middleware/OutputCaching/src/Resources.resx b/src/Middleware/OutputCaching/src/Resources.resx
index 3a19868a7302..122136d3eba6 100644
--- a/src/Middleware/OutputCaching/src/Resources.resx
+++ b/src/Middleware/OutputCaching/src/Resources.resx
@@ -117,7 +117,7 @@
   <resheader name="writer">
     <value>System.Resources.ResXResourceWriter, System.Windows.Forms, Version=4.0.0.0, Culture=neutral, PublicKeyToken=b77a5c561934e089</value>
   </resheader>
-  <data name="Policy_InvalidType" xml:space="preserve">
-    <value>The type '{0}' is not a valid output policy.</value>
+  <data name="TagCannotBeNullOrEmpty" xml:space="preserve">
+    <value>Tag value cannot be null or empty.</value>
   </data>
 </root>
\ No newline at end of file
diff --git a/src/Middleware/OutputCaching/test/MemoryOutputCacheStoreTests.cs b/src/Middleware/OutputCaching/test/MemoryOutputCacheStoreTests.cs
index f3703fc0cd3f..c790370ce9c9 100644
--- a/src/Middleware/OutputCaching/test/MemoryOutputCacheStoreTests.cs
+++ b/src/Middleware/OutputCaching/test/MemoryOutputCacheStoreTests.cs
@@ -197,6 +197,18 @@ public async Task ExpiredEntries_AreRemovedFromTags()
         Assert.Single(tag2s);
     }
 
+    [Theory]
+    [InlineData(null)]
+    [InlineData("")]
+    public async Task Store_Throws_OnInvalidTag(string tag)
+    {
+        var store = new MemoryOutputCacheStore(new MemoryCache(new MemoryCacheOptions()));
+        var value = "abc"u8.ToArray();
+        var key = "abc";
+
+        await Assert.ThrowsAsync<ArgumentException>(async () => await store.SetAsync(key, value, new string[] { tag }, TimeSpan.FromMinutes(1), default));
+    }
+
     private class TestMemoryOptionsClock : Extensions.Internal.ISystemClock
     {
         public DateTimeOffset UtcNow { get; set; }

From 9618e5237ee25bcb0c298b579b6302810e2c328e Mon Sep 17 00:00:00 2001
From: Sebastien Ros <sebastienros@gmail.com>
Date: Thu, 8 Sep 2022 17:43:55 -0700
Subject: [PATCH 09/11] Comment

---
 .../OutputCaching/src/Memory/MemoryOutputCacheStore.cs          | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/Middleware/OutputCaching/src/Memory/MemoryOutputCacheStore.cs b/src/Middleware/OutputCaching/src/Memory/MemoryOutputCacheStore.cs
index a7c3b1f14136..4dfab42c60a2 100644
--- a/src/Middleware/OutputCaching/src/Memory/MemoryOutputCacheStore.cs
+++ b/src/Middleware/OutputCaching/src/Memory/MemoryOutputCacheStore.cs
@@ -32,7 +32,7 @@ public ValueTask EvictByTagAsync(string tag, CancellationToken cancellationToken
             {
                 if (keys != null && keys.Count > 0)
                 {
-                    // If MemoryCache changed to run eviction callbacks run inline in Remove, iterating over keys could throw
+                    // If MemoryCache changed to run eviction callbacks inline in Remove, iterating over keys could throw
                     // To prevent allocating a copy of the keys we check if the eviction callback ran,
                     // and if it did we restart the loop.
 

From 05df3d9e22f633b61876565bde4db6efd98f0204 Mon Sep 17 00:00:00 2001
From: Sebastien Ros <sebastienros@gmail.com>
Date: Fri, 9 Sep 2022 09:55:23 -0700
Subject: [PATCH 10/11] Remove left over exception

---
 .../OutputCaching/src/Memory/MemoryOutputCacheStore.cs      | 6 ++----
 src/Middleware/OutputCaching/src/Resources.resx             | 4 ++--
 2 files changed, 4 insertions(+), 6 deletions(-)

diff --git a/src/Middleware/OutputCaching/src/Memory/MemoryOutputCacheStore.cs b/src/Middleware/OutputCaching/src/Memory/MemoryOutputCacheStore.cs
index 4dfab42c60a2..a75546b6793f 100644
--- a/src/Middleware/OutputCaching/src/Memory/MemoryOutputCacheStore.cs
+++ b/src/Middleware/OutputCaching/src/Memory/MemoryOutputCacheStore.cs
@@ -83,13 +83,11 @@ public ValueTask SetAsync(string key, byte[] value, string[]? tags, TimeSpan val
             {
                 foreach (var tag in tags)
                 {
-                    if (String.IsNullOrEmpty(tag))
+                    if (tag is null)
                     {
-                        throw new ArgumentException(Resources.TagCannotBeNullOrEmpty);
+                        throw new ArgumentException(Resources.TagCannotBeNull);
                     }
 
-                    ArgumentNullException.ThrowIfNull(tag);
-
                     if (!_taggedEntries.TryGetValue(tag, out var keys))
                     {
                         keys = new HashSet<string>();
diff --git a/src/Middleware/OutputCaching/src/Resources.resx b/src/Middleware/OutputCaching/src/Resources.resx
index 122136d3eba6..621b8c2fc42c 100644
--- a/src/Middleware/OutputCaching/src/Resources.resx
+++ b/src/Middleware/OutputCaching/src/Resources.resx
@@ -117,7 +117,7 @@
   <resheader name="writer">
     <value>System.Resources.ResXResourceWriter, System.Windows.Forms, Version=4.0.0.0, Culture=neutral, PublicKeyToken=b77a5c561934e089</value>
   </resheader>
-  <data name="TagCannotBeNullOrEmpty" xml:space="preserve">
-    <value>Tag value cannot be null or empty.</value>
+  <data name="TagCannotBeNull" xml:space="preserve">
+    <value>A tag value cannot be null.</value>
   </data>
 </root>
\ No newline at end of file

From ec917262de088cca4b4c425c5e85e40043ee4ca6 Mon Sep 17 00:00:00 2001
From: Sebastien Ros <sebastienros@gmail.com>
Date: Fri, 9 Sep 2022 10:53:51 -0700
Subject: [PATCH 11/11] Fix test

---
 src/Middleware/OutputCaching/test/MemoryOutputCacheStoreTests.cs | 1 -
 1 file changed, 1 deletion(-)

diff --git a/src/Middleware/OutputCaching/test/MemoryOutputCacheStoreTests.cs b/src/Middleware/OutputCaching/test/MemoryOutputCacheStoreTests.cs
index c790370ce9c9..9441520b849f 100644
--- a/src/Middleware/OutputCaching/test/MemoryOutputCacheStoreTests.cs
+++ b/src/Middleware/OutputCaching/test/MemoryOutputCacheStoreTests.cs
@@ -199,7 +199,6 @@ public async Task ExpiredEntries_AreRemovedFromTags()
 
     [Theory]
     [InlineData(null)]
-    [InlineData("")]
     public async Task Store_Throws_OnInvalidTag(string tag)
     {
         var store = new MemoryOutputCacheStore(new MemoryCache(new MemoryCacheOptions()));