MECM MemoryCache performance; split string keys into separate collection

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

@@ -2,11 +2,13 @@
 // The .NET Foundation licenses this file to you under the MIT license.
 
 using System;
+using System.Collections;
 using System.Collections.Concurrent;
 using System.Collections.Generic;
 using System.Diagnostics;
 using System.Diagnostics.CodeAnalysis;
 using System.Runtime.CompilerServices;
+using System.Runtime.InteropServices;
 using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.Extensions.Logging;
@@ -81,15 +83,7 @@ public MemoryCache(IOptions<MemoryCacheOptions> optionsAccessor, ILoggerFactory
         /// Gets an enumerable of the all the keys in the <see cref="MemoryCache"/>.
         /// </summary>
         public IEnumerable<object> Keys
-        {
-            get
-            {
-                foreach (KeyValuePair<object, CacheEntry> pairs in _coherentState._entries)
-                {
-                    yield return pairs.Key;
-                }
-            }
-        }
+            => _coherentState.GetAllKeys();
 
         /// <summary>
         /// Internal accessor for Size for testing only.
@@ -141,7 +135,7 @@ internal void SetEntry(CacheEntry entry)
             entry.LastAccessed = utcNow;
 
             CoherentState coherentState = _coherentState; // Clear() can update the reference in the meantime
-            if (coherentState._entries.TryGetValue(entry.Key, out CacheEntry? priorEntry))
+            if (coherentState.TryGetValue(entry.Key, out CacheEntry? priorEntry))
             {
                 priorEntry.SetExpired(EvictionReason.Replaced);
             }
@@ -160,19 +154,19 @@ internal void SetEntry(CacheEntry entry)
                 if (priorEntry == null)
                 {
                     // Try to add the new entry if no previous entries exist.
-                    entryAdded = coherentState._entries.TryAdd(entry.Key, entry);
+                    entryAdded = coherentState.TryAdd(entry.Key, entry);
                 }
                 else
                 {
                     // Try to update with the new entry if a previous entries exist.
-                    entryAdded = coherentState._entries.TryUpdate(entry.Key, entry, priorEntry);
+                    entryAdded = coherentState.TryUpdate(entry.Key, entry, priorEntry);
 
                     if (!entryAdded)
                     {
                         // The update will fail if the previous entry was removed after retrieval.
                         // 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 = coherentState._entries.TryAdd(entry.Key, entry);
+                        entryAdded = coherentState.TryAdd(entry.Key, entry);
                     }
                 }
 
@@ -217,7 +211,7 @@ public bool TryGetValue(object key, out object? result)
             DateTime utcNow = UtcNow;
 
             CoherentState coherentState = _coherentState; // Clear() can update the reference in the meantime
-            if (coherentState._entries.TryGetValue(key, out CacheEntry? tmp))
+            if (coherentState.TryGetValue(key, out CacheEntry? tmp))
             {
                 CacheEntry entry = tmp;
                 // Check if expired due to expiration tokens, timers, etc. and if so, remove it.
@@ -276,7 +270,8 @@ public void Remove(object key)
             CheckDisposed();
 
             CoherentState coherentState = _coherentState; // Clear() can update the reference in the meantime
-            if (coherentState._entries.TryRemove(key, out CacheEntry? entry))
+
+            if (coherentState.TryRemove(key, out CacheEntry? entry))
             {
                 if (_options.HasSizeLimit)
                 {
@@ -298,10 +293,10 @@ public void Clear()
             CheckDisposed();
 
             CoherentState oldState = Interlocked.Exchange(ref _coherentState, new CoherentState());
-            foreach (KeyValuePair<object, CacheEntry> entry in oldState._entries)
+            foreach (CacheEntry entry in oldState.GetAllValues())
             {
-                entry.Value.SetExpired(EvictionReason.Removed);
-                entry.Value.InvokeEvictionCallbacks();
+                entry.SetExpired(EvictionReason.Removed);
+                entry.InvokeEvictionCallbacks();
             }
         }
 
@@ -422,10 +417,9 @@ private void ScanForExpiredItems()
             DateTime utcNow = _lastExpirationScan = UtcNow;
 
             CoherentState coherentState = _coherentState; // Clear() can update the reference in the meantime
-            foreach (KeyValuePair<object, CacheEntry> item in coherentState._entries)
-            {
-                CacheEntry entry = item.Value;
 
+            foreach (CacheEntry entry in coherentState.GetAllValues())
+            {
                 if (entry.CheckExpired(utcNow))
                 {
                     coherentState.RemoveEntry(entry, _options);
@@ -547,9 +541,8 @@ private void Compact(long removalSizeTarget, Func<CacheEntry, long> computeEntry
 
             // Sort items by expired & priority status
             DateTime utcNow = UtcNow;
-            foreach (KeyValuePair<object, CacheEntry> item in coherentState._entries)
+            foreach (CacheEntry entry in coherentState.GetAllValues())
             {
-                CacheEntry entry = item.Value;
                 if (entry.CheckExpired(utcNow))
                 {
                     entriesToRemove.Add(entry);
@@ -676,18 +669,71 @@ private static void ValidateCacheKey(object key)
         /// </summary>
         private sealed class CoherentState
         {
-            internal ConcurrentDictionary<object, CacheEntry> _entries = new ConcurrentDictionary<object, CacheEntry>();
+            private readonly ConcurrentDictionary<string, CacheEntry> _stringEntries = new ConcurrentDictionary<string, CacheEntry>(StringKeyComparer.Instance);
+            private readonly ConcurrentDictionary<object, CacheEntry> _nonStringEntries = new ConcurrentDictionary<object, CacheEntry>();
             internal long _cacheSize;
 
-            private ICollection<KeyValuePair<object, CacheEntry>> EntriesCollection => _entries;
+            internal bool TryGetValue(object key, [NotNullWhen(true)] out CacheEntry? entry)
+                => key is string s ? _stringEntries.TryGetValue(s, out entry) : _nonStringEntries.TryGetValue(key, out entry);
+
+            internal bool TryRemove(object key, [NotNullWhen(true)] out CacheEntry? entry)
+                => key is string s ? _stringEntries.TryRemove(s, out entry) : _nonStringEntries.TryRemove(key, out entry);
+
+            internal bool TryAdd(object key, CacheEntry entry)
+                => key is string s ? _stringEntries.TryAdd(s, entry) : _nonStringEntries.TryAdd(key, entry);
+
+            internal bool TryUpdate(object key, CacheEntry entry, CacheEntry comparison)
+                => key is string s ? _stringEntries.TryUpdate(s, entry, comparison) : _nonStringEntries.TryUpdate(key, entry, comparison);
+
+            public IEnumerable<CacheEntry> GetAllValues()
+            {
+                // note this mimics the outgoing code in that we don't just access
+                // .Values, which has additional overheads; this is only used for rare
+                // calls - compaction, clear, etc - so the additional overhead of a
+                // generated enumerator is not alarming
+                foreach (KeyValuePair<string, CacheEntry> entry in _stringEntries)
+                {
+                    yield return entry.Value;
+                }
+                foreach (KeyValuePair<object, CacheEntry> entry in _nonStringEntries)
+                {
+                    yield return entry.Value;
+                }
+            }
+
+            public IEnumerable<object> GetAllKeys()
+            {
+                foreach (KeyValuePair<string, CacheEntry> pairs in _stringEntries)
+                {
+                    yield return pairs.Key;
+                }
+                foreach (KeyValuePair<object, CacheEntry> pairs in _nonStringEntries)
+                {
+                    yield return pairs.Key;
+                }
+            }
+
+            private ICollection<KeyValuePair<string, CacheEntry>> StringEntriesCollection => _stringEntries;
+            private ICollection<KeyValuePair<object, CacheEntry>> NonStringEntriesCollection => _nonStringEntries;
 
-            internal int Count => _entries.Count;
+            internal int Count => _stringEntries.Count + _nonStringEntries.Count;
 
             internal long Size => Volatile.Read(ref _cacheSize);
 
             internal void RemoveEntry(CacheEntry entry, MemoryCacheOptions options)
             {
-                if (EntriesCollection.Remove(new KeyValuePair<object, CacheEntry>(entry.Key, entry)))
+                if (entry.Key is string s)
+                {
+                    if (StringEntriesCollection.Remove(new KeyValuePair<string, CacheEntry>(s, entry)))
+                    {
+                        if (options.SizeLimit.HasValue)
+                        {
+                            Interlocked.Add(ref _cacheSize, -entry.Size);
+                        }
+                        entry.InvokeEvictionCallbacks();
+                    }
+                }
+                else if (NonStringEntriesCollection.Remove(new KeyValuePair<object, CacheEntry>(entry.Key, entry)))
                 {
                     if (options.SizeLimit.HasValue)
                     {
@@ -696,6 +742,35 @@ internal void RemoveEntry(CacheEntry entry, MemoryCacheOptions options)
                     entry.InvokeEvictionCallbacks();
                 }
             }
+
+#if NETCOREAPP
+            // on .NET Core, the inbuilt comparer has Marvin built in; no need to intercept
+            private static class StringKeyComparer
+            {
+                internal static IEqualityComparer<string> Instance => EqualityComparer<string>.Default;
+            }
+#else
+            // otherwise, we need a custom comparer that manually implements Marvin
+            private sealed class StringKeyComparer : IEqualityComparer<string>, IEqualityComparer
+            {
+                private StringKeyComparer() { }
+
+                internal static readonly IEqualityComparer<string> Instance = new StringKeyComparer();
+
+                // special-case string keys and use Marvin hashing
+                public int GetHashCode(string? s) => s is null ? 0
+                    : Marvin.ComputeHash32(MemoryMarshal.AsBytes(s.AsSpan()), Marvin.DefaultSeed);
+
+                public bool Equals(string? x, string? y)
+                    => string.Equals(x, y);
+
+                bool IEqualityComparer.Equals(object x, object y)
+                    => object.Equals(x, y);
+
+                int IEqualityComparer.GetHashCode(object obj)
+                    => obj is string s ? GetHashCode(s) : 0;
+            }
+#endif
         }
     }
 }

src/libraries/Microsoft.Extensions.Caching.Memory/src/Microsoft.Extensions.Caching.Memory.csproj

@@ -5,6 +5,7 @@
     <EnableDefaultItems>true</EnableDefaultItems>
     <IsPackable>true</IsPackable>
     <PackageDescription>In-memory cache implementation of Microsoft.Extensions.Caching.Memory.IMemoryCache.</PackageDescription>
+    <AllowUnsafeBlocks>true</AllowUnsafeBlocks>
   </PropertyGroup>
 
   <ItemGroup>
@@ -21,4 +22,8 @@
     <PackageReference Include="System.ValueTuple" Version="$(SystemValueTupleVersion)" />
   </ItemGroup>
 
+  <ItemGroup Condition="!$([MSBuild]::IsTargetFrameworkCompatible('$(TargetFramework)', 'netstandard2.1'))">
+    <Compile Include="$(CoreLibSharedDir)System\Marvin.cs" Link="Common\System\Marvin.cs" />
+  </ItemGroup>
+
 </Project>

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

@@ -850,6 +850,21 @@ public async Task GetOrCreateAsyncWithCacheEntryOptions()
             Assert.False(cache.TryGetValue(cacheKey, out _));
         }
 
+        [Fact]
+        public void MixedKeysUsage()
+        {
+            // keys are split internally into 2 separate chunks
+            var cache = CreateCache();
+            var typed = Assert.IsType<MemoryCache>(cache);
+            object key0 = 123.45M, key1 = "123.45";
+            cache.Set(key0, "string value");
+            cache.Set(key1, "decimal value");
+
+            Assert.Equal(2, typed.Count);
+            Assert.Equal("string value", cache.Get(key0));
+            Assert.Equal("decimal value", cache.Get(key1));
+        }
+
         private class TestKey
         {
             public override bool Equals(object obj) => true;

src/libraries/System.Private.CoreLib/src/System/Marvin.cs

@@ -6,6 +6,12 @@
 using System.Runtime.CompilerServices;
 using System.Runtime.InteropServices;
 
+#if SYSTEM_PRIVATE_CORELIB
+using static System.Numerics.BitOperations;
+#else
+using System.Security.Cryptography;
+#endif
+
 namespace System
 {
     internal static partial class Marvin
@@ -204,7 +210,7 @@ public static int ComputeHash32(ref byte data, uint count, uint p0, uint p1)
                 else
                 {
                     partialResult |= (uint)Unsafe.ReadUnaligned<ushort>(ref data);
-                    partialResult = BitOperations.RotateLeft(partialResult, 16);
+                    partialResult = RotateLeft(partialResult, 16);
                 }
             }
 
@@ -221,16 +227,16 @@ private static void Block(ref uint rp0, ref uint rp1)
             uint p1 = rp1;
 
             p1 ^= p0;
-            p0 = BitOperations.RotateLeft(p0, 20);
+            p0 = RotateLeft(p0, 20);
 
             p0 += p1;
-            p1 = BitOperations.RotateLeft(p1, 9);
+            p1 = RotateLeft(p1, 9);
 
             p1 ^= p0;
-            p0 = BitOperations.RotateLeft(p0, 27);
+            p0 = RotateLeft(p0, 27);
 
             p0 += p1;
-            p1 = BitOperations.RotateLeft(p1, 19);
+            p1 = RotateLeft(p1, 19);
 
             rp0 = p0;
             rp1 = p1;
@@ -241,8 +247,29 @@ private static void Block(ref uint rp0, ref uint rp1)
         private static unsafe ulong GenerateSeed()
         {
             ulong seed;
+#if SYSTEM_PRIVATE_CORELIB
             Interop.GetRandomBytes((byte*)&seed, sizeof(ulong));
+#else
+            byte[] seedBytes = new byte[sizeof(ulong)];
+            using (RandomNumberGenerator rng = RandomNumberGenerator.Create())
+            {
+                rng.GetBytes(seedBytes);
+                fixed (byte* b = seedBytes)
+                {
+                    seed = *(ulong*)b;
+                }
+            }
+#endif
             return seed;
         }
+
+#if !SYSTEM_PRIVATE_CORELIB
+        [MethodImpl(MethodImplOptions.AggressiveInlining)]
+        private static uint RotateLeft(uint value, int shift)
+        {
+            // This is expected to be optimized into a single rol (or ror with negated shift value) instruction
+            return (value << shift) | (value >> (32 - shift));
+        }
+#endif
     }
 }