[release/7.0] Remove tagged keys as entries are evicted (#43728)

* Remove tagged keys as entries are evicted

* Revert middleware tests changes

* Fix unit test

* Add braces

* Fix flaky tests

* Fix tests

* Feedback

* Feedback

* Comment

* Remove left over exception

* Fix test

Author: sebastienros

src/Middleware/OutputCaching/src/Memory/MemoryOutputCacheStore.cs

@@ -1,23 +1,27 @@
 // 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 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);
 
         _cache = cache;
     }
 
+    // For testing
+    internal Dictionary<string, HashSet<string>> TaggedEntries => _taggedEntries;
+
     public ValueTask EvictByTagAsync(string tag, CancellationToken cancellationToken)
     {
         ArgumentNullException.ThrowIfNull(tag);
@@ -26,12 +30,29 @@ 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);
-                }
+                    // 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.
 
-                _taggedEntries.Remove(tag);
+                    var i = keys.Count;
+                    while (i > 0)
+                    {
+                        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;
+                            }
+                        }
+                    }
+                }
             }
         }
 
@@ -62,35 +83,75 @@ public ValueTask SetAsync(string key, byte[] value, string[]? tags, TimeSpan val
             {
                 foreach (var tag in tags)
                 {
+                    if (tag is null)
+                    {
+                        throw new ArgumentException(Resources.TagCannotBeNull);
+                    }
+
                     if (!_taggedEntries.TryGetValue(tag, out var keys))
                     {
                         keys = new HashSet<string>();
                         _taggedEntries[tag] = keys;
                     }
 
+                    Debug.Assert(keys != null);
+
                     keys.Add(key);
                 }
 
-                SetEntry();
+                SetEntry(key, value, tags, validFor);
             }
         }
         else
         {
-            SetEntry();
+            SetEntry(key, value, tags, validFor);
         }
 
-        void SetEntry()
+        return ValueTask.CompletedTask;
+    }
+
+    void SetEntry(string key, byte[] value, string[]? tags, TimeSpan validFor)
+    {
+        Debug.Assert(key != null);
+
+        var options = new MemoryCacheEntryOptions
         {
-            _cache.Set(
-            key,
-            value,
-            new MemoryCacheEntryOptions
-            {
-                AbsoluteExpirationRelativeToNow = validFor,
-                Size = value.Length
-            });
+            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);
         }
 
-        return ValueTask.CompletedTask;
+        _cache.Set(key, value, options);
+    }
+
+    void RemoveFromTags(object key, object? value, EvictionReason reason, object? state)
+    {
+        var tags = state as string[];
+
+        Debug.Assert(tags != null);
+        Debug.Assert(tags.Length > 0);
+        Debug.Assert(key is string);
+
+        lock (_tagsLock)
+        {
+            foreach (var tag in tags)
+            {
+                if (_taggedEntries.TryGetValue(tag, out var tagged))
+                {
+                    tagged.Remove((string)key);
+
+                    // Remove the collection if there is no more keys in it
+                    if (tagged.Count == 0)
+                    {
+                        _taggedEntries.Remove(tag);
+                    }
+                }
+            }
+        }
     }
 }

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="TagCannotBeNull" xml:space="preserve">
+    <value>A tag value cannot be null.</value>
   </data>
 </root>

src/Middleware/OutputCaching/test/MemoryOutputCacheStoreTests.cs

@@ -71,7 +71,19 @@ public async Task EvictByTag_SingleTag_SingleEntry()
         await store.EvictByTagAsync("tag1", default);
         var result = await store.GetAsync(key, default);
 
+        HashSet<string> tag1s;
+
+        // Wait for the hashset to be removed as it happens on a separate thread
+
+        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);
     }
 
     [Fact]
@@ -140,6 +152,62 @@ public async Task EvictByTag_MultipleTags_MultipleEntries()
         Assert.Null(result2);
     }
 
+    [Fact]
+    public async Task ExpiredEntries_AreRemovedFromTags()
+    {
+        var testClock = new TestMemoryOptionsClock { UtcNow = DateTimeOffset.UtcNow };
+        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);
+        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));
+
+        // 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
+
+        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);
+    }
+
+    [Theory]
+    [InlineData(null)]
+    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; }