Linked cache entry thread safety fix

src/libraries/Microsoft.Extensions.Caching.Memory/src/CacheEntry.CacheEntryTokens.cs

@@ -17,23 +17,51 @@ internal sealed partial class CacheEntry
         // which typically is not using expiration tokens or callbacks
         private sealed class CacheEntryTokens
         {
-            private List<IChangeToken>? _expirationTokens;
+            // ExpirationTokensList is copy-on-write, so CheckForExpiredTokens - which runs on every
+            // cache hit - can scan the tokens without taking a lock while another thread adds to them.
+            private ExpirationTokensList? _expirationTokens;
             private List<IDisposable>? _expirationTokenRegistrations;
             private List<PostEvictionCallbackRegistration>? _postEvictionCallbacks; // this is not really related to tokens, but was moved here to shrink typical CacheEntry size
 
-            internal List<IChangeToken> ExpirationTokens => _expirationTokens ??= new List<IChangeToken>();
-            internal List<PostEvictionCallbackRegistration> PostEvictionCallbacks => _postEvictionCallbacks ??= new List<PostEvictionCallbackRegistration>();
+            internal ExpirationTokensList ExpirationTokens
+            {
+                get
+                {
+                    ExpirationTokensList? expirationTokens = _expirationTokens;
+                    if (expirationTokens is not null)
+                    {
+                        return expirationTokens;
+                    }
+
+                    expirationTokens = new ExpirationTokensList();
+                    return Interlocked.CompareExchange(ref _expirationTokens, expirationTokens, null) ?? expirationTokens;
+                }
+            }
+
+            internal List<PostEvictionCallbackRegistration> PostEvictionCallbacks
+            {
+                get
+                {
+                    List<PostEvictionCallbackRegistration>? postEvictionCallbacks = _postEvictionCallbacks;
+                    if (postEvictionCallbacks is not null)
+                    {
+                        return postEvictionCallbacks;
+                    }
+
+                    postEvictionCallbacks = new List<PostEvictionCallbackRegistration>();
+                    return Interlocked.CompareExchange(ref _postEvictionCallbacks, postEvictionCallbacks, null) ?? postEvictionCallbacks;
+                }
+            }
 
             internal void AttachTokens(CacheEntry cacheEntry)
             {
-                List<IChangeToken>? expirationTokens = _expirationTokens;
+                ExpirationTokensList? expirationTokens = _expirationTokens;
                 if (expirationTokens is not null)
                 {
                     lock (this)
                     {
-                        for (int i = 0; i < expirationTokens.Count; i++)
+                        foreach (IChangeToken expirationToken in expirationTokens.Snapshot)
                         {
-                            IChangeToken expirationToken = expirationTokens[i];
                             if (expirationToken.ActiveChangeCallbacks)
                             {
                                 _expirationTokenRegistrations ??= new List<IDisposable>(1);
@@ -47,12 +75,11 @@ internal void AttachTokens(CacheEntry cacheEntry)
 
             internal bool CheckForExpiredTokens(CacheEntry cacheEntry)
             {
-                List<IChangeToken>? expirationTokens = _expirationTokens;
+                ExpirationTokensList? expirationTokens = _expirationTokens;
                 if (expirationTokens is not null)
                 {
-                    for (int i = 0; i < expirationTokens.Count; i++)
+                    foreach (IChangeToken expiredToken in expirationTokens.Snapshot)
                     {
-                        IChangeToken expiredToken = expirationTokens[i];
                         if (expiredToken.HasChanged)
                         {
                             cacheEntry.SetExpired(EvictionReason.TokenExpired);
@@ -67,16 +94,10 @@ internal bool CheckForExpiredTokens(CacheEntry cacheEntry)
 
             internal void PropagateTokens(CacheEntry parentEntry)
             {
-                if (_expirationTokens != null)
+                ExpirationTokensList? expirationTokens = _expirationTokens;
+                if (expirationTokens is not null)
                 {
-                    lock (this)
-                    {
-                        CacheEntryTokens parentTokens = parentEntry.GetOrCreateTokens();
-                        lock (parentTokens)
-                        {
-                            parentTokens.ExpirationTokens.AddRange(_expirationTokens);
-                        }
-                    }
+                    parentEntry.GetOrCreateTokens().ExpirationTokens.AddRange(expirationTokens.Snapshot);
                 }
             }
 

src/libraries/Microsoft.Extensions.Caching.Memory/src/CacheEntry.ExpirationTokensList.cs

@@ -0,0 +1,142 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+
+using System;
+using System.Collections;
+using System.Collections.Generic;
+using Microsoft.Extensions.Primitives;
+
+namespace Microsoft.Extensions.Caching.Memory
+{
+    internal sealed partial class CacheEntry
+    {
+        /// <summary>
+        /// The list behind <see cref="ICacheEntry.ExpirationTokens"/>. Mutations rebuild the backing
+        /// array under a lock and publish it with a volatile write, so a reader can take the array
+        /// and walk it without synchronizing and without observing a half-written state.
+        /// </summary>
+        /// <remarks>
+        /// The tokens are scanned on every cache hit but written to very rarely, and the entry is
+        /// already published in the cache by the time a linked child entry propagates its tokens into
+        /// it - from whichever thread committed the child. Locking the reader instead would penalize
+        /// the hot path, and would not help anyway: callers mutate this list directly through the
+        /// public <see cref="ICacheEntry.ExpirationTokens"/> property, so the writes can only be made
+        /// safe from inside the list itself.
+        /// </remarks>
+        private sealed class ExpirationTokensList : IList<IChangeToken>
+        {
+            private readonly object _lock = new object();
+            private volatile IChangeToken[] _items = Array.Empty<IChangeToken>();
+
+            /// <summary>
+            /// Gets the current contents. The array is shared with concurrent readers, so callers
+            /// must only read from it.
+            /// </summary>
+            internal IChangeToken[] Snapshot => _items;
+
+            public int Count => _items.Length;
+
+            public bool IsReadOnly => false;
+
+            public IChangeToken this[int index]
+            {
+                get
+                {
+                    IChangeToken[] items = _items;
+                    if ((uint)index >= (uint)items.Length)
+                    {
+                        throw new ArgumentOutOfRangeException(nameof(index));
+                    }
+
+                    return items[index];
+                }
+                set => Mutate(list => list[index] = value);
+            }
+
+            public void Add(IChangeToken item)
+            {
+                lock (_lock)
+                {
+                    IChangeToken[] items = _items;
+                    var updated = new IChangeToken[items.Length + 1];
+                    Array.Copy(items, updated, items.Length);
+                    updated[items.Length] = item;
+                    _items = updated;
+                }
+            }
+
+            internal void AddRange(IChangeToken[] source)
+            {
+                if (source.Length == 0)
+                {
+                    return;
+                }
+
+                lock (_lock)
+                {
+                    IChangeToken[] items = _items;
+                    var updated = new IChangeToken[items.Length + source.Length];
+                    Array.Copy(items, updated, items.Length);
+                    Array.Copy(source, 0, updated, items.Length, source.Length);
+                    _items = updated;
+                }
+            }
+
+            // The remaining mutations are not used by the cache itself and are expected to be rare,
+            // so they favour borrowing List<T>'s behaviour (including its argument validation) over
+            // hand-rolling the array manipulation.
+            public void Insert(int index, IChangeToken item) => Mutate(list => list.Insert(index, item));
+
+            public void RemoveAt(int index) => Mutate(list => list.RemoveAt(index));
+
+            public void Clear()
+            {
+                lock (_lock)
+                {
+                    _items = Array.Empty<IChangeToken>();
+                }
+            }
+
+            public bool Remove(IChangeToken item)
+            {
+                lock (_lock)
+                {
+                    int index = Array.IndexOf(_items, item);
+                    if (index < 0)
+                    {
+                        return false;
+                    }
+
+                    List<IChangeToken> updated = new List<IChangeToken>(_items);
+                    updated.RemoveAt(index);
+                    _items = updated.ToArray();
+                    return true;
+                }
+            }
+
+            public bool Contains(IChangeToken item) => Array.IndexOf(_items, item) >= 0;
+
+            public int IndexOf(IChangeToken item) => Array.IndexOf(_items, item);
+
+            public void CopyTo(IChangeToken[] array, int arrayIndex)
+            {
+                IChangeToken[] items = _items;
+                Array.Copy(items, 0, array, arrayIndex, items.Length);
+            }
+
+            public IEnumerator<IChangeToken> GetEnumerator() => ((IEnumerable<IChangeToken>)_items).GetEnumerator();
+
+            IEnumerator IEnumerable.GetEnumerator() => _items.GetEnumerator();
+
+            private void Mutate(Action<List<IChangeToken>> mutation)
+            {
+                lock (_lock)
+                {
+                    List<IChangeToken> updated = new List<IChangeToken>(_items);
+                    mutation(updated);
+                    _items = updated.ToArray();
+                }
+            }
+        }
+    }
+}

src/libraries/Microsoft.Extensions.Caching.Memory/tests/CacheEntryScopeExpirationTests.cs

@@ -509,6 +509,46 @@ await Task.WhenAll(
             Assert.Null(cache.Get(key4));
         }
 
+        [Fact]
+        public async Task PropagatingTokensToTheParentIsSafeWhileTheParentIsRead()
+        {
+            const int Workers = 4;
+            const int ChildrenPerWorker = 250;
+
+            var cache = CreateCache(trackLinkedCacheEntries: true);
+            string parentKey = "parent";
+
+            ICacheEntry parent = cache.CreateEntry(parentKey);
+            parent.SetValue(new object());
+
+            // Started inside the parent's scope, so every task inherits it as the ambient entry and
+            // each child they create propagates its expiration token into the parent - from several
+            // threads at once, and long after the parent itself has been committed to the cache and
+            // become visible to readers.
+            Task[] workers = Enumerable.Range(0, Workers)
+                .Select(worker => Task.Run(() =>
+                {
+                    for (int i = 0; i < ChildrenPerWorker; i++)
+                    {
+                        using ICacheEntry child = cache.CreateEntry($"child {worker}.{i}");
+                        child.SetValue(i);
+                        child.AddExpirationToken(new TestExpirationToken());
+                    }
+                }))
+                .ToArray();
+
+            parent.Dispose();
+
+            Task allWorkers = Task.WhenAll(workers);
+            while (!allWorkers.IsCompleted)
+            {
+                Assert.True(cache.TryGetValue(parentKey, out _));
+            }
+
+            await allWorkers;
+            Assert.Equal(Workers * ChildrenPerWorker, parent.ExpirationTokens.Count);
+        }
+
         [Fact]
         public async Task OnceExpiredIsSetToTrueItRemainsTrue()
         {

src/libraries/Microsoft.Extensions.Caching.Memory/tests/TokenExpirationTests.cs

@@ -2,6 +2,7 @@
 // The .NET Foundation licenses this file to you under the MIT license.
 
 using System;
+using System.Collections.Generic;
 using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.Extensions.Caching.Memory.Infrastructure;
@@ -259,6 +260,81 @@ public void PostEvictionCallbacksGetInvokedWhenMemoryCacheEntriesExpireWithAnAct
             Assert.False(cache.TryGetValue(key, out _));
         }
 
+        [Fact]
+        public async Task AddingExpirationTokensIsSafeWhileTheEntryIsRead()
+        {
+            const int TokenCount = 2_000;
+
+            var cache = CreateCache();
+            string key = "myKey";
+
+            ICacheEntry entry = cache.CreateEntry(key);
+            entry.SetValue(new object());
+            entry.Dispose(); // commits the entry, making it visible to readers
+
+            Task writer = Task.Run(() =>
+            {
+                for (int i = 0; i < TokenCount; i++)
+                {
+                    entry.AddExpirationToken(new TestExpirationToken());
+                }
+            });
+
+            // Every read scans the entry's expiration tokens; it must never observe a torn list.
+            while (!writer.IsCompleted)
+            {
+                Assert.True(cache.TryGetValue(key, out _));
+            }
+
+            await writer;
+            Assert.Equal(TokenCount, entry.ExpirationTokens.Count);
+        }
+
+        [Fact]
+        public void ExpirationTokensBehaveLikeAList()
+        {
+            var cache = CreateCache();
+            using ICacheEntry entry = cache.CreateEntry("myKey");
+            IList<IChangeToken> tokens = entry.ExpirationTokens;
+
+            var first = new TestExpirationToken();
+            var second = new TestExpirationToken();
+            var third = new TestExpirationToken();
+
+            Assert.Empty(tokens);
+            Assert.False(tokens.IsReadOnly);
+
+            tokens.Add(first);
+            tokens.Add(third);
+            tokens.Insert(1, second);
+            Assert.Equal(new[] { first, second, third }, tokens);
+            Assert.Equal(3, tokens.Count);
+
+            Assert.Same(second, tokens[1]);
+            Assert.Equal(2, tokens.IndexOf(third));
+            Assert.True(tokens.Contains(second));
+            Assert.Throws<ArgumentOutOfRangeException>(() => tokens[3]);
+            Assert.Throws<ArgumentOutOfRangeException>(() => tokens.Insert(4, first));
+            Assert.Throws<ArgumentOutOfRangeException>(() => tokens.RemoveAt(3));
+
+            var target = new IChangeToken[4];
+            tokens.CopyTo(target, 1);
+            Assert.Equal(new IChangeToken[] { null, first, second, third }, target);
+
+            tokens[0] = third;
+            Assert.Same(third, tokens[0]);
+
+            Assert.True(tokens.Remove(second));
+            Assert.False(tokens.Remove(second));
+            Assert.Equal(new[] { third, third }, tokens);
+
+            tokens.RemoveAt(0);
+            Assert.Same(third, Assert.Single(tokens));
+
+            tokens.Clear();
+            Assert.Empty(tokens);
+        }
+
         internal class TestToken : IChangeToken
         {
             private bool _hasChanged;