[release/7.0] Remove tagged keys as entries are evicted

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;

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; }