Update ConcurrentDictionary implementation for MemoryCache

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 @@ namespace Microsoft.Extensions.Caching.Memory
         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 @@ namespace Microsoft.Extensions.Caching.Memory
 
         internal DateTimeOffset LastAccessed { get; set; }
 
+        internal EvictionReason EvictionReason { get; private set; }
+
         public void Dispose()
         {
             if (!_added)
@@ -184,11 +179,11 @@ namespace Microsoft.Extensions.Caching.Memory
 
         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 @@ namespace Microsoft.Extensions.Caching.Memory
 
                 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 @@ namespace Microsoft.Extensions.Caching.Memory
             get { return _entries.Count; }
         }
 
+        private ICollection<KeyValuePair<object, CacheEntry>> EntriesCollection => _entries;
+
         /// <inheritdoc />
         public ICacheEntry CreateEntry(object key)
         {
@@ -86,6 +88,12 @@ namespace Microsoft.Extensions.Caching.Memory
 
         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 @@ namespace Microsoft.Extensions.Caching.Memory
             // Initialize the last access timestamp at the time the entry is added
             entry.LastAccessed = utcNow;
 
-            var added = false;
             CacheEntry priorEntry;
-
+            if (_entries.TryGetValue(entry.Key, out priorEntry))
-            if (_entries.TryRemove(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();
                 }
             }
-
+            else
-            if (priorEntry != null)
-            {
-                priorEntry.InvokeEvictionCallbacks();
-            }
-
-            if (!added)
             {
                 entry.InvokeEvictionCallbacks();
+                if (priorEntry != null)
+                {
+                    RemoveEntry(priorEntry);
+                }
             }
 
             StartScanForExpiredItems();
@@ -150,18 +184,21 @@ namespace Microsoft.Extensions.Caching.Memory
                 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 @@ namespace Microsoft.Extensions.Caching.Memory
                 }
             }
 
-            if (expiredEntry != null)
-            {
-                // TODO: For efficiency queue this up for batch removal
-                RemoveEntry(expiredEntry);
-            }
-
             StartScanForExpiredItems();
 
             return found;
@@ -196,15 +227,10 @@ namespace Microsoft.Extensions.Caching.Memory
 
             CheckDisposed();
             CacheEntry entry;
-            if (_entries.TryGetValue(key, out entry))
+            if (_entries.TryRemove(key, out entry))
             {
                 entry.SetExpired(EvictionReason.Removed);
-            }
+                entry.InvokeEvictionCallbacks();
-
-            if (entry != null)
-            {
-                // TODO: For efficiency consider processing these removals in batches.
-                RemoveEntry(entry);
             }
 
             StartScanForExpiredItems();
@@ -212,30 +238,7 @@ namespace Microsoft.Extensions.Caching.Memory
 
         private void RemoveEntry(CacheEntry entry)
         {
-            // Only remove it if someone hasn't modified it since our lookup
+            if (EntriesCollection.Remove(new KeyValuePair<object, CacheEntry>(entry.Key, entry)))
-            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)
             {
                 entry.InvokeEvictionCallbacks();
             }
@@ -263,18 +266,14 @@ namespace Microsoft.Extensions.Caching.Memory
 
         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 @@ namespace Microsoft.Extensions.Caching.Memory
         /// 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>();
+            var entriesToRemove = new List<CacheEntry>();
-            List<CacheEntry> lowPriEntries = new List<CacheEntry>();
+            var lowPriEntries = new List<CacheEntry>();
-            List<CacheEntry> normalPriEntries = new List<CacheEntry>();
+            var normalPriEntries = new List<CacheEntry>();
-            List<CacheEntry> highPriEntries = new List<CacheEntry>();
+            var highPriEntries = new List<CacheEntry>();
-            List<CacheEntry> neverRemovePriEntries = 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);
+                            throw new NotSupportedException("Not implemented: " + entry.Priority);
-                            break;
                     }
                 }
             }
 
-            int totalEntries = expiredEntries.Count + lowPriEntries.Count + normalPriEntries.Count + highPriEntries.Count + neverRemovePriEntries.Count;
+            int removalCountTarget = (int)(_entries.Count * percentage);
-            int removalCountTarget = (int)(totalEntries * percentage);
 
-            ExpirePriorityBucket(removalCountTarget, expiredEntries, lowPriEntries);
+            ExpirePriorityBucket(removalCountTarget, entriesToRemove, lowPriEntries);
-            ExpirePriorityBucket(removalCountTarget, expiredEntries, normalPriEntries);
+            ExpirePriorityBucket(removalCountTarget, entriesToRemove, normalPriEntries);
-            ExpirePriorityBucket(removalCountTarget, expiredEntries, highPriEntries);
+            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 @@ namespace Microsoft.Extensions.Caching.Memory
             foreach (var entry in priorityEntries.OrderBy(entry => entry.LastAccessed))
             {
                 entry.SetExpired(EvictionReason.Capacity);
-                expiredEntries.Add(entry);
+                entriesToRemove.Add(entry);
-                if (removalCountTarget <= expiredEntries.Count)
+                if (removalCountTarget <= entriesToRemove.Count)
                 {
                     break;
                 }

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

@@ -376,8 +376,8 @@ namespace Microsoft.Extensions.Caching.Memory
 
             Assert.Equal(1, ((CacheEntry)entry1)._expirationTokens.Count());
             Assert.Null(((CacheEntry)entry1)._absoluteExpiration);
-            Assert.Equal(1, ((CacheEntry)entry1)._expirationTokens.Count());
+            Assert.Equal(1, ((CacheEntry)entry)._expirationTokens.Count());
-            Assert.Null(((CacheEntry)entry1)._absoluteExpiration);
+            Assert.Null(((CacheEntry)entry)._absoluteExpiration);
         }
 
         [Fact]

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

@@ -2,16 +2,18 @@
 // Licensed under the Apache License, Version 2.0. See License.txt in the project root for license information.
 
 using System;
+using Microsoft.Extensions.Internal;
 using Xunit;
 
 namespace Microsoft.Extensions.Caching.Memory
 {
     public class CompactTests
     {
-        private MemoryCache CreateCache()
+        private MemoryCache CreateCache(ISystemClock clock = null)
         {
             return new MemoryCache(new MemoryCacheOptions()
             {
+                Clock = clock,
                 CompactOnMemoryPressure = false,
             });
         }
@@ -67,10 +69,14 @@ namespace Microsoft.Extensions.Caching.Memory
         [Fact]
         public void CompactPrioritizesLRU()
         {
-            var cache = CreateCache();
+            var testClock = new TestClock();
+            var cache = CreateCache(testClock);
             cache.Set("key1", "value1");
+            testClock.Add(TimeSpan.FromSeconds(1));
             cache.Set("key2", "value2");
+            testClock.Add(TimeSpan.FromSeconds(1));
             cache.Set("key3", "value3");
+            testClock.Add(TimeSpan.FromSeconds(1));
             cache.Set("key4", "value4");
             Assert.Equal(4, cache.Count);
             cache.Compact(0.90);

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

@@ -4,7 +4,6 @@
 using System;
 using System.Threading;
 using System.Threading.Tasks;
-using Microsoft.Extensions.Primitives;
 using Xunit;
 
 namespace Microsoft.Extensions.Caching.Memory
@@ -410,38 +409,56 @@ namespace Microsoft.Extensions.Caching.Memory
         }
 
         [Fact]
-        public void GetAndSet_AreThreadSafe()
+        public void GetAndSet_AreThreadSafe_AndUpdatesNeverLeavesNullValues()
         {
             var cache = CreateCache();
             string key = "myKey";
             var cts = new CancellationTokenSource();
-            var cts2 = new CancellationTokenSource();
+            var readValueIsNull = false;
 
-            cache.Set(key, new Guid(), new MemoryCacheEntryOptions().AddExpirationToken(new CancellationChangeToken(cts.Token)));
+            cache.Set(key, new Guid());
+
+            var task0 = Task.Run(() =>
+            {
+                while (!cts.IsCancellationRequested)
+                {
+                    cache.Set(key, Guid.NewGuid());
+                }
+            });
 
             var task1 = Task.Run(() =>
             {
-                while (!cts2.IsCancellationRequested)
+                while (!cts.IsCancellationRequested)
                 {
-                    cache.Set(key, new Guid());
+                    cache.Set(key, Guid.NewGuid());
                 }
             });
 
             var task2 = Task.Run(() =>
             {
-                while (!cts2.IsCancellationRequested)
+                while (!cts.IsCancellationRequested)
                 {
-                    cache.Get(key);
+                    if (cache.Get(key) == null)
+                    {
+                        // Stop this task and update flag for assertion
+                        readValueIsNull = true;
+                        break;
+                    }
                 }
             });
 
-            var task3 = Task.Run(async () =>
+            var task3 = Task.Delay(TimeSpan.FromSeconds(10));
-            {
-                await Task.Delay(TimeSpan.FromSeconds(1));
-                cts2.Cancel();
-            });
 
-            Task.WaitAll(task1, task2, task3);
+            Task.WaitAny(task0, task1, task2, task3);
+
+            Assert.False(readValueIsNull);
+            Assert.Equal(TaskStatus.Running, task0.Status);
+            Assert.Equal(TaskStatus.Running, task1.Status);
+            Assert.Equal(TaskStatus.Running, task2.Status);
+            Assert.Equal(TaskStatus.RanToCompletion, task3.Status);
+
+            cts.Cancel();
+            Task.WaitAll(task0, task1, task2, task3);
         }
 
 #if NET451

test/Microsoft.Extensions.Caching.Redis.Tests/Microsoft.Extensions.Caching.Redis.Tests.xproj

@@ -13,5 +13,8 @@
   <PropertyGroup>
     <SchemaVersion>2.0</SchemaVersion>
   </PropertyGroup>
+  <ItemGroup>
+    <Service Include="{82a7f48d-3b50-4b1e-b82e-3ada8210c358}" />
+  </ItemGroup>
   <Import Project="$(VSToolsPath)\DNX\Microsoft.DNX.targets" Condition="'$(VSToolsPath)' != ''" />
 </Project>