Update ConcurrentDictionary implementation for MemoryCache

Author: JunTaoLuo

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

@@ -1,12 +1,6 @@
 // Copyright (c) .NET Foundation. All rights reserved.
 // Licensed under the Apache License, Version 2.0. See License.txt in the project root for license information.
 
-#if NETSTANDARD1_3
-#else
-using System.Runtime.Remoting;
-using System.Runtime.Remoting.Messaging;
-#endif
-
 using System;
 using System.Collections.Generic;
 using System.Threading;
@@ -22,7 +16,6 @@ internal class CacheEntry : ICacheEntry
         private readonly Action<CacheEntry> _notifyCacheOfExpiration;
         private readonly Action<CacheEntry> _notifyCacheEntryDisposed;
         private IList<IDisposable> _expirationTokenRegistrations;
-        private EvictionReason _evictionReason;
         private IList<PostEvictionCallbackRegistration> _postEvictionCallbacks;
         private bool _isExpired;
 
@@ -166,6 +159,8 @@ public IList<PostEvictionCallbackRegistration> PostEvictionCallbacks
 
         internal DateTimeOffset LastAccessed { get; set; }
 
+        internal EvictionReason EvictionReason { get; private set; }
+
         public void Dispose()
         {
             if (!_added)
@@ -184,11 +179,11 @@ internal bool CheckExpired(DateTimeOffset now)
 
         internal void SetExpired(EvictionReason reason)
         {
-            _isExpired = true;
-            if (_evictionReason == EvictionReason.None)
+            if (EvictionReason == EvictionReason.None)
             {
-                _evictionReason = reason;
+                EvictionReason = reason;
             }
+            _isExpired = true;
             DetachTokens();
         }
 
@@ -302,7 +297,7 @@ private static void InvokeCallbacks(CacheEntry entry)
 
                 try
                 {
-                    registration.EvictionCallback?.Invoke(entry.Key, entry.Value, entry._evictionReason, registration.State);
+                    registration.EvictionCallback?.Invoke(entry.Key, entry.Value, entry.EvictionReason, registration.State);
                 }
                 catch (Exception)
                 {

src/Microsoft.Extensions.Caching.Memory/MemoryCache.cs

@@ -72,6 +72,8 @@ public int Count
             get { return _entries.Count; }
         }
 
+        private ICollection<KeyValuePair<object, CacheEntry>> EntriesCollection => _entries;
+
         /// <inheritdoc />
         public ICacheEntry CreateEntry(object key)
         {
@@ -86,6 +88,12 @@ public ICacheEntry CreateEntry(object key)
 
         private void SetEntry(CacheEntry entry)
         {
+            if (_disposed)
+            {
+                // No-op instead of throwing since this is called during CacheEntry.Dispose
+                return;
+            }
+
             var utcNow = _clock.UtcNow;
 
             DateTimeOffset? absoluteExpiration = null;
@@ -112,31 +120,57 @@ private void SetEntry(CacheEntry entry)
             // Initialize the last access timestamp at the time the entry is added
             entry.LastAccessed = utcNow;
 
-            var added = false;
             CacheEntry priorEntry;
-
-            if (_entries.TryRemove(entry.Key, out priorEntry))
+            if (_entries.TryGetValue(entry.Key, out priorEntry))
             {
                 priorEntry.SetExpired(EvictionReason.Replaced);
             }
 
             if (!entry.CheckExpired(utcNow))
             {
-                if (_entries.TryAdd(entry.Key, entry))
+                var entryAdded = false;
+
+                if (priorEntry == null)
+                {
+                    // Try to add the new entry if no previous entries exist.
+                    entryAdded = _entries.TryAdd(entry.Key, entry);
+                }
+                else
+                {
+                    // Try to update with the new entry if a previous entries exist.
+                    entryAdded = _entries.TryUpdate(entry.Key, entry, priorEntry);
+
+                    if (!entryAdded)
+                    {
+                        // The update will fail if the previous entry was removed after retrival.
+                        // Adding the new entry will succeed only if no entry has been added since.
+                        // This guarantees removing an old entry does not prevent adding a new entry.
+                        entryAdded = _entries.TryAdd(entry.Key, entry);
+                    }
+                }
+
+                if (entryAdded)
                 {
                     entry.AttachTokens();
-                    added = true;
                 }
-            }
+                else
+                {
+                    entry.SetExpired(EvictionReason.Replaced);
+                    entry.InvokeEvictionCallbacks();
+                }
 
-            if (priorEntry != null)
-            {
-                priorEntry.InvokeEvictionCallbacks();
+                if (priorEntry != null)
+                {
+                    priorEntry.InvokeEvictionCallbacks();
+                }
             }
-
-            if (!added)
+            else
             {
                 entry.InvokeEvictionCallbacks();
+                if (priorEntry != null)
+                {
+                    RemoveEntry(priorEntry);
+                }
             }
 
             StartScanForExpiredItems();
@@ -150,18 +184,21 @@ public bool TryGetValue(object key, out object result)
                 throw new ArgumentNullException(nameof(key));
             }
 
-            var utcNow = _clock.UtcNow;
-            result = null;
-            bool found = false;
-            CacheEntry expiredEntry = null;
             CheckDisposed();
+
+            result = null;
+            var utcNow = _clock.UtcNow;
+            var found = false;
+
             CacheEntry entry;
             if (_entries.TryGetValue(key, out entry))
             {
                 // Check if expired due to expiration tokens, timers, etc. and if so, remove it.
-                if (entry.CheckExpired(utcNow))
+                // Allow a stale Replaced value to be returned due to a race with SetExpired during SetEntry.
+                if (entry.CheckExpired(utcNow) && entry.EvictionReason != EvictionReason.Replaced)
                 {
-                    expiredEntry = entry;
+                    // TODO: For efficiency queue this up for batch removal
+                    RemoveEntry(entry);
                 }
                 else
                 {
@@ -175,12 +212,6 @@ public bool TryGetValue(object key, out object result)
                 }
             }
 
-            if (expiredEntry != null)
-            {
-                // TODO: For efficiency queue this up for batch removal
-                RemoveEntry(expiredEntry);
-            }
-
             StartScanForExpiredItems();
 
             return found;
@@ -196,46 +227,18 @@ public void Remove(object key)
 
             CheckDisposed();
             CacheEntry entry;
-            if (_entries.TryGetValue(key, out entry))
+            if (_entries.TryRemove(key, out entry))
             {
                 entry.SetExpired(EvictionReason.Removed);
-            }
-
-            if (entry != null)
-            {
-                // TODO: For efficiency consider processing these removals in batches.
-                RemoveEntry(entry);
+                entry.InvokeEvictionCallbacks();
             }
 
             StartScanForExpiredItems();
         }
 
         private void RemoveEntry(CacheEntry entry)
         {
-            // Only remove it if someone hasn't modified it since our lookup
-            CacheEntry currentEntry;
-            if (_entries.TryGetValue(entry.Key, out currentEntry)
-                && object.ReferenceEquals(currentEntry, entry))
-            {
-                _entries.TryRemove(entry.Key, out currentEntry);
-            }
-            entry.InvokeEvictionCallbacks();
-        }
-
-        private void RemoveEntries(List<CacheEntry> entries)
-        {
-            foreach (var entry in entries)
-            {
-                // Only remove it if someone hasn't modified it since our lookup
-                CacheEntry currentEntry;
-                if (_entries.TryGetValue(entry.Key, out currentEntry)
-                    && object.ReferenceEquals(currentEntry, entry))
-                {
-                    _entries.TryRemove(entry.Key, out currentEntry);
-                }
-            }
-
-            foreach (var entry in entries)
+            if (EntriesCollection.Remove(new KeyValuePair<object, CacheEntry>(entry.Key, entry)))
             {
                 entry.InvokeEvictionCallbacks();
             }
@@ -263,18 +266,14 @@ private void StartScanForExpiredItems()
 
         private static void ScanForExpiredItems(MemoryCache cache)
         {
-            List<CacheEntry> expiredEntries = new List<CacheEntry>();
-
             var now = cache._clock.UtcNow;
-            foreach (var entry in cache._entries)
+            foreach (var entry in cache._entries.Values)
             {
-                if (entry.Value.CheckExpired(now))
+                if (entry.CheckExpired(now))
                 {
-                    expiredEntries.Add(entry.Value);
+                    cache.RemoveEntry(entry);
                 }
             }
-
-            cache.RemoveEntries(expiredEntries);
         }
 
         /// This is called after a Gen2 garbage collection. We assume this means there was memory pressure.
@@ -294,80 +293,79 @@ private bool DoMemoryPreassureCollection(object state)
         /// Remove at least the given percentage (0.10 for 10%) of the total entries (or estimated memory?), according to the following policy:
         /// 1. Remove all expired items.
         /// 2. Bucket by CacheItemPriority.
-        /// ?. Least recently used objects.
+        /// 3. Least recently used objects.
         /// ?. Items with the soonest absolute expiration.
         /// ?. Items with the soonest sliding expiration.
         /// ?. Larger objects - estimated by object graph size, inaccurate.
         public void Compact(double percentage)
         {
-            List<CacheEntry> expiredEntries = new List<CacheEntry>();
-            List<CacheEntry> lowPriEntries = new List<CacheEntry>();
-            List<CacheEntry> normalPriEntries = new List<CacheEntry>();
-            List<CacheEntry> highPriEntries = new List<CacheEntry>();
-            List<CacheEntry> neverRemovePriEntries = new List<CacheEntry>();
+            var entriesToRemove = new List<CacheEntry>();
+            var lowPriEntries = new List<CacheEntry>();
+            var normalPriEntries = new List<CacheEntry>();
+            var highPriEntries = new List<CacheEntry>();
 
             // Sort items by expired & priority status
             var now = _clock.UtcNow;
-            foreach (var entry in _entries)
+            foreach (var entry in _entries.Values)
             {
-                if (entry.Value.CheckExpired(now))
+                if (entry.CheckExpired(now))
                 {
-                    expiredEntries.Add(entry.Value);
+                    entriesToRemove.Add(entry);
                 }
                 else
                 {
-                    switch (entry.Value.Priority)
+                    switch (entry.Priority)
                     {
                         case CacheItemPriority.Low:
-                            lowPriEntries.Add(entry.Value);
+                            lowPriEntries.Add(entry);
                             break;
                         case CacheItemPriority.Normal:
-                            normalPriEntries.Add(entry.Value);
+                            normalPriEntries.Add(entry);
                             break;
                         case CacheItemPriority.High:
-                            highPriEntries.Add(entry.Value);
+                            highPriEntries.Add(entry);
                             break;
                         case CacheItemPriority.NeverRemove:
-                            neverRemovePriEntries.Add(entry.Value);
                             break;
                         default:
-                            System.Diagnostics.Debug.Assert(false, "Not implemented: " + entry.Value.Priority);
-                            break;
+                            throw new NotSupportedException("Not implemented: " + entry.Priority);
                     }
                 }
             }
 
-            int totalEntries = expiredEntries.Count + lowPriEntries.Count + normalPriEntries.Count + highPriEntries.Count + neverRemovePriEntries.Count;
-            int removalCountTarget = (int)(totalEntries * percentage);
+            int removalCountTarget = (int)(_entries.Count * percentage);
 
-            ExpirePriorityBucket(removalCountTarget, expiredEntries, lowPriEntries);
-            ExpirePriorityBucket(removalCountTarget, expiredEntries, normalPriEntries);
-            ExpirePriorityBucket(removalCountTarget, expiredEntries, highPriEntries);
+            ExpirePriorityBucket(removalCountTarget, entriesToRemove, lowPriEntries);
+            ExpirePriorityBucket(removalCountTarget, entriesToRemove, normalPriEntries);
+            ExpirePriorityBucket(removalCountTarget, entriesToRemove, highPriEntries);
 
-            RemoveEntries(expiredEntries);
+            foreach (var entry in entriesToRemove)
+            {
+                RemoveEntry(entry);
+            }
         }
 
         /// Policy:
-        /// ?. Least recently used objects.
+        /// 1. Least recently used objects.
         /// ?. Items with the soonest absolute expiration.
         /// ?. Items with the soonest sliding expiration.
         /// ?. Larger objects - estimated by object graph size, inaccurate.
-        private void ExpirePriorityBucket(int removalCountTarget, List<CacheEntry> expiredEntries, List<CacheEntry> priorityEntries)
+        private void ExpirePriorityBucket(int removalCountTarget, List<CacheEntry> entriesToRemove, List<CacheEntry> priorityEntries)
         {
             // Do we meet our quota by just removing expired entries?
-            if (removalCountTarget <= expiredEntries.Count)
+            if (removalCountTarget <= entriesToRemove.Count)
             {
                 // No-op, we've met quota
                 return;
             }
-            if (expiredEntries.Count + priorityEntries.Count <= removalCountTarget)
+            if (entriesToRemove.Count + priorityEntries.Count <= removalCountTarget)
             {
                 // Expire all of the entries in this bucket
                 foreach (var entry in priorityEntries)
                 {
                     entry.SetExpired(EvictionReason.Capacity);
                 }
-                expiredEntries.AddRange(priorityEntries);
+                entriesToRemove.AddRange(priorityEntries);
                 return;
             }
 
@@ -378,8 +376,8 @@ private void ExpirePriorityBucket(int removalCountTarget, List<CacheEntry> expir
             foreach (var entry in priorityEntries.OrderBy(entry => entry.LastAccessed))
             {
                 entry.SetExpired(EvictionReason.Capacity);
-                expiredEntries.Add(entry);
-                if (removalCountTarget <= expiredEntries.Count)
+                entriesToRemove.Add(entry);
+                if (removalCountTarget <= entriesToRemove.Count)
                 {
                     break;
                 }

test/Microsoft.Extensions.Caching.Memory.Tests/CacheEntryScopeExpirationTests.cs

@@ -376,8 +376,8 @@ public void LinkContextsCanNest()
 
             Assert.Equal(1, ((CacheEntry)entry1)._expirationTokens.Count());
             Assert.Null(((CacheEntry)entry1)._absoluteExpiration);
-            Assert.Equal(1, ((CacheEntry)entry1)._expirationTokens.Count());
-            Assert.Null(((CacheEntry)entry1)._absoluteExpiration);
+            Assert.Equal(1, ((CacheEntry)entry)._expirationTokens.Count());
+            Assert.Null(((CacheEntry)entry)._absoluteExpiration);
         }
 
         [Fact]